Add GraphNode identity cache for stable object round-trips (#1853)

* Reorganize graph test files for clarity

Rename test files to reflect what they actually test:
- test_basic -> test_graph_builder (stream capture tests)
- test_conditional -> test_graph_builder_conditional
- test_advanced -> test_graph_update (moved child_graph and
  stream_lifetime tests into test_graph_builder)
- test_capture_alloc -> test_graph_memory_resource
- test_explicit* -> test_graphdef*

Made-with: Cursor

* Enhance Graph.update() and add whole-graph update tests

- Extend Graph.update() to accept both GraphBuilder and GraphDef sources
- Surface CUgraphExecUpdateResultInfo details on update failure instead
  of a generic CUDA_ERROR_GRAPH_EXEC_UPDATE_FAILURE message
- Release the GIL during cuGraphExecUpdate via nogil block
- Add parametrized happy-path test covering both GraphBuilder and GraphDef
- Add error-case tests: unfinished builder, topology mismatch, wrong type

Made-with: Cursor

* Add AdjacencySet proxy for pred/succ and GraphNode.remove()

Replace cached tuple-based pred/succ with mutable AdjacencySet backed
by direct CUDA driver calls. Add GraphNode.remove() wrapping
cuGraphDestroyNode.

Made-with: Cursor

* Add edge mutation support and MutableSet interface for GraphNode adjacencies

Enable adding/removing edges between graph nodes via AdjacencySet (a
MutableSet proxy on GraphNode.pred/succ), node removal via discard(),
and property setters for bulk edge replacement. Includes comprehensive
mutation and interface tests.

Closes part of #1330 (step 2: edge mutation on GraphDef).

Made-with: Cursor

* Use requires_module mark for numpy version checks in mutation tests

Replace inline skipif version check with requires_module(np, "2.1")
from the shared test helpers, consistent with other test files.

Made-with: Cursor

* Fix empty-graph return type: return set() instead of () for nodes/edges

Made-with: Cursor

* Rename AdjacencySet to AdjacencySetProxy, add bulk ops and safety guards

Rename class and file to AdjacencySetProxy to clarify write-through
semantics. Add bulk-efficient clear(), __isub__(), __ior__() overrides
and remove_edges() on the Cython core. Guard GraphNode.discard() against
double-destroy via membership check. Filter duplicates in update(). Add
error-path tests for wrong types, cross-graph edges, and self-edges.

Made-with: Cursor

* Add destroy() method with handle invalidation, remove GRAPH_NODE_SENTINEL

Replace discard() with destroy() which calls cuGraphDestroyNode and then
zeroes the CUgraphNode resource in the handle box via
invalidate_graph_node_handle. This prevents stale memory access on
destroyed nodes. Properties (type, pred, succ, handle) degrade gracefully
to None/empty for destroyed nodes.

Remove the GRAPH_NODE_SENTINEL (0x1) approach in favor of using NULL for
both sentinels and destroyed nodes, which is simpler and avoids the risk
of passing 0x1 to driver APIs that treat it as a valid pointer.

Made-with: Cursor

* Add GraphNode identity cache for stable Python object round-trips

Nodes retrieved via GraphDef.nodes(), edges(), or pred/succ traversal
now return the same Python object that was originally created, enabling
identity checks with `is`. A C++ HandleRegistry deduplicates
CUgraphNode handles, and a Cython WeakValueDictionary caches the
Python wrapper objects.

Made-with: Cursor

* Purge node cache on destroy to prevent stale identity lookups

Made-with: Cursor

* Skip NULL nodes in graph_node_registry to fix sentinel identity collision

Sentinel (entry) nodes use NULL as their CUgraphNode, so caching them
under a NULL key caused all sentinels across different graphs to share
the same handle. This made nodes built from the wrong graph's entry
point, causing CUDA_ERROR_INVALID_VALUE for conditional nodes and hash
collisions in equality tests.

Made-with: Cursor

* Unregister destroyed nodes from C++ graph_node_registry

When a node is destroyed, the driver may reuse its CUgraphNode pointer
for a new node. Without unregistering the old entry, the registry
returns a stale handle pointing to the wrong node type and graph.

Made-with: Cursor

* Add dedicated test for node identity preservation through round-trips

Made-with: Cursor

* Rename _node_cache/_cached to _node_registry/_registered

Aligns Python-side terminology with the C++ graph_node_registry.

Made-with: Cursor

* Fix unregister_handle and rename invalidate_graph_node_handle

unregister_handle: remove the expired() guard that prevented erasure
when the shared_ptr was still alive. This caused stale registry
entries after destroy(), leading to CUDA_ERROR_INVALID_VALUE when
the driver reused CUgraphNode pointer values.

Rename invalidate_graph_node_handle -> invalidate_graph_node for
consistency with the rest of the graph node API.

Made-with: Cursor

Author: Andy-Jost

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -174,13 +174,8 @@ class HandleRegistry {
     }
 
     void unregister_handle(const Key& key) noexcept {
-        try {
-            std::lock_guard<std::mutex> lock(mutex_);
-            auto it = map_.find(key);
-            if (it != map_.end() && it->second.expired()) {
-                map_.erase(it);
-            }
-        } catch (...) {}
+        std::lock_guard<std::mutex> lock(mutex_);
+        map_.erase(key);
     }
 
     Handle lookup(const Key& key) {
@@ -969,17 +964,32 @@ static const GraphNodeBox* get_box(const GraphNodeHandle& h) {
     );
 }
 
+static HandleRegistry<CUgraphNode, GraphNodeHandle> graph_node_registry;
+
 GraphNodeHandle create_graph_node_handle(CUgraphNode node, const GraphHandle& h_graph) {
+    if (node) {
+        if (auto h = graph_node_registry.lookup(node)) {
+            return h;
+        }
+    }
     auto box = std::make_shared<const GraphNodeBox>(GraphNodeBox{node, h_graph});
-    return GraphNodeHandle(box, &box->resource);
+    GraphNodeHandle h(box, &box->resource);
+    if (node) {
+        graph_node_registry.register_handle(node, h);
+    }
+    return h;
 }
 
 GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept {
     return h ? get_box(h)->h_graph : GraphHandle{};
 }
 
-void invalidate_graph_node_handle(const GraphNodeHandle& h) noexcept {
+void invalidate_graph_node(const GraphNodeHandle& h) noexcept {
     if (h) {
+        CUgraphNode node = get_box(h)->resource;
+        if (node) {
+            graph_node_registry.unregister_handle(node);
+        }
         get_box(h)->resource = nullptr;
     }
 }

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -416,7 +416,7 @@ GraphNodeHandle create_graph_node_handle(CUgraphNode node, const GraphHandle& h_
 GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept;
 
 // Zero the CUgraphNode resource inside the handle, marking it invalid.
-void invalidate_graph_node_handle(const GraphNodeHandle& h) noexcept;
+void invalidate_graph_node(const GraphNodeHandle& h) noexcept;
 
 // ============================================================================
 // Graphics resource handle functions

cuda_core/cuda/core/_graph/_graph_def/_graph_node.pyx

@@ -48,7 +48,7 @@ from cuda.core._resource_handles cimport (
     create_graph_handle_ref,
     create_graph_node_handle,
     graph_node_get_graph,
-    invalidate_graph_node_handle,
+    invalidate_graph_node,
 )
 from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value
 
@@ -57,10 +57,19 @@ from cuda.core._graph._utils cimport (
     _attach_user_object,
 )
 
+import weakref
+
 from cuda.core import Device
 from cuda.core._graph._graph_def._adjacency_set_proxy import AdjacencySetProxy
 from cuda.core._utils.cuda_utils import driver, handle_return
 
+_node_registry = weakref.WeakValueDictionary()
+
+
+cdef inline GraphNode _registered(GraphNode n):
+    _node_registry[<uintptr_t>n._h_node.get()] = n
+    return n
+
 
 cdef class GraphNode:
     """Base class for all graph nodes.
@@ -144,7 +153,8 @@ cdef class GraphNode:
             return
         with nogil:
             HANDLE_RETURN(cydriver.cuGraphDestroyNode(node))
-        invalidate_graph_node_handle(self._h_node)
+        _node_registry.pop(<uintptr_t>self._h_node.get(), None)
+        invalidate_graph_node(self._h_node)
 
     @property
     def pred(self):
@@ -522,18 +532,29 @@ cdef inline ConditionalNode _make_conditional_node(
     n._cond_type = cond_type
     n._branches = branches
 
-    return n
+    return _registered(n)
 
 cdef inline GraphNode GN_create(GraphHandle h_graph, cydriver.CUgraphNode node):
+    cdef GraphNodeHandle h_node = create_graph_node_handle(node, h_graph)
+
+    # Sentinel: virtual node to represent the graph entry point.
     if node == NULL:
         n = GraphNode.__new__(GraphNode)
-        (<GraphNode>n)._h_node = create_graph_node_handle(node, h_graph)
+        (<GraphNode>n)._h_node = h_node
         return n
 
-    cdef GraphNodeHandle h_node = create_graph_node_handle(node, h_graph)
+    # Return a registered object or create and register a new one.
+    registered = _node_registry.get(<uintptr_t>h_node.get())
+    if registered is not None:
+        return <GraphNode>registered
+    else:
+        return _registered(GN_create_impl(h_node))
+
+
+cdef inline GraphNode GN_create_impl(GraphNodeHandle h_node):
     cdef cydriver.CUgraphNodeType node_type
     with nogil:
-        HANDLE_RETURN(cydriver.cuGraphNodeGetType(node, &node_type))
+        HANDLE_RETURN(cydriver.cuGraphNodeGetType(as_cu(h_node), &node_type))
 
     if node_type == cydriver.CU_GRAPH_NODE_TYPE_EMPTY:
         return EmptyNode._create_impl(h_node)
@@ -595,10 +616,10 @@ cdef inline KernelNode GN_launch(GraphNode self, LaunchConfig conf, Kernel ker,
     _attach_user_object(as_cu(h_graph), <void*>new KernelHandle(ker._h_kernel),
                         <cydriver.CUhostFn>_destroy_kernel_handle_copy)
 
-    return KernelNode._create_with_params(
+    return _registered(KernelNode._create_with_params(
         create_graph_node_handle(new_node, h_graph),
         conf.grid, conf.block, conf.shmem_size,
-        ker._h_kernel)
+        ker._h_kernel))
 
 
 cdef inline EmptyNode GN_join(GraphNode self, tuple nodes):
@@ -624,7 +645,7 @@ cdef inline EmptyNode GN_join(GraphNode self, tuple nodes):
         HANDLE_RETURN(cydriver.cuGraphAddEmptyNode(
             &new_node, as_cu(h_graph), deps_ptr, num_deps))
 
-    return EmptyNode._create_impl(create_graph_node_handle(new_node, h_graph))
+    return _registered(EmptyNode._create_impl(create_graph_node_handle(new_node, h_graph)))
 
 
 cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object options):
@@ -700,9 +721,9 @@ cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object options):
         HANDLE_RETURN(cydriver.cuGraphAddMemAllocNode(
             &new_node, as_cu(h_graph), deps, num_deps, &alloc_params))
 
-    return AllocNode._create_with_params(
+    return _registered(AllocNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), alloc_params.dptr, size,
-        device_id, memory_type, tuple(peer_ids))
+        device_id, memory_type, tuple(peer_ids)))
 
 
 cdef inline FreeNode GN_free(GraphNode self, cydriver.CUdeviceptr c_dptr):
@@ -720,7 +741,7 @@ cdef inline FreeNode GN_free(GraphNode self, cydriver.CUdeviceptr c_dptr):
         HANDLE_RETURN(cydriver.cuGraphAddMemFreeNode(
             &new_node, as_cu(h_graph), deps, num_deps, c_dptr))
 
-    return FreeNode._create_with_params(create_graph_node_handle(new_node, h_graph), c_dptr)
+    return _registered(FreeNode._create_with_params(create_graph_node_handle(new_node, h_graph), c_dptr))
 
 
 cdef inline MemsetNode GN_memset(
@@ -755,9 +776,9 @@ cdef inline MemsetNode GN_memset(
             &new_node, as_cu(h_graph), deps, num_deps,
             &memset_params, ctx))
 
-    return MemsetNode._create_with_params(
+    return _registered(MemsetNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), c_dst,
-        val, elem_size, width, height, pitch)
+        val, elem_size, width, height, pitch))
 
 
 cdef inline MemcpyNode GN_memcpy(
@@ -816,9 +837,9 @@ cdef inline MemcpyNode GN_memcpy(
         HANDLE_RETURN(cydriver.cuGraphAddMemcpyNode(
             &new_node, as_cu(h_graph), deps, num_deps, &params, ctx))
 
-    return MemcpyNode._create_with_params(
+    return _registered(MemcpyNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), c_dst, c_src, size,
-        c_dst_type, c_src_type)
+        c_dst_type, c_src_type))
 
 
 cdef inline ChildGraphNode GN_embed(GraphNode self, GraphDef child_def):
@@ -843,8 +864,8 @@ cdef inline ChildGraphNode GN_embed(GraphNode self, GraphDef child_def):
 
     cdef GraphHandle h_embedded = create_graph_handle_ref(embedded_graph, h_graph)
 
-    return ChildGraphNode._create_with_params(
-        create_graph_node_handle(new_node, h_graph), h_embedded)
+    return _registered(ChildGraphNode._create_with_params(
+        create_graph_node_handle(new_node, h_graph), h_embedded))
 
 
 cdef inline EventRecordNode GN_record_event(GraphNode self, Event ev):
@@ -865,8 +886,8 @@ cdef inline EventRecordNode GN_record_event(GraphNode self, Event ev):
     _attach_user_object(as_cu(h_graph), <void*>new EventHandle(ev._h_event),
                         <cydriver.CUhostFn>_destroy_event_handle_copy)
 
-    return EventRecordNode._create_with_params(
-        create_graph_node_handle(new_node, h_graph), ev._h_event)
+    return _registered(EventRecordNode._create_with_params(
+        create_graph_node_handle(new_node, h_graph), ev._h_event))
 
 
 cdef inline EventWaitNode GN_wait_event(GraphNode self, Event ev):
@@ -887,8 +908,8 @@ cdef inline EventWaitNode GN_wait_event(GraphNode self, Event ev):
     _attach_user_object(as_cu(h_graph), <void*>new EventHandle(ev._h_event),
                         <cydriver.CUhostFn>_destroy_event_handle_copy)
 
-    return EventWaitNode._create_with_params(
-        create_graph_node_handle(new_node, h_graph), ev._h_event)
+    return _registered(EventWaitNode._create_with_params(
+        create_graph_node_handle(new_node, h_graph), ev._h_event))
 
 
 cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_data):
@@ -914,6 +935,6 @@ cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_
             &new_node, as_cu(h_graph), deps, num_deps, &node_params))
 
     cdef object callable_obj = fn if not isinstance(fn, ct._CFuncPtr) else None
-    return HostCallbackNode._create_with_params(
+    return _registered(HostCallbackNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), callable_obj,
-        node_params.fn, node_params.userData)
+        node_params.fn, node_params.userData))

cuda_core/cuda/core/_resource_handles.pxd

@@ -186,7 +186,7 @@ cdef GraphHandle create_graph_handle_ref(cydriver.CUgraph graph, const GraphHand
 # Graph node handles
 cdef GraphNodeHandle create_graph_node_handle(cydriver.CUgraphNode node, const GraphHandle& h_graph) except+ nogil
 cdef GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept nogil
-cdef void invalidate_graph_node_handle(const GraphNodeHandle& h) noexcept nogil
+cdef void invalidate_graph_node(const GraphNodeHandle& h) noexcept nogil
 
 # Graphics resource handles
 cdef GraphicsResourceHandle create_graphics_resource_handle(

cuda_core/cuda/core/_resource_handles.pyx

@@ -159,7 +159,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
         cydriver.CUgraphNode node, const GraphHandle& h_graph) except+ nogil
     GraphHandle graph_node_get_graph "cuda_core::graph_node_get_graph" (
         const GraphNodeHandle& h) noexcept nogil
-    void invalidate_graph_node_handle "cuda_core::invalidate_graph_node_handle" (
+    void invalidate_graph_node "cuda_core::invalidate_graph_node" (
         const GraphNodeHandle& h) noexcept nogil
 
     # Graphics resource handles

cuda_core/tests/graph/test_graphdef.py

@@ -661,6 +661,7 @@ def test_node_type_preserved_by_nodes(node_spec):
     matched = [n for n in all_nodes if n == node]
     assert len(matched) == 1
     assert isinstance(matched[0], spec.roundtrip_class)
+    assert matched[0] is node
 
 
 def test_node_type_preserved_by_pred_succ(node_spec):
@@ -670,6 +671,7 @@ def test_node_type_preserved_by_pred_succ(node_spec):
         matched = [s for s in predecessor.succ if s == node]
         assert len(matched) == 1
         assert isinstance(matched[0], spec.roundtrip_class)
+        assert matched[0] is node
 
 
 def test_node_attrs(node_spec):
@@ -697,6 +699,31 @@ def test_node_attrs_preserved_by_nodes(node_spec):
         assert getattr(retrieved, attr) == getattr(node, attr), f"{spec.name}.{attr} not preserved by nodes()"
 
 
+def test_identity_preservation(init_cuda):
+    """Round-trips through nodes(), edges(), and pred/succ return extant
+    objects rather than duplicates."""
+    g = GraphDef()
+    a = g.join()
+    b = a.join()
+
+    # nodes()
+    assert any(x is a for x in g.nodes())
+    assert any(x is b for x in g.nodes())
+
+    # succ/pred
+    a.succ = {b}
+    (b2,) = a.succ
+    assert b2 is b
+
+    (a2,) = b.pred
+    assert a2 is a
+
+    # edges()
+    ((a2, b2),) = g.edges()
+    assert a2 is a
+    assert b2 is b
+
+
 # =============================================================================
 # GraphDef basics
 # =============================================================================

cuda_core/tests/graph/test_graphdef_mutation.py

@@ -380,7 +380,7 @@ def test_convert_linear_to_fan_in(init_cuda):
     for node in g.nodes():
         if isinstance(node, MemsetNode):
             node.pred.clear()
-        elif isinstance(node, KernelNode) and node != reduce_node:
+        elif isinstance(node, KernelNode) and node is not reduce_node:
             node.succ.add(reduce_node)
 
     assert len(g.edges()) == 8