From a4ad50e0242029e9cba01d1d891dcf07b29879e6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 19 Jun 2026 02:29:42 -0700 Subject: [PATCH] fix(prebuilt): harden nvidia-smi GPU detection against WSL GPU-PV slowness On Windows-on-ARM + NVIDIA the Studio install runs inside WSL2, where nvidia-smi is served over GPU-PV and can take far longer than its usual sub-second response when the host is under heavy CPU load (the concurrent pip / frontend / cmake work during install). detect_host probed nvidia-smi with a single 20s timeout; under that load it raised TimeoutExpired, the GPU was treated as ABSENT, and the host was misrouted to the ggml-org CPU prebuilt -> rejected on an NVIDIA host -> slow (and on thermal-limited laptops, risky) CUDA source build, even though a usable arm64 CUDA prebuilt was published. Add _nvidia_smi_capture(): retry the three detect_host nvidia-smi probes with a generous 60s per-attempt timeout. It is only reachable when nvidia-smi exists on PATH, so CPU-only hosts incur no extra wait. Measured: nvidia-smi took 42-59s under a -j20 build on an N1X; with the fix the probe rides it out and detect_host correctly reports has_usable_nvidia + compute_cap, so the CUDA prebuilt is selected (no source build). --- studio/install_llama_prebuilt.py | 34 ++++++++++++++++++++++++++++---- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 4edab0b8ab..e89b142cb1 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -2855,6 +2855,33 @@ def _pick_rocm_gfx_target(out: str) -> str | None: return _tokens[0] +def _nvidia_smi_capture( + command: list[str], + *, + attempts: int = 2, + timeout: int = 60, +) -> subprocess.CompletedProcess[str]: + """run_capture for nvidia-smi probes, hardened against transient slowness. + + nvidia-smi normally answers in well under a second, but under WSL2 GPU-PV it + can take far longer when the host is under heavy CPU load -- e.g. the + concurrent pip / frontend / cmake work during an `unsloth studio` install. + A single short timeout then raises TimeoutExpired, detect_host treats the + GPU as ABSENT, and the host is misrouted to a CPU prebuilt / slow source + build instead of the CUDA bundle it can actually use. Retry with a generous + per-attempt timeout. Only ever reached when nvidia-smi exists on PATH, so + CPU-only hosts never incur this wait. + """ + last_exc: Exception | None = None + for _attempt in range(max(1, attempts)): + try: + return run_capture(command, timeout = timeout) + except subprocess.TimeoutExpired as exc: + last_exc = exc + time.sleep(2) + raise last_exc if last_exc is not None else RuntimeError("nvidia-smi capture failed") + + def detect_host() -> HostInfo: system = platform.system() machine = platform.machine().lower() @@ -2880,7 +2907,7 @@ def detect_host() -> HostInfo: # container leftovers), which would otherwise misclassify an AMD # ROCm host as NVIDIA and short-circuit the ROCm path. try: - listing = run_capture([nvidia_smi, "-L"], timeout = 20) + listing = _nvidia_smi_capture([nvidia_smi, "-L"]) gpu_lines = [line for line in listing.stdout.splitlines() if line.startswith("GPU ")] if gpu_lines: has_physical_nvidia = True @@ -2889,7 +2916,7 @@ def detect_host() -> HostInfo: pass try: - result = run_capture([nvidia_smi], timeout = 20) + result = _nvidia_smi_capture([nvidia_smi]) merged = "\n".join(part for part in (result.stdout, result.stderr) if part) # Newer NVIDIA drivers (e.g. 610.x on Windows) print # "CUDA UMD Version: X.Y" instead of the legacy @@ -2907,13 +2934,12 @@ def detect_host() -> HostInfo: pass try: - caps = run_capture( + caps = _nvidia_smi_capture( [ nvidia_smi, "--query-gpu=index,uuid,compute_cap", "--format=csv,noheader", ], - timeout = 20, ) visible_gpu_rows: list[tuple[str, str, str]] = [] for raw in caps.stdout.splitlines():