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

Author: Andy-Jost

cuda_core/cuda/core/_graph/_graph_builder.pyx

@@ -5,6 +5,8 @@
 import weakref
 from dataclasses import dataclass
 
+from libc.stdint cimport intptr_t
+
 from cuda.bindings cimport cydriver
 
 from cuda.core._graph._utils cimport _attach_host_callback_to_graph
@@ -14,6 +16,7 @@ from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
 from cuda.core._utils.version cimport cy_binding_version, cy_driver_version
 
 from cuda.core._utils.cuda_utils import (
+    CUDAError,
     driver,
     handle_return,
 )
@@ -783,24 +786,41 @@ class Graph:
         """
         return self._mnff.graph
 
-    def update(self, builder: GraphBuilder):
-        """Update the graph using new build configuration from the builder.
+    def update(self, source: "GraphBuilder | GraphDef") -> None:
+        """Update the graph using a new graph definition.
 
-        The topology of the provided builder must be identical to this graph.
+        The topology of the provided source must be identical to this graph.
 
         Parameters
         ----------
-        builder : :obj:`~_graph.GraphBuilder`
-            The builder to update the graph with.
+        source : :obj:`~_graph.GraphBuilder` or :obj:`~_graph._graphdef.GraphDef`
+            The graph definition to update from. A GraphBuilder must have
+            finished building.
 
         """
-        if not builder._building_ended:
-            raise ValueError("Graph has not finished building.")
+        from cuda.core._graph._graphdef import GraphDef
+
+        cdef cydriver.CUgraph cu_graph
+        cdef cydriver.CUgraphExec cu_exec = <cydriver.CUgraphExec><intptr_t>int(self._mnff.graph)
 
-        # Update the graph with the new nodes from the builder
-        exec_update_result = handle_return(driver.cuGraphExecUpdate(self._mnff.graph, builder._mnff.graph))
-        if exec_update_result.result != driver.CUgraphExecUpdateResult.CU_GRAPH_EXEC_UPDATE_SUCCESS:
-            raise RuntimeError(f"Failed to update graph: {exec_update_result.result()}")
+        if isinstance(source, GraphBuilder):
+            if not source._building_ended:
+                raise ValueError("Graph has not finished building.")
+            cu_graph = <cydriver.CUgraph><intptr_t>int(source._mnff.graph)
+        elif isinstance(source, GraphDef):
+            cu_graph = <cydriver.CUgraph><intptr_t>int(source.handle)
+        else:
+            raise TypeError(
+                f"expected GraphBuilder or GraphDef, got {type(source).__name__}")
+
+        cdef cydriver.CUgraphExecUpdateResultInfo result_info
+        cdef cydriver.CUresult err
+        with nogil:
+            err = cydriver.cuGraphExecUpdate(cu_exec, cu_graph, &result_info)
+        if err != cydriver.CUresult.CUDA_SUCCESS:
+            reason = driver.CUgraphExecUpdateResult(result_info.result)
+            msg = f"Graph update failed: {reason.__doc__.strip()} ({reason.name})"
+            raise CUDAError(msg)
 
     def upload(self, stream: Stream):
         """Uploads the graph in a stream.

cuda_core/tests/graph/test_graph_update.py

@@ -5,17 +5,66 @@
 
 import numpy as np
 import pytest
-from helpers.graph_kernels import compile_conditional_kernels
+from helpers.graph_kernels import compile_common_kernels, compile_conditional_kernels
 
 from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, launch
+from cuda.core._graph._graphdef import GraphDef
+from cuda.core._utils.cuda_utils import CUDAError
 
 
+@pytest.mark.parametrize("builder", ["GraphBuilder", "GraphDef"])
 @pytest.mark.skipif(tuple(int(i) for i in np.__version__.split(".")[:2]) < (2, 1), reason="need numpy 2.1.0+")
-def test_graph_update(init_cuda):
+def test_graph_update_kernel_args(init_cuda, builder):
+    """Update redirects a kernel to write to a different pointer."""
+    mod = compile_common_kernels()
+    add_one = mod.get_kernel("add_one")
+
+    launch_stream = Device().create_stream()
+    mr = LegacyPinnedMemoryResource()
+    b = mr.allocate(8)
+    arr = np.from_dlpack(b).view(np.int32)
+    arr[0] = 0
+    arr[1] = 0
+
+    if builder == "GraphBuilder":
+
+        def build(ptr):
+            gb = Device().create_graph_builder().begin_building()
+            launch(gb, LaunchConfig(grid=1, block=1), add_one, ptr)
+            launch(gb, LaunchConfig(grid=1, block=1), add_one, ptr)
+            finished = gb.end_building()
+            return finished.complete(), finished
+    elif builder == "GraphDef":
+
+        def build(ptr):
+            g = GraphDef()
+            g.launch(LaunchConfig(grid=1, block=1), add_one, ptr)
+            g.launch(LaunchConfig(grid=1, block=1), add_one, ptr)
+            return g.instantiate(), g
+
+    graph, _ = build(arr[0:].ctypes.data)
+    _, source1 = build(arr[1:].ctypes.data)
+
+    graph.launch(launch_stream)
+    launch_stream.sync()
+    assert arr[0] == 2
+    assert arr[1] == 0
+
+    graph.update(source1)
+    graph.launch(launch_stream)
+    launch_stream.sync()
+    assert arr[0] == 2
+    assert arr[1] == 2
+
+    b.close()
+
+
+@pytest.mark.skipif(tuple(int(i) for i in np.__version__.split(".")[:2]) < (2, 1), reason="need numpy 2.1.0+")
+def test_graph_update_conditional(init_cuda):
+    """Update swaps conditional switch graphs with matching topology."""
     mod = compile_conditional_kernels(int)
     add_one = mod.get_kernel("add_one")
 
-    # Allocate memory
     launch_stream = Device().create_stream()
     mr = LegacyPinnedMemoryResource()
     b = mr.allocate(12)
@@ -72,9 +121,6 @@ def build_graph(condition_value):
         pytest.skip("Driver does not support conditional switch")
 
     # Launch the first graph
-    assert arr[0] == 0
-    assert arr[1] == 0
-    assert arr[2] == 0
     graph = graph_variants[0].complete()
     graph.launch(launch_stream)
     launch_stream.sync()
@@ -98,4 +144,65 @@ def build_graph(condition_value):
     assert arr[1] == 3
     assert arr[2] == 3
 
+    # Close the memory resource now because the garbage collected might
+    # de-allocate it during the next graph builder process
     b.close()
+
+
+# =============================================================================
+# Error cases
+# =============================================================================
+
+
+def test_graph_update_unfinished_builder(init_cuda):
+    """Update with an unfinished GraphBuilder raises ValueError."""
+    mod = compile_common_kernels()
+    empty_kernel = mod.get_kernel("empty_kernel")
+
+    gb_finished = Device().create_graph_builder().begin_building()
+    launch(gb_finished, LaunchConfig(grid=1, block=1), empty_kernel)
+    graph = gb_finished.end_building().complete()
+
+    gb_unfinished = Device().create_graph_builder().begin_building()
+    launch(gb_unfinished, LaunchConfig(grid=1, block=1), empty_kernel)
+
+    with pytest.raises(ValueError, match="Graph has not finished building"):
+        graph.update(gb_unfinished)
+
+    gb_unfinished.end_building()
+
+
+def test_graph_update_topology_mismatch(init_cuda):
+    """Update with a different topology raises CUDAError."""
+    mod = compile_common_kernels()
+    empty_kernel = mod.get_kernel("empty_kernel")
+
+    # Two-node graph
+    gb1 = Device().create_graph_builder().begin_building()
+    launch(gb1, LaunchConfig(grid=1, block=1), empty_kernel)
+    launch(gb1, LaunchConfig(grid=1, block=1), empty_kernel)
+    graph = gb1.end_building().complete()
+
+    # Three-node graph (different topology)
+    gb2 = Device().create_graph_builder().begin_building()
+    launch(gb2, LaunchConfig(grid=1, block=1), empty_kernel)
+    launch(gb2, LaunchConfig(grid=1, block=1), empty_kernel)
+    launch(gb2, LaunchConfig(grid=1, block=1), empty_kernel)
+    gb2.end_building()
+
+    expected = r"Graph update failed: The update failed because the topology changed \(CU_GRAPH_EXEC_UPDATE_ERROR_TOPOLOGY_CHANGED\)"
+    with pytest.raises(CUDAError, match=expected):
+        graph.update(gb2)
+
+
+def test_graph_update_wrong_type(init_cuda):
+    """Update with an invalid type raises TypeError."""
+    mod = compile_common_kernels()
+    empty_kernel = mod.get_kernel("empty_kernel")
+
+    gb = Device().create_graph_builder().begin_building()
+    launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
+    graph = gb.end_building().complete()
+
+    with pytest.raises(TypeError, match="expected GraphBuilder or GraphDef"):
+        graph.update("not a graph")