Skip to content

Commit 182be13

Browse files
cpcloudcursoragent
andcommitted
feat(pathfinder): support "cuda"/"nvml" driver libs and reject unsupported libnames
Add support for loading NVIDIA driver libraries ("cuda", "nvml") via load_nvidia_dynamic_lib(). These are part of the display driver (not the CTK) and use a simplified system-search-only path, skipping site-packages, conda, CUDA_HOME, and canary probe steps. Also reject unrecognized libnames with a ValueError instead of silently falling through to system search, which could produce surprising results for unsupported libs like "cupti". Closes #1288, closes #1564. Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 0921933 commit 182be13

File tree

4 files changed

+219
-11
lines changed

4 files changed

+219
-11
lines changed

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py

Lines changed: 59 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,11 @@
66
import sys
77

88
from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib
9-
from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL, load_dependencies
9+
from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL, load_dependencies
10+
from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
11+
SUPPORTED_LINUX_SONAMES,
12+
SUPPORTED_WINDOWS_DLLS,
13+
)
1014
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
1115

1216
if IS_WINDOWS:
@@ -22,8 +26,44 @@
2226
load_with_system_search,
2327
)
2428

29+
# All libnames recognized by load_nvidia_dynamic_lib, across all categories
30+
# (CTK, third-party, driver). Built from the platform-appropriate soname/DLL
31+
# registry so that platform-specific libs (e.g. cufile on Linux) are included
32+
# only where they apply.
33+
_ALL_SUPPORTED_LIBNAMES: frozenset[str] = frozenset(
34+
(SUPPORTED_WINDOWS_DLLS if IS_WINDOWS else SUPPORTED_LINUX_SONAMES).keys()
35+
)
36+
37+
# Driver libraries: shipped with the NVIDIA display driver, always on the
38+
# system linker path. These skip all CTK search steps (site-packages,
39+
# conda, CUDA_HOME, canary) and go straight to system search.
40+
_DRIVER_ONLY_LIBNAMES = frozenset(("cuda", "nvml"))
41+
42+
43+
def _load_driver_lib_no_cache(libname: str) -> LoadedDL:
44+
"""Load an NVIDIA driver library (system-search only).
45+
46+
Driver libs (libcuda, libnvidia-ml) are part of the display driver, not
47+
the CUDA Toolkit. They are always on the system linker path, so the
48+
full CTK search cascade (site-packages, conda, CUDA_HOME, canary) is
49+
unnecessary.
50+
"""
51+
loaded = check_if_already_loaded_from_elsewhere(libname, False)
52+
if loaded is not None:
53+
return loaded
54+
loaded = load_with_system_search(libname)
55+
if loaded is not None:
56+
return loaded
57+
raise DynamicLibNotFoundError(
58+
f'"{libname}" is an NVIDIA driver library and can only be found via'
59+
f" system search. Ensure the NVIDIA display driver is installed."
60+
)
61+
2562

2663
def _load_lib_no_cache(libname: str) -> LoadedDL:
64+
if libname in _DRIVER_ONLY_LIBNAMES:
65+
return _load_driver_lib_no_cache(libname)
66+
2767
finder = _FindNvidiaDynamicLib(libname)
2868
abs_path = finder.try_site_packages()
2969
if abs_path is not None:
@@ -83,6 +123,7 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
83123
https://github.com/NVIDIA/cuda-python/issues/1011
84124
85125
Raises:
126+
ValueError: If ``libname`` is not a recognized library name.
86127
DynamicLibNotFoundError: If the library cannot be found or loaded.
87128
RuntimeError: If Python is not 64-bit.
88129
@@ -123,6 +164,18 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
123164
124165
- If set, use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
125166
167+
**Driver libraries** (``"cuda"``, ``"nvml"``):
168+
169+
These are part of the NVIDIA display driver (not the CUDA Toolkit) and
170+
are always on the system linker path. For these libraries the search
171+
is simplified to:
172+
173+
0. Already loaded in the current process
174+
1. OS default mechanisms (``dlopen`` / ``LoadLibraryW``)
175+
176+
The CTK-specific steps (site-packages, conda, ``CUDA_HOME``, canary
177+
probe) are skipped entirely.
178+
126179
Notes:
127180
The search is performed **per library**. There is currently no mechanism to
128181
guarantee that multiple libraries are all resolved from the same location.
@@ -135,4 +188,9 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
135188
f" Currently running: {pointer_size_bits}-bit Python"
136189
f" {sys.version_info.major}.{sys.version_info.minor}"
137190
)
191+
if libname not in _ALL_SUPPORTED_LIBNAMES:
192+
raise ValueError(
193+
f"Unsupported library name: {libname!r}."
194+
f" Supported names: {sorted(_ALL_SUPPORTED_LIBNAMES)}"
195+
)
138196
return _load_lib_no_cache(libname)

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,15 @@
214214
"nvpl_fftw": ("libnvpl_fftw.so.0",),
215215
"nvshmem_host": ("libnvshmem_host.so.3",),
216216
}
217-
SUPPORTED_LINUX_SONAMES = SUPPORTED_LINUX_SONAMES_CTK | SUPPORTED_LINUX_SONAMES_OTHER
217+
# Driver libraries: shipped with the NVIDIA driver, always on the system
218+
# linker path. Only system search is needed (no site-packages / conda /
219+
# CUDA_HOME). Note the non-standard naming: "cuda" → libcuda.so.1,
220+
# "nvml" → libnvidia-ml.so.1.
221+
SUPPORTED_LINUX_SONAMES_DRIVER = {
222+
"cuda": ("libcuda.so.1",),
223+
"nvml": ("libnvidia-ml.so.1",),
224+
}
225+
SUPPORTED_LINUX_SONAMES = SUPPORTED_LINUX_SONAMES_CTK | SUPPORTED_LINUX_SONAMES_OTHER | SUPPORTED_LINUX_SONAMES_DRIVER
218226

219227
# Based on these files:
220228
# cuda_12.0.1_528.33_windows.exe
@@ -338,7 +346,11 @@
338346
"cutensor": ("cutensor.dll",),
339347
"cutensorMg": ("cutensorMg.dll",),
340348
}
341-
SUPPORTED_WINDOWS_DLLS = SUPPORTED_WINDOWS_DLLS_CTK | SUPPORTED_WINDOWS_DLLS_OTHER
349+
SUPPORTED_WINDOWS_DLLS_DRIVER = {
350+
"cuda": ("nvcuda.dll",),
351+
"nvml": ("nvml.dll",),
352+
}
353+
SUPPORTED_WINDOWS_DLLS = SUPPORTED_WINDOWS_DLLS_CTK | SUPPORTED_WINDOWS_DLLS_OTHER | SUPPORTED_WINDOWS_DLLS_DRIVER
342354

343355
LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY = (
344356
"cufft",
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
"""Tests for NVIDIA driver library loading ("cuda", "nvml").
5+
6+
These libraries are part of the display driver, not the CUDA Toolkit.
7+
They use a simplified system-search-only path, skipping site-packages,
8+
conda, CUDA_HOME, and the canary probe.
9+
"""
10+
11+
import pytest
12+
13+
from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL
14+
from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import (
15+
_DRIVER_ONLY_LIBNAMES,
16+
_load_driver_lib_no_cache,
17+
_load_lib_no_cache,
18+
)
19+
20+
_MODULE = "cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib"
21+
22+
23+
def _make_loaded_dl(path, found_via):
24+
return LoadedDL(path, False, 0xDEAD, found_via)
25+
26+
27+
# ---------------------------------------------------------------------------
28+
# _DRIVER_ONLY_LIBNAMES registry
29+
# ---------------------------------------------------------------------------
30+
31+
32+
def test_driver_only_libnames_contains_cuda():
33+
assert "cuda" in _DRIVER_ONLY_LIBNAMES
34+
35+
36+
def test_driver_only_libnames_contains_nvml():
37+
assert "nvml" in _DRIVER_ONLY_LIBNAMES
38+
39+
40+
def test_driver_only_libnames_does_not_contain_ctk_libs():
41+
for ctk_lib in ("cudart", "nvrtc", "cublas", "nvvm"):
42+
assert ctk_lib not in _DRIVER_ONLY_LIBNAMES
43+
44+
45+
# ---------------------------------------------------------------------------
46+
# _load_driver_lib_no_cache
47+
# ---------------------------------------------------------------------------
48+
49+
50+
def test_driver_lib_returns_already_loaded(mocker):
51+
already = LoadedDL("/usr/lib/libcuda.so.1", True, 0xBEEF, "was-already-loaded-from-elsewhere")
52+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=already)
53+
mocker.patch(f"{_MODULE}.load_with_system_search")
54+
55+
result = _load_driver_lib_no_cache("cuda")
56+
57+
assert result is already
58+
# system search should not have been called
59+
from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as mod
60+
61+
mod.load_with_system_search.assert_not_called()
62+
63+
64+
def test_driver_lib_falls_through_to_system_search(mocker):
65+
loaded = _make_loaded_dl("/usr/lib/libcuda.so.1", "system-search")
66+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None)
67+
mocker.patch(f"{_MODULE}.load_with_system_search", return_value=loaded)
68+
69+
result = _load_driver_lib_no_cache("cuda")
70+
71+
assert result is loaded
72+
assert result.found_via == "system-search"
73+
74+
75+
def test_driver_lib_raises_when_not_found(mocker):
76+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None)
77+
mocker.patch(f"{_MODULE}.load_with_system_search", return_value=None)
78+
79+
with pytest.raises(DynamicLibNotFoundError, match="NVIDIA driver library"):
80+
_load_driver_lib_no_cache("nvml")
81+
82+
83+
def test_driver_lib_does_not_search_site_packages(mocker):
84+
"""Driver libs must not go through the CTK search cascade."""
85+
loaded = _make_loaded_dl("/usr/lib/libcuda.so.1", "system-search")
86+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None)
87+
mocker.patch(f"{_MODULE}.load_with_system_search", return_value=loaded)
88+
89+
from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib
90+
91+
spy = mocker.spy(_FindNvidiaDynamicLib, "try_site_packages")
92+
_load_driver_lib_no_cache("cuda")
93+
spy.assert_not_called()
94+
95+
96+
# ---------------------------------------------------------------------------
97+
# _load_lib_no_cache dispatches driver libs correctly
98+
# ---------------------------------------------------------------------------
99+
100+
101+
@pytest.mark.parametrize("libname", sorted(_DRIVER_ONLY_LIBNAMES))
102+
def test_load_lib_no_cache_dispatches_to_driver_path(libname, mocker):
103+
loaded = _make_loaded_dl(f"/usr/lib/fake_{libname}.so", "system-search")
104+
mock_driver = mocker.patch(f"{_MODULE}._load_driver_lib_no_cache", return_value=loaded)
105+
106+
result = _load_lib_no_cache(libname)
107+
108+
assert result is loaded
109+
mock_driver.assert_called_once_with(libname)
110+
111+
112+
def test_load_lib_no_cache_does_not_dispatch_ctk_lib_to_driver_path(mocker):
113+
"""Ensure regular CTK libs don't take the driver shortcut."""
114+
mock_driver = mocker.patch(f"{_MODULE}._load_driver_lib_no_cache")
115+
# Let the normal path run far enough to prove the driver path wasn't used.
116+
# We'll make it fail quickly at check_if_already_loaded_from_elsewhere.
117+
from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib
118+
119+
mocker.patch.object(_FindNvidiaDynamicLib, "try_site_packages", return_value=None)
120+
mocker.patch.object(_FindNvidiaDynamicLib, "try_with_conda_prefix", return_value=None)
121+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None)
122+
mocker.patch(f"{_MODULE}.load_dependencies")
123+
mocker.patch(
124+
f"{_MODULE}.load_with_system_search",
125+
return_value=_make_loaded_dl("/usr/lib/libcudart.so.13", "system-search"),
126+
)
127+
128+
_load_lib_no_cache("cudart")
129+
130+
mock_driver.assert_not_called()

cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,6 @@
44
import json
55
import os
66
import platform
7-
from unittest.mock import patch
8-
97
import pytest
108
import spawned_process_runner
119
from child_load_nvidia_dynamic_lib_helper import build_child_process_failed_for_libname_message, child_process_func
@@ -62,12 +60,22 @@ def test_supported_libnames_windows_libnames_requiring_os_add_dll_directory_cons
6260
)
6361

6462

65-
def test_runtime_error_on_non_64bit_python():
66-
with (
67-
patch("struct.calcsize", return_value=3), # fake 24-bit pointer
68-
pytest.raises(RuntimeError, match=r"requires 64-bit Python\. Currently running: 24-bit Python"),
69-
):
70-
load_nvidia_dynamic_lib("not_used")
63+
def test_runtime_error_on_non_64bit_python(mocker):
64+
mocker.patch("struct.calcsize", return_value=3) # fake 24-bit pointer
65+
with pytest.raises(RuntimeError, match=r"requires 64-bit Python\. Currently running: 24-bit Python"):
66+
load_nvidia_dynamic_lib("cudart")
67+
68+
69+
def test_unsupported_libname_raises_value_error():
70+
with pytest.raises(ValueError, match=r"Unsupported library name: 'bogus'"):
71+
load_nvidia_dynamic_lib("bogus")
72+
73+
74+
def test_unsupported_libname_error_lists_supported_names():
75+
with pytest.raises(ValueError, match=r"Supported names:") as exc_info:
76+
load_nvidia_dynamic_lib("not_a_real_lib")
77+
# The error message should mention at least one well-known lib
78+
assert "cudart" in str(exc_info.value)
7179

7280

7381
IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES = {

0 commit comments

Comments
 (0)