cuda.core: consolidate driver version query functions (#2040)

* cuda.core: consolidate get_driver_version and get_driver_version_full

Merge the two separate driver version query functions into a single
get_driver_version() that returns (umd_version, kmd_version) — a pair
of version tuples. UMD is a 2-tuple (major, minor) and KMD is a
3-tuple (major, minor, patch). The function now requires NVML.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Remove unused driver import from _system.pyx

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Fix KMD version test to allow 2-tuple on WSL

Windows driver version strings only have two components, so
nvmlSystemGetDriverVersion returns X.Y on WSL instead of X.Y.Z.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Split into get_user_mode_driver_version and get_kernel_mode_driver_version

Per review feedback, split the consolidated get_driver_version() into
two separate functions so each returns a simple tuple[int, ...]:

- get_user_mode_driver_version(): works with or without NVML (falls
  back to cuDriverGetVersion), returns (major, minor)
- get_kernel_mode_driver_version(): requires NVML, returns
  (major, minor, patch) or (major, minor) on WSL

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Add example version strings to driver version docstrings

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Remove NVML implementation detail from UMD docstring

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Sort __all__ alphabetically in system/__init__.py

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: leofang

cuda_core/cuda/core/system/__init__.py

@@ -11,10 +11,10 @@
 
 __all__ = [
     "CUDA_BINDINGS_NVML_IS_COMPATIBLE",
-    "get_driver_version",
-    "get_driver_version_full",
+    "get_kernel_mode_driver_version",
     "get_num_devices",
     "get_process_name",
+    "get_user_mode_driver_version",
 ]
 
 

cuda_core/cuda/core/system/_system.pyx

@@ -29,52 +29,48 @@ else:
     from cuda.core._utils.cuda_utils import driver, handle_return, runtime
 
 
-def get_driver_version(kernel_mode: bool = False) -> tuple[int, int]:
+def get_user_mode_driver_version() -> tuple[int, ...]:
     """
-    Get the driver version.
+    Get the user-mode (UMD / CUDA) driver version.
 
-    Parameters
-    ----------
-    kernel_mode: bool
-        When `True`, return the kernel-mode driver version, e.g. 580.65.06.
-        Otherwise, return the user-mode driver version, e.g. 13.0.1.
+    This is the most commonly needed version when checking CUDA driver
+    compatibility.  It works with all ``cuda-bindings`` versions.
 
     Returns
     -------
-    version: tuple[int, int]
-        Tuple in the format `(MAJOR, MINOR)`.
+    version : tuple[int, ...]
+        A 2-tuple ``(MAJOR, MINOR)``, e.g. ``(13, 0)`` for CUDA 13.0.
     """
-    return get_driver_version_full(kernel_mode)[:2]
+    cdef int v
+    if CUDA_BINDINGS_NVML_IS_COMPATIBLE:
+        initialize()
+        v = nvml.system_get_cuda_driver_version()
+    else:
+        v = handle_return(driver.cuDriverGetVersion())
+    return (v // 1000, (v // 10) % 100)
 
 
-def get_driver_version_full(kernel_mode: bool = False) -> tuple[int, int, int]:
+def get_kernel_mode_driver_version() -> tuple[int, ...]:
     """
-    Get the full driver version.
-
-    Parameters
-    ----------
-    kernel_mode: bool
-        When `True`, return the kernel-mode driver version, e.g. 580.65.06.
-        Otherwise, return the user-mode driver version, e.g. 13.0.1.
+    Get the kernel-mode (KMD / GPU) driver version, e.g. 580.65.06.
 
     Returns
     -------
-    version: tuple[int, int, int]
-        Tuple in the format `(MAJOR, MINOR, PATCH)`.
+    version : tuple[int, ...]
+        Typically a 3-tuple ``(MAJOR, MINOR, PATCH)``
+        (2-tuple on WSL), e.g. ``(580, 65, 6)``.
+
+    Raises
+    ------
+    RuntimeError
+        If the NVML library is not available.
     """
-    cdef int v
-    if kernel_mode:
-        if not CUDA_BINDINGS_NVML_IS_COMPATIBLE:
-            raise ValueError("Kernel-mode driver version requires NVML support")
-        initialize()
-        return tuple(int(v) for v in nvml.system_get_driver_version().split("."))
-    else:
-        if CUDA_BINDINGS_NVML_IS_COMPATIBLE:
-            initialize()
-            v = nvml.system_get_cuda_driver_version()
-        else:
-            v = handle_return(driver.cuDriverGetVersion())
-        return (v // 1000, (v // 10) % 100, v % 10)
+    if not CUDA_BINDINGS_NVML_IS_COMPATIBLE:
+        raise RuntimeError(
+            "get_kernel_mode_driver_version requires NVML support"
+        )
+    initialize()
+    return tuple(int(x) for x in nvml.system_get_driver_version().split("."))
 
 
 def get_nvml_version() -> tuple[int, ...]:
@@ -137,8 +133,8 @@ def get_process_name(pid: int) -> str:
 
 __all__ = [
     "get_driver_branch",
-    "get_driver_version",
-    "get_driver_version_full",
+    "get_kernel_mode_driver_version",
+    "get_user_mode_driver_version",
     "get_nvml_version",
     "get_num_devices",
     "get_process_name",

cuda_core/docs/source/api.rst

@@ -274,8 +274,8 @@ Basic functions
 .. autosummary::
    :toctree: generated/
 
-   system.get_driver_version
-   system.get_driver_version_full
+   system.get_user_mode_driver_version
+   system.get_kernel_mode_driver_version
    system.get_driver_branch
    system.get_num_devices
    system.get_nvml_version

cuda_core/docs/source/release/1.0.0-notes.rst

@@ -163,6 +163,12 @@ Breaking changes
   :mod:`cuda.core.utils` module.
   (`#2028 <https://github.com/NVIDIA/cuda-python/issues/2028>`__)
 
+- Replaced ``system.get_driver_version()`` and
+  ``system.get_driver_version_full()`` with
+  :func:`system.get_user_mode_driver_version` (works with or without
+  NVML) and :func:`system.get_kernel_mode_driver_version` (requires
+  NVML).  Each returns a ``tuple[int, ...]``.
+
 Fixes and enhancements
 -----------------------
 

cuda_core/tests/system/test_system_system.py

@@ -19,11 +19,25 @@
 from .conftest import skip_if_nvml_unsupported
 
 
-def test_driver_version():
-    driver_version = system.get_driver_version()
+def test_user_mode_driver_version():
+    umd = system.get_user_mode_driver_version()
+    assert isinstance(umd, tuple)
+    assert len(umd) == 2
     version = handle_return(driver.cuDriverGetVersion())
-    expected_driver_version = (version // 1000, (version % 1000) // 10)
-    assert driver_version == expected_driver_version, "Driver version does not match expected value"
+    expected = (version // 1000, (version % 1000) // 10)
+    assert umd == expected, "UMD driver version does not match expected value"
+
+
+@skip_if_nvml_unsupported
+def test_kernel_mode_driver_version():
+    kmd = system.get_kernel_mode_driver_version()
+    assert isinstance(kmd, tuple)
+    assert len(kmd) in (2, 3)
+    ver_maj, ver_min, *ver_patch = kmd
+    assert 400 <= ver_maj < 1000
+    assert ver_min >= 0
+    if ver_patch:
+        assert 0 <= ver_patch[0] <= 99
 
 
 def test_num_devices():
@@ -41,28 +55,11 @@ def test_devices():
         assert device.device_id == expected_device.device_id, "Device ID does not match expected value"
 
 
-def test_cuda_driver_version():
-    cuda_driver_version = system.get_driver_version_full()
-    assert isinstance(cuda_driver_version, tuple)
-    assert len(cuda_driver_version) == 3
-
-    ver_maj, ver_min, ver_patch = cuda_driver_version
-    assert ver_maj >= 10
-    assert 0 <= ver_min <= 99
-    assert 0 <= ver_patch <= 9
-
-
-@skip_if_nvml_unsupported
-def test_gpu_driver_version():
-    driver_version = system.get_driver_version(kernel_mode=True)
-    assert isinstance(driver_version, tuple)
-    assert len(driver_version) in (2, 3)
-
-    (ver_maj, ver_min, *ver_patch) = driver_version
-    assert 400 <= ver_maj < 1000
-    assert ver_min >= 0
-    if ver_patch:
-        assert 0 <= ver_patch[0] <= 99
+def test_kernel_mode_driver_version_requires_nvml():
+    if system.CUDA_BINDINGS_NVML_IS_COMPATIBLE:
+        pytest.skip("NVML is available, cannot test the error path")
+    with pytest.raises(RuntimeError, match="requires NVML support"):
+        system.get_kernel_mode_driver_version()
 
 
 @skip_if_nvml_unsupported