install: fix second-order edge cases in pinned torch-index ROCm/CUDA handling

Parse the ROCm torch probe positionally so an empty HIP marker is kept:
CPU/CUDA torch no longer reads as HIP, so the ROCm reinstall is not skipped.
Emit one "<marker>|<version>" line (like the CUDA probe) for a robust parse.

Limit the gfx torch 2.11 expectation to the install allowlist
(gfx120X-all/gfx1151/gfx1150). A pinned gfx110X-all/gfx90a/gfx908 index stays
on the default <2.11 specs, so a correct 2.10+rocm wheel is no longer judged a
mismatch and force-reinstalled every update.

Distinguish an AMD per-arch wheel (three-part +rocmA.B.C) from a generic
pytorch.org wheel (two-part +rocmA.B): a gfx per-arch pin over a generic 2.11
wheel now reinstalls the per-arch wheel, while an already-installed per-arch
wheel is not re-flagged (no reinstall loop).

Mirror all of the above in setup.ps1 via new Test-RocmGfx211Leaf /
Test-CudaFamilyLeaf / Get-RocmPinStaleTags helpers, reused by both the
install-spec path and the stale-venv check so they cannot diverge again.
Require a digit after "cu" (^cu[0-9]) in setup.ps1, install.ps1 and install.sh
so a mirror leaf like /custom or /current is not branded CUDA and does not
rebuild the venv every run.

Add tests: CPU/CUDA probe -> has_hip_torch False; gfx110X-all pin + 2.10 wheel
not stale; gfx1151 pin + generic 2.11 wheel stale; gfx1151 pin + per-arch wheel
not stale; /custom and /current not CUDA; plus cross-language allowlist and
cu-digit parity guards, and a PowerShell unit test for the new setup.ps1 helpers.
This commit is contained in:
Daniel Han 2026-07-06 00:38:33 +00:00
commit e1f05c2c0e
7 changed files with 462 additions and 68 deletions

View file

@ -255,3 +255,81 @@ class TestTorchIndexOverrideParity:
"setup.ps1 stale check should restrict the ROCm expected-tag to the "
"supported gfx wheel arches"
)
class TestGfx211AllowlistParity:
"""The gfx per-arch leaves that carry the torch 2.11 floor (gfx120X-all /
gfx1151 / gfx1150) must be the SAME set in every installer AND in each
installer's stale/mismatch check. When these diverged, a pinned gfx110X-all /
gfx90a / gfx908 wheel (which stays <2.11) was force-reinstalled every update."""
EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150"}
def test_install_sh_allowlist(self):
text = INSTALL_SH.read_text(encoding = "utf-8").lower()
# install.sh: the TORCH_CONSTRAINT case (rocm7.2|gfx120x-all|gfx1151|gfx1150).
m = re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150", text)
assert m, "install.sh gfx-2.11 allowlist case not found / changed"
def test_install_ps1_allowlist(self):
text = INSTALL_PS1.read_text(encoding = "utf-8").lower()
m = re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text)
assert m, "install.ps1 $_pinGfx211 allowlist not found / changed"
def test_setup_ps1_defines_single_allowlist_helper(self):
# setup.ps1 must define the allowlist once (Test-RocmGfx211Leaf) and the
# install-spec path must reuse it, so the stale check and install spec can
# never disagree again.
text = SETUP_PS1.read_text(encoding = "utf-8")
assert "function Test-RocmGfx211Leaf" in text, (
"setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper"
)
assert re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower()), (
"Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist"
)
assert "$_pinGfx211 = Test-RocmGfx211Leaf" in text, (
"setup.ps1 install-spec path should reuse Test-RocmGfx211Leaf, not "
"re-hardcode the allowlist (they must not diverge)"
)
def test_stack_py_allowlist(self):
text = STACK_PY.read_text(encoding = "utf-8").lower()
assert '"gfx120x-all", "gfx1151", "gfx1150"' in text, (
"install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed"
)
class TestCudaLeafDigitParity:
"""A wheel-family leaf is CUDA only when it is "cu" + digits (cu118/cu128/...).
A bare cu* glob wrongly catches mirror leaves like /custom or /current; when
that happened the venv was marked stale and rebuilt on every run. Every
installer must require a digit after "cu" in its family/CUDA classification."""
def test_stack_py_requires_cu_digit(self):
text = STACK_PY.read_text(encoding = "utf-8")
assert re.search(r'r"\^cu\[0-9\]"', text), (
"install_python_stack.py _is_cuda_family_leaf must match ^cu[0-9]"
)
def test_setup_ps1_requires_cu_digit(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
assert re.search(r"'\^cu\[0-9\]'", text), (
"setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9], not a bare cu* glob"
)
# The stale-venv branch must go through the digit-guarded helper.
assert "Test-CudaFamilyLeaf $_pinLeaf" in text, (
"setup.ps1 stale check should classify CUDA via Test-CudaFamilyLeaf"
)
def test_install_ps1_requires_cu_digit_in_gpu_branch(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
assert re.search(r"'\^cu\[0-9\]'", text), (
"install.ps1 Get-TauriGpuBranch must require a digit after cu"
)
def test_install_sh_requires_cu_digit_in_gpu_branch(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The _tauri_gpu_branch cuda case must be cu[0-9]*, not a bare cu*.
assert re.search(r"cu\[0-9\]\*\)\s*echo \"cuda\"", text), (
"install.sh _tauri_gpu_branch cuda case must be cu[0-9]*, not cu*"
)