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