From 3e262ab173fc0f5890e8247dd1a4f865cee14919 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 7 May 2026 11:33:37 +0000 Subject: [PATCH] Studio: add torch's pip nvidia DLL dirs to PATH on Windows Studio's install_python_stack bundles torch with matching CUDA wheels (nvidia-cuda-runtime-cu13, nvidia-cublas-cu13, etc.) which ship cudart64_X.dll, cublas64_X.dll, and cublasLt64_X.dll under the prefix's Lib/site-packages/nvidia//(bin|Library/bin)/ tree. The Linux runtime env block in start_llama_server already pulls the equivalent nvidia/cu*/lib paths into LD_LIBRARY_PATH, but the Windows block did not do this, so the prebuilt llama-server.exe could not resolve cudart64_X.dll at runtime unless the user had a matching system CUDA toolkit on PATH. That is the root cause of the Windows reports in unslothai/unsloth#5106 ("GPU detected but model loaded entirely on RAM/CPU"), and matches Roland's repeated workaround in that issue: install matching CUDA toolkit version. Brings the Windows env block in line with the Linux pattern: * New LlamaCppBackend._windows_pip_nvidia_dll_dirs resolver globs /Lib/site-packages/nvidia//bin and /Lib/site-packages/nvidia//Library/bin. Both layouts are seen in the wild across cuda_runtime / cublas / cudnn / nvjitlink wheels. * The Windows env block now extends path_dirs with the resolver's output before falling back to CUDA_PATH/bin, so pip-installed wheels are the canonical source (mirroring the Linux LD_LIBRARY_PATH ordering). System CUDA toolkit remains a valid fallback. Tests: 7 new cases in studio/backend/tests/test_llama_cpp_windows_nvidia_path.py: * empty resolver when no nvidia wheels installed * nvidia//bin layout resolved * nvidia//Library/bin layout resolved * mixed bin and Library/bin layouts both resolved * unrelated site-packages contents not walked * non-directory entries skipped * missing prefix does not raise 110 backend tests pass. No regressions. Refs #5106 --- studio/backend/core/inference/llama_cpp.py | 30 +++- .../test_llama_cpp_windows_nvidia_path.py | 134 ++++++++++++++++++ 2 files changed, 162 insertions(+), 2 deletions(-) create mode 100644 studio/backend/tests/test_llama_cpp_windows_nvidia_path.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 38c0261f5a..bac82e4d39 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -956,6 +956,27 @@ class LlamaCppBackend: logger.debug(f"torch GPU probe failed: {e}") return [] + @staticmethod + def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: + """Return DLL dirs from pip-installed nvidia wheels under + ``/Lib/site-packages/nvidia/`` so llama-server.exe can + load cudart64_X.dll / cublas64_X.dll without a system CUDA + toolkit. Mirrors the Linux nvidia/cu*/lib LD_LIBRARY_PATH + block. Wheel layouts vary, so we cover the two seen patterns: + ``nvidia//bin`` and ``nvidia//Library/bin``.""" + import glob as _glob + + nvidia_root = os.path.join(prefix, "Lib", "site-packages", "nvidia") + out: list[str] = [] + for pattern in ( + os.path.join(nvidia_root, "*", "bin"), + os.path.join(nvidia_root, "*", "Library", "bin"), + ): + for nv_dir in _glob.glob(pattern): + if os.path.isdir(nv_dir): + out.append(nv_dir) + return out + @staticmethod def _select_gpus( model_size_bytes: int, @@ -2319,9 +2340,14 @@ class LlamaCppBackend: binary_dir = str(Path(binary).parent) if sys.platform == "win32": - # On Windows, CUDA DLLs (cublas64_12.dll, cudart64_12.dll, etc.) - # must be on PATH. Add CUDA_PATH\bin if available. + # 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") diff --git a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py new file mode 100644 index 0000000000..cb1eed8725 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py @@ -0,0 +1,134 @@ +# 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 the Windows pip-nvidia DLL dir resolver. + +Studio installs torch with bundled CUDA wheels (nvidia-cuda-runtime-cu13, +nvidia-cublas-cu13, etc.) and the prebuilt llama-server.exe must find +those DLLs at runtime to load CUDA. Mirrors the Linux LD_LIBRARY_PATH +block. See unslothai/unsloth#5106. +""" +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +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 before importing the module under test. +_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")) + +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", "TimeoutException", "ReadTimeout", + "ReadError", "RemoteProtocolError", "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", (), { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + + +def _make_nvidia_layout(prefix: Path, pkgs_with_layout: dict[str, str]): + """Build a fake /Lib/site-packages/nvidia//{bin|Library/bin} + tree with a stub DLL inside each leaf so isdir() picks them up.""" + nv = prefix / "Lib" / "site-packages" / "nvidia" + for pkg, layout in pkgs_with_layout.items(): + if layout == "bin": + d = nv / pkg / "bin" + elif layout == "library_bin": + d = nv / pkg / "Library" / "bin" + else: + raise ValueError(layout) + d.mkdir(parents = True, exist_ok = True) + (d / "stub.dll").write_bytes(b"") + + +class TestWindowsPipNvidiaDllDirs: + def test_returns_empty_when_no_nvidia_wheels(self, tmp_path): + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert result == [] + + def test_picks_up_bin_layout(self, tmp_path): + _make_nvidia_layout(tmp_path, { + "cuda_runtime": "bin", + "cublas": "bin", + "cudnn": "bin", + }) + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert len(result) == 3 + assert all(Path(p).is_dir() for p in result) + assert all(Path(p).name == "bin" for p in result) + names = {Path(p).parent.name for p in result} + assert names == {"cuda_runtime", "cublas", "cudnn"} + + def test_picks_up_library_bin_layout(self, tmp_path): + _make_nvidia_layout(tmp_path, { + "cuda_runtime": "library_bin", + "cublas": "library_bin", + }) + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert len(result) == 2 + for p in result: + assert Path(p).is_dir() + assert Path(p).parent.name == "Library" + assert Path(p).parent.parent.name in {"cuda_runtime", "cublas"} + + def test_mixed_layouts_all_resolved(self, tmp_path): + _make_nvidia_layout(tmp_path, { + "cuda_runtime": "bin", + "cublas": "library_bin", + "cudnn": "bin", + "nvjitlink": "library_bin", + }) + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert len(result) == 4 + + def test_does_not_walk_outside_nvidia(self, tmp_path): + # Ensure unrelated site-packages contents are not picked up. + site = tmp_path / "Lib" / "site-packages" + (site / "torch" / "lib").mkdir(parents = True) + (site / "torch" / "lib" / "stub.dll").write_bytes(b"") + (site / "numpy").mkdir(parents = True) + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert result == [] + + def test_skips_non_directories(self, tmp_path): + nv = tmp_path / "Lib" / "site-packages" / "nvidia" + (nv / "cuda_runtime").mkdir(parents = True) + # Create a regular file at the path where 'bin' would normally be a dir + (nv / "cuda_runtime" / "bin").write_bytes(b"not a dir") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert result == [] + + def test_missing_prefix_does_not_raise(self): + # If sys.prefix points to a path that doesn't exist (unusual, + # but possible during test setup), the resolver must just + # return [] rather than raising. + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs( + "/this/path/does/not/exist/anywhere" + ) + assert result == []