Fix ROCm/gfx pin case normalization, ROCm-tag requirement, and CUDA-leaf classification

Normalize torch-index leaves to lowercase before the gfx*/rocm*/cu* allowlist
matches so the canonical gfx120X-all (capital X) gets the torch 2.11 floor in
install.sh (leaf, flavor and repairable helpers). Require an installed +rocm
local tag before a rocmX.Y or non-2.11 gfx pin is judged satisfied in
setup.ps1 Get-RocmPinStaleTags and the Python _rocm_pin_family_mismatch, so an
untagged CPU/CUDA wheel never leaves the pin unapplied. Classify a leaf as CUDA
only via ^cu[0-9]: the Python _TORCH_BACKEND derivation now uses
_is_cuda_family_leaf, and install.sh brands cuda only on cu[0-9]* (unset on an
unknown /current /custom mirror leaf) so the stack probes the GPU instead of
skipping ROCm repair. Add bash, Python and PowerShell tests for capital
gfx120X-all floor, current/custom not-cuda, and untagged-wheel ROCm pins.
This commit is contained in:
Daniel Han 2026-07-06 01:38:07 +00:00
commit 3b17c9bb0b
8 changed files with 325 additions and 114 deletions

View file

@ -47,9 +47,9 @@ class TestNoTorchBackendAutoInInstallSh:
def test_fallback_uses_torch_backend_auto(self):
"""The fallback branch should use --torch-backend=auto as recovery."""
text = INSTALL_SH.read_text(encoding = "utf-8")
assert (
"GPU detection failed" in text
), "install.sh should have a fallback branch for when GPU detection fails"
assert "GPU detection failed" in text, (
"install.sh should have a fallback branch for when GPU detection fails"
)
class TestInstallShHasGpuDetection:
@ -57,15 +57,15 @@ class TestInstallShHasGpuDetection:
def test_function_exists(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
assert (
"get_torch_index_url()" in text
), "install.sh is missing the get_torch_index_url() function"
assert "get_torch_index_url()" in text, (
"install.sh is missing the get_torch_index_url() function"
)
def test_torch_index_url_assigned(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
assert (
"TORCH_INDEX_URL=$(get_torch_index_url)" in text
), "install.sh should assign TORCH_INDEX_URL from get_torch_index_url()"
assert "TORCH_INDEX_URL=$(get_torch_index_url)" in text, (
"install.sh should assign TORCH_INDEX_URL from get_torch_index_url()"
)
class TestCudaMappingParity:
@ -133,15 +133,15 @@ class TestPyTorchMirrorEnvVar:
def test_install_sh_has_mirror_var(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
assert (
"UNSLOTH_PYTORCH_MIRROR" in text
), "install.sh should reference UNSLOTH_PYTORCH_MIRROR"
assert "UNSLOTH_PYTORCH_MIRROR" in text, (
"install.sh should reference UNSLOTH_PYTORCH_MIRROR"
)
def test_install_ps1_has_mirror_var(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
assert (
"UNSLOTH_PYTORCH_MIRROR" in text
), "install.ps1 should reference UNSLOTH_PYTORCH_MIRROR"
assert "UNSLOTH_PYTORCH_MIRROR" in text, (
"install.ps1 should reference UNSLOTH_PYTORCH_MIRROR"
)
class TestUvBytecodeCompileTimeout:
@ -167,21 +167,21 @@ class TestUvBytecodeCompileTimeout:
def test_install_sh_preserves_timeout_override(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
assert (
': "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"' in text
), "install.sh should default UV_COMPILE_BYTECODE_TIMEOUT without overwriting callers"
assert (
"export UV_COMPILE_BYTECODE_TIMEOUT" in text
), "install.sh should export UV_COMPILE_BYTECODE_TIMEOUT for uv subprocesses"
assert ': "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"' in text, (
"install.sh should default UV_COMPILE_BYTECODE_TIMEOUT without overwriting callers"
)
assert "export UV_COMPILE_BYTECODE_TIMEOUT" in text, (
"install.sh should export UV_COMPILE_BYTECODE_TIMEOUT for uv subprocesses"
)
def test_install_ps1_preserves_timeout_override(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
assert (
"if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT)" in text
), "install.ps1 should preserve caller UV_COMPILE_BYTECODE_TIMEOUT overrides"
assert (
'$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text
), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT"
assert "if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT)" in text, (
"install.ps1 should preserve caller UV_COMPILE_BYTECODE_TIMEOUT overrides"
)
assert '$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text, (
"install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT"
)
class TestTorchIndexOverrideParity:
@ -207,9 +207,9 @@ class TestTorchIndexOverrideParity:
# The AMD ROCm reroute must be skipped when the index is explicitly pinned,
# so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten.
text = path.read_text(encoding = "utf-8")
assert (
"TorchIndexPinned" in text
), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
assert "TorchIndexPinned" in text, (
f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
)
def test_cuda_pin_overrides_cvd_hide_gate(self):
# A pinned cu* index skips ALL host-GPU probing (parity with install.sh's
@ -226,9 +226,9 @@ class TestTorchIndexOverrideParity:
"_ensure_cuda_torch should compute a CUDA-pin flag so the pin can "
"override the CVD hide gate"
)
assert re.search(
r"if not _cuda_pinned and _cvd is not None", body
), "the CVD hide gate must be bypassed when a CUDA index is pinned"
assert re.search(r"if not _cuda_pinned and _cvd is not None", body), (
"the CVD hide gate must be bypassed when a CUDA index is pinned"
)
def test_cpu_repair_pins_supported_torch_range(self):
# The explicit-CPU repair must not install a bare torch trio: the /cpu
@ -281,12 +281,12 @@ class TestGfx211AllowlistParity:
# 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 "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)"
@ -294,9 +294,9 @@ class TestGfx211AllowlistParity:
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"
assert '"gfx120x-all", "gfx1151", "gfx1150"' in text, (
"install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed"
)
class TestCudaLeafDigitParity:
@ -307,29 +307,52 @@ class TestCudaLeafDigitParity:
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]"
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"
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"
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"
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*"
)
def test_install_sh_backend_export_requires_cu_digit(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The UNSLOTH_TORCH_BACKEND export must brand CUDA only on cu[0-9]* -- a
# bare catch-all *) -> cuda would mis-brand /current, /custom mirror pins
# as CUDA and make the stack skip ROCm repair on AMD hosts (comment #2's
# bug via install.sh instead of standalone studio update).
assert re.search(r'cu\[0-9\]\*\)\s*export UNSLOTH_TORCH_BACKEND="cuda"', text), (
"install.sh backend export must brand cuda only on cu[0-9]*"
)
# An unknown leaf must NOT commit a cuda backend (it unsets instead).
assert re.search(r"\*\)\s*unset UNSLOTH_TORCH_BACKEND", text), (
"install.sh backend export must unset (not force cuda) on an unknown leaf"
)
def test_install_sh_lowercases_backend_leaf(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The leaf feeding both the backend case and the 2.11 floor case must be
# lowercased so the canonical gfx120X-all (capital X) matches.
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*"
r"_torch_index_leaf=\$\(printf '%s' \"\$_torch_index_leaf\" \| tr '\[:upper:\]' '\[:lower:\]'\)",
text,
), "install.sh must lowercase _torch_index_leaf before the gfx/rocm/cu case matches"