Finalize the pathfinder CUDA driver version query API.

Expose DriverCudaVersion, QueryDriverCudaVersionError, and query_driver_cuda_version publicly, and align the internal naming, caching, docs, and test coverage around the CUDA-specific driver version query.

Made-with: Cursor

Author: rwgk

cuda_pathfinder/cuda/pathfinder/__init__.py

@@ -59,6 +59,9 @@
 from cuda.pathfinder._static_libs.find_static_lib import (
     locate_static_lib as locate_static_lib,
 )
+from cuda.pathfinder._utils.driver_info import DriverCudaVersion as DriverCudaVersion
+from cuda.pathfinder._utils.driver_info import QueryDriverCudaVersionError as QueryDriverCudaVersionError
+from cuda.pathfinder._utils.driver_info import query_driver_cuda_version as query_driver_cuda_version
 from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home as get_cuda_path_or_home
 
 from cuda.pathfinder._version import __version__  # isort: skip

cuda_pathfinder/cuda/pathfinder/_utils/driver_info.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import ctypes
+import functools
 from collections.abc import Callable
 from dataclasses import dataclass
 
@@ -13,6 +14,10 @@
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
 
+class QueryDriverCudaVersionError(RuntimeError):
+    """Raised when ``query_driver_cuda_version()`` cannot determine the CUDA driver version."""
+
+
 @dataclass(frozen=True, slots=True)
 class DriverCudaVersion:
     """
@@ -38,17 +43,21 @@ class DriverCudaVersion:
     minor: int
 
 
-def query_driver_version() -> DriverCudaVersion:
+@functools.cache
+def query_driver_cuda_version() -> DriverCudaVersion:
     """Return the CUDA driver version parsed into its major/minor components."""
-    encoded = _query_driver_version_int()
-    return DriverCudaVersion(
-        encoded=encoded,
-        major=encoded // 1000,
-        minor=(encoded % 1000) // 10,
-    )
+    try:
+        encoded = _query_driver_cuda_version_int()
+        return DriverCudaVersion(
+            encoded=encoded,
+            major=encoded // 1000,
+            minor=(encoded % 1000) // 10,
+        )
+    except Exception as exc:
+        raise QueryDriverCudaVersionError("Failed to query the CUDA driver version.") from exc
 
 
-def _query_driver_version_int() -> int:
+def _query_driver_cuda_version_int() -> int:
     """Return the encoded CUDA driver version from ``cuDriverGetVersion()``."""
     loaded_cuda = _load_nvidia_dynamic_lib("cuda")
     if IS_WINDOWS:

cuda_pathfinder/docs/source/api.rst

@@ -18,6 +18,10 @@ CUDA bitcode and static libraries.
 
    get_cuda_path_or_home
 
+   DriverCudaVersion
+   QueryDriverCudaVersionError
+   query_driver_cuda_version
+
    SUPPORTED_NVIDIA_LIBNAMES
    load_nvidia_dynamic_lib
    LoadedDL

cuda_pathfinder/tests/test_driver_lib_loading.py

@@ -160,17 +160,19 @@ def raise_child_process_failed():
         assert os.path.isfile(abs_path)
 
 
-def test_real_query_driver_version(info_summary_append):
+def test_real_query_driver_cuda_version(info_summary_append):
     driver_info._load_nvidia_dynamic_lib.cache_clear()
+    driver_info.query_driver_cuda_version.cache_clear()
     try:
-        version = driver_info.query_driver_version()
-    except Exception as exc:
+        version = driver_info.query_driver_cuda_version()
+    except driver_info.QueryDriverCudaVersionError as exc:
         if STRICTNESS == "all_must_work":
             raise
         info_summary_append(f"driver version unavailable: {exc.__class__.__name__}: {exc}")
         return
     finally:
         driver_info._load_nvidia_dynamic_lib.cache_clear()
+        driver_info.query_driver_cuda_version.cache_clear()
 
     info_summary_append(f"driver_version={version.major}.{version.minor} (encoded={version.encoded})")
     assert version.encoded > 0

cuda_pathfinder/tests/test_utils_driver_info.py

@@ -9,6 +9,13 @@
 from cuda.pathfinder._utils import driver_info
 
 
+@pytest.fixture(autouse=True)
+def _clear_driver_cuda_version_query_cache():
+    driver_info.query_driver_cuda_version.cache_clear()
+    yield
+    driver_info.query_driver_cuda_version.cache_clear()
+
+
 class _FakeCuDriverGetVersion:
     def __init__(self, *, status: int, version: int):
         self.argtypes = None
@@ -35,7 +42,7 @@ def _loaded_cuda(abs_path: str) -> LoadedDL:
     )
 
 
-def test_query_driver_version_uses_windll_on_windows(monkeypatch):
+def test_query_driver_cuda_version_uses_windll_on_windows(monkeypatch):
     fake_driver_lib = _FakeDriverLib(status=0, version=12080)
     loaded_paths: list[str] = []
 
@@ -52,26 +59,43 @@ def fake_windll(abs_path: str):
 
     monkeypatch.setattr(driver_info.ctypes, "WinDLL", fake_windll, raising=False)
 
-    assert driver_info._query_driver_version_int() == 12080
+    assert driver_info._query_driver_cuda_version_int() == 12080
     assert loaded_paths == [r"C:\Windows\System32\nvcuda.dll"]
 
 
-def test_query_driver_version_returns_parsed_dataclass(monkeypatch):
-    monkeypatch.setattr(driver_info, "_query_driver_version_int", lambda: 12080)
+def test_query_driver_cuda_version_returns_parsed_dataclass(monkeypatch):
+    monkeypatch.setattr(driver_info, "_query_driver_cuda_version_int", lambda: 12080)
 
-    assert driver_info.query_driver_version() == driver_info.DriverCudaVersion(
+    assert driver_info.query_driver_cuda_version() == driver_info.DriverCudaVersion(
         encoded=12080,
         major=12,
         minor=8,
     )
 
 
-def test_query_driver_version_int_raises_when_cuda_call_fails(monkeypatch):
+def test_query_driver_cuda_version_wraps_internal_failures(monkeypatch):
+    root_cause = RuntimeError("low-level query failed")
+
+    def fail_query_driver_cuda_version_int() -> int:
+        raise root_cause
+
+    monkeypatch.setattr(driver_info, "_query_driver_cuda_version_int", fail_query_driver_cuda_version_int)
+
+    with pytest.raises(
+        driver_info.QueryDriverCudaVersionError,
+        match="Failed to query the CUDA driver version",
+    ) as exc_info:
+        driver_info.query_driver_cuda_version()
+
+    assert exc_info.value.__cause__ is root_cause
+
+
+def test_query_driver_cuda_version_int_raises_when_cuda_call_fails(monkeypatch):
     fake_driver_lib = _FakeDriverLib(status=1, version=0)
 
     monkeypatch.setattr(driver_info, "IS_WINDOWS", False)
     monkeypatch.setattr(driver_info, "_load_nvidia_dynamic_lib", lambda _libname: _loaded_cuda("/usr/lib/libcuda.so.1"))
     monkeypatch.setattr(driver_info.ctypes, "CDLL", lambda _abs_path: fake_driver_lib)
 
     with pytest.raises(RuntimeError, match=r"cuDriverGetVersion\(\) \(status=1\)"):
-        driver_info._query_driver_version_int()
+        driver_info._query_driver_cuda_version_int()