ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Installer: GPU detection follow-ups after #6174 (poisoned venv repair, llama.cpp routing, probe bounds) (#6183)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-06-11 05:06:02 -07:00 • committed by GitHub
commit 2db9fad4b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 1439 additions and 74 deletions
Download patch file Download diff file Expand all files Collapse all files

55
install.ps1
View file

@ -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 <n>:" 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.

51
install.sh
View file

@ -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.

24
studio/install_llama_prebuilt.py
View file

@ -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: <number>' 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

155
studio/install_python_stack.py
View file

@ -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)

81
studio/setup.ps1
View file

@ -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 <n>:" 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])"

166
studio/setup.sh
View file

@ -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)"

51
tests/sh/test_get_torch_index_url.sh
View file

@ -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"

248
tests/studio/install/test_cuda_repair.py Normal file
View file

@ -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"]))

480
tests/studio/install/test_gpu_detection_followups.py Normal file
View file

@ -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

202
tests/studio/install/test_probe_timeouts.py Normal file
View file

@ -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 <n>:' 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)