cuda.core: defer graph user-object payload cleanup (NVIDIA#2371)

* cuda.core: defer CUDA user-object payload reclamation

Move graph attachment destruction off CUDA's internal callback thread through a coalesced Py_AddPendingCall drain. Queue payloads lock-free, retry scheduling failures from later safe entries, and leak intact payloads once Python finalization begins.

* test(cuda.core): stress deferred graph reclamation

Verify more than 32 CUDA callbacks coalesce onto Python's main thread, recover after Py_AddPendingCall queue saturation, and remain safe during interpreter shutdown.

* docs(cuda.core): clarify deferred reclaimer roles

Add concise descriptions for each reclaimer operation and state field, and explain the typed synchronous helper versus CUDA's void-pointer callback ABI.

* refactor(cuda.core): simplify deferred queue envelopes

Use the process-wide reclaimer directly instead of storing the same pointer in every envelope, publish the singleton atomically, and make the intrusive next link non-atomic under the queue's release/acquire handoff.

* refactor(cuda.core): unify deferred cleanup terminology

Use payload, cleanup item, cleanup queue, enqueue, drain, and cleanup consistently across implementation, documentation, and tests.

* test(cuda.core): isolate shutdown subprocess

Run the interpreter-shutdown probe outside the source tree so CI imports the installed package with generated modules.

* test(cuda.core): remove pending-call timing race

Validate queue-saturation safety without assuming when CUDA invokes its asynchronous user-object destructor.

Author: Andy-Jost

cuda_core/cuda/core/_cpp/DESIGN.md

@@ -217,6 +217,21 @@ Handle destructors may run from any thread. The implementation includes RAII gua
 The handle API functions are safe to call with or without the GIL held. They
 will release the GIL (if necessary) before calling CUDA driver API functions.
 
+### CUDA User-Object Cleanup
+
+CUDA user-object destructors may run on an internal driver thread where CUDA
+API calls are forbidden. Graph attachment payloads can release handles whose
+deleters call CUDA or run Python finalizers, so their CUDA callback must not
+delete the payload directly.
+
+The callback enqueues a preallocated `DeferredCleanupItem` on a
+process-lifetime lock-free `DeferredCleanupQueue` and schedules one coalesced
+`Py_AddPendingCall`. The pending callback drains all queued payloads from
+Python's main thread. If scheduling fails, payloads remain intact for a later
+retry; they are intentionally leaked once interpreter finalization begins. A
+single pending call is shared by all queued payloads because CPython's
+main-thread pending-call queue is bounded.
+
 ### Static Initialization and Deadlock Hazards
 
 When writing C++ code that interacts with Python, a subtle deadlock can occur

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -6,6 +6,7 @@
 
 #include "resource_handles.hpp"
 #include <cuda.h>
+#include <atomic>
 #include <array>
 #include <cstdint>
 #include <cstdlib>
@@ -185,6 +186,162 @@ class GILAcquireGuard {
 
 }  // namespace
 
+// ============================================================================
+// CUDA user-object deferred cleanup
+//
+// CUDA may invoke a user-object destructor on an internal thread where CUDA
+// calls are forbidden. Payload cleanup can release resource handles whose
+// deleters call CUDA, so the callback only transfers a preallocated intrusive
+// node to this process-lifetime queue. One coalesced pending call drains all
+// queued payloads from Python's main thread.
+// ============================================================================
+
+namespace {
+
+class DeferredCleanupQueue;
+
+using CleanupFn = void (*)(void*) noexcept;
+
+// Preallocated queue item that transfers one payload out of CUDA's callback.
+struct DeferredCleanupItem {
+    DeferredCleanupItem* next = nullptr;
+    void* payload;
+    CleanupFn cleanup;
+
+    DeferredCleanupItem(void* payload_, CleanupFn cleanup_) noexcept
+        : payload(payload_), cleanup(cleanup_) {}
+};
+
+// Process-lifetime MPSC queue that drains payloads from Python's main thread.
+class DeferredCleanupQueue {
+public:
+    // Transfer one preallocated cleanup item from a producer to the queue.
+    void enqueue(DeferredCleanupItem* item) noexcept {
+        DeferredCleanupItem* head = head_.load(std::memory_order_relaxed);
+        do {
+            item->next = head;
+        } while (!head_.compare_exchange_weak(
+            head, item, std::memory_order_release, std::memory_order_relaxed));
+        schedule();
+    }
+
+    // Permanently disable pending-call scheduling during interpreter shutdown.
+    void stop() noexcept {
+        accepting_.store(false, std::memory_order_release);
+    }
+
+    // Reattempt scheduling for payloads left queued after an earlier failure.
+    void retry_schedule() noexcept {
+        schedule();
+    }
+
+private:
+    // Adapt queue draining to CPython's int (*)(void*) callback ABI.
+    static int pending_call(void* arg) noexcept {
+        static_cast<DeferredCleanupQueue*>(arg)->drain();
+        return 0;
+    }
+
+    // Coalesce all queued work behind at most one CPython pending call.
+    void schedule() noexcept {
+        if (!accepting_.load(std::memory_order_acquire) ||
+            !head_.load(std::memory_order_acquire)) {
+            return;
+        }
+        bool expected = false;
+        if (!scheduled_.compare_exchange_strong(
+                expected, true, std::memory_order_acq_rel,
+                std::memory_order_relaxed)) {
+            return;
+        }
+        if (Py_AddPendingCall(&DeferredCleanupQueue::pending_call, this) != 0) {
+            // Keep every payload queued. A later enqueue or safe cuda-core
+            // entry can retry without blocking CUDA's callback thread.
+            scheduled_.store(false, std::memory_order_release);
+        }
+    }
+
+    // Detach and destroy all queued payloads from Python's main thread.
+    void drain() noexcept {
+        if (!Py_IsInitialized() || py_is_finalizing()) {
+            stop();
+            scheduled_.store(false, std::memory_order_release);
+            return;  // Intentionally leak intact payloads during shutdown.
+        }
+
+        while (DeferredCleanupItem* list =
+                   head_.exchange(nullptr, std::memory_order_acquire)) {
+            while (list) {
+                DeferredCleanupItem* next = list->next;
+                list->cleanup(list->payload);
+                delete list;
+                list = next;
+            }
+        }
+
+        scheduled_.store(false, std::memory_order_release);
+        if (head_.load(std::memory_order_acquire)) {
+            schedule();
+        }
+    }
+
+    // Head of the intrusive multi-producer, single-consumer payload stack.
+    std::atomic<DeferredCleanupItem*> head_{nullptr};
+    // True while one cuda-core drain callback is pending or executing.
+    std::atomic<bool> scheduled_{false};
+    // False once shutdown begins, causing later payloads to be leaked safely.
+    std::atomic<bool> accepting_{true};
+};
+
+// Published once at module initialization and intentionally never freed.
+std::atomic<DeferredCleanupQueue*> deferred_cleanup_queue{nullptr};
+
+void stop_deferred_cleanup() {
+    if (DeferredCleanupQueue* queue =
+            deferred_cleanup_queue.load(std::memory_order_acquire)) {
+        queue->stop();
+    }
+}
+
+DeferredCleanupItem* make_cleanup_item(
+        void* payload, CleanupFn cleanup) {
+    DeferredCleanupQueue* queue =
+        deferred_cleanup_queue.load(std::memory_order_acquire);
+    if (!queue) {
+        throw std::runtime_error("deferred cleanup is not initialized");
+    }
+    queue->retry_schedule();
+    return new DeferredCleanupItem(payload, cleanup);
+}
+
+// Execute an item's cleanup synchronously before CUDA has acquired ownership.
+void execute_cleanup_now(DeferredCleanupItem* item) noexcept {
+    item->cleanup(item->payload);
+    delete item;
+}
+
+// CUDA's CUhostFn ABI is void (*)(void*); recover and enqueue the cleanup item.
+void enqueue_cleanup(void* item) noexcept {
+    auto* cleanup = static_cast<DeferredCleanupItem*>(item);
+    if (DeferredCleanupQueue* queue =
+            deferred_cleanup_queue.load(std::memory_order_acquire)) {
+        queue->enqueue(cleanup);
+    }
+}
+
+}  // namespace
+
+void initialize_deferred_cleanup() {
+    if (deferred_cleanup_queue.load(std::memory_order_acquire)) {
+        return;
+    }
+    auto* queue = new DeferredCleanupQueue();
+    deferred_cleanup_queue.store(queue, std::memory_order_release);
+    if (Py_AtExit(stop_deferred_cleanup) != 0) {
+        queue->stop();
+    }
+}
+
 // ============================================================================
 // Handle reverse-lookup registry
 //
@@ -1146,7 +1303,7 @@ void free_deleter(const void* p) noexcept {
     std::free(const_cast<void*>(p));
 }
 
-void destroy_graph_slot_table(void* table) noexcept {
+void cleanup_graph_slot_table(void* table) noexcept {
     delete static_cast<GraphSlotTable*>(table);
 }
 
@@ -1176,13 +1333,20 @@ GraphSlotTable* ensure_slot_table(const GraphBox* box) {
         return nullptr;
     }
     auto* table = new GraphSlotTable();
+    DeferredCleanupItem* cleanup_item = nullptr;
+    try {
+        cleanup_item = make_cleanup_item(table, cleanup_graph_slot_table);
+    } catch (...) {
+        delete table;
+        throw;
+    }
     CUuserObject user_obj = nullptr;
     {
         GILReleaseGuard gil;
-        if (p_cuUserObjectCreate(&user_obj, table,
-                                 reinterpret_cast<CUhostFn>(destroy_graph_slot_table),
+        if (p_cuUserObjectCreate(&user_obj, cleanup_item,
+                                 reinterpret_cast<CUhostFn>(enqueue_cleanup),
                                  1, CU_USER_OBJECT_NO_DESTRUCTOR_SYNC) != CUDA_SUCCESS) {
-            delete table;  // no user object created; nothing else owns the table
+            execute_cleanup_now(cleanup_item);
             return nullptr;
         }
         if (p_cuGraphRetainUserObject(box->resource, user_obj, 1,

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -256,6 +256,10 @@ StreamHandle create_stream_handle_with_owner(CUstream stream, PyObject* owner);
 // If Python is finalized or finalizing, the object is intentionally leaked.
 void py_object_user_object_destroy(void* py_object) noexcept;
 
+// Initialize the process-lifetime CUDA user-object cleanup queue. Called once
+// from module initialization while Python is fully initialized.
+void initialize_deferred_cleanup();
+
 // Return the context dependency associated with a stream handle, if any.
 ContextHandle get_stream_context(const StreamHandle& h) noexcept;
 

cuda_core/cuda/core/_resource_handles.pyx

@@ -61,6 +61,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
         cydriver.CUstream stream, object owner) except+ nogil
     void py_object_user_object_destroy "cuda_core::py_object_user_object_destroy" (
         void* py_object) noexcept nogil
+    void initialize_deferred_cleanup "cuda_core::initialize_deferred_cleanup" () except+
     ContextHandle get_stream_context "cuda_core::get_stream_context" (
         const StreamHandle& h) noexcept nogil
     StreamHandle get_legacy_stream "cuda_core::get_legacy_stream" () except+ nogil
@@ -461,6 +462,7 @@ cdef void _init_driver_fn_pointers() noexcept:
     p_cuDevSmResourceSplit = _get_optional_driver_fn("cuDevSmResourceSplit")
 
 _init_driver_fn_pointers()
+initialize_deferred_cleanup()
 
 # =============================================================================
 # NVRTC function pointer initialization