Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow to switch breeze to use uv internally to create virtualenvs #43587

Merged
merged 1 commit into from
Nov 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -642,6 +642,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' ||
Expand Down
7 changes: 7 additions & 0 deletions .github/workflows/k8s-tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
24 changes: 14 additions & 10 deletions dev/breeze/doc/images/output_setup_config.svg
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
2 changes: 1 addition & 1 deletion dev/breeze/doc/images/output_setup_config.txt
Original file line number Diff line number Diff line change
@@ -1 +1 @@
235af93483ea83592052476479757683
f49dbd1127c59b472db1c92d7362c9e1
Original file line number Diff line number Diff line change
Expand Up @@ -461,7 +461,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

Expand Down
31 changes: 25 additions & 6 deletions dev/breeze/src/airflow_breeze/commands/setup_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,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).",
Expand All @@ -201,6 +207,7 @@ def version():
def change_config(
python: str,
backend: str,
use_uv: bool,
postgres_version: str,
mysql_version: str,
cheatsheet: bool,
Expand All @@ -213,14 +220,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)
Expand All @@ -236,23 +251,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()


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@
"--backend",
"--postgres-version",
"--mysql-version",
"--use-uv",
"--cheatsheet",
"--asciiart",
"--colour",
Expand Down
1 change: 1 addition & 0 deletions dev/breeze/src/airflow_breeze/global_constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,7 @@
ALLOWED_INSTALL_MYSQL_CLIENT_TYPES = ["mariadb", "mysql"]

PIP_VERSION = "24.3.1"
UV_VERSION = "0.4.29"

DEFAULT_UV_HTTP_TIMEOUT = 300
DEFAULT_WSL2_HTTP_TIMEOUT = 900
Expand Down
27 changes: 22 additions & 5 deletions dev/breeze/src/airflow_breeze/utils/kubernetes_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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()),
Expand Down Expand Up @@ -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,
)
Expand All @@ -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():
Expand Down
6 changes: 4 additions & 2 deletions dev/breeze/src/airflow_breeze/utils/run_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, TESTS_PROVIDERS_ROOT
Expand Down Expand Up @@ -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,
Expand Down
30 changes: 27 additions & 3 deletions dev/breeze/src/airflow_breeze/utils/virtualenv_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from collections.abc import Generator
from pathlib import Path

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

Expand All @@ -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()
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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]:
Expand All @@ -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,
)
6 changes: 5 additions & 1 deletion scripts/ci/pre_commit/update_installers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]+)"')

Expand Down Expand Up @@ -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)
Expand Down