Use uv as packaging tool used in CI builds (apache#37692)

The `uv` tool released in Feb 2024 by ruff creators provides a
way faster drop-in replacement to `pip` and we are using it now
in our CI, when it can bring significant speed improvements (and
soon possibly more features).

This PR replaces `pip install` and `pip uninstall` with equivalent
`uv pip install` and `uv pip uninstall` commands, controlled by
a single `AIRFLOW_USE_UV` ARG. In CI images it is set to "true" so
CI images are prepared using UV, but PROD images (which are also
used during CI tests) are built using pip. This way we can get
both - stability and compliance for user-facing `pip` installation
and speed and new features coming from `uv`.

.github/workflows/ci.yml

@@ -175,7 +175,7 @@ jobs:
   # We only push CI cache as PROD cache usually does not gain as much from fresh cache because
   # it uses prepared airflow and provider packages that invalidate the cache anyway most of the time
   push-early-buildx-cache-to-github-registry:
-    timeout-minutes: 50
+    timeout-minutes: 110
     name: "Push Early Image Cache"
     runs-on: ${{fromJSON(needs.build-info.outputs.runs-on)}}
     needs:

dev/breeze/doc/ci/02_images.md

@@ -428,6 +428,7 @@ can be used for CI images:
 | `DEPENDENCIES_EPOCH_NUMBER`       | `2`                                                                     | increasing this number will reinstall all apt dependencies                                                                                                 |
 | `ADDITIONAL_PIP_INSTALL_FLAGS`    |                                                                         | additional `pip` flags passed to the installation commands (except when reinstalling `pip` itself)                                                         |
 | `PIP_NO_CACHE_DIR`                | `true`                                                                  | if true, then no pip cache will be stored                                                                                                                  |
+| `UV_NO_CACHE`                     | `true`                                                                  | if true, then no uv cache will be stored                                                                                                                   |
 | `HOME`                            | `/root`                                                                 | Home directory of the root user (CI image has root user as default)                                                                                        |
 | `AIRFLOW_HOME`                    | `/root/airflow`                                                         | Airflow's HOME (that's where logs and sqlite databases are stored)                                                                                         |
 | `AIRFLOW_SOURCES`                 | `/opt/airflow`                                                          | Mounted sources of Airflow                                                                                                                                 |
@@ -447,6 +448,8 @@ can be used for CI images:
 | `ADDITIONAL_DEV_APT_DEPS`         |                                                                         | Additional apt dev dependencies installed in the first part of the image                                                                                   |
 | `ADDITIONAL_DEV_APT_ENV`          |                                                                         | Additional env variables defined when installing dev deps                                                                                                  |
 | `AIRFLOW_PIP_VERSION`             | `24.0`                                                                  | PIP version used.                                                                                                                                          |
+| `AIRFLOW_UV_VERSION`              | `0.1.10`                                                                | UV version used.                                                                                                                                           |
+| `AIRFLOW_USE_UV`                  | `true`                                                                  | Whether to use UV for installation.                                                                                                                        |
 | `PIP_PROGRESS_BAR`                | `on`                                                                    | Progress bar for PIP installation                                                                                                                          |
 
 Here are some examples of how CI images can built manually. CI is always

dev/breeze/doc/images/output_ci-image_build.txt

@@ -1 +1 @@
-d2ef2733519d945c8cfd4fed63a43f24
+f535999147ac00393852eb3b28d7125b

dev/breeze/doc/images/output_prod-image_build.txt

@@ -1 +1 @@
-e3ed45363576b9bb3eb1cf5e3a2c08df
+07693b2597b00fdb156949c753dae783

dev/breeze/src/airflow_breeze/commands/ci_image_commands.py

@@ -252,6 +252,14 @@ def get_exitcode(status: int) -> int:
     "(see `breeze ci find-backtracking-candidates`).",
 )
 
+option_use_uv_ci = click.option(
+    "--use-uv/--no-use-uv",
+    is_flag=True,
+    default=True,
+    help="Use uv instead of pip as packaging tool.",
+    envvar="USE_UV",
+)
+
 option_upgrade_to_newer_dependencies = click.option(
     "-u",
     "--upgrade-to-newer-dependencies",
@@ -319,6 +327,7 @@ def get_exitcode(status: int) -> int:
 @option_tag_as_latest
 @option_upgrade_on_failure
 @option_upgrade_to_newer_dependencies
+@option_use_uv_ci
 @option_verbose
 @option_version_suffix_for_pypi_ci
 def build(
@@ -359,6 +368,7 @@ def build(
     tag_as_latest: bool,
     upgrade_on_failure: bool,
     upgrade_to_newer_dependencies: bool,
+    use_uv: bool,
     version_suffix_for_pypi: str,
 ):
     """Build CI image. Include building multiple images for all python versions."""
@@ -425,6 +435,7 @@ def run_build(ci_image_params: BuildCiParams) -> None:
         tag_as_latest=tag_as_latest,
         upgrade_on_failure=upgrade_on_failure,
         upgrade_to_newer_dependencies=upgrade_to_newer_dependencies,
+        use_uv=use_uv,
         version_suffix_for_pypi=version_suffix_for_pypi,
     )
     if platform:

dev/breeze/src/airflow_breeze/commands/ci_image_commands_config.py

@@ -59,6 +59,7 @@
                 "--debian-version",
                 "--install-mysql-client-type",
                 "--python-image",
+                "--use-uv",
             ],
         },
         {

dev/breeze/src/airflow_breeze/commands/production_image_commands_config.py

@@ -53,10 +53,10 @@
         {
             "name": "Advanced build options (for power users)",
             "options": [
+                "--additional-pip-install-flags",
+                "--commit-sha",
                 "--debian-version",
                 "--python-image",
-                "--commit-sha",
-                "--additional-pip-install-flags",
             ],
         },
         {

dev/breeze/src/airflow_breeze/commands/release_management_commands.py

@@ -214,6 +214,8 @@ class VersionedFile(NamedTuple):
 
 
 AIRFLOW_PIP_VERSION = "24.0"
+AIRFLOW_UV_VERSION = "0.1.10"
+AIRFLOW_USE_UV = False
 WHEEL_VERSION = "0.36.2"
 GITPYTHON_VERSION = "3.1.40"
 RICH_VERSION = "13.7.0"

dev/breeze/src/airflow_breeze/params/build_ci_params.py

@@ -41,6 +41,7 @@ class BuildCiParams(CommonBuildParams):
     eager_upgrade_additional_requirements: str | None = None
     skip_provider_dependencies_check: bool = False
     skip_image_upgrade_check: bool = False
+    use_uv: bool = True
     warn_image_upgrade_needed: bool = False
 
     @property
@@ -65,6 +66,7 @@ def prepare_arguments_for_docker_build_command(self) -> list[str]:
         self._req_arg("AIRFLOW_IMAGE_DATE_CREATED", self.airflow_image_date_created)
         self._req_arg("AIRFLOW_IMAGE_REPOSITORY", self.airflow_image_repository)
         self._req_arg("AIRFLOW_PRE_CACHED_PIP_PACKAGES", self.airflow_pre_cached_pip_packages)
+        self._req_arg("AIRFLOW_USE_UV", self.use_uv)
         self._req_arg("AIRFLOW_VERSION", self.airflow_version)
         self._req_arg("BUILD_ID", self.build_id)
         self._req_arg("CONSTRAINTS_GITHUB_REPOSITORY", self.constraints_github_repository)

dev/breeze/src/airflow_breeze/params/build_prod_params.py

@@ -51,10 +51,10 @@ class BuildProdParams(CommonBuildParams):
     install_airflow_reference: str | None = None
     install_airflow_version: str | None = None
     install_packages_from_context: bool = False
-    use_constraints_for_context_packages: bool = False
     installation_method: str = "."
     runtime_apt_command: str | None = None
     runtime_apt_deps: str | None = None
+    use_constraints_for_context_packages: bool = False
 
     @property
     def airflow_version(self) -> str:

docs/docker-stack/build-arg-ref.rst

@@ -47,6 +47,10 @@ Those are the most common arguments that you use when you want to build a custom
 +------------------------------------------+------------------------------------------+---------------------------------------------+
 | ``AIRFLOW_PIP_VERSION``                  | ``24.0``                                 |  PIP version used.                          |
 +------------------------------------------+------------------------------------------+---------------------------------------------+
+| ``AIRFLOW_UV_VERSION``                   | ``0.1.10``                               |  UV version used.                           |
++------------------------------------------+------------------------------------------+---------------------------------------------+
+| ``AIRFLOW_USE_UV``                       | ``false``                                |  Whether to use UV.                         |
++------------------------------------------+------------------------------------------+---------------------------------------------+
 | ``ADDITIONAL_PIP_INSTALL_FLAGS``         |                                          | additional ``pip`` flags passed to the      |
 |                                          |                                          | installation commands (except when          |
 |                                          |                                          | reinstalling ``pip`` itself)                |

scripts/docker/common.sh

@@ -18,6 +18,10 @@
 # shellcheck shell=bash
 set -euo pipefail
 
+: "${AIRFLOW_PIP_VERSION:?Should be set}"
+: "${AIRFLOW_UV_VERSION:?Should be set}"
+: "${AIRFLOW_USE_UV:?Should be set}"
+
 function common::get_colors() {
     COLOR_BLUE=$'\e[34m'
     COLOR_GREEN=$'\e[32m'
@@ -31,6 +35,40 @@ function common::get_colors() {
     export COLOR_YELLOW
 }
 
+function common::get_packaging_tool() {
+    ## IMPORTANT: IF YOU MODIFY THIS FUNCTION YOU SHOULD ALSO MODIFY CORRESPONDING FUNCTION IN
+    ## `scripts/in_container/_in_container_utils.sh`
+    if [[ ${AIRFLOW_USE_UV} == "true" ]]; then
+        echo
+        echo "${COLOR_BLUE}Using 'uv' to install Airflow${COLOR_RESET}"
+        echo
+        export PACKAGING_TOOL="uv"
+        export PACKAGING_TOOL_CMD="uv pip"
+        export EXTRA_INSTALL_FLAGS=""
+        export EXTRA_UNINSTALL_FLAGS=""
+        export RESOLUTION_HIGHEST_FLAG="--resolution highest"
+        export RESOLUTION_LOWEST_DIRECT_FLAG="--resolution lowest-direct"
+        # We need to lie about VIRTUAL_ENV to make uv works
+        # Until https://github.com/astral-sh/uv/issues/1396 is fixed
+        # In case we are running user installation, we need to set VIRTUAL_ENV to user's home + .local
+        if [[ ${PIP_USER=} == "true" ]]; then
+            VIRTUAL_ENV="${HOME}/.local"
+        else
+            VIRTUAL_ENV=$(python -c "import sys; print(sys.prefix)")
+        fi
+        export VIRTUAL_ENV
+    else
+        echo
+        echo "${COLOR_BLUE}Using 'pip' to install Airflow${COLOR_RESET}"
+        echo
+        export PACKAGING_TOOL="pip"
+        export PACKAGING_TOOL_CMD="pip"
+        export EXTRA_INSTALL_FLAGS="--root-user-action ignore"
+        export EXTRA_UNINSTALL_FLAGS="--yes"
+        export RESOLUTION_HIGHEST_FLAG="--upgrade-strategy eager"
+        export RESOLUTION_LOWEST_DIRECT_FLAG="--upgrade --upgrade-strategy only-if-needed"
+    fi
+}
 
 function common::get_airflow_version_specification() {
     if [[ -z ${AIRFLOW_VERSION_SPECIFICATION=}
@@ -66,20 +104,41 @@ function common::get_constraints_location() {
     fi
 }
 
-function common::show_pip_version_and_location() {
+function common::show_packaging_tool_version_and_location() {
    echo "PATH=${PATH}"
-   echo "pip on path: $(which pip)"
-   echo "Using pip: $(pip --version)"
+   if [[ ${PACKAGING_TOOL} == "pip" ]]; then
+       echo "${COLOR_BLUE}Using 'pip' to install Airflow${COLOR_RESET}"
+       echo "pip on path: $(which pip)"
+       echo "Using pip: $(pip --version)"
+   else
+       echo "${COLOR_BLUE}Using 'uv' to install Airflow${COLOR_RESET}"
+       echo "uv on path: $(which uv)"
+       echo "Using uv: $(uv --version)"
+   fi
 }
 
-function common::install_pip_version() {
+function common::install_packaging_tool() {
     echo
     echo "${COLOR_BLUE}Installing pip version ${AIRFLOW_PIP_VERSION}${COLOR_RESET}"
     echo
     if [[ ${AIRFLOW_PIP_VERSION} =~ .*https.* ]]; then
-        pip install --disable-pip-version-check "pip @ ${AIRFLOW_PIP_VERSION}"
+        # shellcheck disable=SC2086
+        pip install --root-user-action ignore --disable-pip-version-check "pip @ ${AIRFLOW_PIP_VERSION}"
     else
-        pip install --disable-pip-version-check "pip==${AIRFLOW_PIP_VERSION}"
+        # shellcheck disable=SC2086
+        pip install --root-user-action ignore --disable-pip-version-check "pip==${AIRFLOW_PIP_VERSION}"
+    fi
+    if [[ ${AIRFLOW_USE_UV} == "true" ]]; then
+        echo
+        echo "${COLOR_BLUE}Installing uv version ${AIRFLOW_UV_VERSION}${COLOR_RESET}"
+        echo
+        if [[ ${AIRFLOW_UV_VERSION} =~ .*https.* ]]; then
+            # shellcheck disable=SC2086
+            pip install --root-user-action ignore --disable-pip-version-check "uv @ ${AIRFLOW_UV_VERSION}"
+        else
+            # shellcheck disable=SC2086
+            pip install --root-user-action ignore --disable-pip-version-check "uv==${AIRFLOW_UV_VERSION}"
+        fi
     fi
     mkdir -p "${HOME}/.local/bin"
 }

scripts/docker/entrypoint_ci.sh

@@ -221,8 +221,12 @@ function check_boto_upgrade() {
     echo
     echo "${COLOR_BLUE}Upgrading boto3, botocore to latest version to run Amazon tests with them${COLOR_RESET}"
     echo
-    pip uninstall --root-user-action ignore aiobotocore s3fs -y || true
-    pip install --root-user-action ignore --upgrade boto3 botocore
+    # shellcheck disable=SC2086
+    ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} aiobotocore s3fs || true
+    # We need to include oss2 as dependency as otherwise jmespath will be bumped and it will not pass
+    # the pip check test, Similarly gcloud-aio-auth limit is needed to be included as it bumps cryptography
+    # shellcheck disable=SC2086
+    ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade boto3 botocore "oss2>=2.14.0" "gcloud-aio-auth>=4.0.0,<5.0.0"
     pip check
 }
 
@@ -232,25 +236,31 @@ function check_pydantic() {
         echo
         echo "${COLOR_YELLOW}Reinstalling airflow from local sources to account for pyproject.toml changes${COLOR_RESET}"
         echo
-        pip install --root-user-action ignore -e .
+        # shellcheck disable=SC2086
+        ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} -e .
         echo
         echo "${COLOR_YELLOW}Remove pydantic and 3rd party libraries that depend on it${COLOR_RESET}"
         echo
-        pip uninstall --root-user-action ignore pydantic aws-sam-translator openai pyiceberg qdrant-client cfn-lint -y
+        # shellcheck disable=SC2086
+        ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} pydantic aws-sam-translator openai \
+           pyiceberg qdrant-client cfn-lint weaviate-client
         pip check
     elif [[ ${PYDANTIC=} == "v1" ]]; then
         echo
         echo "${COLOR_YELLOW}Reinstalling airflow from local sources to account for pyproject.toml changes${COLOR_RESET}"
         echo
-        pip install --root-user-action ignore -e .
+        # shellcheck disable=SC2086
+        ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} -e .
         echo
-        echo "${COLOR_YELLOW}Uninstalling pyicberg which is not compatible with Pydantic 1${COLOR_RESET}"
+        echo "${COLOR_YELLOW}Uninstalling dependencies which are not compatible with Pydantic 1${COLOR_RESET}"
         echo
-        pip uninstall pyiceberg -y
+        # shellcheck disable=SC2086
+        ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} pyiceberg waeviate-client
         echo
         echo "${COLOR_YELLOW}Downgrading Pydantic to < 2${COLOR_RESET}"
         echo
-        pip install --upgrade "pydantic<2.0.0"
+        # shellcheck disable=SC2086
+        ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade "pydantic<2.0.0"
         pip check
     else
         echo
@@ -269,7 +279,8 @@ function check_download_sqlalchemy() {
     echo
     echo "${COLOR_BLUE}Downgrading sqlalchemy to minimum supported version: ${min_sqlalchemy_version}${COLOR_RESET}"
     echo
-    pip install --root-user-action ignore "sqlalchemy==${min_sqlalchemy_version}"
+    # shellcheck disable=SC2086
+    ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} "sqlalchemy==${min_sqlalchemy_version}"
     pip check
 }
 
@@ -282,7 +293,8 @@ function check_download_pendulum() {
     echo
     echo "${COLOR_BLUE}Downgrading pendulum to minimum supported version: ${min_pendulum_version}${COLOR_RESET}"
     echo
-    pip install --root-user-action ignore "pendulum==${min_pendulum_version}"
+    # shellcheck disable=SC2086
+    ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} "pendulum==${min_pendulum_version}"
     pip check
 }
 

scripts/docker/install_additional_dependencies.sh

@@ -20,7 +20,6 @@ set -euo pipefail
 
 : "${UPGRADE_TO_NEWER_DEPENDENCIES:?Should be true or false}"
 : "${ADDITIONAL_PYTHON_DEPS:?Should be set}"
-: "${AIRFLOW_PIP_VERSION:?Should be set}"
 
 # shellcheck source=scripts/docker/common.sh
 . "$( dirname "${BASH_SOURCE[0]}" )/common.sh"
@@ -32,10 +31,10 @@ function install_additional_dependencies() {
         echo "${COLOR_BLUE}Installing additional dependencies while upgrading to newer dependencies${COLOR_RESET}"
         echo
         set -x
-        pip install --root-user-action ignore --upgrade --upgrade-strategy eager \
+        ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade ${RESOLUTION_HIGHEST_FLAG} \
             ${ADDITIONAL_PIP_INSTALL_FLAGS} \
             ${ADDITIONAL_PYTHON_DEPS} ${EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS=}
-        common::install_pip_version
+        common::install_packaging_tool
         set +x
         echo
         echo "${COLOR_BLUE}Running 'pip check'${COLOR_RESET}"
@@ -46,10 +45,10 @@ function install_additional_dependencies() {
         echo "${COLOR_BLUE}Installing additional dependencies upgrading only if needed${COLOR_RESET}"
         echo
         set -x
-        pip install --root-user-action ignore --upgrade --upgrade-strategy only-if-needed \
+        ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade "${RESOLUTION_LOWEST_DIRECT_FLAG}" \
             ${ADDITIONAL_PIP_INSTALL_FLAGS} \
             ${ADDITIONAL_PYTHON_DEPS}
-        common::install_pip_version
+        common::install_packaging_tool
         set +x
         echo
         echo "${COLOR_BLUE}Running 'pip check'${COLOR_RESET}"
@@ -59,9 +58,10 @@ function install_additional_dependencies() {
 }
 
 common::get_colors
+common::get_packaging_tool
 common::get_airflow_version_specification
 common::override_pip_version_if_needed
 common::get_constraints_location
-common::show_pip_version_and_location
+common::show_packaging_tool_version_and_location
 
 install_additional_dependencies