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

1# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5  

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. 

9  

10  

11CUDA_BINDINGS_NVML_IS_COMPATIBLE: bool 

12  

13  

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 

23  

24  

25cdef bint _IS_WSL = _detect_wsl() 

26  

27  

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 

37  

38  

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)) 

45  

46  

47if CUDA_BINDINGS_NVML_IS_COMPATIBLE: 

48 try: 

49 from cuda.bindings import nvml 

50 except ImportError: 

51 CUDA_BINDINGS_NVML_IS_COMPATIBLE = False 

52  

53 from cuda.core.system._nvml_context import initialize 

54else: 

55 from cuda.core._utils.cuda_utils import driver, handle_return, runtime 

56  

57  

58def get_user_mode_driver_version() -> tuple[int, ...]: 

59 """ 

60 Get the user-mode (UMD / CUDA) driver version. 

61  

62 This is the most commonly needed version when checking CUDA driver 

63 compatibility. It works with all ``cuda-bindings`` versions. 

64  

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

77  

78  

79def get_kernel_mode_driver_version() -> tuple[int, ...]: 

80 """ 

81 Get the kernel-mode (KMD / GPU) driver version, e.g. 580.65.06. 

82  

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)``. 

88  

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

100  

101  

102def get_nvml_version() -> tuple[int, ...]: 

103 """ 

104 The version of the NVML library. 

105  

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

114  

115  

116def get_driver_branch() -> str: 

117 """ 

118 Retrieves the driver branch of the NVIDIA driver installed on the system. 

119  

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

129  

130  

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()) 

140  

141  

142def get_process_name(pid: int) -> str: 

143 """ 

144 The name of process with given PID. 

145  

146 Parameters 

147 ---------- 

148 pid: int 

149 The PID of the process for which to get the name. 

150  

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

168  

169 initialize() 1b

170 if not _IS_WSL: 1b

171 return _get_process_name(pid) 1b

172  

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) 

187  

188  

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]