Studio: add torch's pip nvidia DLL dirs to PATH on Windows (#5324)

* 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/<pkg>/(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 <prefix>/Lib/site-packages/nvidia/<pkg>/bin and <prefix>/Lib/site-packages/nvidia/<pkg>/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/<pkg>/bin layout resolved * nvidia/<pkg>/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 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: also scan torch/lib in Windows pip nvidia DLL resolver PyTorch's Windows CUDA wheels frequently bundle cudart64_X.dll and cublas64_X.dll directly under Lib/site-packages/torch/lib/ instead of shipping separate nvidia-cuda-runtime-cuXX / nvidia-cublas-cuXX wheels. On those installs _windows_pip_nvidia_dll_dirs previously returned nothing useful, and llama-server.exe fell back to needing a system CUDA toolkit on PATH -- the original #5106 failure mode. The install-side equivalent python_runtime_dirs in install_llama_prebuilt.py already treats torch/lib as a Python runtime DLL source for the same reason. Bring the runtime resolver in parity so torch-bundled-CUDA installs find their cudart at llama-server start. Updates the existing test that codified the bug (asserted torch/lib was excluded), and adds three new cases: pickup, combined-with-nvidia, and the must-be-a-directory guard. * Studio: cover cu13 bin/x86_64 layout in Windows DLL resolver Three follow-ups from a 12-reviewer batch over c1c8a074 (PR #5324): 1. The current nvidia-cuda-runtime (unsuffixed) 13.2.75 and nvidia-cublas 13.4.0.1 Windows wheels on PyPI ship under nvidia/cu13/bin/x86_64/cudart64_13.dll etc, not under nvidia/PKG/bin/. The previous resolver matched only one directory level past nvidia/PKG/ and silently missed the actual cu13 DLL location, leaving CUDA 13 users on the same failure mode as before #5106. Verified against: pip download nvidia-cuda-runtime --platform win_amd64 which produces nvidia/cu13/bin/x86_64/cudart64_13.dll. 2. glob.glob over sys.prefix interprets [ and ] as a character class. Valid Windows usernames / install paths can contain those characters (for example C:\Users\alice[work]\studio), so the previous resolver silently returned an empty list for such prefixes even when DLL dirs were present. 3. The resolver only ever returned nvidia/PKG/bin -- if both bin and bin/x86_64 exist (current wheels do), Windows DLL search should land on the arch-specific subdir first so the explicit cudart64_X.dll location wins. Rewritten as a pathlib.Path.iterdir walk to fix all three: no glob escaping needed, arch-specific subdirs added explicitly, and ordering puts bin/x86_64 before bin. Conda-style Library/bin/x86_64 and Library/bin/x64 are also covered for parity. A seen set dedupes when wheels happen to expose the same directory through multiple layouts. New tests: - test_picks_up_cu13_bin_x86_64_layout (the actual real-world cu13 case) - test_picks_up_bin_x64_layout - test_mixed_cu12_and_cu13_layouts - test_glob_meta_in_prefix_is_safe (bracket repro) - test_arch_subdir_listed_before_parent_bin (ordering) Verified empirically against PyPI: nvidia-cuda-runtime 13.2.75 -> nvidia/cu13/bin/x86_64/cudart64_13.dll nvidia-cublas 13.4.0.1 -> nvidia/cu13/bin/x86_64/cublas64_13.dll nvidia/cu13/bin/x86_64/cublasLt64_13.dll nvidia-cudnn-cu13 9.22.0.52 -> nvidia/cudnn/bin/cudnn64_9.dll (already covered) Refs #5106 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-11 05:42:09 -07:00 · 2026-05-11 05:42:09 -07:00 · 379f5a5aa6
commit 379f5a5aa6
parent e346193ae8
2 changed files with 326 additions and 2 deletions
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@ -956,6 +956,66 @@ 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 CUDA wheels under
+        ``<prefix>/Lib/site-packages/`` 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, with parity for the Windows-specific wheel layouts seen
+        in the wild. Covered patterns:
+          * ``nvidia/<pkg>/bin`` -- legacy modular wheels
+            (``nvidia-cuda-runtime-cu12``, ``nvidia-cublas-cu12``, etc.).
+          * ``nvidia/<pkg>/bin/x86_64`` and ``.../bin/x64`` -- current
+            CUDA 13 wheel layout used by the unsuffixed
+            ``nvidia-cuda-runtime`` / ``nvidia-cublas`` packages, which
+            ship under ``nvidia/cu13/bin/x86_64/`` (#5106).
+          * ``nvidia/<pkg>/Library/bin`` (and arch subdirs) -- conda-
+            style wheel repacks.
+          * ``torch/lib`` -- PyTorch's own CUDA-bundled Windows wheel,
+            which can ship ``cudart64_*.dll`` directly here instead of
+            as separate ``nvidia-*`` wheels. The install-side helper
+            ``python_runtime_dirs`` in ``install_llama_prebuilt.py``
+            covers this path for the same reason.
+
+        Walks the tree with ``Path.iterdir`` rather than ``glob.glob``
+        so the resolver is safe against Windows paths containing
+        ``[`` or ``]`` (valid in usernames; would otherwise be
+        interpreted as a glob character class and silently miss
+        existing dirs)."""
+        site_packages = Path(prefix) / "Lib" / "site-packages"
+        out: list[str] = []
+        seen: set[str] = set()
+
+        def _add(path: Path) -> None:
+            if not path.is_dir():
+                return
+            key = os.path.normcase(os.path.abspath(str(path)))
+            if key in seen:
+                return
+            seen.add(key)
+            out.append(str(path))
+
+        nvidia_root = site_packages / "nvidia"
+        if nvidia_root.is_dir():
+            for pkg_dir in nvidia_root.iterdir():
+                if not pkg_dir.is_dir():
+                    continue
+                # Order matters for PATH search: arch-specific subdirs
+                # first so the explicit cudart64_X.dll location wins
+                # over a sibling ``bin`` that might be empty.
+                for sub in (
+                    pkg_dir / "bin" / "x86_64",
+                    pkg_dir / "bin" / "x64",
+                    pkg_dir / "bin",
+                    pkg_dir / "Library" / "bin" / "x86_64",
+                    pkg_dir / "Library" / "bin" / "x64",
+                    pkg_dir / "Library" / "bin",
+                ):
+                    _add(sub)
+        _add(site_packages / "torch" / "lib")
+        return out
+
    @staticmethod
    def _select_gpus(
        model_size_bytes: int,
@ -2319,9 +2379,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")
--- a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
+++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
@ -0,0 +1,259 @@
+# 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 <prefix>/Lib/site-packages/nvidia/<pkg>/{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_known_paths(self, tmp_path):
+        # Only nvidia/<pkg>/{bin,Library/bin} and torch/lib are picked
+        # up. Unrelated site-packages contents (numpy, scipy, ...) must
+        # be ignored.
+        site = tmp_path / "Lib" / "site-packages"
+        (site / "numpy").mkdir(parents = True)
+        (site / "scipy" / "linalg").mkdir(parents = True)
+        result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path))
+        assert result == []
+
+    def test_picks_up_torch_lib(self, tmp_path):
+        # PyTorch's Windows CUDA wheel bundles cudart64_X.dll /
+        # cublas64_X.dll directly under Lib/site-packages/torch/lib/
+        # instead of as separate nvidia-* wheels. Without this, users
+        # on torch-bundled-CUDA installs still hit #5106.
+        torch_lib = tmp_path / "Lib" / "site-packages" / "torch" / "lib"
+        torch_lib.mkdir(parents = True)
+        (torch_lib / "cudart64_12.dll").write_bytes(b"")
+        result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path))
+        assert len(result) == 1
+        assert Path(result[0]) == torch_lib
+
+    def test_torch_lib_combined_with_nvidia_wheels(self, tmp_path):
+        # Both modular nvidia-* wheels and torch/lib are returned when
+        # present together.
+        _make_nvidia_layout(
+            tmp_path,
+            {
+                "cuda_runtime": "bin",
+                "cublas": "bin",
+            },
+        )
+        torch_lib = tmp_path / "Lib" / "site-packages" / "torch" / "lib"
+        torch_lib.mkdir(parents = True)
+        (torch_lib / "cudart64_13.dll").write_bytes(b"")
+        result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path))
+        assert len(result) == 3
+        names = {Path(p).name for p in result}
+        assert names == {"bin", "lib"}
+        assert any(Path(p) == torch_lib for p in result)
+
+    def test_torch_lib_must_be_a_directory(self, tmp_path):
+        # If torch/lib exists as a file (broken install), it is
+        # ignored, not returned.
+        site = tmp_path / "Lib" / "site-packages" / "torch"
+        site.mkdir(parents = True)
+        (site / "lib").write_bytes(b"not a dir")
+        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 == []
+
+    def test_picks_up_cu13_bin_x86_64_layout(self, tmp_path):
+        # Current ``nvidia-cuda-runtime`` 13.x and ``nvidia-cublas``
+        # 13.x Windows wheels ship DLLs under
+        # ``nvidia/cu13/bin/x86_64/`` instead of ``nvidia/<pkg>/bin/``.
+        # Without this, users on the new CUDA 13 wheel generation hit
+        # the original #5106 failure mode.
+        dll_dir = (
+            tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64"
+        )
+        dll_dir.mkdir(parents = True)
+        for name in ("cudart64_13.dll", "cublas64_13.dll", "cublasLt64_13.dll"):
+            (dll_dir / name).write_bytes(b"")
+        result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path))
+        assert str(dll_dir) in result, f"cu13 bin/x86_64 not in {result}"
+
+    def test_picks_up_bin_x64_layout(self, tmp_path):
+        # Some repackaged wheels use ``bin/x64`` (Windows-x64 convention)
+        # instead of ``bin/x86_64`` (NVIDIA-internal convention).
+        dll_dir = tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x64"
+        dll_dir.mkdir(parents = True)
+        (dll_dir / "cudart64_13.dll").write_bytes(b"")
+        result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path))
+        assert str(dll_dir) in result
+
+    def test_mixed_cu12_and_cu13_layouts(self, tmp_path):
+        # A venv could have both the modular cu12 wheels (legacy) and
+        # the unsuffixed cu13 wheel installed side by side. Both must
+        # be reachable.
+        site = tmp_path / "Lib" / "site-packages"
+        cu12_bin = site / "nvidia" / "cuda_runtime" / "bin"
+        cu13_arch = site / "nvidia" / "cu13" / "bin" / "x86_64"
+        cu12_bin.mkdir(parents = True)
+        cu13_arch.mkdir(parents = True)
+        result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path))
+        result_set = {Path(p) for p in result}
+        assert cu12_bin in result_set
+        assert cu13_arch in result_set
+
+    def test_glob_meta_in_prefix_is_safe(self, tmp_path):
+        # Windows usernames / install paths can contain ``[`` or ``]``.
+        # A glob-based resolver would interpret these as a character
+        # class and silently return [] even when DLL dirs exist. The
+        # iterdir-based implementation must work on such paths.
+        prefix = tmp_path / "studio_[gpu]_install"
+        dll_dir = prefix / "Lib" / "site-packages" / "nvidia" / "cuda_runtime" / "bin"
+        dll_dir.mkdir(parents = True)
+        (dll_dir / "cudart64_12.dll").write_bytes(b"")
+        result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix))
+        assert str(dll_dir) in result, f"bracket-prefixed path returned empty: {result}"
+
+    def test_arch_subdir_listed_before_parent_bin(self, tmp_path):
+        # When both ``nvidia/<pkg>/bin/`` and
+        # ``nvidia/<pkg>/bin/x86_64/`` exist, the arch-specific subdir
+        # must be listed first so Windows DLL search picks up the
+        # cudart64_X.dll location even if the parent ``bin`` is empty.
+        site = tmp_path / "Lib" / "site-packages"
+        outer_bin = site / "nvidia" / "cu13" / "bin"
+        arch_bin = outer_bin / "x86_64"
+        arch_bin.mkdir(parents = True)
+        (arch_bin / "cudart64_13.dll").write_bytes(b"")
+        result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path))
+        # outer_bin exists as a directory (it contains arch_bin); the
+        # arch-specific subdir should come first in the list.
+        result_paths = [Path(p) for p in result]
+        assert arch_bin in result_paths
+        assert outer_bin in result_paths
+        assert result_paths.index(arch_bin) < result_paths.index(outer_bin)