Handle std::bad_alloc in make_deallocation_stream / set_deallocation_stream

cuda_core/cuda/core/_cpp/rt/api.hpp

@@ -239,8 +239,9 @@ StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept;
 // Set the deallocation stream for a device pointer handle.
 // Returns CUDA_ERROR_INVALID_CONTEXT when a default-stream token cannot be
 // bound because no CUDA context is current.
+// May throw std::bad_alloc.
 CUresult set_deallocation_stream(
-    const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept;
+    const DevicePtrHandle& h, const StreamHandle& h_stream);
 
 // ============================================================================
 // Virtual memory management (VMM_DESIGN.md)

cuda_core/cuda/core/_cpp/rt/internal.hpp

@@ -59,8 +59,9 @@ struct DeallocationStream {
 
 // Implemented in stream.cpp
 ContextHandle deallocation_context(const DeallocationStream& stream) noexcept;
-// Implemented in stream.cpp
-bool make_deallocation_stream(const StreamHandle& h, DeallocationStream& out) noexcept;
+// Implemented in stream.cpp. Not noexcept: binding a default-stream token
+// allocates. May throw std::bad_alloc.
+bool make_deallocation_stream(const StreamHandle& h, DeallocationStream& out);
 
 // Implemented in virtual_memory.cpp. Synchronize the stream a VMM buffer
 // recorded before its mappings are released. Skips the sync, with a report

cuda_core/cuda/core/_cpp/rt/memory.cpp

@@ -11,6 +11,7 @@
 #include <atomic>
 #include <cstddef>
 #include <cstring>
+#include <exception>
 #include <memory>
 #include <mutex>
 #include <stdexcept>
@@ -171,7 +172,7 @@ StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept {
 }
 
 // Replace the stream that orders a device pointer's deallocation.
-CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept {
+CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) {
     if (!h) {
         return CUDA_ERROR_INVALID_VALUE;
     }
@@ -183,21 +184,51 @@ CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h
     return CUDA_SUCCESS;
 }
 
+// Build the box for a freshly acquired device pointer. Binding the deallocation
+// stream and allocating the box can fail: make_deallocation_stream returns
+// false (err set) when no context is current, and either step can throw
+// std::bad_alloc. On any failure run undo() to release the pointer and return
+// nullptr; a thrown exception is left in `failure` for the caller to rethrow
+// once it holds no lock (undo may not report under a lock, see
+// deviceptr_import_ipc). Never throws.
+template <typename Undo>
+static DevicePtrBox* make_deviceptr_box(
+        CUdeviceptr ptr, const StreamHandle& h_stream, Undo&& undo,
+        std::exception_ptr& failure) noexcept {
+    DevicePtrBox* box = nullptr;
+    try {
+        DeallocationStream ds;
+        if (make_deallocation_stream(h_stream, ds)) {
+            box = new DevicePtrBox{ptr, std::move(ds)};
+        }
+    } catch (...) {
+        failure = std::current_exception();
+    }
+    if (!box) {
+        undo();
+    }
+    return box;
+}
+
 DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) {
     GILReleaseGuard gil;
     CUdeviceptr ptr;
     if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemAllocFromPoolAsync, &ptr, size, *h_pool, as_cu(h_stream)))) {
         return {};
     }
 
-    DeallocationStream ds;
-    if (!make_deallocation_stream(h_stream, ds)) {
-        pw_cuMemFreeAsync(ptr, as_cu(h_stream));
+    std::exception_ptr failure;
+    DevicePtrBox* raw = make_deviceptr_box(
+        ptr, h_stream, [&]() noexcept { pw_cuMemFreeAsync(ptr, as_cu(h_stream)); }, failure);
+    if (!raw) {
+        if (failure) {
+            std::rethrow_exception(failure);
+        }
         return {};
     }
 
     auto box = std::shared_ptr<DevicePtrBox>(
-        new DevicePtrBox{ptr, std::move(ds)},
+        raw,
         [h_pool](DevicePtrBox* b) {
             GILReleaseGuard gil;
             const DeallocationStream& stream = b->deallocation;
@@ -220,14 +251,18 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream)
         return {};
     }
 
-    DeallocationStream ds;
-    if (!make_deallocation_stream(h_stream, ds)) {
-        pw_cuMemFreeAsync(ptr, as_cu(h_stream));
+    std::exception_ptr failure;
+    DevicePtrBox* raw = make_deviceptr_box(
+        ptr, h_stream, [&]() noexcept { pw_cuMemFreeAsync(ptr, as_cu(h_stream)); }, failure);
+    if (!raw) {
+        if (failure) {
+            std::rethrow_exception(failure);
+        }
         return {};
     }
 
     auto box = std::shared_ptr<DevicePtrBox>(
-        new DevicePtrBox{ptr, std::move(ds)},
+        raw,
         [](DevicePtrBox* b) {
             GILReleaseGuard gil;
             const DeallocationStream& stream = b->deallocation;
@@ -488,6 +523,7 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
         GILReleaseGuard gil;
         CUresult discard_status = CUDA_SUCCESS;
         CUdeviceptr discarded = 0;
+        std::exception_ptr failure;
         {
             std::lock_guard<std::mutex> lock(ipc_import_mutex);
 
@@ -500,10 +536,19 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
                 return {};
             }
 
-            DeallocationStream ds;
-            if (make_deallocation_stream(h_stream, ds)) {
+            // No deallocation stream could be recorded (or the box could not be
+            // allocated): discard the import with the raw call (a pw_ report would
+            // acquire the GIL under the mutex).
+            DevicePtrBox* raw = make_deviceptr_box(
+                ptr, h_stream,
+                [&]() noexcept {
+                    discard_status = p_cuMemFreeAsync(ptr, as_cu(h_stream));  // raw: under ipc_import_mutex
+                    discarded = ptr;
+                },
+                failure);
+            if (raw) {
                 auto box = std::shared_ptr<DevicePtrBox>(
-                    new DevicePtrBox{ptr, std::move(ds)},
+                    raw,
                     [h_pool, key](DevicePtrBox* b) {
                         // Release the GIL first (the GIL is the outermost lock), then hold the
                         // mutex across unregister + free: a concurrent import that finds this
@@ -530,11 +575,6 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
                 ipc_ptr_cache.register_handle(key, h);
                 return h;
             }
-
-            // No deallocation stream could be recorded: discard the import with
-            // the raw call (a pw_ report would acquire the GIL under the mutex).
-            discard_status = p_cuMemFreeAsync(ptr, as_cu(h_stream));  // raw: under ipc_import_mutex
-            discarded = ptr;
         }
         if (discard_status != CUDA_SUCCESS) {
             char operation[160];
@@ -543,6 +583,9 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
                               "failed while discarding an IPC import that could not record a "
                               "deallocation stream; the mapping leaked");
         }
+        if (failure) {
+            std::rethrow_exception(failure);
+        }
         return {};
 
     } else {
@@ -552,14 +595,18 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
             return {};
         }
 
-        DeallocationStream ds;
-        if (!make_deallocation_stream(h_stream, ds)) {
-            pw_cuMemFreeAsync(ptr, as_cu(h_stream));
+        std::exception_ptr failure;
+        DevicePtrBox* raw = make_deviceptr_box(
+            ptr, h_stream, [&]() noexcept { pw_cuMemFreeAsync(ptr, as_cu(h_stream)); }, failure);
+        if (!raw) {
+            if (failure) {
+                std::rethrow_exception(failure);
+            }
             return {};
         }
 
         auto box = std::shared_ptr<DevicePtrBox>(
-            new DevicePtrBox{ptr, std::move(ds)},
+            raw,
             [h_pool](DevicePtrBox* b) {
                 GILReleaseGuard gil;
                 const DeallocationStream& stream = b->deallocation;

cuda_core/cuda/core/_cpp/rt/stream.cpp

@@ -189,9 +189,11 @@ StreamHandle create_context_bound_legacy_stream(const ContextHandle& h_context)
 namespace detail {
 // Real streams are copied unchanged. Default-stream tokens without an embedded
 // context are bound to the current context. Returns false (and sets err) when a
-// default-stream token cannot be bound because no context is current.
+// default-stream token cannot be bound because no context is current. Binding
+// a default-stream token allocates (the bound StreamBox and the current-context
+// handle), so this may throw std::bad_alloc and must not be noexcept.
 bool make_deallocation_stream(
-        const StreamHandle& h, DeallocationStream& out) noexcept {
+        const StreamHandle& h, DeallocationStream& out) {
     out = {};
     if (!h) {
         return true;

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -316,7 +316,13 @@ cdef class Buffer:
             record_stream = mr.is_device_accessible
             self._h_ptr = deviceptr_create_with_mr(c_ptr, size, mr)
             if record_stream:
-                _ds_status = set_deallocation_stream(self._h_ptr, s._h_stream)
+                try:
+                    _ds_status = set_deallocation_stream(self._h_ptr, s._h_stream)
+                except MemoryError:
+                    # Reset for the same reason as below, inside the handler so the
+                    # destructor cannot swallow the MemoryError that is propagating.
+                    self._h_ptr.reset()
+                    raise
                 if _ds_status != cydriver.CUresult.CUDA_SUCCESS:
                     # Reset before raising: the DevicePtrHandle destructor would otherwise
                     # invoke _mr_dealloc_callback, which catches any inner exception and

cuda_core/cuda/core/_rt.pxd

@@ -292,7 +292,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 cydriver.CUresult 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) except+ nogil
 
 # Virtual memory management (VMM_DESIGN.md)
 cdef MemAllocationHandle create_mem_allocation_handle(

cuda_core/cuda/core/_rt.pyx

@@ -172,7 +172,7 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt":
     StreamHandle deallocation_stream "cuda_core::rt::deallocation_stream" (
         const DevicePtrHandle& h) noexcept nogil
     cydriver.CUresult set_deallocation_stream "cuda_core::rt::set_deallocation_stream" (
-        const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil
+        const DevicePtrHandle& h, const StreamHandle& h_stream) except+ nogil
 
     # Virtual memory management (VMM_DESIGN.md)
     MemAllocationHandle create_mem_allocation_handle "cuda_core::rt::create_mem_allocation_handle" (

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

@@ -104,6 +104,10 @@ New features
 Fixes and enhancements
 ----------------------
 
+- :meth:`Buffer.set_deallocation_stream` and :meth:`Buffer.from_handle` now raise
+  :class:`MemoryError` when the host runs out of memory while binding a default
+  stream, instead of aborting the process with ``std::terminate``.
+
 - :attr:`system.Device.uuid_without_prefix` preserves UUIDs that NVML already
   reports without a prefix, such as on Orin. This also allows
   :meth:`system.Device.to_cuda_device` to match those devices correctly.