Coverage for cuda/core/_memory/_peer_access_utils.pyx: 22.84%
232 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 01:23 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
5from __future__ import annotations
7from collections.abc import Callable, Iterable, Iterator, MutableSet, Set
8from dataclasses import dataclass
9from typing import TYPE_CHECKING, Any, TypeVar
11_S = TypeVar("_S")
13from cuda.bindings cimport cydriver
14from cuda.core._memory._device_memory_resource cimport DeviceMemoryResource
15from cuda.core._memory._location cimport cumemlocation_from_id
16from cuda.core._memory._memory_pool cimport MP_check_open
17from cuda.core._rt cimport as_cu
18from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
19from cpython.mem cimport PyMem_Malloc, PyMem_Free
20from libcpp.vector cimport vector
22if TYPE_CHECKING:
23 from cuda.core._device import Device
26@dataclass(frozen=True)
27class PeerAccessPlan:
28 """Normalized peer-access target state and the driver updates it requires."""
30 target_ids: tuple[int, ...]
31 to_add: tuple[int, ...]
32 to_remove: tuple[int, ...]
35def normalize_peer_access_targets(
36 owner_device_id: int,
37 requested_devices: Iterable[object],
38 *,
39 resolve_device_id: Callable[[object], int],
40) -> tuple[int, ...]:
41 """Return sorted, unique peer device IDs, excluding the owner device."""
43 target_ids = {resolve_device_id(device) for device in requested_devices} 1acde
44 target_ids.discard(owner_device_id) 1acde
45 return tuple(sorted(target_ids)) 1acde
48def plan_peer_access_update(
49 owner_device_id: int,
50 current_peer_ids: Iterable[int],
51 requested_devices: Iterable[object],
52 *,
53 resolve_device_id: Callable[[object], int],
54 can_access_peer: Callable[[int], bool],
55) -> PeerAccessPlan:
56 """Compute the peer-access target state and add/remove deltas."""
58 target_ids = normalize_peer_access_targets( 1acde
59 owner_device_id,
60 requested_devices,
61 resolve_device_id=resolve_device_id, 1acde
62 )
63 bad = tuple(dev_id for dev_id in target_ids if not can_access_peer(dev_id)) 1acde
64 if bad: 1acde
65 bad_ids = ", ".join(str(dev_id) for dev_id in bad) 1e
66 raise ValueError(f"Device {owner_device_id} cannot access peer(s): {bad_ids}") 1e
68 current_ids = set(current_peer_ids) 1acd
69 target_id_set = set(target_ids) 1acd
70 return PeerAccessPlan( 1acd
71 target_ids=target_ids,
72 to_add=tuple(sorted(target_id_set - current_ids)), 1acd
73 to_remove=tuple(sorted(current_ids - target_id_set)), 1acd
74 )
77def _resolve_peer_device_id(value: Device | int | None) -> int:
78 """Coerce ``Device | int`` into a device-ordinal int."""
79 from cuda.core._device import Device 1b
81 return Device(value).device_id 1b
84# ---- driver-touching helpers (cdef inline, called from .pyx code) -----------
86cdef inline DeviceMemoryResource _check_peer_access_open(DeviceMemoryResource mr):
87 MP_check_open(mr) 1fba
88 return mr 1ba
91cdef inline tuple _query_peer_access_ids(DeviceMemoryResource mr):
92 """Return the current peer device IDs as a sorted tuple of ints.
94 The full driver loop runs inside a single ``nogil`` block. Because
95 ``range(total)`` ascends, the result is already sorted.
96 """
97 MP_check_open(mr) 1ba
98 cdef int total
99 cdef int dev_id
100 cdef int owner_id = mr._dev_id 1ba
101 cdef cydriver.CUmemAccess_flags flags
102 cdef cydriver.CUmemLocation location
103 cdef cydriver.CUmemoryPool h_pool = as_cu(mr._h_pool) 1ba
104 cdef vector[int] peers
105 cdef size_t i
107 location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 1ba
109 with nogil: 1ba
110 HANDLE_RETURN(cydriver.cuDeviceGetCount(&total)) 1ba
111 for dev_id in range(total): 1ba
112 if dev_id == owner_id: 1ba
113 continue 1ba
114 location.id = dev_id
115 HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, h_pool, &location))
116 if flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE:
117 peers.push_back(dev_id)
119 cdef size_t n = peers.size() 1ba
120 return tuple(peers[i] for i in range(n)) 1ba
123cdef inline bint _peer_access_includes(DeviceMemoryResource mr, int dev_id):
124 """Return True if peer access from ``dev_id`` is currently granted."""
125 MP_check_open(mr)
126 cdef cydriver.CUmemAccess_flags flags
127 cdef cydriver.CUmemLocation location = cumemlocation_from_id(
128 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, dev_id)
129 with nogil:
130 HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(mr._h_pool), &location))
131 return flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
134def _set_pool_access(DeviceMemoryResource mr, to_add: tuple[int, ...], to_remove: tuple[int, ...]) -> None:
135 """Issue one ``cuMemPoolSetAccess`` for the given add/remove deltas.
137 The thin Python-callable layer that wraps the actual driver call: building
138 the ``CUmemAccessDesc`` array and invoking ``cuMemPoolSetAccess`` happens
139 in here. Tests monkeypatch this on the module to spy on real driver work
140 without intercepting earlier no-op paths.
142 Preconditions: ``len(to_add) + len(to_remove) > 0`` (the caller is
143 responsible for skipping empty diffs).
144 """
145 MP_check_open(mr)
146 cdef size_t count = len(to_add) + len(to_remove)
147 cdef cydriver.CUmemAccessDesc* access_desc = NULL
148 cdef size_t i = 0
150 access_desc = <cydriver.CUmemAccessDesc*>PyMem_Malloc(count * sizeof(cydriver.CUmemAccessDesc))
151 if access_desc == NULL:
152 raise MemoryError("Failed to allocate memory for access descriptors")
154 try:
155 for dev_id in to_add:
156 access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
157 access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
158 access_desc[i].location.id = dev_id
159 i += 1
160 for dev_id in to_remove:
161 access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_NONE
162 access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
163 access_desc[i].location.id = dev_id
164 i += 1
166 with nogil:
167 HANDLE_RETURN(cydriver.cuMemPoolSetAccess(as_cu(mr._h_pool), access_desc, count))
168 finally:
169 if access_desc != NULL:
170 PyMem_Free(access_desc)
173def _apply_peer_access_diff(DeviceMemoryResource mr not None: DeviceMemoryResource, to_add: Iterable[int], to_remove: Iterable[int]) -> None:
174 """Apply a peer-access diff in at most one driver call.
176 Every write path on :class:`PeerAccessibleBySetProxy` and the
177 ``peer_accessible_by`` setter routes through this function. Empty diffs
178 short-circuit here so the driver-level helper :func:`_set_pool_access` is
179 only invoked when there is actual work for ``cuMemPoolSetAccess`` to do.
180 """
181 MP_check_open(mr) 1a
182 add_tuple = tuple(to_add) 1a
183 remove_tuple = tuple(to_remove) 1a
184 if not add_tuple and not remove_tuple: 1a
185 return 1a
186 _set_pool_access(mr, add_tuple, remove_tuple)
189cpdef void replace_peer_accessible_by(DeviceMemoryResource mr, object devices):
190 """Replace the full peer-access set in a single batched driver call.
192 Backs the ``mr.peer_accessible_by = [...]`` setter. Uses the same planner
193 as the proxy's bulk ops; the only difference is that adds and removes are
194 derived from the symmetric difference between current driver state and the
195 requested target set.
196 """
197 MP_check_open(mr) 1a
198 from cuda.core._device import Device 1a
200 this_dev = Device(mr._dev_id) 1a
201 plan = plan_peer_access_update( 1a
202 owner_device_id=mr._dev_id, 1a
203 current_peer_ids=_query_peer_access_ids(mr), 1a
204 requested_devices=devices,
205 resolve_device_id=_resolve_peer_device_id, 1a
206 can_access_peer=this_dev.can_access_peer, 1a
207 )
208 _apply_peer_access_diff(mr, plan.to_add, plan.to_remove) 1a
211# ---- Python MutableSet proxy ------------------------------------------------
213class PeerAccessibleBySetProxy(MutableSet["Device"]):
214 """Live driver-backed view of the peer devices granted access to a memory pool.
216 Reads (``__contains__``, ``__iter__``, ``len(...)``) call ``cuMemPoolGetAccess``;
217 writes (``add``, ``discard``, and bulk ops) call ``cuMemPoolSetAccess``. There
218 is no in-memory mirror, so the view always reflects the current driver state
219 and stays consistent across multiple wrappers around the same pool.
221 Iteration yields :class:`~cuda.core.Device` objects. ``add``, ``discard``, and
222 ``__contains__`` accept either a :class:`~cuda.core.Device` or a device-ordinal
223 ``int``; the owner device is silently ignored when supplied.
225 All bulk operations (``update``, ``|=``, ``&=``, ``-=``, ``^=``, ``clear``)
226 issue exactly one ``cuMemPoolSetAccess`` call. This matters: peer-access
227 transitions can take seconds per pool because every existing memory mapping
228 is updated, so coalescing into a single driver call lets the toolkit handle
229 the mappings in parallel.
230 """
232 __slots__ = ("_mr",)
234 def __init__(self, DeviceMemoryResource mr not None: DeviceMemoryResource) -> None:
235 self._mr = mr 1fba
237 @classmethod
238 def _from_iterable(cls, it: Iterable[_S]) -> set[_S]:
239 # Binary set operators (&, |, -, ^) collect their result through
240 # _from_iterable. Returning a plain set lets the user reason about
241 # the result independently of any pool's driver state.
242 return set(it)
244 # --- abstract MutableSet methods ---
246 def __contains__(self, value: object) -> bool:
247 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr) 1b
248 try: 1b
249 dev_id = _resolve_peer_device_id(value) 1b
250 except (TypeError, ValueError):
251 return False
252 if dev_id == mr._dev_id: 1b
253 return False 1b
254 return _peer_access_includes(mr, dev_id)
256 def __iter__(self) -> Iterator[Device]:
257 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr) 1ba
258 from cuda.core._device import Device 1ba
260 return iter(Device(dev_id) for dev_id in _query_peer_access_ids(mr)) 1ba
262 def __len__(self) -> int:
263 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr) 1f
264 return len(_query_peer_access_ids(mr))
266 def add(self, value: Device | int) -> None:
267 """Grant peer access from ``value`` to allocations in this pool."""
268 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr)
269 dev_id = _resolve_peer_device_id(value)
270 if dev_id == mr._dev_id:
271 return
272 if _peer_access_includes(mr, dev_id):
273 return
274 from cuda.core._device import Device
275 if not Device(mr._dev_id).can_access_peer(dev_id):
276 raise ValueError(f"Device {mr._dev_id} cannot access peer: {dev_id}")
277 _apply_peer_access_diff(mr, (dev_id,), ())
279 def discard(self, value: Device | int) -> None:
280 """Revoke peer access from ``value`` to allocations in this pool."""
281 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr)
282 try:
283 dev_id = _resolve_peer_device_id(value)
284 except (TypeError, ValueError):
285 return
286 if dev_id == mr._dev_id:
287 return
288 if not _peer_access_includes(mr, dev_id):
289 return
290 _apply_peer_access_diff(mr, (), (dev_id,))
292 # --- bulk overrides: one driver call per op ---
294 def clear(self) -> None:
295 """Revoke all peer access in a single driver call."""
296 _check_peer_access_open(self._mr)
297 self._apply((), _query_peer_access_ids(self._mr))
299 def update(self, *others: Iterable[Device | int]) -> None:
300 """Grant peer access to every device in ``others`` in one driver call."""
301 _check_peer_access_open(self._mr)
302 to_add = []
303 for other in others:
304 to_add.extend(other)
305 if to_add:
306 self._apply(to_add, ())
308 def difference_update(self, *others: Iterable[Device | int]) -> None:
309 """Revoke peer access for every device in ``others`` in one driver call."""
310 _check_peer_access_open(self._mr)
311 revoke_ids = set()
312 for other in others:
313 for value in other:
314 try:
315 revoke_ids.add(_resolve_peer_device_id(value))
316 except (TypeError, ValueError):
317 continue
318 current = set(_query_peer_access_ids(self._mr))
319 to_remove = revoke_ids & current
320 if to_remove:
321 self._apply((), to_remove)
323 def intersection_update(self, *others: Iterable[Device | int]) -> None:
324 """Restrict peer access to the intersection in a single driver call."""
325 _check_peer_access_open(self._mr)
326 keep_ids = None
327 for other in others:
328 ids = set()
329 for value in other:
330 try:
331 ids.add(_resolve_peer_device_id(value))
332 except (TypeError, ValueError):
333 continue
334 keep_ids = ids if keep_ids is None else keep_ids & ids
335 if keep_ids is None:
336 return # ``set.intersection_update()`` with no args is a no-op
337 current = set(_query_peer_access_ids(self._mr))
338 to_remove = current - keep_ids
339 if to_remove:
340 self._apply((), to_remove)
342 def symmetric_difference_update(self, other: Iterable[Device | int]) -> None:
343 """Toggle peer access for every device in ``other`` in one driver call."""
344 _check_peer_access_open(self._mr)
345 toggle_ids = set()
346 for value in other:
347 try:
348 toggle_ids.add(_resolve_peer_device_id(value))
349 except (TypeError, ValueError):
350 continue
351 current = set(_query_peer_access_ids(self._mr))
352 to_add = toggle_ids - current
353 to_remove = toggle_ids & current
354 if to_add or to_remove:
355 self._apply(to_add, to_remove)
357 def __ior__(self, other: Set[Any]) -> PeerAccessibleBySetProxy: # type: ignore[misc]
358 self.update(other)
359 return self
361 def __iand__(self, other: Set[Any]) -> PeerAccessibleBySetProxy:
362 self.intersection_update(other)
363 return self
365 def __isub__(self, other: Set[Any]) -> PeerAccessibleBySetProxy:
366 if other is self:
367 self.clear()
368 else:
369 self.difference_update(other)
370 return self
372 def __ixor__(self, other: Set[Any]) -> PeerAccessibleBySetProxy: # type: ignore[misc]
373 self.symmetric_difference_update(other)
374 return self
376 def __repr__(self) -> str:
377 return f"PeerAccessibleBySetProxy({set(self)!r})"
379 # --- internal: route every write through one batched driver call ---
381 def _apply(self, additions, removals) -> None:
382 """Compute the diff and issue a single ``cuMemPoolSetAccess``.
384 ``additions`` and ``removals`` are user-supplied (``Device | int``);
385 only the owner device is filtered out. Adds are validated through
386 :meth:`Device.can_access_peer` via :func:`plan_peer_access_update`;
387 removals bypass that check (revoking is always permitted).
388 """
389 from cuda.core._device import Device
391 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr)
392 owner_id = mr._dev_id
393 owner = Device(owner_id)
394 current = _query_peer_access_ids(mr)
396 # Plan additions through the existing helper (validates can_access_peer).
397 plan = plan_peer_access_update(
398 owner_device_id=owner_id,
399 current_peer_ids=current,
400 # union of (current set + requested adds) so the planner emits
401 # exactly the to_add deltas for these additions, no removals.
402 requested_devices=[*current, *additions],
403 resolve_device_id=_resolve_peer_device_id,
404 can_access_peer=owner.can_access_peer,
405 )
406 to_add = plan.to_add
408 # Removals: resolve, drop owner and unknowns, intersect with current.
409 current_set = set(current)
410 revoke_ids = set()
411 for value in removals:
412 try:
413 dev_id = _resolve_peer_device_id(value)
414 except (TypeError, ValueError):
415 continue
416 if dev_id == owner_id:
417 continue
418 if dev_id in current_set:
419 revoke_ids.add(dev_id)
420 to_remove = tuple(sorted(revoke_ids))
422 if not to_add and not to_remove:
423 return
424 _apply_peer_access_diff(mr, to_add, to_remove)