From 00781b7695ff9aaf5da3fbdcbaa11715abe514e3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 13:42:17 +0000 Subject: [PATCH] install: honor pinned torch index over CVD/GPU gates and fix leaf-based ROCm classification - install_python_stack.py: an explicit cu* pin now clears the CUDA_VISIBLE_DEVICES empty/-1 hide gate as well as the NVIDIA-presence gate, so CVD=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update repairs to CUDA wheels (parity with install.sh's get_torch_index_url override, which skips all GPU probing). Unpinned CVD=-1 still skips. - install_python_stack.py: _ensure_cpu_torch installs the bounded _CPU_TORCH_PKG_SPEC instead of a bare torch/torchvision/torchaudio trio; the /cpu index now also serves torch 2.11+, which is outside the supported <2.11 range. - install.sh: the torch>=2.11 constraint case matches the index leaf (rocm7.2|gfx*) instead of the whole URL, so a mirror base path containing a gfx/rocm7.2 segment with a cu*/cpu family is not false-matched onto the 2.11 line. - setup.ps1: the stale-venv check expects rocm torch only for arches the install path maps to a repo.amd.com wheel index; an unmapped/unreadable arch installs CPU, so a correct CPU venv is no longer marked stale. - Tests for each of the above. --- install.sh | 9 +++- studio/install_python_stack.py | 35 +++++++++++----- studio/setup.ps1 | 19 ++++++++- tests/python/test_cross_platform_parity.py | 49 ++++++++++++++++++++++ tests/studio/install/test_cuda_repair.py | 16 +++++++ tests/studio/install/test_rocm_support.py | 17 ++++++++ 6 files changed, 132 insertions(+), 13 deletions(-) diff --git a/install.sh b/install.sh index 84e40e2fb9..e07cf22c30 100755 --- a/install.sh +++ b/install.sh @@ -2482,8 +2482,13 @@ esac # torchvision/torchaudio independently and a bare name can resolve a 2.12-built # wheel (ABI mismatch). Matches setup.ps1's *FloorMap and _ROCM_TORCH_PKG_SPECS. # All other ROCm tags and CUDA stay within <2.11.0. -case "$TORCH_INDEX_URL" in - */rocm7.2|*/gfx*) +# Match on the FINAL path segment ($_torch_index_leaf, computed above), NOT the +# whole URL: a custom UNSLOTH_PYTORCH_MIRROR whose base path contains a "gfx" or +# "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. +case "$_torch_index_leaf" in + rocm7.2|gfx*) 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" diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 4b50fece75..f9c99eb8cb 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -115,6 +115,12 @@ _CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = ( "torchaudio>=2.4,<2.11.0", ) +# CPU torch repair specs (see _ensure_cpu_torch). Same bounds/reasoning as the +# CUDA spec above: the /cpu index now also publishes torch 2.11+, so a bare trio +# from the exclusive --index-url would resolve outside the repo's supported +# <2.11 range (and can pick a torchvision built against a different torch major). +_CPU_TORCH_PKG_SPEC: tuple[str, str, str] = _CUDA_TORCH_PKG_SPEC + # torchao's C++ extensions are built against ONE exact torch release; a newer # torch makes torchao skip its cpp kernels ("Skipping import of cpp extensions # due to incompatible torch version ...") and fall back to slow Python. Because @@ -1126,18 +1132,23 @@ def _ensure_cuda_torch() -> None: # Never undo a deliberate ROCm install (setup.ps1 sets this marker). if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": return + # An explicit CUDA wheel-index pin (headless / container / CI cross-install) + # commits to CUDA wheels and, like install.sh's get_torch_index_url override, + # skips ALL host-GPU probing -- so it must clear BOTH the CUDA_VISIBLE_DEVICES + # hide gate and the NVIDIA-presence gate below, not just the latter. Otherwise + # `CUDA_VISIBLE_DEVICES=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update` + # (the exact GPU-less CI case this override targets) would still bail here. + _cuda_pinned = _explicit_cuda_torch_index_url() is not None # CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU (for # example a mixed AMD+NVIDIA host that runs ROCm torch on the AMD card); - # never force CUDA wheels over that choice. + # never force CUDA wheels over that choice unless a CUDA index is pinned. _cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if _cvd is not None and _cvd.strip() in ("", "-1"): + if not _cuda_pinned and _cvd is not None and _cvd.strip() in ("", "-1"): return # Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu() # covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent. - # An explicit CUDA wheel-index pin (headless / container / CI cross-install) - # commits to CUDA wheels regardless of whether a GPU is visible here, so it - # overrides the GPU-presence gate. - if not _has_usable_nvidia_gpu() and _explicit_cuda_torch_index_url() is None: + # The explicit CUDA pin overrides the GPU-presence gate too. + if not _cuda_pinned and not _has_usable_nvidia_gpu(): return # Classify the installed torch: "hip" (ROCm build -- the poisoning @@ -1267,14 +1278,18 @@ def _ensure_cpu_torch() -> None: f" torch is a GPU build but an explicit CPU index is pinned -- " f"reinstalling CPU torch from {pin}" ) - # The pytorch.org /cpu index is curated, so a bare trio resolves consistently. + # Pin to the supported torch<2.11 family (same bounds as the CUDA/ROCm repair + # specs). The /cpu index now also serves torch 2.11+, so a bare trio off the + # exclusive --index-url could resolve outside the supported range or drag in + # an ABI-mismatched torchvision/torchaudio. + _torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC pip_install( "CPU torch repair", "--force-reinstall", "--no-cache-dir", - "torch", - "torchvision", - "torchaudio", + _torch_pkg, + _vision_pkg, + _audio_pkg, "--index-url", pin, constrain = False, diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 617f7f8d31..c80ea18d4f 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2618,7 +2618,24 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode # false -- name-inferred Adrenalin hosts still get ROCm torch below.) # Without this an unpinned ROCm venv compares "rocm" != "cpu" and is # needlessly rebuilt, and an installer-managed setup exits as stale. - $expectedTorchTag = "rocm" + # But only the arches the install path below maps to a repo.amd.com + # wheel index (the $archFamilyMap set) actually get ROCm torch; an + # unmapped/unknown arch (e.g. name-inferred RDNA 2 gfx103X) or an + # unreadable arch -- even with $HasROCm -- falls back to CPU torch + # there (see the $archFamily null branch below). For those, expect + # "cpu" so a correct CPU venv is not marked stale and rebuilt on every + # update (or aborted under installer-managed setup). + $_rocmWheelArches = @( + "gfx1201", "gfx1200", # RDNA 4 + "gfx1151", "gfx1150", # RDNA 3.5 (Strix Halo/Point) + "gfx1103", "gfx1102", "gfx1101", "gfx1100", # RDNA 3 + "gfx90a", "gfx908" # MI200 / MI100 + ) + if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) { + $expectedTorchTag = "rocm" + } else { + $expectedTorchTag = "cpu" + } } else { $expectedTorchTag = "cpu" } diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index ae1e247634..a2674a52a3 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -210,3 +210,52 @@ class TestTorchIndexOverrideParity: 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 + # get_torch_index_url override), so the Python CUDA repair must let the pin + # clear the CUDA_VISIBLE_DEVICES hide gate, not just the NVIDIA-presence + # gate. Otherwise CVD=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update + # (the GPU-less CI case) would bail before repairing. + text = STACK_PY.read_text(encoding = "utf-8") + m = re.search( + r"def _ensure_cuda_torch\(\).*?(?=\ndef )", text, re.DOTALL + ) + assert m, "could not locate _ensure_cuda_torch" + body = m.group(0) + # The CVD hide-gate return must be guarded by the CUDA-pin flag. + assert "_cuda_pinned" in body, ( + "_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" + + def test_cpu_repair_pins_supported_torch_range(self): + # The explicit-CPU repair must not install a bare torch trio: the /cpu + # index now also serves torch 2.11+, so a bare install off the exclusive + # --index-url can resolve outside the repo's supported <2.11 range or pull + # an ABI-mismatched companion. It must use the bounded CPU/CUDA spec. + text = STACK_PY.read_text(encoding = "utf-8") + m = re.search( + r"def _ensure_cpu_torch\(\).*?(?=\ndef )", text, re.DOTALL + ) + assert m, "could not locate _ensure_cpu_torch" + body = m.group(0) + assert "_CPU_TORCH_PKG_SPEC" in body, ( + "_ensure_cpu_torch should install the bounded _CPU_TORCH_PKG_SPEC, " + "not a bare torch/torchvision/torchaudio trio" + ) + + def test_setup_ps1_stale_check_gates_rocm_on_supported_arch(self): + # The stale-venv check must only expect ROCm torch for arches the install + # path actually maps to a repo.amd.com wheel index. An unmapped arch + # (name-inferred RDNA 2 gfx103X) or an unreadable arch installs CPU torch, + # so expecting "rocm" there marks a correct CPU venv stale and rebuilds it + # every update (or aborts under installer-managed setup). + text = SETUP_PS1.read_text(encoding = "utf-8") + assert "_rocmWheelArches" in text, ( + "setup.ps1 stale check should restrict the ROCm expected-tag to the " + "supported gfx wheel arches" + ) diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index 38e9764ea4..b311e00bb7 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -148,6 +148,22 @@ class TestCudaRepairFires: assert mock_pip.call_count == 1 assert "cu128" in _index_url(mock_pip) + def test_cvd_hidden_but_explicit_cuda_pin_repairs(self): + # CUDA_VISIBLE_DEVICES=-1 hides the GPU, but an explicit cu* pin skips ALL + # host-GPU probing (like install.sh's get_torch_index_url override), so the + # CVD hide gate must not suppress the repair. This is the exact GPU-less CI + # case the override targets: CVD=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128. + for _cvd in ("-1", ""): + mock_pip = _run_cuda_repair( + nvidia = False, + backend = "cuda", + cvd = _cvd, + index_family = "cu128", + torch_state = "hip", + ) + assert mock_pip.call_count == 1 + assert "cu128" in _index_url(mock_pip) + # No-op cases. diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 8312300c6d..c276599291 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2876,6 +2876,23 @@ class TestStrixRocm71Override: source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") assert "TORCH_CONSTRAINT" in source and "2.11" in source + def test_torch_constraint_211_matches_leaf_not_whole_url(self): + """The 2.11 constraint case must match the index LEAF, not the whole URL. + + A custom UNSLOTH_PYTORCH_MIRROR whose base path contains a gfx/rocm7.2 + segment (e.g. https://mirror.local/gfx-cache) with a cu*/cpu family must + not be pushed to the torch 2.11 line -- same leaf-only reasoning the + UNSLOTH_TORCH_BACKEND classification uses. + """ + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + # The 2.11 constraint block must switch on $_torch_index_leaf, not on the + # full $TORCH_INDEX_URL (which the earlier, buggy version matched with + # */gfx* and would false-positive on a mirror base path). + assert 'case "$_torch_index_leaf" in\n rocm7.2|gfx*)' in source, ( + "the torch>=2.11 constraint must match the index leaf (rocm7.2|gfx*), " + "not the whole URL" + ) + def test_amd_rocm_mirror_env_var_respected(self): """install.sh must honour UNSLOTH_AMD_ROCM_MIRROR for air-gapped installs.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")