pathfinder: add display-driver release metadata and NVML query helper

rwgk · cursoragent · rwgk · commit 10ac5b5d8e72 · 2026-05-03T20:18:54.000-07:00
Introduce ``DriverReleaseVersion`` and ``query_driver_release_version()``
in ``cuda.pathfinder._utils.driver_info`` so callers can read the
display-driver release shown as ``Driver Version`` in ``nvidia-smi``
(for example ``595.58.03``, branch ``595``). The helper queries
``nvmlSystemGetDriverVersion()`` via ``ctypes`` against the NVML library
loaded through pathfinder's existing dynamic-lib loader, and parses the
result into a frozen dataclass exposing both the raw text and the
branch number used by NVIDIA's published minor-version compatibility
tables.

Tests cover the dataclass parser, the cache-clear lifecycle, the
``QueryDriverReleaseVersionError`` wrapping, and a fake-NVML
end-to-end path that asserts the shutdown call always runs.

Co-authored-by: Cursor &lt;cursoragent@cursor.com&gt;
diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/driver_info.py b/cuda_pathfinder/cuda/pathfinder/_utils/driver_info.py
@@ -5,6 +5,7 @@
 
 import ctypes
 import functools
+import re
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import cast
@@ -15,11 +16,19 @@
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 from cuda.pathfinder._utils.toolkit_info import EncodedCudaVersion
 
+_NVML_SUCCESS = 0
+_NVML_SYSTEM_DRIVER_VERSION_BUFFER_LENGTH = 80
+_DRIVER_RELEASE_VERSION_RE = re.compile(r"^\d+(?:\.\d+){1,2}$")
+
 
 class QueryDriverCudaVersionError(RuntimeError):
     """Raised when ``query_driver_cuda_version()`` cannot determine the CUDA driver version."""
 
 
+class QueryDriverReleaseVersionError(RuntimeError):
+    """Raised when ``query_driver_release_version()`` cannot determine the display-driver release version."""
+
+
 @dataclass(frozen=True, slots=True)
 class DriverCudaVersion(EncodedCudaVersion):
     """
@@ -44,6 +53,37 @@ class DriverCudaVersion(EncodedCudaVersion):
     """
 
 
+@dataclass(frozen=True, slots=True)
+class DriverReleaseVersion:
+    """
+    Display-driver release version shown as ``Driver Version`` in ``nvidia-smi``.
+
+    Example ``nvidia-smi`` output::
+
+        +---------------------------------------------------------------------+
+        | NVIDIA-SMI 595.58.03  Driver Version: 595.58.03  CUDA Version: 13.2 |
+        +---------------------------------------------------------------------+
+
+    For the example above, ``DriverReleaseVersion(text="595.58.03",
+    components=(595, 58, 3), branch=595)`` corresponds to ``Driver Version:
+    595.58.03``. The ``branch`` field is the first numeric component because
+    NVIDIA's compatibility docs publish minimum display-driver requirements in
+    branch form such as ``>= 580`` for CUDA 13.x minor-version compatibility.
+    """
+
+    text: str
+    components: tuple[int, ...]
+    branch: int
+
+    @classmethod
+    def from_text(cls, text: str) -> DriverReleaseVersion:
+        normalized_text = text.strip()
+        if not _DRIVER_RELEASE_VERSION_RE.fullmatch(normalized_text):
+            raise ValueError(f"Invalid driver release version text: {text!r}")
+        components = tuple(int(component) for component in normalized_text.split("."))
+        return cls(text=normalized_text, components=components, branch=components[0])
+
+
 @functools.cache
 def query_driver_cuda_version() -> DriverCudaVersion:
     """Return the CUDA driver version parsed into its major/minor components."""
@@ -54,6 +94,15 @@ def query_driver_cuda_version() -> DriverCudaVersion:
         raise QueryDriverCudaVersionError("Failed to query the CUDA driver version.") from exc
 
 
+@functools.cache
+def query_driver_release_version() -> DriverReleaseVersion:
+    """Return the display-driver release version parsed into branch/components."""
+    try:
+        return DriverReleaseVersion.from_text(_query_driver_release_version_text())
+    except Exception as exc:
+        raise QueryDriverReleaseVersionError("Failed to query the display-driver release version.") from exc
+
+
 def _query_driver_cuda_version_int() -> int:
     """Return the encoded CUDA driver version from ``cuDriverGetVersion()``."""
     loaded_cuda = _load_nvidia_dynamic_lib("cuda")
@@ -72,3 +121,44 @@ def _query_driver_cuda_version_int() -> int:
     if status != 0:
         raise RuntimeError(f"Failed to query CUDA driver version via cuDriverGetVersion() (status={status}).")
     return version.value
+
+
+def _query_driver_release_version_text() -> str:
+    """Return the display-driver release version from ``nvmlSystemGetDriverVersion()``."""
+    loaded_nvml = _load_nvidia_dynamic_lib("nvml")
+    nvml_lib = ctypes.CDLL(loaded_nvml.abs_path)
+
+    nvml_init_v2 = nvml_lib.nvmlInit_v2
+    nvml_init_v2.argtypes = []
+    nvml_init_v2.restype = ctypes.c_int
+
+    nvml_system_get_driver_version = nvml_lib.nvmlSystemGetDriverVersion
+    nvml_system_get_driver_version.argtypes = [ctypes.POINTER(ctypes.c_char), ctypes.c_uint]
+    nvml_system_get_driver_version.restype = ctypes.c_int
+
+    nvml_shutdown = nvml_lib.nvmlShutdown
+    nvml_shutdown.argtypes = []
+    nvml_shutdown.restype = ctypes.c_int
+
+    init_status = nvml_init_v2()
+    if init_status != _NVML_SUCCESS:
+        raise RuntimeError(f"Failed to initialize NVML via nvmlInit_v2() (status={init_status}).")
+
+    try:
+        version_buffer = ctypes.create_string_buffer(_NVML_SYSTEM_DRIVER_VERSION_BUFFER_LENGTH)
+        status = nvml_system_get_driver_version(version_buffer, _NVML_SYSTEM_DRIVER_VERSION_BUFFER_LENGTH)
+        if status != _NVML_SUCCESS:
+            raise RuntimeError(
+                f"Failed to query driver release version via nvmlSystemGetDriverVersion() (status={status})."
+            )
+        release_version = version_buffer.value.decode()
+    except BaseException as exc:
+        shutdown_status = nvml_shutdown()
+        if shutdown_status != _NVML_SUCCESS:
+            raise RuntimeError(f"Failed to shut down NVML via nvmlShutdown() (status={shutdown_status}).") from exc
+        raise
+
+    shutdown_status = nvml_shutdown()
+    if shutdown_status != _NVML_SUCCESS:
+        raise RuntimeError(f"Failed to shut down NVML via nvmlShutdown() (status={shutdown_status}).")
+    return release_version
diff --git a/cuda_pathfinder/tests/test_utils_driver_info.py b/cuda_pathfinder/tests/test_utils_driver_info.py
@@ -12,8 +12,10 @@
 @pytest.fixture(autouse=True)
 def _clear_driver_cuda_version_query_cache():
     driver_info.query_driver_cuda_version.cache_clear()
+    driver_info.query_driver_release_version.cache_clear()
     yield
     driver_info.query_driver_cuda_version.cache_clear()
+    driver_info.query_driver_release_version.cache_clear()
 
 
 class _FakeCuDriverGetVersion:
@@ -33,6 +35,47 @@ def __init__(self, *, status: int, version: int):
         self.cuDriverGetVersion = _FakeCuDriverGetVersion(status=status, version=version)
 
 
+class _FakeNvmlFunction:
+    def __init__(self, func):
+        self.argtypes = None
+        self.restype = None
+        self._func = func
+
+    def __call__(self, *args):
+        return self._func(*args)
+
+
+class _FakeNvmlLib:
+    def __init__(
+        self,
+        *,
+        init_status: int = 0,
+        driver_release_version: str = "595.58.03",
+        query_status: int = 0,
+        shutdown_statuses: tuple[int, ...] = (0,),
+    ):
+        self.shutdown_calls = 0
+        remaining_shutdown_statuses = list(shutdown_statuses)
+
+        self.nvmlInit_v2 = _FakeNvmlFunction(lambda: init_status)
+
+        def nvml_system_get_driver_version(version_buffer, _buffer_length) -> int:
+            if query_status != 0:
+                return query_status
+            version_buffer.value = driver_release_version.encode()
+            return 0
+
+        self.nvmlSystemGetDriverVersion = _FakeNvmlFunction(nvml_system_get_driver_version)
+
+        def nvml_shutdown() -> int:
+            self.shutdown_calls += 1
+            if remaining_shutdown_statuses:
+                return remaining_shutdown_statuses.pop(0)
+            return 0
+
+        self.nvmlShutdown = _FakeNvmlFunction(nvml_shutdown)
+
+
 def _loaded_cuda(abs_path: str) -> LoadedDL:
     return LoadedDL(
         abs_path=abs_path,
@@ -42,6 +85,86 @@ def _loaded_cuda(abs_path: str) -> LoadedDL:
     )
 
 
+def _loaded_nvml(abs_path: str) -> LoadedDL:
+    return LoadedDL(
+        abs_path=abs_path,
+        was_already_loaded_from_elsewhere=False,
+        _handle_uint=0xCAFE,
+        found_via="system-search",
+    )
+
+
+def test_driver_release_version_from_text_parses_branch():
+    assert driver_info.DriverReleaseVersion.from_text("595.58.03") == driver_info.DriverReleaseVersion(
+        text="595.58.03",
+        components=(595, 58, 3),
+        branch=595,
+    )
+
+
+def test_query_driver_release_version_returns_parsed_dataclass(monkeypatch):
+    monkeypatch.setattr(driver_info, "_query_driver_release_version_text", lambda: "595.58.03")
+
+    assert driver_info.query_driver_release_version() == driver_info.DriverReleaseVersion(
+        text="595.58.03",
+        components=(595, 58, 3),
+        branch=595,
+    )
+
+
+def test_query_driver_release_version_wraps_internal_failures(monkeypatch):
+    root_cause = RuntimeError("low-level release query failed")
+
+    def fail_query_driver_release_version_text() -> str:
+        raise root_cause
+
+    monkeypatch.setattr(driver_info, "_query_driver_release_version_text", fail_query_driver_release_version_text)
+
+    with pytest.raises(
+        driver_info.QueryDriverReleaseVersionError,
+        match="Failed to query the display-driver release version",
+    ) as exc_info:
+        driver_info.query_driver_release_version()
+
+    assert exc_info.value.__cause__ is root_cause
+
+
+def test_query_driver_release_version_text_uses_nvml(monkeypatch):
+    fake_nvml_lib = _FakeNvmlLib(driver_release_version="595.58.03")
+    loaded_paths: list[str] = []
+
+    monkeypatch.setattr(
+        driver_info,
+        "_load_nvidia_dynamic_lib",
+        lambda _libname: _loaded_nvml("/usr/lib/libnvidia-ml.so.1"),
+    )
+
+    def fake_cdll(abs_path: str):
+        loaded_paths.append(abs_path)
+        return fake_nvml_lib
+
+    monkeypatch.setattr(driver_info.ctypes, "CDLL", fake_cdll)
+
+    assert driver_info._query_driver_release_version_text() == "595.58.03"
+    assert loaded_paths == ["/usr/lib/libnvidia-ml.so.1"]
+    assert fake_nvml_lib.shutdown_calls == 1
+
+
+def test_query_driver_release_version_text_raises_when_nvml_call_fails(monkeypatch):
+    fake_nvml_lib = _FakeNvmlLib(query_status=1)
+
+    monkeypatch.setattr(
+        driver_info,
+        "_load_nvidia_dynamic_lib",
+        lambda _libname: _loaded_nvml("/usr/lib/libnvidia-ml.so.1"),
+    )
+    monkeypatch.setattr(driver_info.ctypes, "CDLL", lambda _abs_path: fake_nvml_lib)
+
+    with pytest.raises(RuntimeError, match=r"nvmlSystemGetDriverVersion\(\) \(status=1\)"):
+        driver_info._query_driver_release_version_text()
+    assert fake_nvml_lib.shutdown_calls == 1
+
+
 def test_query_driver_cuda_version_uses_windll_on_windows(monkeypatch):
     fake_driver_lib = _FakeDriverLib(status=0, version=12080)
     loaded_paths: list[str] = []
