diff --git a/machine-learning/immich_ml/sessions/ort.py b/machine-learning/immich_ml/sessions/ort.py index 374cfba6d..805d19205 100644 --- a/machine-learning/immich_ml/sessions/ort.py +++ b/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]: diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index 93d8821c5..b0e31ec3c 100644 --- a/machine-learning/test_main.py +++ b/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)]