Coverage for cuda/core/_memory/_peer_access_utils.pyx: 38.89%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from __future__ import annotations 

6  

7from collections.abc import Callable, Iterable, Iterator, MutableSet, Set (empty)

8from dataclasses import dataclass (empty)

9from typing import TYPE_CHECKING, Any, TypeVar (empty)

10  

11_S = TypeVar("_S") (empty)

12  

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._resource_handles 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 

21  

22if TYPE_CHECKING: (empty)

23 from cuda.core._device import Device 

24  

25  

26@dataclass(frozen=True) (empty)

27class PeerAccessPlan: (empty)

28 """Normalized peer-access target state and the driver updates it requires.""" 

29  

30 target_ids: tuple[int, ...] 

31 to_add: tuple[int, ...] 

32 to_remove: tuple[int, ...] 

33  

34  

35def normalize_peer_access_targets( (empty)

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.""" 

42  

43 target_ids = {resolve_device_id(device) for device in requested_devices} 5 ctx1badef

44 target_ids.discard(owner_device_id) 4 ctx1adef

45 return tuple(sorted(target_ids)) 4 ctx1adef

46  

47  

48def plan_peer_access_update( (empty)

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.""" 

57  

58 target_ids = normalize_peer_access_targets( 4 ctx1adef

59 owner_device_id, 

60 requested_devices, 

61 resolve_device_id=resolve_device_id, 4 ctx1adef

62 ) 

63 bad = tuple(dev_id for dev_id in target_ids if not can_access_peer(dev_id)) 4 ctx1adef

64 if bad: 4 ctx1adef

65 bad_ids = ", ".join(str(dev_id) for dev_id in bad) 1 ctx1f

66 raise ValueError(f"Device {owner_device_id} cannot access peer(s): {bad_ids}") 1 ctx1f

67  

68 current_ids = set(current_peer_ids) 3 ctx1ade

69 target_id_set = set(target_ids) 3 ctx1ade

70 return PeerAccessPlan( 3 ctx1ade

71 target_ids=target_ids, 

72 to_add=tuple(sorted(target_id_set - current_ids)), 3 ctx1ade

73 to_remove=tuple(sorted(current_ids - target_id_set)), 3 ctx1ade

74 ) 

75  

76  

77def _resolve_peer_device_id(value: Device | int | None) -> int: (empty)

78 """Coerce ``Device | int`` into a device-ordinal int.""" 

79 from cuda.core._device import Device 1 ctx1c

80  

81 return Device(value).device_id 1 ctx1c

82  

83  

84# ---- driver-touching helpers (cdef inline, called from .pyx code) ----------- 

85  

86cdef inline DeviceMemoryResource _check_peer_access_open(object mr): (empty)

87 cdef DeviceMemoryResource mr_typed = <DeviceMemoryResource>mr 3 ctx1gca

88 MP_check_open(mr_typed) 3 ctx1gca

89 return mr_typed 2 ctx1ca

90  

91  

92cdef inline tuple _query_peer_access_ids(DeviceMemoryResource mr): (empty)

93 """Return the current peer device IDs as a sorted tuple of ints. 

94  

95 The full driver loop runs inside a single ``nogil`` block. Because 

96 ``range(total)`` ascends, the result is already sorted. 

97 """ 

98 MP_check_open(mr) 2 ctx1ca

99 cdef int total 

100 cdef int dev_id 

101 cdef int owner_id = mr._dev_id 2 ctx1ca

102 cdef cydriver.CUmemAccess_flags flags 

103 cdef cydriver.CUmemLocation location 

104 cdef cydriver.CUmemoryPool h_pool = as_cu(mr._h_pool) 2 ctx1ca

105 cdef vector[int] peers 

106 cdef size_t i 

107  

108 location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 2 ctx1ca

109  

110 with nogil: 2 ctx1ca

111 HANDLE_RETURN(cydriver.cuDeviceGetCount(&total)) 2 ctx1ca

112 for dev_id in range(total): 2 ctx1ca

113 if dev_id == owner_id: 2 ctx1ca

114 continue 2 ctx1ca

115 location.id = dev_id (empty)

116 HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, h_pool, &location)) 

117 if flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE: 

118 peers.push_back(dev_id) 

119  

120 cdef size_t n = peers.size() 2 ctx1ca

121 return tuple(peers[i] for i in range(n)) 2 ctx1ca

122  

123  

124cdef inline bint _peer_access_includes(DeviceMemoryResource mr, int dev_id): (empty)

125 """Return True if peer access from ``dev_id`` is currently granted.""" 

126 MP_check_open(mr) 

127 cdef cydriver.CUmemAccess_flags flags 

128 cdef cydriver.CUmemLocation location = cumemlocation_from_id( 

129 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, dev_id) 

130 with nogil: 

131 HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(mr._h_pool), &location)) 

132 return flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE 

133  

134  

135def _set_pool_access(mr: object, to_add: tuple[int, ...], to_remove: tuple[int, ...]) -> None: (empty)

136 """Issue one ``cuMemPoolSetAccess`` for the given add/remove deltas. 

137  

138 The thin Python-callable layer that wraps the actual driver call: building 

139 the ``CUmemAccessDesc`` array and invoking ``cuMemPoolSetAccess`` happens 

140 in here. Tests monkeypatch this on the module to spy on real driver work 

141 without intercepting earlier no-op paths. 

142  

143 Preconditions: ``len(to_add) + len(to_remove) > 0`` (the caller is 

144 responsible for skipping empty diffs). 

145 """ 

146 cdef DeviceMemoryResource mr_typed = <DeviceMemoryResource>mr 

147 MP_check_open(mr_typed) 

148 cdef size_t count = len(to_add) + len(to_remove) 

149 cdef cydriver.CUmemAccessDesc* access_desc = NULL 

150 cdef size_t i = 0 

151  

152 access_desc = <cydriver.CUmemAccessDesc*>PyMem_Malloc(count * sizeof(cydriver.CUmemAccessDesc)) 

153 if access_desc == NULL: 

154 raise MemoryError("Failed to allocate memory for access descriptors") 

155  

156 try: 

157 for dev_id in to_add: 

158 access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE 

159 access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 

160 access_desc[i].location.id = dev_id 

161 i += 1 

162 for dev_id in to_remove: 

163 access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_NONE 

164 access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 

165 access_desc[i].location.id = dev_id 

166 i += 1 

167  

168 with nogil: 

169 HANDLE_RETURN(cydriver.cuMemPoolSetAccess(as_cu(mr_typed._h_pool), access_desc, count)) 

170 finally: 

171 if access_desc != NULL: 

172 PyMem_Free(access_desc) 

173  

174  

175def _apply_peer_access_diff(mr: DeviceMemoryResource, to_add: Iterable[int], to_remove: Iterable[int]) -> None: (empty)

176 """Apply a peer-access diff in at most one driver call. 

177  

178 Every write path on :class:`PeerAccessibleBySetProxy` and the 

179 ``peer_accessible_by`` setter routes through this function. Empty diffs 

180 short-circuit here so the driver-level helper :func:`_set_pool_access` is 

181 only invoked when there is actual work for ``cuMemPoolSetAccess`` to do. 

182 """ 

183 MP_check_open(mr) 1 ctx1a

184 add_tuple = tuple(to_add) 1 ctx1a

185 remove_tuple = tuple(to_remove) 1 ctx1a

186 if not add_tuple and not remove_tuple: 1 ctx1a

187 return 1 ctx1a

188 _set_pool_access(mr, add_tuple, remove_tuple) 

189  

190  

191cpdef void replace_peer_accessible_by(DeviceMemoryResource mr, object devices): (empty)

192 """Replace the full peer-access set in a single batched driver call. 

193  

194 Backs the ``mr.peer_accessible_by = [...]`` setter. Uses the same planner 

195 as the proxy's bulk ops; the only difference is that adds and removes are 

196 derived from the symmetric difference between current driver state and the 

197 requested target set. 

198 """ 

199 MP_check_open(mr) 1 ctx1a

200 from cuda.core._device import Device 1 ctx1a

201  

202 this_dev = Device(mr._dev_id) 1 ctx1a

203 plan = plan_peer_access_update( 1 ctx1a

204 owner_device_id=mr._dev_id, 1 ctx1a

205 current_peer_ids=_query_peer_access_ids(mr), 1 ctx1a

206 requested_devices=devices, 

207 resolve_device_id=_resolve_peer_device_id, 1 ctx1a

208 can_access_peer=this_dev.can_access_peer, 1 ctx1a

209 ) 

210 _apply_peer_access_diff(mr, plan.to_add, plan.to_remove) 1 ctx1a

211  

212  

213# ---- Python MutableSet proxy ------------------------------------------------ 

214  

215class PeerAccessibleBySetProxy(MutableSet["Device"]): (empty)

216 """Live driver-backed view of the peer devices granted access to a memory pool. 

217  

218 Reads (``__contains__``, ``__iter__``, ``len(...)``) call ``cuMemPoolGetAccess``; 

219 writes (``add``, ``discard``, and bulk ops) call ``cuMemPoolSetAccess``. There 

220 is no in-memory mirror, so the view always reflects the current driver state 

221 and stays consistent across multiple wrappers around the same pool. 

222  

223 Iteration yields :class:`~cuda.core.Device` objects. ``add``, ``discard``, and 

224 ``__contains__`` accept either a :class:`~cuda.core.Device` or a device-ordinal 

225 ``int``; the owner device is silently ignored when supplied. 

226  

227 All bulk operations (``update``, ``|=``, ``&=``, ``-=``, ``^=``, ``clear``) 

228 issue exactly one ``cuMemPoolSetAccess`` call. This matters: peer-access 

229 transitions can take seconds per pool because every existing memory mapping 

230 is updated, so coalescing into a single driver call lets the toolkit handle 

231 the mappings in parallel. 

232 """ 

233  

234 __slots__ = ("_mr",) (empty)

235  

236 def __init__(self, mr: DeviceMemoryResource) -> None: (empty)

237 self._mr = mr 3 ctx1gca

238  

239 @classmethod (empty)

240 def _from_iterable(cls, it: Iterable[_S]) -> set[_S]: 

241 # Binary set operators (&, |, -, ^) collect their result through 

242 # _from_iterable. Returning a plain set lets the user reason about 

243 # the result independently of any pool's driver state. 

244 return set(it) 

245  

246 # --- abstract MutableSet methods --- 

247  

248 def __contains__(self, value: object) -> bool: (empty)

249 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr) 1 ctx1c

250 try: 1 ctx1c

251 dev_id = _resolve_peer_device_id(value) 1 ctx1c

252 except (TypeError, ValueError): 

253 return False 

254 if dev_id == mr._dev_id: 1 ctx1c

255 return False 1 ctx1c

256 return _peer_access_includes(mr, dev_id) 

257  

258 def __iter__(self) -> Iterator[Device]: (empty)

259 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr) 2 ctx1ca

260 from cuda.core._device import Device 2 ctx1ca

261  

262 return iter(Device(dev_id) for dev_id in _query_peer_access_ids(mr)) 2 ctx1ca

263  

264 def __len__(self) -> int: (empty)

265 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr) 1 ctx1g

266 return len(_query_peer_access_ids(mr)) 

267  

268 def add(self, value: Device | int) -> None: (empty)

269 """Grant peer access from ``value`` to allocations in this pool.""" 

270 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr) 

271 dev_id = _resolve_peer_device_id(value) 

272 if dev_id == mr._dev_id: 

273 return 

274 if _peer_access_includes(mr, dev_id): 

275 return 

276 from cuda.core._device import Device 

277 if not Device(mr._dev_id).can_access_peer(dev_id): 

278 raise ValueError(f"Device {mr._dev_id} cannot access peer: {dev_id}") 

279 _apply_peer_access_diff(mr, (dev_id,), ()) 

280  

281 def discard(self, value: Device | int) -> None: (empty)

282 """Revoke peer access from ``value`` to allocations in this pool.""" 

283 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr) 

284 try: 

285 dev_id = _resolve_peer_device_id(value) 

286 except (TypeError, ValueError): 

287 return 

288 if dev_id == mr._dev_id: 

289 return 

290 if not _peer_access_includes(mr, dev_id): 

291 return 

292 _apply_peer_access_diff(mr, (), (dev_id,)) 

293  

294 # --- bulk overrides: one driver call per op --- 

295  

296 def clear(self) -> None: (empty)

297 """Revoke all peer access in a single driver call.""" 

298 _check_peer_access_open(self._mr) 

299 self._apply((), _query_peer_access_ids(self._mr)) 

300  

301 def update(self, *others: Iterable[Device | int]) -> None: (empty)

302 """Grant peer access to every device in ``others`` in one driver call.""" 

303 _check_peer_access_open(self._mr) 

304 to_add = [] 

305 for other in others: 

306 to_add.extend(other) 

307 if to_add: 

308 self._apply(to_add, ()) 

309  

310 def difference_update(self, *others: Iterable[Device | int]) -> None: (empty)

311 """Revoke peer access for every device in ``others`` in one driver call.""" 

312 _check_peer_access_open(self._mr) 

313 revoke_ids = set() 

314 for other in others: 

315 for value in other: 

316 try: 

317 revoke_ids.add(_resolve_peer_device_id(value)) 

318 except (TypeError, ValueError): 

319 continue 

320 current = set(_query_peer_access_ids(self._mr)) 

321 to_remove = revoke_ids & current 

322 if to_remove: 

323 self._apply((), to_remove) 

324  

325 def intersection_update(self, *others: Iterable[Device | int]) -> None: (empty)

326 """Restrict peer access to the intersection in a single driver call.""" 

327 _check_peer_access_open(self._mr) 

328 keep_ids = None 

329 for other in others: 

330 ids = set() 

331 for value in other: 

332 try: 

333 ids.add(_resolve_peer_device_id(value)) 

334 except (TypeError, ValueError): 

335 continue 

336 keep_ids = ids if keep_ids is None else keep_ids & ids 

337 if keep_ids is None: 

338 return # ``set.intersection_update()`` with no args is a no-op 

339 current = set(_query_peer_access_ids(self._mr)) 

340 to_remove = current - keep_ids 

341 if to_remove: 

342 self._apply((), to_remove) 

343  

344 def symmetric_difference_update(self, other: Iterable[Device | int]) -> None: (empty)

345 """Toggle peer access for every device in ``other`` in one driver call.""" 

346 _check_peer_access_open(self._mr) 

347 toggle_ids = set() 

348 for value in other: 

349 try: 

350 toggle_ids.add(_resolve_peer_device_id(value)) 

351 except (TypeError, ValueError): 

352 continue 

353 current = set(_query_peer_access_ids(self._mr)) 

354 to_add = toggle_ids - current 

355 to_remove = toggle_ids & current 

356 if to_add or to_remove: 

357 self._apply(to_add, to_remove) 

358  

359 def __ior__(self, other: Set[Any]) -> PeerAccessibleBySetProxy: # type: ignore[misc] (empty)

360 self.update(other) 

361 return self 

362  

363 def __iand__(self, other: Set[Any]) -> PeerAccessibleBySetProxy: (empty)

364 self.intersection_update(other) 

365 return self 

366  

367 def __isub__(self, other: Set[Any]) -> PeerAccessibleBySetProxy: (empty)

368 if other is self: 

369 self.clear() 

370 else: 

371 self.difference_update(other) 

372 return self 

373  

374 def __ixor__(self, other: Set[Any]) -> PeerAccessibleBySetProxy: # type: ignore[misc] (empty)

375 self.symmetric_difference_update(other) 

376 return self 

377  

378 def __repr__(self) -> str: (empty)

379 return f"PeerAccessibleBySetProxy({set(self)!r})" 

380  

381 # --- internal: route every write through one batched driver call --- 

382  

383 def _apply(self, additions, removals) -> None: (empty)

384 """Compute the diff and issue a single ``cuMemPoolSetAccess``. 

385  

386 ``additions`` and ``removals`` are user-supplied (``Device | int``); 

387 only the owner device is filtered out. Adds are validated through 

388 :meth:`Device.can_access_peer` via :func:`plan_peer_access_update`; 

389 removals bypass that check (revoking is always permitted). 

390 """ 

391 from cuda.core._device import Device 

392  

393 cdef DeviceMemoryResource mr = _check_peer_access_open(self._mr) 

394 owner_id = mr._dev_id 

395 owner = Device(owner_id) 

396 current = _query_peer_access_ids(mr) 

397  

398 # Plan additions through the existing helper (validates can_access_peer). 

399 plan = plan_peer_access_update( 

400 owner_device_id=owner_id, 

401 current_peer_ids=current, 

402 # union of (current set + requested adds) so the planner emits 

403 # exactly the to_add deltas for these additions, no removals. 

404 requested_devices=[*current, *additions], 

405 resolve_device_id=_resolve_peer_device_id, 

406 can_access_peer=owner.can_access_peer, 

407 ) 

408 to_add = plan.to_add 

409  

410 # Removals: resolve, drop owner and unknowns, intersect with current. 

411 current_set = set(current) 

412 revoke_ids = set() 

413 for value in removals: 

414 try: 

415 dev_id = _resolve_peer_device_id(value) 

416 except (TypeError, ValueError): 

417 continue 

418 if dev_id == owner_id: 

419 continue 

420 if dev_id in current_set: 

421 revoke_ids.add(dev_id) 

422 to_remove = tuple(sorted(revoke_ids)) 

423  

424 if not to_add and not to_remove: 

425 return 

426 _apply_peer_access_diff(mr, to_add, to_remove)