diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 2883f30b20..65530a3f00 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1162,6 +1162,11 @@ def _has_usable_nvidia_gpu() -> bool: timeout, driver initialisation race). If either probe confirms an NVIDIA GPU the function returns True so _has_rocm_gpu() is blocked. + On Windows nvidia-smi.exe is not always on PATH, so fall back to the two + fixed driver locations install.ps1 / setup.ps1 already probe. Without this + an AMD iGPU + NVIDIA dGPU host whose PATH misses nvidia-smi looks + NVIDIA-free here and gets routed to the ROCm wheels. + CUDA_VISIBLE_DEVICES set to "" or "-1" hides every NVIDIA device (mixed AMD+NVIDIA hosts steering work to the AMD card); neither probe honours that env var, so check it first and report the GPU as not usable. Unset @@ -1170,8 +1175,8 @@ def _has_usable_nvidia_gpu() -> bool: cvd = os.environ.get("CUDA_VISIBLE_DEVICES") if cvd is not None and cvd.strip() in ("", "-1"): return False - exe = shutil.which("nvidia-smi") - if exe: + + def _lists_a_gpu(exe: str) -> bool: try: result = subprocess.run( [exe, "-L"], @@ -1180,10 +1185,42 @@ def _has_usable_nvidia_gpu() -> bool: text = True, timeout = 10, ) - if result.returncode == 0 and "GPU " in result.stdout: - return True except Exception: - pass + return False + return result.returncode == 0 and "GPU " in result.stdout + + # Try every candidate until one lists a GPU, rather than committing to the + # first executable found. A stale or driverless nvidia-smi on PATH exits + # non-zero listing nothing; stopping there would report the host as + # NVIDIA-free and route it into _ensure_rocm_torch() even though a working + # driver binary sits at a fixed location. install.ps1 and setup.ps1 both + # gate their fallback on the GPU check failing, not on the PATH lookup + # missing, so mirror that. + candidates = [] + _path_exe = shutil.which("nvidia-smi") + if _path_exe: + candidates.append(_path_exe) + if IS_WINDOWS: + candidates.extend( + ( + os.path.join( + os.environ.get("ProgramFiles", r"C:\Program Files"), + "NVIDIA Corporation", + "NVSMI", + "nvidia-smi.exe", + ), + os.path.join( + os.environ.get("SystemRoot", r"C:\Windows"), + "System32", + "nvidia-smi.exe", + ), + ) + ) + for _candidate in candidates: + if _candidate != _path_exe and not os.path.isfile(_candidate): + continue + if _lists_a_gpu(_candidate): + return True # Fallback: the NVIDIA driver exposes one subdirectory per GPU under # /proc/driver/nvidia/gpus/ on Linux regardless of nvidia-smi state. if sys.platform != "win32": diff --git a/tests/studio/install/test_nvidia_smi_candidate_probing.py b/tests/studio/install/test_nvidia_smi_candidate_probing.py new file mode 100644 index 0000000000..c33c551417 --- /dev/null +++ b/tests/studio/install/test_nvidia_smi_candidate_probing.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""_has_usable_nvidia_gpu must keep probing after an unusable nvidia-smi. + +A stale or driverless nvidia-smi on PATH exits non-zero and lists no GPU. +Treating the first executable found as the answer reports a mixed AMD+NVIDIA +Windows host as NVIDIA-free, which routes it into _ensure_rocm_torch() and +replaces a working CUDA stack with ROCm wheels. install.ps1 and setup.ps1 both +gate their fixed-location fallback on the GPU check failing rather than on the +PATH lookup missing; this pins the Python helper to the same rule. + +The stubs are real executables run through the real subprocess call, so the +test exercises the actual control flow rather than a mocked return value. +""" + +import importlib.util +import os +import pathlib +import sys +import types + +import pytest + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +_STUDIO = _REPO_ROOT / "studio" +_SRC = _STUDIO / "install_python_stack.py" + +_STALE = 'echo "No devices were found"; exit 9' +_WORKING = 'echo "GPU 0: NVIDIA H100 (UUID: GPU-abc)"; exit 0' + + +def _load_module(): + # install_python_stack imports backend.utils.wheel_utils, which resolves + # only with studio/ on sys.path. That is how the installer invokes it. + if str(_STUDIO) not in sys.path: + sys.path.insert(0, str(_STUDIO)) + spec = importlib.util.spec_from_file_location("_ips_probe_under_test", _SRC) + module = importlib.util.module_from_spec(spec) + sys.modules["_ips_probe_under_test"] = module + spec.loader.exec_module(module) + return module + + +def _write_stub(path: pathlib.Path, body: str) -> None: + path.parent.mkdir(parents = True, exist_ok = True) + # Not /usr/bin/env: PATH is narrowed to the stub directory below, so env + # would not find an interpreter. + path.write_text("#!/bin/bash\n" + body + "\n") + path.chmod(0o755) + + +@pytest.fixture +def probe(tmp_path, monkeypatch): + """Run _has_usable_nvidia_gpu as if on Windows, with stubbed nvidia-smi.""" + + def _run( + path_smi: str | None, + fixed_smi: str | None, + cuda_visible_devices: str | None = None, + ) -> bool: + path_dir = tmp_path / "pathbin" + path_dir.mkdir(exist_ok = True) + if path_smi is not None: + _write_stub(path_dir / "nvidia-smi", path_smi) + program_files = tmp_path / "ProgramFiles" + if fixed_smi is not None: + _write_stub( + program_files / "NVIDIA Corporation" / "NVSMI" / "nvidia-smi.exe", + fixed_smi, + ) + monkeypatch.setenv("PATH", str(path_dir)) + monkeypatch.setenv("ProgramFiles", str(program_files)) + monkeypatch.setenv("SystemRoot", str(tmp_path / "Windows")) + if cuda_visible_devices is None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + else: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda_visible_devices) + module = _load_module() + monkeypatch.setattr(module, "IS_WINDOWS", True) + # A real NVIDIA host has /proc/driver/nvidia/gpus, and the helper's + # Linux-only fallback would then answer True for every case and mask + # what the Windows path did. Present as win32 to isolate it. + monkeypatch.setattr(module, "sys", types.SimpleNamespace(platform = "win32")) + return module._has_usable_nvidia_gpu() + + return _run + + +def test_stale_path_nvidia_smi_still_reaches_the_fixed_locations(probe): + # The regression: a driverless nvidia-smi on PATH used to end the search. + assert probe(_STALE, _WORKING) is True + + +def test_absent_path_nvidia_smi_reaches_the_fixed_locations(probe): + assert probe(None, _WORKING) is True + + +def test_working_path_nvidia_smi_is_enough(probe): + assert probe(_WORKING, None) is True + + +def test_no_nvidia_smi_anywhere_reports_no_gpu(probe): + assert probe(None, None) is False + + +def test_stale_everywhere_reports_no_gpu(probe): + # Every candidate answering "no GPU" must stay False, or an AMD-only host + # with a leftover nvidia-smi would be denied the ROCm wheels. + assert probe(_STALE, _STALE) is False + + +@pytest.mark.parametrize("hidden", ["", "-1", " "]) +def test_cuda_visible_devices_hidden_wins_over_a_working_probe(probe, hidden): + assert probe(_WORKING, _WORKING, cuda_visible_devices = hidden) is False + + +def test_cuda_visible_devices_listing_a_device_does_not_block_detection(probe): + assert probe(_WORKING, None, cuda_visible_devices = "0") is True