diff --git a/install.sh b/install.sh index e3fdf10cdd..53b09cffab 100755 --- a/install.sh +++ b/install.sh @@ -2431,10 +2431,13 @@ case "$_torch_index_leaf" in *) export UNSLOTH_TORCH_BACKEND="cuda" ;; esac -# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. +# rocm7.2 and the AMD per-gfx indexes (repo.amd.com/.../gfxNNNN) ship torch +# 2.11.0 -- adjust the constraint to allow it. This also covers a pinned full-URL +# or family override (e.g. UNSLOTH_TORCH_INDEX_URL=.../gfx1151) that returns early +# above and so never hits the Strix reroute that otherwise raises this constraint. # All other ROCm tags and CUDA stay within <2.11.0. case "$TORCH_INDEX_URL" in - */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; + */rocm7.2|*/gfx*) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; esac # Auto-detect GPU for AMD ROCm based diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index a2c1e44556..37c3c7d449 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -981,6 +981,30 @@ def _detect_cuda_torch_index_url() -> str: return f"{_PYTORCH_WHL_BASE}/{tag}" +def _explicit_torch_index_url() -> "str | None": + """The wheel index URL pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY, else None. + + Lets the CUDA/ROCm repair helpers honour the exact pinned family/URL instead + of re-probing the GPU. Mirrors install.sh::get_torch_index_url's override. + """ + url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + if url: + return url.rstrip("/") + family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + if family: + return f"{_PYTORCH_WHL_BASE}/{family.strip('/')}" + return None + + +def _explicit_rocm_torch_index_url() -> "str | None": + """The pinned wheel index URL when it names a ROCm family (rocm*/gfx*), else None.""" + url = _explicit_torch_index_url() + if url is None: + return None + leaf = url.rstrip("/").rsplit("/", 1)[-1].lower() + return url if leaf.startswith(("rocm", "gfx")) else None + + def _ensure_cuda_torch() -> None: """Repair a venv whose torch is a ROCm build on an NVIDIA host. @@ -1013,7 +1037,10 @@ def _ensure_cuda_torch() -> None: return # Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu() # covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent. - if not _has_usable_nvidia_gpu(): + # 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_torch_index_url() is None: return # Classify the installed torch: "hip" (ROCm build -- the poisoning @@ -1256,7 +1283,9 @@ def _ensure_rocm_torch() -> None: # an incompatible wheel. Use HIP_VISIBLE_DEVICES for the runtime target. _strix_override_url: "str | None" = None _strix_override_pkgs: "tuple[str, str, str] | None" = None - if ver < (7, 2): + # An explicit ROCm wheel-index pin is authoritative: never auto-reroute it to + # the AMD per-gfx index (the caller already chose the family/URL). + if ver < (7, 2) and _explicit_rocm_torch_index_url() is None: gfx_codes = _detect_amd_gfx_codes() _strix_gfx = {"gfx1151", "gfx1150"} _detected_strix = _strix_gfx.intersection(gfx_codes) @@ -1319,23 +1348,35 @@ def _ensure_rocm_torch() -> None: ) rocm_torch_ready = True elif not has_hip_torch: - # Select best matching wheel tag (newest ROCm version <= installed) - tag = next( - ( - t - for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) - if ver >= (maj, mn) - ), - None, - ) + # Honour an explicit ROCm wheel-index pin verbatim instead of re-detecting + # the host ROCm version; otherwise select the best wheel tag (newest ROCm + # version <= installed). gfx*/rocm7.2 indexes serve torch 2.11+, so match + # the constraints to the index leaf when overridden. + _override_idx = _explicit_rocm_torch_index_url() + if _override_idx is not None: + index_url = _override_idx + tag = index_url.rstrip("/").rsplit("/", 1)[-1].lower() + else: + tag = next( + ( + t + for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) + if ver >= (maj, mn) + ), + None, + ) if tag is None: print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " f"skipping torch reinstall") else: - index_url = f"{_PYTORCH_WHL_BASE}/{tag}" - print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}") - _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get( - tag, _ROCM_TORCH_PKG_SPECS["_default"] - ) + if _override_idx is None: + index_url = f"{_PYTORCH_WHL_BASE}/{tag}" + print(f" ROCm torch -- installing from {index_url}") + if tag.startswith("gfx"): + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["rocm7.2"] + else: + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get( + tag, _ROCM_TORCH_PKG_SPECS["_default"] + ) pip_install( f"ROCm torch ({tag})", "--force-reinstall", diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index cea4383268..38e9764ea4 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -64,15 +64,19 @@ def _run_cuda_repair( rocm_marker = False, smi_path = "/usr/bin/nvidia-smi", cvd = None, + index_family = None, ): """Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock. - cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it.""" + cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it. + index_family sets UNSLOTH_TORCH_INDEX_FAMILY (the explicit wheel-index pin).""" env = {} if rocm_marker: env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1" if cvd is not None: env["CUDA_VISIBLE_DEVICES"] = cvd + if index_family is not None: + env["UNSLOTH_TORCH_INDEX_FAMILY"] = index_family def _which(name, *a, **k): if name == "nvidia-smi": @@ -99,6 +103,9 @@ def _run_cuda_repair( stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None) if cvd is None: stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None) + if index_family is None: + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None) _ensure_cuda_torch() return mock_pip @@ -128,6 +135,19 @@ class TestCudaRepairFires: mock_pip = _run_cuda_repair(torch_state = "hip") assert mock_pip.call_count == 1 + def test_no_gpu_but_explicit_cuda_pin_repairs(self): + # Headless / container / CI cross-install: an explicit cu* index pin + # commits to CUDA wheels even though no NVIDIA GPU is visible here, so a + # ROCm-poisoned venv is still repaired (to the pinned family). + mock_pip = _run_cuda_repair( + nvidia = False, + backend = "cuda", + 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 4c50fa1c8e..b46967dffe 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -680,6 +680,34 @@ class TestEnsureRocmTorch: torch_call = mock_pip.call_args_list[0] assert "rocm7.2" in str(torch_call) + @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_explicit_gfx_index_honored_and_skips_strix_reroute( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """An explicit gfx wheel-index pin is authoritative: install from it verbatim + with torch 2.11, and never re-probe gfx codes to second-guess it (host ROCm 6.4 + would otherwise pick the rocm6.4 wheel / trigger the Strix re-route).""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # cpu torch -> reinstall + 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): + # Would raise if the Strix block ran (it is skipped on an explicit pin). + with patch.object(stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError): + _ensure_rocm_torch() + assert mock_pip.call_count == 1 + torch_call = str(mock_pip.call_args_list[0]) + assert "gfx1151" in torch_call + assert "torch>=2.11.0,<2.12.0" in torch_call + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install")