install: converge torch-index pin detection via a per-venv marker
Introduce a torch-index MARKER that records the exact wheel --index-url used
after each successful torch install, so `unsloth studio update` / repair makes
the "did the pinned index change?" decision by an EXACT string compare rather
than inferring it from the wheel +rocm/+cu version tag. The tag cannot encode
the AMD per-arch gfx family (two 2.11 gfx indexes both install +rocm7.13.0), so
the tag heuristic missed a gfx1151 -> gfx120X-all switch and a custom-URL swap.
Marker path is per-venv (.unsloth-torch-index), one line = the resolved index
URL, written atomically (temp + rename). Path, format and normalization are
shared across all four installers (install.sh, install_python_stack.py,
setup.ps1, install.ps1).
- Reapply gfx pins on a per-arch target change: the marker's exact compare
reinstalls when the pinned index differs, even when both wheels share a tag.
- Honor custom ROCm URL pins during repair: an explicit index whose leaf is not
rocm/gfx/cu/cpu (e.g. simple, current) now reinstalls torch VERBATIM from the
pin when it differs from the marker ("URL wins verbatim").
- Align the KNOWN-2.11 rocm/gfx set to exactly rocm7.2 plus the gfx allowlist
gfx120x-all/gfx1151/gfx1150 in every language; stop treating an unknown newer
rocm (rocm7.3, which does not exist) as the 2.11 line speculatively.
Backward compatible: with no marker (old venvs, torch installed out-of-band) the
existing +rocm/version-tag heuristics still decide, and a matching marker never
reinstall-loops. A cu128 CUDA pin stays a CUDA pin; custom and current leaves are
not CUDA. Adds marker tests (py/sh/ps) plus cross-installer parity checks.
This commit is contained in:
Daniel Han
parent
1213135cd1
commit
753ef2f5cf
10 changed files with 1149 additions and 125 deletions
|
|
@ -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")
|
||||
|
|
@ -356,3 +356,108 @@ class TestCudaLeafDigitParity:
|
|||
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"
|
||||
|
||||
|
||||
class TestTorchIndexMarkerParity:
|
||||
"""All four installers must agree on the torch-index marker (PR #6692):
|
||||
the same filename, the same write points, and the same read helpers."""
|
||||
|
||||
MARKER = ".unsloth-torch-index"
|
||||
|
||||
def test_all_installers_use_same_marker_filename(self):
|
||||
# The exact marker filename must appear in every installer so bash / py /
|
||||
# ps write and read the same per-venv path.
|
||||
for path, label in (
|
||||
(INSTALL_SH, "install.sh"),
|
||||
(INSTALL_PS1, "install.ps1"),
|
||||
(SETUP_PS1, "setup.ps1"),
|
||||
(STACK_PY, "install_python_stack.py"),
|
||||
):
|
||||
text = path.read_text(encoding = "utf-8")
|
||||
assert self.MARKER in text, f"{label} must reference the marker '{self.MARKER}'"
|
||||
|
||||
def test_all_installers_write_the_marker(self):
|
||||
# install.sh + PowerShell use a _write_torch_index_marker / Write-TorchIndexMarker
|
||||
# helper; the Python stack calls _write_torch_index_marker at its install sites.
|
||||
assert "_write_torch_index_marker" in INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert "Write-TorchIndexMarker" in INSTALL_PS1.read_text(encoding = "utf-8")
|
||||
assert "Write-TorchIndexMarker" in SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "_write_torch_index_marker" in STACK_PY.read_text(encoding = "utf-8")
|
||||
|
||||
def test_setup_ps1_and_python_read_the_marker(self):
|
||||
# The repair/update side (setup.ps1 stale check, python _ensure_rocm_torch)
|
||||
# must READ the marker to make the exact pin-change decision.
|
||||
assert "Read-TorchIndexMarker" in SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "Test-MarkerPinMismatch" in SETUP_PS1.read_text(encoding = "utf-8")
|
||||
stack = STACK_PY.read_text(encoding = "utf-8")
|
||||
assert "_read_torch_index_marker" in stack
|
||||
assert "_marker_pin_mismatch" in stack
|
||||
|
||||
def test_all_installers_normalize_index_url(self):
|
||||
# The exact-compare normalization (trim, strip trailing slash, lowercase
|
||||
# leaf) must exist in every language so the compare is identical.
|
||||
assert "_normalize_index_url" in INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert "Get-NormalizedIndexUrl" in SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "_normalize_index_url" in STACK_PY.read_text(encoding = "utf-8")
|
||||
|
||||
def test_marker_written_atomically(self):
|
||||
# bash uses a temp file + mv; PowerShell a temp file + Move-Item; Python
|
||||
# tempfile + os.replace. Confirm the atomic-write intent in each.
|
||||
assert re.search(r"mv -f \"\$_wm_tmp\"", INSTALL_SH.read_text(encoding = "utf-8"))
|
||||
for ps in (INSTALL_PS1, SETUP_PS1):
|
||||
assert "Move-Item" in ps.read_text(encoding = "utf-8")
|
||||
assert "os.replace" in STACK_PY.read_text(encoding = "utf-8")
|
||||
|
||||
|
||||
class TestKnown211SetParity:
|
||||
"""The KNOWN-2.11 rocm/gfx set must be identical across all four installers:
|
||||
exactly {rocm7.2} plus the gfx allowlist {gfx120x-all, gfx1151, gfx1150}.
|
||||
rocm7.3 / torch 2.12 do not exist, so no side may floor them speculatively."""
|
||||
|
||||
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"
|
||||
)
|
||||
# No speculative rocm7.3 anywhere.
|
||||
assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3"
|
||||
|
||||
def test_python_known_211_versions_is_only_rocm72(self):
|
||||
text = STACK_PY.read_text(encoding = "utf-8")
|
||||
assert "_ROCM_KNOWN_TORCH211_VERSIONS" in text
|
||||
# The frozenset literal is exactly {(7, 2)}.
|
||||
m = re.search(r"_ROCM_KNOWN_TORCH211_VERSIONS[^=]*=\s*frozenset\(\{([^}]*)\}\)", text)
|
||||
assert m is not None, "install_python_stack.py must define _ROCM_KNOWN_TORCH211_VERSIONS"
|
||||
assert "(7, 2)" in m.group(1)
|
||||
assert "7, 3" not in m.group(1) and "7, 1" not in m.group(1)
|
||||
|
||||
def test_setup_ps1_known_211_helper_is_only_rocm72(self):
|
||||
text = SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "Test-RocmKnown211Version" in text
|
||||
# The predicate is Major -eq 7 -and Minor -eq 2 (only rocm7.2).
|
||||
assert re.search(
|
||||
r"Test-RocmKnown211Version[\s\S]{0,400}\$Major -eq 7 -and \$Minor -eq 2", text
|
||||
), "setup.ps1 Test-RocmKnown211Version must accept only rocm7.2"
|
||||
|
||||
def test_install_ps1_pin_floor_is_only_rocm72(self):
|
||||
text = INSTALL_PS1.read_text(encoding = "utf-8")
|
||||
# The pinned-ROCm install-spec floor must be Major -eq 7 -and Minor -eq 2,
|
||||
# not the speculative >= 2 that would floor a non-existent rocm7.3.
|
||||
assert re.search(
|
||||
r"\$_pinRocm211 = \(\[int\]\$Matches\[1\] -eq 7 -and \[int\]\$Matches\[2\] -eq 2\)",
|
||||
text,
|
||||
), "install.ps1 pinned-ROCm floor must be rocm7.2 only (no speculative >= 2)"
|
||||
|
||||
def test_gfx_allowlist_matches_across_installers(self):
|
||||
# The gfx 2.11 allowlist {gfx120x-all, gfx1151, gfx1150} must appear in each.
|
||||
gfx = ("gfx120x-all", "gfx1151", "gfx1150")
|
||||
for path, label in (
|
||||
(INSTALL_SH, "install.sh"),
|
||||
(INSTALL_PS1, "install.ps1"),
|
||||
(SETUP_PS1, "setup.ps1"),
|
||||
(STACK_PY, "install_python_stack.py"),
|
||||
):
|
||||
low = path.read_text(encoding = "utf-8").lower()
|
||||
for g in gfx:
|
||||
assert g in low, f"{label} missing gfx 2.11 allowlist member {g}"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue