cuda.core: capture complete Buffer deallocation recipe at creation

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -67,6 +67,7 @@ void clear_last_error() noexcept;
 extern decltype(&cuDevicePrimaryCtxRetain) p_cuDevicePrimaryCtxRetain;
 extern decltype(&cuDevicePrimaryCtxRelease) p_cuDevicePrimaryCtxRelease;
 extern decltype(&cuCtxGetCurrent) p_cuCtxGetCurrent;
+extern decltype(&cuCtxSetCurrent) p_cuCtxSetCurrent;
 extern decltype(&cuGreenCtxCreate) p_cuGreenCtxCreate;
 extern decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy;
 extern decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx;
@@ -423,7 +424,10 @@ DevicePtrHandle deviceptr_import_ipc(
 StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept;
 
 // Set the deallocation stream for a device pointer handle.
-void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept;
+// Returns CUDA_ERROR_INVALID_CONTEXT when a default-stream token cannot be
+// bound because no CUDA context is current.
+CUresult set_deallocation_stream(
+    const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept;
 
 // ============================================================================
 // Library handle functions

cuda_core/cuda/core/_memory/_buffer.pyi

@@ -39,12 +39,15 @@ class Buffer:
         ...
 
     @classmethod
-    def _init(cls, ptr: DevicePointerType, size: int, mr: MemoryResource | None=None, ipc_descriptor: IPCBufferDescriptor | None=None, owner: object | None=None) -> Buffer:
+    def _init(cls, ptr: DevicePointerType, size: int, mr: MemoryResource | None=None, ipc_descriptor: IPCBufferDescriptor | None=None, owner: object | None=None, *, stream: Stream | GraphBuilder | None=None) -> Buffer:
         """Create a Buffer from a raw pointer.
 
         When ``mr`` is provided, the buffer takes ownership: ``mr.deallocate()``
         is called when the buffer is closed or garbage collected.  When ``owner``
         is provided, the owner is kept alive but no deallocation is performed.
+        When ``mr`` is provided, a deallocation stream is recorded at creation
+        (``stream`` if given, otherwise ``default_stream()``). Recording a
+        default-stream token requires a CUDA context to be current.
         """
 
     @staticmethod
@@ -55,7 +58,7 @@ class Buffer:
         ...
 
     @staticmethod
-    def from_handle(ptr: DevicePointerType, size: int, mr: MemoryResource | None=None, owner: object | None=None) -> Buffer:
+    def from_handle(ptr: DevicePointerType, size: int, mr: MemoryResource | None=None, owner: object | None=None, *, stream: Stream | GraphBuilder | None=None) -> Buffer:
         """Create a new :class:`Buffer` object from a pointer.
 
         Parameters
@@ -72,6 +75,13 @@ class Buffer:
             An object holding external allocation that the ``ptr`` points to.
             The reference is kept as long as the buffer is alive.
             The ``owner`` and ``mr`` cannot be specified together.
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
+            Keyword-only. The stream used to order the buffer's deallocation
+            when ``mr`` owns the pointer. Defaults to ``default_stream()``.
+            Recording a default-stream token requires a CUDA context to be
+            current. If the buffer may be freed from a different host thread,
+            pass a stream other than the per-thread default stream, which
+            refers to a different stream on each thread.
 
         Note
         ----
@@ -264,7 +274,12 @@ class MemoryResource:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
             Keyword-only. The stream on which to perform the allocation
             asynchronously. Must be passed explicitly; pass
-            ``device.default_stream`` to use the default stream.
+            ``device.default_stream`` to use the default stream. For subclasses
+            that support stream-ordered deallocation, this stream also orders
+            the buffer's eventual deallocation, so if the buffer may be freed
+            from a different host thread, prefer a stream other than the
+            per-thread default stream, which refers to a different stream on
+            each thread.
 
         Returns
         -------

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -16,11 +16,14 @@ from cuda.core._memory cimport _ipc
 from cuda.core._resource_handles cimport (
     DevicePtrHandle,
     StreamHandle,
+    ContextHandle,
     deviceptr_create_with_owner,
     deviceptr_create_with_mr,
     register_mr_dealloc_callback,
     as_intptr,
     as_cu,
+    get_current_context,
+    get_stream_context,
     set_deallocation_stream,
 )
 from cuda.core.typing import DevicePointerType
@@ -49,23 +52,21 @@ cdef void _mr_dealloc_callback(
     size_t size,
     const StreamHandle& h_stream,
 ) noexcept:
-    """Called by the C++ deleter to deallocate via MemoryResource.deallocate.
-
-    This is the C++ teardown path: there is no Python caller frame from
-    which to obtain a stream. If the device-pointer handle was created
-    without ``set_deallocation_stream`` being called (e.g. buffers minted
-    via ``Buffer.from_handle(ptr, size, mr=mr)`` from DLPack import,
-    third-party adapters, or other foreign sources), ``h_stream`` is
-    empty here. Stream-ordered MR ``deallocate`` overrides reject
-    ``stream=None`` (issue #2001), so without a fallback the destructor
-    would print a warning and leak the allocation. Fall back to the
-    legacy/per-thread default stream so the free still happens; this is
-    the unique exception to the "no implicit default-stream fallback"
-    policy because the teardown has no other source of truth.
-    """
+    """Called by the C++ deleter to deallocate via MemoryResource.deallocate."""
     cdef Stream stream
     try:
-        stream = Stream._from_handle(Stream, h_stream) if h_stream else default_stream()
+        if not h_stream:
+            print(
+                "Warning: no deallocation stream was recorded; falling back to "
+                "the default stream for mr.deallocate() during Buffer "
+                "destruction. This is an internal cuda-core error; please "
+                "report it with your CUDA driver, CUDA Toolkit, and "
+                "cuda-python versions.",
+                file=sys.stderr,
+            )
+            stream = default_stream()
+        else:
+            stream = Stream._from_handle(Stream, h_stream)
         mr.deallocate(int(ptr), size, stream=stream)
     except Exception as exc:
         print(f"Warning: mr.deallocate() failed during Buffer destruction: {exc}",
@@ -74,6 +75,21 @@ cdef void _mr_dealloc_callback(
 register_mr_dealloc_callback(_mr_dealloc_callback)
 
 
+cdef inline void _require_deallocation_stream_context(Stream s) except *:
+    """Default-stream tokens need a current context to pin into the free recipe."""
+    cdef ContextHandle h_ctx
+    if get_stream_context(s._h_stream):
+        return
+    h_ctx = get_current_context()
+    if h_ctx:
+        return
+    raise RuntimeError(
+        "Cannot record a default deallocation stream when no CUDA context is "
+        "current. Call Device.set_current() first, or pass stream= with a "
+        "non-default Stream."
+    )
+
+
 __all__ = ['Buffer', 'MemoryResource']
 
 
@@ -176,20 +192,31 @@ cdef class Buffer:
     def _init(
         cls, ptr: DevicePointerType, size_t size, mr: MemoryResource | None = None,
         ipc_descriptor: IPCBufferDescriptor | None = None,
-        owner : object | None = None
+        owner : object | None = None,
+        *,
+        stream: Stream | GraphBuilder | None = None,
     ) -> Buffer:
         """Create a Buffer from a raw pointer.
 
         When ``mr`` is provided, the buffer takes ownership: ``mr.deallocate()``
         is called when the buffer is closed or garbage collected.  When ``owner``
         is provided, the owner is kept alive but no deallocation is performed.
+        When ``mr`` is provided, a deallocation stream is recorded at creation
+        (``stream`` if given, otherwise ``default_stream()``). Recording a
+        default-stream token requires a CUDA context to be current.
         """
         if mr is not None and owner is not None:
             raise ValueError("owner and memory resource cannot be both specified together")
+        if stream is not None and mr is None:
+            raise ValueError("stream requires a memory resource (mr)")
         cdef Buffer self = Buffer.__new__(cls)
         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)
+            _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))
         else:
             self._h_ptr = deviceptr_create_with_owner(c_ptr, owner)
         self._size = size
@@ -217,6 +244,8 @@ cdef class Buffer:
     def from_handle(
         ptr: DevicePointerType, size_t size, mr: MemoryResource | None = None,
         owner: object | None = None,
+        *,
+        stream: Stream | GraphBuilder | None = None,
     ) -> Buffer:
         """Create a new :class:`Buffer` object from a pointer.
 
@@ -234,14 +263,21 @@ cdef class Buffer:
             An object holding external allocation that the ``ptr`` points to.
             The reference is kept as long as the buffer is alive.
             The ``owner`` and ``mr`` cannot be specified together.
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
+            Keyword-only. The stream used to order the buffer's deallocation
+            when ``mr`` owns the pointer. Defaults to ``default_stream()``.
+            Recording a default-stream token requires a CUDA context to be
+            current. If the buffer may be freed from a different host thread,
+            pass a stream other than the per-thread default stream, which
+            refers to a different stream on each thread.
 
         Note
         ----
         When neither ``mr`` nor ``owner`` is specified, this creates a
         non-owning reference.  The pointer will NOT be freed when the
         :class:`Buffer` is closed or garbage collected.
         """
-        return Buffer._init(ptr, size, mr=mr, owner=owner)
+        return Buffer._init(ptr, size, mr=mr, owner=owner, stream=stream)
 
     @classmethod
     def from_ipc_descriptor(
@@ -546,7 +582,12 @@ cdef class MemoryResource:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
             Keyword-only. The stream on which to perform the allocation
             asynchronously. Must be passed explicitly; pass
-            ``device.default_stream`` to use the default stream.
+            ``device.default_stream`` to use the default stream. For subclasses
+            that support stream-ordered deallocation, this stream also orders
+            the buffer's eventual deallocation, so if the buffer may be freed
+            from a different host thread, prefer a stream other than the
+            per-thread default stream, which refers to a different stream on
+            each thread.
 
         Returns
         -------
@@ -627,7 +668,8 @@ cdef inline void Buffer_close(Buffer self, object stream):
     # Update deallocation stream if provided
     if stream is not None:
         s = Stream_accept(stream)
-        set_deallocation_stream(self._h_ptr, s._h_stream)
+        _require_deallocation_stream_context(s)
+        HANDLE_RETURN(set_deallocation_stream(self._h_ptr, s._h_stream))
     # Reset handle - RAII deleter will free the memory (and release owner ref in C++)
     self._h_ptr.reset()
     self._size = 0

cuda_core/cuda/core/_memory/_graph_memory_resource.pyx

@@ -225,7 +225,7 @@ cdef inline Buffer GMR_allocate(cyGraphMemoryResource self, size_t size, Stream
     return Buffer_from_deviceptr_handle(h_ptr, size, self, None)
 
 
-cdef inline void GMR_deallocate(intptr_t ptr, size_t size, Stream stream) noexcept:
+cdef inline void GMR_deallocate(intptr_t ptr, size_t size, Stream stream) except *:
     cdef cydriver.CUstream s = as_cu(stream._h_stream)
     cdef cydriver.CUdeviceptr devptr = <cydriver.CUdeviceptr>ptr
     with nogil:

cuda_core/cuda/core/_memory/_managed_buffer.py

@@ -154,6 +154,8 @@ def from_handle(
         size: int,
         mr: MemoryResource | None = None,
         owner: object | None = None,
+        *,
+        stream: Stream | GraphBuilder | None = None,
     ) -> Buffer:
         """Wrap an existing managed-memory pointer in a :class:`ManagedBuffer`.
 
@@ -173,8 +175,15 @@ def from_handle(
         owner : object, optional
             An object that keeps the underlying allocation alive.
             ``owner`` and ``mr`` cannot both be specified.
+        stream : Stream | GraphBuilder, optional
+            Keyword-only. The stream used to order the buffer's deallocation
+            when ``mr`` owns the pointer. Defaults to ``default_stream()``.
+            Recording a default-stream token requires a CUDA context to be
+            current. If the buffer may be freed from a different host thread,
+            pass a stream other than the per-thread default stream, which
+            refers to a different stream on each thread.
         """
-        return cls._init(ptr, size, mr=mr, owner=owner)
+        return cls._init(ptr, size, mr=mr, owner=owner, stream=stream)
 
     @property
     def read_mostly(self) -> bool:

cuda_core/cuda/core/_memory/_memory_pool.pyx

@@ -347,14 +347,11 @@ cdef Buffer _MP_allocate(_MemPool self, size_t size, Stream stream, type cls = B
 
 cdef inline void _MP_deallocate(
     _MemPool self, uintptr_t ptr, size_t size, Stream stream
-) noexcept nogil:
+) except *:
     cdef cydriver.CUstream s = as_cu(stream._h_stream)
     cdef cydriver.CUdeviceptr devptr = <cydriver.CUdeviceptr>ptr
-    cdef cydriver.CUresult r
     with nogil:
-        r = cydriver.cuMemFreeAsync(devptr, s)
-        if r != cydriver.CUDA_ERROR_INVALID_CONTEXT:
-            HANDLE_RETURN(r)
+        HANDLE_RETURN(cydriver.cuMemFreeAsync(devptr, s))
 
 
 cdef inline _MP_close(_MemPool self):

cuda_core/cuda/core/_resource_handles.pxd

@@ -239,7 +239,7 @@ cdef void register_mr_dealloc_callback(MRDeallocCallback cb) noexcept
 cdef DevicePtrHandle deviceptr_import_ipc(
     const MemoryPoolHandle& h_pool, const void* export_data, const StreamHandle& h_stream) except+ nogil
 cdef StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept nogil
-cdef void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil
+cdef cydriver.CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil
 
 # Library handles
 cdef LibraryHandle create_library_handle_from_file(const char* path) except+ nogil

cuda_core/cuda/core/_resource_handles.pyx

@@ -128,7 +128,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
         const MemoryPoolHandle& h_pool, const void* export_data, const StreamHandle& h_stream) except+ nogil
     StreamHandle deallocation_stream "cuda_core::deallocation_stream" (
         const DevicePtrHandle& h) noexcept nogil
-    void set_deallocation_stream "cuda_core::set_deallocation_stream" (
+    cydriver.CUresult set_deallocation_stream "cuda_core::set_deallocation_stream" (
         const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil
 
     # Library handles
@@ -293,6 +293,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     void* p_cuDevicePrimaryCtxRetain "reinterpret_cast<void*&>(cuda_core::p_cuDevicePrimaryCtxRetain)"
     void* p_cuDevicePrimaryCtxRelease "reinterpret_cast<void*&>(cuda_core::p_cuDevicePrimaryCtxRelease)"
     void* p_cuCtxGetCurrent "reinterpret_cast<void*&>(cuda_core::p_cuCtxGetCurrent)"
+    void* p_cuCtxSetCurrent "reinterpret_cast<void*&>(cuda_core::p_cuCtxSetCurrent)"
     void* p_cuGreenCtxCreate "reinterpret_cast<void*&>(cuda_core::p_cuGreenCtxCreate)"
     void* p_cuGreenCtxDestroy "reinterpret_cast<void*&>(cuda_core::p_cuGreenCtxDestroy)"
     void* p_cuCtxFromGreenCtx "reinterpret_cast<void*&>(cuda_core::p_cuCtxFromGreenCtx)"
@@ -397,6 +398,7 @@ cdef void* _get_optional_driver_fn(str name):
 
 cdef void _init_driver_fn_pointers() noexcept:
     global p_cuDevicePrimaryCtxRetain, p_cuDevicePrimaryCtxRelease, p_cuCtxGetCurrent
+    global p_cuCtxSetCurrent
     global p_cuGreenCtxCreate, p_cuGreenCtxDestroy, p_cuCtxFromGreenCtx
     global p_cuDevResourceGenerateDesc, p_cuGreenCtxStreamCreate
     global p_cuStreamCreateWithPriority, p_cuStreamDestroy
@@ -425,6 +427,7 @@ cdef void _init_driver_fn_pointers() noexcept:
     p_cuDevicePrimaryCtxRetain = _get_driver_fn("cuDevicePrimaryCtxRetain")
     p_cuDevicePrimaryCtxRelease = _get_driver_fn("cuDevicePrimaryCtxRelease")
     p_cuCtxGetCurrent = _get_driver_fn("cuCtxGetCurrent")
+    p_cuCtxSetCurrent = _get_driver_fn("cuCtxSetCurrent")
     p_cuGreenCtxCreate = _get_optional_driver_fn("cuGreenCtxCreate")
     p_cuGreenCtxDestroy = _get_optional_driver_fn("cuGreenCtxDestroy")
     p_cuCtxFromGreenCtx = _get_optional_driver_fn("cuCtxFromGreenCtx")

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

@@ -9,6 +9,25 @@
 Fixes and enhancements
 ----------------------
 
+- A :class:`Buffer` is now freed correctly even when the CUDA context current
+  at teardown is not the one it was allocated in, or when no context is current
+  at all. This happens routinely when a buffer is released by the garbage
+  collector on another thread or by deferred CUDA graph cleanup; previously the
+  free could fail or be skipped, leaking the allocation.
+  (`#2497 <https://github.com/NVIDIA/cuda-python/issues/2497>`__)
+
+- :meth:`Buffer.from_handle` and :meth:`ManagedBuffer.from_handle` accept a
+  keyword-only ``stream`` that records the stream used to order the buffer's
+  deallocation when the memory resource owns the pointer. It defaults to
+  ``default_stream()``, which requires a CUDA context to be current so the
+  free recipe can pin that context.
+  (`#2497 <https://github.com/NVIDIA/cuda-python/issues/2497>`__)
+
+- 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
+  buffer cleanup remains non-raising and reports failures as warnings.
+
 - Graph node resources are now retained independently across graph clones,
   executable graphs, updates, node deletion, and in-flight launches. Previously,
   modifying a graph definition could release resources still used by an