This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4b848993911 Speed up constraints-version-check by resolving the
baseline once (#70559)
4b848993911 is described below
commit 4b84899391110aa2166a3a57dbc9b98c9c3aaa76
Author: Jarek Potiuk <[email protected]>
AuthorDate: Tue Jul 28 14:01:41 2026 +0200
Speed up constraints-version-check by resolving the baseline once (#70559)
The Deps matrix in finalize-tests is a serial tail on every canary run: it
starts only after the rest of the run has finished and then adds another
40-70 minutes to the wall clock. Roughly half of that time was spent
recomputing an answer the command already had.
Explaining why a package cannot be upgraded needs two resolutions: the
unpinned workspace, and the workspace with the package pinned to its
latest version. Only the second depends on the package. The first is the
same for every package in a run, because each explanation restores
pyproject.toml and uv.lock before the next one starts and the remaining
inputs are fixed for the whole command. It was nevertheless re-run for
each outdated package, so a scheduled run doing 63 explanations paid for
126 full `uv sync --refresh` invocations where 64 would have done.
Resolving the baseline once also removes a latent source of
non-determinism: every baseline sync passed --refresh, so an index change
part-way through a run could leave two packages compared against
different baselines.
---
.../utils/constraints_version_check.py | 211 +++++++++++++--------
dev/breeze/tests/test_constraints_version_check.py | 169 +++++++++++++++++
2 files changed, 296 insertions(+), 84 deletions(-)
diff --git a/dev/breeze/src/airflow_breeze/utils/constraints_version_check.py
b/dev/breeze/src/airflow_breeze/utils/constraints_version_check.py
index 16e0cdb51a4..d5ee2f6cee8 100755
--- a/dev/breeze/src/airflow_breeze/utils/constraints_version_check.py
+++ b/dev/breeze/src/airflow_breeze/utils/constraints_version_check.py
@@ -446,14 +446,6 @@ def process_packages(
github_repository: str | None,
cooldown_days: int = 4,
) -> tuple[int, int, list[str], dict[str, int]]:
- @contextmanager
- def preserve_pyproject_file(pyproject_path: Path):
- original_content = pyproject_path.read_text()
- try:
- yield
- finally:
- pyproject_path.write_text(original_content)
-
def fetch_pypi_data(pkg: str) -> dict:
pypi_url = f"https://pypi.org/pypi/{pkg}/json";
with urllib.request.urlopen(pypi_url) as resp:
@@ -472,6 +464,9 @@ def process_packages(
skipped_count = 0
explanations = []
status_counts: dict[str, int] = {"ok": 0, "new": 0, "warning": 0,
"critical": 0}
+ # Resolved lazily on the first package that needs an explanation, then
shared by all of
+ # them — see resolve_baseline_versions() for why one resolution is enough.
+ baseline: tuple[str, dict[str, str]] | None = None
for pkg, pinned_version in packages:
try:
@@ -505,6 +500,13 @@ def process_packages(
skipped_count += 1
if explain_why and not is_latest_version:
+ if baseline is None:
+ baseline = resolve_baseline_versions(
+ python_version=python_version,
+ airflow_constraints_mode=airflow_constraints_mode,
+ github_repository=github_repository,
+ )
+ baseline_text, baseline_versions = baseline
explanation = explain_package_upgrade(
pkg=pkg,
pinned_version=pinned_version,
@@ -512,6 +514,8 @@ def process_packages(
python_version=python_version,
airflow_constraints_mode=airflow_constraints_mode,
github_repository=github_repository,
+ baseline_text=baseline_text,
+ baseline_versions=baseline_versions,
)
explanations.append(explanation)
except HTTPError as e:
@@ -629,92 +633,122 @@ def find_downgrades(
return sorted(downgrades)
-def explain_package_upgrade(
- pkg: str,
- pinned_version: str,
- latest_version: str,
- python_version: str,
- airflow_constraints_mode: str,
- github_repository: str | None,
-) -> str:
- explanation = (
- f"[bold blue]\n--- Explaining for {pkg} (current: {pinned_version},
latest: {latest_version}) ---[/]"
- )
+@contextmanager
+def preserve_files(*paths: Path):
+ """Restore the given files' contents on exit — ``uv sync`` rewrites
``uv.lock``."""
+ originals = {path: path.read_text() for path in paths}
+ try:
+ yield
+ finally:
+ for path, content in originals.items():
+ path.write_text(content)
- @contextmanager
- def preserve_pyproject_file(pyproject_path: Path):
- original_content = pyproject_path.read_text()
- try:
- yield pyproject_path
- finally:
- pyproject_path.write_text(original_content)
- additional_args = []
+def get_additional_sync_args(airflow_constraints_mode: str) -> list[str]:
if airflow_constraints_mode == "constraints-source-providers":
# In case of source constraints we also need to add all development
dependencies
# to reflect exactly what is installed in the CI image by default. The
``ci-image``
# group aggregates dev/docs/docs-gen plus any hard-to-install provider
extras
# (see root pyproject.toml).
- additional_args.extend(["--group", "ci-image"])
- with (
- preserve_pyproject_file(AIRFLOW_ROOT_PATH / "pyproject.toml") as
airflow_pyproject,
- preserve_pyproject_file(AIRFLOW_ROOT_PATH / "uv.lock"),
- ):
- from packaging.utils import canonicalize_name
+ return ["--group", "ci-image"]
+ return []
+
+
+# Marker echoed between ``uv sync`` and ``uv pip freeze`` so the freeze output
can be
+# sliced out of the combined shell log.
+FREEZE_MARKER = "===BREEZE_RESOLVED_FREEZE==="
- canonical_pkg = str(canonicalize_name(pkg))
- shell_params = ShellParams(
+def sync_and_freeze(
+ *,
+ python_version: str,
+ airflow_constraints_mode: str,
+ github_repository: str | None,
+ title: str,
+):
+ """Resolve at --resolution highest and, in the *same* shell, freeze the
result.
+
+ Each ``execute_command_in_shell`` call is a fresh ``docker compose run
--rm``
+ container, so running ``uv pip freeze`` as a separate call would not
reliably see
+ the environment the sync just populated. Chaining both in one ``bash -c``
keeps
+ the freeze in the same shell/venv as the sync. ``&&`` ensures the freeze
only runs
+ when the sync succeeds and that a sync failure is still reflected in the
return
+ code. Returns ``(result, combined_output_text, {canonical_name:
version})``.
+ """
+ sync = shlex.join(
+ [
+ "uv",
+ "sync",
+ "--all-packages",
+ *get_additional_sync_args(airflow_constraints_mode),
+ "--resolution",
+ "highest",
+ "--refresh",
+ "--python",
+ python_version,
+ ]
+ )
+ output = Output(title=title, file_name=get_temp_file_name())
+ result = execute_command_in_shell(
+ ShellParams(
github_repository=github_repository,
python=python_version,
mount_sources=MOUNT_SELECTED,
+ ),
+ project_name="breeze-constraints",
+ command=shlex.join(["bash", "-c", f"{sync} && echo {FREEZE_MARKER} &&
uv pip freeze"]),
+ output=output,
+ signal_error=False,
+ )
+ text = Path(output.file_name).read_text()
+ versions = parse_freeze(text.split(FREEZE_MARKER, 1)[1]) if FREEZE_MARKER
in text else {}
+ return result, text, versions
+
+
+def resolve_baseline_versions(
+ *,
+ python_version: str,
+ airflow_constraints_mode: str,
+ github_repository: str | None,
+) -> tuple[str, dict[str, str]]:
+ """Resolve the unpinned workspace at --resolution highest, once.
+
+ This is the resolution that actually generates the constraints, so it is
the ground
+ truth for "what would the constraints pick". It depends only on the
workspace and the
+ command's own arguments — never on which package is being explained — and
every
+ explanation restores ``pyproject.toml``/``uv.lock`` before the next one
starts. So one
+ resolution is enough for the whole run, and recomputing it per package
would repeat an
+ identical several-minute ``uv sync`` dozens of times.
+ """
+ with preserve_files(AIRFLOW_ROOT_PATH / "pyproject.toml",
AIRFLOW_ROOT_PATH / "uv.lock"):
+ _, text, versions = sync_and_freeze(
+ python_version=python_version,
+ airflow_constraints_mode=airflow_constraints_mode,
+ github_repository=github_repository,
+ title="output_baseline",
)
+ return text, versions
- # Marker echoed between ``uv sync`` and ``uv pip freeze`` so the
freeze output can be
- # sliced out of the combined shell log.
- freeze_marker = "===BREEZE_RESOLVED_FREEZE==="
-
- def sync_and_freeze(title: str):
- """Resolve at --resolution highest and, in the *same* shell,
freeze the result.
-
- Each ``execute_command_in_shell`` call is a fresh ``docker compose
run --rm``
- container, so running ``uv pip freeze`` as a separate call would
not reliably see
- the environment the sync just populated. Chaining both in one
``bash -c`` keeps
- the freeze in the same shell/venv as the sync. ``&&`` ensures the
freeze only runs
- when the sync succeeds and that a sync failure is still reflected
in the return
- code. Returns ``(result, combined_output_text, {canonical_name:
version})``.
- """
- sync = shlex.join(
- [
- "uv",
- "sync",
- "--all-packages",
- *additional_args,
- "--resolution",
- "highest",
- "--refresh",
- "--python",
- python_version,
- ]
- )
- output = Output(title=title, file_name=get_temp_file_name())
- result = execute_command_in_shell(
- shell_params,
- project_name="breeze-constraints",
- command=shlex.join(["bash", "-c", f"{sync} && echo
{freeze_marker} && uv pip freeze"]),
- output=output,
- signal_error=False,
- )
- text = Path(output.file_name).read_text()
- versions = parse_freeze(text.split(freeze_marker, 1)[1]) if
freeze_marker in text else {}
- return result, text, versions
- # Baseline: resolve the workspace at --resolution highest *without*
any pin and
- # record what version that resolution naturally selects for the
package. This is
- # the resolution that actually generates the constraints, so it is the
ground truth
- # for "what would the constraints pick".
- _, before_text, before_versions = sync_and_freeze("output_before")
- baseline_version = before_versions.get(canonical_pkg)
+def explain_package_upgrade(
+ pkg: str,
+ pinned_version: str,
+ latest_version: str,
+ python_version: str,
+ airflow_constraints_mode: str,
+ github_repository: str | None,
+ baseline_text: str,
+ baseline_versions: dict[str, str],
+) -> str:
+ explanation = (
+ f"[bold blue]\n--- Explaining for {pkg} (current: {pinned_version},
latest: {latest_version}) ---[/]"
+ )
+ with preserve_files(AIRFLOW_ROOT_PATH / "pyproject.toml",
AIRFLOW_ROOT_PATH / "uv.lock"):
+ from packaging.utils import canonicalize_name
+
+ airflow_pyproject = AIRFLOW_ROOT_PATH / "pyproject.toml"
+ canonical_pkg = str(canonicalize_name(pkg))
+ baseline_version = baseline_versions.get(canonical_pkg)
update_pyproject_dependency(airflow_pyproject, pkg, latest_version,
python_version)
if get_verbose():
@@ -722,7 +756,12 @@ def explain_package_upgrade(
airflow_pyproject.read_text(), "toml", theme="monokai",
line_numbers=True, word_wrap=False
)
explanation += "\n" + str(syntax)
- after_result, after_text, after_versions =
sync_and_freeze("output_after")
+ after_result, after_text, after_versions = sync_and_freeze(
+ python_version=python_version,
+ airflow_constraints_mode=airflow_constraints_mode,
+ github_repository=github_repository,
+ title="output_after",
+ )
# A zero exit code only proves that *some* valid resolution exists
with the pin — not
# that --resolution highest would ever select it. Inspect what was
actually resolved:
@@ -730,7 +769,7 @@ def explain_package_upgrade(
# resolution (i.e. the constraints) keeps the package at its lower
version, so this is
# NOT a clean upgrade.
resolved_version = after_versions.get(canonical_pkg)
- downgrades = find_downgrades(before_versions, after_versions,
exclude=canonical_pkg)
+ downgrades = find_downgrades(baseline_versions, after_versions,
exclude=canonical_pkg)
if after_result.returncode != 0:
# Forcing the package to its latest version produced no valid
resolution at all:
@@ -743,7 +782,7 @@ def explain_package_upgrade(
conflict = extract_uv_conflict(after_text)
if conflict:
explanation += f"\n\n[bold yellow]Conflict as reported by
uv:[/]\n{conflict}"
- elif not before_versions or not after_versions:
+ elif not baseline_versions or not after_versions:
# Without the resolved version lists we cannot tell a clean
upgrade apart from one
# that only works by downgrading other packages — never silently
claim success.
explanation += (
@@ -781,7 +820,11 @@ def explain_package_upgrade(
printf_cmd = "printf '%s\\n' " + " ".join(shlex.quote(pin) for pin
in conflict_pins)
probe_output = Output(title="conflict_probe",
file_name=get_temp_file_name())
execute_command_in_shell(
- shell_params,
+ ShellParams(
+ github_repository=github_repository,
+ python=python_version,
+ mount_sources=MOUNT_SELECTED,
+ ),
project_name="breeze-constraints",
command=shlex.join(
[
@@ -811,7 +854,7 @@ def explain_package_upgrade(
# Full resolver logs of both phases — only when explicitly
requested, since they
# are very long (each is a complete uv sync plus freeze).
explanation += (
- f"\n\n[yellow]--- uv resolver output: phase 1, baseline (no
pin) ---[/]\n{before_text}"
+ f"\n\n[yellow]--- uv resolver output: phase 1, baseline (no
pin) ---[/]\n{baseline_text}"
f"\n[yellow]--- uv resolver output: phase 2, with
{pkg}=={latest_version} pinned ---[/]"
f"\n{after_text}"
)
diff --git a/dev/breeze/tests/test_constraints_version_check.py
b/dev/breeze/tests/test_constraints_version_check.py
new file mode 100644
index 00000000000..c61f0d3f1ce
--- /dev/null
+++ b/dev/breeze/tests/test_constraints_version_check.py
@@ -0,0 +1,169 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import contextlib
+import json
+from pathlib import Path
+from unittest import mock
+
+import pytest
+
+from airflow_breeze.utils.constraints_version_check import (
+ explain_package_upgrade,
+ get_table_format,
+ process_packages,
+)
+
+MODULE = "airflow_breeze.utils.constraints_version_check"
+
+# Old enough that the default 4-day cooldown never filters these releases out.
+OLD_UPLOAD_TIME = "2020-01-01T00:00:00.000000Z"
+
+
+def _pypi_payload(latest: str, *versions: str) -> bytes:
+ releases = {v: [{"upload_time_iso_8601": OLD_UPLOAD_TIME, "yanked":
False}] for v in versions}
+ return json.dumps({"info": {"version": latest}, "releases":
releases}).encode()
+
+
[email protected]
+def pypi(monkeypatch):
+ """Serve every package the same two-release history: pinned 1.0.0, latest
2.0.0."""
+
+ def fake_urlopen(_url):
+ response = mock.MagicMock()
+ response.read.return_value = _pypi_payload("2.0.0", "1.0.0", "2.0.0")
+ return contextlib.nullcontext(response)
+
+ monkeypatch.setattr(f"{MODULE}.urllib.request.urlopen", fake_urlopen)
+
+
+def _run_process_packages(packages, explain_why=True):
+ col_widths, format_str, _, _ = get_table_format(packages)
+ return process_packages(
+ packages=packages,
+ constraints_date=None,
+ mode="full",
+ explain_why=explain_why,
+ col_widths=col_widths,
+ format_str=format_str,
+ python_version="3.11",
+ airflow_constraints_mode="constraints",
+ github_repository="apache/airflow",
+ )
+
+
[email protected](f"{MODULE}.explain_package_upgrade", return_value="explanation")
[email protected](f"{MODULE}.resolve_baseline_versions", return_value=("baseline
log", {"pkg-a": "1.0.0"}))
+def test_baseline_is_resolved_once_for_all_outdated_packages(mock_baseline,
mock_explain, pypi):
+ packages = [("pkg-a", "1.0.0"), ("pkg-b", "1.0.0"), ("pkg-c", "1.0.0")]
+
+ _, _, explanations, _ = _run_process_packages(packages)
+
+ assert mock_explain.call_count == 3
+ assert len(explanations) == 3
+ mock_baseline.assert_called_once_with(
+ python_version="3.11",
+ airflow_constraints_mode="constraints",
+ github_repository="apache/airflow",
+ )
+ for call in mock_explain.call_args_list:
+ assert call.kwargs["baseline_text"] == "baseline log"
+ assert call.kwargs["baseline_versions"] == {"pkg-a": "1.0.0"}
+
+
[email protected](f"{MODULE}.explain_package_upgrade", return_value="explanation")
[email protected](f"{MODULE}.resolve_baseline_versions")
+def test_baseline_is_not_resolved_when_nothing_needs_explaining(mock_baseline,
mock_explain, pypi):
+ # Already at the latest version, so no package triggers an explanation.
+ _run_process_packages([("pkg-a", "2.0.0"), ("pkg-b", "2.0.0")])
+
+ mock_explain.assert_not_called()
+ mock_baseline.assert_not_called()
+
+
[email protected](f"{MODULE}.explain_package_upgrade")
[email protected](f"{MODULE}.resolve_baseline_versions")
+def test_baseline_is_not_resolved_without_explain_why(mock_baseline,
mock_explain, pypi):
+ _run_process_packages([("pkg-a", "1.0.0")], explain_why=False)
+
+ mock_explain.assert_not_called()
+ mock_baseline.assert_not_called()
+
+
[email protected](f"{MODULE}.update_pyproject_dependency")
[email protected](f"{MODULE}.preserve_files")
[email protected](f"{MODULE}.sync_and_freeze")
+def test_explain_package_upgrade_syncs_only_the_pinned_resolution(
+ mock_sync, mock_preserve, mock_update_pyproject
+):
+ mock_preserve.return_value = contextlib.nullcontext()
+ mock_sync.return_value = (mock.MagicMock(returncode=0), "after log",
{"pkg-a": "2.0.0"})
+
+ explanation = explain_package_upgrade(
+ pkg="pkg-a",
+ pinned_version="1.0.0",
+ latest_version="2.0.0",
+ python_version="3.11",
+ airflow_constraints_mode="constraints",
+ github_repository="apache/airflow",
+ baseline_text="baseline log",
+ baseline_versions={"pkg-a": "1.0.0"},
+ )
+
+ mock_sync.assert_called_once()
+ assert mock_sync.call_args.kwargs["title"] == "output_after"
+ assert "can be upgraded from 1.0.0 to 2.0.0" in explanation
+
+
[email protected](f"{MODULE}.update_pyproject_dependency")
[email protected](f"{MODULE}.preserve_files")
[email protected](f"{MODULE}.execute_command_in_shell")
[email protected](f"{MODULE}.sync_and_freeze")
+def test_explain_package_upgrade_reads_baseline_for_downgrade_detection(
+ mock_sync, mock_conflict_probe, mock_preserve, mock_update_pyproject
+):
+ mock_preserve.return_value = contextlib.nullcontext()
+
+ def write_conflict_narrative(*_args, **kwargs):
+ # The downgrade branch reruns uv from scratch to capture the resolver
narrative.
+ Path(kwargs["output"].file_name).write_text(
+ "No solution found\nBecause pkg-a==2.0.0 depends on other-pkg<5"
+ )
+ return mock.MagicMock(returncode=1)
+
+ mock_conflict_probe.side_effect = write_conflict_narrative
+ # Reaching pkg-a 2.0.0 pushed other-pkg back from 5.0.0 to 4.0.0.
+ mock_sync.return_value = (
+ mock.MagicMock(returncode=0),
+ "after log",
+ {"pkg-a": "2.0.0", "other-pkg": "4.0.0"},
+ )
+
+ explanation = explain_package_upgrade(
+ pkg="pkg-a",
+ pinned_version="1.0.0",
+ latest_version="2.0.0",
+ python_version="3.11",
+ airflow_constraints_mode="constraints",
+ github_repository="apache/airflow",
+ baseline_text="baseline log",
+ baseline_versions={"pkg-a": "1.0.0", "other-pkg": "5.0.0"},
+ )
+
+ assert "only by DOWNGRADING" in explanation
+ assert "other-pkg: 5.0.0 -> 4.0.0" in explanation
