# 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 nvidia//{bin|Library/bin} tree with a stub DLL per leaf.""" 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//{bin,Library/bin} and torch/lib are picked up. # Unrelated site-packages contents (numpy, scipy, ...) are 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/cublas64 DLLs under # torch/lib/ rather than as nvidia-* wheels; else still hits #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 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. 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) # Regular file 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): # Nonexistent sys.prefix: resolver must return [], not raise. 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): # nvidia 13.x Windows wheels ship DLLs under nvidia/cu13/bin/x86_64/ # not nvidia//bin/; else the new CUDA 13 wheels hit #5106. 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) # rather than ``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 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 paths can contain ``[``/``]``; a glob-based resolver would # read these as a character class. The iterdir impl must handle them. 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 bin/ and bin/x86_64/ exist, the arch subdir must come first # so the Windows DLL search finds cudart64_X.dll if 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 dir (it holds 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)