Coverage for cuda/core/system/_system_events.pyx: 84.44%
45 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 02:41 +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 1efghc
31 bus = (gpu_id >> 8) & 0xFF 1efghc
32 device = gpu_id & 0xFF 1efghc
33 return f"{domain:08X}:{bus:02X}:{device:02X}.0" 1efghc
36cdef class SystemEvent:
37 """
38 Data about a collection of system events.
39 """
40 cdef object _event_data
42 def __init__(self, event_data: nvml.SystemEventData_v1):
43 assert len(event_data) == 1 1cd
44 self._event_data = event_data 1cd
46 @property
47 def event_type(self) -> SystemEventType:
48 """
49 The :obj:`~SystemEventType` that was triggered.
50 """
51 return _SYSTEM_EVENT_TYPE_MAPPING[self._event_data.event_type] 1d
53 @property
54 def gpu_id(self) -> int:
55 """
56 The GPU ID in PCI ID format.
57 """
58 return self._event_data.gpu_id 1cd
60 @property
61 def device(self) -> _device.Device:
62 """
63 The :obj:`~_device.Device` associated with this event.
64 """
65 return _device.Device(pci_bus_id=_pci_bus_id_from_gpu_id(self.gpu_id)) 1c
68cdef class SystemEvents:
69 """
70 Data about a collection of system events.
71 """
72 cdef object _event_data
74 def __init__(self, event_data: nvml.SystemEventData_v1):
75 self._event_data = event_data 1d
77 def __len__(self) -> int:
78 return len(self._event_data) 1d
80 def __getitem__(self, idx: int) -> SystemEvent:
81 """
82 Get the :obj:`~_system_events.SystemEvent` at the specified index.
83 """
84 return SystemEvent(self._event_data[idx]) 1d
87cdef class RegisteredSystemEvents:
88 """
89 Represents a set of events that can be waited on for a specific device.
90 """
91 cdef intptr_t _event_set
93 def __init__(self, events: SystemEventType | str | list[SystemEventType | str]):
94 cdef unsigned long long event_bitmask
95 if isinstance(events, (str, SystemEventType)): 1a
96 events = [events]
98 if isinstance(events, list): 1a
99 event_bitmask = 0 1a
100 for ev in events: 1a
101 try: 1a
102 ev_enum = _SYSTEM_EVENT_TYPE_INV_MAPPING[ev] 1a
103 except KeyError:
104 raise ValueError(
105 f"Invalid event type: {ev}. "
106 f"Must be one of {list(SystemEventType.__members__.values())}"
107 ) from None
108 event_bitmask |= <unsigned long long>int(ev_enum) 1a
109 else:
110 raise TypeError("events must be an SystemEventType, str, or list of SystemEventType or str")
112 initialize() 1a
114 self._event_set = 0 1a
115 self._event_set = nvml.system_event_set_create() 1a
116 # If this raises, the event needs to be freed and this is handled by
117 # this class's __dealloc__ method.
118 nvml.system_register_events(event_bitmask, self._event_set) 1a
120 def __dealloc__(self) -> None:
121 if self._event_set != 0: 1a
122 nvml.system_event_set_free(self._event_set) 1a
124 def wait(self, timeout_ms: int = 0, buffer_size: int = 1) -> SystemEvents:
125 """
126 Wait for events in the system event set.
128 For Fermi™ or newer fully supported devices.
130 If some events are ready to be delivered at the time of the call,
131 function returns immediately. If there are no events ready to be
132 delivered, function sleeps till event arrives but not longer than
133 specified timeout. If timeout passes, a
134 :class:`cuda.core.system.TimeoutError` is raised. This function in
135 certain conditions can return before specified timeout passes (e.g. when
136 interrupt arrives)
138 Parameters
139 ----------
140 timeout_ms: int
141 The timeout in milliseconds. A value of 0 means to wait indefinitely.
142 buffer_size: int
143 The maximum number of events to retrieve. Must be at least 1.
145 Returns
146 -------
147 :obj:`~_system_events.SystemEvents`
148 A set of events that were received. The number of events returned may
149 be less than the specified buffer size if fewer events were available.
151 Raises
152 ------
153 :class:`cuda.core.system.TimeoutError`
154 If the timeout expires before an event is received.
155 :class:`cuda.core.system.GpuIsLostError`
156 If the GPU has fallen off the bus or is otherwise inaccessible.
157 """
158 return SystemEvents(nvml.system_event_set_wait(self._event_set, timeout_ms, buffer_size)) 1a
161def register_events(events: SystemEventType | str | list[SystemEventType | str]) -> RegisteredSystemEvents:
162 """
163 Starts recording of events on test system.
165 For Linux only.
167 All events that occurred before this call are not recorded. Wait for events
168 using the :meth:`RegisteredSystemEvents.wait` method on the result.
170 Examples
171 --------
172 >>> from cuda.core import system
173 >>> events = system.register_events([SystemEventType.UNBIND])
174 >>> while event := events.wait(timeout_ms=10000):
175 ... print(f"Event {event.event_type} occurred.")
177 Parameters
178 ----------
179 events: SystemEventType, str, or list of SystemEventType or str
180 The event type or list of event types to register for this device.
182 Returns
183 -------
184 :obj:`~_system_events.RegisteredSystemEvents`
185 An object representing the registered events. Call
186 :meth:`~_system_events.RegisteredSystemEvents.wait` on this object to wait for events.
188 Raises
189 ------
190 :class:`cuda.core.system.NotSupportedError`
191 None of the requested event types are registered.
192 """
193 return RegisteredSystemEvents(events) 1a
196__all__ = [
197 "register_events",
198]