Guard against memory allocation in graph capture (#24192)

frigate/detectors/detection_runners.py

@@ -199,15 +199,20 @@ class CudaGraphRunner(BaseModelRunner):
             EnrichmentModelTypeEnum.yolov9_license_plate.value,
         ]
 
+    # ORT performs two regular runs before it starts capturing, but on some
+    # driver / cuDNN combinations the arena still has to extend on the run that
+    # captures, and cudaMalloc is not allowed during capture. Running with
+    # capture disabled first keeps those allocations outside of the capture.
+    GRAPH_FREE_WARMUP_RUNS = 2
+
     def __init__(self, session: ort.InferenceSession, cuda_device_id: int):
         self._session = session
         self._cuda_device_id = cuda_device_id
-        self._captured = False
+        self._prepared = False
         self._io_binding: ort.IOBinding | None = None
         self._input_name: str | None = None
         self._output_names: list[str] | None = None
         self._input_ortvalue: ort.OrtValue | None = None
-        self._output_ortvalues: ort.OrtValue | None = None
 
     def get_input_names(self) -> list[str]:
         """Get input names for the model."""
@@ -217,35 +222,41 @@ class CudaGraphRunner(BaseModelRunner):
         """Get the input width of the model."""
         return self._session.get_inputs()[0].shape[3]
 
+    def _prepare(self, input_name: str, tensor_input: np.ndarray) -> None:
+        """Bind CUDA buffers and warm the session up with capture disabled."""
+        self._io_binding = self._session.io_binding()
+        self._input_name = input_name
+        self._output_names = [o.name for o in self._session.get_outputs()]
+
+        self._input_ortvalue = ort.OrtValue.ortvalue_from_numpy(
+            tensor_input, "cuda", self._cuda_device_id
+        )
+        self._io_binding.bind_ortvalue_input(self._input_name, self._input_ortvalue)
+
+        for name in self._output_names:
+            # Bind outputs to CUDA and allow ORT to allocate appropriately
+            self._io_binding.bind_output(name, "cuda", self._cuda_device_id)
+
+        # gpu_graph_id -1 disables capture and replay for the run
+        warmup_options = ort.RunOptions()
+        warmup_options.add_run_config_entry("gpu_graph_id", "-1")
+
+        for _ in range(self.GRAPH_FREE_WARMUP_RUNS):
+            self._session.run_with_iobinding(self._io_binding, warmup_options)
+
+        self._prepared = True
+
     def run(self, input: dict[str, Any]):
         # Extract the single tensor input (assuming one input)
         input_name = list(input.keys())[0]
-        tensor_input = input[input_name]
-        tensor_input = np.ascontiguousarray(tensor_input)
+        tensor_input = np.ascontiguousarray(input[input_name])
 
-        if not self._captured:
-            # Prepare IOBinding with CUDA buffers and let ORT allocate outputs on device
-            self._io_binding = self._session.io_binding()
-            self._input_name = input_name
-            self._output_names = [o.name for o in self._session.get_outputs()]
+        if not self._prepared:
+            self._prepare(input_name, tensor_input)
+        else:
+            # Replay using updated input
+            self._input_ortvalue.update_inplace(tensor_input)
 
-            self._input_ortvalue = ort.OrtValue.ortvalue_from_numpy(
-                tensor_input, "cuda", self._cuda_device_id
-            )
-            self._io_binding.bind_ortvalue_input(self._input_name, self._input_ortvalue)
-
-            for name in self._output_names:
-                # Bind outputs to CUDA and allow ORT to allocate appropriately
-                self._io_binding.bind_output(name, "cuda", self._cuda_device_id)
-
-            # First IOBinding run to allocate, execute, and capture CUDA Graph
-            ro = ort.RunOptions()
-            self._session.run_with_iobinding(self._io_binding, ro)
-            self._captured = True
-            return self._io_binding.copy_outputs_to_cpu()
-
-        # Replay using updated input, copy results to CPU
-        self._input_ortvalue.update_inplace(tensor_input)
         ro = ort.RunOptions()
         self._session.run_with_iobinding(self._io_binding, ro)
         return self._io_binding.copy_outputs_to_cpu()

frigate/test/test_detection_runners.py

@@ -1,10 +1,15 @@
 """Tests for ONNX Runtime session option selection."""
 
 import unittest
+from unittest.mock import MagicMock, patch
 
+import numpy as np
 import onnxruntime as ort
 
-from frigate.detectors.detection_runners import get_ort_session_options
+from frigate.detectors.detection_runners import (
+    CudaGraphRunner,
+    get_ort_session_options,
+)
 from frigate.detectors.detector_config import ModelTypeEnum
 from frigate.embeddings.types import EnrichmentModelTypeEnum
 
@@ -39,3 +44,45 @@ class TestGetOrtSessionOptions(unittest.TestCase):
         ]:
             with self.subTest(model_type=model_type):
                 self.assertIsNone(get_ort_session_options(model_type))
+
+
+class TestCudaGraphRunner(unittest.TestCase):
+    """CUDA graph capture fails if the arena has to allocate during capture, so
+    the session is warmed up with capture disabled before the first real run."""
+
+    def setUp(self):
+        self.session = MagicMock()
+        self.session.get_outputs.return_value = [MagicMock(name="output")]
+        self.io_binding = self.session.io_binding.return_value
+        self.input = {"images": np.zeros((1, 3, 320, 320), np.float32)}
+
+    def _annotations(self) -> list[str | None]:
+        """Graph annotation id passed with each run, None when unset."""
+        annotations = []
+
+        for call in self.session.run_with_iobinding.call_args_list:
+            try:
+                annotations.append(call.args[1].get_run_config_entry("gpu_graph_id"))
+            except RuntimeError:
+                annotations.append(None)
+
+        return annotations
+
+    def test_first_run_warms_up_with_capture_disabled(self):
+        with patch.object(ort.OrtValue, "ortvalue_from_numpy"):
+            CudaGraphRunner(self.session, 0).run(self.input)
+
+        self.assertEqual(
+            self._annotations(),
+            ["-1"] * CudaGraphRunner.GRAPH_FREE_WARMUP_RUNS + [None],
+        )
+
+    def test_later_runs_allow_capture(self):
+        with patch.object(ort.OrtValue, "ortvalue_from_numpy"):
+            runner = CudaGraphRunner(self.session, 0)
+            runner.run(self.input)
+            self.session.run_with_iobinding.reset_mock()
+            runner.run(self.input)
+
+        self.assertEqual(self._annotations(), [None])
+        runner._input_ortvalue.update_inplace.assert_called_once()