Coverage for cuda/core/system/_system.pyx: 73.44%
64 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-13 01:38 +0000
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-13 01:38 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
6# This file needs to either use NVML exclusively, or when `cuda.bindings.nvml`
7# isn't available, fall back to non-NVML-based methods for backward
8# compatibility.
11CUDA_BINDINGS_NVML_IS_COMPATIBLE: bool
14# Please keep in sync with the equivalent implementation in
15# cuda_python_test_helpers/cuda_python_test_helpers/__init__.py.
16cdef bint _detect_wsl():
17 try:
18 with open("/proc/sys/kernel/osrelease") as f:
19 data = f.read().lower()
20 except OSError:
21 return False
22 return "microsoft" in data or "wsl" in data
25cdef bint _IS_WSL = _detect_wsl()
28# The WSL locale guard lives in a separate module that is only compiled on
29# Linux (build_hooks.py excludes it on Windows), because it relies on POSIX
30# per-thread locale APIs that MSVC does not provide. On non-Linux platforms
31# the import fails and we fall back to a no-op guard; _IS_WSL is then False
32# so the guard is never entered anyway.
33if _IS_WSL:
34 from cuda.core._utils._wsl_locale import c_locale_guard
35else:
36 c_locale_guard = None
39try:
40 from cuda.bindings._version import __version_tuple__ as _BINDINGS_VERSION
41except ImportError:
42 CUDA_BINDINGS_NVML_IS_COMPATIBLE = False
43else:
44 CUDA_BINDINGS_NVML_IS_COMPATIBLE = _BINDINGS_VERSION >= (13, 2, 0) or (_BINDINGS_VERSION[0] == 12 and _BINDINGS_VERSION[1:3] >= (9, 6))
47if CUDA_BINDINGS_NVML_IS_COMPATIBLE:
48 try:
49 from cuda.bindings import nvml
50 except ImportError:
51 CUDA_BINDINGS_NVML_IS_COMPATIBLE = False
53 from cuda.core.system._nvml_context import initialize
54else:
55 from cuda.core._utils.cuda_utils import driver, handle_return, runtime
58def get_user_mode_driver_version() -> tuple[int, ...]:
59 """
60 Get the user-mode (UMD / CUDA) driver version.
62 This is the most commonly needed version when checking CUDA driver
63 compatibility. It works with all ``cuda-bindings`` versions.
65 Returns
66 -------
67 version : tuple[int, ...]
68 A 2-tuple ``(MAJOR, MINOR)``, e.g. ``(13, 0)`` for CUDA 13.0.
69 """
70 cdef int v
71 if CUDA_BINDINGS_NVML_IS_COMPATIBLE: 1c
72 initialize() 1c
73 v = nvml.system_get_cuda_driver_version() 1c
74 else:
75 v = handle_return(driver.cuDriverGetVersion())
76 return (v // 1000, (v // 10) % 100) 1c
79def get_kernel_mode_driver_version() -> tuple[int, ...]:
80 """
81 Get the kernel-mode (KMD / GPU) driver version, e.g. 580.65.06.
83 Returns
84 -------
85 version : tuple[int, ...]
86 Typically a 3-tuple ``(MAJOR, MINOR, PATCH)``
87 (2-tuple on WSL), e.g. ``(580, 65, 6)``.
89 Raises
90 ------
91 RuntimeError
92 If the NVML library is not available.
93 """
94 if not CUDA_BINDINGS_NVML_IS_COMPATIBLE: 1d
95 raise RuntimeError(
96 "get_kernel_mode_driver_version requires NVML support"
97 )
98 initialize() 1d
99 return tuple(int(x) for x in nvml.system_get_driver_version().split(".")) 1d
102def get_nvml_version() -> tuple[int, ...]:
103 """
104 The version of the NVML library.
106 Returns
107 -------
108 version: tuple[int, ...]
109 Tuple of integers representing the NVML version components.
110 """
111 if not CUDA_BINDINGS_NVML_IS_COMPATIBLE: 19
112 raise RuntimeError("NVML library is not available")
113 return tuple(int(v) for v in nvml.system_get_nvml_version().split(".")) 19
116def get_driver_branch() -> str:
117 """
118 Retrieves the driver branch of the NVIDIA driver installed on the system.
120 Returns
121 -------
122 branch: str
123 The driver branch string (e.g., ``"560"``, ``"open"``, etc.).
124 """
125 if not CUDA_BINDINGS_NVML_IS_COMPATIBLE: 1e
126 raise RuntimeError("NVML library is not available")
127 initialize() 1e
128 return nvml.system_get_driver_branch() 1e
131def get_num_devices() -> int:
132 """
133 Return the number of devices in the system.
134 """
135 if CUDA_BINDINGS_NVML_IS_COMPATIBLE: 1fghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ012345678
136 initialize() 1fghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ012345678
137 return nvml.device_get_count_v2() 1fghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ012345678
138 else:
139 return handle_return(runtime.cudaGetDeviceCount())
142def get_process_name(pid: int) -> str:
143 """
144 The name of process with given PID.
146 Parameters
147 ----------
148 pid: int
149 The PID of the process for which to get the name.
151 Returns
152 -------
153 name: str
154 The process name.
155 """
156 def _get_process_name(pid) -> str: 1b
157 # NVML caches process names on a per-PID basis when queried via
158 # nvmlSystemGetProcessName, and the cache is populated when enumerating
159 # running processes on devices. To ensure the name is cached for the
160 # requested PID, we walk all devices and query their running processes.
161 for i in range(nvml.device_get_count_v2()): 1b
162 try: 1b
163 dev_h = nvml.device_get_handle_by_index_v2(i) 1b
164 nvml.device_get_compute_running_processes_v3(dev_h) 1b
165 except nvml.NvmlError:
166 continue
167 return nvml.system_get_process_name(pid) 1b
169 initialize() 1b
170 if not _IS_WSL: 1b
171 return _get_process_name(pid) 1b
173 # WSL workaround: nvmlSystemGetProcessName on WSL takes a wide-char
174 # conversion path when the calling thread's locale is non-"C". That path
175 # walks a UTF-16LE source buffer with a 4-byte stride (as if it were
176 # UTF-32LE) and emits 5-byte UTF-8 sequences that look like garbage
177 # preceding the trailing basename of /proc/<pid>/exe. CPython's startup
178 # unconditionally calls setlocale(LC_ALL, ""), so essentially every
179 # cuda.core caller hits this. The cached entry for the PID is set the
180 # first time NVML resolves it (typically inside
181 # nvmlDeviceGetComputeRunningProcesses_v3), so to recover a correct value
182 # we re-prime the cache under the "C" locale before reading the name.
183 # c_locale_guard uses POSIX per-thread locale APIs (see _wsl_locale.pyx)
184 # so other threads' view of the locale is unaffected.
185 with c_locale_guard(): # no-cython-lint
186 return _get_process_name(pid)
189__all__ = [
190 "get_driver_branch",
191 "get_kernel_mode_driver_version",
192 "get_user_mode_driver_version",
193 "get_nvml_version",
194 "get_num_devices",
195 "get_process_name",
196 "CUDA_BINDINGS_NVML_IS_COMPATIBLE",
197]