Coverage for cuda/core/system/_system_events.pyx: 73.91%
46 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-05 01:11 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
6from libc.stdint cimport intptr_t
8from cuda.bindings import nvml
10from ._nvml_context cimport initialize
12from . import _device
13from cuda.core.system.typing import SystemEventType
16_SYSTEM_EVENT_TYPE_MAPPING = {
17 nvml.SystemEventType.GPU_DRIVER_UNBIND: SystemEventType.UNBIND,
18 nvml.SystemEventType.GPU_DRIVER_BIND: SystemEventType.BIND,
19}
22_SYSTEM_EVENT_TYPE_INV_MAPPING = {v: k for k, v in _SYSTEM_EVENT_TYPE_MAPPING.items()}
25def _pci_bus_id_from_gpu_id(gpu_id: int) -> str:
26 """
27 Decode an NVML System Event packed ``gpu_id`` into an NVML-style PCI bus ID
28 string.
29 """
30 domain = (gpu_id >> 16) & 0xFFFF 1fghic
31 bus = (gpu_id >> 8) & 0xFF 1fghic
32 device = gpu_id & 0xFF 1fghic
33 return f"{domain:08X}:{bus:02X}:{device:02X}.0" 1fghic
36cdef class SystemEvent:
37 """
38 Data about a collection of system events.
39 """
40 cdef object _event_data
42 def __init__(self, event_data not None: nvml.SystemEventData_v1):
43 if len(event_data) != 1: 1cjd
44 raise ValueError(f"SystemEvent expects exactly one event, got {len(event_data)}") 1j
45 self._event_data = event_data 1cd
47 @property
48 def event_type(self) -> SystemEventType:
49 """
50 The :obj:`~SystemEventType` that was triggered.
51 """
52 return _SYSTEM_EVENT_TYPE_MAPPING[self._event_data.event_type] 1d
54 @property
55 def gpu_id(self) -> int:
56 """
57 The GPU ID in PCI ID format.
58 """
59 return self._event_data.gpu_id 1cd
61 @property
62 def device(self) -> _device.Device:
63 """
64 The :obj:`~_device.Device` associated with this event.
65 """
66 return _device.Device(pci_bus_id=_pci_bus_id_from_gpu_id(self.gpu_id)) 1c
69cdef class SystemEvents:
70 """
71 Data about a collection of system events.
72 """
73 cdef object _event_data
75 def __init__(self, event_data not None: nvml.SystemEventData_v1):
76 self._event_data = event_data 1d
78 def __len__(self) -> int:
79 return len(self._event_data) 1d
81 def __getitem__(self, Py_ssize_t idx) -> SystemEvent:
82 """
83 Get the :obj:`~_system_events.SystemEvent` at the specified index.
84 """
85 return SystemEvent(self._event_data[idx]) 1d
88cdef class RegisteredSystemEvents:
89 """
90 Represents a set of events that can be waited on for a specific device.
91 """
92 cdef intptr_t _event_set
94 def __init__(self, events: SystemEventType | str | list[SystemEventType | str]):
95 cdef unsigned long long event_bitmask
96 if isinstance(events, (str, SystemEventType)): 1bea
97 events = [events]
99 if isinstance(events, list): 1bea
100 event_bitmask = 0 1ba
101 for ev in events: 1ba
102 try: 1ba
103 ev_enum = _SYSTEM_EVENT_TYPE_INV_MAPPING[ev] 1ba
104 except KeyError: 1b
105 raise ValueError( 1b
106 f"Invalid event type: {ev}. " 1b
107 f"Must be one of {list(SystemEventType.__members__.values())}" 1b
108 ) from None 1b
109 event_bitmask |= <unsigned long long>int(ev_enum) 1a
110 else:
111 raise TypeError("events must be an SystemEventType, str, or list of SystemEventType or str") 1e
113 initialize() 1a
115 self._event_set = 0 1a
116 self._event_set = nvml.system_event_set_create() 1a
117 # If this raises, the event needs to be freed and this is handled by
118 # this class's __dealloc__ method.
119 nvml.system_register_events(event_bitmask, self._event_set) 1a
121 def __dealloc__(self) -> None:
122 if self._event_set != 0: 1bea
123 nvml.system_event_set_free(self._event_set) 1a
125 def wait(self, timeout_ms: int = 0, buffer_size: int = 1) -> SystemEvents:
126 """
127 Wait for events in the system event set.
129 For Fermi™ or newer fully supported devices.
131 If some events are ready to be delivered at the time of the call,
132 function returns immediately. If there are no events ready to be
133 delivered, function sleeps till event arrives but not longer than
134 specified timeout. If timeout passes, a
135 :class:`cuda.core.system.TimeoutError` is raised. This function in
136 certain conditions can return before specified timeout passes (e.g. when
137 interrupt arrives)
139 Parameters
140 ----------
141 timeout_ms: int
142 The timeout in milliseconds. A value of 0 means to wait indefinitely.
143 buffer_size: int
144 The maximum number of events to retrieve. Must be at least 1.
146 Returns
147 -------
148 :obj:`~_system_events.SystemEvents`
149 A set of events that were received. The number of events returned may
150 be less than the specified buffer size if fewer events were available.
152 Raises
153 ------
154 :class:`cuda.core.system.TimeoutError`
155 If the timeout expires before an event is received.
156 :class:`cuda.core.system.GpuIsLostError`
157 If the GPU has fallen off the bus or is otherwise inaccessible.
158 """
159 return SystemEvents(nvml.system_event_set_wait(self._event_set, timeout_ms, buffer_size)) 1a
162def register_events(events: SystemEventType | str | list[SystemEventType | str]) -> RegisteredSystemEvents:
163 """
164 Starts recording of events on test system.
166 For Linux only.
168 All events that occurred before this call are not recorded. Wait for events
169 using the :meth:`RegisteredSystemEvents.wait` method on the result.
171 Examples
172 --------
173 >>> from cuda.core import system
174 >>> events = system.register_events([SystemEventType.UNBIND])
175 >>> while event := events.wait(timeout_ms=10000):
176 ... print(f"Event {event.event_type} occurred.")
178 Parameters
179 ----------
180 events: SystemEventType, str, or list of SystemEventType or str
181 The event type or list of event types to register for this device.
183 Returns
184 -------
185 :obj:`~_system_events.RegisteredSystemEvents`
186 An object representing the registered events. Call
187 :meth:`~_system_events.RegisteredSystemEvents.wait` on this object to wait for events.
189 Raises
190 ------
191 :class:`cuda.core.system.NotSupportedError`
192 None of the requested event types are registered.
193 """
194 return RegisteredSystemEvents(events) 1bea
197__all__ = [
198 "register_events",
199]