ci: check aarch64 dependency resolution in uv-lock-checks

[stellarfeline]

Summary

Follow-up to #9095, as offered there: a resolution-only CI guard for linux/aarch64 that needs no ARM hardware.

Nothing in CI exercises aarch64 dependency resolution, so a future dependency bump could silently re-break the PyPI fallback markers that #9095 adds to pyproject.toml (this exact regression class already happened once, when the ROCm 7.1 bump added an environments filter that excluded ARM).

This adds one step to the existing uv lock checks workflow, running under the same change-detection condition and uv version (0.6.10):

for extra in cpu cuda rocm; do
  uv pip compile pyproject.toml --extra $extra --python-platform aarch64-manylinux_2_28 --python-version 3.12 > /dev/null
done

All three torch extras are checked because each carries its own aarch64 fallback pins. aarch64-manylinux_2_28 (rather than the bare aarch64-unknown-linux-gnu alias) is required with uv 0.6.10: the alias resolves to a manylinux_2_17 baseline, below the 2_27/2_28 glibc floor of the onnxruntime and torch aarch64 wheels. 2_28 matches what the wheels themselves target.

Stacked on #9095 — the new check only passes with that PR's pyproject.toml markers, so this branch includes its commits. Once #9095 merges, this diff reduces to the single workflow file.

QA

Verified locally with the exact CI uv version (0.6.10, aarch64 binary):

• Positive: all three extras resolve successfully against this branch's pyproject.toml (also cross-checked with uv 0.11.8).
• Negative: the same command against current main's pyproject.toml fails with torchvision==0.22.1+cpu has no wheels with a matching platform tag — i.e. the guard catches precisely the regression it targets.

The uv lock checks workflow run on this PR exercises the step for real.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable) (CI check)
• Documentation added / updated (if applicable)

[lstein]

Thanks for following up on this — a resolution-only aarch64 guard is exactly the right shape for the problem, and the manylinux_2_28 detail is a genuinely non-obvious bit of legwork (I confirmed the bare aarch64-unknown-linux-gnu alias fails on onnxruntime==1.19.2 under uv 0.6.10, so that choice is load-bearing and well-explained).

I reproduced everything below with CI's exact uv (0.6.10, aarch64-equivalent resolution from an x86_64 host), plus real uv lock runs to see which breakages the lockfile actually encodes. My conclusion is that the check as written misses the regression it was written for, so I'd like to see it re-aimed at uv.lock before this lands.

---

0. The branch is stale

#9095 merged on 2026-07-17, so everything here except the workflow file is already in main. As of now git diff origin/main HEAD is a set of reverts: diffusers 0.39.0 → 0.37.0, removal of imageio[ffmpeg] / psutil, removal of the invokeai.backend.qwen3 / invokeai.backend.t5 package-data entries, and three deleted rows in system-requirements.mdx. A merge won't apply those, but CI on this PR is currently exercising a three-week-old dependency set, so the run isn't evidence about the merged result.

Good news: I ran the new step against current main's pyproject.toml with uv 0.6.10 and all three extras pass, so a rebase won't turn it red.

1. (blocking) The guard doesn't catch the environments regression it targets

The description names the motivating incident: "the ROCm 7.1 bump added an environments filter that excluded ARM". uv pip compile ignores tool.uv.environments.

Taking current main's pyproject.toml and narrowing only that one line back to x86_64-only:

environments = ["sys_platform == 'win32' or sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'x86_64')"]

• all three uv pip compile ... --python-platform aarch64-manylinux_2_28 invocations exit 0 (uv 0.6.10 and 0.11.23 alike);

• uv lock succeeds, and the regenerated lock's cpu extra contains only

  { name = "torch", version = "2.7.1+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" },
    marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }
  

  i.e. no torch at all on linux/aarch64;

• uv lock --locked passes.

Entirely green CI, aarch64 silently dead again. That's the exact historical failure replayed against the new guard.

2. (blocking) It also misses removal of the extras' aarch64 fallback pins

pyproject.toml says those entries are load-bearing:

Without these explicit entries the extra's conflict-universe would contain no torch at all on aarch64 (uv partitions the base declarations into the no-extra universe).

That comment is correct. I deleted the six torch==2.7.1; sys_platform == 'linux' and platform_machine == 'aarch64' / torchvision==0.22.1; ... lines and ran uv lock: the cpu, cuda and rocm extras again resolve torch for x86_64/darwin/win32 only. The new step exits 0 on all three extras for that same file.

So the two mechanisms whose own comments warn "break this and aarch64 gets no torch" are both invisible to the guard. (The third — the { index = "pypi", marker = ... } source entries — is fine: removing those makes uv lock itself fail with No solution found ... for split (platform_machine == 'aarch64' ...), so it's already covered.)

3. The one class it does catch is already covered by uv lock --locked

Now that environments includes aarch64, a leaked WHL-index pin makes resolution outright impossible. Restoring the pre-#9095 unconditional pins on top of current main:

"cpu" = ["torch==2.7.1+cpu", "torchvision==0.22.1+cpu"]

uv lock exits 1 (No solution found when resolving dependencies for split (platform_machine == 'aarch64' and sys_platform == 'linux')), so the existing step at line 70 already fails without the new one. The new step's marginal value here is a clearer error that names the offending extra — real, but small.

Combined with §1 and §2: the step duplicates what's covered and skips what isn't.

4. It's non-hermetic — it validates a dependency set nobody installs

uv pip compile re-resolves from the live indexes rather than reading uv.lock. Comparing today's aarch64 compile of main's pyproject.toml against main's uv.lock, 20 of 188 packages differ — bitsandbytes 0.50.0 vs 0.49.2, huggingface-hub 1.25.1 vs 1.24.0, websockets 17.0 vs 16.1.1, tqdm 4.70.0 vs 4.69.0, and so on.

Two consequences: an unrelated upstream release that drops aarch64 wheels turns this red on a PR that only touched a comment in pyproject.toml; and conversely a locked version lacking aarch64 wheels stays green as long as some newer version has them.

5. Suggested replacement: assert against the lockfile

This checks the artifact we actually ship. I verified it passes on main's real uv.lock and fails on both §1 and §2 (which the current step passes):

      - name: check aarch64 resolution
        if: ${{ steps.changed-files.outputs.uvlock-pyprojecttoml_any_changed == 'true' || inputs.always_run == true }}
        # torch/torchvision must fall back to PyPI on linux/aarch64 (the PyTorch WHL
        # indexes have no aarch64 torchvision wheels). Three separate things in
        # pyproject.toml make that work -- tool.uv.environments, the extra-scoped
        # index markers, and the aarch64 fallback pins -- and breaking any of them
        # leaves a lockfile that resolves no torch at all on aarch64 while every
        # other check stays green. Assert directly on the lock instead.
        # --no-project matters: without it `uv run` would sync the whole project.
        run: uv run --no-project --with packaging --python 3.12 scripts/check_aarch64_lock.py
        shell: bash

# scripts/check_aarch64_lock.py
"""Assert that uv.lock resolves torch and torchvision for linux/aarch64.

See the comment in .github/workflows/uv-lock-checks.yml. Run from the repo root.
"""

import sys
import tomllib

from packaging.markers import Marker

AARCH64 = {
    "sys_platform": "linux",
    "platform_machine": "aarch64",
    "os_name": "posix",
    "platform_system": "Linux",
    "python_version": "3.12",
    "python_full_version": "3.12.0",
    "implementation_name": "cpython",
    "platform_python_implementation": "CPython",
}

with open("uv.lock", "rb") as f:
    lock = tomllib.load(f)

root = next(p for p in lock["package"] if p["name"] == "invokeai")
problems: list[str] = []

for extra in ("cpu", "cuda", "rocm"):
    deps = root.get("optional-dependencies", {}).get(extra, [])
    # uv encodes conflicting extras into markers as `extra-<len(name)>-<name>-<extra>`.
    key = f"extra-{len('invokeai')}-invokeai-{extra}"
    env = {**AARCH64, "extra": key}
    for want in ("torch", "torchvision"):
        resolved = [
            d
            for d in deps
            if d["name"] == want and ("marker" not in d or Marker(d["marker"]).evaluate(env))
        ]
        if not resolved:
            problems.append(f"  [{extra}] no {want} resolves on linux/aarch64")
        else:
            src = resolved[0].get("source", {}).get("registry", "?")
            print(f"  [{extra}] {want}=={resolved[0].get('version')} from {src}")

if problems:
    print("\nuv.lock does not resolve torch/torchvision on linux/aarch64:")
    print("\n".join(problems))
    print("\nCheck tool.uv.environments, the [tool.uv.sources] aarch64 markers, and the")
    print("aarch64 fallback pins in the cpu/cuda/rocm extras, then re-run `uv lock`.")
    sys.exit(1)

print("\naarch64 OK")

Output on main today:

  [cpu] torch==2.7.1 from https://pypi.org/simple
  [cpu] torchvision==0.22.1 from https://pypi.org/simple
  [cuda] torch==2.7.1 from https://pypi.org/simple
  [cuda] torchvision==0.22.1 from https://pypi.org/simple
  [rocm] torch==2.7.1 from https://pypi.org/simple
  [rocm] torchvision==0.22.1 from https://pypi.org/simple

aarch64 OK

Two caveats worth knowing: the extra-8-invokeai-* key shape is uv's internal encoding for conflicting extras and could change across uv versions (it would fail loudly, not silently — and the pinned uv makes that a deliberate-upgrade event); and it asserts resolution, not wheel availability for the pinned versions. Because of that second gap, keeping your uv pip compile loop alongside it is defensible — I'd just not have it be the only guard. If you keep it, --exclude-newer would take the edge off §4.

6. Smaller things

• manylinux_2_28 is an undocumented magic number. It tracks what today's torch/onnxruntime aarch64 wheels target. When some dependency raises its aarch64 glibc floor to 2_31/2_34, CI fails on something that works for every real user. A comment saying "this is the floor of the wheels, bump it when they bump" — or deriving it from the Docker base image, which is what actually constrains us — would save the next person a confusing debug session.
• Only Python 3.12 is checked, though requires-python = ">=3.11, <3.13". 3.11 resolves fine today (I checked), so it's a coverage gap rather than a bug; for py in 3.11 3.12 is nearly free.
• timeout-minutes: 5 # expected run time: <1 min (line 41) now covers four full resolutions against download.pytorch.org plus PyPI. Probably still fine with enable-cache: true, but the comment is stale.

Things I attacked that held up

• shell: bash supplies -eo pipefail, so a failure in the first loop iteration does abort — the loop doesn't mask exit codes.
• --python-version 3.12 doesn't need a matching interpreter on the runner; no setup-python step is required.
• The pin-leak detection isn't uv-version-fragile — pre-#9095 pyproject.toml fails under both 0.6.10 and 0.11.23.
• Checking all three extras isn't redundant even though they collapse to identical aarch64 resolutions: a leaked pin in one extra shows up only in that extra (my cpu-only regression failed cpu and passed cuda/rocm).
• The new step's if: condition matches the existing steps exactly.

---

To summarise what I'd need to approve: rebase onto main (the diff should reduce to the workflow file, as you predicted), and re-aim the check at uv.lock so it covers the environments and fallback-pin regressions. Happy to push the lock-based check to this branch myself if that's easier — just say the word.

[stellarfeline]

Thanks — you're right on both blocking points, and I reproduced them before changing anything (uv 0.6.10, real uv lock runs on current main): narrowing tool.uv.environments and deleting the extras' aarch64 fallback pins each leave uv lock, uv lock --locked and the old compile loop all exiting 0 with no torch on aarch64. §3 also confirmed — uv lock alone already fails there, which was the loop's only real catch. So I've re-aimed the check at uv.lock and dropped the compile loop rather than keeping it with --exclude-newer.

Rebased; as you predicted the diff is now just the workflow plus scripts/check_aarch64_lock.py.

Changes from your sketch:

• supported-markers checked first, so the environments regression gets a precise diagnostic instead of six "no torch" lines.
• Extras read from the lockfile's conflicts table rather than hardcoded, so a future fourth torch extra is covered.
• Asserts wheel availability, not just resolution — uv.lock lists wheel filenames, so this closes your second caveat, hermetically and for the pinned version (strictly stronger than the loop, which was happy if some version had wheels).
• 3.11 and 3.12, derived from requires-python, so the 3.13 bump extends coverage on its own. The manylinux_2_28 magic number is gone with the loop, and the step no longer touches download.pytorch.org at all, so the timeout-minutes comment is accurate again.
• Your extra-8-invokeai-* caveat is an explicit guard now. Worth knowing: a per-extra version of that guard is wrong — when an extra's conflict collapses (your §2 case) uv drops the token from that extra's markers entirely, so it reports "encoding changed" for a real regression. It scans the whole lockfile instead.

One more thing your review surfaced indirectly: with the rebase the PR no longer touches pyproject.toml/uv.lock, so the step was skipped on my first push — the old PR only appeared to exercise it because it was stacked on #9095. I've added a changed-files group on the script itself (and extended setup uv to match, or it would fail on command not found), so edits to the check re-run it. It now runs green in CI for real, with check lockfile correctly skipping.

Known gap, stated plainly: this asserts torch/torchvision, not that every other dependency has aarch64 wheels. The loop nominally covered that but non-hermetically, with the false-positive behaviour in your §4; honest coverage needs an ARM runner installing the lock. Happy to open a follow-up if you want it tracked.

Thanks for the offer to push it yourself — no need, but the shape of it is yours.

[lstein]

Re-reviewed at 923625a635 with CI's uv 0.6.10. Both blocking points are properly fixed — I re-ran my original regression locks against the new script rather than taking the description's word for it:

lockfile	result
current main's real uv.lock	aarch64 OK, exit 0
environments narrowed to x86_64 (§1)	exit 1 — "uv.lock excludes linux/aarch64 entirely", with the offending marker printed
aarch64 fallback pins deleted (§2)	exit 1 — per-extra "no torch resolves on linux/aarch64"

The wheel assertion isn't vacuous either: stripping the aarch64 wheels from the PyPI torch 2.7.1 package entry yields torch==2.7.1 from https://pypi.org/simple has no linux/aarch64 wheel for all three extras. That is genuinely stronger than the old compile loop, and hermetic — nice.

I also confirmed the step actually executes rather than skipping: the uv-lock-checks log on this head shows all twelve [extra] pyX.Y: ... lines and aarch64 OK in 0.4s. Catching that the previous version had been silently skipped, and fixing the setup uv gate along with it, was a good save — that failure mode would have been invisible.

The three follow-ups from the reply all check out: supported-markers is the key uv writes from tool.uv.environments, extras discovered from the conflicts table resolve to cpu/cuda/rocm, and supported_python_versions yields exactly 3.11, 3.12 from the lock's own requires-python. The reasoning about scanning the whole lockfile for the extra-token encoding rather than per-extra is right, and I verified it: neither regression lock trips that guard falsely.

Keeping this as changes-requested for one defect below, which is a small fix.

---

1. Percent-encoded wheel URLs crash the script — reachable via the regression it guards [blocking, one-line fix]

The PyTorch WHL index percent-encodes + in wheel URLs. has_aarch64_wheel hands the raw filename to parse_wheel_filename (scripts/check_aarch64_lock.py:63,66), which rejects it:

packaging.utils.InvalidWheelFilename: Invalid wheel filename (invalid version):
  'torch-2.7.1%2Bcpu-cp311-cp311-manylinux_2_28_aarch64'

Triggering sequence: a marker regression lets the +cpu (or +cu128) pin apply on aarch64. Because torch — unlike torchvision — does ship aarch64 wheels on the index, uv lock still succeeds and uv lock --locked stays green; the script then selects that package for aarch64 and CI gets an unhandled traceback instead of a verdict. I reproduced it by repointing the cpu extra's aarch64 torch at 2.7.1+cpu in main's lock.

So this fires in precisely the situation the check exists for: a torch marker regression, in a state that still locks cleanly.

The fix is unquote(), not a bare try/except — swallowing the error and continuing would report "has no linux/aarch64 wheel" for a wheel that plainly is one:

from urllib.parse import unquote

filename = unquote(wheel.get("url", wheel.get("path", "")).rsplit("/", 1)[-1])

Verified: with that change the same lock correctly prints aarch64 OK (torch+cpu genuinely has an aarch64 wheel), and main's lock is unchanged. Worth adding a try/except InvalidWheelFilename underneath that fails loudly and names the filename, so a future unparseable name is a clear error rather than a raw traceback — but the unquote is the actual fix.

2. Version-less dependency entries report torch==None not in uv.lock [non-blocking]

packages is keyed on (name, version, source), but uv omits version/source from a dependency entry when the package resolves to a single version across the whole lock. I built such a lockfile — all three extras resolving PyPI torch==2.7.1, uv lock clean, aarch64 fully installable — and got twelve lines of torch==None not in uv.lock, plus remediation text pointing at markers that are perfectly fine (scripts/check_aarch64_lock.py:137-140).

It takes a layout change to reach (all three extras coming off the WHL indexes) and it fails loudly rather than silently, so I don't consider it blocking. But it's a few lines while the file is open:

if dep.get("version") is not None:
    package = packages.get((dep["name"], dep["version"], str(dep.get("source"))))
else:  # uv omits version/source when a package resolves to a single version
    same_name = [p for p in lock["package"] if p["name"] == dep["name"]]
    package = same_name[0] if len(same_name) == 1 else None

Nits

• --with packaging is unpinned, so the verdict depends on whatever packaging releases next — and parse_wheel_filename strictness is exactly what has tightened across its releases. --with 'packaging<26' restores hermeticity for about zero cost, and it's the last non-hermetic edge left in the check.
• Only resolved[0] is inspected when several dependency entries match the same aarch64 environment. uv shouldn't emit overlapping markers, so this is theoretical; checking all matches is a one-word change.

Attacks that held up

The abi3 comparison correctly rejects a cp312 abi3 wheel for py3.11; --no-project is present, so uv run doesn't sync the project just to run the check; the changed-files gate covers both uv.lock/pyproject.toml and the script itself; deriving the Python versions from the lock means the eventual 3.13 bump extends coverage for free; and the root is None / empty-conflicts / encoding-changed guards all fail closed with actionable messages rather than passing vacuously.

The stated known gap — this asserts torch/torchvision, not that every other dependency has aarch64 wheels — is the honest framing, and a follow-up issue is the right home for it. Please do open one; I'll take that over an ARM runner for now.

---

To clear this review: the unquote fix in §1. §2 and the nits are yours to take or leave.

[stellarfeline]

All four items taken, in 55734b999e:

• §1 fixed — unquote() on the filename, plus InvalidWheelFilename now exits with a one-line message naming the file instead of a traceback. Reproduced your repro first (cpu extra's aarch64 torch repointed at 2.7.1+cpu): old script dies with the exact traceback you quoted; fixed script prints torch==2.7.1+cpu from https://download.pytorch.org/whl/cpu and aarch64 OK, and a deliberately mangled filename yields the clean one-liner.

• §2 taken, with one addition to your sketch: the reported version/registry are backfilled from the package entry when the dependency omits them, so the output shows torch==2.7.1 rather than torch==None on such lockfiles. Verified against a synthetic single-version lock: previously twelve torch==None not in uv.lock lines, now aarch64 OK.

• Both nits taken — --with 'packaging<26' (verdicts verified identical under 25.x and 26.0), and every dependency entry matching the aarch64 environment is checked, not just the first.

Also re-ran the §1/§2 regression locks from the last round against the changed script — both still fail with the same diagnostics, so nothing regressed.

On the follow-up for full-dependency coverage: I'm going to leave that one on the table rather than open an issue I won't be driving — the gap is documented in the script's docstring if anyone wants to pick it up. Thanks for two genuinely rigorous reviews; the check is much better for them.

[lstein]

Re-reviewed at c8ec9fb8b0. Everything from the last round is fixed, and I re-ran the fixtures rather than reading the diff for it:

check	result
§1 repro (cpu extra's aarch64 torch repointed at 2.7.1+cpu)	[cpu] py3.11: torch==2.7.1+cpu from https://download.pytorch.org/whl/cpu … aarch64 OK — previous script died on InvalidWheelFilename at the same input
deliberately mangled wheel filename	one-line unparseable wheel filename in uv.lock (...), exit 1 — no traceback
§2 repro (single-version lock, dep entries with no version/source)	aarch64 OK with versions backfilled — previous script emitted torchvision==None not in uv.lock
environments narrowed to x86_64	exit 1, uv.lock excludes linux/aarch64 entirely
aarch64 fallback pins deleted	exit 1, per-extra no torch resolves on linux/aarch64
aarch64 wheels stripped from the torch entry	exit 1, has no linux/aarch64 wheel
current main's real uv.lock	aarch64 OK, exit 0

Backfilling version/registry from the package entry rather than printing None is a better call than my sketch — the success line stays informative on those lockfiles instead of merely not-failing. Both nits are in too: I confirmed packaging 26.3 and packaging<26 give identical verdicts on the real lock and on the §1 repro, so the cap is insurance rather than a load-bearing pin, which is the right shape for it.

I also confirmed the check is really running at this head rather than passing by being skipped — the uv-lock-checks log shows all twelve [extra] pyX.Y lines and aarch64 OK in 0.4s, off the aarch64check changed-files group (uv.lock itself is untouched by this PR, so the uv lock --locked step correctly skips).

On the merge in c8ec9fb8b0: I checked it isn't hiding a resolution mistake. The merged tree's pyproject.toml and uv.lock are byte-identical to main at 22d83498ed, and the net diff against main is exactly the two new files — nothing from main got clobbered on the way through.

Approving. One non-blocking find below, plus a nit; neither needs another round from me.

---

Conflicting dependency groups crash the extras discovery [non-blocking]

scripts/check_aarch64_lock.py:117 assumes every entry in the lock's conflicts table has an extra key:

extras = sorted({e["extra"] for group in lock.get("conflicts", []) for e in group if e["package"] == root_name})

tool.uv.conflicts also accepts dependency groups, and uv serialises those into the same table without an extra key. Built with the pinned uv 0.6.10:

conflicts = [[
    { package = "invokeai", extra = "cpu" },
    { package = "invokeai", extra = "cuda" },
], [
    { package = "invokeai", group = "ga" },
    { package = "invokeai", group = "gb" },
]]

Against that lockfile the script raises KeyError: 'extra' — a raw traceback again, in the same class as the two you just fixed. Triggering sequence is a pyproject.toml change rather than a lockfile regression: someone adds a pair of conflicting dependency groups (a plausible enough thing to want alongside the torch extras), uv lock succeeds, and this check dies without a verdict.

Filtering instead of subscripting is enough:

extras = sorted({
    e["extra"] for group in lock.get("conflicts", []) for e in group
    if e.get("package") == root_name and "extra" in e
})

Verified: with it, the group-conflicts lockfile reaches the encoding guard and exits with the intended message; main's lock still prints aarch64 OK, and the environments-narrowed regression still exits 1.

Nit

• has_aarch64_wheel accepts any platform tag containing aarch64, which includes musllinux_*_aarch64. The docstring is careful to exclude macOS arm64 as "a different platform", and musl is one too — a lockfile whose only aarch64 wheel were musl-only would pass while glibc aarch64 (what the Docker image and every install doc target) has nothing to install. Entirely theoretical for torch, and a tag.platform.startswith(("manylinux", "linux")) guard alongside the existing check closes it if you think it's worth the line.

Attacks that held up

The version-less lookup falls back to the single package entry and to None when the name is ambiguous, so it can't silently pick the wrong one; iterating all resolved entries doesn't double-report on the real lock (uv emits disjoint markers); the InvalidWheelFilename handler exits rather than continuing, so an unparseable name can't masquerade as "no aarch64 wheel"; the abi3 comparison still rejects a cp312 abi3 wheel for py3.11; supported_python_versions still derives 3.11, 3.12 from the lock's own requires-python, so the eventual 3.13 bump widens coverage for free; and scripts/ is inside ruff's scope, so the new file is linted by python-checks rather than sitting outside CI.

Nice piece of work — this is a materially better guard than what it replaced, and the two rounds of fixes landed with their own repros each time.