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}"
|
||||
|
|
|
|||
92
tests/sh/test_torch_index_marker.sh
Executable file
92
tests/sh/test_torch_index_marker.sh
Executable file
|
|
@ -0,0 +1,92 @@
|
|||
#!/bin/bash
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
# Unit tests for install.sh's torch-index MARKER helpers (_normalize_index_url,
|
||||
# _write_torch_index_marker). These converge the ROCm/gfx pin-change detection
|
||||
# across install.sh / install_python_stack.py / setup.ps1 / install.ps1 by
|
||||
# recording the resolved wheel --index-url so a later update can compare exactly.
|
||||
# Helpers are extracted from install.sh and sourced (parity with test_torch_flavor.sh).
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
INSTALL_SH="$SCRIPT_DIR/../../install.sh"
|
||||
PASS=0
|
||||
FAIL=0
|
||||
|
||||
_TORCH_INDEX_MARKER_NAME=".unsloth-torch-index"
|
||||
|
||||
# Extract the marker helpers from install.sh and source them.
|
||||
_FUNC_FILE=$(mktemp)
|
||||
{
|
||||
sed -n '/^_normalize_index_url()/,/^}/p' "$INSTALL_SH"
|
||||
echo ""
|
||||
sed -n '/^_write_torch_index_marker()/,/^}/p' "$INSTALL_SH"
|
||||
} > "$_FUNC_FILE"
|
||||
# shellcheck disable=SC1090
|
||||
. "$_FUNC_FILE"
|
||||
rm -f "$_FUNC_FILE"
|
||||
|
||||
assert_eq() {
|
||||
_label="$1"; _expected="$2"; _actual="$3"
|
||||
if [ "$_actual" = "$_expected" ]; then
|
||||
echo " PASS: $_label"; PASS=$((PASS + 1))
|
||||
else
|
||||
echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1))
|
||||
fi
|
||||
}
|
||||
|
||||
echo "=== _normalize_index_url ==="
|
||||
assert_eq "trailing slashes stripped + leaf lowered" \
|
||||
"https://repo.amd.com/rocm/whl/gfx120x-all" \
|
||||
"$(_normalize_index_url 'https://repo.amd.com/rocm/whl/gfx120X-all///')"
|
||||
assert_eq "whitespace trimmed" \
|
||||
"https://download.pytorch.org/whl/cu128" \
|
||||
"$(_normalize_index_url ' https://download.pytorch.org/whl/cu128 ')"
|
||||
assert_eq "host case preserved, only leaf lowered" \
|
||||
"https://Mirror.Local/simple" \
|
||||
"$(_normalize_index_url 'https://Mirror.Local/Simple/')"
|
||||
# gfx120X-all (capital X) and AMD's lowercase pip leaf normalise equal.
|
||||
assert_eq "gfx120X-all == gfx120x-all after normalize" \
|
||||
"$(_normalize_index_url 'https://repo.amd.com/rocm/whl/gfx120x-all')" \
|
||||
"$(_normalize_index_url 'https://repo.amd.com/rocm/whl/gfx120X-all')"
|
||||
assert_eq "empty -> empty" "" "$(_normalize_index_url ' ')"
|
||||
# rocm7.2 KNOWN-2.11 leaf normalises to itself.
|
||||
assert_eq "rocm7.2 unchanged" \
|
||||
"https://download.pytorch.org/whl/rocm7.2" \
|
||||
"$(_normalize_index_url 'https://download.pytorch.org/whl/rocm7.2/')"
|
||||
|
||||
echo "=== _write_torch_index_marker ==="
|
||||
_VD=$(mktemp -d)
|
||||
_write_torch_index_marker "$_VD" "https://download.pytorch.org/whl/rocm7.2"
|
||||
assert_eq "marker written verbatim (single line)" \
|
||||
"https://download.pytorch.org/whl/rocm7.2" \
|
||||
"$(cat "$_VD/$_TORCH_INDEX_MARKER_NAME" 2>/dev/null)"
|
||||
|
||||
# Overwrite (per-arch switch gfx1151 -> gfx120X-all): the marker is replaced.
|
||||
_write_torch_index_marker "$_VD" "https://repo.amd.com/rocm/whl/gfx120X-all"
|
||||
assert_eq "marker overwritten on re-install" \
|
||||
"https://repo.amd.com/rocm/whl/gfx120X-all" \
|
||||
"$(cat "$_VD/$_TORCH_INDEX_MARKER_NAME" 2>/dev/null)"
|
||||
|
||||
# Blank URL is ignored (nothing meaningful to record) -- existing marker kept.
|
||||
_write_torch_index_marker "$_VD" " "
|
||||
assert_eq "blank url leaves prior marker intact" \
|
||||
"https://repo.amd.com/rocm/whl/gfx120X-all" \
|
||||
"$(cat "$_VD/$_TORCH_INDEX_MARKER_NAME" 2>/dev/null)"
|
||||
|
||||
# No stray temp files left behind by the atomic write.
|
||||
_leftover=$(find "$_VD" -maxdepth 1 -name "$_TORCH_INDEX_MARKER_NAME.*.tmp" 2>/dev/null | wc -l | tr -d ' ')
|
||||
assert_eq "no stray temp file left" "0" "$_leftover"
|
||||
|
||||
# Missing venv dir -> no-op, no error, no file created.
|
||||
_MISSING="$_VD/does_not_exist_dir"
|
||||
_write_torch_index_marker "$_MISSING" "https://x/cu128"
|
||||
assert_eq "missing venv dir -> no marker" \
|
||||
"absent" \
|
||||
"$([ -e "$_MISSING/$_TORCH_INDEX_MARKER_NAME" ] && echo present || echo absent)"
|
||||
|
||||
rm -rf "$_VD"
|
||||
|
||||
echo ""
|
||||
echo "Results: $PASS passed, $FAIL failed"
|
||||
[ "$FAIL" -eq 0 ]
|
||||
|
|
@ -87,6 +87,13 @@ def _run_cuda_repair(
|
|||
return smi_path
|
||||
return None
|
||||
|
||||
import tempfile as _tempfile
|
||||
|
||||
# Isolate the torch-index marker so _ensure_cuda_torch's post-reinstall marker
|
||||
# write lands in a throwaway dir, never the real venv (and never leaks state).
|
||||
_marker_dir = _tempfile.mkdtemp(prefix = "unsloth-marker-test-")
|
||||
_marker_path = Path(_marker_dir) / ".unsloth-torch-index"
|
||||
|
||||
with (
|
||||
patch.object(stack_mod, "_TORCH_BACKEND", backend),
|
||||
patch.object(stack_mod, "IS_MACOS", is_macos),
|
||||
|
|
@ -95,6 +102,7 @@ def _run_cuda_repair(
|
|||
patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = nvidia),
|
||||
patch.object(stack_mod.shutil, "which", side_effect = _which),
|
||||
patch.object(stack_mod.os.path, "isfile", return_value = bool(smi_path)),
|
||||
patch.object(stack_mod, "_torch_index_marker_path", return_value = _marker_path),
|
||||
patch.object(stack_mod, "pip_install") as mock_pip,
|
||||
patch.object(
|
||||
stack_mod.subprocess,
|
||||
|
|
@ -351,12 +359,12 @@ class TestTorchBackendDerivationFromPin:
|
|||
def test_source_uses_helper_not_bare_startswith(self):
|
||||
# Guard against a regression back to elif _idx_leaf.startswith("cu").
|
||||
src = _STACK_PATH.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"elif _is_cuda_family_leaf(_idx_leaf):" in src
|
||||
), "_TORCH_BACKEND derivation must classify CUDA via _is_cuda_family_leaf"
|
||||
assert (
|
||||
'elif _idx_leaf.startswith("cu"):' not in src
|
||||
), "_TORCH_BACKEND derivation must not use a bare startswith('cu')"
|
||||
assert "elif _is_cuda_family_leaf(_idx_leaf):" in src, (
|
||||
"_TORCH_BACKEND derivation must classify CUDA via _is_cuda_family_leaf"
|
||||
)
|
||||
assert 'elif _idx_leaf.startswith("cu"):' not in src, (
|
||||
"_TORCH_BACKEND derivation must not use a bare startswith('cu')"
|
||||
)
|
||||
|
||||
|
||||
# CUDA index ladder.
|
||||
|
|
|
|||
|
|
@ -557,6 +557,17 @@ class TestDetectRocmVersion:
|
|||
class TestEnsureRocmTorch:
|
||||
"""Verify ROCm torch reinstall logic."""
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _isolate_torch_index_marker(self, tmp_path):
|
||||
"""Point the torch-index marker at a per-test tmp path so _ensure_rocm_torch's
|
||||
marker writes never touch the real venv and stale markers never leak between
|
||||
tests. The file is ABSENT by default -> these tests exercise the no-marker
|
||||
fallback to the +rocm/version-tag heuristic (backward compatibility). The
|
||||
path is exposed as self._marker_path for tests that seed a marker."""
|
||||
self._marker_path = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = self._marker_path):
|
||||
yield
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
def test_no_rocm_skips(self, mock_nvidia, mock_pip):
|
||||
|
|
@ -814,9 +825,9 @@ class TestEnsureRocmTorch:
|
|||
_args = [str(a) for a in _call.args]
|
||||
if "--index-url" in _args:
|
||||
_url = _args[_args.index("--index-url") + 1]
|
||||
assert "rocm7.2" not in _url or "torch" not in " ".join(
|
||||
_args
|
||||
), "torch must not be reinstalled when the pin already matches"
|
||||
assert "rocm7.2" not in _url or "torch" not in " ".join(_args), (
|
||||
"torch must not be reinstalled when the pin already matches"
|
||||
)
|
||||
# A torch reinstall would pass torch>=... as a positional; assert none did.
|
||||
assert not any(
|
||||
any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
|
||||
|
|
@ -993,6 +1004,267 @@ class TestEnsureRocmTorch:
|
|||
mock_pip.assert_not_called()
|
||||
|
||||
|
||||
# TEST: install_python_stack.py -- torch-index MARKER mechanism (PR #6692)
|
||||
|
||||
|
||||
class TestTorchIndexMarkerHelpers:
|
||||
"""Pure marker helpers: normalization, read/write round-trip, exact compare."""
|
||||
|
||||
def test_normalize_index_url_lowercases_only_leaf(self):
|
||||
f = stack_mod._normalize_index_url
|
||||
# Trailing slashes stripped; ONLY the final segment lowercased.
|
||||
assert f("https://repo.amd.com/rocm/whl/gfx120X-all///") == (
|
||||
"https://repo.amd.com/rocm/whl/gfx120x-all"
|
||||
)
|
||||
# Host case preserved (only the leaf is lowered).
|
||||
assert f("https://Mirror.Local/Simple/") == "https://Mirror.Local/simple"
|
||||
# Whitespace trimmed.
|
||||
assert f(" https://download.pytorch.org/whl/cu128 ") == (
|
||||
"https://download.pytorch.org/whl/cu128"
|
||||
)
|
||||
# gfx120X-all (capital X) and AMD's lowercase pip leaf compare equal.
|
||||
assert f("https://repo.amd.com/rocm/whl/gfx120X-all") == (
|
||||
f("https://repo.amd.com/rocm/whl/gfx120x-all")
|
||||
)
|
||||
# Empty / whitespace-only -> None.
|
||||
assert f(" ") is None
|
||||
assert f(None) is None
|
||||
|
||||
def test_marker_write_read_round_trip(self, tmp_path):
|
||||
marker = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker):
|
||||
stack_mod._write_torch_index_marker("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
# Recorded verbatim (single stripped line).
|
||||
assert marker.read_text().strip() == ("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
assert stack_mod._read_torch_index_marker() == ("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
|
||||
def test_marker_write_is_atomic_and_ignores_blank(self, tmp_path):
|
||||
marker = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker):
|
||||
# Blank / whitespace-only is ignored (nothing to record).
|
||||
stack_mod._write_torch_index_marker("")
|
||||
stack_mod._write_torch_index_marker(" ")
|
||||
assert not marker.exists()
|
||||
# No stray temp files left behind by the atomic write.
|
||||
stack_mod._write_torch_index_marker("https://x/cu128")
|
||||
leftovers = [p.name for p in tmp_path.iterdir() if p.name != ".unsloth-torch-index"]
|
||||
assert leftovers == []
|
||||
|
||||
def test_missing_or_corrupt_marker_reads_as_absent(self, tmp_path):
|
||||
marker = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker):
|
||||
# Absent.
|
||||
assert stack_mod._read_torch_index_marker() is None
|
||||
# Empty file -> absent.
|
||||
marker.write_text("")
|
||||
assert stack_mod._read_torch_index_marker() is None
|
||||
# Whitespace-only -> absent.
|
||||
marker.write_text(" \n")
|
||||
assert stack_mod._read_torch_index_marker() is None
|
||||
|
||||
def test_marker_pin_mismatch_exact_compare(self, tmp_path):
|
||||
marker = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker):
|
||||
# No marker -> None (caller falls back to the heuristic).
|
||||
assert stack_mod._marker_pin_mismatch("https://repo.amd.com/rocm/whl/gfx1151") is None
|
||||
# Marker gfx1151, pin gfx120X-all -> mismatch (True). This is the exact
|
||||
# per-arch switch the version-tag heuristic cannot see (#2543).
|
||||
marker.write_text("https://repo.amd.com/rocm/whl/gfx1151\n")
|
||||
assert (
|
||||
stack_mod._marker_pin_mismatch("https://repo.amd.com/rocm/whl/gfx120X-all") is True
|
||||
)
|
||||
# Marker matches the pin (case-insensitive leaf, trailing slash) -> False.
|
||||
assert stack_mod._marker_pin_mismatch("https://repo.amd.com/rocm/whl/gfx1151/") is False
|
||||
|
||||
def test_known_211_versions_only_rocm72(self):
|
||||
# KNOWN-2.11 rocm set is exactly {rocm7.2} -- rocm7.1/rocm7.3 are NOT in it.
|
||||
known = stack_mod._ROCM_KNOWN_TORCH211_VERSIONS
|
||||
assert (7, 2) in known
|
||||
assert (7, 1) not in known
|
||||
assert (7, 3) not in known
|
||||
assert (8, 0) not in known
|
||||
|
||||
def test_rocm_pin_family_mismatch_known_211_set(self):
|
||||
# The rocmX.Y unreadable-installed fallback uses the KNOWN-2.11 set, so a
|
||||
# (hypothetical) rocm7.3 pin is NOT treated as the 2.11 line speculatively.
|
||||
f = stack_mod._rocm_pin_family_mismatch
|
||||
base = "https://download.pytorch.org/whl"
|
||||
# rocm7.3 pin (unknown -> <2.11 line) over an unreadable-version +rocm wheel
|
||||
# that is <2.11: not a mismatch (both on the non-2.11 line).
|
||||
assert f(f"{base}/rocm7.3", "2.10.0+rocm") is False
|
||||
# rocm7.2 pin (KNOWN-2.11) over the same <2.11 unreadable wheel: mismatch.
|
||||
assert f(f"{base}/rocm7.2", "2.10.0+rocm") is True
|
||||
|
||||
|
||||
class TestEnsureRocmTorchMarker:
|
||||
"""_ensure_rocm_torch marker integration (the #2543 per-arch switch fix)."""
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _marker(self, tmp_path):
|
||||
self._marker_path = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = self._marker_path):
|
||||
yield
|
||||
|
||||
def _seed(self, url):
|
||||
self._marker_path.write_text(url + "\n")
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install_try", return_value = True)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
|
||||
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
|
||||
def test_marker_gfx_switch_reinstalls_despite_same_wheel_tag(
|
||||
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
|
||||
):
|
||||
"""Marker records gfx1151; user re-pins to gfx120X-all. Both indexes install
|
||||
a +rocm7.13.0 per-arch wheel, so the version-tag heuristic sees NO difference
|
||||
and would leave the old arch in place. The marker's exact compare catches it
|
||||
and reinstalls from the new gfx index (#2543)."""
|
||||
self._seed("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
# has_hip_torch True + an already-installed AMD per-arch (three-part) wheel:
|
||||
# _rocm_pin_family_mismatch(gfx120X-all, ...+rocm7.13.0) would return False.
|
||||
mock_probe.stdout = b"7.13.0|2.11.0+rocm7.13.0\n"
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx120X-all"}
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
with patch("os.path.isdir", return_value = True):
|
||||
with patch("subprocess.run", return_value = mock_probe):
|
||||
with patch.object(
|
||||
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
|
||||
):
|
||||
_ensure_rocm_torch()
|
||||
torch_call = str(mock_pip.call_args_list[0])
|
||||
assert "gfx120X-all" in torch_call or "gfx120x-all" in torch_call
|
||||
# Marker rewritten to the new index.
|
||||
assert "gfx120X-all" in self._marker_path.read_text()
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install_try", return_value = True)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
|
||||
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
|
||||
def test_marker_matches_pin_no_reinstall(
|
||||
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
|
||||
):
|
||||
"""Marker matches the current pin exactly -> NO torch reinstall (no loop),
|
||||
even when the wheel tag would otherwise look ambiguous."""
|
||||
self._seed("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
mock_probe.stdout = b"7.13.0|2.11.0+rocm7.13.0\n"
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"}
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
with patch("os.path.isdir", return_value = True):
|
||||
with patch("subprocess.run", return_value = mock_probe):
|
||||
with patch.object(
|
||||
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
|
||||
):
|
||||
_ensure_rocm_torch()
|
||||
# No torch reinstall (marker matches). bnb-only calls may still happen.
|
||||
assert not any(
|
||||
any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
|
||||
)
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install_try", return_value = True)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
|
||||
@patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4))
|
||||
def test_no_marker_falls_back_to_heuristic_no_forced_reinstall(
|
||||
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
|
||||
):
|
||||
"""NO marker + a healthy ROCm wheel that already satisfies the pin -> the
|
||||
heuristic (_rocm_pin_family_mismatch) decides, and a correct venv is NOT
|
||||
force-reinstalled (backward compatibility for old venvs)."""
|
||||
# No marker seeded (autouse fixture leaves the file absent).
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
# rocm6.4 pin over an installed +rocm6.4 wheel -> heuristic says no mismatch.
|
||||
mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n"
|
||||
env = {"UNSLOTH_TORCH_INDEX_FAMILY": "rocm6.4"}
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
|
||||
with patch("os.path.isdir", return_value = True):
|
||||
with patch("subprocess.run", return_value = mock_probe):
|
||||
_ensure_rocm_torch()
|
||||
assert not any(
|
||||
any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
|
||||
)
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_verbatim_custom_url_reinstall_on_marker_mismatch(self, mock_pip):
|
||||
"""Marker records .../simple; user pins a custom .../current index (leaf is
|
||||
neither rocm/gfx/cu/cpu). _ensure_verbatim_torch_index reinstalls torch
|
||||
VERBATIM from that URL -- "URL wins verbatim" (#2544)."""
|
||||
self._seed("https://mirror.local/simple")
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/current"}
|
||||
with patch.object(stack_mod, "NO_TORCH", False):
|
||||
with patch.object(stack_mod, "IS_MACOS", False):
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
stack_mod._ensure_verbatim_torch_index()
|
||||
assert mock_pip.call_count == 1
|
||||
call = str(mock_pip.call_args_list[0])
|
||||
assert "https://mirror.local/current" in call
|
||||
assert "torch" in call
|
||||
# Marker rewritten to the pinned custom URL.
|
||||
assert "current" in self._marker_path.read_text()
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_verbatim_custom_url_no_marker_is_noop(self, mock_pip):
|
||||
"""No marker + a custom-URL pin -> _ensure_verbatim_torch_index does NOTHING
|
||||
(an old venv must not be blindly force-reinstalled from an unverified URL)."""
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/current"}
|
||||
with patch.object(stack_mod, "NO_TORCH", False):
|
||||
with patch.object(stack_mod, "IS_MACOS", False):
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
stack_mod._ensure_verbatim_torch_index()
|
||||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_verbatim_custom_url_matching_marker_no_reinstall(self, mock_pip):
|
||||
"""Marker matches the custom URL pin -> no reinstall (idempotent, no loop)."""
|
||||
self._seed("https://mirror.local/current")
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/current/"}
|
||||
with patch.object(stack_mod, "NO_TORCH", False):
|
||||
with patch.object(stack_mod, "IS_MACOS", False):
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
stack_mod._ensure_verbatim_torch_index()
|
||||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_verbatim_skips_known_family_pins(self, mock_pip):
|
||||
"""A known-family pin (rocm/gfx/cu/cpu) is NOT handled by the verbatim path --
|
||||
the dedicated _ensure_{rocm,cuda,cpu} helpers own those. cu128 stays CUDA."""
|
||||
self._seed("https://download.pytorch.org/whl/cu126")
|
||||
for pin in (
|
||||
"https://download.pytorch.org/whl/cu128",
|
||||
"https://download.pytorch.org/whl/rocm7.2",
|
||||
"https://repo.amd.com/rocm/whl/gfx1151",
|
||||
"https://download.pytorch.org/whl/cpu",
|
||||
):
|
||||
mock_pip.reset_mock()
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": pin}
|
||||
with patch.object(stack_mod, "NO_TORCH", False):
|
||||
with patch.object(stack_mod, "IS_MACOS", False):
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
stack_mod._ensure_verbatim_torch_index()
|
||||
mock_pip.assert_not_called()
|
||||
|
||||
|
||||
# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard
|
||||
|
||||
|
||||
|
|
@ -1022,9 +1294,9 @@ class TestHasRocmGpuKfdVendorGuard:
|
|||
|
||||
src = self._src()
|
||||
# Word boundary so "vendor_id 41098" doesn't match "vendor_id 4098".
|
||||
assert (
|
||||
_re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src
|
||||
), "_has_rocm_gpu vendor_id check should use word boundary anchors"
|
||||
assert _re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src, (
|
||||
"_has_rocm_gpu vendor_id check should use word boundary anchors"
|
||||
)
|
||||
|
||||
def test_sysfs_fallback_guarded_by_non_win32(self):
|
||||
"""KFD sysfs fallback must be Linux-only (guarded by sys.platform != 'win32')."""
|
||||
|
|
@ -1034,9 +1306,9 @@ class TestHasRocmGpuKfdVendorGuard:
|
|||
def test_cpu_node_excluded(self):
|
||||
"""gpu_id == '0' must be excluded (CPU topology nodes)."""
|
||||
src = self._src()
|
||||
assert (
|
||||
'!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src
|
||||
), "_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)"
|
||||
assert '!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src, (
|
||||
"_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)"
|
||||
)
|
||||
|
||||
def test_install_sh_has_vendor_check(self):
|
||||
"""_has_amd_rocm_gpu in install.sh sysfs fallback must also check vendor_id 4098."""
|
||||
|
|
@ -1069,12 +1341,12 @@ class TestHasRocmGpuKfdVendorGuard:
|
|||
func_start = source.find("_has_amd_rocm_gpu()")
|
||||
func_end = source.find("\n}", func_start)
|
||||
func_body = source[func_start:func_end]
|
||||
assert (
|
||||
"_has_usable_nvidia_gpu" in func_body
|
||||
), "_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts"
|
||||
assert (
|
||||
"return 1" in func_body
|
||||
), "_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected"
|
||||
assert "_has_usable_nvidia_gpu" in func_body, (
|
||||
"_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts"
|
||||
)
|
||||
assert "return 1" in func_body, (
|
||||
"_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected"
|
||||
)
|
||||
|
||||
def test_has_usable_nvidia_gpu_proc_fallback_present(self):
|
||||
"""`_has_usable_nvidia_gpu` must have a /proc/driver/nvidia fallback."""
|
||||
|
|
@ -1290,12 +1562,12 @@ class TestInstallShStructure:
|
|||
rocm_call = body.find("_has_amd_rocm_gpu")
|
||||
assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu"
|
||||
assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia branch"
|
||||
assert (
|
||||
rocm_call > no_nvidia_branch
|
||||
), "ROCm detection should sit inside the 'no NVIDIA' branch"
|
||||
assert (
|
||||
nvidia_call < no_nvidia_branch
|
||||
), "NVIDIA detection should run before the no-NVIDIA branch"
|
||||
assert rocm_call > no_nvidia_branch, (
|
||||
"ROCm detection should sit inside the 'no NVIDIA' branch"
|
||||
)
|
||||
assert nvidia_call < no_nvidia_branch, (
|
||||
"NVIDIA detection should run before the no-NVIDIA branch"
|
||||
)
|
||||
|
||||
def test_bitsandbytes_amd_install(self):
|
||||
"""install.sh should install bitsandbytes for AMD when ROCm detected."""
|
||||
|
|
@ -1361,9 +1633,9 @@ class TestInstallShStructure:
|
|||
stripped = line.lstrip()
|
||||
if stripped.startswith("#"):
|
||||
continue
|
||||
assert (
|
||||
"((" not in line or "))" not in line or "$(()" in line
|
||||
), f"get_torch_index_url line {i} may use non-POSIX (( ))"
|
||||
assert "((" not in line or "))" not in line or "$(()" in line, (
|
||||
f"get_torch_index_url line {i} may use non-POSIX (( ))"
|
||||
)
|
||||
|
||||
def test_macos_returns_cpu_before_rocm_check(self):
|
||||
"""macOS should return CPU immediately (before any ROCm check)."""
|
||||
|
|
@ -1382,9 +1654,9 @@ class TestInstallShStructure:
|
|||
torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)")
|
||||
backend_pos = source.find("UNSLOTH_TORCH_BACKEND")
|
||||
assert backend_pos > 0, "UNSLOTH_TORCH_BACKEND must be set in install.sh"
|
||||
assert (
|
||||
backend_pos > torch_url_pos
|
||||
), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved"
|
||||
assert backend_pos > torch_url_pos, (
|
||||
"UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved"
|
||||
)
|
||||
assert '"cuda"' in source[backend_pos : backend_pos + 500]
|
||||
assert '"rocm"' in source[backend_pos : backend_pos + 500]
|
||||
assert '"cpu"' in source[backend_pos : backend_pos + 500]
|
||||
|
|
@ -1398,12 +1670,12 @@ class TestInstallShStructure:
|
|||
func_start = source.find("_has_amd_rocm_gpu()")
|
||||
func_end = source.find("\n}", func_start)
|
||||
func_body = source[func_start:func_end]
|
||||
assert (
|
||||
"vendor_id" in func_body
|
||||
), "_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes"
|
||||
assert (
|
||||
"4098" in func_body
|
||||
), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)"
|
||||
assert "vendor_id" in func_body, (
|
||||
"_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes"
|
||||
)
|
||||
assert "4098" in func_body, (
|
||||
"_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)"
|
||||
)
|
||||
|
||||
def test_kfd_awk_resets_state_per_file(self):
|
||||
"""KFD sysfs awk must reset gpu/amd state per file (FNR==1) to avoid Ryzen+NVIDIA false positives."""
|
||||
|
|
@ -1428,9 +1700,9 @@ class TestInstallShStructure:
|
|||
"get_torch_index_url must use a _nvidia_detected flag (separate from "
|
||||
"_smi) so that proc-only NVIDIA detection still selects CUDA wheels"
|
||||
)
|
||||
assert (
|
||||
'_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body
|
||||
), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1"
|
||||
assert '_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body, (
|
||||
"get_torch_index_url AMD branch must be skipped when _nvidia_detected=1"
|
||||
)
|
||||
|
||||
|
||||
# TEST: Live regression on current host (NVIDIA B200 expected)
|
||||
|
|
@ -2419,9 +2691,9 @@ class TestRuntimeBnbRocmSourceGuards:
|
|||
"""A failed redetect must not downgrade a persisted suffix to '72'."""
|
||||
for path in (self._MAIN_PATH, self._TRAINING_WORKER_PATH):
|
||||
source = path.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
'_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"' in source
|
||||
), path.name
|
||||
assert '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"' in source, (
|
||||
path.name
|
||||
)
|
||||
|
||||
def test_main_requires_found_rocm_dll(self):
|
||||
"""HIP_PATH/ROCM_PATH alone (HIP SDK on a CUDA/CPU box) must not force
|
||||
|
|
@ -2870,9 +3142,9 @@ class TestStrixHaloGfxArchDetection:
|
|||
"""Both files must use the gfx\\d+[a-z]? regex to parse arch from amd-smi output."""
|
||||
for path in (_SETUP_PS1_PATH, _INSTALL_PS1_PATH):
|
||||
source = path.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"gfx\\d+" in source or r"gfx\d+" in source
|
||||
), f"gfx arch regex not found in {path.name}"
|
||||
assert "gfx\\d+" in source or r"gfx\d+" in source, (
|
||||
f"gfx arch regex not found in {path.name}"
|
||||
)
|
||||
|
||||
|
||||
# TEST: HIP SDK tool path resolution via HIP_PATH / ROCM_PATH env vars
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ $tokens = $null; $errors = $null
|
|||
$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors)
|
||||
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" }
|
||||
|
||||
foreach ($name in @("Test-RocmGfx211Leaf", "Test-CudaFamilyLeaf", "Get-RocmPinStaleTags")) {
|
||||
foreach ($name in @("Test-RocmGfx211Leaf", "Test-RocmKnown211Version", "Test-CudaFamilyLeaf", "Get-RocmPinStaleTags")) {
|
||||
$fn = $ast.FindAll({ param($n)
|
||||
$n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
|
||||
}, $true)
|
||||
|
|
@ -85,6 +85,17 @@ Check "gfx90a pin + 2.10.0 (untagged) -> stale" (IsStale "gfx90a" "2.10.0
|
|||
Check "gfx120x-all pin + 2.11.0+rocm7.2 (generic) -> stale" (IsStale "gfx120x-all" "2.11.0+rocm7.2")
|
||||
Check "gfx120x-all pin + 2.10.0 (untagged) -> stale" (IsStale "gfx120x-all" "2.10.0")
|
||||
|
||||
Write-Host "Test-RocmKnown211Version + KNOWN-2.11 fallback (rocm7.2 only; no speculative rocm7.3)"
|
||||
Check "rocm7.2 -> known 2.11" (Test-RocmKnown211Version -Major 7 -Minor 2)
|
||||
Check "rocm7.1 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 1))
|
||||
Check "rocm7.3 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 3))
|
||||
Check "rocm8.0 -> not known" (-not (Test-RocmKnown211Version -Major 8 -Minor 0))
|
||||
# Unreadable-installed fallback: a rocm7.3 pin (unknown -> <2.11 line) over a <2.11
|
||||
# +rocm wheel with an unreadable version is NOT stale (both on the non-2.11 line);
|
||||
# rocm7.2 (KNOWN-2.11) over the same wheel IS stale. This is the #2534 alignment.
|
||||
Check "rocm7.3 pin + 2.10.0+rocm (unreadable ver) -> not stale" (-not (IsStale "rocm7.3" "2.10.0+rocm"))
|
||||
Check "rocm7.2 pin + 2.10.0+rocm (unreadable ver) -> stale" (IsStale "rocm7.2" "2.10.0+rocm")
|
||||
|
||||
Write-Host ""
|
||||
if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
|
||||
Write-Host "All checks passed" -ForegroundColor Green
|
||||
|
|
|
|||
101
tests/studio/test_torch_index_marker.ps1
Normal file
101
tests/studio/test_torch_index_marker.ps1
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
#!/usr/bin/env pwsh
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
# Unit tests for studio/setup.ps1's torch-index MARKER helpers
|
||||
# (Get-NormalizedIndexUrl, Read-TorchIndexMarker, Write-TorchIndexMarker,
|
||||
# Test-MarkerPinMismatch, Test-RocmKnown211Version). These converge the ROCm/gfx
|
||||
# pin-change detection across install.sh / install_python_stack.py / setup.ps1 /
|
||||
# install.ps1. Pure helpers, AST-extracted and run in-process -- no GPU/venv.
|
||||
# Run: pwsh -NoProfile -File tests/studio/test_torch_index_marker.ps1
|
||||
|
||||
$ErrorActionPreference = "Stop"
|
||||
$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1")
|
||||
$setupPath = (Resolve-Path $setupPath).Path
|
||||
|
||||
# Parse setup.ps1 (also a syntax gate) and extract the helpers.
|
||||
$tokens = $null; $errors = $null
|
||||
$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors)
|
||||
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" }
|
||||
|
||||
# The marker filename constant the helpers close over.
|
||||
$TorchIndexMarkerName = ".unsloth-torch-index"
|
||||
|
||||
foreach ($name in @(
|
||||
"Get-NormalizedIndexUrl", "Get-TorchIndexMarkerPath", "Read-TorchIndexMarker",
|
||||
"Write-TorchIndexMarker", "Test-MarkerPinMismatch", "Test-RocmKnown211Version"
|
||||
)) {
|
||||
$fn = $ast.FindAll({ param($n)
|
||||
$n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
|
||||
}, $true)
|
||||
if ($fn.Count -ne 1) { throw "expected exactly one $name in setup.ps1, found $($fn.Count)" }
|
||||
Invoke-Expression $fn[0].Extent.Text
|
||||
}
|
||||
|
||||
$failures = 0
|
||||
function Check($name, $cond) {
|
||||
if ($cond) { Write-Host " PASS $name" }
|
||||
else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ }
|
||||
}
|
||||
|
||||
Write-Host "Get-NormalizedIndexUrl (trim / strip trailing slash / lowercase leaf)"
|
||||
Check "trailing slashes + leaf lowered" `
|
||||
((Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all///") -eq "https://repo.amd.com/rocm/whl/gfx120x-all")
|
||||
Check "whitespace trimmed" `
|
||||
((Get-NormalizedIndexUrl " https://download.pytorch.org/whl/cu128 ") -eq "https://download.pytorch.org/whl/cu128")
|
||||
Check "host case preserved, leaf lowered" `
|
||||
((Get-NormalizedIndexUrl "https://Mirror.Local/Simple/") -eq "https://Mirror.Local/simple")
|
||||
Check "gfx120X-all == gfx120x-all after normalize" `
|
||||
((Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all") -eq (Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120x-all"))
|
||||
Check "empty -> null" ($null -eq (Get-NormalizedIndexUrl " "))
|
||||
|
||||
Write-Host "Write/Read-TorchIndexMarker (round trip, atomic, blank ignored)"
|
||||
$venv = Join-Path ([System.IO.Path]::GetTempPath()) ("unsloth-marker-" + [System.Guid]::NewGuid().ToString("N"))
|
||||
New-Item -ItemType Directory -Path $venv | Out-Null
|
||||
try {
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://repo.amd.com/rocm/whl/gfx1151"
|
||||
Check "marker written verbatim" `
|
||||
((Read-TorchIndexMarker -VenvDir $venv) -eq "https://repo.amd.com/rocm/whl/gfx1151")
|
||||
# Per-arch switch overwrites the marker.
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all"
|
||||
Check "marker overwritten on re-install" `
|
||||
((Read-TorchIndexMarker -VenvDir $venv) -eq "https://repo.amd.com/rocm/whl/gfx120X-all")
|
||||
# Blank URL ignored -- prior marker kept.
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl " "
|
||||
Check "blank url leaves prior marker intact" `
|
||||
((Read-TorchIndexMarker -VenvDir $venv) -eq "https://repo.amd.com/rocm/whl/gfx120X-all")
|
||||
# No stray temp file left behind.
|
||||
$tmpLeft = @(Get-ChildItem -LiteralPath $venv -Filter "$TorchIndexMarkerName.*.tmp" -ErrorAction SilentlyContinue).Count
|
||||
Check "no stray temp file left" ($tmpLeft -eq 0)
|
||||
|
||||
Write-Host "Test-MarkerPinMismatch (exact compare; null when no marker)"
|
||||
# Marker gfx120X-all now recorded. A gfx1151 pin differs -> mismatch (#2543).
|
||||
Check "gfx marker vs gfx1151 pin -> mismatch" `
|
||||
((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://repo.amd.com/rocm/whl/gfx1151") -eq $true)
|
||||
# Same pin (trailing slash, case) -> not a mismatch (no reinstall loop).
|
||||
Check "same pin (slash/case) -> no mismatch" `
|
||||
((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://repo.amd.com/rocm/whl/gfx120x-all/") -eq $false)
|
||||
# Custom URL change /simple -> /current is a mismatch (#2544).
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://mirror.local/simple"
|
||||
Check "custom /simple marker vs /current pin -> mismatch" `
|
||||
((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://mirror.local/current") -eq $true)
|
||||
|
||||
Write-Host "Read-TorchIndexMarker (missing / empty -> null)"
|
||||
Remove-Item -LiteralPath (Get-TorchIndexMarkerPath -VenvDir $venv) -Force
|
||||
Check "missing marker -> null" ($null -eq (Read-TorchIndexMarker -VenvDir $venv))
|
||||
Set-Content -LiteralPath (Get-TorchIndexMarkerPath -VenvDir $venv) -Value "" -NoNewline
|
||||
Check "empty marker -> null" ($null -eq (Read-TorchIndexMarker -VenvDir $venv))
|
||||
Check "no marker -> Test-MarkerPinMismatch null" `
|
||||
($null -eq (Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://x/rocm7.2"))
|
||||
} finally {
|
||||
Remove-Item -LiteralPath $venv -Recurse -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
|
||||
Write-Host "Test-RocmKnown211Version (KNOWN-2.11 set == rocm7.2 only)"
|
||||
Check "rocm7.2 -> true" (Test-RocmKnown211Version -Major 7 -Minor 2)
|
||||
Check "rocm7.1 -> false" (-not (Test-RocmKnown211Version -Major 7 -Minor 1))
|
||||
Check "rocm7.3 -> false" (-not (Test-RocmKnown211Version -Major 7 -Minor 3))
|
||||
Check "rocm8.0 -> false" (-not (Test-RocmKnown211Version -Major 8 -Minor 0))
|
||||
|
||||
Write-Host ""
|
||||
if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
|
||||
Write-Host "All checks passed" -ForegroundColor Green
|
||||
Loading…
Add table
Add a link
Reference in a new issue