"""Tests for CUDA torch repair on poisoned NVIDIA venvs. Verifies _ensure_cuda_torch (studio/install_python_stack.py) reinstalls CUDA torch when a venv on an NVIDIA host carries a ROCm torch build (the pre-fix KFD gpu_id false positive), without touching healthy CUDA, deliberate CPU wheels, ROCm hosts, macOS, or Windows. All tests use mocks -- no GPU required. """ import importlib.util import sys from pathlib import Path from unittest.mock import MagicMock, patch import pytest # ── Load module under test (mirrors test_rocm_support.py) ──────────────────── PACKAGE_ROOT = Path(__file__).resolve().parents[3] _STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" _STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH) assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None stack_mod = importlib.util.module_from_spec(_STACK_SPEC) sys.modules[_STACK_SPEC.name] = stack_mod _STACK_SPEC.loader.exec_module(stack_mod) _ensure_cuda_torch = stack_mod._ensure_cuda_torch _detect_cuda_torch_index_url = stack_mod._detect_cuda_torch_index_url # ── Helpers ────────────────────────────────────────────────────────────────── def _make_run( torch_state = "hip", cuda_version = "12.8", torch_rc = 0, smi_rc = 0, ): """Build a subprocess.run side_effect. The torch-classify probe runs sys.executable and reads bytes stdout; the nvidia-smi version probe runs the smi path with text=True. Distinguish by the executable. """ def _run(cmd, *args, **kwargs): result = MagicMock() exe = str(cmd[0]) if cmd else "" if exe == sys.executable: result.returncode = torch_rc result.stdout = (torch_state + "\n").encode() return result # nvidia-smi version probe (text = True) result.returncode = smi_rc out = f"CUDA Version: {cuda_version}\n" if cuda_version else "No devices found\n" result.stdout = out if kwargs.get("text") else out.encode() return result return _run def _run_cuda_repair( *, backend = "", nvidia = True, torch_state = "hip", cuda_version = "12.8", torch_rc = 0, smi_rc = 0, is_macos = False, is_windows = False, no_torch = False, rocm_marker = False, smi_path = "/usr/bin/nvidia-smi", cvd = None, ): """Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock. cvd controls CUDA_VISIBLE_DEVICES: None removes it from the environment (the host machine may export one), any string sets it explicitly. """ env = {} if rocm_marker: env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1" if cvd is not None: env["CUDA_VISIBLE_DEVICES"] = cvd def _which(name, *a, **k): if name == "nvidia-smi": return smi_path return None with ( patch.object(stack_mod, "_TORCH_BACKEND", backend), patch.object(stack_mod, "IS_MACOS", is_macos), patch.object(stack_mod, "IS_WINDOWS", is_windows), patch.object(stack_mod, "NO_TORCH", no_torch), 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, "pip_install") as mock_pip, patch.object( stack_mod.subprocess, "run", side_effect = _make_run(torch_state, cuda_version, torch_rc, smi_rc), ), patch.dict(stack_mod.os.environ, env, clear = False), ): if not rocm_marker: stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None) if cvd is None: stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None) _ensure_cuda_torch() return mock_pip def _index_url(mock_pip) -> str: """Return the --index-url value from the recorded pip_install call.""" args = [str(a) for a in mock_pip.call_args.args] return args[args.index("--index-url") + 1] # ── Repair fires only on the poisoning signature ───────────────────────────── class TestCudaRepairFires: def test_hip_build_on_nvidia_triggers_repair(self): mock_pip = _run_cuda_repair(torch_state = "hip", cuda_version = "12.8") assert mock_pip.call_count == 1 call_args = [str(a) for a in mock_pip.call_args.args] assert "--force-reinstall" in call_args assert "--no-cache-dir" in call_args assert "cu128" in _index_url(mock_pip) assert mock_pip.call_args.kwargs["constrain"] is False def test_rocm_in_version_string_triggers_repair(self): # AMD SDK / Radeon wheels may not set torch.version.hip but encode # rocm in __version__; the probe prints "hip" for both. mock_pip = _run_cuda_repair(torch_state = "hip") assert mock_pip.call_count == 1 # ── No-op cases ────────────────────────────────────────────────────────────── class TestCudaRepairSkips: def test_healthy_cuda_torch_no_repair(self): mock_pip = _run_cuda_repair(torch_state = "cuda") mock_pip.assert_not_called() def test_deliberate_cpu_wheel_no_repair(self): mock_pip = _run_cuda_repair(torch_state = "cpu") mock_pip.assert_not_called() def test_backend_rocm_skips(self): mock_pip = _run_cuda_repair(backend = "rocm", torch_state = "hip") mock_pip.assert_not_called() def test_backend_cpu_skips(self): mock_pip = _run_cuda_repair(backend = "cpu", torch_state = "hip") mock_pip.assert_not_called() def test_unknown_backend_skips(self): mock_pip = _run_cuda_repair(backend = "auto", torch_state = "hip") mock_pip.assert_not_called() def test_no_nvidia_gpu_skips(self): mock_pip = _run_cuda_repair(nvidia = False, torch_state = "hip") mock_pip.assert_not_called() def test_torch_missing_skips(self): # Non-zero probe exit = torch missing / un-importable. mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1) mock_pip.assert_not_called() def test_macos_skips(self): mock_pip = _run_cuda_repair(is_macos = True, torch_state = "hip") mock_pip.assert_not_called() def test_windows_skips(self): mock_pip = _run_cuda_repair(is_windows = True, torch_state = "hip") mock_pip.assert_not_called() def test_no_torch_mode_skips(self): mock_pip = _run_cuda_repair(no_torch = True, torch_state = "hip") mock_pip.assert_not_called() def test_rocm_install_marker_skips(self): mock_pip = _run_cuda_repair(rocm_marker = True, torch_state = "hip") mock_pip.assert_not_called() def test_cvd_minus_one_skips(self): # CUDA_VISIBLE_DEVICES=-1 deliberately hides the NVIDIA GPU (mixed # AMD+NVIDIA host running ROCm torch on the AMD card). mock_pip = _run_cuda_repair(cvd = "-1", torch_state = "hip") mock_pip.assert_not_called() def test_cvd_empty_skips(self): mock_pip = _run_cuda_repair(cvd = "", torch_state = "hip") mock_pip.assert_not_called() def test_cvd_explicit_device_still_repairs(self): mock_pip = _run_cuda_repair(cvd = "0", torch_state = "hip") assert mock_pip.call_count == 1 # ── CUDA index ladder ──────────────────────────────────────────────────────── class TestCudaIndexResolution: def test_cuda_128_selects_cu128(self): assert "cu128" in _index_url(_run_cuda_repair(cuda_version = "12.8")) def test_cuda_130_selects_cu130(self): assert "cu130" in _index_url(_run_cuda_repair(cuda_version = "13.0")) def test_cuda_126_selects_cu126(self): assert "cu126" in _index_url(_run_cuda_repair(cuda_version = "12.6")) def test_cuda_124_selects_cu124(self): assert "cu124" in _index_url(_run_cuda_repair(cuda_version = "12.4")) def test_cuda_118_selects_cu118(self): assert "cu118" in _index_url(_run_cuda_repair(cuda_version = "11.8")) def test_unreadable_version_defaults_cu126(self): # nvidia-smi runs but prints no CUDA version line (or fails). mock_pip = _run_cuda_repair(cuda_version = "", smi_rc = 1) assert "cu126" in _index_url(mock_pip) def test_proc_fallback_no_smi_defaults_cu126(self): # NVIDIA usable via /proc fallback, nvidia-smi absent entirely. mock_pip = _run_cuda_repair(smi_path = None) assert "cu126" in _index_url(mock_pip) def test_detect_index_url_uses_pytorch_base(self): with ( patch.object(stack_mod.shutil, "which", return_value = None), patch.object(stack_mod.os.path, "isfile", return_value = False), ): url = _detect_cuda_torch_index_url() assert url == f"{stack_mod._PYTORCH_WHL_BASE}/cu126" if __name__ == "__main__": sys.exit(pytest.main([__file__, "-q"]))