cuda.core: convert peer_accessible_by to a live MutableSet view

DeviceMemoryResource.peer_accessible_by previously returned a sorted
tuple[int, ...] backed by a Python-level cache, which was prone to
divergence from driver state across multiple wrappers around the same
memory pool. The setter accepted Device | int and emitted a single
batched cuMemPoolSetAccess covering the diff against the cache.

This commit replaces the property with a live driver-backed view:

- Adds PeerAccessibleBySetProxy in _memory/_peer_access_utils.py, a
  collections.abc.MutableSet whose reads call cuMemPoolGetAccess and
  whose writes call cuMemPoolSetAccess. Iteration yields Device
  objects; add, discard, and __contains__ accept either a Device or a
  device-ordinal int. The proxy is constructed fresh on every property
  access, so there is nothing to cache or pickle.

- Drops the _peer_accessible_by cache field (and its initializations
  in __cinit__, _DMR_init, and from_allocation_handle), eliminating
  the owned/non-owned read split. All pools now share the same code
  path and always query the driver.

- All bulk operations on the proxy (update, |=, &=, -=, ^=, clear,
  pop) issue exactly one cuMemPoolSetAccess call. Peer-access
  transitions can take seconds per pool because every existing memory
  mapping is updated, so coalescing into a single driver call lets the
  toolkit handle the mappings in parallel. The property setter
  (mr.peer_accessible_by = [...]) preserves its original single-call
  behavior via the same shared planner path.

- Single-element add validates can_access_peer through
  plan_peer_access_update, matching the existing setter contract.

This is a breaking change captured in the v1.0.0 release notes.
Callers comparing against tuples must update to set comparisons
(mr.peer_accessible_by == {Device(0)}). Existing tests are migrated;
new tests for set-interface conformance are intentionally deferred to
a follow-up.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: Andy-Jost

cuda_core/cuda/core/_memory/_device_memory_resource.pxd

@@ -9,7 +9,6 @@ from cuda.core._memory._ipc cimport IPCDataForMR
 cdef class DeviceMemoryResource(_MemPool):
     cdef:
         int _dev_id
-        object _peer_accessible_by
 
 
 cpdef DMR_mempool_get_access(DeviceMemoryResource, int)

cuda_core/cuda/core/_memory/_device_memory_resource.pyx

@@ -25,7 +25,11 @@ import multiprocessing
 import platform  # no-cython-lint
 import uuid
 
-from ._peer_access_utils import plan_peer_access_update
+from ._peer_access_utils import (
+    PeerAccessibleBySetProxy,
+    _resolve_peer_device_id,
+    plan_peer_access_update,
+)
 from cuda.core._utils.cuda_utils import check_multiprocessing_start_method
 
 __all__ = ['DeviceMemoryResource', 'DeviceMemoryResourceOptions']
@@ -131,7 +135,6 @@ cdef class DeviceMemoryResource(_MemPool):
 
     def __cinit__(self, *args, **kwargs):
         self._dev_id = cydriver.CU_DEVICE_INVALID
-        self._peer_accessible_by = None
 
     def __init__(self, device_id: Device | int, options=None):
         _DMR_init(self, device_id, options)
@@ -191,7 +194,6 @@ cdef class DeviceMemoryResource(_MemPool):
             _ipc.MP_from_allocation_handle(cls, alloc_handle))
         from .._device import Device
         mr._dev_id = Device(device_id).device_id
-        mr._peer_accessible_by = ()
         return mr
 
     @property
@@ -213,34 +215,60 @@ cdef class DeviceMemoryResource(_MemPool):
     @property
     def peer_accessible_by(self):
         """
-        Get or set the devices that can access allocations from this memory
-        pool. Access can be modified at any time and affects all allocations
+        Live driver-backed set view of the devices that can access allocations
         from this memory pool.
 
-        Returns a tuple of sorted device IDs that currently have peer access to
-        allocations from this memory pool.
+        Returns a :class:`PeerAccessibleBySetProxy` (a
+        :class:`collections.abc.MutableSet`) whose reads call
+        ``cuMemPoolGetAccess`` and whose writes call ``cuMemPoolSetAccess``.
+        Iteration yields :class:`Device` objects; ``add``, ``discard``, and
+        ``__contains__`` accept either a :class:`Device` or a device-ordinal
+        ``int``. There is no in-memory cache, so the view always reflects the
+        current driver state and stays consistent across multiple wrappers
+        around the same pool.
 
-        When setting, accepts a sequence of :obj:`~_device.Device` objects or device IDs.
-        Setting to an empty sequence revokes all peer access.
+        When setting, accepts an iterable of :obj:`~_device.Device` objects or
+        device IDs. Setting replaces the full set in a single batched driver call.
 
-        For non-owned pools (the default or current device pool), the state
-        is always queried from the driver to reflect changes made by other
-        wrappers or direct driver calls.
+        Bulk operations (``update``, ``|=``, ``&=``, ``-=``, ``^=``, ``clear``,
+        and the property setter) each issue exactly one ``cuMemPoolSetAccess``
+        call so the toolkit can update existing memory mappings in parallel.
 
         Examples
         --------
         >>> dmr = DeviceMemoryResource(0)
-        >>> dmr.peer_accessible_by = [1]  # Grant access to device 1
-        >>> assert dmr.peer_accessible_by == (1,)
-        >>> dmr.peer_accessible_by = []  # Revoke access
+        >>> dmr.peer_accessible_by.add(1)            # grant access to device 1
+        >>> assert dmr.peer_accessible_by == {Device(1)}
+        >>> dmr.peer_accessible_by |= {Device(2)}    # batched grant via |=
+        >>> dmr.peer_accessible_by = []              # revoke all in one call
         """
-        if not self._mempool_owned:
-            _DMR_query_peer_access(self)
-        return self._peer_accessible_by
+        return PeerAccessibleBySetProxy(self)
 
     @peer_accessible_by.setter
     def peer_accessible_by(self, devices):
-        _DMR_set_peer_accessible_by(self, devices)
+        _DMR_replace_peer_accessible_by(self, devices)
+
+    def _query_peer_access_ids(self):
+        """Return the current peer device IDs as a sorted tuple of ints.
+
+        Always queries the driver via ``cuMemPoolGetAccess`` for every visible
+        device. Used by :class:`PeerAccessibleBySetProxy` for ``__iter__`` and
+        ``__len__``.
+        """
+        return _DMR_query_peer_access_ids(self)
+
+    def _peer_access_includes(self, int dev_id) -> bool:
+        """Return True if peer access from ``dev_id`` is currently granted."""
+        return _DMR_peer_access_includes(self, dev_id)
+
+    def _apply_peer_access_diff(self, to_add, to_remove):
+        """Issue a single ``cuMemPoolSetAccess`` for the given add/remove deltas.
+
+        ``to_add`` and ``to_remove`` are iterables of device-ordinal ints.
+        Both must already be filtered (no owner, no overlap, no duplicates).
+        Used by :class:`PeerAccessibleBySetProxy` for batched writes.
+        """
+        _DMR_apply_peer_access_diff(self, tuple(to_add), tuple(to_remove))
 
     @property
     def is_device_accessible(self) -> bool:
@@ -253,8 +281,8 @@ cdef class DeviceMemoryResource(_MemPool):
         return False
 
 
-cdef inline _DMR_query_peer_access(DeviceMemoryResource self):
-    """Query the driver for the actual peer access state of this pool."""
+cdef inline tuple _DMR_query_peer_access_ids(DeviceMemoryResource self):
+    """Return the current peer device IDs as a sorted tuple of ints."""
     cdef int total
     cdef cydriver.CUmemAccess_flags flags
     cdef cydriver.CUmemLocation location
@@ -273,59 +301,74 @@ cdef inline _DMR_query_peer_access(DeviceMemoryResource self):
         if flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE:
             peers.append(dev_id)
 
-    self._peer_accessible_by = tuple(sorted(peers))
+    return tuple(sorted(peers))
 
 
-cdef inline _DMR_set_peer_accessible_by(DeviceMemoryResource self, devices):
-    from .._device import Device
+cdef inline bint _DMR_peer_access_includes(DeviceMemoryResource self, int dev_id):
+    """Return True if peer access from ``dev_id`` is currently granted."""
+    cdef cydriver.CUmemAccess_flags flags
+    cdef cydriver.CUmemLocation location
 
-    this_dev = Device(self._dev_id)
-    cdef object resolve_device_id = lambda dev: Device(dev).device_id
-    cdef object plan
-    cdef tuple target_ids
-    cdef tuple to_add
-    cdef tuple to_rm
-    if not self._mempool_owned:
-        _DMR_query_peer_access(self)
-    plan = plan_peer_access_update(
-        owner_device_id=self._dev_id,
-        current_peer_ids=self._peer_accessible_by,
-        requested_devices=devices,
-        resolve_device_id=resolve_device_id,
-        can_access_peer=this_dev.can_access_peer,
-    )
-    target_ids = plan.target_ids
-    to_add = plan.to_add
-    to_rm = plan.to_remove
-    cdef size_t count = len(to_add) + len(to_rm)
+    location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+    location.id = dev_id
+    with nogil:
+        HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(self._h_pool), &location))
+    return flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
+
+
+cdef inline _DMR_apply_peer_access_diff(
+    DeviceMemoryResource self, tuple to_add, tuple to_remove
+):
+    """Issue one ``cuMemPoolSetAccess`` for the given add/remove deltas."""
+    cdef size_t count = len(to_add) + len(to_remove)
     cdef cydriver.CUmemAccessDesc* access_desc = NULL
     cdef size_t i = 0
 
-    if count > 0:
-        access_desc = <cydriver.CUmemAccessDesc*>PyMem_Malloc(count * sizeof(cydriver.CUmemAccessDesc))
-        if access_desc == NULL:
-            raise MemoryError("Failed to allocate memory for access descriptors")
+    if count == 0:
+        return
+
+    access_desc = <cydriver.CUmemAccessDesc*>PyMem_Malloc(count * sizeof(cydriver.CUmemAccessDesc))
+    if access_desc == NULL:
+        raise MemoryError("Failed to allocate memory for access descriptors")
+
+    try:
+        for dev_id in to_add:
+            access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
+            access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+            access_desc[i].location.id = dev_id
+            i += 1
+        for dev_id in to_remove:
+            access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_NONE
+            access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+            access_desc[i].location.id = dev_id
+            i += 1
+
+        with nogil:
+            HANDLE_RETURN(cydriver.cuMemPoolSetAccess(as_cu(self._h_pool), access_desc, count))
+    finally:
+        if access_desc != NULL:
+            PyMem_Free(access_desc)
 
-        try:
-            for dev_id in to_add:
-                access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
-                access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
-                access_desc[i].location.id = dev_id
-                i += 1
 
-            for dev_id in to_rm:
-                access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_NONE
-                access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
-                access_desc[i].location.id = dev_id
-                i += 1
+cdef inline _DMR_replace_peer_accessible_by(DeviceMemoryResource self, devices):
+    """Replace the full peer-access set in a single batched driver call.
 
-            with nogil:
-                HANDLE_RETURN(cydriver.cuMemPoolSetAccess(as_cu(self._h_pool), access_desc, count))
-        finally:
-            if access_desc != NULL:
-                PyMem_Free(access_desc)
+    Backs the ``mr.peer_accessible_by = [...]`` setter. Uses the same planner
+    as the proxy's bulk ops; the only difference is that adds and removes are
+    derived from the symmetric difference between current driver state and the
+    requested target set.
+    """
+    from .._device import Device
 
-        self._peer_accessible_by = tuple(target_ids)
+    this_dev = Device(self._dev_id)
+    plan = plan_peer_access_update(
+        owner_device_id=self._dev_id,
+        current_peer_ids=_DMR_query_peer_access_ids(self),
+        requested_devices=devices,
+        resolve_device_id=_resolve_peer_device_id,
+        can_access_peer=this_dev.can_access_peer,
+    )
+    _DMR_apply_peer_access_diff(self, plan.to_add, plan.to_remove)
 
 
 cdef inline _DMR_init(DeviceMemoryResource self, device_id, options):
@@ -351,7 +394,6 @@ cdef inline _DMR_init(DeviceMemoryResource self, device_id, options):
         self._mempool_owned = False
         MP_raise_release_threshold(self)
     else:
-        self._peer_accessible_by = ()
         MP_init_create_pool(
             self,
             cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,