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.
This commit is contained in:
danielhanchen 2026-05-18 00:10:18 +00:00
commit 5a6ddc343d
2 changed files with 78 additions and 155 deletions

View file

@ -1059,16 +1059,11 @@ class LlamaCppBackend:
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 to the inherited env so
llama-server.exe can resolve cudart / cublas DLLs. Extracted as
a staticmethod so the test in
``studio/backend/tests/test_5106_windows_gpu_detection_mock.py``
asserts against the exact production logic instead of a local
reconstruction. Order: binary_dir first (Windows DLL search
step 1, application directory), then pip nvidia wheels (mirrors
Linux LD_LIBRARY_PATH), then optional system CUDA toolkit
(``CUDA_PATH/bin`` and ``CUDA_PATH/bin/x64``). #5106."""
"""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:
@ -2518,9 +2513,7 @@ 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. See _build_windows_path_dirs for the
# ordering rationale (#5106).
# See _build_windows_path_dirs for ordering. #5106.
path_dirs = self._build_windows_path_dirs(
binary_dir,
sys.prefix,

View file

@ -1,35 +1,21 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""End-to-end Windows GPU-detection regression test for #5106.
"""Windows GPU-detection regression test on a synthetic layout.
The bug: on Windows hosts without a system CUDA toolkit, Studio's
prebuilt llama-server.exe could not resolve ``cudart64_X.dll`` /
``cublas64_X.dll`` / ``cublasLt64_X.dll`` at LoadLibrary time, so
``ggml-cuda.dll`` (which has a static PE import on ``cublas64_X.dll``)
failed to load and llama-server silently fell back to CPU even when
``nvidia-smi`` reported the GPU.
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 lands in two halves:
* PR #5322 (install-time): downloads upstream's paired
``cudart-llama-bin-win-cuda-X.Y-x64.zip`` and overlays its three
DLLs into ``install_dir/build/bin/Release/`` next to
``llama-server.exe``. Windows DLL search resolves them from step
(1) -- the application directory.
* PR #5324 (launch-time): prepends pip-installed
``nvidia/<pkg>/{bin,bin/x86_64,Library/bin}`` and ``torch/lib``
directories to ``PATH`` when spawning ``llama-server.exe``.
Windows DLL search resolves the DLLs from step (3) -- the ``PATH``
environment variable -- even on existing installs that pre-date
PR #5322.
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.
This test exercises both halves on a synthetic Windows layout that
mirrors the real artifact contents (verified empirically against
upstream b9103 ``cudart-llama-bin-win-cuda-13.1-x64.zip`` and the
``nvidia-cuda-runtime`` / ``nvidia-cublas`` PyPI win_amd64 wheels).
CI runners have no GPUs, so we mock the ``nvidia-smi`` probe directly
to assert Studio detects the synthetic GPU AND ends up with cudart
reachable through both the binary-directory and the PATH path.
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
@ -48,22 +34,13 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Stub heavy deps the rest of the studio backend pulls in IFF they
# fail to import here. Unconditionally installing a stub would shadow
# the real module for every subsequent test in this dir
# (test_anthropic_messages.py, test_training_*, etc) and break their
# ``from httpx import HTTPError, Response`` imports.
#
# Why try-import instead of ``importlib.util.find_spec``: ``find_spec``
# only checks discoverability, not import success. ``studio/backend/
# loggers/__init__.py`` re-exports ``handlers.get_logger``, and
# ``handlers.py`` does ``from fastapi import Request, Response`` at
# module load. In a lightweight env without fastapi, ``find_spec
# ("loggers")`` returns a spec but the actual import raises during
# ``from core.inference.llama_cpp import LlamaCppBackend`` collection.
# Calling ``import_module`` here surfaces that failure and falls back
# to the local stub. CI has fastapi installed so this is purely a
# developer-machine ergonomics fix.
# 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
@ -122,19 +99,17 @@ _maybe_stub("httpx", _build_httpx_stub)
from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
# Real upstream b9103 cudart bundle contents (verified by direct unzip
# of the GitHub release artifact). Exactly these three filenames, no
# executables, no subdirectories.
# 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"),
}
# Real win_amd64 wheel layouts observed on PyPI (verified by
# ``pip download nvidia-cuda-runtime --platform win_amd64`` and
# ``unzip -l`` against the resulting wheel).
# 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 modular wheels (cu-suffixed)
# Legacy cu-suffixed wheels
"nvidia/cuda_runtime/bin": ["cudart64_12.dll"],
"nvidia/cublas/bin": [
"cublas64_12.dll",
@ -146,7 +121,7 @@ REAL_PIP_NVIDIA_WHEEL_LAYOUTS = {
"cudnn_adv64_9.dll",
"cudnn_ops64_9.dll",
],
# New unsuffixed cu13 wheel layout: nvidia/cu13/bin/x86_64/
# Unsuffixed cu13 wheels
"nvidia/cu13/bin/x86_64": [
"cudart64_13.dll",
"cublas64_13.dll",
@ -157,29 +132,26 @@ REAL_PIP_NVIDIA_WHEEL_LAYOUTS = {
def _populate_studio_venv(prefix: Path) -> None:
"""Drop fake wheel files matching the real PyPI win_amd64 layouts
seen on actual nvidia-cuda-runtime / nvidia-cublas / nvidia-cudnn
wheels. The resolver doesn't care about file contents, only
presence + directory structure."""
"""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")
# Studio's install_python_stack always installs torch alongside
# the nvidia wheels.
# 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:
"""Drop a Windows-style install_dir/build/bin/Release/ tree
populated as PR #5322 would after the paired cudart overlay."""
"""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)
# Main archive payload
for fn in (
"llama-server.exe",
"llama-quantize.exe",
@ -191,7 +163,7 @@ def _populate_studio_install(install_dir: Path, runtime: str = "13.1") -> None:
"mtmd.dll",
):
(rel / fn).write_bytes(b"PE-stub")
# Paired cudart bundle payload (this is what #5322 adds)
# The cudart overlay #5322 contributes.
for fn in REAL_UPSTREAM_CUDART_BUNDLE[runtime]:
(rel / fn).write_bytes(b"PE-stub")
@ -199,22 +171,17 @@ def _populate_studio_install(install_dir: Path, runtime: str = "13.1") -> None:
def _build_path_dirs_like_start_llama_server(
binary_dir: Path, prefix: Path, cuda_path: str = ""
) -> list[str]:
"""Thin wrapper around the production
``LlamaCppBackend._build_windows_path_dirs`` helper, kept so the
test reads with ``Path`` arguments. Asserting against the real
staticmethod (rather than a hand-copy of its body) is the whole
point: if the production PATH order ever drops or reorders
``_windows_pip_nvidia_dll_dirs``, these tests fail."""
"""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 in
LlamaCppBackend._get_gpu_free_memory returns the supplied CSV.
Other subprocess.run calls (if any in this test process) pass
through to the real subprocess.run."""
"""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):
@ -231,21 +198,15 @@ def _mock_nvidia_smi_run(fake_output: str, returncode: int = 0) -> "mock._patch"
# Tests
# --------------------------------------------------------------------- #
class TestWindowsGpuDetectionAfter5106Fix:
"""Validates the end-to-end #5106 fix on a synthetic Windows
layout. CI runners have no GPU, so we mock nvidia-smi but exercise
every other layer (resolver, PATH builder, install layout) for
real."""
"""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):
"""Sanity: the production nvidia-smi probe parses CSV output
and returns (index, free_mib) tuples. This is the entry point
Studio uses to decide whether a GPU is reachable at all."""
# Clear inherited visibility masks so the synthetic CSV is not
# filtered or shadowed by the parent shell (e.g. on a shared
# GPU runner with CUDA_VISIBLE_DEVICES=1 set).
"""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)
# noahterbest's exact #5106 reproducer: RTX 4090, 22805 MiB free.
# 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()
@ -254,7 +215,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
], f"GPU probe failed to parse mocked nvidia-smi output: {gpus}"
def test_nvidia_smi_probe_respects_cuda_visible_devices(self, monkeypatch):
"""A user with CUDA_VISIBLE_DEVICES=1 should only see GPU 1."""
"""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):
@ -262,10 +223,8 @@ class TestWindowsGpuDetectionAfter5106Fix:
assert gpus == [(1, 24576)], gpus
def test_windows_install_dir_has_all_three_cudart_dlls(self, tmp_path):
"""After PR #5322, install_dir/build/bin/Release/ must contain
all three DLLs from the upstream cudart bundle. Without all
three, ggml-cuda.dll's static PE import on cublas64_X.dll
cannot resolve and the CUDA backend fails to register."""
"""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"
@ -275,16 +234,11 @@ class TestWindowsGpuDetectionAfter5106Fix:
assert (rel / "ggml-cuda.dll").exists()
def test_resolver_finds_real_pypi_wheel_layouts(self, tmp_path):
"""The launch-time PATH resolver must pick up every layout
used by real pip-installed CUDA wheels on Windows today:
* nvidia/<pkg>/bin (legacy cu-suffixed wheels)
* nvidia/<pkg>/bin/x86_64 (new cu13 unsuffixed wheels)
* torch/lib (some torch builds bundle CUDA DLLs here)
"""
"""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))
# All four real-world layouts must be present:
site = prefix / "Lib" / "site-packages"
for expected in (
site / "nvidia" / "cuda_runtime" / "bin",
@ -298,13 +252,9 @@ class TestWindowsGpuDetectionAfter5106Fix:
), f"resolver missed {expected.relative_to(prefix)}: {out}"
def test_path_assembly_makes_cudart_reachable_without_toolkit(self, tmp_path):
"""The exact #5106 scenario: Windows host with the GPU
detected, pip nvidia wheels installed, but NO system CUDA
toolkit (no CUDA_PATH). After the fix, the PATH that
``start_llama_server`` prepends to the llama-server.exe
subprocess env must make ``cudart64_*.dll`` reachable from at
least one entry. We verify directly by walking each PATH
entry on the live filesystem."""
"""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)
@ -313,11 +263,10 @@ class TestWindowsGpuDetectionAfter5106Fix:
path_dirs = _build_path_dirs_like_start_llama_server(
binary_dir, prefix, cuda_path = ""
)
# binary_dir is first (Windows DLL search step 1).
# 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 MUST be findable from at least one PATH entry.
cudart_locations = []
for entry in path_dirs:
for cudart_name in ("cudart64_12.dll", "cudart64_13.dll"):
@ -327,23 +276,18 @@ class TestWindowsGpuDetectionAfter5106Fix:
f"cudart unreachable from any PATH entry -- #5106 not fixed.\n"
f"PATH entries searched: {path_dirs}"
)
# Confirm cudart is reachable from BOTH the install dir (PR
# #5322's contribution) AND a pip nvidia dir (PR #5324's
# contribution). Defense in depth.
# 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"PR #5322's cudart drop not reachable: {cudart_locations}"
), f"#5322's cudart drop not reachable: {cudart_locations}"
assert (
"studio_venv" in sources
), f"PR #5324's pip nvidia dir not contributing cudart: {cudart_locations}"
), f"#5324's pip nvidia dir not contributing cudart: {cudart_locations}"
def test_cublas_and_cublasLt_also_reachable(self, tmp_path):
"""ggml-cuda.dll has a static PE import on cublas64_X.dll
(verified by ``objdump -p`` on the upstream b9103 build).
cublas64_X.dll has a static PE import on cublasLt64_X.dll.
All three must be reachable or LoadLibrary("ggml-cuda.dll")
returns NULL."""
"""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)
@ -358,37 +302,31 @@ class TestWindowsGpuDetectionAfter5106Fix:
)
def test_no_pip_nvidia_wheels_still_works_via_install_dir(self, tmp_path):
"""A user with no pip nvidia wheels (CPU-only torch install,
``unsloth run`` standalone, custom torch builds) should still
get cudart via PR #5322's paired download alone -- binary_dir
is enough."""
"""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()
# No pip nvidia / torch wheels installed
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)
# Only binary_dir should be in PATH.
assert path_dirs == [
str(binary_dir)
], f"bare venv produced unexpected PATH: {path_dirs}"
# binary_dir has all three cudart DLLs.
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):
"""A user on an existing pre-#5322 Studio install (binary_dir
lacks cudart) should still get cudart via PR #5324's pip
wheel directories on 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 missing.
# Main archive payload only; cudart bundle absent.
for fn in (
"llama-server.exe",
"llama.dll",
@ -396,7 +334,6 @@ class TestWindowsGpuDetectionAfter5106Fix:
"ggml-base.dll",
):
(rel / fn).write_bytes(b"PE-stub")
# cudart NOT in binary_dir on this scenario.
path_dirs = _build_path_dirs_like_start_llama_server(rel, prefix)
cudart_reachable = any(
(Path(d) / "cudart64_12.dll").exists()
@ -404,7 +341,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
for d in path_dirs
)
assert cudart_reachable, (
"PR #5324 pip wheel fallback failed: cudart unreachable from PATH "
"#5324 pip wheel fallback failed: cudart unreachable from PATH "
f"on cudart-less install. PATH entries: {path_dirs}"
)
cublas_reachable = any(
@ -415,21 +352,17 @@ class TestWindowsGpuDetectionAfter5106Fix:
assert cublas_reachable, "cublas unreachable on cudart-less install"
def test_pre_pr_scenario_would_have_failed(self, tmp_path):
"""Negative control: reconstruct the pre-#5322 + pre-#5324
world (cudart NOT dropped by installer, PATH not augmented by
launcher) and assert that cudart is unreachable -- the
original #5106 failure mode. This ensures the test would
actually catch a regression."""
"""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)
# Main archive only; no cudart bundle.
for fn in ("llama-server.exe", "llama.dll", "ggml-cuda.dll"):
(rel / fn).write_bytes(b"PE-stub")
# Pre-PR PATH: binary_dir + CUDA_PATH/bin only. Pip nvidia
# dirs NOT added. No system CUDA toolkit (the #5106 scenario).
# 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()
@ -443,20 +376,17 @@ class TestWindowsGpuDetectionAfter5106Fix:
class TestWindowsSysPlatformMocked:
"""Validate that the win32 branch in start_llama_server is the
branch we test, not the linux fallback. We can't easily call the
full start_llama_server method (it constructs a llama-server
subprocess), but we can patch sys.platform and re-import the
"""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)
# On win32, the resolver should be called and return non-empty.
out = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix))
assert out, f"resolver returned empty under sys.platform=win32: {out}"
# Specifically, the cu13 arch dir must be in the output.
# cu13 arch dir must be in the output.
cu13_arch = (
prefix / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64"
)