fix: Stop add_cpu_torch_hashes failing on a universal torch split

`make lock-python-dependencies-all` aborts before it writes a single requirements file: ValueError: Missing CPU hashes for torch==2.14.1+cpu make: *** [lock-python-dependencies-all] Error 1 `uv pip compile --universal --torch-backend cpu` splits torch by marker, so the lock it produces holds two pins per package: torch==2.14.1 ; sys_platform == 'darwin' torch==2.14.1+cpu ; sys_platform != 'darwin' `add_hashes` looked each entry up by appending `+cpu` to its version. For the second pin that asks for `2.14.1+cpu+cpu`, which cannot exist, so the lookup missed and the script raised. The `+cpu` pin is already the CPU wheel and so needs nothing added, and it is now returned untouched. Returning it rather than normalising the lookup key keeps its hash set exactly as the resolver wrote it, so this cannot alter any artifact a lock already accepts. The existing tests passed because their input carries only a bare `torch==2.13.0`, never the split a real `--universal` resolve emits. Two tests are added for that shape, including the torch and torchvision pair the lock actually contains, and both assert idempotency. Checked against all three committed `py3.*-ci-requirements.txt`: each now runs clean and is idempotent, where each raised before. The genuine mismatch case still raises, which the untouched third test covers. Signed-off-by: hao-xu5 <hxu44@apple.com>
haoxu0 · Oct 7, 2026 · Oct 8, 2026 · Oct 8, 2026 · Oct 7, 2026 · 0140f54cb4a33b61ad698036364431c9857b5fd0
commit 0140f54cb4a33b61ad698036364431c9857b5fd0
@@ -22,6 +22,17 @@ def add_cpu_hashes(requirements: str, cpu_requirements: str) -> str:
 
     def add_hashes(match: re.Match[str]) -> str:
         name, version = match[1], match[2]
+        if version.endswith("+cpu"):
+            # A universal resolve splits torch by marker, so the lock already
+            # holds a `+cpu` pin of its own alongside the PyPI one:
+            #
+            #   torch==2.14.1 ; sys_platform == 'darwin'
+            #   torch==2.14.1+cpu ; sys_platform != 'darwin'
+            #
+            # That entry is the CPU wheel, so it needs nothing added. Appending
+            # `+cpu` to its version to look it up would ask for `2.14.1+cpu+cpu`
+            # and fail, which is why this returns the entry untouched.
+            return match[0]
         hashes = cpu_hashes.get((name, f"{version}+cpu"))
         if not hashes:
             raise ValueError(f"Missing CPU hashes for {name}=={version}")

@@ -34,3 +34,60 @@ def test_cpu_hashes_reject_different_package_version() -> None:
     cpu_requirements = f"torch==2.12.0+cpu \\\n    --hash=sha256:{'b' * 64}\n"
     with pytest.raises(ValueError, match="Missing CPU hashes for torch==2.13.0"):
         add_hashes(requirements, cpu_requirements)
+
+
+def test_cpu_hashes_leave_an_existing_cpu_pin_untouched() -> None:
+    """A universal resolve emits a `+cpu` pin of its own, which needs nothing.
+
+    `uv pip compile --universal --torch-backend cpu` splits torch by marker, so
+    the lock holds both a PyPI pin for darwin and a `+cpu` pin for everything
+    else. The `+cpu` entry is already the CPU wheel; looking it up by appending
+    `+cpu` to its version would ask for `2.13.0+cpu+cpu` and fail.
+    """
+    script = (
+        Path(__file__).resolve().parents[6] / "infra/scripts/add_cpu_torch_hashes.py"
+    )
+    add_hashes = runpy.run_path(str(script))["add_cpu_hashes"]
+    pypi_hash = "a" * 64
+    cpu_hash = "b" * 64
+    requirements = (
+        f"torch==2.13.0 ; sys_platform == 'darwin' \\\n"
+        f"    --hash=sha256:{pypi_hash}\n"
+        f"torch==2.13.0+cpu ; sys_platform != 'darwin' \\\n"
+        f"    --hash=sha256:{cpu_hash}\n"
+    )
+    cpu_requirements = f"torch==2.13.0+cpu \\\n    --hash=sha256:{cpu_hash}\n"
+    expected = (
+        f"torch==2.13.0 ; sys_platform == 'darwin' \\\n"
+        f"    --hash=sha256:{pypi_hash} \\\n"
+        f"    --hash=sha256:{cpu_hash}\n"
+        f"torch==2.13.0+cpu ; sys_platform != 'darwin' \\\n"
+        f"    --hash=sha256:{cpu_hash}\n"
+    )
+    assert add_hashes(requirements, cpu_requirements) == expected
+    assert add_hashes(expected, cpu_requirements) == expected
+
+
+def test_cpu_hashes_handle_a_marker_split_for_torch_and_torchvision() -> None:
+    """The real lock splits both packages, which is what broke the script."""
+    script = (
+        Path(__file__).resolve().parents[6] / "infra/scripts/add_cpu_torch_hashes.py"
+    )
+    add_hashes = runpy.run_path(str(script))["add_cpu_hashes"]
+    requirements = (
+        f"torch==2.14.1 ; sys_platform == 'darwin' \\\n    --hash=sha256:{'a' * 64}\n"
+        f"torch==2.14.1+cpu ; sys_platform != 'darwin' \\\n"
+        f"    --hash=sha256:{'b' * 64}\n"
+        f"torchvision==0.29.1 ; sys_platform == 'darwin' \\\n"
+        f"    --hash=sha256:{'c' * 64}\n"
+        f"torchvision==0.29.1+cpu ; sys_platform != 'darwin' \\\n"
+        f"    --hash=sha256:{'d' * 64}\n"
+    )
+    cpu_requirements = (
+        f"torch==2.14.1+cpu \\\n    --hash=sha256:{'b' * 64}\n"
+        f"torchvision==0.29.1+cpu \\\n    --hash=sha256:{'d' * 64}\n"
+    )
+    result = add_hashes(requirements, cpu_requirements)
+    assert f"    --hash=sha256:{'b' * 64}" in result.split("torch==2.14.1+cpu")[0]
+    assert f"    --hash=sha256:{'d' * 64}" in result.split("torchvision==0.29.1+cpu")[0]
+    assert result.count("torch==2.14.1+cpu ; sys_platform != 'darwin'") == 1