Allow to switch breeze to use uv internally to create virtualenvs (#43587) (#43624)

potiuk · utkarsharma2 · commit f593d8b07386 · 2024-12-04T14:51:18.000+05:30
Breeze sometimes creates "internal" virtualenvs in local ".build" directory when it needs - for example in order to run k8s tests or for release management commands. This PR adds capability to switch breeze to use `uv` instead of `pip` to install depdendencies in those envs. You can now switch breeze to use uv by `breeze setup config --use-uv` and switch back to pip by `breeze setup config --no-use-uv`. (cherry picked from commit a2a0ef0)
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -622,6 +622,7 @@ jobs:
       kubernetes-versions-list-as-string: ${{ needs.build-info.outputs.kubernetes-versions-list-as-string }}
       kubernetes-combos-list-as-string: ${{ needs.build-info.outputs.kubernetes-combos-list-as-string }}
       include-success-outputs: ${{ needs.build-info.outputs.include-success-outputs }}
+      use-uv: ${{ needs.build-info.outputs.force-pip && 'false' || 'true' }}
       debug-resources: ${{ needs.build-info.outputs.debug-resources }}
     if: >
       ( needs.build-info.outputs.run-kubernetes-tests == 'true' ||
diff --git a/.github/workflows/k8s-tests.yml b/.github/workflows/k8s-tests.yml
@@ -44,6 +44,10 @@ on:  # yamllint disable-line rule:truthy
         description: "Whether to include success outputs"
         required: true
         type: string
+      use-uv:
+        description: "Whether to use uv"
+        required: true
+        type: string
       debug-resources:
         description: "Whether to debug resources"
         required: true
@@ -96,6 +100,9 @@ jobs:
           key: "\
             k8s-env-${{ steps.breeze.outputs.host-python-version }}-\
             ${{ hashFiles('scripts/ci/kubernetes/k8s_requirements.txt','hatch_build.py') }}"
+      - name: "Switch breeze to use uv"
+        run: breeze setup-config --use-uv
+        if: inputs.use-uv == 'true'
       - name: Run complete K8S tests ${{ inputs.kubernetes-combos-list-as-string }}
         run: breeze k8s run-complete-tests --run-in-parallel --upgrade --no-copy-local-sources
         env:
diff --git a/Dockerfile b/Dockerfile
@@ -49,8 +49,8 @@ ARG AIRFLOW_VERSION="2.9.3"
 
 ARG PYTHON_BASE_IMAGE="python:3.8-slim-bookworm"
 
-ARG AIRFLOW_PIP_VERSION=24.2
-ARG AIRFLOW_UV_VERSION=0.4.1
+ARG AIRFLOW_PIP_VERSION=24.3.1
+ARG AIRFLOW_UV_VERSION=0.4.29
 ARG AIRFLOW_USE_UV="false"
 ARG UV_HTTP_TIMEOUT="300"
 ARG AIRFLOW_IMAGE_REPOSITORY="https://github.com/apache/airflow"
diff --git a/Dockerfile.ci b/Dockerfile.ci
@@ -1297,8 +1297,8 @@ ARG DEFAULT_CONSTRAINTS_BRANCH="constraints-main"
 # It can also be overwritten manually by setting the AIRFLOW_CI_BUILD_EPOCH environment variable.
 ARG AIRFLOW_CI_BUILD_EPOCH="10"
 ARG AIRFLOW_PRE_CACHED_PIP_PACKAGES="true"
-ARG AIRFLOW_PIP_VERSION=24.2
-ARG AIRFLOW_UV_VERSION=0.4.1
+ARG AIRFLOW_PIP_VERSION=24.3.1
+ARG AIRFLOW_UV_VERSION=0.4.29
 ARG AIRFLOW_USE_UV="true"
 # Setup PIP
 # By default PIP install run without cache to make image smaller
@@ -1321,8 +1321,8 @@ ARG AIRFLOW_VERSION=""
 # Additional PIP flags passed to all pip install commands except reinstalling pip itself
 ARG ADDITIONAL_PIP_INSTALL_FLAGS=""
 
-ARG AIRFLOW_PIP_VERSION=24.2
-ARG AIRFLOW_UV_VERSION=0.4.1
+ARG AIRFLOW_PIP_VERSION=24.3.1
+ARG AIRFLOW_UV_VERSION=0.4.29
 ARG AIRFLOW_USE_UV="true"
 
 ENV AIRFLOW_REPO=${AIRFLOW_REPO}\
diff --git a/dev/breeze/doc/ci/02_images.md b/dev/breeze/doc/ci/02_images.md
@@ -447,8 +447,8 @@ can be used for CI images:
 | `DEV_APT_DEPS`                    |                            | Dev APT dependencies installed in the first part of the image                                                                                              |
 | `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_PIP_VERSION`             | `24.3.1`                   | PIP version used.                                                                                                                                          |
+| `AIRFLOW_UV_VERSION`              | `0.4.29`                   | UV version used.                                                                                                                                           |
 | `AIRFLOW_USE_UV`                  | `true`                     | Whether to use UV for installation.                                                                                                                        |
 | `PIP_PROGRESS_BAR`                | `on`                       | Progress bar for PIP installation                                                                                                                          |
 
diff --git a/dev/breeze/doc/images/output-commands.svg b/dev/breeze/doc/images/output-commands.svg
diff --git a/dev/breeze/doc/images/output_setup_config.svg b/dev/breeze/doc/images/output_setup_config.svg
diff --git a/dev/breeze/doc/images/output_setup_config.txt b/dev/breeze/doc/images/output_setup_config.txt
@@ -1 +1 @@
-422c8c524b557fcf5924da4c8590935d
+96e10564034b282769a2c48ebf7176e2
diff --git a/dev/breeze/src/airflow_breeze/commands/release_management_commands.py b/dev/breeze/src/airflow_breeze/commands/release_management_commands.py
@@ -226,8 +226,8 @@ class VersionedFile(NamedTuple):
     file_name: str
 
 
-AIRFLOW_PIP_VERSION = "24.0"
-AIRFLOW_UV_VERSION = "0.1.10"
+AIRFLOW_PIP_VERSION = "24.3.1"
+AIRFLOW_UV_VERSION = "0.4.29"
 AIRFLOW_USE_UV = False
 WHEEL_VERSION = "0.36.2"
 GITPYTHON_VERSION = "3.1.40"
@@ -451,7 +451,11 @@ def _check_sdist_to_wheel_dists(dists_info: tuple[DistributionPackageInfo, ...])
                 continue
 
             if not venv_created:
-                python_path = create_venv(Path(tmp_dir_name) / ".venv", pip_version=AIRFLOW_PIP_VERSION)
+                python_path = create_venv(
+                    Path(tmp_dir_name) / ".venv",
+                    pip_version=AIRFLOW_PIP_VERSION,
+                    uv_version=AIRFLOW_UV_VERSION,
+                )
                 pip_command = create_pip_command(python_path)
                 venv_created = True
 
diff --git a/dev/breeze/src/airflow_breeze/commands/setup_commands.py b/dev/breeze/src/airflow_breeze/commands/setup_commands.py
@@ -192,6 +192,12 @@ def version():
 @option_mysql_version
 @click.option("-C/-c", "--cheatsheet/--no-cheatsheet", help="Enable/disable cheatsheet.", default=None)
 @click.option("-A/-a", "--asciiart/--no-asciiart", help="Enable/disable ASCIIart.", default=None)
+@click.option(
+    "-U/-u",
+    "--use-uv/--no-use-uv",
+    help="Enable/disable using uv for creating venvs by breeze.",
+    default=None,
+)
 @click.option(
     "--colour/--no-colour",
     help="Enable/disable Colour mode (useful for colour blind-friendly communication).",
@@ -200,6 +206,7 @@ def version():
 def change_config(
     python: str,
     backend: str,
+    use_uv: bool,
     postgres_version: str,
     mysql_version: str,
     cheatsheet: bool,
@@ -212,14 +219,22 @@ def change_config(
     asciiart_file = "suppress_asciiart"
     cheatsheet_file = "suppress_cheatsheet"
     colour_file = "suppress_colour"
+    use_uv_file = "use_uv"
 
+    if use_uv is not None:
+        if use_uv:
+            touch_cache_file(use_uv_file)
+            get_console().print("[info]Enable using uv[/]")
+        else:
+            delete_cache(use_uv_file)
+            get_console().print("[info]Disable using uv[/]")
     if asciiart is not None:
         if asciiart:
             delete_cache(asciiart_file)
-            get_console().print("[info]Enable ASCIIART![/]")
+            get_console().print("[info]Enable ASCIIART[/]")
         else:
             touch_cache_file(asciiart_file)
-            get_console().print("[info]Disable ASCIIART![/]")
+            get_console().print("[info]Disable ASCIIART[/]")
     if cheatsheet is not None:
         if cheatsheet:
             delete_cache(cheatsheet_file)
@@ -235,23 +250,27 @@ def change_config(
             touch_cache_file(colour_file)
             get_console().print("[info]Disable Colour[/]")
 
-    def get_status(file: str):
+    def get_supress_status(file: str):
         return "disabled" if check_if_cache_exists(file) else "enabled"
 
+    def get_status(file: str):
+        return "enabled" if check_if_cache_exists(file) else "disabled"
+
     get_console().print()
     get_console().print("[info]Current configuration:[/]")
     get_console().print()
     get_console().print(f"[info]* Python: {python}[/]")
     get_console().print(f"[info]* Backend: {backend}[/]")
+    get_console().print(f"[info]* Use uv: {get_status(use_uv_file)}[/]")
     get_console().print()
     get_console().print(f"[info]* Postgres version: {postgres_version}[/]")
     get_console().print(f"[info]* MySQL version: {mysql_version}[/]")
     get_console().print()
-    get_console().print(f"[info]* ASCIIART: {get_status(asciiart_file)}[/]")
-    get_console().print(f"[info]* Cheatsheet: {get_status(cheatsheet_file)}[/]")
+    get_console().print(f"[info]* ASCIIART: {get_supress_status(asciiart_file)}[/]")
+    get_console().print(f"[info]* Cheatsheet: {get_supress_status(cheatsheet_file)}[/]")
     get_console().print()
     get_console().print()
-    get_console().print(f"[info]* Colour: {get_status(colour_file)}[/]")
+    get_console().print(f"[info]* Colour: {get_supress_status(colour_file)}[/]")
     get_console().print()
 
 
diff --git a/dev/breeze/src/airflow_breeze/commands/setup_commands_config.py b/dev/breeze/src/airflow_breeze/commands/setup_commands_config.py
@@ -63,6 +63,7 @@
                 "--backend",
                 "--postgres-version",
                 "--mysql-version",
+                "--use-uv",
                 "--cheatsheet",
                 "--asciiart",
                 "--colour",
diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py
@@ -155,7 +155,8 @@
 
 ALLOWED_INSTALL_MYSQL_CLIENT_TYPES = ["mariadb", "mysql"]
 
-PIP_VERSION = "24.0"
+PIP_VERSION = "24.3.1"
+UV_VERSION = "0.4.29"
 
 DEFAULT_UV_HTTP_TIMEOUT = 300
 DEFAULT_WSL2_HTTP_TIMEOUT = 900
diff --git a/dev/breeze/src/airflow_breeze/utils/kubernetes_utils.py b/dev/breeze/src/airflow_breeze/utils/kubernetes_utils.py
@@ -41,12 +41,15 @@
     HELM_VERSION,
     KIND_VERSION,
     PIP_VERSION,
+    UV_VERSION,
 )
+from airflow_breeze.utils.cache import check_if_cache_exists
 from airflow_breeze.utils.console import Output, get_console
 from airflow_breeze.utils.host_info_utils import Architecture, get_host_architecture, get_host_os
 from airflow_breeze.utils.path_utils import AIRFLOW_SOURCES_ROOT, BUILD_CACHE_DIR
 from airflow_breeze.utils.run_utils import RunCommandResult, run_command
 from airflow_breeze.utils.shared_options import get_dry_run, get_verbose
+from airflow_breeze.utils.virtualenv_utils import create_pip_command, create_uv_command
 
 K8S_ENV_PATH = BUILD_CACHE_DIR / ".k8s-env"
 K8S_CLUSTERS_PATH = BUILD_CACHE_DIR / ".k8s-clusters"
@@ -301,10 +304,12 @@ def _requirements_changed() -> bool:
 
 
 def _install_packages_in_k8s_virtualenv():
+    if check_if_cache_exists("use_uv"):
+        command = create_uv_command(PYTHON_BIN_PATH)
+    else:
+        command = create_pip_command(PYTHON_BIN_PATH)
     install_command_no_constraints = [
-        str(PYTHON_BIN_PATH),
-        "-m",
-        "pip",
+        *command,
         "install",
         "-r",
         str(K8S_REQUIREMENTS_PATH.resolve()),
@@ -405,8 +410,9 @@ def create_virtualenv(force_venv_setup: bool) -> RunCommandResult:
         )
         return venv_command_result
     get_console().print(f"[info]Reinstalling PIP version in {K8S_ENV_PATH}")
+    command = create_pip_command(PYTHON_BIN_PATH)
     pip_reinstall_result = run_command(
-        [str(PYTHON_BIN_PATH), "-m", "pip", "install", f"pip=={PIP_VERSION}"],
+        [*command, "install", f"pip=={PIP_VERSION}"],
         check=False,
         capture_output=True,
     )
@@ -416,8 +422,19 @@ def create_virtualenv(force_venv_setup: bool) -> RunCommandResult:
             f"{pip_reinstall_result.stdout}\n{pip_reinstall_result.stderr}"
         )
         return pip_reinstall_result
-    get_console().print(f"[info]Installing necessary packages in {K8S_ENV_PATH}")
+    uv_reinstall_result = run_command(
+        [*command, "install", f"uv=={UV_VERSION}"],
+        check=False,
+        capture_output=True,
+    )
+    if uv_reinstall_result.returncode != 0:
+        get_console().print(
+            f"[error]Error when updating uv to {UV_VERSION}:[/]\n"
+            f"{uv_reinstall_result.stdout}\n{uv_reinstall_result.stderr}"
+        )
+        return uv_reinstall_result
 
+    get_console().print(f"[info]Installing necessary packages in {K8S_ENV_PATH}")
     install_packages_result = _install_packages_in_k8s_virtualenv()
     if install_packages_result.returncode == 0:
         if get_dry_run():
diff --git a/dev/breeze/src/airflow_breeze/utils/run_tests.py b/dev/breeze/src/airflow_breeze/utils/run_tests.py
@@ -22,7 +22,7 @@
 from itertools import chain
 from subprocess import DEVNULL
 
-from airflow_breeze.global_constants import PIP_VERSION
+from airflow_breeze.global_constants import PIP_VERSION, UV_VERSION
 from airflow_breeze.utils.console import Output, get_console
 from airflow_breeze.utils.packages import get_excluded_provider_folders, get_suspended_provider_folders
 from airflow_breeze.utils.path_utils import AIRFLOW_SOURCES_ROOT
@@ -59,7 +59,9 @@ def verify_an_image(
     env["DOCKER_IMAGE"] = image_name
     if slim_image:
         env["TEST_SLIM_IMAGE"] = "true"
-    with create_temp_venv(pip_version=PIP_VERSION, requirements_file=DOCKER_TESTS_REQUIREMENTS) as py_exe:
+    with create_temp_venv(
+        pip_version=PIP_VERSION, uv_version=UV_VERSION, requirements_file=DOCKER_TESTS_REQUIREMENTS
+    ) as py_exe:
         command_result = run_command(
             [py_exe, "-m", "pytest", str(test_path), *pytest_args, *extra_pytest_args],
             env=env,
diff --git a/dev/breeze/src/airflow_breeze/utils/virtualenv_utils.py b/dev/breeze/src/airflow_breeze/utils/virtualenv_utils.py
@@ -23,6 +23,7 @@
 from pathlib import Path
 from typing import Generator
 
+from airflow_breeze.utils.cache import check_if_cache_exists
 from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.run_utils import run_command
 
@@ -31,10 +32,15 @@ def create_pip_command(python: str | Path) -> list[str]:
     return [python.as_posix() if hasattr(python, "as_posix") else str(python), "-m", "pip"]
 
 
+def create_uv_command(python: str | Path) -> list[str]:
+    return [python.as_posix() if hasattr(python, "as_posix") else str(python), "-m", "uv", "pip"]
+
+
 def create_venv(
     venv_path: str | Path,
     python: str | None = None,
     pip_version: str | None = None,
+    uv_version: str | None = None,
     requirements_file: str | Path | None = None,
 ) -> str:
     venv_path = Path(venv_path).resolve().absolute()
@@ -53,10 +59,13 @@ def create_venv(
     if not python_path.exists():
         get_console().print(f"\n[errors]Python interpreter is not exist in path {python_path}. Exiting!\n")
         sys.exit(1)
-    pip_command = create_pip_command(python_path)
+    if check_if_cache_exists("use_uv"):
+        command = create_uv_command(python_path)
+    else:
+        command = create_pip_command(python_path)
     if pip_version:
         result = run_command(
-            [*pip_command, "install", f"pip=={pip_version}", "-q"],
+            [*command, "install", f"pip=={pip_version}", "-q"],
             check=False,
             capture_output=False,
             text=True,
@@ -67,10 +76,23 @@ def create_venv(
                 f"{result.stdout}\n{result.stderr}"
             )
             sys.exit(result.returncode)
+    if uv_version:
+        result = run_command(
+            [*command, "install", f"uv=={uv_version}", "-q"],
+            check=False,
+            capture_output=False,
+            text=True,
+        )
+        if result.returncode != 0:
+            get_console().print(
+                f"[error]Error when installing uv in {venv_path.as_posix()}[/]\n"
+                f"{result.stdout}\n{result.stderr}"
+            )
+            sys.exit(result.returncode)
     if requirements_file:
         requirements_file = Path(requirements_file).absolute().as_posix()
         result = run_command(
-            [*pip_command, "install", "-r", requirements_file, "-q"],
+            [*command, "install", "-r", requirements_file, "-q"],
             check=True,
             capture_output=False,
             text=True,
@@ -88,6 +110,7 @@ def create_venv(
 def create_temp_venv(
     python: str | None = None,
     pip_version: str | None = None,
+    uv_version: str | None = None,
     requirements_file: str | Path | None = None,
     prefix: str | None = None,
 ) -> Generator[str, None, None]:
@@ -96,5 +119,6 @@ def create_temp_venv(
             Path(tmp_dir_name) / ".venv",
             python=python,
             pip_version=pip_version,
+            uv_version=uv_version,
             requirements_file=requirements_file,
         )
diff --git a/scripts/ci/install_breeze.sh b/scripts/ci/install_breeze.sh
@@ -25,7 +25,7 @@ if [[ ${PYTHON_VERSION=} != "" ]]; then
     PYTHON_ARG="--python=$(which python"${PYTHON_VERSION}") "
 fi
 
-python -m pip install --upgrade pip==24.0
+python -m pip install --upgrade pip==24.3.1
 python -m pip install "pipx>=1.4.1"
 python -m pipx uninstall apache-airflow-breeze >/dev/null 2>&1 || true
 # shellcheck disable=SC2086
diff --git a/scripts/ci/pre_commit/update_installers.py b/scripts/ci/pre_commit/update_installers.py
@@ -65,6 +65,7 @@ def get_latest_pypi_version(package_name: str) -> str:
 
 AIRFLOW_UV_PATTERN = re.compile(r"(AIRFLOW_UV_VERSION=)([0-9.]+)")
 AIRFLOW_UV_QUOTED_PATTERN = re.compile(r"(AIRFLOW_UV_VERSION = )(\"[0-9.]+\")")
+UV_QUOTED_PATTERN = re.compile(r"(UV_VERSION = )(\"[0-9.]+\")")
 AIRFLOW_UV_DOC_PATTERN = re.compile(r"(\| *`AIRFLOW_UV_VERSION` *\| *)(`[0-9.]+`)( *\|)")
 UV_GREATER_PATTERN = re.compile(r'"(uv>=)([0-9]+)"')
 
@@ -118,11 +119,14 @@ def replacer(match):
             new_content = replace_group_2_while_keeping_total_length(
                 AIRFLOW_UV_PATTERN, uv_version, new_content
             )
+            new_content = replace_group_2_while_keeping_total_length(
+                UV_GREATER_PATTERN, uv_version, new_content
+            )
             new_content = replace_group_2_while_keeping_total_length(
                 AIRFLOW_UV_QUOTED_PATTERN, f'"{uv_version}"', new_content
             )
             new_content = replace_group_2_while_keeping_total_length(
-                UV_GREATER_PATTERN, uv_version, new_content
+                UV_QUOTED_PATTERN, f'"{uv_version}"', new_content
             )
         if new_content != file_content:
             file.write_text(new_content)

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-422c8c524b557fcf5924da4c8590935d`
	`1`	`+96e10564034b282769a2c48ebf7176e2`