|
12 | 12 | ) |
13 | 13 | from cuda.pathfinder._utils.platform_aware import IS_WINDOWS |
14 | 14 |
|
| 15 | +_NVML_SUCCESS = 0 |
| 16 | +_NVML_SYSTEM_DRIVER_VERSION_BUFFER_LENGTH = 80 |
| 17 | + |
15 | 18 |
|
16 | 19 | @dataclass(frozen=True, slots=True) |
17 | 20 | class DriverCudaVersion: |
@@ -48,3 +51,44 @@ def _query_driver_version_int() -> int: |
48 | 51 | if status != 0: |
49 | 52 | raise RuntimeError(f"Failed to query CUDA driver version via cuDriverGetVersion() (status={status}).") |
50 | 53 | 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 |
0 commit comments