Skip to content

Commit 4eade17

Browse files
committed
Add a pathfinder NVML driver release version helper.
Query nvmlSystemGetDriverVersion() through pathfinder's driver library loading path and add a minimal real test so the implementation is preserved as a future reference. Made-with: Cursor
1 parent f72c963 commit 4eade17

2 files changed

Lines changed: 58 additions & 0 deletions

File tree

‎cuda_pathfinder/cuda/pathfinder/_utils/driver_info.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
)
1313
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
1414

15+
_NVML_SUCCESS = 0
16+
_NVML_SYSTEM_DRIVER_VERSION_BUFFER_LENGTH = 80
17+
1518

1619
@dataclass(frozen=True, slots=True)
1720
class DriverCudaVersion:
@@ -48,3 +51,44 @@ def _query_driver_version_int() -> int:
4851
if status != 0:
4952
raise RuntimeError(f"Failed to query CUDA driver version via cuDriverGetVersion() (status={status}).")
5053
return version.value
54+
55+
56+
def _query_driver_release_version_text() -> str:
57+
"""Return the graphics driver release version from ``nvmlSystemGetDriverVersion()``."""
58+
loaded_nvml = _load_nvidia_dynamic_lib("nvml")
59+
nvml_lib = ctypes.CDLL(loaded_nvml.abs_path)
60+
61+
nvml_init_v2 = nvml_lib.nvmlInit_v2
62+
nvml_init_v2.argtypes = []
63+
nvml_init_v2.restype = ctypes.c_int
64+
65+
nvml_system_get_driver_version = nvml_lib.nvmlSystemGetDriverVersion
66+
nvml_system_get_driver_version.argtypes = [ctypes.POINTER(ctypes.c_char), ctypes.c_uint]
67+
nvml_system_get_driver_version.restype = ctypes.c_int
68+
69+
nvml_shutdown = nvml_lib.nvmlShutdown
70+
nvml_shutdown.argtypes = []
71+
nvml_shutdown.restype = ctypes.c_int
72+
73+
init_status = nvml_init_v2()
74+
if init_status != _NVML_SUCCESS:
75+
raise RuntimeError(f"Failed to initialize NVML via nvmlInit_v2() (status={init_status}).")
76+
77+
try:
78+
version_buffer = ctypes.create_string_buffer(_NVML_SYSTEM_DRIVER_VERSION_BUFFER_LENGTH)
79+
status = nvml_system_get_driver_version(version_buffer, _NVML_SYSTEM_DRIVER_VERSION_BUFFER_LENGTH)
80+
if status != _NVML_SUCCESS:
81+
raise RuntimeError(
82+
f"Failed to query driver release version via nvmlSystemGetDriverVersion() (status={status})."
83+
)
84+
release_version = version_buffer.value.decode()
85+
except BaseException as exc:
86+
shutdown_status = nvml_shutdown()
87+
if shutdown_status != _NVML_SUCCESS:
88+
raise RuntimeError(f"Failed to shut down NVML via nvmlShutdown() (status={shutdown_status}).") from exc
89+
raise
90+
91+
shutdown_status = nvml_shutdown()
92+
if shutdown_status != _NVML_SUCCESS:
93+
raise RuntimeError(f"Failed to shut down NVML via nvmlShutdown() (status={shutdown_status}).")
94+
return release_version

‎cuda_pathfinder/tests/test_utils_driver_info.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,20 @@ def _loaded_cuda(abs_path: str) -> LoadedDL:
3535
)
3636

3737

38+
def test_query_driver_release_version_text():
39+
driver_info._load_nvidia_dynamic_lib.cache_clear()
40+
try:
41+
release_version = driver_info._query_driver_release_version_text()
42+
finally:
43+
driver_info._load_nvidia_dynamic_lib.cache_clear()
44+
45+
components = tuple(int(component) for component in release_version.split("."))
46+
assert len(components) in (2, 3)
47+
assert 400 <= components[0] < 1000
48+
for component in components[1:]:
49+
assert component >= 0
50+
51+
3852
def test_query_driver_version_uses_windll_on_windows(monkeypatch):
3953
fake_driver_lib = _FakeDriverLib(status=0, version=12080)
4054
loaded_paths: list[str] = []

0 commit comments

Comments
 (0)