diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index e252955cde..d27fff0d82 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1175,9 +1175,32 @@ 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 not exe and IS_WINDOWS: - for _candidate in ( + def _lists_a_gpu(exe: str) -> bool: + try: + result = subprocess.run( + [exe, "-L"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 10, + ) + except Exception: + 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", @@ -1189,23 +1212,12 @@ def _has_usable_nvidia_gpu() -> bool: "System32", "nvidia-smi.exe", ), - ): - if os.path.isfile(_candidate): - exe = _candidate - break - if exe: - try: - result = subprocess.run( - [exe, "-L"], - stdout = subprocess.PIPE, - stderr = subprocess.DEVNULL, - text = True, - timeout = 10, - ) - if result.returncode == 0 and "GPU " in result.stdout: - return True - except Exception: - pass + )) + 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