From fb3ffdffa53d10e350dc88ffce370d7215153a9a Mon Sep 17 00:00:00 2001 From: Stella Wang Date: Thu, 30 Jul 2026 05:16:41 -0400 Subject: [PATCH 1/2] ci: assert uv.lock keeps torch installable on linux/aarch64 Nothing in CI covered aarch64, so a dependency bump could silently re-break the PyPI fallback that linux/aarch64 relies on for torch and torchvision. That already happened once, when a ROCm bump narrowed tool.uv.environments. Check the lockfile rather than re-resolving: tool.uv.environments must still admit aarch64, every torch extra must resolve torch and torchvision there, and the pinned versions must ship linux/aarch64 wheels for each supported Python. Narrowing tool.uv.environments or dropping the extras' aarch64 fallback pins both leave uv lock --locked green, so neither was caught before. --- .github/workflows/uv-lock-checks.yml | 18 ++- scripts/check_aarch64_lock.py | 161 +++++++++++++++++++++++++++ 2 files changed, 178 insertions(+), 1 deletion(-) create mode 100644 scripts/check_aarch64_lock.py diff --git a/.github/workflows/uv-lock-checks.yml b/.github/workflows/uv-lock-checks.yml index d57163165fb..0c96ad80b9f 100644 --- a/.github/workflows/uv-lock-checks.yml +++ b/.github/workflows/uv-lock-checks.yml @@ -1,6 +1,9 @@ # Check the `uv` lockfile for consistency with `pyproject.toml`. # # If this check fails, you should run `uv lock` to update the lockfile. +# +# Also checks that the lockfile keeps working for linux/aarch64, which no test job covers. +# This reads the lockfile only - no ARM hardware or extra resolution needed. name: 'uv lock checks' @@ -54,9 +57,11 @@ jobs: uvlock-pyprojecttoml: - 'pyproject.toml' - 'uv.lock' + aarch64check: + - 'scripts/check_aarch64_lock.py' - name: setup uv - if: ${{ steps.changed-files.outputs.uvlock-pyprojecttoml_any_changed == 'true' || inputs.always_run == true }} + if: ${{ steps.changed-files.outputs.uvlock-pyprojecttoml_any_changed == 'true' || steps.changed-files.outputs.aarch64check_any_changed == 'true' || inputs.always_run == true }} uses: astral-sh/setup-uv@v8.1.0 with: version: '0.6.10' @@ -66,3 +71,14 @@ jobs: if: ${{ steps.changed-files.outputs.uvlock-pyprojecttoml_any_changed == 'true' || inputs.always_run == true }} run: uv lock --locked # this will exit with 1 if the lockfile is not consistent with pyproject.toml shell: bash + + - name: check aarch64 support in lockfile + # Also runs when the check itself changes, so edits to it are exercised. + if: ${{ steps.changed-files.outputs.uvlock-pyprojecttoml_any_changed == 'true' || steps.changed-files.outputs.aarch64check_any_changed == 'true' || inputs.always_run == true }} + # On linux/aarch64 torch and torchvision must resolve from PyPI rather than from the PyTorch WHL + # indexes, which have no aarch64 torchvision wheel. Several things in pyproject.toml have to hold + # for that, and breaking most of them leaves a lockfile with no torch at all on aarch64 while + # `uv lock --locked` above still passes. See the script's docstring. + # `--no-project` keeps `uv run` from syncing the whole project just to run this. + run: uv run --no-project --with packaging --python 3.12 scripts/check_aarch64_lock.py ./uv.lock + shell: bash diff --git a/scripts/check_aarch64_lock.py b/scripts/check_aarch64_lock.py new file mode 100644 index 00000000000..ff7b2f0c937 --- /dev/null +++ b/scripts/check_aarch64_lock.py @@ -0,0 +1,161 @@ +"""Assert that `uv.lock` gives linux/aarch64 an installable torch and torchvision. + +Nothing else in CI covers aarch64. Three separate mechanisms in `pyproject.toml` conspire to make torch +resolve from PyPI there instead of from the PyTorch WHL indexes (which ship no aarch64 torchvision wheel): +`tool.uv.environments`, the `[tool.uv.sources]` platform markers, and the aarch64 fallback pins in each +torch extra. Break the sources markers and `uv lock` fails loudly, but break either of the other two and +the lockfile simply stops mentioning torch on aarch64 while `uv lock --locked` stays green -- which is how +aarch64 support was silently lost once before, when a ROCm bump narrowed `tool.uv.environments`. + +So assert against the lockfile itself, which is the artifact users actually install from. + +Usage: check_aarch64_lock.py [path/to/uv.lock] (needs `packaging`) +""" + +import sys +import tomllib +from pathlib import Path +from typing import Any + +from packaging.markers import Marker +from packaging.specifiers import SpecifierSet +from packaging.utils import parse_wheel_filename + +# A linux/aarch64 interpreter, for evaluating the lockfile's environment markers. `python_version` is +# filled in per supported minor version below. +AARCH64_ENV = { + "sys_platform": "linux", + "platform_machine": "aarch64", + "platform_system": "Linux", + "os_name": "posix", + "implementation_name": "cpython", + "platform_python_implementation": "CPython", +} + +REQUIRED = ("torch", "torchvision") + + +def env_for(python_version: str, extra: str | None = None) -> dict[str, str]: + env = {**AARCH64_ENV, "python_version": python_version, "python_full_version": f"{python_version}.0"} + if extra is not None: + env["extra"] = extra + return env + + +def matches(marker: str | None, env: dict[str, str]) -> bool: + return marker is None or Marker(marker).evaluate(env) + + +def supported_python_versions(requires_python: str) -> list[str]: + """The `3.x` versions admitted by the lockfile's `requires-python` (e.g. ">=3.11, <3.13" -> 3.11, 3.12).""" + spec = SpecifierSet(requires_python) + return [f"3.{minor}" for minor in range(8, 30) if spec.contains(f"3.{minor}.0")] + + +def has_aarch64_wheel(package: dict[str, Any], python_version: str) -> bool: + """Whether `package` ships a linux/aarch64 wheel usable by CPython `python_version`. + + Only linux aarch64 counts -- macOS arm64 wheels (`macosx_11_0_arm64`) are a different platform. + """ + minor = int(python_version.split(".")[1]) + accepted = {f"cp3{minor}", "py3", f"py3{minor}"} + for wheel in package.get("wheels", []): + filename = wheel.get("url", wheel.get("path", "")).rsplit("/", 1)[-1] + if not filename.endswith(".whl"): + continue + for tag in parse_wheel_filename(filename)[3]: + if "aarch64" not in tag.platform: + continue + if tag.interpreter in accepted: + return True + # An abi3 wheel built for an older CPython also works on this one. + if tag.abi == "abi3" and tag.interpreter.startswith("cp3") and tag.interpreter[3:].isdigit(): + if int(tag.interpreter[3:]) <= minor: + return True + return False + + +def main() -> int: + lock_path = Path(sys.argv[1] if len(sys.argv) > 1 else "uv.lock") + if not lock_path.is_file(): + print(f"File not found: {lock_path}", file=sys.stderr) + return 1 + + text = lock_path.read_text() + lock = tomllib.loads(text) + + python_versions = supported_python_versions(lock["requires-python"]) + packages = {(p["name"], p.get("version"), str(p.get("source"))): p for p in lock["package"]} + problems: list[str] = [] + + # 1. `tool.uv.environments` must still admit aarch64 -- if it doesn't, there is no aarch64 resolution + # to inspect and the per-extra checks below would report a confusing pile of missing torch. + # An empty `supported-markers` means uv was given no restriction, which is fine. + markers = lock.get("supported-markers", []) + if markers and not any(matches(m, env_for(v)) for m in markers for v in python_versions): + print("uv.lock excludes linux/aarch64 entirely. supported-markers:") + for m in markers: + print(f" {m}") + print("\nWiden `tool.uv.environments` in pyproject.toml to include") + print("(sys_platform == 'linux' and platform_machine == 'aarch64'), then re-run `uv lock`.") + return 1 + + # 2. Every conflicting torch extra must resolve torch and torchvision on aarch64, from a registry that + # actually has aarch64 wheels for them. + root_name = "invokeai" + root = next((p for p in lock["package"] if p["name"] == root_name), None) + if root is None: + print(f"no {root_name!r} package in {lock_path} -- has the project been renamed?") + return 1 + extras = sorted({e["extra"] for group in lock.get("conflicts", []) for e in group if e["package"] == root_name}) + if not extras: + print("no conflicting extras found in uv.lock -- has the cpu/cuda/rocm extra layout changed?") + return 1 + + # uv encodes conflicting extras into markers as `extra---`. That is uv's + # internal spelling, so confirm it is still what the lockfile uses -- otherwise the marker evaluation + # below would silently match nothing and we would report a torch problem that isn't real. (A broken + # extra can legitimately drop the token from its own markers, so look at the whole lockfile.) + keys = {extra: f"extra-{len(root_name)}-{root_name}-{extra}" for extra in extras} + unknown = sorted(k for k in keys.values() if k not in text) + if unknown: + print(f"uv.lock never mentions {', '.join(repr(k) for k in unknown)}.") + print("uv's encoding of conflicting extras has changed; update scripts/check_aarch64_lock.py.") + return 1 + + for extra in extras: + deps = root.get("optional-dependencies", {}).get(extra, []) + key = keys[extra] + for name in REQUIRED: + for python_version in python_versions: + env = env_for(python_version, extra=key) + resolved = [d for d in deps if d["name"] == name and matches(d.get("marker"), env)] + if not resolved: + problems.append(f" [{extra}] py{python_version}: no {name} resolves on linux/aarch64") + continue + dep = resolved[0] + package = packages.get((dep["name"], dep.get("version"), str(dep.get("source")))) + registry = (dep.get("source") or {}).get("registry", "?") + if package is None: + problems.append(f" [{extra}] py{python_version}: {name}=={dep.get('version')} not in uv.lock") + elif not has_aarch64_wheel(package, python_version): + problems.append( + f" [{extra}] py{python_version}: {name}=={dep.get('version')} from {registry}" + " has no linux/aarch64 wheel" + ) + else: + print(f" [{extra}] py{python_version}: {name}=={dep.get('version')} from {registry}") + + if problems: + print("\nuv.lock does not give linux/aarch64 an installable torch/torchvision:") + print("\n".join(problems)) + print("\nCheck `tool.uv.environments`, the `[tool.uv.sources]` aarch64 markers and the aarch64") + print("fallback pins in the cpu/cuda/rocm extras in pyproject.toml, then re-run `uv lock`.") + return 1 + + print("\naarch64 OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 55734b999e354d91b012eec6e528064d2ce3de99 Mon Sep 17 00:00:00 2001 From: Stella Wang Date: Thu, 30 Jul 2026 22:41:10 -0400 Subject: [PATCH 2/2] fix(ci): survive percent-encoded wheel URLs in aarch64 lock check The PyTorch WHL indexes percent-encode the `+` of local versions in wheel URLs, so a marker regression that pins aarch64 torch to a WHL index crashed `parse_wheel_filename` with a raw traceback instead of a verdict -- in exactly the still-locks-cleanly state the check exists for. Unquote the filename, and turn any remaining unparseable name into a one-line failure instead of a traceback. Also from review: look up version-less dependency entries (uv omits version/source when a package resolves to a single version across the lockfile), check every dependency entry matching the aarch64 environment rather than the first, and cap `--with packaging` below 26 so a parsing-strictness bump can't change the verdict on its own. --- .github/workflows/uv-lock-checks.yml | 5 +-- scripts/check_aarch64_lock.py | 46 +++++++++++++++++++--------- 2 files changed, 34 insertions(+), 17 deletions(-) diff --git a/.github/workflows/uv-lock-checks.yml b/.github/workflows/uv-lock-checks.yml index 0c96ad80b9f..1d4246ef319 100644 --- a/.github/workflows/uv-lock-checks.yml +++ b/.github/workflows/uv-lock-checks.yml @@ -79,6 +79,7 @@ jobs: # indexes, which have no aarch64 torchvision wheel. Several things in pyproject.toml have to hold # for that, and breaking most of them leaves a lockfile with no torch at all on aarch64 while # `uv lock --locked` above still passes. See the script's docstring. - # `--no-project` keeps `uv run` from syncing the whole project just to run this. - run: uv run --no-project --with packaging --python 3.12 scripts/check_aarch64_lock.py ./uv.lock + # `--no-project` keeps `uv run` from syncing the whole project just to run this. `packaging` is + # capped so a future release tightening wheel-filename parsing can't change the verdict on its own. + run: uv run --no-project --with 'packaging<26' --python 3.12 scripts/check_aarch64_lock.py ./uv.lock shell: bash diff --git a/scripts/check_aarch64_lock.py b/scripts/check_aarch64_lock.py index ff7b2f0c937..2fdb5a0ed02 100644 --- a/scripts/check_aarch64_lock.py +++ b/scripts/check_aarch64_lock.py @@ -16,10 +16,11 @@ import tomllib from pathlib import Path from typing import Any +from urllib.parse import unquote from packaging.markers import Marker from packaging.specifiers import SpecifierSet -from packaging.utils import parse_wheel_filename +from packaging.utils import InvalidWheelFilename, parse_wheel_filename # A linux/aarch64 interpreter, for evaluating the lockfile's environment markers. `python_version` is # filled in per supported minor version below. @@ -60,10 +61,16 @@ def has_aarch64_wheel(package: dict[str, Any], python_version: str) -> bool: minor = int(python_version.split(".")[1]) accepted = {f"cp3{minor}", "py3", f"py3{minor}"} for wheel in package.get("wheels", []): - filename = wheel.get("url", wheel.get("path", "")).rsplit("/", 1)[-1] + # The PyTorch WHL indexes percent-encode the `+` of local versions in wheel URLs (`torch-2.7.1%2Bcpu-...`). + filename = unquote(wheel.get("url", wheel.get("path", "")).rsplit("/", 1)[-1]) if not filename.endswith(".whl"): continue - for tag in parse_wheel_filename(filename)[3]: + try: + tags = parse_wheel_filename(filename)[3] + except InvalidWheelFilename as e: + # The exception message names the offending filename. + sys.exit(f"unparseable wheel filename in uv.lock ({e}) -- update scripts/check_aarch64_lock.py") + for tag in tags: if "aarch64" not in tag.platform: continue if tag.interpreter in accepted: @@ -133,18 +140,27 @@ def main() -> int: if not resolved: problems.append(f" [{extra}] py{python_version}: no {name} resolves on linux/aarch64") continue - dep = resolved[0] - package = packages.get((dep["name"], dep.get("version"), str(dep.get("source")))) - registry = (dep.get("source") or {}).get("registry", "?") - if package is None: - problems.append(f" [{extra}] py{python_version}: {name}=={dep.get('version')} not in uv.lock") - elif not has_aarch64_wheel(package, python_version): - problems.append( - f" [{extra}] py{python_version}: {name}=={dep.get('version')} from {registry}" - " has no linux/aarch64 wheel" - ) - else: - print(f" [{extra}] py{python_version}: {name}=={dep.get('version')} from {registry}") + # uv shouldn't emit overlapping markers, but if it ever does, check every match. + for dep in resolved: + if dep.get("version") is not None: + package = packages.get((name, dep["version"], str(dep.get("source")))) + else: + # uv omits version/source from a dependency entry when the package resolves to a + # single version across the whole lockfile. + candidates = [p for p in lock["package"] if p["name"] == name] + package = candidates[0] if len(candidates) == 1 else None + version = dep.get("version") or (package or {}).get("version") + source = dep.get("source") or (package or {}).get("source") or {} + registry = source.get("registry", "?") + if package is None: + problems.append(f" [{extra}] py{python_version}: {name}=={version} not in uv.lock") + elif not has_aarch64_wheel(package, python_version): + problems.append( + f" [{extra}] py{python_version}: {name}=={version} from {registry}" + " has no linux/aarch64 wheel" + ) + else: + print(f" [{extra}] py{python_version}: {name}=={version} from {registry}") if problems: print("\nuv.lock does not give linux/aarch64 an installable torch/torchvision:")