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

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 (empty)

10  

11from cuda.bindings cimport cydriver 

12from cuda.core._memory._device_memory_resource cimport DeviceMemoryResource 

13from cuda.core._memory._location cimport cumemlocation_from_id 

14from cuda.core._memory._memory_pool cimport MP_check_open 

15from cuda.core._resource_handles cimport as_cu 

16from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

17from cpython.mem cimport PyMem_Malloc, PyMem_Free 

18from libcpp.vector cimport vector 

19  

20if TYPE_CHECKING: (empty)

21 from cuda.core._device import Device (empty)

22  

23  

24@dataclass(frozen=True) (empty)

25class PeerAccessPlan: (empty)

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

27  

28 target_ids: tuple[int, ...] 

29 to_add: tuple[int, ...] 

30 to_remove: tuple[int, ...] 

31  

32  

33def normalize_peer_access_targets( (empty)

34 owner_device_id: int, 

35 requested_devices: Iterable[object], 

36 *, 

37 resolve_device_id: Callable[[object], int], 

38) -> tuple[int, ...]: 

39 """Return sorted, unique peer device IDs, excluding the owner device.""" 

40  

41 target_ids = {resolve_device_id(device) for device in requested_devices} 4 ctx1adef

42 target_ids.discard(owner_device_id) 4 ctx1adef

43 return tuple(sorted(target_ids)) 5 ctx1badef

44  

45  

46def plan_peer_access_update( (empty)

47 owner_device_id: int, 

48 current_peer_ids: Iterable[int], 

49 requested_devices: Iterable[object], 

50 *, 

51 resolve_device_id: Callable[[object], int], 

52 can_access_peer: Callable[[int], bool], 

53) -> PeerAccessPlan: 

54 """Compute the peer-access target state and add/remove deltas.""" 

55  

56 target_ids = normalize_peer_access_targets( 4 ctx1adef

57 owner_device_id, 

58 requested_devices, 

59 resolve_device_id=resolve_device_id, 4 ctx1adef

60 ) 

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

62 if bad: 4 ctx1adef

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

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

65  

66 current_ids = set(current_peer_ids) 3 ctx1ade

67 target_id_set = set(target_ids) 3 ctx1ade

68 return PeerAccessPlan( 3 ctx1ade

69 target_ids=target_ids, 

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

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

72 ) 

73  

74  

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

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

77 from cuda.core._device import Device 1 ctx1c

78  

79 return Device(value).device_id 1 ctx1c

80  

81  

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

83  

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

85 cdef DeviceMemoryResource mr_typed = <DeviceMemoryResource>mr 3 ctx1gca

86 MP_check_open(mr_typed) 3 ctx1gca

87 return mr_typed 2 ctx1ca

88  

89  

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

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

92  

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

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

95 """ 

96 MP_check_open(mr) 2 ctx1ca

97 cdef int total 

98 cdef int dev_id 

99 cdef int owner_id = mr._dev_id 2 ctx1ca

100 cdef cydriver.CUmemAccess_flags flags 

101 cdef cydriver.CUmemLocation location 

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

103 cdef vector[int] peers 

104 cdef size_t i 

105  

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

107  

108 with nogil: 2 ctx1ca

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

110 for dev_id in range(total): 2 ctx1ca

111 if dev_id == owner_id: 2 ctx1ca

112 continue 2 ctx1ca

113 location.id = dev_id 

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

115 if flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE: (empty)

116 peers.push_back(dev_id) 

117  

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

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

120  

121  

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

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

124 MP_check_open(mr) 

125 cdef cydriver.CUmemAccess_flags flags 

126 cdef cydriver.CUmemLocation location = cumemlocation_from_id( 

127 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, dev_id) 

128 with nogil: 

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

130 return flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE 

131  

132  

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

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

135  

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

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

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

139 without intercepting earlier no-op paths. 

140  

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

142 responsible for skipping empty diffs). 

143 """ 

144 cdef DeviceMemoryResource mr_typed = <DeviceMemoryResource>mr 

145 MP_check_open(mr_typed) 

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

147 cdef cydriver.CUmemAccessDesc* access_desc = NULL 

148 cdef size_t i = 0 

149  

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

153  

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 

165  

166 with nogil: 

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

168 finally: 

169 if access_desc != NULL: 

170 PyMem_Free(access_desc) 

171  

172  

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

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

175  

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) 1 ctx1a

182 add_tuple = tuple(to_add) 1 ctx1a

183 remove_tuple = tuple(to_remove) 1 ctx1a

184 if not add_tuple and not remove_tuple: 1 ctx1a

185 return 1 ctx1a

186 _set_pool_access(mr, add_tuple, remove_tuple) 

187  

188  

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

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

191  

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) 1 ctx1a

198 from cuda.core._device import Device 1 ctx1a

199  

200 this_dev = Device(mr._dev_id) 1 ctx1a

201 plan = plan_peer_access_update( 1 ctx1a

202 owner_device_id=mr._dev_id, 1 ctx1a

203 current_peer_ids=_query_peer_access_ids(mr), 1 ctx1a

204 requested_devices=devices, 

205 resolve_device_id=_resolve_peer_device_id, 1 ctx1a

206 can_access_peer=this_dev.can_access_peer, 1 ctx1a

207 ) 

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

209  

210  

211# ---- Python MutableSet proxy ------------------------------------------------ 

212  

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

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

215  

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. 

220  

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. 

224  

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

231  

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

233  

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

235 self._mr = mr 3 ctx1gca

236  

237 @classmethod (empty)

238 def _from_iterable(cls, it: Iterable[Device]) -> set[Device]: # type: ignore[override] 

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) 

243  

244 # --- abstract MutableSet methods --- 

245  

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

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

248 try: 1 ctx1c

249 dev_id = _resolve_peer_device_id(value) 1 ctx1c

250 except (TypeError, ValueError): 

251 return False 

252 if dev_id == mr._dev_id: 1 ctx1c

253 return False 1 ctx1c

254 return _peer_access_includes(mr, dev_id) 

255  

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

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

258 from cuda.core._device import Device 2 ctx1ca

259  

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

261  

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

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

264 return len(_query_peer_access_ids(mr)) 

265  

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

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,), ()) 

278  

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

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,)) 

291  

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

293  

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

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)) 

298  

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

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, ()) 

307  

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

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) 

322  

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

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) 

341  

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

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) 

356  

357 def __ior__(self, other: Set[Any]) -> PeerAccessibleBySetProxy: (empty)

358 self.update(other) 

359 return self 

360  

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

362 self.intersection_update(other) 

363 return self 

364  

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

366 if other is self: 

367 self.clear() 

368 else: 

369 self.difference_update(other) 

370 return self 

371  

372 def __ixor__(self, other: Set[Any]) -> PeerAccessibleBySetProxy: (empty)

373 self.symmetric_difference_update(other) 

374 return self 

375  

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

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

378  

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

380  

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

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

383  

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 

390  

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) 

395  

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 

407  

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)) 

421  

422 if not to_add and not to_remove: 

423 return 

424 _apply_peer_access_diff(mr, to_add, to_remove)