From 2db9fad4b5450daffce6c56dbf37292562ee602d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 05:06:02 -0700 Subject: [PATCH] Installer: GPU detection follow-ups after #6174 (poisoned venv repair, llama.cpp routing, probe bounds) (#6183) * Installer: harden GPU detection follow-ups after #6174 Ports the NVIDIA-priority and /proc/driver/nvidia/gpus hardening from #6174 to the remaining pathways and adds recovery for already-poisoned venvs: - install_python_stack.py: add _ensure_cuda_torch so 'unsloth studio update' force-reinstalls CUDA torch when the venv carries a ROCm build on an NVIDIA Linux host (the pre-#6174 poisoning signature). Honors UNSLOTH_TORCH_BACKEND, UNSLOTH_ROCM_TORCH_INSTALLED, and CUDA_VISIBLE_DEVICES=-1/'' opt-outs; never touches healthy CUDA, deliberate CPU wheels, macOS, or Windows. - install_llama_prebuilt.py: detect_host gains the /proc NVIDIA fallback and skips ROCm probes when NVIDIA is usable; forwarded --rocm-gfx/--has-rocm overrides still win. - setup.sh: GPU summary classifies NVIDIA first through a timeout-bounded probe with the /proc fallback; AMD probes are bounded and gain a KFD vendor_id 4098 fallback; the llama.cpp source build only selects GGML_CUDA/GGML_HIP when the matching GPU is actually detected. - install.sh: bound both nvidia-smi calls with a 10s timeout (no behavior change when healthy or when the timeout binary is absent); classify the exported UNSLOTH_TORCH_BACKEND on the final index path segment so custom mirrors containing 'rocm'/'gfx' in their base path are not mislabeled. - install.ps1 + setup.ps1: NVIDIA probes now require a real 'GPU N:' row from nvidia-smi -L under a 10s bound instead of bare exit code 0; later CUDA version and compute_cap queries are bounded too. Tests: 3 new test files (50+ tests), suite at 788 passed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Resolve-CudaToolkit driver probe for extracted-function unit test tests/studio/test_resolve_cuda_toolkit.ps1 extracts Resolve-CudaToolkit alone into a child pwsh and stubs nvidia-smi with a .ps1 script. The bounded runner is not in scope there (and ProcessStartInfo cannot dispatch .ps1 stubs), so the DriverMaxCuda parse silently returned nothing and the major-mismatch scenarios failed. Fall back to direct invocation when Invoke-NvidiaSmiBounded is unavailable; production setup.ps1 always has it defined and keeps the 10s bound. * Treat CUDA_VISIBLE_DEVICES empty or -1 as hidden in NVIDIA-first guards The NVIDIA-first guards added in this branch only special-cased CUDA_VISIBLE_DEVICES=-1 at two setup.sh gates and ignored the empty-string form entirely, while the Python detector (install_llama_prebuilt.py) already treats both as hidden. On a mixed AMD+NVIDIA host steered to the AMD card via CUDA_VISIBLE_DEVICES, the guards suppressed the AMD probes, so setup.sh fell to a CPU llama.cpp build and install.sh picked CUDA wheels instead of ROCm. Move the policy into the helpers so every consumer agrees: - install.sh: new _cvd_hides_nvidia checked first in _has_usable_nvidia_gpu - studio/setup.sh: same via _setup_cvd_hides_nvidia; the two ad-hoc CUDA_VISIBLE_DEVICES=-1 gate conditions are now redundant and removed - studio/install_python_stack.py: _has_usable_nvidia_gpu returns False when CUDA_VISIBLE_DEVICES is set to or -1 (whitespace tolerated) Tests: 5 new sh scenarios (hidden via , -1, padded -1, visible device, and mixed host with hidden NVIDIA restoring the ROCm route) plus a pytest class covering all three implementations behaviourally. Addresses the review comment on the NVIDIA-first setup.sh block. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Retrigger CI after PyPI 503 outage during the previous run --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.ps1 | 55 +- install.sh | 51 +- studio/install_llama_prebuilt.py | 24 +- studio/install_python_stack.py | 155 ++++++ studio/setup.ps1 | 81 ++- studio/setup.sh | 166 ++++-- tests/sh/test_get_torch_index_url.sh | 51 +- tests/studio/install/test_cuda_repair.py | 248 +++++++++ .../install/test_gpu_detection_followups.py | 480 ++++++++++++++++++ tests/studio/install/test_probe_timeouts.py | 202 ++++++++ 10 files changed, 1439 insertions(+), 74 deletions(-) create mode 100644 tests/studio/install/test_cuda_repair.py create mode 100644 tests/studio/install/test_gpu_detection_followups.py create mode 100644 tests/studio/install/test_probe_timeouts.py diff --git a/install.ps1 b/install.ps1 index cbc7b206c8..2938d4334f 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1406,14 +1406,58 @@ shell.Run cmd, 0, False } } + # ── Helper: run nvidia-smi under a timeout ── + # A wedged NVIDIA driver can make nvidia-smi block during init or after a + # reset; WaitForExit bounds it (mirrors Invoke-AmdSmiNoElevate) so detection + # cannot hang the installer. No RunAsInvoker compat layer: nvidia-smi does + # not auto-elevate. Returns combined stdout+stderr; "" on timeout/failure. + function Invoke-NvidiaSmiBounded { + param( + [Parameter(Mandatory = $true, Position = 0)][string]$Exe, + [Parameter(Position = 1)][string[]]$SmiArgs = @(), + [int]$TimeoutSec = 10 + ) + try { + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $Exe + $psi.Arguments = ($SmiArgs -join ' ') + $psi.UseShellExecute = $false + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + if (-not $proc.WaitForExit($TimeoutSec * 1000)) { + try { $proc.Kill() } catch {} + $global:LASTEXITCODE = 124 + return "" + } + $global:LASTEXITCODE = $proc.ExitCode + return ($outTask.Result + "`n" + $errTask.Result) + } catch { + $global:LASTEXITCODE = 1 + return "" + } + } + + # ── Helper: nvidia-smi -L lists at least one real GPU ── + # Exit code 0 alone is not enough: a stale/driverless nvidia-smi can exit 0 + # while listing no GPU, which would mark an AMD host NVIDIA and suppress + # ROCm detection. Require a "GPU :" data row. + function Test-NvidiaSmiHasGpu { + param([Parameter(Mandatory = $true)][string]$Exe) + $out = Invoke-NvidiaSmiBounded $Exe @('-L') + return ($LASTEXITCODE -eq 0 -and $out -match '(?m)^GPU\s+\d+:') + } + # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ── $HasNvidiaSmi = $false $NvidiaSmiExe = $null try { $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue - if ($nvSmiCmd) { - & $nvSmiCmd.Source *> $null - if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source } + if ($nvSmiCmd -and (Test-NvidiaSmiHasGpu $nvSmiCmd.Source)) { + $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source } } catch {} if (-not $HasNvidiaSmi) { @@ -1423,8 +1467,7 @@ shell.Run cmd, 0, False )) { if (Test-Path $p) { try { - & $p *> $null - if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break } + if (Test-NvidiaSmiHasGpu $p) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break } } catch {} } } @@ -1694,7 +1737,7 @@ shell.Run cmd, 0, False $baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { - $output = & $NvidiaSmiExe 2>&1 | Out-String + $output = Invoke-NvidiaSmiBounded $NvidiaSmiExe # Newer NVIDIA drivers (e.g. 610.x on Windows) print # "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y". # Accept both spellings so we don't fall through to the cu126 default. diff --git a/install.sh b/install.sh index 44d1c39d02..df9eca65c7 100755 --- a/install.sh +++ b/install.sh @@ -1745,13 +1745,42 @@ _has_amd_rocm_gpu() { return 1 } +# ── Bounded command runner ── +# Runs a command under a 10s timeout when the `timeout` binary is available, +# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during +# driver init or after a reset) from hanging the installer: a timed-out probe +# exits nonzero and is treated exactly like a failed probe. No-op semantics on +# hosts without `timeout` (e.g. macOS) or when the probe is healthy. +_run_bounded() { + if command -v timeout >/dev/null 2>&1; then + timeout 10 "$@" + else + "$@" + fi +} + +# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every +# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to +# the AMD card). Unset means all devices visible. nvidia-smi ignores this env +# var, so the probes below cannot see the distinction on their own. +_cvd_hides_nvidia() { + [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 + _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') + [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] +} + # ── NVIDIA usable-GPU helper ── # Returns 0 (true) if an NVIDIA GPU is present and usable. # Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, # which the NVIDIA driver populates on Linux regardless of nvidia-smi state # -- handles PATH gaps, subprocess timeouts, and driver init races that # could otherwise cause nvidia-smi to fail and silence NVIDIA detection. +# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches +# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. _has_usable_nvidia_gpu() { + if _cvd_hides_nvidia; then + return 1 + fi _nvsmi="" if command -v nvidia-smi >/dev/null 2>&1; then _nvsmi="nvidia-smi" @@ -1759,7 +1788,7 @@ _has_usable_nvidia_gpu() { _nvsmi="/usr/bin/nvidia-smi" fi if [ -n "$_nvsmi" ]; then - if "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then return 0 fi fi @@ -1871,7 +1900,11 @@ get_torch_index_url() { # of the legacy "CUDA Version: X.Y"; accept both with two BRE expressions # (POSIX sed does not support "?" without -E). The two patterns are # mutually exclusive per line, so head -1 picks the first emitted match. - _cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \ + # Bound the call (a wedged nvidia-smi would otherwise hang here) and force + # the C locale for stable parsing. LC_ALL is exported inside this command + # substitution subshell so it reaches nvidia-smi through _run_bounded + # without depending on `env`; the export is scoped to the subshell. + _cuda_ver=$(export LC_ALL=C; _run_bounded "$_smi" 2>/dev/null \ | sed -n \ -e 's/.*CUDA UMD Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ -e 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ @@ -2125,10 +2158,16 @@ TORCH_INDEX_URL=$(get_torch_index_url) # Export the resolved torch backend ("cuda", "rocm", or "cpu") so that # downstream scripts (setup.sh -> install_python_stack.py) know what was # chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts. -case "$TORCH_INDEX_URL" in - */rocm*|*/gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; - */cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; - *) export UNSLOTH_TORCH_BACKEND="cuda" ;; +# Classify on the FINAL path segment only: a custom UNSLOTH_PYTORCH_MIRROR +# whose base path happens to contain "rocm" or "gfx" must not mislabel a +# cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix +# overrides in gfxNNNN/, so the trailing slash is stripped first). +_torch_index_leaf="${TORCH_INDEX_URL%/}" +_torch_index_leaf="${_torch_index_leaf##*/}" +case "$_torch_index_leaf" in + rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; + cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; + *) export UNSLOTH_TORCH_BACKEND="cuda" ;; esac # rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index f9425f5c4f..22fab82923 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -2890,7 +2890,27 @@ def detect_host() -> HostInfo: except Exception: pass + # Linux /proc/driver/nvidia/gpus fallback: the NVIDIA driver exposes one + # subdir per GPU here regardless of nvidia-smi state, so a host whose + # nvidia-smi is absent from PATH, wedged, or failing is still recognised as + # NVIDIA. Mirrors the fallback added to install.sh / install_python_stack.py + # in PR 6174 so the prebuilt installer does not misroute such hosts to ROCm + # or CPU. driver_cuda_version / compute_caps stay unset here; downstream + # CUDA asset selection treats unknown SMs as "prefer portable" and an + # unknown driver runtime line as "no published CUDA match" (returns None, + # no crash), so planning falls back to a source build with GGML_CUDA=ON. + if is_linux and not has_physical_nvidia: + try: + proc_gpu_dir = "/proc/driver/nvidia/gpus" + if os.path.isdir(proc_gpu_dir) and os.listdir(proc_gpu_dir): + has_physical_nvidia = True + has_usable_nvidia = visible_device_tokens != [] + except OSError: + pass + # Detect AMD ROCm (HIP) -- require actual GPU, not just tools installed + # NVIDIA takes precedence: when an NVIDIA GPU is usable, skip ROCm probing + # entirely so co-installed ROCm tools cannot misroute the host (PR 6174). def _amd_smi_has_gpu(stdout: str) -> bool: """Check for 'GPU: ' data rows, not just a table header.""" @@ -2898,7 +2918,7 @@ def detect_host() -> HostInfo: has_rocm = False rocm_gfx_target: str | None = None - if is_linux: + if is_linux and not has_usable_nvidia: # WSL2 ROCDXG: the system rocminfo enumerates the GPU over /dev/dxg # only when HSA_ENABLE_DXG_DETECTION=1 (a no-op on bare metal), and # rocminfo can live only under /opt/rocm/bin (the profile.d PATH @@ -2937,7 +2957,7 @@ def detect_host() -> HostInfo: has_rocm = True rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout) break - elif is_windows: + elif is_windows and not has_usable_nvidia: # Windows: prefer active probes that validate GPU presence. # hipinfo / amd-smi are often NOT on PATH -- the HIP SDK installer # sets HIP_PATH / ROCM_PATH but does not always add the bin dir to diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 0460580922..e540aac305 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -90,6 +90,18 @@ _PYTORCH_WHL_BASE = ( os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl" ).rstrip("/") +# CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are +# pinned to the torch<2.11 family rather than left bare: the install uses an +# exclusive --index-url (no PyPI fallback), so a bare name could resolve a +# torchvision built against a different torch major (e.g. 0.27 for torch 2.12) +# and fail at runtime with an ABI mismatch. Same bounds as the _default ROCm +# spec above, which targets the same torch family. +_CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = ( + "torch>=2.4,<2.11.0", + "torchvision>=0.19,<0.26.0", + "torchaudio>=2.4,<2.11.0", +) + # AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/). # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs. _ROCM_WINDOWS_INDEX_BASE = ( @@ -654,7 +666,15 @@ def _has_usable_nvidia_gpu() -> bool: case where nvidia-smi is present but the subprocess fails (PATH gap, timeout, driver initialisation race). If either probe confirms an NVIDIA GPU the function returns True so _has_rocm_gpu() is blocked. + + 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 + means all devices visible. """ + 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: try: @@ -786,6 +806,139 @@ def _install_bnb_windows_rocm() -> bool: return True +def _detect_cuda_torch_index_url() -> str: + """Return the pytorch.org CUDA wheel index URL for the host's NVIDIA driver. + + Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update` + repairs to the same wheel family a fresh `curl | sh` install would pick. + Probes nvidia-smi (PATH, then /usr/bin/nvidia-smi) and parses both the + legacy "CUDA Version:" and the newer "CUDA UMD Version:" spellings. + Defaults to cu126 when nvidia-smi is missing or the version is unreadable + (e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback). + """ + exe = shutil.which("nvidia-smi") + if not exe and os.path.isfile("/usr/bin/nvidia-smi"): + exe = "/usr/bin/nvidia-smi" + tag = "cu126" # default when the driver CUDA version cannot be read + if exe: + try: + result = subprocess.run( + [exe], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 10, + ) + if result.returncode == 0: + m = re.search(r"CUDA(?: UMD)? Version:\s*(\d+)\.(\d+)", result.stdout) + if m: + major, minor = int(m.group(1)), int(m.group(2)) + if major >= 13: + tag = "cu130" + elif major == 12 and minor >= 8: + tag = "cu128" + elif major == 12 and minor >= 6: + tag = "cu126" + elif major >= 12: + tag = "cu124" + elif major >= 11: + tag = "cu118" + else: + tag = "cpu" # ancient driver: no usable CUDA wheels + except Exception: + pass + return f"{_PYTORCH_WHL_BASE}/{tag}" + + +def _ensure_cuda_torch() -> None: + """Repair a venv whose torch is a ROCm build on an NVIDIA host. + + Counterpart to _ensure_rocm_torch. A venv poisoned by the pre-fix KFD + gpu_id false positive (ROCm torch installed on an NVIDIA-only machine) + keeps that broken torch on `studio update`, because a torch+rocm wheel + satisfies the version constraint and nothing force-reinstalls it. This + detects that exact case and reinstalls CUDA torch. + + Only repairs when torch actually links against HIP/ROCm. Healthy CUDA + torch and deliberate CPU-only torch are left untouched. + """ + # Respect an explicit backend choice from install.sh: only "" (standalone + # `studio update`) or "cuda" should ever force CUDA wheels. "rocm"/"cpu" + # (or any unrecognised value) are deliberate and must not be overridden. + if _TORCH_BACKEND not in ("", "cuda"): + return + # No CUDA torch on macOS; Windows venv/torch lifecycle is owned by + # install.ps1 (and the KFD poisoning bug is Linux-only), so skip both. + if IS_MACOS or IS_WINDOWS or NO_TORCH: + return + # Never undo a deliberate ROCm install (setup.ps1 sets this marker). + if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": + return + # 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. + _cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if _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. + if not _has_usable_nvidia_gpu(): + return + + # Classify the installed torch: "hip" (ROCm build -- the poisoning + # signature), "cuda" (healthy), or "cpu" (deliberate CPU wheel). A + # non-zero exit means torch is missing or un-importable; the base install + # step handles that, so leave it alone. + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import torch; " + "hip = getattr(torch.version, 'hip', '') or ''; " + "cuda = getattr(torch.version, 'cuda', '') or ''; " + "ver = getattr(torch, '__version__', '').lower(); " + "print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 90, + ) + except (OSError, subprocess.TimeoutExpired): + return + if probe.returncode != 0: + return + # Take the last non-empty stdout line: stray output from sitecustomize or + # an import hook must not mask the marker (fail-closed either way). + _marker_lines = [ + line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip() + ] + if not _marker_lines or _marker_lines[-1] != "hip": + return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is + + index_url = _detect_cuda_torch_index_url() + _torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC + print( + f" torch is a ROCm build on an NVIDIA host -- reinstalling " + f"CUDA torch from {index_url}\n" + f" (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch " + f"on a mixed AMD+NVIDIA host)" + ) + pip_install( + "CUDA torch repair", + "--force-reinstall", + "--no-cache-dir", + _torch_pkg, + _vision_pkg, + _audio_pkg, + "--index-url", + index_url, + constrain = False, + ) + + def _ensure_rocm_torch() -> None: """Reinstall torch with ROCm wheels when the venv received CPU-only torch. @@ -1848,6 +2001,7 @@ def install_python_stack() -> int: # Must follow base packages so torch is present for inspection. if not IS_MACOS and not NO_TORCH: _progress(_torch_step_label("check")) + _ensure_cuda_torch() _ensure_rocm_torch() # Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python @@ -2033,6 +2187,7 @@ def install_python_stack() -> int: # whichever intermediate step clobbered it. if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: _progress(_torch_step_label("final")) + _ensure_cuda_torch() _ensure_rocm_torch() # 14. Final check (silent; third-party conflicts are expected) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index e707cfda7a..29350e3932 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -299,7 +299,10 @@ function Get-CudaComputeCapability { if (-not $smiExe) { return $null } try { - $raw = & $smiExe --query-gpu=compute_cap --format=csv,noheader 2>$null + # Bounded: a wedged nvidia-smi must not hang setup after the initial + # -L probe succeeded (the helper merges stderr after stdout, so the + # first line is still the compute_cap value). + $raw = Invoke-NvidiaSmiBounded $smiExe @('--query-gpu=compute_cap', '--format=csv,noheader') if ($LASTEXITCODE -ne 0 -or -not $raw) { return $null } # nvidia-smi may return multiple GPUs; take the first one @@ -363,10 +366,10 @@ function Get-PytorchCudaTag { if (-not $smiExe) { return "cu126" } try { - # 2>&1 | Out-String merges stderr into stdout then converts to a single - # string. Plain 2>$null doesn't fully suppress stderr in PS 5.1 -- - # ErrorRecord objects leak into $output and break the -match. - $output = & $smiExe 2>&1 | Out-String + # Bounded: a wedged nvidia-smi must not hang setup. The helper merges + # stderr into the returned string, matching the old 2>&1 | Out-String + # shape (plain 2>$null leaks ErrorRecord objects in PS 5.1). + $output = Invoke-NvidiaSmiBounded $smiExe # Newer NVIDIA drivers (e.g. 610.x on Windows) print # "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y". # Accept both spellings so we don't fall through to the cu126 default. @@ -667,16 +670,58 @@ try { # ============================================ # 1a. GPU detection # ============================================ +# ── Helper: run nvidia-smi under a timeout ── +# A wedged NVIDIA driver can make nvidia-smi block during init or after a reset; +# WaitForExit bounds it (mirrors Invoke-AmdSmiNoElevate below) so detection +# cannot hang setup. No RunAsInvoker compat layer: nvidia-smi does not +# auto-elevate. Returns combined stdout+stderr; "" on timeout/failure. +function Invoke-NvidiaSmiBounded { + param( + [Parameter(Mandatory = $true, Position = 0)][string]$Exe, + [Parameter(Position = 1)][string[]]$SmiArgs = @(), + [int]$TimeoutSec = 10 + ) + try { + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $Exe + $psi.Arguments = ($SmiArgs -join ' ') + $psi.UseShellExecute = $false + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + if (-not $proc.WaitForExit($TimeoutSec * 1000)) { + try { $proc.Kill() } catch {} + $global:LASTEXITCODE = 124 + return "" + } + $global:LASTEXITCODE = $proc.ExitCode + return ($outTask.Result + "`n" + $errTask.Result) + } catch { + $global:LASTEXITCODE = 1 + return "" + } +} + +# ── Helper: nvidia-smi -L lists at least one real GPU ── +# Exit code 0 alone is not enough: a stale/driverless nvidia-smi can exit 0 +# while listing no GPU, which would mark an AMD host NVIDIA and suppress ROCm +# detection. Require a "GPU :" data row. +function Test-NvidiaSmiHasGpu { + param([Parameter(Mandatory = $true)][string]$Exe) + $out = Invoke-NvidiaSmiBounded $Exe @('-L') + return ($LASTEXITCODE -eq 0 -and $out -match '(?m)^GPU\s+\d+:') +} + $HasNvidiaSmi = $false $NvidiaSmiExe = $null # Absolute path -- survives Refresh-Environment try { $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue - if ($nvSmiCmd) { - & $nvSmiCmd.Source *> $null - if ($LASTEXITCODE -eq 0) { - $HasNvidiaSmi = $true - $NvidiaSmiExe = $nvSmiCmd.Source - } + if ($nvSmiCmd -and (Test-NvidiaSmiHasGpu $nvSmiCmd.Source)) { + $HasNvidiaSmi = $true + $NvidiaSmiExe = $nvSmiCmd.Source } } catch {} # Fallback: nvidia-smi may not be on PATH even though a GPU + driver exist. @@ -689,8 +734,7 @@ if (-not $HasNvidiaSmi) { foreach ($p in $nvSmiDefaults) { if (Test-Path $p) { try { - & $p *> $null - if ($LASTEXITCODE -eq 0) { + if (Test-NvidiaSmiHasGpu $p) { $HasNvidiaSmi = $true $NvidiaSmiExe = $p Write-Host " Found nvidia-smi at $(Split-Path $p -Parent)" -ForegroundColor Gray @@ -1151,7 +1195,16 @@ function Resolve-CudaToolkit { $DriverMaxCuda = $null try { - $smiOut = & $NvidiaSmiExe 2>&1 | Out-String + # Bounded: source-build toolkit resolution must not hang on a wedged smi. + # test_resolve_cuda_toolkit.ps1 extracts this function alone into a child + # pwsh (no Invoke-NvidiaSmiBounded in scope) and stubs nvidia-smi with a + # .ps1 script, so fall back to direct invocation when the bounded runner + # is unavailable; production setup.ps1 always has it defined. + $smiOut = if (Get-Command Invoke-NvidiaSmiBounded -ErrorAction SilentlyContinue) { + Invoke-NvidiaSmiBounded $NvidiaSmiExe + } else { + & $NvidiaSmiExe 2>&1 | Out-String + } # Newer drivers report "CUDA UMD Version: X.Y" instead of "CUDA Version: X.Y"; accept both. if ($smiOut -match "CUDA(?: UMD)? Version:\s+([\d]+)\.([\d]+)") { $DriverMaxCuda = "$($Matches[1]).$($Matches[2])" diff --git a/studio/setup.sh b/studio/setup.sh index 66c03da391..a8603bd0da 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -155,9 +155,60 @@ _nvcc_meets_llama_minimum() { echo "$_raw" } +# Run a GPU probe under a 10s timeout when `timeout` is available so a wedged +# NVIDIA driver cannot hang setup; fall back to a bare call where it is not. +_setup_run_smi() { + if command -v timeout >/dev/null 2>&1; then + timeout 10 "$@" + else + "$@" + fi +} + +# Returns 0 when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every NVIDIA +# device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to the +# AMD card). Unset means all devices visible. nvidia-smi ignores this env var, +# so the probes below cannot see the distinction on their own. +_setup_cvd_hides_nvidia() { + [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 + _setup_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') + [ -z "$_setup_cvd_trim" ] || [ "$_setup_cvd_trim" = "-1" ] +} + +# Returns 0 when an NVIDIA GPU is present and usable. Primary probe is +# `nvidia-smi -L` (timeout-bounded). Fallback is /proc/driver/nvidia/gpus, +# which the driver populates per GPU regardless of nvidia-smi state -- handles +# PATH gaps and driver init races. Mirrors install.sh _has_usable_nvidia_gpu +# (PR 6174) so setup routes the same way as the torch installer. A GPU hidden +# via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches +# install_llama_prebuilt.py has_usable_nvidia), so the AMD probes still run +# and a mixed host steered to its AMD card keeps the ROCm route. +_setup_has_usable_nvidia_gpu() { + if _setup_cvd_hides_nvidia; then + return 1 + fi + _setup_nvsmi="" + if command -v nvidia-smi >/dev/null 2>&1; then + _setup_nvsmi="nvidia-smi" + elif [ -x "/usr/bin/nvidia-smi" ]; then + _setup_nvsmi="/usr/bin/nvidia-smi" + fi + if [ -n "$_setup_nvsmi" ]; then + if _setup_run_smi "$_setup_nvsmi" -L 2>/dev/null \ + | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + return 0 + fi + fi + if [ -d /proc/driver/nvidia/gpus ] && \ + [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then + return 0 + fi + return 1 +} + _cuda_driver_max_version() { command -v nvidia-smi >/dev/null 2>&1 || return 0 - nvidia-smi 2>/dev/null \ + _setup_run_smi nvidia-smi 2>/dev/null \ | sed -nE 's/.*CUDA( UMD)? Version:[[:space:]]*([0-9]+)\.([0-9]+).*/\2.\3/p' \ | head -1 || true } @@ -815,25 +866,42 @@ _setup_amd_detected=false _setup_nvidia_usable=false _setup_gfx_all="" _setup_mkt="" -if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then - _setup_amd_detected=true - _setup_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) - _setup_mkt=$(rocminfo 2>/dev/null | awk -F': ' \ - '/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) -elif command -v amd-smi >/dev/null 2>&1 && \ - amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then - _setup_amd_detected=true - _setup_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) - [ -z "$_setup_gfx_all" ] && \ - _setup_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) - _setup_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \ - '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) +# NVIDIA priority: classify NVIDIA first and skip the AMD probes entirely on +# a usable-NVIDIA host (mirrors _has_rocm_gpu in install_python_stack.py). +# This also keeps a wedged rocminfo/amd-smi from hanging setup before the +# host is classified; the AMD probes themselves run under _setup_run_smi. +if _setup_has_usable_nvidia_gpu; then + _setup_nvidia_usable=true +fi +if [ "$_setup_nvidia_usable" != true ]; then + if command -v rocminfo >/dev/null 2>&1 && \ + _setup_run_smi rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then + _setup_amd_detected=true + _setup_gfx_all=$(_setup_run_smi rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + _setup_mkt=$(_setup_run_smi rocminfo 2>/dev/null | awk -F': ' \ + '/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) + elif command -v amd-smi >/dev/null 2>&1 && \ + _setup_run_smi amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then + _setup_amd_detected=true + _setup_gfx_all=$(_setup_run_smi amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + [ -z "$_setup_gfx_all" ] && \ + _setup_gfx_all=$(_setup_run_smi amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + _setup_mkt=$(_setup_run_smi amd-smi static --asic 2>/dev/null | awk -F'[:|]' \ + '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) + elif [ -e /dev/kfd ] && \ + awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \ + gpu && amd { found=1 } END{ exit !found }' \ + /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then + # KFD sysfs fallback, AMD vendor_id 4098 only (mirrors install.sh + # _has_amd_rocm_gpu): covers AMD hosts where rocminfo/amd-smi are + # missing but the kernel exposes the GPU, so the source-build gate + # below does not drop them to a CPU llama.cpp build. No gfx arch is + # available from this path; name-based inference handles it. + _setup_amd_detected=true + fi fi -if command -v nvidia-smi >/dev/null 2>&1 && \ - nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then - _setup_nvidia_usable=true +if [ "$_setup_nvidia_usable" = true ]; then step "gpu" "NVIDIA GPU detected" elif [ "$_setup_amd_detected" = true ]; then _setup_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}" @@ -918,15 +986,15 @@ _HOST_MACHINE="$(uname -m 2>/dev/null || true)" # use unslothai. _LINUX_HAS_GPU=false # Route to the fork only for a usable GPU. NVIDIA counts only when a device is -# actually enumerated (_setup_nvidia_usable, from the nvidia-smi -L probe above) -# AND not hidden via CUDA_VISIBLE_DEVICES=-1 -- mirroring install_llama_prebuilt.py's -# has_usable_nvidia. Mere nvidia-smi presence (CPU-only CUDA-toolkit containers, -# broken drivers) or a hidden GPU therefore takes the ggml-org CPU prebuilt -# instead of a slow source build. AMD is deliberately left on tooling presence, -# not usability: an unusable NVIDIA host has a good CPU prebuilt to fall back to, -# whereas tightening AMD would regress ROCm hosts exposing only hipconfig/hipinfo -# into an unnecessary CPU build. -if [ "$_setup_nvidia_usable" = true ] && [ "${CUDA_VISIBLE_DEVICES:-}" != "-1" ]; then +# actually enumerated and not hidden via CUDA_VISIBLE_DEVICES=""/-1 +# (_setup_nvidia_usable, from _setup_has_usable_nvidia_gpu above) -- mirroring +# install_llama_prebuilt.py's has_usable_nvidia. Mere nvidia-smi presence +# (CPU-only CUDA-toolkit containers, broken drivers) or a hidden GPU therefore +# takes the ggml-org CPU prebuilt instead of a slow source build. AMD is +# deliberately left on tooling presence, not usability: an unusable NVIDIA host +# has a good CPU prebuilt to fall back to, whereas tightening AMD would regress +# ROCm hosts exposing only hipconfig/hipinfo into an unnecessary CPU build. +if [ "$_setup_nvidia_usable" = true ]; then _LINUX_HAS_GPU=true else for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo; do @@ -1271,23 +1339,35 @@ else GPU_BACKEND="" NVCC_PATH="" - if command -v nvcc &>/dev/null; then - NVCC_PATH="$(command -v nvcc)" - GPU_BACKEND="cuda" - elif [ -x /usr/local/cuda/bin/nvcc ]; then - NVCC_PATH="/usr/local/cuda/bin/nvcc" - export PATH="/usr/local/cuda/bin:$PATH" - GPU_BACKEND="cuda" - elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then - # Pick the newest cuda-XX.X directory - NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)" - export PATH="$(dirname "$NVCC_PATH"):$PATH" - GPU_BACKEND="cuda" + # Gate the CUDA toolkit search on an actually-usable NVIDIA GPU + # (_setup_nvidia_usable, computed in the GPU summary block above; + # already false when hidden via CUDA_VISIBLE_DEVICES=""/-1). + # A CUDA toolkit alone (CPU-only build container, leftover packages) + # is not proof of a GPU: building with -DGGML_CUDA=ON there yields a + # binary that fails at runtime, so fall through to the CPU build. + if [ "$_setup_nvidia_usable" = true ]; then + if command -v nvcc &>/dev/null; then + NVCC_PATH="$(command -v nvcc)" + GPU_BACKEND="cuda" + elif [ -x /usr/local/cuda/bin/nvcc ]; then + NVCC_PATH="/usr/local/cuda/bin/nvcc" + export PATH="/usr/local/cuda/bin:$PATH" + GPU_BACKEND="cuda" + elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then + # Pick the newest cuda-XX.X directory + NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)" + export PATH="$(dirname "$NVCC_PATH"):$PATH" + GPU_BACKEND="cuda" + fi fi - # Check for ROCm (AMD) only if CUDA was not already selected + # Check for ROCm (AMD) only if CUDA was not already selected, and + # only when an AMD GPU was actually detected (_setup_amd_detected). + # hipcc presence alone (HIP SDK, no GPU) must not select a HIP build. + # NVIDIA-usable hosts never build HIP (defense in depth: the AMD + # probes above are already skipped when NVIDIA is usable). ROCM_HIPCC="" - if [ -z "$GPU_BACKEND" ]; then + if [ -z "$GPU_BACKEND" ] && [ "$_setup_nvidia_usable" != true ] && [ "$_setup_amd_detected" = true ]; then if command -v hipcc &>/dev/null; then ROCM_HIPCC="$(command -v hipcc)" GPU_BACKEND="rocm" @@ -1349,7 +1429,7 @@ else CUDA_ARCHS="" if command -v nvidia-smi &>/dev/null; then - _raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true) + _raw_caps=$(_setup_run_smi nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true) while IFS= read -r _cap; do _cap=$(echo "$_cap" | tr -d '[:space:]') if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then @@ -1455,7 +1535,7 @@ else CMAKE_ARGS="$CMAKE_ARGS -DGPU_TARGETS=${GPU_TARGETS}" _BUILD_DESC="building (ROCm, ${GPU_TARGETS//;/+})" fi - elif [ -d /usr/local/cuda ] || nvidia-smi &>/dev/null; then + elif [ -d /usr/local/cuda ] || _setup_run_smi nvidia-smi &>/dev/null; then _BUILD_DESC="building (CPU, CUDA driver found but nvcc missing)" elif [ -d /opt/rocm ] || command -v rocm-smi &>/dev/null; then _BUILD_DESC="building (CPU, ROCm driver found but hipcc missing)" diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index e20fd1ca86..3dece1f69a 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -13,6 +13,10 @@ FAIL=0 _FUNC_FILE=$(mktemp) _FAKE_SMI_DIR=$(mktemp -d) { + sed -n '/^_run_bounded()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_cvd_hides_nvidia()/,/^}/p' "$INSTALL_SH" + echo "" sed -n '/^_has_amd_rocm_gpu()/,/^}/p' "$INSTALL_SH" echo "" sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH" @@ -107,7 +111,7 @@ MOCK # Build a minimal tools directory with symlinks to essential commands # (uname, grep, head, etc.) but WITHOUT nvidia-smi or amd-smi. _TOOLS_DIR=$(mktemp -d) -for _cmd in uname grep sed head sh bash cat awk printf; do +for _cmd in uname grep sed head sh bash cat awk printf tr; do _real=$(command -v "$_cmd" 2>/dev/null || true) [ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd" done @@ -116,12 +120,19 @@ done # $1 = directory with mock nvidia-smi (prepended to PATH), or "none" for no-GPU test run_func() { _mock_dir="$1" + # Default: strip CUDA_VISIBLE_DEVICES so the host environment cannot leak + # in; a second argument sets it explicitly (hidden-GPU scenarios). + if [ "$#" -ge 2 ]; then + _cvd_setup="export CUDA_VISIBLE_DEVICES='$2'" + else + _cvd_setup="unset CUDA_VISIBLE_DEVICES" + fi if [ "$_mock_dir" = "none" ]; then # Minimal PATH with only basic tools, no nvidia-smi anywhere - PATH="$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null + PATH="$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null else # Put mock nvidia-smi dir first, then basic tools - PATH="$_mock_dir:$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null + PATH="$_mock_dir:$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null fi } @@ -332,6 +343,40 @@ _result=$(run_func "$_dir") assert_eq "CUDA Version 13.7 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result" rm -rf "$_dir" +# 34) CUDA_VISIBLE_DEVICES="" hides the NVIDIA GPU -> cpu (no AMD present) +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir" "") +assert_eq "CVD='' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + +# 35) CUDA_VISIBLE_DEVICES=-1 hides the NVIDIA GPU -> cpu (no AMD present) +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir" "-1") +assert_eq "CVD=-1 hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + +# 36) Mixed AMD+NVIDIA host with NVIDIA hidden -> ROCm route is restored +_cuda_dir=$(make_mock_smi "12.6") +_amd_dir=$(make_mock_amd_smi "6.4") +_combined_dir=$(mktemp -d) +ln -sf "$_cuda_dir/nvidia-smi" "$_combined_dir/nvidia-smi" +ln -sf "$_amd_dir/amd-smi" "$_combined_dir/amd-smi" +_result=$(run_func "$_combined_dir" "-1") +assert_eq "CUDA+ROCm with CVD=-1 -> rocm6.4" "https://download.pytorch.org/whl/rocm6.4" "$_result" +rm -rf "$_cuda_dir" "$_amd_dir" "$_combined_dir" + +# 37) CUDA_VISIBLE_DEVICES=0 (a visible device) must NOT hide the GPU +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir" "0") +assert_eq "CVD=0 keeps NVIDIA -> cu128" "https://download.pytorch.org/whl/cu128" "$_result" +rm -rf "$_dir" + +# 38) Whitespace-padded "-1" still hides the GPU +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir" " -1 ") +assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + rm -f "$_FUNC_FILE" rm -rf "$_FAKE_SMI_DIR" rm -rf "$_TOOLS_DIR" diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py new file mode 100644 index 0000000000..83c2d962e8 --- /dev/null +++ b/tests/studio/install/test_cuda_repair.py @@ -0,0 +1,248 @@ +"""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"])) diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py new file mode 100644 index 0000000000..983969a4d3 --- /dev/null +++ b/tests/studio/install/test_gpu_detection_followups.py @@ -0,0 +1,480 @@ +"""Tests for the GPU-detection follow-ups to PR 6174. + +PR 6174 made NVIDIA take precedence and added a /proc/driver/nvidia/gpus +fallback in install.sh and studio/install_python_stack.py. These tests cover the +same hardening ported to the llama.cpp prebuilt installer +(studio/install_llama_prebuilt.py) and the Studio shell setup (studio/setup.sh): + + * detect_host() recognises NVIDIA via /proc/driver/nvidia/gpus when nvidia-smi + is unavailable, and skips ROCm probing when NVIDIA is usable. + * setup.sh routes through a timeout-bounded NVIDIA probe with a /proc fallback + and only selects a CUDA/ROCm source build when the matching GPU is detected. + +All tests use mocks or source-level assertions -- no GPU, network, or real +nvidia-smi/rocminfo invocation. +""" + +import importlib.util +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] + +# Load studio/install_llama_prebuilt.py the same way the sibling suite does. +_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" +_SPEC = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt_followups", _MODULE_PATH +) +assert _SPEC is not None and _SPEC.loader is not None +prebuilt_mod = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = prebuilt_mod +_SPEC.loader.exec_module(prebuilt_mod) + +detect_host = prebuilt_mod.detect_host +_apply_host_overrides = prebuilt_mod._apply_host_overrides + +SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh" + + +def _make_run_capture(rocminfo_stdout: str = ""): + """Return a fake run_capture: rocminfo reports rocminfo_stdout, everything + else (nvidia-smi, amd-smi) returns empty so only the patched probes matter.""" + + def _run_capture(cmd, *args, **kwargs): + exe = str(cmd[0]) if cmd else "" + result = MagicMock() + if exe.endswith("rocminfo"): + result.returncode = 0 + result.stdout = rocminfo_stdout + else: + result.returncode = 1 + result.stdout = "" + result.stderr = "" + return result + + return _run_capture + + +def _run_detect_host( + *, + machine: str = "x86_64", + system: str = "Linux", + which_map: dict | None = None, + proc_dir_entries: list | None = None, + rocminfo_stdout: str = "", + env: dict | None = None, +): + """Drive detect_host() against a fully synthetic host.""" + which_map = which_map or {} + proc_dir_entries = proc_dir_entries if proc_dir_entries is not None else [] + + real_isdir = prebuilt_mod.os.path.isdir + real_listdir = prebuilt_mod.os.listdir + proc_path = "/proc/driver/nvidia/gpus" + + def fake_isdir(p): + if str(p) == proc_path: + return bool(proc_dir_entries) + return real_isdir(p) + + def fake_listdir(p): + if str(p) == proc_path: + if not proc_dir_entries: + raise OSError("no such dir") + return list(proc_dir_entries) + return real_listdir(p) + + patches = [ + patch.object(prebuilt_mod.platform, "system", return_value = system), + patch.object(prebuilt_mod.platform, "machine", return_value = machine), + patch.object(prebuilt_mod.platform, "mac_ver", return_value = ("", ("", "", ""), "")), + patch.object(prebuilt_mod.shutil, "which", side_effect = lambda n: which_map.get(n)), + patch.object(prebuilt_mod, "run_capture", side_effect = _make_run_capture(rocminfo_stdout)), + patch.object(prebuilt_mod.os.path, "isdir", side_effect = fake_isdir), + patch.object(prebuilt_mod.os, "listdir", side_effect = fake_listdir), + patch.object(prebuilt_mod.os, "access", return_value = False), + patch.dict(prebuilt_mod.os.environ, env or {}, clear = False), + ] + for p in patches: + p.start() + try: + # Ensure CUDA_VISIBLE_DEVICES does not leak in from the test host unless + # the scenario sets it explicitly. + if env is None or "CUDA_VISIBLE_DEVICES" not in env: + prebuilt_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None) + return detect_host() + finally: + for p in patches: + p.stop() + + +# ── install_llama_prebuilt.detect_host(): /proc NVIDIA fallback ────────────── + + +class TestDetectHostProcFallback: + def test_proc_fallback_marks_physical_nvidia_when_smi_absent(self): + """No nvidia-smi, but /proc/driver/nvidia/gpus is populated -> NVIDIA.""" + host = _run_detect_host( + which_map = {}, # nvidia-smi resolves to None + proc_dir_entries = ["0000:01:00.0"], + ) + assert host.has_physical_nvidia is True + + def test_proc_fallback_has_usable_nvidia_when_devices_visible(self): + """Default CUDA_VISIBLE_DEVICES (unset) -> visible tokens non-empty -> usable.""" + host = _run_detect_host( + which_map = {}, + proc_dir_entries = ["0000:01:00.0"], + ) + assert host.has_usable_nvidia is True + + def test_proc_fallback_not_usable_when_devices_hidden(self): + """CUDA_VISIBLE_DEVICES='' hides all GPUs -> physical yes, usable no.""" + host = _run_detect_host( + which_map = {}, + proc_dir_entries = ["0000:01:00.0"], + env = {"CUDA_VISIBLE_DEVICES": ""}, + ) + assert host.has_physical_nvidia is True + assert host.has_usable_nvidia is False + + def test_empty_proc_dir_does_not_mark_nvidia(self): + """A driver dir that exists but is empty must not assert a GPU.""" + host = _run_detect_host(which_map = {}, proc_dir_entries = []) + assert host.has_physical_nvidia is False + + def test_proc_fallback_is_linux_only(self): + """The /proc fallback must not run on Windows (path is Linux-only).""" + host = _run_detect_host( + system = "Windows", + machine = "amd64", + which_map = {}, + proc_dir_entries = ["0000:01:00.0"], + ) + assert host.has_physical_nvidia is False + + +# ── install_llama_prebuilt.detect_host(): NVIDIA precedence over ROCm ──────── + + +class TestDetectHostNvidiaPrecedence: + def test_rocm_probe_skipped_when_proc_nvidia_present(self): + """rocminfo reports gfx1100, but a proc-detected NVIDIA GPU wins.""" + host = _run_detect_host( + which_map = {"rocminfo": "/usr/bin/rocminfo"}, + proc_dir_entries = ["0000:01:00.0"], + rocminfo_stdout = " Name: gfx1100\n", + ) + assert host.has_usable_nvidia is True + assert host.has_rocm is False + + def test_rocm_detected_when_no_nvidia(self): + """With no NVIDIA signal at all, rocminfo gfx1100 -> has_rocm True.""" + host = _run_detect_host( + which_map = {"rocminfo": "/usr/bin/rocminfo"}, + proc_dir_entries = [], + rocminfo_stdout = " Name: gfx1100\n", + ) + assert host.has_usable_nvidia is False + assert host.has_rocm is True + + +# ── _apply_host_overrides: forwarded --rocm-gfx / --has-rocm still win ─────── + + +class TestOverridesStillWin: + def test_forwarded_gfx_forces_rocm_on_non_nvidia_host(self): + host = _run_detect_host(which_map = {}, proc_dir_entries = []) + assert host.has_rocm is False + overridden = _apply_host_overrides(host, override_rocm_gfx = "gfx1100") + assert overridden.has_rocm is True + assert overridden.rocm_gfx_target == "gfx1100" + + def test_override_has_rocm_forces_rocm(self): + host = _run_detect_host(which_map = {}, proc_dir_entries = []) + overridden = _apply_host_overrides(host, override_has_rocm = True) + assert overridden.has_rocm is True + + def test_force_cpu_drops_nvidia_attributes(self): + host = _run_detect_host(which_map = {}, proc_dir_entries = ["0000:01:00.0"]) + assert host.has_usable_nvidia is True + overridden = _apply_host_overrides(host, force_cpu = True) + assert overridden.has_usable_nvidia is False + assert overridden.has_physical_nvidia is False + assert overridden.has_rocm is False + + +# ── setup.sh source-level guarantees ──────────────────────────────────────── + + +class TestSetupShHardening: + @pytest.fixture(scope = "class") + def setup_src(self) -> str: + return SETUP_SH.read_text(encoding = "utf-8") + + def test_has_usable_nvidia_helper_exists(self, setup_src): + assert "_setup_has_usable_nvidia_gpu()" in setup_src + + def test_helper_uses_proc_fallback(self, setup_src): + start = setup_src.find("_setup_has_usable_nvidia_gpu()") + end = setup_src.find("\n}", start) + body = setup_src[start:end] + assert ( + "/proc/driver/nvidia/gpus" in body + ), "_setup_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus" + + def test_gpu_summary_uses_helper(self, setup_src): + assert "if _setup_has_usable_nvidia_gpu; then" in setup_src + + def test_timeout_wrapper_exists(self, setup_src): + start = setup_src.find("_setup_run_smi()") + assert start >= 0, "_setup_run_smi timeout wrapper must exist" + end = setup_src.find("\n}", start) + body = setup_src[start:end] + assert "timeout 10" in body + assert "command -v timeout" in body + + def test_cuda_source_build_gated_on_usable_nvidia(self, setup_src): + """The nvcc source-build search must be gated on _setup_nvidia_usable. + + The hidden-GPU policy (CUDA_VISIBLE_DEVICES=""/-1) lives inside + _setup_has_usable_nvidia_gpu, so the gate itself only needs the flag. + """ + anchor = setup_src.find('NVCC_PATH=""\n') + assert anchor >= 0 + window = setup_src[anchor : anchor + 700] + assert ( + 'if [ "$_setup_nvidia_usable" = true ]' in window + ), "CUDA toolkit search must require a usable NVIDIA GPU, not just nvcc" + + def test_nvidia_helper_honours_hidden_cvd(self, setup_src): + """_setup_has_usable_nvidia_gpu must consult the hidden-CVD helper so + CUDA_VISIBLE_DEVICES=""/-1 suppresses NVIDIA before the AMD probes are + gated (mixed hosts steered to the AMD card keep the ROCm route).""" + assert "_setup_cvd_hides_nvidia()" in setup_src + start = setup_src.find("_setup_has_usable_nvidia_gpu() {") + end = setup_src.find("\n}", start) + body = setup_src[start:end] + assert "_setup_cvd_hides_nvidia" in body + + def test_rocm_source_build_gated_on_amd_detected(self, setup_src): + """The hipcc source-build search must be gated on _setup_amd_detected.""" + anchor = setup_src.find('ROCM_HIPCC=""') + assert anchor >= 0 + window = setup_src[anchor : anchor + 400] + assert ( + '[ "$_setup_amd_detected" = true ]' in window + ), "ROCm toolkit search must require a detected AMD GPU, not just hipcc" + + def test_compute_cap_probe_timeout_wrapped(self, setup_src): + assert "_setup_run_smi nvidia-smi --query-gpu=compute_cap" in setup_src + + def test_driver_version_probe_timeout_wrapped(self, setup_src): + start = setup_src.find("_cuda_driver_max_version()") + end = setup_src.find("\n}", start) + body = setup_src[start:end] + assert "_setup_run_smi nvidia-smi" in body + + +# TEST: install.sh -- UNSLOTH_TORCH_BACKEND classified on the final path segment + + +class TestBackendExportLeafClassification: + """A custom UNSLOTH_PYTORCH_MIRROR whose base path contains "rocm" or + "gfx" must not mislabel a cu*/cpu index as ROCm; classification uses the + final path segment of TORCH_INDEX_URL only.""" + + @pytest.fixture(scope = "class") + def install_src(self) -> str: + return (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + + def test_export_block_uses_leaf(self, install_src): + anchor = install_src.find("_torch_index_leaf=") + assert anchor >= 0, "backend export must classify on the final path segment" + window = install_src[anchor : anchor + 500] + assert 'export UNSLOTH_TORCH_BACKEND="rocm"' in window + assert 'export UNSLOTH_TORCH_BACKEND="cpu"' in window + assert 'export UNSLOTH_TORCH_BACKEND="cuda"' in window + + def test_leaf_classification_behaviour(self, tmp_path): + import subprocess as sp + + script = tmp_path / "leaf.sh" + src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + anchor = src.find("_torch_index_leaf=") + block = src[anchor : src.find("esac", anchor) + 4] + # Drive the extracted block with adversarial mirror URLs. + script.write_text( + "#!/bin/sh\n" + 'TORCH_INDEX_URL="$1"\n' + block + "\n" + 'printf "%s" "$UNSLOTH_TORCH_BACKEND"\n' + ) + cases = { + "https://download.pytorch.org/whl/cu128": "cuda", + "https://download.pytorch.org/whl/cpu": "cpu", + "https://download.pytorch.org/whl/rocm6.4": "rocm", + "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/": "rocm", + "https://repo.amd.com/rocm/whl/gfx1151/": "rocm", + "https://mirror.local/rocm-cache/cu128": "cuda", + "https://mirror.local/gfx-cache/cpu": "cpu", + } + for url, expected in cases.items(): + out = sp.run( + ["sh", str(script), url], capture_output = True, text = True, timeout = 30 + ).stdout.strip() + assert out == expected, f"{url} classified as {out!r}, expected {expected!r}" + + +# TEST: CUDA_VISIBLE_DEVICES=""/-1 hides NVIDIA in every usable-GPU helper + + +_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" +_STACK_SPEC = importlib.util.spec_from_file_location( + "studio_install_python_stack_followups", _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) + + +def _stack_nvidia_usable(cvd): + """Drive install_python_stack._has_usable_nvidia_gpu with a mocked + nvidia-smi that always reports a GPU; cvd = None removes the env var.""" + + def fake_run(cmd, *args, **kwargs): + result = MagicMock() + result.returncode = 0 + result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n" + return result + + env = {} if cvd is None else {"CUDA_VISIBLE_DEVICES": cvd} + with ( + patch.object( + stack_mod.shutil, + "which", + side_effect = lambda n: "/usr/bin/nvidia-smi" if n == "nvidia-smi" else None, + ), + patch.object(stack_mod.subprocess, "run", side_effect = fake_run), + patch.dict(stack_mod.os.environ, env, clear = False), + ): + if cvd is None: + stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None) + return stack_mod._has_usable_nvidia_gpu() + + +class TestHiddenCvdNotUsable: + """CUDA_VISIBLE_DEVICES set to "" or "-1" deliberately hides every NVIDIA + device (mixed AMD+NVIDIA hosts steering work to the AMD card). All three + _has_usable_nvidia_gpu implementations (install_python_stack.py, install.sh, + setup.sh) must report the GPU as not usable so the AMD/CPU routes run, + matching install_llama_prebuilt.py's has_usable_nvidia.""" + + def test_python_unset_cvd_is_usable(self): + assert _stack_nvidia_usable(None) is True + + def test_python_empty_cvd_not_usable(self): + assert _stack_nvidia_usable("") is False + + def test_python_minus_one_not_usable(self): + assert _stack_nvidia_usable("-1") is False + + def test_python_padded_minus_one_not_usable(self): + assert _stack_nvidia_usable(" -1 ") is False + + def test_python_explicit_device_is_usable(self): + assert _stack_nvidia_usable("0") is True + + def test_python_device_list_is_usable(self): + assert _stack_nvidia_usable("0,1") is True + + def test_hidden_nvidia_restores_rocm_detection(self): + """Mixed host, NVIDIA hidden via CVD=-1, rocminfo reports gfx1100: + _has_rocm_gpu must proceed past the NVIDIA guard and return True + (before this fix the guard ignored CVD and blocked ROCm).""" + + def fake_run(cmd, *args, **kwargs): + result = MagicMock() + result.returncode = 0 + exe = str(cmd[0]) + if exe.endswith("rocminfo"): + result.stdout = " Name: gfx1100\n" + else: + result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n" + return result + + which_map = { + "rocminfo": "/usr/bin/rocminfo", + "nvidia-smi": "/usr/bin/nvidia-smi", + } + with ( + patch.object(stack_mod.shutil, "which", side_effect = which_map.get), + patch.object(stack_mod.subprocess, "run", side_effect = fake_run), + patch.dict(stack_mod.os.environ, {"CUDA_VISIBLE_DEVICES": "-1"}, clear = False), + ): + assert stack_mod._has_rocm_gpu() is True + + @staticmethod + def _run_sh_helper(tmp_path, src: str, fn_names: list, cvd): + """Extract shell functions, run the usable-GPU one against a fake + nvidia-smi, and return "usable"/"not_usable".""" + import os as _os + import subprocess as sp + + blocks = [] + for name in fn_names: + start = src.find(f"{name}() {{") + assert start >= 0, f"{name} missing" + end = src.find("\n}", start) + 2 + blocks.append(src[start:end]) + fake_bin = tmp_path / "bin" + fake_bin.mkdir(exist_ok = True) + smi = fake_bin / "nvidia-smi" + smi.write_text("#!/bin/sh\necho 'GPU 0: NVIDIA Fake (UUID: GPU-x)'\n") + smi.chmod(0o755) + script = tmp_path / "probe.sh" + script.write_text( + "#!/bin/sh\n" + "\n".join(blocks) + "\n" + f"if {fn_names[-1]}; then echo usable; else echo not_usable; fi\n" + ) + env = dict(_os.environ) + env["PATH"] = f"{fake_bin}:{env['PATH']}" + if cvd is None: + env.pop("CUDA_VISIBLE_DEVICES", None) + else: + env["CUDA_VISIBLE_DEVICES"] = cvd + return sp.run( + ["sh", str(script)], capture_output = True, text = True, timeout = 30, env = env + ).stdout.strip() + + @pytest.mark.parametrize( + "cvd, expected", + [(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")], + ) + def test_install_sh_helper_cvd(self, tmp_path, cvd, expected): + src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + out = self._run_sh_helper( + tmp_path, + src, + ["_run_bounded", "_cvd_hides_nvidia", "_has_usable_nvidia_gpu"], + cvd, + ) + assert out == expected + + @pytest.mark.parametrize( + "cvd, expected", + [(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")], + ) + def test_setup_sh_helper_cvd(self, tmp_path, cvd, expected): + src = SETUP_SH.read_text(encoding = "utf-8") + out = self._run_sh_helper( + tmp_path, + src, + ["_setup_run_smi", "_setup_cvd_hides_nvidia", "_setup_has_usable_nvidia_gpu"], + cvd, + ) + assert out == expected diff --git a/tests/studio/install/test_probe_timeouts.py b/tests/studio/install/test_probe_timeouts.py new file mode 100644 index 0000000000..acea0ed34d --- /dev/null +++ b/tests/studio/install/test_probe_timeouts.py @@ -0,0 +1,202 @@ +"""Tests that NVIDIA probes in the installers are bounded by a timeout. + +Covers audit findings 5 and 6: a wedged nvidia-smi must not hang the installer, +and the Windows probe must require a real GPU listing (not just exit code 0). + +Source-level assertions verify the guards are present in install.sh / install.ps1 +/ setup.ps1; one behavioral shell test confirms the bash helper actually returns +within the timeout when nvidia-smi hangs. +""" + +import os +import shutil +import stat +import subprocess +import sys +import tempfile +from pathlib import Path + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] +INSTALL_SH = PACKAGE_ROOT / "install.sh" +INSTALL_PS1 = PACKAGE_ROOT / "install.ps1" +SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1" + + +def _extract_sh_function_body(source: str, name: str) -> str: + """Return a shell function body from `source` by brace matching.""" + needle = f"{name}() {{" + start = source.find(needle) + if start < 0: + return "" + depth = 0 + i = start + len(needle) - 1 + n = len(source) + while i < n: + ch = source[i] + if ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return source[start : i + 1] + i += 1 + return source[start:] + + +# ── install.sh: _run_bounded helper and its use at every nvidia-smi call ── + + +class TestInstallShBoundedProbe: + def _src(self) -> str: + return INSTALL_SH.read_text(encoding = "utf-8") + + def test_run_bounded_helper_defined(self): + body = _extract_sh_function_body(self._src(), "_run_bounded") + assert body, "install.sh must define a _run_bounded helper" + assert ( + "command -v timeout" in body + ), "_run_bounded must check for the `timeout` binary before using it" + assert "timeout 10" in body, "_run_bounded must apply a 10s timeout" + # Must fall back to running unbounded when `timeout` is unavailable + # (e.g. macOS) so semantics are unchanged there. + assert ( + "else" in body and '"$@"' in body + ), "_run_bounded must run the command unbounded when `timeout` is absent" + + def test_nvidia_smi_dash_l_probe_is_bounded(self): + body = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu") + assert body, "install.sh must define _has_usable_nvidia_gpu" + # The -L probe must go through the bounded runner, not call nvidia-smi raw. + assert ( + '_run_bounded "$_nvsmi" -L' in body + ), "_has_usable_nvidia_gpu must run nvidia-smi -L through _run_bounded" + # The /proc fallback from PR 6174 must still be present. + assert "/proc/driver/nvidia" in body + + def test_cuda_version_parse_is_bounded(self): + body = _extract_sh_function_body(self._src(), "get_torch_index_url") + assert body, "install.sh must define get_torch_index_url" + assert ( + "_run_bounded" in body + ), "get_torch_index_url CUDA-version parse must run nvidia-smi through _run_bounded" + # The locale must be forced without depending on `env` being on PATH. + assert "LC_ALL=C" in body + # _nvidia_detected gating from PR 6174 must remain. + assert "_nvidia_detected" in body + + def test_no_unbounded_nvidia_smi_invocation_remains(self): + """Every nvidia-smi *execution* in install.sh goes through _run_bounded. + + `command -v nvidia-smi` and `-x /usr/bin/nvidia-smi` are resolution + checks, not executions, and are allowed. An execution looks like + `"$_nvsmi" ...` / `$_smi ...` / `nvidia-smi -L`. + """ + body_nvidia = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu") + body_torch = _extract_sh_function_body(self._src(), "get_torch_index_url") + # In _has_usable_nvidia_gpu the only execution of $_nvsmi must be bounded. + assert '"$_nvsmi" -L' not in body_nvidia.replace( + '_run_bounded "$_nvsmi" -L', "" + ), "found an unbounded nvidia-smi -L execution in _has_usable_nvidia_gpu" + # In get_torch_index_url the $_smi execution must be bounded. + assert ( + "LC_ALL=C $_smi" not in body_torch + ), "found an unbounded LC_ALL=C $_smi execution in get_torch_index_url" + + +# ── install.ps1 / setup.ps1: bounded, GPU-row-validated Windows probe ── + + +class TestPowerShellBoundedProbe: + @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1]) + def test_bounded_helper_present(self, path): + src = path.read_text(encoding = "utf-8") + assert ( + "function Invoke-NvidiaSmiBounded" in src + ), f"{path.name} must define Invoke-NvidiaSmiBounded" + assert ( + "WaitForExit($TimeoutSec * 1000)" in src + ), f"{path.name} bounded probe must use WaitForExit with a timeout" + # Kill + sentinel on timeout, mirroring Invoke-AmdSmiNoElevate. + assert ( + "$proc.Kill()" in src and "124" in src + ), f"{path.name} must kill nvidia-smi and signal a timeout exit code" + + @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1]) + def test_probe_requires_gpu_row(self, path): + src = path.read_text(encoding = "utf-8") + assert ( + "function Test-NvidiaSmiHasGpu" in src + ), f"{path.name} must define Test-NvidiaSmiHasGpu" + assert "@('-L')" in src, f"{path.name} must probe nvidia-smi with -L" + assert ( + "^GPU\\s+\\d+:" in src + ), f"{path.name} must require a 'GPU :' data row, not just exit code 0" + + @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1]) + def test_detection_uses_validated_probe(self, path): + src = path.read_text(encoding = "utf-8") + # The exit-code-only pattern must be gone from the detection block. + assert ( + "& $nvSmiCmd.Source *> $null" not in src + ), f"{path.name} must not use the exit-code-only nvidia-smi probe" + assert ( + "Test-NvidiaSmiHasGpu $nvSmiCmd.Source" in src + ), f"{path.name} PATH probe must use Test-NvidiaSmiHasGpu" + assert ( + "Test-NvidiaSmiHasGpu $p" in src + ), f"{path.name} hardcoded-path fallback must use Test-NvidiaSmiHasGpu" + + +# ── Behavioral: a hanging nvidia-smi must not hang _has_usable_nvidia_gpu ── + + +def _have_timeout() -> bool: + return shutil.which("timeout") is not None + + +@pytest.mark.skipif(not _have_timeout(), reason = "`timeout` binary not available") +def test_has_usable_nvidia_gpu_returns_under_timeout(): + """Extract _run_bounded + _has_usable_nvidia_gpu, point them at a fake + nvidia-smi that sleeps 30s, and assert the probe returns well under that. + """ + src = INSTALL_SH.read_text(encoding = "utf-8") + helper = _extract_sh_function_body(src, "_run_bounded") + fn = _extract_sh_function_body(src, "_has_usable_nvidia_gpu") + assert helper and fn + + workdir = tempfile.mkdtemp(prefix = "pr6174_timeout_", dir = str(PACKAGE_ROOT.parent)) + try: + fake_dir = Path(workdir, "bin") + fake_dir.mkdir() + fake_smi = fake_dir / "nvidia-smi" + fake_smi.write_text("#!/bin/sh\nsleep 30\n") + fake_smi.chmod(fake_smi.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH) + + # Build a minimal PATH that includes the fake nvidia-smi plus the real + # `timeout`/`awk`/`ls` it needs. Use the fake dir first so it wins. + real_bins = {Path(shutil.which(c)).parent for c in ("timeout", "awk", "ls", "sh")} + path_env = os.pathsep.join([str(fake_dir)] + [str(p) for p in real_bins]) + + # Force the /proc fallback off so the result depends only on the probe, + # and so a host with real NVIDIA does not mask the timeout behaviour. + script = ( + f"{helper}\n{fn}\n" + "if _has_usable_nvidia_gpu; then echo DETECTED; else echo NONE; fi\n" + ) + proc = subprocess.run( + ["sh", "-c", script], + env = {"PATH": path_env}, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 20, # generous: the internal timeout is 10s, sleep is 30s + ) + # The probe must have returned (not hung). On this CI host /proc/driver/ + # nvidia/gpus is absent, so a timed-out smi yields NONE; on a real NVIDIA + # host the /proc fallback yields DETECTED. Either way it must not hang. + assert proc.stdout.strip() in {"NONE", "DETECTED"} + finally: + shutil.rmtree(workdir, ignore_errors = True)