From c3d246b96ef615bd13c257b8565d6e064da9c82d Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Mon, 4 Nov 2024 09:59:29 +0100 Subject: [PATCH] Allow to switch breeze to use uv internally to create virtualenvs (#43587) (#43624) 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 a2a0ef09357f278bde031092e395f13286fd3076) --- .github/workflows/ci.yml | 1 + .github/workflows/k8s-tests.yml | 7 ++++ Dockerfile | 4 +- Dockerfile.ci | 8 ++-- dev/breeze/doc/ci/02_images.md | 4 +- dev/breeze/doc/images/output-commands.svg | 42 +++++++++---------- dev/breeze/doc/images/output_setup_config.svg | 32 +++++++------- dev/breeze/doc/images/output_setup_config.txt | 2 +- .../commands/release_management_commands.py | 10 +++-- .../airflow_breeze/commands/setup_commands.py | 31 +++++++++++--- .../commands/setup_commands_config.py | 1 + .../src/airflow_breeze/global_constants.py | 3 +- .../airflow_breeze/utils/kubernetes_utils.py | 27 +++++++++--- .../src/airflow_breeze/utils/run_tests.py | 6 ++- .../airflow_breeze/utils/virtualenv_utils.py | 30 +++++++++++-- scripts/ci/install_breeze.sh | 2 +- scripts/ci/pre_commit/update_installers.py | 6 ++- 17 files changed, 150 insertions(+), 66 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 866f8f253d401..c94518489d28a 100644 --- 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 index c4b72a9afc924..9a764e88c4e99 100644 --- 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 index cf5226c00086f..4cdf1a8bb3409 100644 --- 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 index d23e810fa3677..e188a7ec39115 100644 --- 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 index f9ea4faaee7c0..1db263f8b3aa0 100644 --- 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 index 08d3dc2a13eea..5888d1fc862eb 100644 --- a/dev/breeze/doc/images/output-commands.svg +++ b/dev/breeze/doc/images/output-commands.svg @@ -298,53 +298,53 @@ Usage:breeze[OPTIONS] COMMAND [ARGS]... ╭─ Execution mode ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--python-pPython major/minor version used in Airflow image for images.│ +│--python-pPython major/minor version used in Airflow image for images.│ │(>3.8< | 3.9 | 3.10 | 3.11 | 3.12)                          │ │[default: 3.8]                                              │ -│--integrationIntegration(s) to enable when running (can be more than one).                       │ +│--integrationIntegration(s) to enable when running (can be more than one).                       │ │(all | all-testable | cassandra | celery | drill | kafka | kerberos | mongo | mssql │ │| openlineage | otel | pinot | qdrant | redis | statsd | trino | ydb)               │ -│--standalone-dag-processorRun standalone dag processor for start-airflow.│ -│--database-isolationRun airflow in database isolation mode.│ +│--standalone-dag-processorRun standalone dag processor for start-airflow.│ +│--database-isolationRun airflow in database isolation mode.│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Docker Compose selection and cleanup ───────────────────────────────────────────────────────────────────────────────╮ -│--project-nameName of the docker-compose project to bring down. The `docker-compose` is for legacy breeze       │ -│project name and you can use `breeze down --project-name docker-compose` to stop all containers   │ +│--project-nameName of the docker-compose project to bring down. The `docker-compose` is for legacy breeze       │ +│project name and you can use `breeze down --project-name docker-compose` to stop all containers   │ │belonging to it.                                                                                  │ │(breeze | pre-commit | docker-compose)                                                            │ │[default: breeze]                                                                                 │ -│--docker-hostOptional - docker host to use when running docker commands. When set, the `--builder` option is   │ +│--docker-hostOptional - docker host to use when running docker commands. When set, the `--builder` option is   │ │ignored when building images.                                                                     │ │(TEXT)                                                                                            │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Database ───────────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--backend-bDatabase backend to use. If 'none' is chosen, Breeze will start with an invalid database    │ +│--backend-bDatabase backend to use. If 'none' is chosen, Breeze will start with an invalid database    │ │configuration, meaning there will be no database available, and any attempts to connect to  │ │the Airflow database will fail.                                                             │ │(>sqlite< | mysql | postgres | none)                                                        │ │[default: sqlite]                                                                           │ -│--postgres-version-PVersion of Postgres used.(>12< | 13 | 14 | 15 | 16)[default: 12]│ -│--mysql-version-MVersion of MySQL used.(>8.0< | 8.4)[default: 8.0]│ -│--db-reset-dReset DB when entering the container.│ +│--postgres-version-PVersion of Postgres used.(>12< | 13 | 14 | 15 | 16)[default: 12]│ +│--mysql-version-MVersion of MySQL used.(>8.0< | 8.4)[default: 8.0]│ +│--db-reset-dReset DB when entering the container.│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Build CI image (before entering shell) ─────────────────────────────────────────────────────────────────────────────╮ -│--github-repository-gGitHub repository used to pull, push run images.(TEXT)[default: apache/airflow]│ -│--builderBuildx builder used to perform `docker buildx build` commands.(TEXT)│ +│--github-repository-gGitHub repository used to pull, push run images.(TEXT)[default: apache/airflow]│ +│--builderBuildx builder used to perform `docker buildx build` commands.(TEXT)│ │[default: autodetect]                                         │ -│--use-uv/--no-use-uvUse uv instead of pip as packaging tool to build the image.[default: use-uv]│ -│--uv-http-timeoutTimeout for requests that UV makes (only used in case of UV builds).(INTEGER RANGE)│ +│--use-uv/--no-use-uvUse uv instead of pip as packaging tool to build the image.[default: use-uv]│ +│--uv-http-timeoutTimeout for requests that UV makes (only used in case of UV builds).(INTEGER RANGE)│ │[default: 300; x>=1]                                                │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Other options ──────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--forward-credentials-fForward local credentials to container when running.│ -│--max-timeMaximum time that the command should take - if it takes longer, the command will fail.│ +│--forward-credentials-fForward local credentials to container when running.│ +│--max-timeMaximum time that the command should take - if it takes longer, the command will fail.│ │(INTEGER RANGE)                                                                       │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--answer-aForce answer to questions.(y | n | q | yes | no | quit)│ -│--dry-run-DIf dry-run is set, commands are only printed, not executed.│ -│--verbose-vPrint verbose information about performed steps.│ -│--help-hShow this message and exit.│ +│--answer-aForce answer to questions.(y | n | q | yes | no | quit)│ +│--dry-run-DIf dry-run is set, commands are only printed, not executed.│ +│--verbose-vPrint verbose information about performed steps.│ +│--help-hShow this message and exit.│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Developer commands ─────────────────────────────────────────────────────────────────────────────────────────────────╮ │start-airflow          Enter breeze environment and starts all Airflow components in the tmux session. Compile    │ diff --git a/dev/breeze/doc/images/output_setup_config.svg b/dev/breeze/doc/images/output_setup_config.svg index 9a42467ea5281..5a44bb20030b9 100644 --- a/dev/breeze/doc/images/output_setup_config.svg +++ b/dev/breeze/doc/images/output_setup_config.svg @@ -1,4 +1,4 @@ - + Show/update configuration (Python, Backend, Cheatsheet, ASCIIART). ╭─ Config flags ───────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--python-pPython major/minor version used in Airflow image for images.│ +│--python-pPython major/minor version used in Airflow image for images.│ │(>3.8< | 3.9 | 3.10 | 3.11 | 3.12)                          │ │[default: 3.8]                                              │ -│--backend-bDatabase backend to use. If 'none' is chosen, Breeze will start with an invalid│ +│--backend-bDatabase backend to use. If 'none' is chosen, Breeze will start with an invalid│ │database configuration, meaning there will be no database available, and any   │ │attempts to connect to the Airflow database will fail.                         │ │(>sqlite< | mysql | postgres | none)                                           │ │[default: sqlite]                                                              │ -│--postgres-version-PVersion of Postgres used.(>12< | 13 | 14 | 15 | 16)[default: 12]│ -│--mysql-version-MVersion of MySQL used.(>8.0< | 8.4)[default: 8.0]│ -│--cheatsheet/--no-cheatsheet-C/-cEnable/disable cheatsheet.│ -│--asciiart/--no-asciiart-A/-aEnable/disable ASCIIart.│ -│--colour/--no-colourEnable/disable Colour mode (useful for colour blind-friendly communication).│ -╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ -╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ -│--help-hShow this message and exit.│ -╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +│--postgres-version-PVersion of Postgres used.(>12< | 13 | 14 | 15 | 16)[default: 12]│ +│--mysql-version-MVersion of MySQL used.(>8.0< | 8.4)[default: 8.0]│ +│--use-uv/--no-use-uv-U/-uEnable/disable using uv for creating venvs by breeze.│ +│--cheatsheet/--no-cheatsheet-C/-cEnable/disable cheatsheet.│ +│--asciiart/--no-asciiart-A/-aEnable/disable ASCIIart.│ +│--colour/--no-colourEnable/disable Colour mode (useful for colour blind-friendly communication).│ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ +│--help-hShow this message and exit.│ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ diff --git a/dev/breeze/doc/images/output_setup_config.txt b/dev/breeze/doc/images/output_setup_config.txt index f47fa38e42c7b..3b2da9a9c043c 100644 --- 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 index 33f3dd338e4d9..648a13d424658 100644 --- 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 index 407ff7f8cdf3f..bc1ac4f1fa56b 100644 --- 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 index 61460f004ec9f..802a41fc273dd 100644 --- 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 index 944870122747d..791e07cfe718f 100644 --- 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 index 69703b4692b50..3aca9d51c130c 100644 --- 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 index 73cbb430817cc..b34fa3b341020 100644 --- 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 index 0288e49b90975..3c6a175a0fcd0 100644 --- 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 index 5ffd604670b0a..aa5a3160060bf 100755 --- 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 index a90e07d38c9f3..f55a937df0cdf 100755 --- 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)