bugfix: handle cuda graph errors

cpp/src/routing/cuda_graph.cuh

@@ -5,9 +5,12 @@
  */
 /* clang-format on */
 
-#include <cuda/stream>
 #include <cuopt/error.hpp>
-#include <utilities/macros.cuh>
+#include <utilities/scope_guard.hpp>
+
+#include <raft/util/cudart_utils.hpp>
+
+#include <cuda/stream>
 
 #pragma once
 
@@ -21,37 +24,51 @@ struct cuda_graph_t {
   {
     // Use ThreadLocal mode to allow multi-threaded batch execution
     // Global mode blocks other streams from performing operations during capture
-    cudaStreamBeginCapture(stream.get(), cudaStreamCaptureModeThreadLocal);
+    RAFT_CUDA_TRY(cudaStreamBeginCapture(stream.get(), cudaStreamCaptureModeThreadLocal));
     capture_started = true;
   }
 
   void end_capture(cuda::stream_ref stream)
   {
-    cuopt_assert(capture_started, "start_capture was not called before end_capture!");
-    cuopt_expects(capture_started, error_type_t::RuntimeError, "A runtime error occurred!");
-    cudaStreamEndCapture(stream.get(), &graph);
+    cuopt_expects(
+      capture_started, error_type_t::RuntimeError, "CUDA graph capture has not started!");
+    auto end_err    = cudaStreamEndCapture(stream.get(), &graph);
     capture_started = false;
+    RAFT_CUDA_TRY(end_err);
+    scope_guard destroy_graph([&] { RAFT_CUDA_TRY_NO_THROW(cudaGraphDestroy(graph)); });
     if (graph_created) {
       // If the graph fails to update, errorNode will be set to the
       // node causing the failure and updateResult will be set to a
       // reason code.
-      cudaGraphExecUpdate(instance, graph, &errorNode, &updateResult);
+      auto update_err = cudaGraphExecUpdate(instance, graph, &errorNode, &updateResult);
+      if (update_err == cudaErrorGraphExecUpdateFailure) {
+        // Expected update failures must not poison later CUDA error checks.
+        cudaGetLastError();
+      } else {
+        RAFT_CUDA_TRY(update_err);
+      }
     }
     // Instantiate during the first iteration or whenever the update
-    // fails for any reason
+    // cannot reuse the existing executable graph.
     if (!graph_created || updateResult != cudaGraphExecUpdateSuccess) {
       // If a previous update failed, destroy the cudaGraphExec_t
       // before re-instantiating it
-      if (graph_created) { cudaGraphExecDestroy(instance); }
+      if (graph_created) {
+        graph_created = false;
+        RAFT_CUDA_TRY(cudaGraphExecDestroy(instance));
+      }
       // Instantiate graphExec from graph. The error node and
       // error message parameters are unused here.
-      cudaGraphInstantiate(&instance, graph);
+      RAFT_CUDA_TRY(cudaGraphInstantiate(&instance, graph));
       graph_created = true;
     }
-    cudaGraphDestroy(graph);
   }
 
-  void launch_graph(cuda::stream_ref stream) { cudaGraphLaunch(instance, stream.get()); }
+  void launch_graph(cuda::stream_ref stream)
+  {
+    cuopt_expects(graph_created, error_type_t::RuntimeError, "CUDA graph is not instantiated!");
+    RAFT_CUDA_TRY(cudaGraphLaunch(instance, stream.get()));
+  }
 
   bool graph_created   = false;
   bool capture_started = false;