diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 5994b6ab4e..acd1d0c00b 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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, diff --git a/studio/backend/tests/test_5106_windows_gpu_detection_mock.py b/studio/backend/tests/test_windows_gpu_detection_mock.py similarity index 57% rename from studio/backend/tests/test_5106_windows_gpu_detection_mock.py rename to studio/backend/tests/test_windows_gpu_detection_mock.py index 9a7583ea87..023630fb9a 100644 --- a/studio/backend/tests/test_5106_windows_gpu_detection_mock.py +++ b/studio/backend/tests/test_windows_gpu_detection_mock.py @@ -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//{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//{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 /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//bin (legacy cu-suffixed wheels) - * nvidia//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//bin, nvidia//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" )