Skip to content

Commit 902a7f5

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 902a7f5

4 files changed

Lines changed: 207 additions & 10 deletions

File tree

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py

Lines changed: 56 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,6 @@ 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(f"Unsupported library name: {libname!r}. Supported names: {sorted(_ALL_SUPPORTED_LIBNAMES)}")
138193
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: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
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+
@pytest.mark.parametrize("libname", ["cuda", "nvml"])
33+
def test_driver_only_libnames_contains(libname):
34+
assert libname in _DRIVER_ONLY_LIBNAMES
35+
36+
37+
@pytest.mark.parametrize("libname", ["cudart", "nvrtc", "cublas", "nvvm"])
38+
def test_driver_only_libnames_excludes_ctk_libs(libname):
39+
assert libname not in _DRIVER_ONLY_LIBNAMES
40+
41+
42+
# ---------------------------------------------------------------------------
43+
# _load_driver_lib_no_cache
44+
# ---------------------------------------------------------------------------
45+
46+
47+
def test_driver_lib_returns_already_loaded(mocker):
48+
already = LoadedDL("/usr/lib/libcuda.so.1", True, 0xBEEF, "was-already-loaded-from-elsewhere")
49+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=already)
50+
mocker.patch(f"{_MODULE}.load_with_system_search")
51+
52+
result = _load_driver_lib_no_cache("cuda")
53+
54+
assert result is already
55+
# system search should not have been called
56+
from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as mod
57+
58+
mod.load_with_system_search.assert_not_called()
59+
60+
61+
def test_driver_lib_falls_through_to_system_search(mocker):
62+
loaded = _make_loaded_dl("/usr/lib/libcuda.so.1", "system-search")
63+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None)
64+
mocker.patch(f"{_MODULE}.load_with_system_search", return_value=loaded)
65+
66+
result = _load_driver_lib_no_cache("cuda")
67+
68+
assert result is loaded
69+
assert result.found_via == "system-search"
70+
71+
72+
def test_driver_lib_raises_when_not_found(mocker):
73+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None)
74+
mocker.patch(f"{_MODULE}.load_with_system_search", return_value=None)
75+
76+
with pytest.raises(DynamicLibNotFoundError, match="NVIDIA driver library"):
77+
_load_driver_lib_no_cache("nvml")
78+
79+
80+
def test_driver_lib_does_not_search_site_packages(mocker):
81+
"""Driver libs must not go through the CTK search cascade."""
82+
loaded = _make_loaded_dl("/usr/lib/libcuda.so.1", "system-search")
83+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None)
84+
mocker.patch(f"{_MODULE}.load_with_system_search", return_value=loaded)
85+
86+
from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib
87+
88+
spy = mocker.spy(_FindNvidiaDynamicLib, "try_site_packages")
89+
_load_driver_lib_no_cache("cuda")
90+
spy.assert_not_called()
91+
92+
93+
# ---------------------------------------------------------------------------
94+
# _load_lib_no_cache dispatches driver libs correctly
95+
# ---------------------------------------------------------------------------
96+
97+
98+
@pytest.mark.parametrize("libname", sorted(_DRIVER_ONLY_LIBNAMES))
99+
def test_load_lib_no_cache_dispatches_to_driver_path(libname, mocker):
100+
loaded = _make_loaded_dl(f"/usr/lib/fake_{libname}.so", "system-search")
101+
mock_driver = mocker.patch(f"{_MODULE}._load_driver_lib_no_cache", return_value=loaded)
102+
103+
result = _load_lib_no_cache(libname)
104+
105+
assert result is loaded
106+
mock_driver.assert_called_once_with(libname)
107+
108+
109+
def test_load_lib_no_cache_does_not_dispatch_ctk_lib_to_driver_path(mocker):
110+
"""Ensure regular CTK libs don't take the driver shortcut."""
111+
mock_driver = mocker.patch(f"{_MODULE}._load_driver_lib_no_cache")
112+
# Let the normal path run far enough to prove the driver path wasn't used.
113+
# We'll make it fail quickly at check_if_already_loaded_from_elsewhere.
114+
from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib
115+
116+
mocker.patch.object(_FindNvidiaDynamicLib, "try_site_packages", return_value=None)
117+
mocker.patch.object(_FindNvidiaDynamicLib, "try_with_conda_prefix", return_value=None)
118+
mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None)
119+
mocker.patch(f"{_MODULE}.load_dependencies")
120+
mocker.patch(
121+
f"{_MODULE}.load_with_system_search",
122+
return_value=_make_loaded_dl("/usr/lib/libcudart.so.13", "system-search"),
123+
)
124+
125+
_load_lib_no_cache("cudart")
126+
127+
mock_driver.assert_not_called()

cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
import json
55
import os
66
import platform
7-
from unittest.mock import patch
87

98
import pytest
109
import spawned_process_runner
@@ -62,12 +61,16 @@ def test_supported_libnames_windows_libnames_requiring_os_add_dll_directory_cons
6261
)
6362

6463

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")
64+
def test_runtime_error_on_non_64bit_python(mocker):
65+
mocker.patch("struct.calcsize", return_value=3) # fake 24-bit pointer
66+
with pytest.raises(RuntimeError, match=r"requires 64-bit Python\. Currently running: 24-bit Python"):
67+
load_nvidia_dynamic_lib("cudart")
68+
69+
70+
@pytest.mark.parametrize("libname", ["bogus", "not_a_real_lib", "cupti"])
71+
def test_unsupported_libname_raises_value_error(libname):
72+
with pytest.raises(ValueError, match=rf"Unsupported library name: '{libname}'.*cudart"):
73+
load_nvidia_dynamic_lib(libname)
7174

7275

7376
IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES = {

0 commit comments

Comments
 (0)