diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index bac82e4d39..e12abe2868 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -958,23 +958,35 @@ class LlamaCppBackend: @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``.""" + """Return DLL dirs from pip-installed CUDA wheels under + ``/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 two additional Windows-specific layouts. + Covered patterns: + * ``nvidia//bin`` -- modular ``nvidia-cuda-runtime-cuXX``, + ``nvidia-cublas-cuXX`` wheels. + * ``nvidia//Library/bin`` -- conda-style wheel repacks. + * ``torch/lib`` -- PyTorch's own CUDA-bundled Windows wheel, + which ships ``cudart64_*.dll`` directly under ``torch/lib`` + instead of as separate ``nvidia-*`` wheels (#5106). The + install-side equivalent ``python_runtime_dirs`` in + ``install_llama_prebuilt.py`` covers this path for the same + reason.""" import glob as _glob - nvidia_root = os.path.join(prefix, "Lib", "site-packages", "nvidia") + site_packages = os.path.join(prefix, "Lib", "site-packages") out: list[str] = [] for pattern in ( - os.path.join(nvidia_root, "*", "bin"), - os.path.join(nvidia_root, "*", "Library", "bin"), + os.path.join(site_packages, "nvidia", "*", "bin"), + os.path.join(site_packages, "nvidia", "*", "Library", "bin"), ): for nv_dir in _glob.glob(pattern): if os.path.isdir(nv_dir): out.append(nv_dir) + torch_lib = os.path.join(site_packages, "torch", "lib") + if os.path.isdir(torch_lib): + out.append(torch_lib) return out @staticmethod diff --git a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py index e9c85abf9c..f7e7a009b5 100644 --- a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py +++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py @@ -123,12 +123,53 @@ class TestWindowsPipNvidiaDllDirs: 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. + 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, ...) must + # be ignored. 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) + (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 == []