cuda.core: allow updating Buffer deallocation streams

Author: Andy-Jost

cuda_core/cuda/core/_memory/_buffer.pyi

@@ -127,6 +127,41 @@ class Buffer:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
             The stream object to use for asynchronous deallocation. If None,
             the deallocation stream stored in the handle is used.
+
+        See Also
+        --------
+        set_deallocation_stream
+            Change the deallocation stream without closing the buffer.
+        """
+
+    def set_deallocation_stream(self, stream: Stream | GraphBuilder) -> None:
+        """Change the stream that orders this buffer's eventual deallocation.
+
+        The buffer remains open and usable. A later :meth:`close` without a
+        stream, garbage collection, or release of the final retained device
+        pointer handle uses the replacement stream.
+
+        This method does not synchronize streams or establish dependencies.
+        The caller must ensure that allocation and all accesses are ordered
+        before the deallocation on ``stream``.
+
+        Parameters
+        ----------
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
+            The stream to use for eventual asynchronous deallocation.
+
+        Raises
+        ------
+        RuntimeError
+            If the buffer is already closed, or if a default-stream token
+            cannot be bound because no CUDA context is current.
+        TypeError
+            If ``stream`` is ``None`` or is not an accepted stream object.
+
+        Notes
+        -----
+        Synchronizing concurrent mutation and destruction of the same buffer
+        is the caller's responsibility.
         """
 
     def __enter__(self):

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -213,7 +213,9 @@ cdef class Buffer:
         cdef uintptr_t c_ptr = <uintptr_t>(int(ptr))
         cdef Stream s
         if mr is not None:
-            s = Stream_accept(default_stream() if stream is None else stream)
+            # Validate before taking ownership so a bad stream does not cause
+            # construction failure to deallocate the caller's pointer.
+            s = default_stream() if stream is None else Stream_accept(stream)
             _require_deallocation_stream_context(s)
             self._h_ptr = deviceptr_create_with_mr(c_ptr, size, mr)
             HANDLE_RETURN(set_deallocation_stream(self._h_ptr, s._h_stream))
@@ -238,7 +240,7 @@ cdef class Buffer:
         # The parent process's stream is not portable across processes, so the
         # pickle path cannot thread an explicit stream through. Seed the
         # imported buffer's deallocation with the current context's default
-        # stream; the receiver can override via buffer.close(stream).
+        # stream; the receiver can override it before or during close.
         return Buffer.from_ipc_descriptor(mr, ipc_descriptor, stream=default_stream())
 
     def __reduce__(self) -> tuple[object, ...]:
@@ -333,9 +335,45 @@ cdef class Buffer:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
             The stream object to use for asynchronous deallocation. If None,
             the deallocation stream stored in the handle is used.
+
+        See Also
+        --------
+        set_deallocation_stream
+            Change the deallocation stream without closing the buffer.
         """
         Buffer_close(self, stream)
 
+    def set_deallocation_stream(self, stream: Stream | GraphBuilder) -> None:
+        """Change the stream that orders this buffer's eventual deallocation.
+
+        The buffer remains open and usable. A later :meth:`close` without a
+        stream, garbage collection, or release of the final retained device
+        pointer handle uses the replacement stream.
+
+        This method does not synchronize streams or establish dependencies.
+        The caller must ensure that allocation and all accesses are ordered
+        before the deallocation on ``stream``.
+
+        Parameters
+        ----------
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
+            The stream to use for eventual asynchronous deallocation.
+
+        Raises
+        ------
+        RuntimeError
+            If the buffer is already closed, or if a default-stream token
+            cannot be bound because no CUDA context is current.
+        TypeError
+            If ``stream`` is ``None`` or is not an accepted stream object.
+
+        Notes
+        -----
+        Synchronizing concurrent mutation and destruction of the same buffer
+        is the caller's responsibility.
+        """
+        Buffer_set_deallocation_stream(self, stream)
+
     def __enter__(self):
         return self
 
@@ -667,16 +705,23 @@ cdef Buffer Buffer_from_deviceptr_handle(
     return buf
 
 
+cdef inline void Buffer_set_deallocation_stream(Buffer self, object stream):
+    """Validate and replace a live buffer's deallocation recipe."""
+    cdef Stream s
+    if not self._h_ptr:
+        raise RuntimeError("Cannot set the deallocation stream on a closed Buffer")
+    s = Stream_accept(stream)
+    _require_deallocation_stream_context(s)
+    HANDLE_RETURN(set_deallocation_stream(self._h_ptr, s._h_stream))
+
+
 cdef inline void Buffer_close(Buffer self, object stream):
     """Close a buffer, freeing its memory."""
-    cdef Stream s
     if not self._h_ptr:
         return
     # Update deallocation stream if provided
     if stream is not None:
-        s = Stream_accept(stream)
-        _require_deallocation_stream_context(s)
-        HANDLE_RETURN(set_deallocation_stream(self._h_ptr, s._h_stream))
+        Buffer_set_deallocation_stream(self, stream)
     # Reset handle - RAII deleter will free the memory (and release owner ref in C++)
     self._h_ptr.reset()
     self._size = 0

cuda_core/docs/source/api.rst

@@ -77,6 +77,12 @@ Memory management
    ManagedMemoryResourceOptions
    VirtualMemoryResourceOptions
 
+A :class:`Buffer` records the stream that will order its eventual deallocation.
+Use :meth:`Buffer.set_deallocation_stream` to replace that stream without
+closing the buffer. Changing the recorded stream does not synchronize streams;
+the caller must order allocation and every access before the deallocation,
+using events or other CUDA synchronization mechanisms as needed.
+
 
 CUDA compilation toolchain
 --------------------------

cuda_core/docs/source/release/1.2.0-notes.rst

@@ -23,6 +23,10 @@ Fixes and enhancements
   free recipe can pin that context.
   (`#2497 <https://github.com/NVIDIA/cuda-python/issues/2497>`__)
 
+- Added :meth:`Buffer.set_deallocation_stream` to change the stream that orders
+  a buffer's eventual deallocation without closing the buffer.
+  (`#2600 <https://github.com/NVIDIA/cuda-python/issues/2600>`__)
+
 - Explicit calls to ``deallocate()`` on pool-backed memory resources and
   :class:`GraphMemoryResource` now propagate errors from the underlying CUDA
   free operation. Previously, these errors could be suppressed. Automatic

cuda_core/tests/test_memory.py

@@ -514,6 +514,85 @@ def deallocate(self, ptr, size, *, stream=None):
     assert received["stream"].handle == stream.handle
 
 
+class _StreamCaptureMemoryResource(MemoryResource):
+    def __init__(self, device):
+        self.device = device
+        self.deallocation_streams = []
+
+    @property
+    def is_device_accessible(self):
+        return True
+
+    @property
+    def is_host_accessible(self):
+        return False
+
+    @property
+    def device_id(self):
+        return self.device.device_id
+
+    def allocate(self, size, *, stream):
+        raise NotImplementedError
+
+    def deallocate(self, ptr, size, *, stream):
+        self.deallocation_streams.append(stream)
+
+
+@pytest.mark.agent_authored(model="gpt-5.6")
+@pytest.mark.parametrize(
+    ("configuration", "destruction"),
+    [
+        ("initialization", "close"),
+        ("initialization", "gc"),
+        ("setter", "close"),
+        ("setter", "gc"),
+        ("close", "close"),
+    ],
+)
+def test_buffer_deallocation_stream_configuration_paths(configuration, destruction):
+    """Creation, mutation, and close overrides use the requested stream."""
+    import gc
+
+    device = Device()
+    device.set_current()
+    initial_stream = device.create_stream()
+    target_stream = device.create_stream()
+    mr = _StreamCaptureMemoryResource(device)
+
+    stream = target_stream if configuration == "initialization" else initial_stream
+    buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream)
+    if configuration == "setter":
+        handle = buf.handle
+        buf.set_deallocation_stream(target_stream)
+        assert buf.handle == handle
+        assert buf.size == 1024
+
+    if destruction == "close":
+        buf.close(stream=target_stream if configuration == "close" else None)
+    else:
+        del buf
+        gc.collect()
+
+    assert len(mr.deallocation_streams) == 1
+    assert mr.deallocation_streams[0].handle == target_stream.handle
+
+
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_set_deallocation_stream_rejects_none_and_closed_buffer():
+    device = Device()
+    device.set_current()
+    stream = device.create_stream()
+    mr = _StreamCaptureMemoryResource(device)
+    buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream)
+
+    with pytest.raises(TypeError, match="stream is required"):
+        buf.set_deallocation_stream(None)
+
+    buf.close()
+    with pytest.raises(RuntimeError, match="closed Buffer"):
+        buf.set_deallocation_stream(stream)
+
+
 def test_from_handle_mr_records_default_stream():
     """When a Buffer is minted via :meth:`Buffer.from_handle` with ``mr`` but
     without an explicit ``stream=``, the deallocation stream is recorded at