diff --git a/frigate/detectors/detection_runners.py b/frigate/detectors/detection_runners.py index 0ad6cfa255..7cdd4a0256 100644 --- a/frigate/detectors/detection_runners.py +++ b/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() diff --git a/frigate/test/test_detection_runners.py b/frigate/test/test_detection_runners.py index fbc71161d4..38ea457290 100644 --- a/frigate/test/test_detection_runners.py +++ b/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()