Merge branch 'main' into style-guide

Author: Andy-Jost

cuda_core/cuda/core/system/__init__.py

@@ -23,16 +23,11 @@
 from ._system import *
 
 if CUDA_BINDINGS_NVML_IS_COMPATIBLE:
-    from ._device import Device, DeviceArchitecture
+    from ._device import *
+    from ._device import __all__ as _device_all
     from .exceptions import *
     from .exceptions import __all__ as _exceptions_all
 
-    __all__.extend(
-        [
-            "get_nvml_version",
-            "Device",
-            "DeviceArchitecture",
-        ]
-    )
-
+    __all__.append("get_nvml_version")
+    __all__.extend(_device_all)
     __all__.extend(_exceptions_all)

cuda_core/cuda/core/system/_device.pyx

@@ -175,7 +175,7 @@ cdef class Device:
     """
     Representation of a device.
 
-    ``cuda.core.system.Device`` provides access to various pieces of metadata
+    :class:`cuda.core.system.Device` provides access to various pieces of metadata
     about devices and their topology, as provided by the NVIDIA Management
     Library (NVML).  To use CUDA with a device, use :class:`cuda.core.Device`.
 
@@ -312,3 +312,12 @@ cdef class Device:
         board serial identifier.
         """
         return nvml.device_get_uuid(self._handle)
+
+
+__all__ = [
+    "BAR1MemoryInfo",
+    "Device",
+    "DeviceArchitecture",
+    "MemoryInfo",
+    "PciInfo",
+]

cuda_core/cuda/core/system/_system.pyx

@@ -82,6 +82,16 @@ def get_nvml_version() -> tuple[int, ...]:
     return tuple(int(v) for v in nvml.system_get_nvml_version().split("."))
 
 
+def get_driver_branch() -> str:
+    """
+    Retrieves the driver branch of the NVIDIA driver installed on the system.
+    """
+    if not CUDA_BINDINGS_NVML_IS_COMPATIBLE:
+        raise RuntimeError("NVML library is not available")
+    initialize()
+    return nvml.system_get_driver_branch()
+
+
 def get_num_devices() -> int:
     """
     Return the number of devices in the system.
@@ -112,6 +122,7 @@ def get_process_name(pid: int) -> str:
 
 
 __all__ = [
+    "get_driver_branch",
     "get_driver_version",
     "get_driver_version_full",
     "get_nvml_version",

cuda_core/docs/source/api.rst

@@ -79,6 +79,10 @@ CUDA system information and NVIDIA Management Library (NVML)
    :template: autosummary/cyclass.rst
 
    system.Device
+   system.DeviceArchitecture
+   system.MemoryInfo
+   system.BAR1MemoryInfo
+   system.PciInfo
 
 
 .. module:: cuda.core.utils

cuda_core/tests/system/test_system_system.py

@@ -95,3 +95,11 @@ def test_device_count():
     device_count = system.get_num_devices()
     assert isinstance(device_count, int)
     assert device_count >= 0
+
+
+@skip_if_nvml_unsupported
+def test_get_driver_branch():
+    driver_branch = system.get_driver_branch()
+    assert isinstance(driver_branch, str)
+    assert len(driver_branch) > 0
+    assert driver_branch[0] == "r"