diff --git a/install.sh b/install.sh index 0120197bbd..2d533ad414 100755 --- a/install.sh +++ b/install.sh @@ -2184,9 +2184,13 @@ _torch_flavor_tag() { # Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* -> # rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. +# Lowercase the leaf first so gfx120X-all (capital X) / any cased mirror leaf +# still classifies, keeping the cuXXX tag comparison against _torch_flavor_tag +# (which emits lowercase cuXXX) case-consistent. _expected_torch_flavor_tag() { _u="${1%/}" _leaf="${_u##*/}" + _leaf=$(printf '%s' "$_leaf" | tr '[:upper:]' '[:lower:]') case "$_leaf" in cu[0-9]*) echo "$_leaf" ;; cpu) echo "cpu" ;; @@ -2200,9 +2204,11 @@ _expected_torch_flavor_tag() { # resolves (torch + every transitive dep) via --index-url -- the same URLs the # fresh-install paths above already use -- so a stale wheel is auto-repairable. # Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. +# Lowercase the leaf first so a cased leaf (e.g. gfx120X-all) is recognised. _torch_index_repairable() { _u="${1%/}" _leaf="${_u##*/}" + _leaf=$(printf '%s' "$_leaf" | tr '[:upper:]' '[:lower:]') case "$_leaf" in cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;; *) echo "no" ;; @@ -2480,12 +2486,26 @@ TORCH_INDEX_URL=$(get_torch_index_url) # whose base path happens to contain "rocm" or "gfx" must not mislabel a # cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix # overrides in gfxNNNN/, so the trailing slash is stripped first). +# Lowercase the leaf once here so every gfx*/rocm*/cu* allowlist below matches +# regardless of case. The canonical AMD RDNA4 leaf is gfx120X-all (capital X, +# from the arch maps in install_python_stack.py / install.ps1); without this a +# pin to gfx120X-all would miss the lowercase gfx120x-all allowlist entries. +# The allowlists stay lowercase; inputs are normalised to lowercase. +# CUDA is branded only on a real cu[0-9]* leaf (^cu[0-9]) -- NOT a bare cu*/catch- +# all -- so a full-override mirror leaf like /current or /custom does NOT commit a +# CUDA backend. An unknown leaf leaves the backend var unset so the stack probes +# the GPU instead of returning early in _ensure_rocm_torch on AMD hosts. +# Matches _is_cuda_family_leaf (Python) / Test-CudaFamilyLeaf (PowerShell). _torch_index_leaf="${TORCH_INDEX_URL%/}" _torch_index_leaf="${_torch_index_leaf##*/}" +_torch_index_leaf=$(printf '%s' "$_torch_index_leaf" | tr '[:upper:]' '[:lower:]') case "$_torch_index_leaf" in rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; - *) export UNSLOTH_TORCH_BACKEND="cuda" ;; + cu[0-9]*) export UNSLOTH_TORCH_BACKEND="cuda" ;; + # Unknown leaf (odd mirror, e.g. /current, /custom): do NOT commit a backend. + # Unset so a stale inherited value can't leak and the stack probes the GPU. + *) unset UNSLOTH_TORCH_BACKEND ;; esac # rocm7.2 and the AMD per-gfx indexes with the torch._C._grouped_mm bug on <2.11 @@ -2505,6 +2525,8 @@ esac # "rocm7.2" segment (e.g. https://mirror.local/gfx-cache) with a cu*/cpu family # must not be treated as an AMD per-arch index and pushed to the 2.11 line. This # mirrors the leaf-only backend classification just above. +# _torch_index_leaf is already lowercased above, so the canonical gfx120X-all +# (capital X, from the arch maps) pins here via the lowercase gfx120x-all entry. case "$_torch_index_leaf" in rocm7.2|gfx120x-all|gfx1151|gfx1150) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 2ac0ad042c..3a9e653c84 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1119,6 +1119,12 @@ def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool: _inst_rocm = re.search(r"\+rocm(\d+)\.(\d+)", installed_ver) _inst_ver = (int(_inst_rocm.group(1)), int(_inst_rocm.group(2))) if _inst_rocm else None _inst_is_perarch = re.search(r"\+rocm\d+\.\d+\.\d+", installed_ver) is not None + # A ROCm build MUST carry a +rocm local tag. An untagged CPU/CUDA wheel (no + # +rocm, e.g. "2.10.0" / "2.11.0") never satisfies a ROCm pin -- always a + # mismatch -- mirroring setup.ps1's Get-RocmPinStaleTags. (In practice + # _ensure_rocm_torch only calls this when has_hip_torch is True, but keep the + # pure function correct for any input so it stays in lockstep with the PS side.) + _inst_has_rocm = re.search(r"\+rocm", installed_ver) is not None # Whether the installed torch RELEASE (before "+") is 2.11+. _inst_rel = re.match(r"^(\d+)\.(\d+)", installed_ver) _inst_is_211 = ( @@ -1135,8 +1141,9 @@ def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool: # a generic rocm wheel or any pre-2.11 build IS a mismatch even at 2.11. return not (_inst_is_211 and _inst_is_perarch) # Non-2.11 gfx leaf: install path uses default <2.11 specs, so a correct - # <2.11 wheel must stay. Mismatch only when the installed torch is 2.11+. - return _inst_is_211 + # <2.11 wheel must stay. An untagged (no +rocm) wheel never satisfies the + # pin -> mismatch; otherwise mismatch only when the installed torch is 2.11+. + return (not _inst_has_rocm) or _inst_is_211 # rocmX.Y pin. _pin_is_211 = _pin_ver >= (7, 2) if _pin_ver is not None else False @@ -1146,7 +1153,10 @@ def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool: # (7, 13) -> mismatch, which correctly reinstalls the generic wheel the # user pinned instead of leaving the per-arch one in place. return _pin_ver != _inst_ver - # rocm pin with an unreadable installed version: compare on the torch 2.11 line. + # rocm pin with an unreadable installed version: compare on the torch 2.11 line, + # but an untagged (no +rocm) wheel never satisfies a rocmX.Y pin -> mismatch. + if not _inst_has_rocm: + return True return _pin_is_211 != _inst_is_211 @@ -1707,7 +1717,7 @@ def _ensure_rocm_torch() -> None: None, ) if tag is None: - print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " f"skipping torch reinstall") + print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- skipping torch reinstall") else: if _override_idx is None: index_url = f"{_PYTORCH_WHL_BASE}/{tag}" @@ -1833,7 +1843,12 @@ if not _TORCH_BACKEND: _TORCH_BACKEND = "rocm" elif _idx_leaf == "cpu": _TORCH_BACKEND = "cpu" - elif _idx_leaf.startswith("cu"): + elif _is_cuda_family_leaf(_idx_leaf): + # Require a digit after "cu" (^cu[0-9]) so a full-override URL ending in + # /current or /custom is NOT branded CUDA. A wrong "cuda" backend makes + # _ensure_rocm_torch() return early on AMD hosts and leaves a CPU/wrong + # torch unrepaired; falling through here keeps _TORCH_BACKEND="" so the + # helpers probe the GPU instead. _TORCH_BACKEND = "cuda" diff --git a/studio/setup.ps1 b/studio/setup.ps1 index a966b82afc..b525bd98c2 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -457,6 +457,9 @@ function Test-CudaFamilyLeaf { # is stale; an already-installed per-arch wheel is NOT (no rebuild loop). # * gfx pin NOT in the allowlist (gfx110X-all/gfx90a/gfx908) -> default <2.11 # specs, so a 2.10+rocm wheel is correct; only a 2.11+ build is stale. +# A ROCm pin (gfx* or rocmX.Y) is satisfied ONLY by an installed wheel carrying a +# +rocm local tag: an untagged CPU/CUDA wheel (e.g. 2.10.0 / 2.11.0) never +# satisfies a ROCm pin, so it is reported stale and the pin is (re)applied. function Get-RocmPinStaleTags { param([string]$PinLeaf, [string]$TorchVersion) $_pinRocm = [regex]::Match($PinLeaf, '^rocm(\d+)\.(\d+)') @@ -465,6 +468,9 @@ function Get-RocmPinStaleTags { $_instRocm = [regex]::Match($TorchVersion, '\+rocm(\d+)\.(\d+)') $_instVer = if ($_instRocm.Success) { "$($_instRocm.Groups[1].Value).$($_instRocm.Groups[2].Value)" } else { $null } $_instPerArch = [regex]::IsMatch($TorchVersion, '\+rocm\d+\.\d+\.\d+') + # A ROCm build MUST carry a +rocm local tag. Without it the wheel is a CPU/CUDA + # build that cannot satisfy any ROCm pin, regardless of its release line. + $_instHasRocm = [regex]::IsMatch($TorchVersion, '\+rocm') $_instRel = [regex]::Match($TorchVersion, '^(\d+)\.(\d+)') $_instIs211 = $false if ($_instRel.Success) { @@ -474,29 +480,36 @@ function Get-RocmPinStaleTags { if ($PinLeaf -like 'gfx*') { if (Test-RocmGfx211Leaf $PinLeaf) { # Expect the AMD per-arch (three-part) 2.11 wheel. Satisfied only when - # BOTH a 2.11 release AND a three-part rocm tag are installed. + # BOTH a 2.11 release AND a three-part rocm tag are installed (the + # three-part tag already implies +rocm). $installed = if ($_instIs211 -and $_instPerArch) { "rocm-perarch(torch>=2.11)" } else { "rocm-generic-or-old" } return @{ Expected = "rocm-perarch(torch>=2.11)"; Installed = $installed } } - # Non-2.11 gfx leaf: default <2.11 spec. Stale only when the build is 2.11+. + # Non-2.11 gfx leaf: default <2.11 spec. An untagged (no +rocm) wheel never + # satisfies the pin -> stale. Otherwise stale only when the build is 2.11+. + $installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } return @{ Expected = "rocm(torch<2.11)" - Installed = if ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } + Installed = $installed } } # rocmX.Y pin. if ($_pinVer -and $_instVer) { - # Both rocm versions readable: exact comparison. + # Both rocm versions readable: exact comparison. A readable $_instVer already + # implies a +rocm tag, so no separate tag check is needed here. return @{ Expected = "rocm$_pinVer"; Installed = "rocm$_instVer" } } $_pinNeeds211 = $false if ($_pinRocm.Success) { $_pinNeeds211 = ([int]$_pinRocm.Groups[1].Value -gt 7) -or ([int]$_pinRocm.Groups[1].Value -eq 7 -and [int]$_pinRocm.Groups[2].Value -ge 2) } + # Fallback (installed rocm version unreadable): compare on the 2.11 line, but an + # untagged (no +rocm) wheel never satisfies a rocmX.Y pin -> report it stale. + $installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } return @{ Expected = if ($_pinNeeds211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } - Installed = if ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } + Installed = $installed } } diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 4b34afad4b..7fbff70cdc 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -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" diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh index 293a709360..7528c8c9db 100644 --- a/tests/sh/test_torch_constraint.sh +++ b/tests/sh/test_torch_constraint.sh @@ -271,6 +271,62 @@ bash -c " _uv_got2=$(cat "$_UV_LOG2" 2>/dev/null || echo "") assert_contains "mock uv arm64+py312 receives torch>=2.4" "$_uv_got2" "torch>=2.4,<2.11.0" +# ====================================================================== +# ROCm 2.11 floor: leaf is lowercased before the gfx*/rocm* allowlist match +# ====================================================================== +echo "" +echo "=== ROCm 2.11 floor case (leaf normalization) ===" + +# Structural: install.sh must lowercase _torch_index_leaf before the floor case, +# so the canonical AMD RDNA4 leaf gfx120X-all (capital X) matches gfx120x-all. +_has_lc=$(grep -c '_torch_index_leaf=$(printf .* | tr .\[:upper:\]. .\[:lower:\].)' "$INSTALL_SH" || true) +_has_lc_ok=$([ "$_has_lc" -ge 1 ] && echo "yes" || echo "no") +assert_eq "install.sh lowercases _torch_index_leaf" "yes" "$_has_lc_ok" + +# Runtime: replicate the exact normalization + floor case from install.sh and +# assert both gfx120X-all (capital X, canonical) and gfx120x-all get the floor, +# while non-2.11 leaves (gfx110X-all, rocm6.4, cu128, cpu) keep the default. +run_floor_case() { + _url="$1" + bash -c ' + TORCH_CONSTRAINT="torch>=2.4,<2.11.0" + TORCHVISION_CONSTRAINT="torchvision" + TORCHAUDIO_CONSTRAINT="torchaudio" + _torch_index_leaf="${1%/}" + _torch_index_leaf="${_torch_index_leaf##*/}" + _torch_index_leaf=$(printf "%s" "$_torch_index_leaf" | tr "[:upper:]" "[:lower:]") + case "$_torch_index_leaf" in + rocm7.2|gfx120x-all|gfx1151|gfx1150) + TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" + ;; + esac + echo "$TORCH_CONSTRAINT" + ' _ "$_url" +} + +assert_eq "gfx120X-all (capital) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all')" +assert_eq "gfx120X-all trailing slash -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all/')" +assert_eq "gfx120x-all (lowercase) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120x-all')" +assert_eq "gfx1151 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1151')" +assert_eq "gfx1150 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1150')" +assert_eq "rocm7.2 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://download.pytorch.org/whl/rocm7.2')" +assert_eq "gfx110X-all -> default (no floor)" "torch>=2.4,<2.11.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx110X-all')" +assert_eq "rocm6.4 -> default (no floor)" "torch>=2.4,<2.11.0" \ + "$(run_floor_case 'https://download.pytorch.org/whl/rocm6.4')" +assert_eq "cu128 -> default (no floor)" "torch>=2.4,<2.11.0" \ + "$(run_floor_case 'https://download.pytorch.org/whl/cu128')" +assert_eq "cpu -> default (no floor)" "torch>=2.4,<2.11.0" \ + "$(run_floor_case 'https://download.pytorch.org/whl/cpu')" + # ====================================================================== # Summary # ====================================================================== diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index c2781443e7..b9cfbc3cdd 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -294,6 +294,71 @@ class TestCudaRepairSkips: stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None) +class TestTorchBackendDerivationFromPin: + """The module-level _TORCH_BACKEND derivation (standalone `studio update` + with no install.sh-set UNSLOTH_TORCH_BACKEND) must classify the pinned index + leaf via _is_cuda_family_leaf (^cu[0-9]), NOT a bare startswith("cu"). A + full-override URL ending in /current or /custom must fall through to backend + "" (probe the GPU) so _ensure_rocm_torch() still repairs a wrong/CPU torch on + AMD hosts, instead of being wrongly branded "cuda" and returning early.""" + + @staticmethod + def _derive(env): + # Re-run the exact derivation the module does at import time, using the + # module's own _is_cuda_family_leaf so this stays in lockstep with it. + idx_override = ( + env.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + or env.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + ) + backend = env.get("UNSLOTH_TORCH_BACKEND", "").lower() + if not backend: + leaf = idx_override.rstrip("/").rsplit("/", 1)[-1].lower() + if leaf.startswith(("rocm", "gfx")): + backend = "rocm" + elif leaf == "cpu": + backend = "cpu" + elif stack_mod._is_cuda_family_leaf(leaf): + backend = "cuda" + return backend + + def test_cu128_pin_is_cuda(self): + assert ( + self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://download.pytorch.org/whl/cu128"}) + == "cuda" + ) + + def test_cu128_family_is_cuda(self): + assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "cu128"}) == "cuda" + + def test_current_leaf_not_cuda(self): + # ^cu[0-9] rejects /current -> backend stays "" (probe GPU), so an AMD + # host still repairs a CPU/wrong torch instead of short-circuiting. + assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/current"}) == "" + + def test_custom_leaf_not_cuda(self): + assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/custom"}) == "" + + def test_rocm_and_gfx_pins_are_rocm(self): + assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"}) == "rocm" + assert ( + self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx120X-all"}) + == "rocm" + ) + + def test_cpu_pin_is_cpu(self): + assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "cpu"}) == "cpu" + + 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')" + ) + + # CUDA index ladder. diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index f54d4415fb..aa3a5bd485 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -814,9 +814,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 @@ -916,9 +916,11 @@ class TestEnsureRocmTorch: # gfx pin (2.11 line) vs installed release line. assert f(f"{amd}/gfx1151", "2.10.0+rocm6.4") is True assert f(f"{amd}/gfx1151", "2.11.0+rocm7.13.0") is False - # rocm7.2 pin vs unreadable installed rocm version -> compare on 2.11 line. + # rocm7.2 pin vs an untagged (no +rocm) wheel: a CPU/CUDA build never + # satisfies a ROCm pin, regardless of its release line -> always a mismatch. assert f(f"{base}/rocm7.2", "2.10.0") is True - assert f(f"{base}/rocm7.2", "2.11.0") is False + assert f(f"{base}/rocm7.2", "2.11.0") is True + assert f(f"{base}/rocm6.4", "2.10.0") is True # A 2.11-allowlist gfx pin over a GENERIC (two-part +rocm7.2) 2.11 wheel is a # mismatch -- the user wants AMD's per-arch (three-part) wheel, not generic. assert f(f"{amd}/gfx1151", "2.11.0+rocm7.2") is True @@ -934,6 +936,10 @@ class TestEnsureRocmTorch: assert f(f"{amd}/gfx90a", "2.10.0+rocm6.3") is False assert f(f"{amd}/gfx908", "2.10.0+rocm7.0") is False assert f(f"{amd}/gfx110X-all", "2.11.0+rocm7.2") is True + # A non-2.11 gfx pin over an untagged (no +rocm) wheel is a mismatch even + # when torch is already <2.11: a CPU/CUDA build never satisfies the ROCm pin. + assert f(f"{amd}/gfx110X-all", "2.10.0") is True + assert f(f"{amd}/gfx90a", "2.10.0") is True @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @@ -1016,9 +1022,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').""" @@ -1028,9 +1034,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.""" @@ -1063,12 +1069,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.""" @@ -1284,12 +1290,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.""" @@ -1355,9 +1361,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).""" @@ -1376,9 +1382,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] @@ -1392,12 +1398,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.""" @@ -1422,9 +1428,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) @@ -2413,9 +2419,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 @@ -2864,9 +2870,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 diff --git a/tests/studio/test_setup_pin_stale.ps1 b/tests/studio/test_setup_pin_stale.ps1 index 6497f487d8..26c8e2382e 100644 --- a/tests/studio/test_setup_pin_stale.ps1 +++ b/tests/studio/test_setup_pin_stale.ps1 @@ -59,9 +59,11 @@ Write-Host "Get-RocmPinStaleTags (mirror of _rocm_pin_family_mismatch)" Check "rocm7.2 pin + 2.11.0+rocm7.2 -> not stale" (-not (IsStale "rocm7.2" "2.11.0+rocm7.2")) Check "rocm7.2 pin + 2.10.0+rocm6.4 -> stale" (IsStale "rocm7.2" "2.10.0+rocm6.4") Check "rocm6.4 pin + 2.10.0+rocm6.4 -> not stale" (-not (IsStale "rocm6.4" "2.10.0+rocm6.4")) -# rocm pin vs unreadable installed rocm version -> compare on the 2.11 line. -Check "rocm7.2 pin + 2.10.0 -> stale" (IsStale "rocm7.2" "2.10.0") -Check "rocm7.2 pin + 2.11.0 -> not stale" (-not (IsStale "rocm7.2" "2.11.0")) +# rocm pin vs an untagged (no +rocm) wheel: a CPU/CUDA build never satisfies a +# ROCm pin, regardless of its release line -> always stale (needs reinstall). +Check "rocm7.2 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm7.2" "2.10.0") +Check "rocm7.2 pin + 2.11.0 (untagged) -> stale" (IsStale "rocm7.2" "2.11.0") +Check "rocm6.4 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm6.4" "2.10.0") # 2.11-allowlist gfx pin: per-arch (three-part) wheel is satisfied, generic is stale. Check "gfx1151 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx1151" "2.11.0+rocm7.13.0")) Check "gfx1150 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx1150" "2.11.0+rocm7.13.0")) @@ -73,6 +75,15 @@ Check "gfx110x-all pin + 2.10.0+rocm6.4 -> not stale" (-not (IsStale "gfx110x-a Check "gfx90a pin + 2.10.0+rocm6.3 -> not stale" (-not (IsStale "gfx90a" "2.10.0+rocm6.3")) Check "gfx908 pin + 2.10.0+rocm7.0 -> not stale" (-not (IsStale "gfx908" "2.10.0+rocm7.0")) Check "gfx110x-all pin + 2.11.0+rocm7.2 -> stale" (IsStale "gfx110x-all" "2.11.0+rocm7.2") +# Non-2.11 gfx pin over an untagged (no +rocm) wheel: never satisfies the pin -> +# stale, so the explicit ROCm index is applied even when torch is already <2.11. +Check "gfx110x-all pin + 2.10.0 (untagged) -> stale" (IsStale "gfx110x-all" "2.10.0") +Check "gfx90a pin + 2.10.0 (untagged) -> stale" (IsStale "gfx90a" "2.10.0") +# Capital gfx120X-all leaf is lowercased by Get-TorchIndexLeaf before this helper; +# the caller passes the normalised leaf, so the 2.11-allowlist branch fires and a +# generic/untagged wheel is stale (the per-arch wheel is not). +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 "" if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }