Self review

Author: vzhurba01

cuda_core/cuda/core/experimental/_graph.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
 #
 # SPDX-License-Identifier: Apache-2.0
 
@@ -84,7 +84,7 @@ class CompleteOptions:
     auto_free_on_launch : bool, optional
         Automatically free memory allocated in a graph before relaunching. (Default to False)
     upload_stream : Stream, optional
-        Automatically upload the graph after instantiation. (Default to None)
+        Stream to use to automatically upload the graph after completion. (Default to None)
     device_launch : bool, optional
         Configure the graph to be launchable from the device. This flag can only
         be used on platforms which support unified addressing. This flag cannot be
@@ -96,7 +96,7 @@ class CompleteOptions:
     """
 
     auto_free_on_launch: bool = False
-    upload_stream: bool = None
+    upload_stream: Optional[Stream] = None
     device_launch: bool = False
     use_node_priority: bool = False
 
@@ -138,7 +138,7 @@ def close(self):
     def __init__(self):
         raise NotImplementedError(
             "directly creating a Graph object can be ambiguous. Please either "
-            "call Device.create_graph() or stream.creating_graph()"
+            "call Device.create_graph_builder() or stream.creating_graph_builder()"
         )
 
     @classmethod
@@ -190,15 +190,15 @@ def is_building(self) -> bool:
         elif capture_status == driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_INVALIDATED:
             self.end_building()
             raise RuntimeError(
-                "Build process encountered an error and has been invalidated. Build process has now been ended."
+                "Build process encountered an error and has been invalidated. Build process has been ended."
             )
         else:
             raise NotImplementedError(f"Unsupported capture status type received: {capture_status}")
 
     def end_building(self) -> GraphBuilder:
         """Ends the building process."""
         if not self.is_building:
-            raise RuntimeError("Graph builder was not building.")
+            raise RuntimeError("Graph builder is not building.")
         self._mnff.graph = handle_return(driver.cuStreamEndCapture(self.stream.handle))
 
         # TODO: Resolving https://github.com/NVIDIA/cuda-python/issues/617 would allow us to
@@ -256,7 +256,7 @@ def complete(self, options: Optional[CompleteOptions] = None) -> Graph:
         elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_CONDITIONAL_HANDLE_UNUSED:
             raise RuntimeError("One or more conditional handles are not associated with conditional builders.")
         elif params.result_out != driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_SUCCESS:
-            raise RuntimeError("Graph instantiation failed for an unexpected reason.")
+            raise RuntimeError(f"Graph instantiation failed with unexpected error code: {params.result_out}")
         return graph
 
     def debug_dot_print(self, path, options: Optional[DebugPrintOptions] = None):
@@ -265,7 +265,7 @@ def debug_dot_print(self, path, options: Optional[DebugPrintOptions] = None):
         Parameters
         ----------
         path : str
-            The path to the file to write the DOT debug output to.
+            File path to use for writting debug DOT output
         options : :obj:`~_graph.DebugPrintOptions`, optional
             Customizable dataclass for the debug print options.
 
@@ -313,7 +313,7 @@ def split(self, count) -> Tuple[GraphBuilder, ...]:
         """Splits the original graph builder into multiple graph builders.
 
         The new builders inherit work dependencies from the original builder.
-        The original builder is reused for the split, returned first in the tuple.
+        The original builder is reused for the split and is returned first in the tuple.
 
         Parameters
         ----------
@@ -407,7 +407,7 @@ def create_conditional_handle(self, default_value=None) -> int:
             The newly created conditional handle.
 
         """
-        if default_value:
+        if default_value != None:
             flags = driver.CU_GRAPH_COND_ASSIGN_DEFAULT
         else:
             default_value = 0
@@ -441,7 +441,7 @@ def _cond_with_params(self, node_params) -> GraphBuilder:
             )
         )
 
-        # Create new graph builders for each conditional
+        # Create new graph builders for each condition
         return tuple(
             [
                 GraphBuilder._init(

cuda_core/tests/test_graph.py

@@ -19,15 +19,6 @@
 from cuda.core.experimental._memory import _DefaultPinnedMemorySource
 
 
-def test_graph_is_building(init_cuda):
-    gb = Device().create_graph_builder()
-    assert gb.is_building is False
-    gb.begin_building()
-    assert gb.is_building is True
-    gb.end_building()
-    assert gb.is_building is False
-
-
 def _common_kernels():
     code = """
     extern "C" __device__ __cudart_builtin__ void CUDARTAPI cudaGraphSetConditional(cudaGraphConditionalHandle handle,
@@ -48,6 +39,15 @@ def _common_kernels():
     return mod
 
 
+def test_graph_is_building(init_cuda):
+    gb = Device().create_graph_builder()
+    assert gb.is_building is False
+    gb.begin_building()
+    assert gb.is_building is True
+    gb.end_building()
+    assert gb.is_building is False
+
+
 def test_graph_straight(init_cuda):
     mod = _common_kernels()
     empty_kernel = mod.get_kernel("empty_kernel")
@@ -75,11 +75,18 @@ def test_graph_fork_join(init_cuda):
     gb = Device().create_graph_builder().begin_building()
     launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
 
+    with pytest.raises(ValueError, match="^Invalid split count: expecting >= 2, got 1"):
+        gb.split(1)
+
     left, right = gb.split(2)
     launch(left, LaunchConfig(grid=1, block=1), empty_kernel)
     launch(left, LaunchConfig(grid=1, block=1), empty_kernel)
     launch(right, LaunchConfig(grid=1, block=1), empty_kernel)
     launch(right, LaunchConfig(grid=1, block=1), empty_kernel)
+
+    with pytest.raises(ValueError, match="^Must join with at least two graph builders"):
+        GraphBuilder.join(left)
+
     gb = GraphBuilder.join(left, right)
 
     launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
@@ -135,7 +142,7 @@ def test_graph_is_join_required(init_cuda):
 
     # Create final node
     launch(gb, LaunchConfig(grid=1, block=1), empty_kernel)
-    gb.end_building()
+    gb.end_building().complete()
 
 
 @pytest.mark.skipif(tuple(int(i) for i in np.__version__.split(".")[:2]) < (2, 1), reason="need numpy 2.1.0+")
@@ -161,17 +168,17 @@ def test_graph_repeat_capture(init_cuda):
     assert arr[0] == 1
 
     # Continue capturing to extend the graph
-    with pytest.raises(RuntimeError):
+    with pytest.raises(RuntimeError, match="^Cannot resume building after building has ended."):
         gb.begin_building()
 
 
 def test_graph_capture_errors(init_cuda):
     gb = Device().create_graph_builder()
-    with pytest.raises(RuntimeError):
+    with pytest.raises(RuntimeError, match="^Graph has not finished building."):
         gb.complete()
 
     gb.begin_building()
-    with pytest.raises(RuntimeError):
+    with pytest.raises(RuntimeError, match="^Graph has not finished building."):
         gb.complete()
     gb.end_building().complete()
 
@@ -286,7 +293,7 @@ def test_graph_conditional_if_else(init_cuda, condition_value):
         assert arr[1] == 3
 
 
-@pytest.mark.parametrize("condition_value", [0, 1, 2])
+@pytest.mark.parametrize("condition_value", [0, 1, 2, 3])
 def test_graph_conditional_switch(init_cuda, condition_value):
     mod = _common_kernels()
     add_one = mod.get_kernel("add_one")
@@ -358,6 +365,11 @@ def test_graph_conditional_switch(init_cuda, condition_value):
         assert arr[0] == 1
         assert arr[1] == 0
         assert arr[2] == 3
+    elif condition_value == 3:
+        # No branch is taken if case index is out of range
+        assert arr[0] == 1
+        assert arr[1] == 0
+        assert arr[2] == 0
 
 
 @pytest.mark.parametrize("condition_value", [True, False])