install: keep the torch-index marker additive to flavor validation

Three narrow fixes in the marker-based stale-venv detection:

- setup.ps1: a matching marker no longer overwrites the detected installed
  flavor. The marker compare is now an additional rebuild trigger, so a stale
  wheel (torch swapped to a +cpu build while the marker still records a cuXXX
  pin) is still caught by the flavor check instead of being masked as up to date.

- setup.ps1: a supported AMD arch carrying CPU torch is no longer marked stale
  and wiped. The downstream AMD Windows ROCm override upgrades CPU torch to ROCm
  in place, so wiping first would delete the venv and abort with "Virtual
  environment not found". Only a genuinely wrong CUDA wheel still rebuilds.

- install.sh: the Radeon --find-links path records its repo.radeon.com base in
  the marker instead of the generic pytorch.org ROCm fallback index, so a later
  pin to that generic family correctly reinstalls rather than comparing equal.
  Mirrors install.ps1/setup.ps1, which already record the real AMD index.
This commit is contained in:
danielhanchen 2026-07-06 04:22:13 +00:00
commit 6bbc692587
5 changed files with 155 additions and 137 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,32 +307,32 @@ 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*"
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")
@ -340,13 +340,13 @@ class TestCudaLeafDigitParity:
# 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]*"
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"
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")
@ -417,9 +417,9 @@ class TestKnown211SetParity:
def test_install_sh_known_211_leaf_is_rocm72_and_gfx_allowlist(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The 2.11 floor case matches exactly rocm7.2 + the three gfx leaves.
assert re.search(
r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text
), "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150"
assert re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text), (
"install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150"
)
# No speculative rocm7.3 anywhere.
assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3"