fix(ml): serialize trt-rtx sessions (#31951)

machine-learning/immich_ml/sessions/ort.py

@@ -6,6 +6,7 @@ import platform
 import subprocess
 import sys
 from collections.abc import Mapping, Sequence
+from contextlib import nullcontext
 from dataclasses import dataclass
 from functools import cache, cached_property
 from pathlib import Path
@@ -53,26 +54,34 @@ DTYPES = {
     "tensor(int64)": np.int64,
 }
 
+# TensorRT-RTX records CUDA graphs in global capture mode, which fails any other thread's synchronizing call, and as its
+# context memory can move between runs, any run may record: so its graphs build and run one at a time
+# TODO: Make this more granular
+global_lock = Lock()
+
 
 class OrtGraph:
     def __init__(self, spec: GraphSpec) -> None:
-        self.session = spec.session(prepared(spec))
-        image = self.session.get_inputs()[0]
+        graph = prepared(spec)
+        self.lock = global_lock if spec.provider == "nv_tensorrt_rtx" else nullcontext()
+        with self.lock:
+            self._session = spec.session(graph)
+        image = self._session.get_inputs()[0]
         # whether the graph scales and shifts the image itself, so the host hands over raw pixels
         self.normalizes_input = image.type == "tensor(uint8)"
         # CoreML's rewrite moves the layout change out of the graph, so whoever applied it performs it
         self.channels_first = self.normalizes_input and len(image.shape) == 4 and image.shape[1] == 3
 
     def get_inputs(self) -> Sequence[SessionNode]:
-        inputs: Sequence[SessionNode] = self.session.get_inputs()
+        inputs: Sequence[SessionNode] = self._session.get_inputs()
         return inputs
 
     def get_outputs(self) -> Sequence[SessionNode]:
-        outputs: Sequence[SessionNode] = self.session.get_outputs()
+        outputs: Sequence[SessionNode] = self._session.get_outputs()
         return outputs
 
     def get_metadata(self) -> dict[str, str]:
-        metadata: dict[str, str] = self.session.get_modelmeta().custom_metadata_map
+        metadata: dict[str, str] = self._session.get_modelmeta().custom_metadata_map
         return metadata
 
     def run(
@@ -83,7 +92,13 @@ class OrtGraph:
     ) -> list[NDArray[np.float32]]:
         if self.channels_first:
             input_feed = {name: frames.transpose(0, 3, 1, 2) for name, frames in input_feed.items()}
-        outputs: list[NDArray[Any]] = self.session.run(output_names, input_feed, run_options)
+        return self.infer(output_names, input_feed, run_options)
+
+    def infer(
+        self, output_names: list[str] | None, input_feed: ModelInput, run_options: Any = None
+    ) -> list[NDArray[Any]]:
+        with self.lock:
+            outputs: list[NDArray[Any]] = self._session.run(output_names, input_feed, run_options)
         return outputs
 
 
@@ -321,15 +336,15 @@ class OrtSession:
         """Opens a graph for every shape and runs it once before opening the next: OpenVINO defers work to the first
         run of a graph it imports from its cache, which importing another graph first can corrupt."""
         for shape in self.policy.dims[:1] if self.dynamic else self.policy.dims:
-            session = self.for_shape(shape).session
+            graph = self.for_shape(shape)
             sizes = dict(_overrides(self.policy, shape.pins))  # for a dim the graph leaves free
             feed = {
                 node.name: np.zeros(
                     [dim if isinstance(dim, int) else sizes.get(dim, 1) for dim in node.shape], DTYPES[node.type]
                 )
-                for node in session.get_inputs()
+                for node in graph._session.get_inputs()
             }
-            session.run(None, feed)
+            graph.infer(None, feed)
 
 
 def _overrides(policy: ShapePolicy, pins: Mapping[str, int]) -> list[tuple[str, int]]:
@@ -425,7 +440,10 @@ def _providers_default() -> list[str]:
 
     available_providers = set(ort.get_available_providers())
     log.debug(f"Available ORT providers: {available_providers}")
-    return [provider for provider in SUPPORTED_PROVIDERS if provider in available_providers]
+    providers = [provider for provider in SUPPORTED_PROVIDERS if provider in available_providers]
+    if "nv_tensorrt_rtx" in providers:
+        providers.remove("CUDAExecutionProvider")  # CUDA EP syncs the legacy stream each run, which a recording fails
+    return providers
 
 
 def _disabled_optimizers_default(providers: list[str]) -> list[str]:

machine-learning/test_main.py

@@ -295,6 +295,7 @@ class TestOrtSessions:
     TRT_EP = ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]
     ROCM_EP = ["MIGraphXExecutionProvider", "CPUExecutionProvider"]
     COREML_EP = ["CoreMLExecutionProvider", "CPUExecutionProvider"]
+    TRT_RTX_EP = ["nv_tensorrt_rtx", "CUDAExecutionProvider", "CPUExecutionProvider"]
 
     @pytest.mark.providers(CPU_EP)
     def test_sets_cpu_provider(self, ort_session: mock.Mock, providers: list[str]) -> None:
@@ -341,6 +342,12 @@ class TestOrtSessions:
 
         assert given_providers(ort_session) == self.COREML_EP
 
+    @pytest.mark.providers(TRT_RTX_EP)
+    def test_leaves_cuda_out_beside_tensorrt_rtx(self, ort_session: mock.Mock, providers: list[str]) -> None:
+        ort_sessions("ViT-B-32__openai")
+
+        assert given_providers(ort_session) == ["nv_tensorrt_rtx", "CPUExecutionProvider"]
+
     def test_leaves_a_dimension_free_when_the_model_feeds_several_sizes(
         self, ort_session: mock.Mock, mocker: MockerFixture
     ) -> None:
@@ -408,7 +415,7 @@ class TestOrtSessions:
             events.append("open")
             return mock.DEFAULT
 
-        def run(output_names: Any, feed: dict[str, np.ndarray]) -> list[np.ndarray]:
+        def run(output_names: Any, feed: dict[str, np.ndarray], run_options: Any = None) -> list[np.ndarray]:
             events.append(f"run at {feed['image'].shape[2]}")  # a dim the graph leaves free takes the shape's size
             return [np.zeros(1)]