diff --git a/machine-learning/conftest.py b/machine-learning/conftest.py index 5668e6a77..c22030ba7 100644 --- a/machine-learning/conftest.py +++ b/machine-learning/conftest.py @@ -112,6 +112,7 @@ def gpus() -> Iterator[None]: with ( mock.patch("immich_ml.sessions.ort._intel_gpu", return_value=Device("12.71.4-128eu", "26.22.38646.4")), mock.patch("immich_ml.sessions.ort._amd_gpu", return_value=Device("gfx1100", "7.2.0")), + mock.patch("immich_ml.sessions.ort._nvidia_gpu", return_value=Device("sm89", "617.14 10601")), ): yield diff --git a/machine-learning/immich_ml/models/constants.py b/machine-learning/immich_ml/models/constants.py index 081541049..2920d8257 100644 --- a/machine-learning/immich_ml/models/constants.py +++ b/machine-learning/immich_ml/models/constants.py @@ -89,6 +89,7 @@ _PADDLE_MODELS = { } SUPPORTED_PROVIDERS = [ + "nv_tensorrt_rtx", "CUDAExecutionProvider", "MIGraphXExecutionProvider", "OpenVINOExecutionProvider", diff --git a/machine-learning/immich_ml/sessions/ort.py b/machine-learning/immich_ml/sessions/ort.py index 4ff1a6857..374cfba6d 100644 --- a/machine-learning/immich_ml/sessions/ort.py +++ b/machine-learning/immich_ml/sessions/ort.py @@ -16,7 +16,7 @@ import numpy as np import onnxruntime as ort from immich_model.runtime import RewriteContext, RewritePlan, plan_rewrites from numpy.typing import NDArray -from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf +from onnxruntime.capi.onnxruntime_pybind11_state import Fail, InvalidProtobuf from pydantic import BaseModel from immich_ml.schemas import ModelInput, SessionNode, Shape @@ -27,6 +27,17 @@ from .policy import ShapePolicy # the one provider that handles a free dim well; the rest miscompile it, recompile inside the run, # or partition around it, so they are handed one graph per shape instead DYNAMIC_PROVIDERS = frozenset({"CPUExecutionProvider"}) +PLANNED_AS = {"nv_tensorrt_rtx": "NvTensorRTRTXExecutionProvider"} + +try: + import onnxruntime_ep_nv_tensorrt_rtx as nv_tensorrt_rtx +except ImportError: + pass +else: + try: + ort.register_execution_provider_library(nv_tensorrt_rtx.get_ep_name(), nv_tensorrt_rtx.get_library_path()) + except Fail as e: + log.info(f"TensorRT-RTX is unavailable: {e}") def _label(pins: Mapping[str, int]) -> str: @@ -109,6 +120,8 @@ class GraphSpec: return _intel_gpu(self.openvino_device) case "MIGraphXExecutionProvider": return _amd_gpu(int(settings.device_id)) + case "nv_tensorrt_rtx": + return _nvidia_gpu(int(settings.device_id)) return None @cached_property @@ -159,6 +172,8 @@ class GraphSpec: options = {"arena_extend_strategy": "kSameAsRequested"} case "CUDAExecutionProvider": options = {"arena_extend_strategy": "kSameAsRequested", "device_id": settings.device_id} + case "nv_tensorrt_rtx": + options = {"device_id": settings.device_id, "nv_runtime_cache_path": self.directory.as_posix()} case "MIGraphXExecutionProvider": options = {"device_id": settings.device_id, "migraphx_model_cache_dir": self.directory.as_posix()} case "OpenVINOExecutionProvider": @@ -249,7 +264,7 @@ def prepared(spec: GraphSpec) -> Path: @cache def _plan(provider: str) -> RewritePlan: version = tuple(int(piece) for piece in ort.__version__.split(".")[:3]) - return plan_rewrites(RewriteContext(target=provider, ort_version=version)) + return plan_rewrites(RewriteContext(target=PLANNED_AS.get(provider, provider), ort_version=version)) @cache @@ -355,6 +370,20 @@ def _intel_gpu(device: str) -> Device: raise LookupError(f"OpenCL has no GPU with the UUID of OpenVINO's {device}") +@cache +def _nvidia_gpu(index: int) -> Device: + cuda, device, major, minor = ctypes.CDLL("libcuda.so.1"), ctypes.c_int(), ctypes.c_int(), ctypes.c_int() + if cuda.cuInit(0) or cuda.cuDeviceGet(ctypes.byref(device), index): # honors CUDA_VISIBLE_DEVICES + raise LookupError(f"CUDA has no device {index}") + cuda.cuDeviceGetAttribute(ctypes.byref(major), 75, device) # CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR + cuda.cuDeviceGetAttribute(ctypes.byref(minor), 76, device) # CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR + nvml, driver = ctypes.CDLL("libnvidia-ml.so.1"), ctypes.create_string_buffer(80) + nvml.nvmlInit_v2() + nvml.nvmlSystemGetDriverVersion(driver, 80) + tensorrt_rtx = ctypes.CDLL("libtensorrt_rtx.so.1").getInferLibVersion() + return Device(f"sm{major.value}{minor.value}", f"{driver.value.decode()} {tensorrt_rtx}") + + @cache def _amd_gpu(index: int) -> Device: """The arch MIGraphX keys its programs by, and the MIGraphX and HIP it loads, which its keys leave out.""" diff --git a/machine-learning/pyproject.toml b/machine-learning/pyproject.toml index 3271fbf96..738b05df0 100644 --- a/machine-learning/pyproject.toml +++ b/machine-learning/pyproject.toml @@ -48,7 +48,10 @@ dev = ["locust>=2.15.1", { include-group = "test" }, { include-group = "lint" }] [project.optional-dependencies] cpu = ["onnxruntime>=1.23.2,<2"] -cuda = ["onnxruntime-gpu>=1.27.1,<1.28"] # TODO: last ORT release supporting CUDA 12.2, also need to handle CUDA 13 +cuda = [ + "onnxruntime-gpu>=1.27.1,<1.28", # TODO: last ORT release supporting CUDA 12.2, also need to handle CUDA 13 + "onnxruntime-ep-nv-tensorrt-rtx-cu12>=0.4.0,<0.5 ; platform_machine == 'x86_64'", +] openvino = ["onnxruntime-openvino>=1.24.1,<2"] armnn = ["onnxruntime>=1.23.2,<2"] rknn = ["onnxruntime>=1.23.2,<2", "rknn-toolkit-lite2>=2.3.0,<3"] diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index d4e2e9360..93d8821c5 100644 --- a/machine-learning/test_main.py +++ b/machine-learning/test_main.py @@ -753,7 +753,11 @@ class TestPreparedGraphs: @pytest.mark.ov_device_ids(["GPU.0", "CPU"]) @pytest.mark.parametrize( ("provider", "reader"), - [("OpenVINOExecutionProvider", "_intel_gpu"), ("MIGraphXExecutionProvider", "_amd_gpu")], + [ + ("OpenVINOExecutionProvider", "_intel_gpu"), + ("MIGraphXExecutionProvider", "_amd_gpu"), + ("nv_tensorrt_rtx", "_nvidia_gpu"), + ], ) def test_prepares_once_per_kind_of_device_and_again_for_another_version( self, provider: str, reader: str, ov_device_ids: mock.Mock, mocker: MockerFixture diff --git a/machine-learning/uv.lock b/machine-learning/uv.lock index 00105a84e..502e848cd 100644 --- a/machine-learning/uv.lock +++ b/machine-learning/uv.lock @@ -705,6 +705,7 @@ cpu = [ { name = "onnxruntime" }, ] cuda = [ + { name = "onnxruntime-ep-nv-tensorrt-rtx-cu12", marker = "platform_machine == 'x86_64'" }, { name = "onnxruntime-gpu" }, ] openvino = [ @@ -769,6 +770,7 @@ requires-dist = [ { name = "onnxruntime", marker = "extra == 'armnn'", specifier = ">=1.23.2,<2" }, { name = "onnxruntime", marker = "extra == 'cpu'", specifier = ">=1.23.2,<2" }, { name = "onnxruntime", marker = "extra == 'rknn'", specifier = ">=1.23.2,<2" }, + { name = "onnxruntime-ep-nv-tensorrt-rtx-cu12", marker = "platform_machine == 'x86_64' and extra == 'cuda'", specifier = ">=0.4.0,<0.5" }, { name = "onnxruntime-gpu", marker = "extra == 'cuda'", specifier = ">=1.27.1,<1.28", index = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/" }, { name = "onnxruntime-migraphx", marker = "extra == 'rocm'", specifier = ">=1.23.2,<2" }, { name = "onnxruntime-openvino", marker = "extra == 'openvino'", specifier = ">=1.24.1,<2" }, @@ -1346,6 +1348,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/f0/40f74b7c00077e1e25627067ed98a70df1fef5c0e21b82849190d312554e/onnxruntime-1.30.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:67ad7f03433b6462c627d0f555dece80e6a26bc71e8542ced35cebd32142d1b7", size = 23579340, upload-time = "2026-09-10T16:32:15.532Z" }, ] +[[package]] +name = "onnxruntime-ep-nv-tensorrt-rtx-cu12" +version = "0.4.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/52/a6/1405069d40e2d6d37b21031e5dc2cc8b8197fbb2e3d4e4f65b9155c57e4f/onnxruntime_ep_nv_tensorrt_rtx_cu12-0.4.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:221807b1797a6270f37dfd7981f66e25ca7483f3883bfe23ef2b6a4198fe2b10", size = 167721016, upload-time = "2026-08-14T08:07:42.23Z" }, + { url = "https://files.pythonhosted.org/packages/8b/e4/f19fdbffcf8faf18d9f81973f81e3d378ef371f857b049a3f355535a015c/onnxruntime_ep_nv_tensorrt_rtx_cu12-0.4.0-py3-none-win_amd64.whl", hash = "sha256:5d9dba6aafd8e0863f34d0e5f51ad7c7ae9a5d6a92a55fe45d9637d3a5f1e1a7", size = 105822356, upload-time = "2026-08-12T11:13:34.401Z" }, +] + [[package]] name = "onnxruntime-gpu" version = "1.27.1"