Coverage for cuda/core/system/_system.pyx: 41.67%
48 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 01:23 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
6# cuda.core.system uses NVML through cuda.bindings.nvml. Every cuda-bindings
7# that cuda.core accepts provides the module. See cuda/core/_bindings_floor.py.
8# initialize() loads the NVML library itself on first use, so this module
9# stays importable without CUDA or NVML installed.
12from typing import TYPE_CHECKING
14# Always True, because the cuda-bindings floor made NVML support unconditional.
15# Kept for callers that read it. The assignment sits in a runtime-only block so
16# that the generated stub keeps the bare annotation the public API had. The API
17# check reports a changed attribute value otherwise.
18CUDA_BINDINGS_NVML_IS_COMPATIBLE: bool
19if not TYPE_CHECKING:
20 CUDA_BINDINGS_NVML_IS_COMPATIBLE = True
23# Please keep in sync with the equivalent implementation in
24# cuda_python_test_helpers/cuda_python_test_helpers/__init__.py.
25cdef bint _detect_wsl():
26 try:
27 with open("/proc/sys/kernel/osrelease") as f:
28 data = f.read().lower()
29 except OSError:
30 return False
31 return "microsoft" in data or "wsl" in data
34cdef bint _IS_WSL = _detect_wsl()
37# The WSL locale guard lives in a separate module that is only compiled on
38# Linux (build_hooks.py excludes it on Windows), because it relies on POSIX
39# per-thread locale APIs that MSVC does not provide. On non-Linux platforms
40# the import fails and we fall back to a no-op guard; _IS_WSL is then False
41# so the guard is never entered anyway.
42if _IS_WSL:
43 from cuda.core._utils._wsl_locale import c_locale_guard
44else:
45 c_locale_guard = None
48from cuda.bindings import nvml
49from cuda.core.system._nvml_context import initialize
52def get_user_mode_driver_version() -> tuple[int, ...]:
53 """
54 Get the user-mode (UMD / CUDA) driver version.
56 This is the most commonly needed version when checking CUDA driver
57 compatibility.
59 Returns
60 -------
61 version : tuple[int, ...]
62 A 2-tuple ``(MAJOR, MINOR)``, e.g. ``(13, 0)`` for CUDA 13.0.
63 """
64 cdef int v
65 initialize() 1b
66 v = nvml.system_get_cuda_driver_version() 1b
67 return (v // 1000, (v // 10) % 100) 1b
70def get_kernel_mode_driver_version() -> tuple[int, ...]:
71 """
72 Get the kernel-mode (KMD / GPU) driver version, e.g. 580.65.06.
74 Returns
75 -------
76 version : tuple[int, ...]
77 Typically a 3-tuple ``(MAJOR, MINOR, PATCH)``
78 (2-tuple on WSL), e.g. ``(580, 65, 6)``.
80 Raises
81 ------
82 RuntimeError
83 If the NVML library is not available.
84 """
85 initialize() 1c
86 return tuple(int(x) for x in nvml.system_get_driver_version().split(".")) 1c
89def get_nvml_version() -> tuple[int, ...]:
90 """
91 The version of the NVML library.
93 Returns
94 -------
95 version: tuple[int, ...]
96 Tuple of integers representing the NVML version components.
97 """
98 initialize() 1d
99 return tuple(int(v) for v in nvml.system_get_nvml_version().split(".")) 1d
102def get_driver_branch() -> str:
103 """
104 Retrieves the driver branch of the NVIDIA driver installed on the system.
106 Returns
107 -------
108 branch: str
109 The driver branch string (e.g., ``"560"``, ``"open"``, etc.).
110 """
111 initialize() 1e
112 return nvml.system_get_driver_branch() 1e
115def get_num_devices() -> int:
116 """
117 Return the number of devices in the system.
118 """
119 initialize() 1fghi
120 return nvml.device_get_count_v2() 1fghi
123def get_process_name(pid: int) -> str:
124 """
125 The name of process with given PID.
127 Parameters
128 ----------
129 pid: int
130 The PID of the process for which to get the name.
132 Returns
133 -------
134 name: str
135 The process name.
136 """
137 def _get_process_name(pid) -> str: 1a
138 # NVML caches process names on a per-PID basis when queried via
139 # nvmlSystemGetProcessName, and the cache is populated when enumerating
140 # running processes on devices. To ensure the name is cached for the
141 # requested PID, we walk all devices and query their running processes.
142 for i in range(nvml.device_get_count_v2()): 1a
143 try: 1a
144 dev_h = nvml.device_get_handle_by_index_v2(i) 1a
145 nvml.device_get_compute_running_processes_v3(dev_h) 1a
146 except nvml.NvmlError:
147 continue
148 return nvml.system_get_process_name(pid) 1a
150 initialize() 1a
151 if not _IS_WSL: 1a
152 return _get_process_name(pid) 1a
154 # WSL workaround: nvmlSystemGetProcessName on WSL takes a wide-char
155 # conversion path when the calling thread's locale is non-"C". That path
156 # walks a UTF-16LE source buffer with a 4-byte stride (as if it were
157 # UTF-32LE) and emits 5-byte UTF-8 sequences that look like garbage
158 # preceding the trailing basename of /proc/<pid>/exe. CPython's startup
159 # unconditionally calls setlocale(LC_ALL, ""), so essentially every
160 # cuda.core caller hits this. The cached entry for the PID is set the
161 # first time NVML resolves it (typically inside
162 # nvmlDeviceGetComputeRunningProcesses_v3), so to recover a correct value
163 # we re-prime the cache under the "C" locale before reading the name.
164 # c_locale_guard uses POSIX per-thread locale APIs (see _wsl_locale.pyx)
165 # so other threads' view of the locale is unaffected.
166 with c_locale_guard(): # no-cython-lint
167 return _get_process_name(pid)
170__all__ = [
171 "get_driver_branch",
172 "get_kernel_mode_driver_version",
173 "get_user_mode_driver_version",
174 "get_nvml_version",
175 "get_num_devices",
176 "get_process_name",
177 "CUDA_BINDINGS_NVML_IS_COMPATIBLE",
178]