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

tests/studio: lock in Windows GPU detection fix (#5106) with a synthetic CI test (#5376)

Browse source 
* tests/studio: end-to-end Windows GPU detection mock test (#5106)

Locks in the combined fix from #5322 + #5324 with a synthetic
Windows scenario that CI runners without GPUs can execute. The
test packs the real PyPI win_amd64 wheel layouts (cu12 modular and
the new unsuffixed cu13 nvidia/cu13/bin/x86_64 layout) plus the
exact filename set of the upstream b9103 cudart-llama-bin-win-cuda
bundles, then mocks nvidia-smi output and asserts that:

 * Studio's nvidia-smi probe parses the CSV and reports the GPU.
 * After PR #5322 the install_dir/build/bin/Release/ tree contains
   all three cudart bundle DLLs alongside llama-server.exe.
 * After PR #5324 the PATH built by start_llama_server's win32
   branch lists pip nvidia + torch/lib dirs in addition to the
   binary_dir.
 * cudart64_X.dll, cublas64_X.dll, and cublasLt64_X.dll are
   each reachable from at least one PATH entry, with cudart
   specifically reachable from BOTH the install dir and a pip
   nvidia dir (defence in depth).
 * Bare venvs without pip nvidia wheels still work via #5322's
   binary_dir drop; pre-#5322 installs still work via #5324's
   PATH augmentation.
 * A reconstructed pre-PR scenario (cudart absent from binary_dir
   and pip dirs not on PATH) leaves cudart unreachable, confirming
   the test would catch a future regression.

Bonus housekeeping in studio/install_llama_prebuilt.py: drop the
pointless f-prefix on the literal "llama-" in the
windows_cuda_attempts pairing guard (no behaviour change; lint
nit flagged in the post-merge review).

The mocks model real artifact contents I verified empirically:
 * pip download nvidia-cuda-runtime --platform win_amd64
   produces nvidia/cu13/bin/x86_64/cudart64_13.dll.
 * unzip on the b9103 cudart-llama-bin-win-cuda-13.1-x64.zip
   produces exactly cudart64_13.dll + cublas64_13.dll +
   cublasLt64_13.dll, no executables.
 * objdump -p on the b9103 ggml-cuda.dll shows a static PE
   import on cublas64_13.dll (the root cause of #5106 when
   cublas64_13.dll is unreachable).

Refs #5106 #5322 #5324

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test_5106_windows_gpu_detection_mock: don't shadow real httpx

This file's name sorts before every other file in studio/backend/tests/
(starts with the digit '5'), so pytest collects it first. The previous
``sys.modules.setdefault("httpx", _httpx_stub)`` ran before any other
test imported real httpx, which meant the stub permanently shadowed
the real module for the rest of the collection. Tests that did
``from httpx import HTTPError, Response`` (test_anthropic_messages,
test_browse_folders_route, test_training_*, etc) then failed at
collection with ``ImportError: cannot import name 'HTTPError'``
because the stub did not define those names. The existing
test_llama_cpp_windows_nvidia_path.py did not trigger the same issue
because it sorts after test_a* / test_b* / etc, by which point the
real httpx has already been imported and setdefault is a no-op.

Switch the stub installation to ``importlib.util.find_spec(name) is
None`` so we only fall back to the stub when the real module truly is
not installed. Backend CI installs httpx, structlog, and the
studio/backend/loggers package is reachable via the sys.path
augmentation a few lines above, so on CI all three find_spec calls
succeed and no stubs are installed at all.

Also add HTTPError and Response to the stub module for the offline
case, so anyone running this test outside CI with httpx absent still
gets a stub that satisfies the broader test suite's imports.

Refs #5106

* test_5106 + llama_cpp: extract win32 PATH helper and harden the regression test

Follow-up to PR #5376's review feedback. Three real findings from the
bot reviewers, plus one stale one.

1. (codex P2 line 201, gemini medium line 209) The regression test's
   _build_path_dirs_like_start_llama_server hand-copied the win32
   branch of LlamaCppBackend.start_llama_server, so a future drop or
   reorder of _windows_pip_nvidia_dll_dirs(sys.prefix) in production
   would have passed the test silently.

   Extract a new staticmethod LlamaCppBackend._build_windows_path_dirs
   (binary_dir, prefix, cuda_path). Production start_llama_server now
   calls this helper. The test's wrapper is reduced to a one-line
   delegate that forwards to the staticmethod, so the regression
   asserts against the exact production logic instead of a parallel
   copy of it.

2. (codex P2 line 245) test_nvidia_smi_probe_reports_synthetic_gpu did
   not clear CUDA_VISIBLE_DEVICES. On a shared GPU runner with the
   variable set in the parent shell, _get_gpu_free_memory() filters
   the mocked CSV and returns [] or falls through to the torch
   fallback. Cleared CUDA_VISIBLE_DEVICES and NVIDIA_VISIBLE_DEVICES
   via monkeypatch.delenv(..., raising=False).

3. (codex P2 line 66) _maybe_stub gated on importlib.util.find_spec
   ("loggers"), which returns a spec because studio/backend/loggers/
   is on sys.path. But the actual import chain loads
   loggers/handlers.py which does `from fastapi import Request,
   Response` at module load. In a lightweight env without fastapi
   installed, the stub never lands and `from core.inference.llama_cpp
   import LlamaCppBackend` raises during collection. Switched
   _maybe_stub to a real import attempt under try / except ImportError
   so the stub falls into place when the package is discoverable but
   not importable. CI has fastapi so this is purely a developer-
   machine ergonomics fix.

The fourth comment (codex P1 line 85 "Keep the httpx stub from leaking
across tests") was already addressed by 7437e735, which replaced the
unconditional sys.modules.setdefault with the find_spec-gated
_maybe_stub. No code change needed.

Production behaviour is unchanged: _build_windows_path_dirs returns
exactly the same ordering start_llama_server used inline
([binary_dir, *pip_dirs, cuda_bin?, cuda_bin_x64?]).

Verification (run inside studio/backend):
  pytest tests/test_5106_windows_gpu_detection_mock.py -v
    -> 10 passed
  pytest tests/test_llama_cpp_*.py tests/test_llama_server_args.py
       tests/test_5106_windows_gpu_detection_mock.py -q
    -> 171 passed
  CUDA_VISIBLE_DEVICES=1 pytest tests/test_5106_windows_gpu_detection_mock.py::TestWindowsGpuDetectionAfter5106Fix::test_nvidia_smi_probe_reports_synthetic_gpu
    -> 1 passed

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Rename Windows GPU detection test to a generic filename and trim comments

- studio/backend/tests/test_5106_windows_gpu_detection_mock.py
  -> studio/backend/tests/test_windows_gpu_detection_mock.py
  The file is the generic regression suite for Windows GPU detection;
  encoding the issue number in the filename is noise.
- Shorten module docstring, helper docstrings, per-test docstrings and
  inline comments in the renamed test file. No behaviour change,
  all 10 cases still pass.
- Shorten the _build_windows_path_dirs docstring in
  studio/backend/core/inference/llama_cpp.py and update the test-path
  reference; trim the win32 call-site comment to one line.

Local verification:
- pytest studio/backend/tests/test_windows_gpu_detection_mock.py -- 10 passed.
- pytest studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
  studio/backend/tests/test_llama_server_args.py
  studio/backend/tests/test_windows_gpu_detection_mock.py -- 110 passed.

* Studio: harden _wait_for_health against transient httpx ReadError

The probe loop in LlamaCppBackend._wait_for_health only caught
ConnectError and TimeoutException. On Windows, when llama-server.exe
accepts the TCP probe and then dies before sending HTTP headers, the
peer process RST closes the socket. httpx maps this to ReadError
("WinError 10054 -- An existing connection was forcibly closed by the
remote host"), which fell through the except clause and bubbled out of
_wait_for_health, the routes/inference.py load_model handler, and back
to /api/inference/load as an opaque 500.

The crash diagnostic Studio actually wants to surface lives on the
self._process.poll() branch at the top of the loop body: "llama-server
exited with code X. Output: ...". We never reached that branch on the
WinError 10054 path because the very first probe blew up.

Expand the except to also swallow ReadError and RemoteProtocolError so
the next 0.5-second iteration runs the poll() branch. Outcomes:
  * Process really died: structured exit-code + last-stdout log line.
  * Single transient probe blip: silently retried; load succeeds.

Adds studio/backend/tests/test_llama_cpp_wait_for_health.py with five
cases covering happy-path 200, transient ReadError + dead process,
RemoteProtocolError + dead process, ConnectError cycling until success,
and dead process before the first probe. The new cases would have
failed against the old except clause -- ReadError / RemoteProtocolError
would have propagated instead of returning False.

Found while triaging the Windows Studio GGUF CI flake on this PR's
5a6ddc34 push: llama-server.exe (b9203 prebuilt) crashed within 2.2 s of
launch on the GPU-less runner, and Studio reported "WinError 10054"
instead of an upstream-tag-attributable exit-code line.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
This commit is contained in:
Daniel Han 2026-05-18 00:06:01 -07:00 • committed by GitHub
commit a09e70e8be
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 576 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

43
studio/backend/core/inference/llama_cpp.py
View file

@ -1103,6 +1103,26 @@ class LlamaCppBackend:
_add(site_packages / "torch" / "lib")
return out
@staticmethod
def _build_windows_path_dirs(
binary_dir: str, prefix: str, cuda_path: str
) -> list[str]:
"""Ordered PATH entries the win32 branch of start_llama_server
prepends so llama-server.exe resolves cudart / cublas DLLs:
binary_dir, pip nvidia wheels, CUDA_PATH/bin, CUDA_PATH/bin/x64.
Extracted so test_windows_gpu_detection_mock asserts against
production logic, not a hand-copy. #5106."""
path_dirs = [binary_dir]
path_dirs.extend(LlamaCppBackend._windows_pip_nvidia_dll_dirs(prefix))
if cuda_path:
cuda_bin = os.path.join(cuda_path, "bin")
if os.path.isdir(cuda_bin):
path_dirs.append(cuda_bin)
cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64")
if os.path.isdir(cuda_bin_x64):
path_dirs.append(cuda_bin_x64)
return path_dirs
@staticmethod
def _select_gpus(
model_size_bytes: int,
@ -2631,23 +2651,12 @@ class LlamaCppBackend:
binary_dir = str(Path(binary).parent)
if sys.platform == "win32":
# CUDA DLLs (cudart64_X.dll, cublas64_X.dll, etc.) must
# be on PATH. Order: binary_dir, torch's pip-installed
# nvidia wheels, then a system CUDA toolkit. Pip wheels
# are the canonical source per Studio's install design
# (mirrors the Linux LD_LIBRARY_PATH block below) and
# CUDA_PATH covers users with a system toolkit. #5106.
path_dirs = [binary_dir]
path_dirs.extend(self._windows_pip_nvidia_dll_dirs(sys.prefix))
cuda_path = os.environ.get("CUDA_PATH", "")
if cuda_path:
cuda_bin = os.path.join(cuda_path, "bin")
if os.path.isdir(cuda_bin):
path_dirs.append(cuda_bin)
# Some CUDA installs put DLLs in bin\x64
cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64")
if os.path.isdir(cuda_bin_x64):
path_dirs.append(cuda_bin_x64)
# See _build_windows_path_dirs for ordering. #5106.
path_dirs = self._build_windows_path_dirs(
binary_dir,
sys.prefix,
os.environ.get("CUDA_PATH", ""),
)
existing_path = env.get("PATH", "")
env["PATH"] = ";".join(path_dirs) + ";" + existing_path
else:

156
studio/backend/tests/test_llama_cpp_wait_for_health.py Normal file
View file

@ -0,0 +1,156 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for LlamaCppBackend._wait_for_health resilience.
The probe loop must swallow transient httpx errors and fall through to
the subprocess.poll() branch so a crashed llama-server surfaces a
structured "exited with code X" log instead of bubbling an opaque
exception up to the /api/inference/load route.
"""
from __future__ import annotations
import sys
import types as _types
from pathlib import Path
from unittest import mock
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Match the stubbing pattern in sibling tests so the module imports in
# a lightweight env without fastapi.
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
sys.modules.setdefault("structlog", _types.ModuleType("structlog"))
import httpx # noqa: E402
from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
# Sibling tests in this directory install lightweight httpx stubs via
# sys.modules.setdefault. When collected together, our `httpx` symbol
# may be one of those stubs, which lacks `get`. Ensure the production
# code finds a working `httpx.get` and the standard exception types
# regardless of collection order by adding the missing attributes.
if not hasattr(httpx, "get"):
httpx.get = None # placeholder; every test below monkeypatches it
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadError",
"RemoteProtocolError",
"WriteError",
):
if not hasattr(httpx, _exc_name):
setattr(httpx, _exc_name, type(_exc_name, (Exception,), {}))
def _make_backend(port: int = 12345) -> LlamaCppBackend:
"""Build a barebones LlamaCppBackend instance with only the
attributes _wait_for_health touches. Bypasses __init__ so we do not
pull in the full subprocess + logging stack."""
b = LlamaCppBackend.__new__(LlamaCppBackend)
b._port = port
b._stdout_thread = None
b._stdout_lines = []
b._process = mock.Mock()
return b
class TestWaitForHealthResilience:
def test_returns_true_on_first_200(self, monkeypatch):
b = _make_backend()
b._process.poll.return_value = None
ok_resp = mock.Mock(status_code = 200)
monkeypatch.setattr(httpx, "get", lambda *a, **kw: ok_resp)
assert b._wait_for_health(timeout = 1.0, interval = 0.01) is True
def test_read_error_loops_to_subprocess_poll(self, monkeypatch):
"""WinError 10054 maps to httpx.ReadError. The loop must swallow
it and the next iteration must detect the dead subprocess via
poll() != None, returning False with a structured exit-code log
instead of bubbling the ReadError."""
b = _make_backend()
# First iteration: process alive (so we reach the httpx probe).
# Second iteration: process has exited (so we hit the structured
# exit-code branch and return False).
b._process.poll.side_effect = [None, 1]
b._process.returncode = 1
b._stdout_lines = ["llama-server: ggml-cuda.dll failed to load"]
def raise_read_error(*a, **kw):
raise httpx.ReadError("WinError 10054")
monkeypatch.setattr(httpx, "get", raise_read_error)
assert b._wait_for_health(timeout = 5.0, interval = 0.01) is False
# Both iterations of the loop ran -- the ReadError did not bubble.
assert b._process.poll.call_count >= 2
def test_remote_protocol_error_also_swallowed(self, monkeypatch):
"""Partial / malformed response on the probe (server crashed
mid-headers) raises RemoteProtocolError -- also non-fatal."""
b = _make_backend()
b._process.poll.side_effect = [None, -1]
b._process.returncode = -1
def raise_rpe(*a, **kw):
raise httpx.RemoteProtocolError("partial response")
monkeypatch.setattr(httpx, "get", raise_rpe)
assert b._wait_for_health(timeout = 5.0, interval = 0.01) is False
assert b._process.poll.call_count >= 2
def test_write_error_also_swallowed(self, monkeypatch):
"""Send-side socket failure mid-request raises WriteError --
same recovery path as ReadError."""
b = _make_backend()
b._process.poll.side_effect = [None, 1]
b._process.returncode = 1
def raise_we(*a, **kw):
raise httpx.WriteError("connection broken on write")
monkeypatch.setattr(httpx, "get", raise_we)
assert b._wait_for_health(timeout = 5.0, interval = 0.01) is False
assert b._process.poll.call_count >= 2
def test_connect_error_swallowed_until_success(self, monkeypatch):
"""Sanity: existing ConnectError swallowing still works -- the
loop retries until llama-server eventually answers 200."""
b = _make_backend()
b._process.poll.return_value = None
calls = {"n": 0}
ok_resp = mock.Mock(status_code = 200)
def cycling(*a, **kw):
calls["n"] += 1
if calls["n"] < 3:
raise httpx.ConnectError("not yet")
return ok_resp
monkeypatch.setattr(httpx, "get", cycling)
assert b._wait_for_health(timeout = 5.0, interval = 0.01) is True
assert calls["n"] >= 3
def test_dead_process_before_probe_returns_false(self, monkeypatch):
"""If poll() != None on entry, _wait_for_health must return
False immediately without calling httpx at all."""
b = _make_backend()
b._process.poll.return_value = 137
b._process.returncode = 137
b._stdout_lines = ["llama-server: out of memory"]
called = {"n": 0}
def should_not_be_called(*a, **kw):
called["n"] += 1
raise AssertionError("httpx.get must not run when subprocess is dead")
monkeypatch.setattr(httpx, "get", should_not_be_called)
assert b._wait_for_health(timeout = 5.0, interval = 0.01) is False
assert called["n"] == 0

393
studio/backend/tests/test_windows_gpu_detection_mock.py Normal file
View file

@ -0,0 +1,393 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Windows GPU-detection regression test on a synthetic layout.
The bug (#5106): on Windows without a system CUDA toolkit, the prebuilt
llama-server.exe could not LoadLibrary cudart64_X / cublas64_X /
cublasLt64_X, so ggml-cuda.dll's static import on cublas64_X.dll failed
and the model fell back to CPU even when nvidia-smi reported the GPU.
The fix:
* #5322 overlays upstream's paired cudart bundle into
install_dir/build/bin/Release/ next to llama-server.exe.
* #5324 prepends pip-installed nvidia/<pkg>/{bin,bin/x86_64,Library/
bin} and torch/lib to PATH when launching llama-server.exe.
CI has no GPU so nvidia-smi is mocked; everything else (resolver, PATH
builder, install layout) runs against a real filesystem.
"""
from __future__ import annotations
import os
import subprocess
import sys
import types as _types
import zipfile
from pathlib import Path
from unittest import mock
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Stub heavy deps only if they actually fail to import -- unconditional
# stubs would shadow the real module for sibling tests in this dir.
# Use try-import rather than find_spec: loggers/__init__.py re-exports
# handlers.get_logger, which does `from fastapi import Request,
# Response` at module load. find_spec("loggers") returns a spec even
# without fastapi, but the import then raises. CI has fastapi, so this
# is dev-machine ergonomics only.
import importlib as _importlib # noqa: E402
def _maybe_stub(name: str, builder):
try:
_importlib.import_module(name)
except ImportError:
sys.modules[name] = builder()
def _build_loggers_stub():
m = _types.ModuleType("loggers")
m.get_logger = lambda name: __import__("logging").getLogger(name)
return m
def _build_structlog_stub():
return _types.ModuleType("structlog")
def _build_httpx_stub():
m = _types.ModuleType("httpx")
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
"HTTPError",
):
setattr(m, _exc_name, type(_exc_name, (Exception,), {}))
m.Response = type("Response", (), {})
class _FakeTimeout:
def __init__(self, *a, **kw):
pass
m.Timeout = _FakeTimeout
m.Client = type(
"Client",
(),
{
"__init__": lambda self, **kw: None,
"__enter__": lambda self: self,
"__exit__": lambda self, *a: None,
},
)
return m
_maybe_stub("loggers", _build_loggers_stub)
_maybe_stub("structlog", _build_structlog_stub)
_maybe_stub("httpx", _build_httpx_stub)
from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
# Upstream b9103 cudart bundle: exactly these three DLLs per CUDA major,
# no executables, no subdirectories. Verified by direct unzip.
REAL_UPSTREAM_CUDART_BUNDLE = {
"12.4": ("cudart64_12.dll", "cublas64_12.dll", "cublasLt64_12.dll"),
"13.1": ("cudart64_13.dll", "cublas64_13.dll", "cublasLt64_13.dll"),
}
# PyPI win_amd64 wheel layouts, verified via `pip download ... --platform
# win_amd64` + `unzip -l`. Resolver only cares about directory structure.
REAL_PIP_NVIDIA_WHEEL_LAYOUTS = {
# Legacy cu-suffixed wheels
"nvidia/cuda_runtime/bin": ["cudart64_12.dll"],
"nvidia/cublas/bin": [
"cublas64_12.dll",
"cublasLt64_12.dll",
"nvblas64_12.dll",
],
"nvidia/cudnn/bin": [
"cudnn64_9.dll",
"cudnn_adv64_9.dll",
"cudnn_ops64_9.dll",
],
# Unsuffixed cu13 wheels
"nvidia/cu13/bin/x86_64": [
"cudart64_13.dll",
"cublas64_13.dll",
"cublasLt64_13.dll",
"nvblas64_13.dll",
],
}
def _populate_studio_venv(prefix: Path) -> None:
"""Lay out fake nvidia + torch wheels in <prefix>/Lib/site-packages
matching the real win_amd64 wheel layouts. Contents are stub bytes;
only directory structure matters."""
site = prefix / "Lib" / "site-packages"
for rel, dlls in REAL_PIP_NVIDIA_WHEEL_LAYOUTS.items():
d = site / Path(rel)
d.mkdir(parents = True, exist_ok = True)
for name in dlls:
(d / name).write_bytes(b"PE-stub")
# install_python_stack always installs torch alongside nvidia.
(site / "torch" / "lib").mkdir(parents = True, exist_ok = True)
for fn in ("c10.dll", "torch.dll", "torch_cpu.dll", "torch_python.dll"):
(site / "torch" / "lib" / fn).write_bytes(b"PE-stub")
def _populate_studio_install(install_dir: Path, runtime: str = "13.1") -> None:
"""Lay out install_dir/build/bin/Release/ as #5322 leaves it: main
archive payload + paired cudart bundle overlay."""
rel = install_dir / "build" / "bin" / "Release"
rel.mkdir(parents = True, exist_ok = True)
for fn in (
"llama-server.exe",
"llama-quantize.exe",
"llama-cli.exe",
"llama.dll",
"ggml.dll",
"ggml-base.dll",
"ggml-cuda.dll",
"mtmd.dll",
):
(rel / fn).write_bytes(b"PE-stub")
# The cudart overlay #5322 contributes.
for fn in REAL_UPSTREAM_CUDART_BUNDLE[runtime]:
(rel / fn).write_bytes(b"PE-stub")
def _build_path_dirs_like_start_llama_server(
binary_dir: Path, prefix: Path, cuda_path: str = ""
) -> list[str]:
"""Path-friendly wrapper around LlamaCppBackend._build_windows_path_dirs.
Asserting against the staticmethod (not a hand-copy) is the point:
if the win32 PATH order drops _windows_pip_nvidia_dll_dirs, tests fail."""
return LlamaCppBackend._build_windows_path_dirs(
str(binary_dir), str(prefix), cuda_path
)
def _mock_nvidia_smi_run(fake_output: str, returncode: int = 0) -> "mock._patch":
"""Patch subprocess.run so the nvidia-smi probe returns fake_output;
other subprocess.run calls pass through."""
real_run = subprocess.run
def fake_run(cmd, *args, **kwargs):
if isinstance(cmd, list) and cmd and "nvidia-smi" in cmd[0]:
return subprocess.CompletedProcess(
args = cmd, returncode = returncode, stdout = fake_output, stderr = ""
)
return real_run(cmd, *args, **kwargs)
return mock.patch("subprocess.run", side_effect = fake_run)
# --------------------------------------------------------------------- #
# Tests
# --------------------------------------------------------------------- #
class TestWindowsGpuDetectionAfter5106Fix:
"""End-to-end #5106 fix on a synthetic Windows layout. nvidia-smi
mocked; resolver, PATH builder and install layout exercised live."""
def test_nvidia_smi_probe_reports_synthetic_gpu(self, monkeypatch):
"""Probe parses CSV output and returns (index, free_mib)."""
# Clear inherited masks so the synthetic CSV is not filtered.
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("NVIDIA_VISIBLE_DEVICES", raising = False)
# The #5106 reporter's exact reproducer: RTX 4090, 22805 MiB.
fake_csv = "0, 22805\n"
with _mock_nvidia_smi_run(fake_csv):
gpus = LlamaCppBackend._get_gpu_free_memory()
assert gpus == [
(0, 22805)
], f"GPU probe failed to parse mocked nvidia-smi output: {gpus}"
def test_nvidia_smi_probe_respects_cuda_visible_devices(self, monkeypatch):
"""CUDA_VISIBLE_DEVICES=1 -> only GPU 1 visible."""
fake_csv = "0, 22805\n1, 24576\n2, 16384\n"
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1")
with _mock_nvidia_smi_run(fake_csv):
gpus = LlamaCppBackend._get_gpu_free_memory()
assert gpus == [(1, 24576)], gpus
def test_windows_install_dir_has_all_three_cudart_dlls(self, tmp_path):
"""All three bundle DLLs must land in install_dir/build/bin/
Release; missing any one breaks ggml-cuda.dll's PE import chain."""
install = tmp_path / "studio_install"
_populate_studio_install(install, runtime = "13.1")
rel = install / "build" / "bin" / "Release"
for fn in REAL_UPSTREAM_CUDART_BUNDLE["13.1"]:
assert (rel / fn).exists(), f"missing {fn} in {rel}"
assert (rel / "llama-server.exe").exists()
assert (rel / "ggml-cuda.dll").exists()
def test_resolver_finds_real_pypi_wheel_layouts(self, tmp_path):
"""Resolver must pick up every real-world wheel layout:
nvidia/<pkg>/bin, nvidia/<pkg>/bin/x86_64, torch/lib."""
prefix = tmp_path / "studio_venv"
_populate_studio_venv(prefix)
out = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix))
site = prefix / "Lib" / "site-packages"
for expected in (
site / "nvidia" / "cuda_runtime" / "bin",
site / "nvidia" / "cublas" / "bin",
site / "nvidia" / "cudnn" / "bin",
site / "nvidia" / "cu13" / "bin" / "x86_64",
site / "torch" / "lib",
):
assert (
str(expected) in out
), f"resolver missed {expected.relative_to(prefix)}: {out}"
def test_path_assembly_makes_cudart_reachable_without_toolkit(self, tmp_path):
"""The #5106 scenario: GPU detected, pip nvidia wheels present,
no system CUDA toolkit. cudart must be reachable from PATH, and
from BOTH binary_dir (#5322) and a pip nvidia dir (#5324)."""
prefix = tmp_path / "studio_venv"
install = tmp_path / "studio_install"
_populate_studio_venv(prefix)
_populate_studio_install(install, runtime = "13.1")
binary_dir = install / "build" / "bin" / "Release"
path_dirs = _build_path_dirs_like_start_llama_server(
binary_dir, prefix, cuda_path = ""
)
# binary_dir first -- Windows DLL search step 1.
assert path_dirs[0] == str(
binary_dir
), f"binary_dir must be first in PATH; got {path_dirs[0]}"
cudart_locations = []
for entry in path_dirs:
for cudart_name in ("cudart64_12.dll", "cudart64_13.dll"):
if (Path(entry) / cudart_name).exists():
cudart_locations.append((entry, cudart_name))
assert cudart_locations, (
f"cudart unreachable from any PATH entry -- #5106 not fixed.\n"
f"PATH entries searched: {path_dirs}"
)
# Defence in depth: both fix paths contribute cudart.
sources = {Path(e).relative_to(tmp_path).parts[0] for e, _ in cudart_locations}
assert (
"studio_install" in sources
), f"#5322's cudart drop not reachable: {cudart_locations}"
assert (
"studio_venv" in sources
), f"#5324's pip nvidia dir not contributing cudart: {cudart_locations}"
def test_cublas_and_cublasLt_also_reachable(self, tmp_path):
"""ggml-cuda imports cublas64; cublas64 imports cublasLt64. All
three must resolve or LoadLibrary returns NULL."""
prefix = tmp_path / "studio_venv"
install = tmp_path / "studio_install"
_populate_studio_venv(prefix)
_populate_studio_install(install, runtime = "13.1")
binary_dir = install / "build" / "bin" / "Release"
path_dirs = _build_path_dirs_like_start_llama_server(binary_dir, prefix)
for required in REAL_UPSTREAM_CUDART_BUNDLE["13.1"]:
reachable = any((Path(d) / required).exists() for d in path_dirs)
assert reachable, (
f"{required} unreachable from PATH; #5106 not fixed.\n"
f"PATH entries: {path_dirs}"
)
def test_no_pip_nvidia_wheels_still_works_via_install_dir(self, tmp_path):
"""No pip nvidia wheels (CPU-only torch / unsloth run standalone):
cudart still resolves via #5322's binary_dir drop."""
prefix = tmp_path / "bare_venv"
prefix.mkdir()
install = tmp_path / "studio_install"
_populate_studio_install(install, runtime = "13.1")
binary_dir = install / "build" / "bin" / "Release"
path_dirs = _build_path_dirs_like_start_llama_server(binary_dir, prefix)
assert path_dirs == [
str(binary_dir)
], f"bare venv produced unexpected PATH: {path_dirs}"
for required in REAL_UPSTREAM_CUDART_BUNDLE["13.1"]:
assert (
binary_dir / required
).exists(), f"{required} missing from binary_dir on bare venv install"
def test_no_install_dir_still_works_via_pip_wheels(self, tmp_path):
"""Pre-#5322 install (binary_dir lacks cudart): #5324's pip
wheel directories on PATH still resolve cudart."""
prefix = tmp_path / "studio_venv"
_populate_studio_venv(prefix)
install = tmp_path / "studio_install_pre5322"
rel = install / "build" / "bin" / "Release"
rel.mkdir(parents = True)
# Main archive payload only; cudart bundle absent.
for fn in (
"llama-server.exe",
"llama.dll",
"ggml-cuda.dll",
"ggml-base.dll",
):
(rel / fn).write_bytes(b"PE-stub")
path_dirs = _build_path_dirs_like_start_llama_server(rel, prefix)
cudart_reachable = any(
(Path(d) / "cudart64_12.dll").exists()
or (Path(d) / "cudart64_13.dll").exists()
for d in path_dirs
)
assert cudart_reachable, (
"#5324 pip wheel fallback failed: cudart unreachable from PATH "
f"on cudart-less install. PATH entries: {path_dirs}"
)
cublas_reachable = any(
(Path(d) / "cublas64_12.dll").exists()
or (Path(d) / "cublas64_13.dll").exists()
for d in path_dirs
)
assert cublas_reachable, "cublas unreachable on cudart-less install"
def test_pre_pr_scenario_would_have_failed(self, tmp_path):
"""Negative control: pre-#5322 + pre-#5324 world leaves cudart
unreachable -- the original failure mode. Confirms the test
actually catches a regression."""
prefix = tmp_path / "studio_venv"
_populate_studio_venv(prefix)
install = tmp_path / "pre_pr_install"
rel = install / "build" / "bin" / "Release"
rel.mkdir(parents = True)
for fn in ("llama-server.exe", "llama.dll", "ggml-cuda.dll"):
(rel / fn).write_bytes(b"PE-stub")
# Pre-PR PATH: binary_dir only. No pip nvidia dirs, no toolkit.
pre_pr_path_dirs = [str(rel)]
cudart_reachable_pre = any(
(Path(d) / "cudart64_12.dll").exists()
or (Path(d) / "cudart64_13.dll").exists()
for d in pre_pr_path_dirs
)
assert not cudart_reachable_pre, (
"Test self-check failed: pre-PR scenario unexpectedly had "
f"cudart reachable. {pre_pr_path_dirs}"
)
class TestWindowsSysPlatformMocked:
"""Confirm the win32 branch in start_llama_server is what we test
(not the linux fallback). Patches sys.platform and re-runs the
branch-selecting helper."""
def test_sys_platform_win32_uses_pip_nvidia_resolver(self, monkeypatch, tmp_path):
monkeypatch.setattr(sys, "platform", "win32")
prefix = tmp_path / "studio_venv"
_populate_studio_venv(prefix)
out = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix))
assert out, f"resolver returned empty under sys.platform=win32: {out}"
# cu13 arch dir must be in the output.
cu13_arch = (
prefix / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64"
)
assert str(cu13_arch) in out

2
studio/install_llama_prebuilt.py
View file

@ -2952,7 +2952,7 @@ def windows_cuda_attempts(
# binary archive, not the cudart archive itself.
runtime_archive_name: str | None = None
runtime_archive_url: str | None = None
if selected_name.startswith(f"llama-"):
if selected_name.startswith("llama-"):
cudart_name = f"cudart-llama-bin-win-cuda-{runtime}-x64.zip"
cudart_url = upstream_assets.get(cudart_name)
if cudart_url and cudart_url != asset_url: