From 4973617d5b6c062f2cbcc54394ee87f3c8dda27d Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 11:16:15 -0700 Subject: [PATCH 01/41] feat(openengine): support image inputs Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 57 ++ tensorrt_llm/grpc/openengine/control.py | 6 +- .../grpc/openengine/request_mapping.py | 118 +++- tensorrt_llm/grpc/openengine/servicer.py | 4 +- tensorrt_llm/inputs/media_io.py | 4 + .../integration/test_lists/test-db/l0_a10.yml | 8 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../unittest/grpc/openengine/test_control.py | 17 + .../grpc/openengine/test_multimodal.py | 531 ++++++++++++++++++ .../openengine/test_multimodal_conformance.py | 93 +++ .../inputs/test_async_media_loading.py | 9 + 11 files changed, 837 insertions(+), 11 deletions(-) create mode 100644 tests/unittest/grpc/openengine/test_multimodal.py create mode 100644 tests/unittest/grpc/openengine/test_multimodal_conformance.py diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 6272d1363855..a1936ebc4c8a 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -7,6 +7,17 @@ never discover a capability that Generate then rejects. """ +from typing import Any + +from tensorrt_llm.inputs.registry import ( + MULTIMODAL_PLACEHOLDER_REGISTRY, + BaseMultimodalInputProcessor, +) + +# OpenEngine support is intentionally narrower than the LLM API today. Add a +# modality here only once request mapping and conformance coverage exist for it. +_OPENENGINE_MEDIA_MODALITIES = ("image",) + # Guides each grammar backend can build, keyed by the `GuidedDecoding` oneof # field name. Control advertises from this table and Generate enforces it, so # the two cannot disagree. llguidance has no structural-tag matcher. @@ -32,7 +43,53 @@ def supported_guides(backend: str | None) -> frozenset[str]: return GUIDE_SUPPORT_BY_BACKEND.get(backend.lower(), frozenset()) +def _loaded_model_type(llm: Any) -> str: + """Return the top-level model type for the loaded input processor.""" + input_processor = getattr(llm, "input_processor", None) + config = ( + input_processor.config + if isinstance(input_processor, BaseMultimodalInputProcessor) + else getattr(llm, "_hf_model_config", None) + ) + if config is None: + return "" + + # Composite Hugging Face configs can delegate the instance attribute to + # their text config. The class attribute retains the top-level registry key. + model_type = getattr(type(config), "model_type", None) or getattr(config, "model_type", "") + return model_type if isinstance(model_type, str) else "" + + +def supported_media_modalities(llm: Any) -> frozenset[str]: + """OpenEngine media modalities supported by the loaded model. + + A registry entry alone describes what this TensorRT-LLM build knows about; + it does not prove that the loaded engine instantiated a multimodal input + processor. Requiring both prevents a text-only model from advertising + media merely because another image model is registered in the process. + """ + input_processor = getattr(llm, "input_processor", None) + if not isinstance(input_processor, BaseMultimodalInputProcessor): + return frozenset() + model_type = _loaded_model_type(llm) + if not model_type: + return frozenset() + return frozenset( + modality + for modality in _OPENENGINE_MEDIA_MODALITIES + if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) + ) + + +def media_placeholder(llm: Any, modality: str, item_number: int) -> str: + """Return the model-specific placeholder for a one-based media item.""" + model_type = _loaded_model_type(llm) + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality).format(item_number) + + __all__ = [ "GUIDE_SUPPORT_BY_BACKEND", + "media_placeholder", "supported_guides", + "supported_media_modalities", ] diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index f020d3ae82a1..1495524d4130 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -35,7 +35,7 @@ openengine_pb2_grpc, server_pb2, ) -from .capabilities import supported_guides +from .capabilities import supported_guides, supported_media_modalities from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError from .errors import AbortFailedError from .kv_events import data_parallel_size, events_config @@ -284,6 +284,7 @@ async def GetModelInfo( ) -> model_pb2.ModelInfo: # Single-model server: like Generate, any non-empty name resolves to the # loaded model. + media_modalities = supported_media_modalities(self._llm) info = model_pb2.ModelInfo( model_id=self._model, served_model_name=self._model, @@ -292,9 +293,10 @@ async def GetModelInfo( supports_token_ids_input=True, # Generate rejects `lora_name` even on an enable_lora build. supports_lora=False, - supports_multimodal=False, + supports_multimodal=bool(media_modalities), reasoning_parser=str(self._llm.args.reasoning_parser or ""), ) + info.extra.update({"trtllm_supported_modalities": sorted(media_modalities)}) max_seq_len = _positive_int(self._llm.args.max_seq_len) if max_seq_len is not None: diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index db6ecf866c01..7e7035e0a1a8 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -3,22 +3,31 @@ """Translate an OpenEngine GenerateRequest into TensorRT-LLM request inputs.""" +import asyncio from collections.abc import Mapping from typing import Any +from urllib.parse import urlparse import grpc from tensorrt_llm.conversation_params import ConversationParams +from tensorrt_llm.inputs.media_io import ImageMediaIO +from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams from .bindings import generation_pb2 -from .capabilities import supported_guides +from .capabilities import media_placeholder, supported_guides, supported_media_modalities from .errors import UnsupportedFeatureError # Cap on the number of stop conditions. Each one costs a KMP automaton per # output index and a scan of every generated token, so an unbounded list lets a # single request tax the event loop for the whole stream. _MAX_STOP_CONDITIONS = 64 +# Media decoding can retain both encoded and decoded image buffers. Keep one +# request from scheduling an unbounded number of fetch/decode tasks or holding +# an unbounded amount of image data while preserving practical multi-image use. +_MAX_MEDIA_ITEMS = 64 +_MAX_CONCURRENT_MEDIA_LOADS = 8 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -244,7 +253,9 @@ def _metadata_from_context( return headers or None, target_dp_rank, affinity_headers -def _input_from_request(request: generation_pb2.GenerateRequest) -> str | dict[str, list[int]]: +def _base_input_from_request( + request: generation_pb2.GenerateRequest, +) -> str | dict[str, list[int]]: input_kind = request.WhichOneof("input") if input_kind == "prompt": return request.prompt @@ -253,6 +264,109 @@ def _input_from_request(request: generation_pb2.GenerateRequest) -> str | dict[s raise ValueError("exactly one of prompt or token_ids must be set") +def _image_io_for_model(llm: Any) -> ImageMediaIO: + """Create image I/O using the loaded model's preferred decode format.""" + input_processor = getattr(llm, "input_processor", None) + defaults = None + if isinstance(input_processor, BaseMultimodalInputProcessor): + preferences = input_processor.get_preferred_media_io_kwargs() or {} + defaults = preferences.get("image") + return ImageMediaIO.create(defaults, None) + + +async def _load_image_item( + item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int +) -> Any: + """Resolve and decode one OpenEngine image item.""" + source = item.WhichOneof("source") + if source is None: + raise ValueError(f"media[{index}] must set exactly one source") + if source == "url": + if urlparse(item.url).scheme.lower() not in ("http", "https"): + raise ValueError(f"media[{index}].url must use http or https") + value = item.url + elif source == "data_uri": + if urlparse(item.data_uri).scheme.lower() != "data": + raise ValueError(f"media[{index}].data_uri must use the data scheme") + value = item.data_uri + elif source == "raw_bytes": + if item.mime_type and not item.mime_type.lower().startswith("image/"): + raise ValueError(f"media[{index}].mime_type must be an image type for raw image bytes") + try: + return await image_io.async_load_bytes(bytes(item.raw_bytes)) + except Exception as error: + raise ValueError(f"media[{index}].raw_bytes could not be decoded: {error}") from error + else: + raise ValueError(f"media[{index}] has an unknown source '{source}'") + + try: + return await image_io.async_load(value) + except Exception as error: + raise ValueError(f"media[{index}].{source} could not be loaded: {error}") from error + + +async def _input_from_request( + request: generation_pb2.GenerateRequest, llm: Any +) -> str | dict[str, Any]: + """Build LLM API inputs, resolving image media off the event loop.""" + inputs = _base_input_from_request(request) + if not request.media: + return inputs + + supported = supported_media_modalities(llm) + for index, item in enumerate(request.media): + if item.modality == generation_pb2.MODALITY_UNSPECIFIED: + raise ValueError(f"media[{index}].modality must be set") + if item.modality != generation_pb2.MODALITY_IMAGE: + modality = generation_pb2.Modality.Name(item.modality).removeprefix("MODALITY_").lower() + supported_text = ", ".join(sorted(supported)) or "none" + raise UnsupportedFeatureError( + f"multimodal modality '{modality}' is not supported by the loaded model; " + f"supported modalities: {supported_text}" + ) + if "image" not in supported: + raise UnsupportedFeatureError("image input is not supported by the loaded model") + if len(request.media) > _MAX_MEDIA_ITEMS: + raise ValueError( + f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " + f"got {len(request.media)}" + ) + + image_io = _image_io_for_model(llm) + semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) + + async def load_bounded(item: generation_pb2.MediaItem, index: int) -> Any: + async with semaphore: + return await _load_image_item(item, image_io, index) + + tasks = [ + asyncio.create_task(load_bounded(item, index)) for index, item in enumerate(request.media) + ] + try: + images = await asyncio.gather(*tasks) + except (Exception, asyncio.CancelledError): + # `gather` propagates the first failure without cancelling siblings. + # Explicitly stop any in-flight URL fetches/decodes before returning + # the request error so their memory and connections do not linger. + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + mapped: dict[str, Any] = {"prompt": inputs} if isinstance(inputs, str) else dict(inputs) + mapped["multi_modal_data"] = {"image": images} + if any(item.uuid for item in request.media): + mapped["multi_modal_uuids"] = {"image": [item.uuid or None for item in request.media]} + mapped["mm_item_order"] = [ + { + "modality": "image", + "index": index, + "placeholder": media_placeholder(llm, "image", index + 1), + } + for index in range(len(request.media)) + ] + return mapped + + __all__ = [ "conversation_params_from_request", "sampling_params_from_request", diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index 6a9d6e9018b9..6bc18ffeb6a4 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -169,15 +169,13 @@ async def _generate( f"request_id '{request_id}' is already active", ) return - if request.media: - raise UnsupportedFeatureError("multimodal media is not supported") if request.lora_name: raise UnsupportedFeatureError("LoRA selection is not supported") if request.HasField("kv"): if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: raise UnsupportedFeatureError("prefix cache bypass is not supported") - inputs = _input_from_request(request) + inputs = await _input_from_request(request, self._llm) sampling_params = sampling_params_from_request(request, self._guided_backend) trace_headers, target_dp_rank, affinity_headers = _metadata_from_context(context) conversation_params = conversation_params_from_request(request) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 86f52f71da3c..e8ab41d4dccf 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -1002,6 +1002,10 @@ async def async_load(self, url: str) -> _MediaT: else: raise ValueError(f"Unsupported URL scheme: {parsed.scheme!r}") + async def async_load_bytes(self, data: bytes) -> _MediaT: + """Decode in-memory media without blocking the event loop.""" + return await self._run_in_executor(self.load_bytes, data) + @staticmethod async def _run_in_executor(fn, *args, **kwargs): """Run a blocking decode callable on the configured executor.""" diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 674ad6a72694..064e9daa8b1f 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -63,11 +63,11 @@ l0_a10: - unittest/others/test_tracing.py - unittest/grpc/test_grpc_optional.py - unittest/grpc/smg/test_smg.py - # The one OpenEngine test needing a GPU: it builds a real LLM, because a - # stubbed engine would accept whatever its fake generate_async accepts and so - # would agree with an advertisement the real engine rejects. Any GPU will do - # -- the model is TinyLlama-1.1B -- so it sits with the other gateway tests. + # OpenEngine capability conformance needs real loaded models: a stub would + # accept whatever its fake generate_async tolerates and could agree with a + # false advertisement by construction. The image case uses Qwen3-VL-2B. - unittest/grpc/openengine/test_capability_conformance.py + - unittest/grpc/openengine/test_multimodal_conformance.py - unittest/disaggregated/test_coordinator_e2e.py - unittest/disaggregated/test_coordinator_worker.py - unittest/disaggregated/test_agent_multi_backends.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index f320c5b2a0dd..fd710d0f1645 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -154,6 +154,7 @@ l0_cpu: - unittest/grpc/openengine/test_coordinator.py - unittest/grpc/openengine/test_frontend_lifecycle.py - unittest/grpc/openengine/test_control.py + - unittest/grpc/openengine/test_multimodal.py - unittest/grpc/openengine/test_node.py - unittest/grpc/openengine/test_openengine.py - unittest/grpc/openengine/test_servicer.py diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 10420b6df6bd..f7de8842a222 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -166,6 +166,23 @@ async def test_model_info_reports_the_context_window_and_capabilities(): assert model_pb2.GUIDED_DECODING_MODE_JSON_SCHEMA in info.generation.guided_decoding.modes assert info.generation.output_logprobs.supported is True assert info.supports_token_ids_input is True + assert info.supports_multimodal is False + assert list(info.extra["trtllm_supported_modalities"]) == [] + + +@pytest.mark.asyncio +async def test_model_info_advertises_loaded_image_support(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm.grpc.openengine.control.supported_media_modalities", + lambda llm: frozenset({"image"}), + ) + + info = await _servicer().GetModelInfo( + model_pb2.GetModelInfoRequest(model=MODEL), FakeServicerContext() + ) + + assert info.supports_multimodal is True + assert list(info.extra["trtllm_supported_modalities"]) == ["image"] @pytest.mark.asyncio diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py new file mode 100644 index 000000000000..a4bd537e29f5 --- /dev/null +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -0,0 +1,531 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for OpenEngine image input mapping and discovery.""" + +import asyncio +from collections.abc import AsyncIterator, Sequence +from types import SimpleNamespace +from typing import Any + +import pytest + +grpc = pytest.importorskip( + "grpc", reason='gRPC runtime not installed (pip install "grpcio>=1.67.1,<2")' +) +from conftest import AbortError, FakeServicerContext # noqa: E402 + +import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 +from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.capabilities import supported_media_modalities # noqa: E402 +from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 +from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.inputs.registry import ( # noqa: E402 + MULTIMODAL_PLACEHOLDER_REGISTRY, + BaseMultimodalInputProcessor, +) + +pytestmark = pytest.mark.cpu_only + + +class _ImageConfig: + # Composite configs can expose their text model type on the instance; the + # top-level class attribute is the key used by the multimodal registry. + model_type = "openengine_test_image" + + def __init__(self) -> None: + self.model_type = "openengine_test_text" + + +class _ImageInputProcessor(BaseMultimodalInputProcessor): + """Minimal concrete processor used only for capability detection.""" + + def __init__(self, config: _ImageConfig) -> None: + self._config = config + + @property + def processor(self): + return None + + @property + def tokenizer(self): + return None + + @property + def config(self): + return self._config + + @property + def dtype(self): + return None + + def call_with_text_prompt(self, inputs, sampling_params): + raise NotImplementedError + + def get_preferred_media_io_kwargs(self): + return {"image": {"format": "np"}} + + +class _Tokenizer: + def convert_ids_to_tokens( + self, token_ids: Sequence[int], skip_special_tokens: bool = False + ) -> list[str]: + del skip_special_tokens + return [f"token-{token_id}" for token_id in token_ids] + + def decode(self, token_ids: Sequence[int], skip_special_tokens: bool = False) -> str: + del skip_special_tokens + return "".join(f"token-{token_id}" for token_id in token_ids) + + +class _ResultHandle: + def __init__(self, result: Any) -> None: + self._result = result + self.finished = False + self.prompt_token_ids = [] + self.outputs = [] + self.cached_tokens = 0 + self.error = None + self.aborted = False + + async def __aiter__(self) -> AsyncIterator[Any]: + self.finished = self._result.finished + self.prompt_token_ids = self._result.prompt_token_ids + self.outputs = self._result.outputs + self.cached_tokens = self._result.cached_tokens + self.error = self._result.error + yield self._result + + def abort(self) -> None: + self.aborted = True + + +class _ImageLlm: + def __init__(self, input_processor, config) -> None: + self.input_processor = input_processor + self._hf_model_config = config + self.args = SimpleNamespace(guided_decoding_backend=None) + self.tokenizer = _Tokenizer() + output = SimpleNamespace( + index=0, + token_ids=[99], + text="caption", + logprobs=[], + prompt_logprobs=[], + finish_reason="length", + stop_reason=None, + ) + self.result_handle = _ResultHandle( + SimpleNamespace( + # Represents expansion of one placeholder into five image tokens. + prompt_token_ids=[1, 2, 3, 4, 5, 6], + outputs=[output], + cached_tokens=0, + error=None, + finished=True, + ) + ) + self.generate_kwargs = None + + def generate_async(self, **kwargs): + self.generate_kwargs = kwargs + return self.result_handle + + +@pytest.fixture +def image_llm(monkeypatch): + config = _ImageConfig() + input_processor = _ImageInputProcessor(config) + seen_model_types = [] + + def is_valid(model_type, modality): + seen_model_types.append(model_type) + return model_type == _ImageConfig.model_type and modality == "image" + + monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", is_valid) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: "", + ) + return ( + SimpleNamespace( + _hf_model_config=config, + input_processor=input_processor, + ), + seen_model_types, + ) + + +def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): + llm, seen_model_types = image_llm + + assert supported_media_modalities(llm) == {"image"} + assert seen_model_types == [_ImageConfig.model_type] + assert ( + supported_media_modalities( + SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) + ) + == set() + ) + + +@pytest.mark.asyncio +async def test_text_only_input_mapping_is_unchanged(): + llm = SimpleNamespace() + + prompt = generation_pb2.GenerateRequest(prompt="hello") + token_ids = generation_pb2.GenerateRequest(token_ids=generation_pb2.TokenIds(ids=[1, 2, 3])) + + assert await request_mapping._input_from_request(prompt, llm) == "hello" + assert await request_mapping._input_from_request(token_ids, llm) == { + "prompt_token_ids": [1, 2, 3] + } + + +@pytest.mark.asyncio +async def test_image_sources_preserve_order_uuid_and_model_decode_defaults(image_llm, monkeypatch): + llm, _ = image_llm + created_with = [] + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append((defaults, runtime)) + return cls() + + async def async_load(self, source): + # Finish in the opposite order to prove gather preserves request order. + if source.startswith("https"): + await asyncio.sleep(0.01) + return f"decoded:{source}" + + async def async_load_bytes(self, data): + return f"decoded-bytes:{data.decode()}" + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt=" describe", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="https://example.com/first.png", + uuid="first", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + data_uri="data:image/png;base64,c2Vjb25k", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"third", + mime_type="image/png", + uuid="third", + ), + ], + ) + + mapped = await request_mapping._input_from_request(request, llm) + + assert created_with == [({"format": "np"}, None)] + assert mapped["prompt"] == request.prompt + assert mapped["multi_modal_data"] == { + "image": [ + "decoded:https://example.com/first.png", + "decoded:data:image/png;base64,c2Vjb25k", + "decoded-bytes:third", + ] + } + assert mapped["multi_modal_uuids"] == {"image": ["first", None, "third"]} + assert mapped["mm_item_order"] == [ + {"modality": "image", "index": 0, "placeholder": ""}, + {"modality": "image", "index": 1, "placeholder": ""}, + {"modality": "image", "index": 2, "placeholder": ""}, + ] + + +@pytest.mark.asyncio +async def test_image_loading_has_a_bounded_concurrency(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + active = 0 + maximum_active = 0 + + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + type(self).active += 1 + type(self).maximum_active = max(type(self).maximum_active, type(self).active) + try: + await asyncio.sleep(0.01) + return data + finally: + type(self).active -= 1 + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=str(index).encode(), + ) + for index in range(request_mapping._MAX_CONCURRENT_MEDIA_LOADS * 2 + 1) + ], + ) + + mapped = await request_mapping._input_from_request(request, llm) + + assert _ImageIo.maximum_active == request_mapping._MAX_CONCURRENT_MEDIA_LOADS + assert mapped["multi_modal_data"]["image"] == [ + str(index).encode() for index in range(request_mapping._MAX_CONCURRENT_MEDIA_LOADS * 2 + 1) + ] + + +@pytest.mark.asyncio +async def test_image_loading_cancels_siblings_after_a_failure(image_llm, monkeypatch): + llm, _ = image_llm + slow_started = asyncio.Event() + slow_cancelled = asyncio.Event() + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + if data == b"slow": + slow_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + slow_cancelled.set() + raise + await slow_started.wait() + raise ValueError("decode failed") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=value, + ) + for value in (b"slow", b"broken") + ], + ) + + with pytest.raises(ValueError, match="decode failed"): + await request_mapping._input_from_request(request, llm) + assert slow_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_image_count_is_bounded_before_decoding(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + raise AssertionError("the media-count guard must run before decoder creation") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="too many images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + for _ in range(request_mapping._MAX_MEDIA_ITEMS + 1) + ], + ) + + with pytest.raises(ValueError, match="at most 64 media items"): + await request_mapping._input_from_request(request, llm) + + +@pytest.mark.asyncio +async def test_image_mapping_supports_token_ids(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return "decoded" + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + token_ids=generation_pb2.TokenIds(ids=[10, 20]), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + mapped = await request_mapping._input_from_request(request, llm) + + assert mapped["prompt_token_ids"] == [10, 20] + assert mapped["multi_modal_data"] == {"image": ["decoded"]} + assert "multi_modal_uuids" not in mapped + + +@pytest.mark.asyncio +async def test_generate_submits_image_and_reports_expanded_prompt_usage(image_llm, monkeypatch): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return "decoded" + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + request = generation_pb2.GenerateRequest( + request_id="image-request", + model="image-model", + token_ids=generation_pb2.TokenIds(ids=[1, 2]), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + responses = [response async for response in servicer.Generate(request, FakeServicerContext())] + + assert llm.generate_kwargs["inputs"]["multi_modal_data"] == {"image": ["decoded"]} + assert responses[-1].usage.prompt_tokens == 6 + assert responses[-1].usage.total_tokens == 7 + + +@pytest.mark.asyncio +async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + + unsupported_context = FakeServicerContext() + video_request = generation_pb2.GenerateRequest( + request_id="video-request", + model="image-model", + prompt="video", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"video", + ) + ], + ) + with pytest.raises(AbortError): + async for _ in servicer.Generate(video_request, unsupported_context): + pass + assert unsupported_context.abort_code == grpc.StatusCode.UNIMPLEMENTED + + malformed_context = FakeServicerContext() + file_request = generation_pb2.GenerateRequest( + request_id="file-request", + model="image-model", + prompt="image", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="file:///tmp/private.png", + ) + ], + ) + with pytest.raises(AbortError): + async for _ in servicer.Generate(file_request, malformed_context): + pass + assert malformed_context.abort_code == grpc.StatusCode.INVALID_ARGUMENT + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("modality", "name"), + [ + (generation_pb2.MODALITY_VIDEO, "video"), + (generation_pb2.MODALITY_AUDIO, "audio"), + ], +) +async def test_video_and_audio_are_explicitly_unsupported(image_llm, modality, name): + llm, _ = image_llm + request = generation_pb2.GenerateRequest( + prompt="media", + media=[generation_pb2.MediaItem(modality=modality, raw_bytes=b"data")], + ) + + with pytest.raises(UnsupportedFeatureError, match=name): + await request_mapping._input_from_request(request, llm) + + +@pytest.mark.asyncio +async def test_image_is_rejected_for_a_text_only_loaded_model(): + request = generation_pb2.GenerateRequest( + prompt="image", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + with pytest.raises(UnsupportedFeatureError, match="loaded model"): + await request_mapping._input_from_request( + request, SimpleNamespace(input_processor=object()) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("item", "message"), + [ + (generation_pb2.MediaItem(raw_bytes=b"image"), "modality must be set"), + ( + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE), + "must set exactly one source", + ), + ( + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="file:///tmp/private.png", + ), + "must use http or https", + ), + ( + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + mime_type="video/mp4", + ), + "must be an image type", + ), + ], +) +async def test_malformed_image_items_are_rejected(image_llm, item, message): + llm, _ = image_llm + request = generation_pb2.GenerateRequest(prompt="image", media=[item]) + + with pytest.raises(ValueError, match=message): + await request_mapping._input_from_request(request, llm) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py new file mode 100644 index 000000000000..e97c9a708aac --- /dev/null +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Image conformance tests against a real TensorRT-LLM VLM engine.""" + +import asyncio +import base64 + +import pytest + +pytest.importorskip("grpc", reason='gRPC runtime not installed (pip install "grpcio>=1.67.1,<2")') + +import torch # noqa: E402 +from conftest import FakeServicerContext # noqa: E402 +from utils.llm_data import llm_models_root # noqa: E402 + +from tensorrt_llm import LLM # noqa: E402 +from tensorrt_llm.grpc.openengine.bindings import generation_pb2, model_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 +from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.llmapi import KvCacheConfig # noqa: E402 + +MODEL_NAME = "Qwen3/Qwen3-VL-2B-Instruct" + +# The engine owns background executor threads for the lifetime of this module. +pytestmark = pytest.mark.threadleak(enabled=False) + + +@pytest.fixture(scope="module") +def vlm(): + if not torch.cuda.is_available(): + pytest.skip("requires a GPU to run a real VLM engine") + model_path = llm_models_root(check=True) / MODEL_NAME + engine = LLM( + model=str(model_path), + load_format="dummy", + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.5), + max_batch_size=1, + max_seq_len=2048, + ) + try: + yield engine + finally: + engine.shutdown() + + +def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm): + inference = OpenEngineInferenceServicer(vlm, MODEL_NAME) + control = OpenEngineControlServicer(vlm, MODEL_NAME, inference) + info = asyncio.run( + control.GetModelInfo( + model_pb2.GetModelInfoRequest(model=MODEL_NAME), + FakeServicerContext(), + ) + ) + + assert info.supports_multimodal is True + assert list(info.extra["trtllm_supported_modalities"]) == ["image"] + + prompt = ( + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + "Describe this image briefly.<|im_end|>\n<|im_start|>assistant\n" + ) + prompt_token_ids = vlm.tokenizer.encode(prompt) + image_path = llm_models_root(check=True) / "multimodals/test_data/seashore.png" + image_bytes = image_path.read_bytes() + data_uri = "data:image/png;base64," + base64.b64encode(image_bytes).decode() + request = generation_pb2.GenerateRequest( + request_id="openengine-image-conformance", + model=MODEL_NAME, + token_ids=generation_pb2.TokenIds(ids=prompt_token_ids), + stopping=generation_pb2.StoppingOptions(max_tokens=4), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + data_uri=data_uri, + uuid="seashore", + ) + ], + ) + + async def collect(): + return [response async for response in inference.Generate(request, FakeServicerContext())] + + responses = asyncio.run(collect()) + assert responses + assert not any(response.HasField("error") for response in responses) + finished = [response for response in responses if response.HasField("finished")] + assert len(finished) == 1 + assert finished[0].usage.prompt_tokens > len(prompt_token_ids) + assert finished[0].usage.total_tokens == ( + finished[0].usage.prompt_tokens + finished[0].usage.completion_tokens + ) diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index 5855605a06ca..997b140f79f7 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -58,6 +58,15 @@ def _make_audio_file(tmp_path: str) -> str: class TestAsyncLoadImage: + @pytest.mark.asyncio + async def test_load_image_from_raw_bytes_pil(self): + encoded = _make_image_data_url().partition(",")[2] + image = await media_io_module.ImageMediaIO(format="pil").async_load_bytes( + base64.b64decode(encoded) + ) + assert isinstance(image, Image.Image) + assert image.mode == "RGB" + @pytest.mark.asyncio async def test_load_image_from_data_url_pil(self): url = _make_image_data_url() From e459d6799c891d963192896c3f55c7b82c9bc964 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 16:55:11 -0700 Subject: [PATCH 02/41] fix(openengine): harden multimodal media handling Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 74 ++++++--- tensorrt_llm/grpc/openengine/control.py | 4 +- tensorrt_llm/grpc/openengine/errors.py | 5 + .../grpc/openengine/request_mapping.py | 157 ++++++++++++++---- tensorrt_llm/grpc/openengine/servicer.py | 14 +- tensorrt_llm/inputs/media_io.py | 41 +++-- .../unittest/grpc/openengine/test_control.py | 9 +- .../grpc/openengine/test_multimodal.py | 138 ++++++++++++++- tests/unittest/inputs/test_url_validation.py | 21 ++- 9 files changed, 383 insertions(+), 80 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index a1936ebc4c8a..6affcf574747 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -7,17 +7,15 @@ never discover a capability that Generate then rejects. """ -from typing import Any +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Collection, Mapping from tensorrt_llm.inputs.registry import ( MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, ) -# OpenEngine support is intentionally narrower than the LLM API today. Add a -# modality here only once request mapping and conformance coverage exist for it. -_OPENENGINE_MEDIA_MODALITIES = ("image",) - # Guides each grammar backend can build, keyed by the `GuidedDecoding` oneof # field name. Control advertises from this table and Generate enforces it, so # the two cannot disagree. llguidance has no structural-tag matcher. @@ -29,6 +27,27 @@ } +@dataclass(frozen=True) +class MediaCapabilities: + """Immutable media capability snapshot shared by Control and Generate.""" + + placeholder_templates: Mapping[str, str] + + def __post_init__(self) -> None: + object.__setattr__( + self, "placeholder_templates", MappingProxyType(dict(self.placeholder_templates)) + ) + + @property + def modalities(self) -> frozenset[str]: + """Modalities the loaded model and request mapper both support.""" + return frozenset(self.placeholder_templates) + + def placeholder(self, modality: str, item_number: int) -> str: + """Format the model placeholder for a one-based media item.""" + return self.placeholder_templates[modality].format(item_number) + + def supported_guides(backend: str | None) -> frozenset[str]: """Guides `backend` can build; empty when none can be. @@ -60,36 +79,47 @@ def _loaded_model_type(llm: Any) -> str: return model_type if isinstance(model_type, str) else "" -def supported_media_modalities(llm: Any) -> frozenset[str]: - """OpenEngine media modalities supported by the loaded model. +def resolve_media_capabilities( + llm: Any, implemented_modalities: Collection[str] +) -> MediaCapabilities: + """Resolve media support implemented by both the mapper and loaded model. A registry entry alone describes what this TensorRT-LLM build knows about; it does not prove that the loaded engine instantiated a multimodal input processor. Requiring both prevents a text-only model from advertising - media merely because another image model is registered in the process. + media merely because another image model is registered in the process. The + request mapper supplies ``implemented_modalities`` from its handler table, + so discovery cannot advertise a modality without a request implementation. """ input_processor = getattr(llm, "input_processor", None) if not isinstance(input_processor, BaseMultimodalInputProcessor): - return frozenset() + return MediaCapabilities({}) model_type = _loaded_model_type(llm) if not model_type: - return frozenset() - return frozenset( - modality - for modality in _OPENENGINE_MEDIA_MODALITIES - if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) - ) + return MediaCapabilities({}) - -def media_placeholder(llm: Any, modality: str, item_number: int) -> str: - """Return the model-specific placeholder for a one-based media item.""" - model_type = _loaded_model_type(llm) - return MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality).format(item_number) + placeholders = { + modality: MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality) + for modality in implemented_modalities + if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) + } + # Validate templates once at startup. A malformed registry entry is a + # server configuration failure, not a client request error discovered only + # after Control has advertised the modality. + for modality, template in placeholders.items(): + try: + template.format(1) + except (IndexError, KeyError, ValueError) as error: + raise RuntimeError( + f"invalid multimodal placeholder for model type '{model_type}' " + f"and modality '{modality}': {error}" + ) from error + return MediaCapabilities(placeholders) __all__ = [ "GUIDE_SUPPORT_BY_BACKEND", - "media_placeholder", + "MediaCapabilities", + "resolve_media_capabilities", "supported_guides", - "supported_media_modalities", ] diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index 1495524d4130..8b4292dd33ae 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -35,7 +35,7 @@ openengine_pb2_grpc, server_pb2, ) -from .capabilities import supported_guides, supported_media_modalities +from .capabilities import supported_guides from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError from .errors import AbortFailedError from .kv_events import data_parallel_size, events_config @@ -284,7 +284,7 @@ async def GetModelInfo( ) -> model_pb2.ModelInfo: # Single-model server: like Generate, any non-empty name resolves to the # loaded model. - media_modalities = supported_media_modalities(self._llm) + media_modalities = self._inference.supported_media_modalities info = model_pb2.ModelInfo( model_id=self._model, served_model_name=self._model, diff --git a/tensorrt_llm/grpc/openengine/errors.py b/tensorrt_llm/grpc/openengine/errors.py index 267d2a508c8a..d639794e0c15 100644 --- a/tensorrt_llm/grpc/openengine/errors.py +++ b/tensorrt_llm/grpc/openengine/errors.py @@ -8,6 +8,10 @@ class UnsupportedFeatureError(ValueError): """Raised when OpenEngine requests a feature this adapter cannot map.""" +class MediaUnavailableError(RuntimeError): + """Raised when fetching remote request media fails transiently.""" + + class AbortFailedError(RuntimeError): """Raised when an active request could not be aborted on the engine.""" @@ -18,5 +22,6 @@ def __init__(self, request_id: str, cause: BaseException) -> None: __all__ = [ "AbortFailedError", + "MediaUnavailableError", "UnsupportedFeatureError", ] diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 7e7035e0a1a8..28c948ba89d4 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -4,20 +4,24 @@ """Translate an OpenEngine GenerateRequest into TensorRT-LLM request inputs.""" import asyncio -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass from typing import Any from urllib.parse import urlparse import grpc +import numpy as np +import torch +from PIL import Image from tensorrt_llm.conversation_params import ConversationParams -from tensorrt_llm.inputs.media_io import ImageMediaIO +from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams from .bindings import generation_pb2 -from .capabilities import media_placeholder, supported_guides, supported_media_modalities -from .errors import UnsupportedFeatureError +from .capabilities import MediaCapabilities, resolve_media_capabilities, supported_guides +from .errors import MediaUnavailableError, UnsupportedFeatureError # Cap on the number of stop conditions. Each one costs a KMP automaton per # output index and a scan of every generated token, so an unbounded list lets a @@ -28,10 +32,26 @@ # an unbounded amount of image data while preserving practical multi-image use. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 +# Encoded bodies are separately bounded by the 64 MiB gRPC message limit and +# the media loader's 200 MiB per-URL response limit. This cap covers expansion +# during decode across the entire request. +_MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" +_MediaLoader = Callable[[generation_pb2.MediaItem, int], Awaitable[Any]] + + +@dataclass(frozen=True) +class _MediaHandler: + """One fully implemented OpenEngine media mapping path.""" + + name: str + loader_factory: Callable[[Any], _MediaLoader] + decoded_bytes: Callable[[Any], int] + + def conversation_affinity_enabled(llm: Any) -> bool: """Whether TensorRT-LLM owns attention-DP placement for conversations.""" args = getattr(llm, "args", None) @@ -301,49 +321,113 @@ async def _load_image_item( try: return await image_io.async_load(value) + except MediaFetchError as error: + raise MediaUnavailableError( + f"media[{index}].{source} is temporarily unavailable: {error}" + ) from error except Exception as error: raise ValueError(f"media[{index}].{source} could not be loaded: {error}") from error +def _image_loader_factory(llm: Any) -> _MediaLoader: + image_io = _image_io_for_model(llm) + + async def load(item: generation_pb2.MediaItem, index: int) -> Any: + return await _load_image_item(item, image_io, index) + + return load + + +def _decoded_image_bytes(image: Any) -> int: + """Return retained bytes for a decoded image representation.""" + if isinstance(image, torch.Tensor): + return image.numel() * image.element_size() + if isinstance(image, np.ndarray): + return image.nbytes + if isinstance(image, Image.Image): + return image.width * image.height * len(image.getbands()) + # Custom model media types cannot be sized generically. The built-in image + # handler returns one of the representations above. + return 0 + + +_MEDIA_HANDLERS: dict[int, _MediaHandler] = { + generation_pb2.MODALITY_IMAGE: _MediaHandler( + "image", _image_loader_factory, _decoded_image_bytes + ), +} + + +def media_capabilities_for_model(llm: Any) -> MediaCapabilities: + """Resolve the loaded model against the request handlers in this module.""" + return resolve_media_capabilities(llm, {handler.name for handler in _MEDIA_HANDLERS.values()}) + + +def _modality_name(modality: int) -> str: + try: + return generation_pb2.Modality.Name(modality).removeprefix("MODALITY_").lower() + except ValueError: + return f"unknown ({modality})" + + async def _input_from_request( - request: generation_pb2.GenerateRequest, llm: Any + request: generation_pb2.GenerateRequest, + llm: Any, + media_capabilities: MediaCapabilities | None = None, ) -> str | dict[str, Any]: - """Build LLM API inputs, resolving image media off the event loop.""" + """Build LLM API inputs, resolving media off the event loop.""" inputs = _base_input_from_request(request) if not request.media: return inputs - supported = supported_media_modalities(llm) + if len(request.media) > _MAX_MEDIA_ITEMS: + raise ValueError( + f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " + f"got {len(request.media)}" + ) + + capabilities = media_capabilities or media_capabilities_for_model(llm) + if not capabilities.modalities: + raise UnsupportedFeatureError("media input is not supported by the loaded model") + + resolved_handlers: list[_MediaHandler] = [] for index, item in enumerate(request.media): if item.modality == generation_pb2.MODALITY_UNSPECIFIED: raise ValueError(f"media[{index}].modality must be set") - if item.modality != generation_pb2.MODALITY_IMAGE: - modality = generation_pb2.Modality.Name(item.modality).removeprefix("MODALITY_").lower() - supported_text = ", ".join(sorted(supported)) or "none" + handler = _MEDIA_HANDLERS.get(item.modality) + if handler is None or handler.name not in capabilities.modalities: + modality = _modality_name(item.modality) + supported_text = ", ".join(sorted(capabilities.modalities)) or "none" raise UnsupportedFeatureError( f"multimodal modality '{modality}' is not supported by the loaded model; " f"supported modalities: {supported_text}" ) - if "image" not in supported: - raise UnsupportedFeatureError("image input is not supported by the loaded model") - if len(request.media) > _MAX_MEDIA_ITEMS: - raise ValueError( - f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " - f"got {len(request.media)}" - ) + resolved_handlers.append(handler) - image_io = _image_io_for_model(llm) + loaders = {handler.name: handler.loader_factory(llm) for handler in set(resolved_handlers)} semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) + decoded_bytes = 0 - async def load_bounded(item: generation_pb2.MediaItem, index: int) -> Any: + async def load_bounded( + item: generation_pb2.MediaItem, index: int, handler: _MediaHandler + ) -> Any: + nonlocal decoded_bytes async with semaphore: - return await _load_image_item(item, image_io, index) + decoded = await loaders[handler.name](item, index) + decoded_bytes += handler.decoded_bytes(decoded) + if decoded_bytes > _MAX_TOTAL_DECODED_MEDIA_BYTES: + raise ValueError( + "decoded media exceeds the per-request limit of " + f"{_MAX_TOTAL_DECODED_MEDIA_BYTES} bytes" + ) + return decoded tasks = [ - asyncio.create_task(load_bounded(item, index)) for index, item in enumerate(request.media) + asyncio.create_task(load_bounded(item, index, handler)) + for index, (item, handler) in enumerate(zip(request.media, resolved_handlers)) ] try: - images = await asyncio.gather(*tasks) + decoded_media = await asyncio.gather(*tasks) except (Exception, asyncio.CancelledError): # `gather` propagates the first failure without cancelling siblings. # Explicitly stop any in-flight URL fetches/decodes before returning @@ -352,18 +436,27 @@ async def load_bounded(item: generation_pb2.MediaItem, index: int) -> Any: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) raise + grouped_media: dict[str, list[Any]] = {} + grouped_uuids: dict[str, list[str | None]] = {} + item_order: list[dict[str, Any]] = [] + for item, handler, decoded in zip(request.media, resolved_handlers, decoded_media): + modality_items = grouped_media.setdefault(handler.name, []) + modality_index = len(modality_items) + modality_items.append(decoded) + grouped_uuids.setdefault(handler.name, []).append(item.uuid or None) + item_order.append( + { + "modality": handler.name, + "index": modality_index, + "placeholder": capabilities.placeholder(handler.name, modality_index + 1), + } + ) + mapped: dict[str, Any] = {"prompt": inputs} if isinstance(inputs, str) else dict(inputs) - mapped["multi_modal_data"] = {"image": images} + mapped["multi_modal_data"] = grouped_media if any(item.uuid for item in request.media): - mapped["multi_modal_uuids"] = {"image": [item.uuid or None for item in request.media]} - mapped["mm_item_order"] = [ - { - "modality": "image", - "index": index, - "placeholder": media_placeholder(llm, "image", index + 1), - } - for index in range(len(request.media)) - ] + mapped["multi_modal_uuids"] = grouped_uuids + mapped["mm_item_order"] = item_order return mapped diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index 6bc18ffeb6a4..f1547b65a112 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -25,13 +25,14 @@ Reservation, ) from .disagg import disaggregated_params_from_request -from .errors import AbortFailedError, UnsupportedFeatureError +from .errors import AbortFailedError, MediaUnavailableError, UnsupportedFeatureError from .formatting import _engine_error_response, _stop_texts from .request_mapping import ( _input_from_request, _metadata_from_context, conversation_affinity_enabled, conversation_params_from_request, + media_capabilities_for_model, sampling_params_from_request, ) from .streaming import ( @@ -62,9 +63,15 @@ def __init__( # context worker uses. The actual transfer is driven by opaque_state. self._kv_transfer_backend = kv_transfer_backend self._guided_backend = llm.args.guided_decoding_backend + self._media_capabilities = media_capabilities_for_model(llm) self.subagent_affinity_auth_key = subagent_affinity_auth_key self._active_requests: dict[str, Any] = {} + @property + def supported_media_modalities(self) -> frozenset[str]: + """Media capabilities cached from the loaded model and mapper.""" + return self._media_capabilities.modalities + def active_request_count(self) -> int: """Requests accepted and not yet completed.""" return len(self._active_requests) @@ -175,7 +182,7 @@ async def _generate( if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: raise UnsupportedFeatureError("prefix cache bypass is not supported") - inputs = await _input_from_request(request, self._llm) + inputs = await _input_from_request(request, self._llm, self._media_capabilities) sampling_params = sampling_params_from_request(request, self._guided_backend) trace_headers, target_dp_rank, affinity_headers = _metadata_from_context(context) conversation_params = conversation_params_from_request(request) @@ -230,6 +237,9 @@ async def _generate( except UnsupportedFeatureError as error: await context.abort(grpc.StatusCode.UNIMPLEMENTED, str(error)) return + except MediaUnavailableError as error: + await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) + return except ValueError as error: await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(error)) return diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index e8ab41d4dccf..484fdf5b6f70 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -45,6 +45,10 @@ from tensorrt_llm.logger import logger +class MediaFetchError(RuntimeError): + """A transient failure while fetching remote media.""" + + def rgba_to_rgb( image: Image.Image, background_color: Union[tuple[int, int, int], list[int]] = (255, 255, 255) ) -> Image.Image: @@ -161,7 +165,7 @@ def _validate_url(url: str) -> None: try: infos = socket.getaddrinfo(hostname, None, proto=socket.IPPROTO_TCP) except socket.gaierror as exc: - raise RuntimeError(f"Could not resolve hostname {hostname!r}") from exc + raise MediaFetchError(f"Could not resolve hostname {hostname!r}") from exc allow_private = _allow_private_urls() for _family, _type, _proto, _canon, sockaddr in infos: @@ -257,18 +261,29 @@ async def _safe_aiohttp_get( async def _fetch(fetch_session: aiohttp.ClientSession) -> bytes: current_url = url for _ in range(_MAX_REDIRECTS + 1): - async with fetch_session.get( - current_url, timeout=timeout, allow_redirects=False - ) as response: - if response.status in _REDIRECT_STATUSES: - redirect_url = response.headers.get("Location", "") - current_url = urljoin(current_url, redirect_url) - await asyncio.to_thread(_validate_url, current_url) - continue - maybe_coro = response.raise_for_status() - if asyncio.iscoroutine(maybe_coro): - await maybe_coro - return await _read_aiohttp_content(response) + try: + async with fetch_session.get( + current_url, timeout=timeout, allow_redirects=False + ) as response: + if response.status in _REDIRECT_STATUSES: + redirect_url = response.headers.get("Location", "") + current_url = urljoin(current_url, redirect_url) + await asyncio.to_thread(_validate_url, current_url) + continue + maybe_coro = response.raise_for_status() + if asyncio.iscoroutine(maybe_coro): + await maybe_coro + return await _read_aiohttp_content(response) + except MediaFetchError: + raise + except aiohttp.ClientResponseError as error: + if error.status in (408, 429) or error.status >= 500: + raise MediaFetchError( + f"remote media server returned HTTP {error.status}" + ) from error + raise + except (aiohttp.ClientError, asyncio.TimeoutError) as error: + raise MediaFetchError(f"remote media fetch failed: {error}") from error raise RuntimeError("Too many redirects") if session is not None: diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index f7de8842a222..79b56c77f9d5 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -19,6 +19,7 @@ model_pb2, server_pb2, ) +from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.llmapi import KVEventsConfig # noqa: E402 @@ -172,12 +173,14 @@ async def test_model_info_reports_the_context_window_and_capabilities(): @pytest.mark.asyncio async def test_model_info_advertises_loaded_image_support(monkeypatch): + inference = _inference() monkeypatch.setattr( - "tensorrt_llm.grpc.openengine.control.supported_media_modalities", - lambda llm: frozenset({"image"}), + inference, + "_media_capabilities", + MediaCapabilities({"image": ""}), ) - info = await _servicer().GetModelInfo( + info = await _servicer(inference=inference).GetModelInfo( model_pb2.GetModelInfoRequest(model=MODEL), FakeServicerContext() ) diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index a4bd537e29f5..43ff2297861b 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -8,6 +8,7 @@ from types import SimpleNamespace from typing import Any +import numpy as np import pytest grpc = pytest.importorskip( @@ -17,9 +18,9 @@ import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 -from tensorrt_llm.grpc.openengine.capabilities import supported_media_modalities # noqa: E402 from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, @@ -160,16 +161,75 @@ def is_valid(model_type, modality): def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): llm, seen_model_types = image_llm - assert supported_media_modalities(llm) == {"image"} + assert request_mapping.media_capabilities_for_model(llm).modalities == {"image"} assert seen_model_types == [_ImageConfig.model_type] assert ( - supported_media_modalities( + request_mapping.media_capabilities_for_model( SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) - ) + ).modalities == set() ) +def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): + llm, _ = image_llm + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: "", + ) + + with pytest.raises(RuntimeError, match="invalid multimodal placeholder"): + request_mapping.media_capabilities_for_model(llm) + + +@pytest.mark.asyncio +async def test_handler_table_drives_capability_and_mapping(image_llm, monkeypatch): + llm, _ = image_llm + + def video_loader_factory(_llm): + async def load(item, index): + return f"video-{index}:{bytes(item.raw_bytes).decode()}" + + return load + + monkeypatch.setitem( + request_mapping._MEDIA_HANDLERS, + generation_pb2.MODALITY_VIDEO, + request_mapping._MediaHandler("video", video_loader_factory, lambda decoded: len(decoded)), + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: model_type == _ImageConfig.model_type + and modality in {"image", "video"}, + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + + capabilities = request_mapping.media_capabilities_for_model(llm) + assert capabilities.modalities == {"image", "video"} + + request = generation_pb2.GenerateRequest( + prompt="video", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"clip", + ) + ], + ) + mapped = await request_mapping._input_from_request(request, llm, capabilities) + + assert mapped["multi_modal_data"] == {"video": ["video-0:clip"]} + assert mapped["mm_item_order"] == [ + {"modality": "video", "index": 0, "placeholder": ""} + ] + + @pytest.mark.asyncio async def test_text_only_input_mapping_is_unchanged(): llm = SimpleNamespace() @@ -338,7 +398,9 @@ def create(cls, defaults, runtime): prompt="too many images", media=[ generation_pb2.MediaItem( - modality=generation_pb2.MODALITY_IMAGE, + # Unsupported on purpose: the count guard must win before the + # per-item modality loop. + modality=generation_pb2.MODALITY_VIDEO, raw_bytes=b"image", ) for _ in range(request_mapping._MAX_MEDIA_ITEMS + 1) @@ -349,6 +411,35 @@ def create(cls, defaults, runtime): await request_mapping._input_from_request(request, llm) +@pytest.mark.asyncio +async def test_total_decoded_image_bytes_are_bounded(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return np.zeros((2, 2, 3), dtype=np.uint8) + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(request_mapping, "_MAX_TOTAL_DECODED_MEDIA_BYTES", 20) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + for _ in range(2) + ], + ) + + with pytest.raises(ValueError, match="decoded media exceeds"): + await request_mapping._input_from_request(request, llm) + + @pytest.mark.asyncio async def test_image_mapping_supports_token_ids(image_llm, monkeypatch): llm, _ = image_llm @@ -460,6 +551,43 @@ async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): assert malformed_context.abort_code == grpc.StatusCode.INVALID_ARGUMENT +@pytest.mark.asyncio +async def test_generate_maps_transient_media_fetch_to_unavailable(image_llm, monkeypatch): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load(self, source): + raise MediaFetchError("connection timed out") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + context = FakeServicerContext() + request = generation_pb2.GenerateRequest( + request_id="remote-image", + model="image-model", + prompt="image", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="https://example.com/image.png", + ) + ], + ) + + with pytest.raises(AbortError): + async for _ in servicer.Generate(request, context): + pass + assert context.abort_code == grpc.StatusCode.UNAVAILABLE + + @pytest.mark.asyncio @pytest.mark.parametrize( ("modality", "name"), diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index 9bff8139754c..c2d0421ee5bc 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -26,6 +26,7 @@ from tensorrt_llm.inputs.media_io import ( _MAX_REDIRECTS, + MediaFetchError, _safe_aiohttp_get, _safe_request_get, _validate_url, @@ -129,7 +130,7 @@ def test_rejects_multicast(self, _): side_effect=socket.gaierror("Name or service not known"), ) def test_rejects_unresolvable_hostname(self, _): - with pytest.raises(RuntimeError, match="Could not resolve"): + with pytest.raises(MediaFetchError, match="Could not resolve"): _validate_url("http://this.does.not.exist.invalid/") @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) @@ -346,6 +347,24 @@ class TestSafeAiohttpGet: def _run(self, coro): return asyncio.run(coro) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) + def test_maps_timeout_to_fetch_error(self, _): + class _FakeSession: + def get(self, url, **kwargs): + raise asyncio.TimeoutError("timed out") + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + pass + + with patch( + "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + ): + with pytest.raises(MediaFetchError, match="fetch failed"): + self._run(_safe_aiohttp_get("http://example.com/file.bin")) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_reads_response_chunks_to_eof(self, _): class _FakeResponse: From 7baf0835ec68f7f446e56fc98d6cabdd028e1fa5 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 16:59:18 -0700 Subject: [PATCH 03/41] style(openengine): tighten media comments Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 28 ++++++------------- .../grpc/openengine/request_mapping.py | 23 ++++----------- tensorrt_llm/grpc/openengine/servicer.py | 2 +- 3 files changed, 14 insertions(+), 39 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 6affcf574747..1a7f76bbd90b 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -1,11 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""What this engine can actually do, defined once for Generate and Control. - -Control advertises from this table and Generate enforces it, so a client can -never discover a capability that Generate then rejects. -""" +"""Shared capability resolution for OpenEngine Control and Generate.""" from dataclasses import dataclass from types import MappingProxyType @@ -29,7 +25,7 @@ @dataclass(frozen=True) class MediaCapabilities: - """Immutable media capability snapshot shared by Control and Generate.""" + """Media support shared by Control and Generate.""" placeholder_templates: Mapping[str, str] @@ -40,7 +36,7 @@ def __post_init__(self) -> None: @property def modalities(self) -> frozenset[str]: - """Modalities the loaded model and request mapper both support.""" + """Modalities supported by both the model and request mapper.""" return frozenset(self.placeholder_templates) def placeholder(self, modality: str, item_number: int) -> str: @@ -73,8 +69,7 @@ def _loaded_model_type(llm: Any) -> str: if config is None: return "" - # Composite Hugging Face configs can delegate the instance attribute to - # their text config. The class attribute retains the top-level registry key. + # The class attribute preserves the top-level type for composite configs. model_type = getattr(type(config), "model_type", None) or getattr(config, "model_type", "") return model_type if isinstance(model_type, str) else "" @@ -82,14 +77,9 @@ def _loaded_model_type(llm: Any) -> str: def resolve_media_capabilities( llm: Any, implemented_modalities: Collection[str] ) -> MediaCapabilities: - """Resolve media support implemented by both the mapper and loaded model. - - A registry entry alone describes what this TensorRT-LLM build knows about; - it does not prove that the loaded engine instantiated a multimodal input - processor. Requiring both prevents a text-only model from advertising - media merely because another image model is registered in the process. The - request mapper supplies ``implemented_modalities`` from its handler table, - so discovery cannot advertise a modality without a request implementation. + """Intersect model support with implemented request handlers. + + A live multimodal processor prevents registry-only false positives. """ input_processor = getattr(llm, "input_processor", None) if not isinstance(input_processor, BaseMultimodalInputProcessor): @@ -103,9 +93,7 @@ def resolve_media_capabilities( for modality in implemented_modalities if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) } - # Validate templates once at startup. A malformed registry entry is a - # server configuration failure, not a client request error discovered only - # after Control has advertised the modality. + # Fail startup before Control advertises a malformed template. for modality, template in placeholders.items(): try: template.format(1) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 28c948ba89d4..0f4ab8603c74 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -27,14 +27,10 @@ # output index and a scan of every generated token, so an unbounded list lets a # single request tax the event loop for the whole stream. _MAX_STOP_CONDITIONS = 64 -# Media decoding can retain both encoded and decoded image buffers. Keep one -# request from scheduling an unbounded number of fetch/decode tasks or holding -# an unbounded amount of image data while preserving practical multi-image use. +# Bound per-request decode work and retained image data. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 -# Encoded bodies are separately bounded by the 64 MiB gRPC message limit and -# the media loader's 200 MiB per-URL response limit. This cap covers expansion -# during decode across the entire request. +# Encoded inputs have separate gRPC and remote-loader caps. _MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -45,8 +41,6 @@ @dataclass(frozen=True) class _MediaHandler: - """One fully implemented OpenEngine media mapping path.""" - name: str loader_factory: Callable[[Any], _MediaLoader] decoded_bytes: Callable[[Any], int] @@ -285,7 +279,6 @@ def _base_input_from_request( def _image_io_for_model(llm: Any) -> ImageMediaIO: - """Create image I/O using the loaded model's preferred decode format.""" input_processor = getattr(llm, "input_processor", None) defaults = None if isinstance(input_processor, BaseMultimodalInputProcessor): @@ -297,7 +290,6 @@ def _image_io_for_model(llm: Any) -> ImageMediaIO: async def _load_image_item( item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int ) -> Any: - """Resolve and decode one OpenEngine image item.""" source = item.WhichOneof("source") if source is None: raise ValueError(f"media[{index}] must set exactly one source") @@ -339,15 +331,13 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: def _decoded_image_bytes(image: Any) -> int: - """Return retained bytes for a decoded image representation.""" if isinstance(image, torch.Tensor): return image.numel() * image.element_size() if isinstance(image, np.ndarray): return image.nbytes if isinstance(image, Image.Image): return image.width * image.height * len(image.getbands()) - # Custom model media types cannot be sized generically. The built-in image - # handler returns one of the representations above. + # ImageMediaIO currently returns only the representations above. return 0 @@ -359,7 +349,6 @@ def _decoded_image_bytes(image: Any) -> int: def media_capabilities_for_model(llm: Any) -> MediaCapabilities: - """Resolve the loaded model against the request handlers in this module.""" return resolve_media_capabilities(llm, {handler.name for handler in _MEDIA_HANDLERS.values()}) @@ -375,7 +364,7 @@ async def _input_from_request( llm: Any, media_capabilities: MediaCapabilities | None = None, ) -> str | dict[str, Any]: - """Build LLM API inputs, resolving media off the event loop.""" + """Build LLM API inputs without blocking the event loop on media.""" inputs = _base_input_from_request(request) if not request.media: return inputs @@ -429,9 +418,7 @@ async def load_bounded( try: decoded_media = await asyncio.gather(*tasks) except (Exception, asyncio.CancelledError): - # `gather` propagates the first failure without cancelling siblings. - # Explicitly stop any in-flight URL fetches/decodes before returning - # the request error so their memory and connections do not linger. + # gather leaves siblings running after the first failure. for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index f1547b65a112..f2dabe368a32 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -69,7 +69,7 @@ def __init__( @property def supported_media_modalities(self) -> frozenset[str]: - """Media capabilities cached from the loaded model and mapper.""" + """Cached media modalities.""" return self._media_capabilities.modalities def active_request_count(self) -> int: From ab07b1649f638924908aed565722f4dd0f8547f8 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 17:25:55 -0700 Subject: [PATCH 04/41] fix(openengine): harden decoded media accounting Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 3 +++ .../grpc/openengine/request_mapping.py | 25 ++++++++++------- .../grpc/openengine/test_multimodal.py | 27 +++++++++++++++++++ 3 files changed, 45 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 1a7f76bbd90b..4f553a4b8bdc 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -34,6 +34,9 @@ def __post_init__(self) -> None: self, "placeholder_templates", MappingProxyType(dict(self.placeholder_templates)) ) + def __hash__(self) -> int: + return hash(frozenset(self.placeholder_templates.items())) + @property def modalities(self) -> frozenset[str]: """Modalities supported by both the model and request mapper.""" diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 0f4ab8603c74..cfd7c17bf2ee 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -10,9 +10,6 @@ from urllib.parse import urlparse import grpc -import numpy as np -import torch -from PIL import Image from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError @@ -30,7 +27,8 @@ # Bound per-request decode work and retained image data. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 -# Encoded inputs have separate gRPC and remote-loader caps. +# Checked after each decode; concurrency bounds peak overshoot. Encoded inputs +# have separate gRPC and remote-loader caps. _MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -331,14 +329,17 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: def _decoded_image_bytes(image: Any) -> int: + import numpy as np + import torch + from PIL import Image + if isinstance(image, torch.Tensor): return image.numel() * image.element_size() if isinstance(image, np.ndarray): return image.nbytes if isinstance(image, Image.Image): return image.width * image.height * len(image.getbands()) - # ImageMediaIO currently returns only the representations above. - return 0 + raise TypeError(f"unsupported decoded image type: {type(image).__name__}") _MEDIA_HANDLERS: dict[int, _MediaHandler] = { @@ -376,7 +377,8 @@ async def _input_from_request( ) capabilities = media_capabilities or media_capabilities_for_model(llm) - if not capabilities.modalities: + modalities = capabilities.modalities + if not modalities: raise UnsupportedFeatureError("media input is not supported by the loaded model") resolved_handlers: list[_MediaHandler] = [] @@ -384,16 +386,19 @@ async def _input_from_request( if item.modality == generation_pb2.MODALITY_UNSPECIFIED: raise ValueError(f"media[{index}].modality must be set") handler = _MEDIA_HANDLERS.get(item.modality) - if handler is None or handler.name not in capabilities.modalities: + if handler is None or handler.name not in modalities: modality = _modality_name(item.modality) - supported_text = ", ".join(sorted(capabilities.modalities)) or "none" + supported_text = ", ".join(sorted(modalities)) or "none" raise UnsupportedFeatureError( f"multimodal modality '{modality}' is not supported by the loaded model; " f"supported modalities: {supported_text}" ) resolved_handlers.append(handler) - loaders = {handler.name: handler.loader_factory(llm) for handler in set(resolved_handlers)} + loaders: dict[str, _MediaLoader] = {} + for handler in resolved_handlers: + if handler.name not in loaders: + loaders[handler.name] = handler.loader_factory(llm) semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) decoded_bytes = 0 diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 43ff2297861b..f9aa7985f4c6 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -18,6 +18,7 @@ import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 @@ -29,6 +30,15 @@ pytestmark = pytest.mark.cpu_only +def _ignore_synthetic_decoded_size(monkeypatch) -> None: + handler = request_mapping._MEDIA_HANDLERS[generation_pb2.MODALITY_IMAGE] + monkeypatch.setitem( + request_mapping._MEDIA_HANDLERS, + generation_pb2.MODALITY_IMAGE, + request_mapping._MediaHandler(handler.name, handler.loader_factory, lambda _decoded: 0), + ) + + class _ImageConfig: # Composite configs can expose their text model type on the instance; the # top-level class attribute is the key used by the multimodal registry. @@ -171,6 +181,19 @@ def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): ) +def test_media_capabilities_are_hashable(): + first = MediaCapabilities({"image": ""}) + second = MediaCapabilities({"image": ""}) + + assert first == second + assert hash(first) == hash(second) + + +def test_unknown_decoded_image_type_is_rejected(): + with pytest.raises(TypeError, match="unsupported decoded image type"): + request_mapping._decoded_image_bytes(object()) + + def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): llm, _ = image_llm monkeypatch.setattr( @@ -264,6 +287,7 @@ async def async_load_bytes(self, data): return f"decoded-bytes:{data.decode()}" monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) request = generation_pb2.GenerateRequest( prompt=" describe", media=[ @@ -326,6 +350,7 @@ async def async_load_bytes(self, data): type(self).active -= 1 monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) request = generation_pb2.GenerateRequest( prompt="images", media=[ @@ -453,6 +478,7 @@ async def async_load_bytes(self, data): return "decoded" monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) request = generation_pb2.GenerateRequest( token_ids=generation_pb2.TokenIds(ids=[10, 20]), media=[ @@ -487,6 +513,7 @@ async def async_load_bytes(self, data): return "decoded" monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) servicer = OpenEngineInferenceServicer(llm, model="image-model") request = generation_pb2.GenerateRequest( request_id="image-request", From b44c58f9e4db4d9c6a63978411fc429ed23cc579 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 20:37:00 -0700 Subject: [PATCH 05/41] feat(openengine): support video inputs Signed-off-by: krishung5 --- .../grpc/openengine/request_mapping.py | 66 ++++++-- .../grpc/openengine/test_multimodal.py | 150 ++++++++++++++++-- .../openengine/test_multimodal_conformance.py | 48 +++++- 3 files changed, 233 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index cfd7c17bf2ee..ba7a89956c94 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -12,7 +12,8 @@ import grpc from tensorrt_llm.conversation_params import ConversationParams -from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError +from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError, VideoMediaIO +from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams @@ -276,17 +277,17 @@ def _base_input_from_request( raise ValueError("exactly one of prompt or token_ids must be set") -def _image_io_for_model(llm: Any) -> ImageMediaIO: +def _media_io_for_model(llm: Any, modality: str, media_io_type: type[Any]) -> Any: input_processor = getattr(llm, "input_processor", None) defaults = None if isinstance(input_processor, BaseMultimodalInputProcessor): preferences = input_processor.get_preferred_media_io_kwargs() or {} - defaults = preferences.get("image") - return ImageMediaIO.create(defaults, None) + defaults = preferences.get(modality) + return media_io_type.create(defaults, None) -async def _load_image_item( - item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int +async def _load_media_item( + item: generation_pb2.MediaItem, media_io: Any, index: int, modality: str ) -> Any: source = item.WhichOneof("source") if source is None: @@ -300,17 +301,19 @@ async def _load_image_item( raise ValueError(f"media[{index}].data_uri must use the data scheme") value = item.data_uri elif source == "raw_bytes": - if item.mime_type and not item.mime_type.lower().startswith("image/"): - raise ValueError(f"media[{index}].mime_type must be an image type for raw image bytes") + if item.mime_type and not item.mime_type.lower().startswith(f"{modality}/"): + raise ValueError( + f"media[{index}].mime_type must start with {modality}/ for raw {modality} bytes" + ) try: - return await image_io.async_load_bytes(bytes(item.raw_bytes)) + return await media_io.async_load_bytes(bytes(item.raw_bytes)) except Exception as error: raise ValueError(f"media[{index}].raw_bytes could not be decoded: {error}") from error else: raise ValueError(f"media[{index}] has an unknown source '{source}'") try: - return await image_io.async_load(value) + return await media_io.async_load(value) except MediaFetchError as error: raise MediaUnavailableError( f"media[{index}].{source} is temporarily unavailable: {error}" @@ -320,10 +323,19 @@ async def _load_image_item( def _image_loader_factory(llm: Any) -> _MediaLoader: - image_io = _image_io_for_model(llm) + image_io = _media_io_for_model(llm, "image", ImageMediaIO) + + async def load(item: generation_pb2.MediaItem, index: int) -> Any: + return await _load_media_item(item, image_io, index, "image") + + return load + + +def _video_loader_factory(llm: Any) -> _MediaLoader: + video_io = _media_io_for_model(llm, "video", VideoMediaIO) async def load(item: generation_pb2.MediaItem, index: int) -> Any: - return await _load_image_item(item, image_io, index) + return await _load_media_item(item, video_io, index, "video") return load @@ -342,10 +354,40 @@ def _decoded_image_bytes(image: Any) -> int: raise TypeError(f"unsupported decoded image type: {type(image).__name__}") +def _decoded_audio_data_bytes(audio: AudioData) -> int: + import numpy as np + import torch + + if isinstance(audio.samples, torch.Tensor): + return audio.samples.numel() * audio.samples.element_size() + if isinstance(audio.samples, np.ndarray): + return audio.samples.nbytes + raise TypeError(f"unsupported decoded audio type: {type(audio.samples).__name__}") + + +def _decoded_video_bytes(video: Any) -> int: + import numpy as np + + if not isinstance(video, VideoData): + raise TypeError(f"unsupported decoded video type: {type(video).__name__}") + if isinstance(video.frames, np.ndarray): + total = video.frames.nbytes + elif isinstance(video.frames, list): + total = sum(_decoded_image_bytes(frame) for frame in video.frames) + else: + raise TypeError(f"unsupported decoded video frames type: {type(video.frames).__name__}") + if video.audio is not None: + total += _decoded_audio_data_bytes(video.audio) + return total + + _MEDIA_HANDLERS: dict[int, _MediaHandler] = { generation_pb2.MODALITY_IMAGE: _MediaHandler( "image", _image_loader_factory, _decoded_image_bytes ), + generation_pb2.MODALITY_VIDEO: _MediaHandler( + "video", _video_loader_factory, _decoded_video_bytes + ), } diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index f9aa7985f4c6..29916c42bb7c 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for OpenEngine image input mapping and discovery.""" +"""Unit tests for OpenEngine multimodal input mapping and discovery.""" import asyncio from collections.abc import AsyncIterator, Sequence @@ -22,6 +22,7 @@ from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 +from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData # noqa: E402 from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, @@ -74,7 +75,10 @@ def call_with_text_prompt(self, inputs, sampling_params): raise NotImplementedError def get_preferred_media_io_kwargs(self): - return {"image": {"format": "np"}} + return { + "image": {"format": "np"}, + "video": {"format": "np", "num_frames": 4}, + } class _Tokenizer: @@ -168,11 +172,33 @@ def is_valid(model_type, modality): ) +@pytest.fixture +def image_video_llm(monkeypatch): + config = _ImageConfig() + input_processor = _ImageInputProcessor(config) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: ( + model_type == _ImageConfig.model_type and modality in {"image", "video"} + ), + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + return SimpleNamespace( + _hf_model_config=config, + input_processor=input_processor, + ) + + def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): llm, seen_model_types = image_llm assert request_mapping.media_capabilities_for_model(llm).modalities == {"image"} - assert seen_model_types == [_ImageConfig.model_type] + assert seen_model_types == [_ImageConfig.model_type] * 2 assert ( request_mapping.media_capabilities_for_model( SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) @@ -194,6 +220,18 @@ def test_unknown_decoded_image_type_is_rejected(): request_mapping._decoded_image_bytes(object()) +def test_decoded_video_bytes_include_frames_and_audio(): + video = VideoData( + frames=np.zeros((2, 3, 4, 3), dtype=np.uint8), + metadata={}, + audio=AudioData(samples=np.zeros(8, dtype=np.float32), sample_rate=16000), + ) + + assert request_mapping._decoded_video_bytes(video) == 2 * 3 * 4 * 3 + 8 * 4 + with pytest.raises(TypeError, match="unsupported decoded video type"): + request_mapping._decoded_video_bytes(object()) + + def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): llm, _ = image_llm monkeypatch.setattr( @@ -224,8 +262,9 @@ async def load(item, index): monkeypatch.setattr( MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", - lambda model_type, modality: model_type == _ImageConfig.model_type - and modality in {"image", "video"}, + lambda model_type, modality: ( + model_type == _ImageConfig.model_type and modality in {"image", "video"} + ), ) monkeypatch.setattr( MULTIMODAL_PLACEHOLDER_REGISTRY, @@ -328,6 +367,75 @@ async def async_load_bytes(self, data): ] +@pytest.mark.asyncio +async def test_image_and_video_sources_preserve_global_order(image_video_llm, monkeypatch): + created_with = [] + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append(("image", defaults, runtime)) + return cls() + + async def async_load_bytes(self, data): + return np.full((1, 1, 3), data[0], dtype=np.uint8) + + class _VideoIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append(("video", defaults, runtime)) + return cls() + + async def async_load(self, source): + return VideoData( + frames=np.zeros((1, 1, 1, 3), dtype=np.uint8), + metadata={"source": source}, + ) + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(request_mapping, "VideoMediaIO", _VideoIo) + request = generation_pb2.GenerateRequest( + prompt="mixed media", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + url="https://example.com/first.mp4", + uuid="video-id", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"i", + mime_type="image/png", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + data_uri="data:video/mp4;base64,dmlkZW8=", + ), + ], + ) + + mapped = await request_mapping._input_from_request(request, image_video_llm) + + assert created_with == [ + ("video", {"format": "np", "num_frames": 4}, None), + ("image", {"format": "np"}, None), + ] + assert [video.metadata["source"] for video in mapped["multi_modal_data"]["video"]] == [ + "https://example.com/first.mp4", + "data:video/mp4;base64,dmlkZW8=", + ] + assert mapped["multi_modal_data"]["image"][0].shape == (1, 1, 3) + assert mapped["multi_modal_uuids"] == { + "video": ["video-id", None], + "image": [None], + } + assert mapped["mm_item_order"] == [ + {"modality": "video", "index": 0, "placeholder": ""}, + {"modality": "image", "index": 0, "placeholder": ""}, + {"modality": "video", "index": 1, "placeholder": ""}, + ] + + @pytest.mark.asyncio async def test_image_loading_has_a_bounded_concurrency(image_llm, monkeypatch): llm, _ = image_llm @@ -616,21 +724,14 @@ async def async_load(self, source): @pytest.mark.asyncio -@pytest.mark.parametrize( - ("modality", "name"), - [ - (generation_pb2.MODALITY_VIDEO, "video"), - (generation_pb2.MODALITY_AUDIO, "audio"), - ], -) -async def test_video_and_audio_are_explicitly_unsupported(image_llm, modality, name): +async def test_audio_is_explicitly_unsupported(image_llm): llm, _ = image_llm request = generation_pb2.GenerateRequest( prompt="media", - media=[generation_pb2.MediaItem(modality=modality, raw_bytes=b"data")], + media=[generation_pb2.MediaItem(modality=generation_pb2.MODALITY_AUDIO, raw_bytes=b"data")], ) - with pytest.raises(UnsupportedFeatureError, match=name): + with pytest.raises(UnsupportedFeatureError, match="audio"): await request_mapping._input_from_request(request, llm) @@ -674,7 +775,7 @@ async def test_image_is_rejected_for_a_text_only_loaded_model(): raw_bytes=b"image", mime_type="video/mp4", ), - "must be an image type", + "must start with image/", ), ], ) @@ -684,3 +785,20 @@ async def test_malformed_image_items_are_rejected(image_llm, item, message): with pytest.raises(ValueError, match=message): await request_mapping._input_from_request(request, llm) + + +@pytest.mark.asyncio +async def test_raw_video_mime_type_is_validated(image_video_llm): + request = generation_pb2.GenerateRequest( + prompt="video", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"video", + mime_type="image/png", + ) + ], + ) + + with pytest.raises(ValueError, match="must start with video/"): + await request_mapping._input_from_request(request, image_video_llm) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index e97c9a708aac..1e88ed330f66 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Image conformance tests against a real TensorRT-LLM VLM engine.""" +"""Multimodal conformance tests against a real TensorRT-LLM VLM engine.""" import asyncio import base64 @@ -44,7 +44,7 @@ def vlm(): engine.shutdown() -def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm): +def test_image_and_video_capabilities_agree_with_generate(vlm, tmp_path): inference = OpenEngineInferenceServicer(vlm, MODEL_NAME) control = OpenEngineControlServicer(vlm, MODEL_NAME, inference) info = asyncio.run( @@ -55,7 +55,7 @@ def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm): ) assert info.supports_multimodal is True - assert list(info.extra["trtllm_supported_modalities"]) == ["image"] + assert list(info.extra["trtllm_supported_modalities"]) == ["image", "video"] prompt = ( "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" @@ -91,3 +91,45 @@ async def collect(): assert finished[0].usage.total_tokens == ( finished[0].usage.prompt_tokens + finished[0].usage.completion_tokens ) + + cv2 = pytest.importorskip("cv2", reason="video conformance requires OpenCV") + import numpy as np + + video_path = tmp_path / "colors.mp4" + writer = cv2.VideoWriter(str(video_path), cv2.VideoWriter_fourcc(*"mp4v"), 2.0, (32, 32)) + if not writer.isOpened(): + pytest.skip("OpenCV has no MP4 encoder") + for value in (0, 64, 128, 255): + writer.write(np.full((32, 32, 3), value, dtype=np.uint8)) + writer.release() + + video_prompt = ( + "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" + "Describe this video briefly.<|im_end|>\n<|im_start|>assistant\n" + ) + video_request = generation_pb2.GenerateRequest( + request_id="openengine-video-conformance", + model=MODEL_NAME, + token_ids=generation_pb2.TokenIds(ids=vlm.tokenizer.encode(video_prompt)), + stopping=generation_pb2.StoppingOptions(max_tokens=4), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + data_uri="data:video/mp4;base64," + + base64.b64encode(video_path.read_bytes()).decode(), + uuid="colors", + ) + ], + ) + + async def collect_video(): + return [ + response async for response in inference.Generate(video_request, FakeServicerContext()) + ] + + video_responses = asyncio.run(collect_video()) + assert video_responses + assert not any(response.HasField("error") for response in video_responses) + video_finished = [response for response in video_responses if response.HasField("finished")] + assert len(video_finished) == 1 + assert video_finished[0].usage.prompt_tokens > len(video_request.token_ids.ids) From be67da06209e1780e5567f71ed67e45b0cc7766b Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 21:04:56 -0700 Subject: [PATCH 06/41] feat(openengine): support audio inputs Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 14 +- .../grpc/openengine/request_mapping.py | 20 ++- .../grpc/openengine/test_multimodal.py | 105 ++++++++++++++- .../openengine/test_multimodal_conformance.py | 126 +++++++++++++++++- 4 files changed, 248 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 4f553a4b8bdc..c1b7cbdd1e27 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -64,11 +64,7 @@ def supported_guides(backend: str | None) -> frozenset[str]: def _loaded_model_type(llm: Any) -> str: """Return the top-level model type for the loaded input processor.""" input_processor = getattr(llm, "input_processor", None) - config = ( - input_processor.config - if isinstance(input_processor, BaseMultimodalInputProcessor) - else getattr(llm, "_hf_model_config", None) - ) + config = getattr(input_processor, "config", None) or getattr(llm, "_hf_model_config", None) if config is None: return "" @@ -82,14 +78,16 @@ def resolve_media_capabilities( ) -> MediaCapabilities: """Intersect model support with implemented request handlers. - A live multimodal processor prevents registry-only false positives. + A live registered processor prevents registry-only false positives. """ input_processor = getattr(llm, "input_processor", None) - if not isinstance(input_processor, BaseMultimodalInputProcessor): - return MediaCapabilities({}) model_type = _loaded_model_type(llm) if not model_type: return MediaCapabilities({}) + if not isinstance(input_processor, BaseMultimodalInputProcessor): + registered_model_type = getattr(type(input_processor), "_registered_model_type", None) + if registered_model_type != model_type or not callable(input_processor): + return MediaCapabilities({}) placeholders = { modality: MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index ba7a89956c94..08300c39ae7c 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -12,7 +12,7 @@ import grpc from tensorrt_llm.conversation_params import ConversationParams -from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError, VideoMediaIO +from tensorrt_llm.inputs.media_io import AudioMediaIO, ImageMediaIO, MediaFetchError, VideoMediaIO from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams @@ -340,6 +340,15 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: return load +def _audio_loader_factory(llm: Any) -> _MediaLoader: + audio_io = _media_io_for_model(llm, "audio", AudioMediaIO) + + async def load(item: generation_pb2.MediaItem, index: int) -> Any: + return await _load_media_item(item, audio_io, index, "audio") + + return load + + def _decoded_image_bytes(image: Any) -> int: import numpy as np import torch @@ -365,6 +374,12 @@ def _decoded_audio_data_bytes(audio: AudioData) -> int: raise TypeError(f"unsupported decoded audio type: {type(audio.samples).__name__}") +def _decoded_audio_bytes(audio: Any) -> int: + if not isinstance(audio, tuple) or len(audio) != 2: + raise TypeError(f"unsupported decoded audio type: {type(audio).__name__}") + return _decoded_audio_data_bytes(AudioData(samples=audio[0], sample_rate=audio[1])) + + def _decoded_video_bytes(video: Any) -> int: import numpy as np @@ -388,6 +403,9 @@ def _decoded_video_bytes(video: Any) -> int: generation_pb2.MODALITY_VIDEO: _MediaHandler( "video", _video_loader_factory, _decoded_video_bytes ), + generation_pb2.MODALITY_AUDIO: _MediaHandler( + "audio", _audio_loader_factory, _decoded_audio_bytes + ), } diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 29916c42bb7c..09c567fc5513 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -81,6 +81,18 @@ def get_preferred_media_io_kwargs(self): } +class _AudioInputProcessor: + """Registered processor that does not use BaseMultimodalInputProcessor.""" + + _registered_model_type = _ImageConfig.model_type + + def __init__(self, config: _ImageConfig) -> None: + self.config = config + + def __call__(self, inputs, sampling_params): + raise NotImplementedError + + class _Tokenizer: def convert_ids_to_tokens( self, token_ids: Sequence[int], skip_special_tokens: bool = False @@ -194,11 +206,33 @@ def image_video_llm(monkeypatch): ) +@pytest.fixture +def image_video_audio_llm(monkeypatch): + config = _ImageConfig() + input_processor = _ImageInputProcessor(config) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: ( + model_type == _ImageConfig.model_type and modality in {"image", "video", "audio"} + ), + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + return SimpleNamespace( + _hf_model_config=config, + input_processor=input_processor, + ) + + def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): llm, seen_model_types = image_llm assert request_mapping.media_capabilities_for_model(llm).modalities == {"image"} - assert seen_model_types == [_ImageConfig.model_type] * 2 + assert seen_model_types == [_ImageConfig.model_type] * 3 assert ( request_mapping.media_capabilities_for_model( SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) @@ -207,6 +241,26 @@ def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): ) +def test_media_capability_accepts_a_registered_non_base_processor(monkeypatch): + config = _ImageConfig() + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: model_type == _ImageConfig.model_type and modality == "audio", + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: "", + ) + llm = SimpleNamespace( + _hf_model_config=config, + input_processor=_AudioInputProcessor(config), + ) + + assert request_mapping.media_capabilities_for_model(llm).modalities == {"audio"} + + def test_media_capabilities_are_hashable(): first = MediaCapabilities({"image": ""}) second = MediaCapabilities({"image": ""}) @@ -232,6 +286,12 @@ def test_decoded_video_bytes_include_frames_and_audio(): request_mapping._decoded_video_bytes(object()) +def test_decoded_audio_bytes_count_samples_and_reject_unknown_types(): + assert request_mapping._decoded_audio_bytes((np.zeros(8, dtype=np.float32), 16000)) == 8 * 4 + with pytest.raises(TypeError, match="unsupported decoded audio type"): + request_mapping._decoded_audio_bytes(object()) + + def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): llm, _ = image_llm monkeypatch.setattr( @@ -368,8 +428,9 @@ async def async_load_bytes(self, data): @pytest.mark.asyncio -async def test_image_and_video_sources_preserve_global_order(image_video_llm, monkeypatch): +async def test_media_sources_preserve_global_order(image_video_audio_llm, monkeypatch): created_with = [] + audio_sources = [] class _ImageIo: @classmethod @@ -392,8 +453,19 @@ async def async_load(self, source): metadata={"source": source}, ) + class _AudioIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append(("audio", defaults, runtime)) + return cls() + + async def async_load(self, source): + audio_sources.append(source) + return np.zeros(8, dtype=np.float32), 16000 + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) monkeypatch.setattr(request_mapping, "VideoMediaIO", _VideoIo) + monkeypatch.setattr(request_mapping, "AudioMediaIO", _AudioIo) request = generation_pb2.GenerateRequest( prompt="mixed media", media=[ @@ -407,6 +479,11 @@ async def async_load(self, source): raw_bytes=b"i", mime_type="image/png", ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + url="https://example.com/audio.wav", + uuid="audio-id", + ), generation_pb2.MediaItem( modality=generation_pb2.MODALITY_VIDEO, data_uri="data:video/mp4;base64,dmlkZW8=", @@ -414,12 +491,15 @@ async def async_load(self, source): ], ) - mapped = await request_mapping._input_from_request(request, image_video_llm) + mapped = await request_mapping._input_from_request(request, image_video_audio_llm) assert created_with == [ ("video", {"format": "np", "num_frames": 4}, None), ("image", {"format": "np"}, None), + ("audio", None, None), ] + assert audio_sources == ["https://example.com/audio.wav"] + assert mapped["multi_modal_data"]["audio"][0][1] == 16000 assert [video.metadata["source"] for video in mapped["multi_modal_data"]["video"]] == [ "https://example.com/first.mp4", "data:video/mp4;base64,dmlkZW8=", @@ -428,10 +508,12 @@ async def async_load(self, source): assert mapped["multi_modal_uuids"] == { "video": ["video-id", None], "image": [None], + "audio": ["audio-id"], } assert mapped["mm_item_order"] == [ {"modality": "video", "index": 0, "placeholder": ""}, {"modality": "image", "index": 0, "placeholder": ""}, + {"modality": "audio", "index": 0, "placeholder": ""}, {"modality": "video", "index": 1, "placeholder": ""}, ] @@ -802,3 +884,20 @@ async def test_raw_video_mime_type_is_validated(image_video_llm): with pytest.raises(ValueError, match="must start with video/"): await request_mapping._input_from_request(request, image_video_llm) + + +@pytest.mark.asyncio +async def test_raw_audio_mime_type_is_validated(image_video_audio_llm): + request = generation_pb2.GenerateRequest( + prompt="audio", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + raw_bytes=b"audio", + mime_type="video/mp4", + ) + ], + ) + + with pytest.raises(ValueError, match="must start with audio/"): + await request_mapping._input_from_request(request, image_video_audio_llm) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index 1e88ed330f66..8abeb05c5c2b 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -1,10 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Multimodal conformance tests against a real TensorRT-LLM VLM engine.""" +"""Multimodal conformance tests against real TensorRT-LLM engines.""" import asyncio import base64 +from pathlib import Path import pytest @@ -18,15 +19,34 @@ from tensorrt_llm.grpc.openengine.bindings import generation_pb2, model_pb2 # noqa: E402 from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 -from tensorrt_llm.llmapi import KvCacheConfig # noqa: E402 +from tensorrt_llm.llmapi import KvCacheConfig, SchedulerConfig # noqa: E402 MODEL_NAME = "Qwen3/Qwen3-VL-2B-Instruct" - -# The engine owns background executor threads for the lifetime of this module. +WHISPER_MODEL_NAME = "openai/whisper-tiny" +_WHISPER_OUTPUT_TOKEN_IDS = [ + 1939, + 613, + 4598, + 8028, + 389, + 3011, + 582, + 259, + 1570, + 365, + 1454, + 813, + 38675, + 3378, + 13, + 50257, +] + +# The engines own background executor threads while each test runs. pytestmark = pytest.mark.threadleak(enabled=False) -@pytest.fixture(scope="module") +@pytest.fixture def vlm(): if not torch.cuda.is_available(): pytest.skip("requires a GPU to run a real VLM engine") @@ -44,6 +64,53 @@ def vlm(): engine.shutdown() +def _whisper_model_path() -> Path: + root = llm_models_root(check=True) + candidates = [ + root / "whisper-models/whisper-tiny", + root / "whisper/whisper-tiny", + root / "whisper-tiny", + ] + for candidate in candidates: + if candidate.exists(): + return candidate + pytest.skip(f"whisper-tiny not found under {root}") + + +def _whisper_audio_path() -> Path: + root = llm_models_root(check=True) + for checkpoint in ("large-v3", "large-v2"): + candidate = root / f"whisper-models/{checkpoint}/1221-135766-0002.wav" + if candidate.exists(): + return candidate + pytest.skip(f"Whisper test audio not found under {root}/whisper-models") + + +@pytest.fixture +def whisper(): + if not torch.cuda.is_available(): + pytest.skip("requires a GPU to run a real Whisper engine") + engine = LLM( + model=str(_whisper_model_path()), + disable_overlap_scheduler=True, + enable_chunked_prefill=False, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + free_gpu_memory_fraction=0.2, + cross_kv_cache_fraction=0.5, + use_kv_cache_manager_v2=True, + ), + max_batch_size=1, + max_input_len=1500, + max_num_tokens=1500, + scheduler_config=SchedulerConfig(use_python_scheduler=True), + ) + try: + yield engine + finally: + engine.shutdown() + + def test_image_and_video_capabilities_agree_with_generate(vlm, tmp_path): inference = OpenEngineInferenceServicer(vlm, MODEL_NAME) control = OpenEngineControlServicer(vlm, MODEL_NAME, inference) @@ -133,3 +200,52 @@ async def collect_video(): video_finished = [response for response in video_responses if response.HasField("finished")] assert len(video_finished) == 1 assert video_finished[0].usage.prompt_tokens > len(video_request.token_ids.ids) + + +def test_audio_capability_agrees_with_generate(whisper): + inference = OpenEngineInferenceServicer(whisper, WHISPER_MODEL_NAME) + control = OpenEngineControlServicer(whisper, WHISPER_MODEL_NAME, inference) + info = asyncio.run( + control.GetModelInfo( + model_pb2.GetModelInfoRequest(model=WHISPER_MODEL_NAME), + FakeServicerContext(), + ) + ) + + assert info.supports_multimodal is True + assert list(info.extra["trtllm_supported_modalities"]) == ["audio"] + + audio_uri = ( + "data:audio/wav;base64," + base64.b64encode(_whisper_audio_path().read_bytes()).decode() + ) + request = generation_pb2.GenerateRequest( + request_id="openengine-audio-conformance", + model=WHISPER_MODEL_NAME, + prompt="", + sampling=generation_pb2.SamplingParams(temperature=0.0), + stopping=generation_pb2.StoppingOptions(max_tokens=96), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + data_uri=audio_uri, + uuid="librispeech-1221-135766-0002", + ) + ], + ) + + async def collect_audio(): + return [response async for response in inference.Generate(request, FakeServicerContext())] + + responses = asyncio.run(collect_audio()) + assert not any(response.HasField("error") for response in responses) + token_ids = [ + token.token_id + for response in responses + if response.HasField("token") + for token in response.token.tokens + ] + assert token_ids[0] in (1939, 10890) + assert token_ids[1:] == _WHISPER_OUTPUT_TOKEN_IDS[1:] + finished = [response for response in responses if response.HasField("finished")] + assert len(finished) == 1 + assert finished[0].usage.completion_tokens == len(token_ids) From 13e0e056d2375406e677f617ebaf486d9500da63 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 11:16:15 -0700 Subject: [PATCH 07/41] feat(openengine): support image inputs Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 57 ++ tensorrt_llm/grpc/openengine/control.py | 6 +- .../grpc/openengine/request_mapping.py | 118 +++- tensorrt_llm/grpc/openengine/servicer.py | 4 +- tensorrt_llm/inputs/media_io.py | 4 + .../integration/test_lists/test-db/l0_a10.yml | 8 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../unittest/grpc/openengine/test_control.py | 17 + .../grpc/openengine/test_multimodal.py | 531 ++++++++++++++++++ .../openengine/test_multimodal_conformance.py | 93 +++ .../inputs/test_async_media_loading.py | 9 + 11 files changed, 837 insertions(+), 11 deletions(-) create mode 100644 tests/unittest/grpc/openengine/test_multimodal.py create mode 100644 tests/unittest/grpc/openengine/test_multimodal_conformance.py diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 6272d1363855..a1936ebc4c8a 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -7,6 +7,17 @@ never discover a capability that Generate then rejects. """ +from typing import Any + +from tensorrt_llm.inputs.registry import ( + MULTIMODAL_PLACEHOLDER_REGISTRY, + BaseMultimodalInputProcessor, +) + +# OpenEngine support is intentionally narrower than the LLM API today. Add a +# modality here only once request mapping and conformance coverage exist for it. +_OPENENGINE_MEDIA_MODALITIES = ("image",) + # Guides each grammar backend can build, keyed by the `GuidedDecoding` oneof # field name. Control advertises from this table and Generate enforces it, so # the two cannot disagree. llguidance has no structural-tag matcher. @@ -32,7 +43,53 @@ def supported_guides(backend: str | None) -> frozenset[str]: return GUIDE_SUPPORT_BY_BACKEND.get(backend.lower(), frozenset()) +def _loaded_model_type(llm: Any) -> str: + """Return the top-level model type for the loaded input processor.""" + input_processor = getattr(llm, "input_processor", None) + config = ( + input_processor.config + if isinstance(input_processor, BaseMultimodalInputProcessor) + else getattr(llm, "_hf_model_config", None) + ) + if config is None: + return "" + + # Composite Hugging Face configs can delegate the instance attribute to + # their text config. The class attribute retains the top-level registry key. + model_type = getattr(type(config), "model_type", None) or getattr(config, "model_type", "") + return model_type if isinstance(model_type, str) else "" + + +def supported_media_modalities(llm: Any) -> frozenset[str]: + """OpenEngine media modalities supported by the loaded model. + + A registry entry alone describes what this TensorRT-LLM build knows about; + it does not prove that the loaded engine instantiated a multimodal input + processor. Requiring both prevents a text-only model from advertising + media merely because another image model is registered in the process. + """ + input_processor = getattr(llm, "input_processor", None) + if not isinstance(input_processor, BaseMultimodalInputProcessor): + return frozenset() + model_type = _loaded_model_type(llm) + if not model_type: + return frozenset() + return frozenset( + modality + for modality in _OPENENGINE_MEDIA_MODALITIES + if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) + ) + + +def media_placeholder(llm: Any, modality: str, item_number: int) -> str: + """Return the model-specific placeholder for a one-based media item.""" + model_type = _loaded_model_type(llm) + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality).format(item_number) + + __all__ = [ "GUIDE_SUPPORT_BY_BACKEND", + "media_placeholder", "supported_guides", + "supported_media_modalities", ] diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index f020d3ae82a1..1495524d4130 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -35,7 +35,7 @@ openengine_pb2_grpc, server_pb2, ) -from .capabilities import supported_guides +from .capabilities import supported_guides, supported_media_modalities from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError from .errors import AbortFailedError from .kv_events import data_parallel_size, events_config @@ -284,6 +284,7 @@ async def GetModelInfo( ) -> model_pb2.ModelInfo: # Single-model server: like Generate, any non-empty name resolves to the # loaded model. + media_modalities = supported_media_modalities(self._llm) info = model_pb2.ModelInfo( model_id=self._model, served_model_name=self._model, @@ -292,9 +293,10 @@ async def GetModelInfo( supports_token_ids_input=True, # Generate rejects `lora_name` even on an enable_lora build. supports_lora=False, - supports_multimodal=False, + supports_multimodal=bool(media_modalities), reasoning_parser=str(self._llm.args.reasoning_parser or ""), ) + info.extra.update({"trtllm_supported_modalities": sorted(media_modalities)}) max_seq_len = _positive_int(self._llm.args.max_seq_len) if max_seq_len is not None: diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index db6ecf866c01..7e7035e0a1a8 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -3,22 +3,31 @@ """Translate an OpenEngine GenerateRequest into TensorRT-LLM request inputs.""" +import asyncio from collections.abc import Mapping from typing import Any +from urllib.parse import urlparse import grpc from tensorrt_llm.conversation_params import ConversationParams +from tensorrt_llm.inputs.media_io import ImageMediaIO +from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams from .bindings import generation_pb2 -from .capabilities import supported_guides +from .capabilities import media_placeholder, supported_guides, supported_media_modalities from .errors import UnsupportedFeatureError # Cap on the number of stop conditions. Each one costs a KMP automaton per # output index and a scan of every generated token, so an unbounded list lets a # single request tax the event loop for the whole stream. _MAX_STOP_CONDITIONS = 64 +# Media decoding can retain both encoded and decoded image buffers. Keep one +# request from scheduling an unbounded number of fetch/decode tasks or holding +# an unbounded amount of image data while preserving practical multi-image use. +_MAX_MEDIA_ITEMS = 64 +_MAX_CONCURRENT_MEDIA_LOADS = 8 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -244,7 +253,9 @@ def _metadata_from_context( return headers or None, target_dp_rank, affinity_headers -def _input_from_request(request: generation_pb2.GenerateRequest) -> str | dict[str, list[int]]: +def _base_input_from_request( + request: generation_pb2.GenerateRequest, +) -> str | dict[str, list[int]]: input_kind = request.WhichOneof("input") if input_kind == "prompt": return request.prompt @@ -253,6 +264,109 @@ def _input_from_request(request: generation_pb2.GenerateRequest) -> str | dict[s raise ValueError("exactly one of prompt or token_ids must be set") +def _image_io_for_model(llm: Any) -> ImageMediaIO: + """Create image I/O using the loaded model's preferred decode format.""" + input_processor = getattr(llm, "input_processor", None) + defaults = None + if isinstance(input_processor, BaseMultimodalInputProcessor): + preferences = input_processor.get_preferred_media_io_kwargs() or {} + defaults = preferences.get("image") + return ImageMediaIO.create(defaults, None) + + +async def _load_image_item( + item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int +) -> Any: + """Resolve and decode one OpenEngine image item.""" + source = item.WhichOneof("source") + if source is None: + raise ValueError(f"media[{index}] must set exactly one source") + if source == "url": + if urlparse(item.url).scheme.lower() not in ("http", "https"): + raise ValueError(f"media[{index}].url must use http or https") + value = item.url + elif source == "data_uri": + if urlparse(item.data_uri).scheme.lower() != "data": + raise ValueError(f"media[{index}].data_uri must use the data scheme") + value = item.data_uri + elif source == "raw_bytes": + if item.mime_type and not item.mime_type.lower().startswith("image/"): + raise ValueError(f"media[{index}].mime_type must be an image type for raw image bytes") + try: + return await image_io.async_load_bytes(bytes(item.raw_bytes)) + except Exception as error: + raise ValueError(f"media[{index}].raw_bytes could not be decoded: {error}") from error + else: + raise ValueError(f"media[{index}] has an unknown source '{source}'") + + try: + return await image_io.async_load(value) + except Exception as error: + raise ValueError(f"media[{index}].{source} could not be loaded: {error}") from error + + +async def _input_from_request( + request: generation_pb2.GenerateRequest, llm: Any +) -> str | dict[str, Any]: + """Build LLM API inputs, resolving image media off the event loop.""" + inputs = _base_input_from_request(request) + if not request.media: + return inputs + + supported = supported_media_modalities(llm) + for index, item in enumerate(request.media): + if item.modality == generation_pb2.MODALITY_UNSPECIFIED: + raise ValueError(f"media[{index}].modality must be set") + if item.modality != generation_pb2.MODALITY_IMAGE: + modality = generation_pb2.Modality.Name(item.modality).removeprefix("MODALITY_").lower() + supported_text = ", ".join(sorted(supported)) or "none" + raise UnsupportedFeatureError( + f"multimodal modality '{modality}' is not supported by the loaded model; " + f"supported modalities: {supported_text}" + ) + if "image" not in supported: + raise UnsupportedFeatureError("image input is not supported by the loaded model") + if len(request.media) > _MAX_MEDIA_ITEMS: + raise ValueError( + f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " + f"got {len(request.media)}" + ) + + image_io = _image_io_for_model(llm) + semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) + + async def load_bounded(item: generation_pb2.MediaItem, index: int) -> Any: + async with semaphore: + return await _load_image_item(item, image_io, index) + + tasks = [ + asyncio.create_task(load_bounded(item, index)) for index, item in enumerate(request.media) + ] + try: + images = await asyncio.gather(*tasks) + except (Exception, asyncio.CancelledError): + # `gather` propagates the first failure without cancelling siblings. + # Explicitly stop any in-flight URL fetches/decodes before returning + # the request error so their memory and connections do not linger. + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + mapped: dict[str, Any] = {"prompt": inputs} if isinstance(inputs, str) else dict(inputs) + mapped["multi_modal_data"] = {"image": images} + if any(item.uuid for item in request.media): + mapped["multi_modal_uuids"] = {"image": [item.uuid or None for item in request.media]} + mapped["mm_item_order"] = [ + { + "modality": "image", + "index": index, + "placeholder": media_placeholder(llm, "image", index + 1), + } + for index in range(len(request.media)) + ] + return mapped + + __all__ = [ "conversation_params_from_request", "sampling_params_from_request", diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index 6a9d6e9018b9..6bc18ffeb6a4 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -169,15 +169,13 @@ async def _generate( f"request_id '{request_id}' is already active", ) return - if request.media: - raise UnsupportedFeatureError("multimodal media is not supported") if request.lora_name: raise UnsupportedFeatureError("LoRA selection is not supported") if request.HasField("kv"): if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: raise UnsupportedFeatureError("prefix cache bypass is not supported") - inputs = _input_from_request(request) + inputs = await _input_from_request(request, self._llm) sampling_params = sampling_params_from_request(request, self._guided_backend) trace_headers, target_dp_rank, affinity_headers = _metadata_from_context(context) conversation_params = conversation_params_from_request(request) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 86f52f71da3c..e8ab41d4dccf 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -1002,6 +1002,10 @@ async def async_load(self, url: str) -> _MediaT: else: raise ValueError(f"Unsupported URL scheme: {parsed.scheme!r}") + async def async_load_bytes(self, data: bytes) -> _MediaT: + """Decode in-memory media without blocking the event loop.""" + return await self._run_in_executor(self.load_bytes, data) + @staticmethod async def _run_in_executor(fn, *args, **kwargs): """Run a blocking decode callable on the configured executor.""" diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 2b5b6f887860..6ff220686f73 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -63,11 +63,11 @@ l0_a10: - unittest/others/test_tracing.py - unittest/grpc/test_grpc_optional.py - unittest/grpc/smg/test_smg.py - # The one OpenEngine test needing a GPU: it builds a real LLM, because a - # stubbed engine would accept whatever its fake generate_async accepts and so - # would agree with an advertisement the real engine rejects. Any GPU will do - # -- the model is TinyLlama-1.1B -- so it sits with the other gateway tests. + # OpenEngine capability conformance needs real loaded models: a stub would + # accept whatever its fake generate_async tolerates and could agree with a + # false advertisement by construction. The image case uses Qwen3-VL-2B. - unittest/grpc/openengine/test_capability_conformance.py + - unittest/grpc/openengine/test_multimodal_conformance.py - unittest/disaggregated/test_coordinator_e2e.py - unittest/disaggregated/test_coordinator_worker.py - unittest/disaggregated/test_agent_multi_backends.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 58ebb7ca4eba..3fca6014da5c 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -160,6 +160,7 @@ l0_cpu: - unittest/grpc/openengine/test_coordinator.py - unittest/grpc/openengine/test_frontend_lifecycle.py - unittest/grpc/openengine/test_control.py + - unittest/grpc/openengine/test_multimodal.py - unittest/grpc/openengine/test_node.py - unittest/grpc/openengine/test_openengine.py - unittest/grpc/openengine/test_servicer.py diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 10420b6df6bd..f7de8842a222 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -166,6 +166,23 @@ async def test_model_info_reports_the_context_window_and_capabilities(): assert model_pb2.GUIDED_DECODING_MODE_JSON_SCHEMA in info.generation.guided_decoding.modes assert info.generation.output_logprobs.supported is True assert info.supports_token_ids_input is True + assert info.supports_multimodal is False + assert list(info.extra["trtllm_supported_modalities"]) == [] + + +@pytest.mark.asyncio +async def test_model_info_advertises_loaded_image_support(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm.grpc.openengine.control.supported_media_modalities", + lambda llm: frozenset({"image"}), + ) + + info = await _servicer().GetModelInfo( + model_pb2.GetModelInfoRequest(model=MODEL), FakeServicerContext() + ) + + assert info.supports_multimodal is True + assert list(info.extra["trtllm_supported_modalities"]) == ["image"] @pytest.mark.asyncio diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py new file mode 100644 index 000000000000..a4bd537e29f5 --- /dev/null +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -0,0 +1,531 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for OpenEngine image input mapping and discovery.""" + +import asyncio +from collections.abc import AsyncIterator, Sequence +from types import SimpleNamespace +from typing import Any + +import pytest + +grpc = pytest.importorskip( + "grpc", reason='gRPC runtime not installed (pip install "grpcio>=1.67.1,<2")' +) +from conftest import AbortError, FakeServicerContext # noqa: E402 + +import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 +from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.capabilities import supported_media_modalities # noqa: E402 +from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 +from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.inputs.registry import ( # noqa: E402 + MULTIMODAL_PLACEHOLDER_REGISTRY, + BaseMultimodalInputProcessor, +) + +pytestmark = pytest.mark.cpu_only + + +class _ImageConfig: + # Composite configs can expose their text model type on the instance; the + # top-level class attribute is the key used by the multimodal registry. + model_type = "openengine_test_image" + + def __init__(self) -> None: + self.model_type = "openengine_test_text" + + +class _ImageInputProcessor(BaseMultimodalInputProcessor): + """Minimal concrete processor used only for capability detection.""" + + def __init__(self, config: _ImageConfig) -> None: + self._config = config + + @property + def processor(self): + return None + + @property + def tokenizer(self): + return None + + @property + def config(self): + return self._config + + @property + def dtype(self): + return None + + def call_with_text_prompt(self, inputs, sampling_params): + raise NotImplementedError + + def get_preferred_media_io_kwargs(self): + return {"image": {"format": "np"}} + + +class _Tokenizer: + def convert_ids_to_tokens( + self, token_ids: Sequence[int], skip_special_tokens: bool = False + ) -> list[str]: + del skip_special_tokens + return [f"token-{token_id}" for token_id in token_ids] + + def decode(self, token_ids: Sequence[int], skip_special_tokens: bool = False) -> str: + del skip_special_tokens + return "".join(f"token-{token_id}" for token_id in token_ids) + + +class _ResultHandle: + def __init__(self, result: Any) -> None: + self._result = result + self.finished = False + self.prompt_token_ids = [] + self.outputs = [] + self.cached_tokens = 0 + self.error = None + self.aborted = False + + async def __aiter__(self) -> AsyncIterator[Any]: + self.finished = self._result.finished + self.prompt_token_ids = self._result.prompt_token_ids + self.outputs = self._result.outputs + self.cached_tokens = self._result.cached_tokens + self.error = self._result.error + yield self._result + + def abort(self) -> None: + self.aborted = True + + +class _ImageLlm: + def __init__(self, input_processor, config) -> None: + self.input_processor = input_processor + self._hf_model_config = config + self.args = SimpleNamespace(guided_decoding_backend=None) + self.tokenizer = _Tokenizer() + output = SimpleNamespace( + index=0, + token_ids=[99], + text="caption", + logprobs=[], + prompt_logprobs=[], + finish_reason="length", + stop_reason=None, + ) + self.result_handle = _ResultHandle( + SimpleNamespace( + # Represents expansion of one placeholder into five image tokens. + prompt_token_ids=[1, 2, 3, 4, 5, 6], + outputs=[output], + cached_tokens=0, + error=None, + finished=True, + ) + ) + self.generate_kwargs = None + + def generate_async(self, **kwargs): + self.generate_kwargs = kwargs + return self.result_handle + + +@pytest.fixture +def image_llm(monkeypatch): + config = _ImageConfig() + input_processor = _ImageInputProcessor(config) + seen_model_types = [] + + def is_valid(model_type, modality): + seen_model_types.append(model_type) + return model_type == _ImageConfig.model_type and modality == "image" + + monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", is_valid) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: "", + ) + return ( + SimpleNamespace( + _hf_model_config=config, + input_processor=input_processor, + ), + seen_model_types, + ) + + +def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): + llm, seen_model_types = image_llm + + assert supported_media_modalities(llm) == {"image"} + assert seen_model_types == [_ImageConfig.model_type] + assert ( + supported_media_modalities( + SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) + ) + == set() + ) + + +@pytest.mark.asyncio +async def test_text_only_input_mapping_is_unchanged(): + llm = SimpleNamespace() + + prompt = generation_pb2.GenerateRequest(prompt="hello") + token_ids = generation_pb2.GenerateRequest(token_ids=generation_pb2.TokenIds(ids=[1, 2, 3])) + + assert await request_mapping._input_from_request(prompt, llm) == "hello" + assert await request_mapping._input_from_request(token_ids, llm) == { + "prompt_token_ids": [1, 2, 3] + } + + +@pytest.mark.asyncio +async def test_image_sources_preserve_order_uuid_and_model_decode_defaults(image_llm, monkeypatch): + llm, _ = image_llm + created_with = [] + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append((defaults, runtime)) + return cls() + + async def async_load(self, source): + # Finish in the opposite order to prove gather preserves request order. + if source.startswith("https"): + await asyncio.sleep(0.01) + return f"decoded:{source}" + + async def async_load_bytes(self, data): + return f"decoded-bytes:{data.decode()}" + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt=" describe", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="https://example.com/first.png", + uuid="first", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + data_uri="data:image/png;base64,c2Vjb25k", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"third", + mime_type="image/png", + uuid="third", + ), + ], + ) + + mapped = await request_mapping._input_from_request(request, llm) + + assert created_with == [({"format": "np"}, None)] + assert mapped["prompt"] == request.prompt + assert mapped["multi_modal_data"] == { + "image": [ + "decoded:https://example.com/first.png", + "decoded:data:image/png;base64,c2Vjb25k", + "decoded-bytes:third", + ] + } + assert mapped["multi_modal_uuids"] == {"image": ["first", None, "third"]} + assert mapped["mm_item_order"] == [ + {"modality": "image", "index": 0, "placeholder": ""}, + {"modality": "image", "index": 1, "placeholder": ""}, + {"modality": "image", "index": 2, "placeholder": ""}, + ] + + +@pytest.mark.asyncio +async def test_image_loading_has_a_bounded_concurrency(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + active = 0 + maximum_active = 0 + + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + type(self).active += 1 + type(self).maximum_active = max(type(self).maximum_active, type(self).active) + try: + await asyncio.sleep(0.01) + return data + finally: + type(self).active -= 1 + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=str(index).encode(), + ) + for index in range(request_mapping._MAX_CONCURRENT_MEDIA_LOADS * 2 + 1) + ], + ) + + mapped = await request_mapping._input_from_request(request, llm) + + assert _ImageIo.maximum_active == request_mapping._MAX_CONCURRENT_MEDIA_LOADS + assert mapped["multi_modal_data"]["image"] == [ + str(index).encode() for index in range(request_mapping._MAX_CONCURRENT_MEDIA_LOADS * 2 + 1) + ] + + +@pytest.mark.asyncio +async def test_image_loading_cancels_siblings_after_a_failure(image_llm, monkeypatch): + llm, _ = image_llm + slow_started = asyncio.Event() + slow_cancelled = asyncio.Event() + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + if data == b"slow": + slow_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + slow_cancelled.set() + raise + await slow_started.wait() + raise ValueError("decode failed") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=value, + ) + for value in (b"slow", b"broken") + ], + ) + + with pytest.raises(ValueError, match="decode failed"): + await request_mapping._input_from_request(request, llm) + assert slow_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_image_count_is_bounded_before_decoding(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + raise AssertionError("the media-count guard must run before decoder creation") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="too many images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + for _ in range(request_mapping._MAX_MEDIA_ITEMS + 1) + ], + ) + + with pytest.raises(ValueError, match="at most 64 media items"): + await request_mapping._input_from_request(request, llm) + + +@pytest.mark.asyncio +async def test_image_mapping_supports_token_ids(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return "decoded" + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + token_ids=generation_pb2.TokenIds(ids=[10, 20]), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + mapped = await request_mapping._input_from_request(request, llm) + + assert mapped["prompt_token_ids"] == [10, 20] + assert mapped["multi_modal_data"] == {"image": ["decoded"]} + assert "multi_modal_uuids" not in mapped + + +@pytest.mark.asyncio +async def test_generate_submits_image_and_reports_expanded_prompt_usage(image_llm, monkeypatch): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return "decoded" + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + request = generation_pb2.GenerateRequest( + request_id="image-request", + model="image-model", + token_ids=generation_pb2.TokenIds(ids=[1, 2]), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + responses = [response async for response in servicer.Generate(request, FakeServicerContext())] + + assert llm.generate_kwargs["inputs"]["multi_modal_data"] == {"image": ["decoded"]} + assert responses[-1].usage.prompt_tokens == 6 + assert responses[-1].usage.total_tokens == 7 + + +@pytest.mark.asyncio +async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + + unsupported_context = FakeServicerContext() + video_request = generation_pb2.GenerateRequest( + request_id="video-request", + model="image-model", + prompt="video", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"video", + ) + ], + ) + with pytest.raises(AbortError): + async for _ in servicer.Generate(video_request, unsupported_context): + pass + assert unsupported_context.abort_code == grpc.StatusCode.UNIMPLEMENTED + + malformed_context = FakeServicerContext() + file_request = generation_pb2.GenerateRequest( + request_id="file-request", + model="image-model", + prompt="image", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="file:///tmp/private.png", + ) + ], + ) + with pytest.raises(AbortError): + async for _ in servicer.Generate(file_request, malformed_context): + pass + assert malformed_context.abort_code == grpc.StatusCode.INVALID_ARGUMENT + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("modality", "name"), + [ + (generation_pb2.MODALITY_VIDEO, "video"), + (generation_pb2.MODALITY_AUDIO, "audio"), + ], +) +async def test_video_and_audio_are_explicitly_unsupported(image_llm, modality, name): + llm, _ = image_llm + request = generation_pb2.GenerateRequest( + prompt="media", + media=[generation_pb2.MediaItem(modality=modality, raw_bytes=b"data")], + ) + + with pytest.raises(UnsupportedFeatureError, match=name): + await request_mapping._input_from_request(request, llm) + + +@pytest.mark.asyncio +async def test_image_is_rejected_for_a_text_only_loaded_model(): + request = generation_pb2.GenerateRequest( + prompt="image", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + with pytest.raises(UnsupportedFeatureError, match="loaded model"): + await request_mapping._input_from_request( + request, SimpleNamespace(input_processor=object()) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("item", "message"), + [ + (generation_pb2.MediaItem(raw_bytes=b"image"), "modality must be set"), + ( + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE), + "must set exactly one source", + ), + ( + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="file:///tmp/private.png", + ), + "must use http or https", + ), + ( + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + mime_type="video/mp4", + ), + "must be an image type", + ), + ], +) +async def test_malformed_image_items_are_rejected(image_llm, item, message): + llm, _ = image_llm + request = generation_pb2.GenerateRequest(prompt="image", media=[item]) + + with pytest.raises(ValueError, match=message): + await request_mapping._input_from_request(request, llm) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py new file mode 100644 index 000000000000..e97c9a708aac --- /dev/null +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Image conformance tests against a real TensorRT-LLM VLM engine.""" + +import asyncio +import base64 + +import pytest + +pytest.importorskip("grpc", reason='gRPC runtime not installed (pip install "grpcio>=1.67.1,<2")') + +import torch # noqa: E402 +from conftest import FakeServicerContext # noqa: E402 +from utils.llm_data import llm_models_root # noqa: E402 + +from tensorrt_llm import LLM # noqa: E402 +from tensorrt_llm.grpc.openengine.bindings import generation_pb2, model_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 +from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.llmapi import KvCacheConfig # noqa: E402 + +MODEL_NAME = "Qwen3/Qwen3-VL-2B-Instruct" + +# The engine owns background executor threads for the lifetime of this module. +pytestmark = pytest.mark.threadleak(enabled=False) + + +@pytest.fixture(scope="module") +def vlm(): + if not torch.cuda.is_available(): + pytest.skip("requires a GPU to run a real VLM engine") + model_path = llm_models_root(check=True) / MODEL_NAME + engine = LLM( + model=str(model_path), + load_format="dummy", + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.5), + max_batch_size=1, + max_seq_len=2048, + ) + try: + yield engine + finally: + engine.shutdown() + + +def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm): + inference = OpenEngineInferenceServicer(vlm, MODEL_NAME) + control = OpenEngineControlServicer(vlm, MODEL_NAME, inference) + info = asyncio.run( + control.GetModelInfo( + model_pb2.GetModelInfoRequest(model=MODEL_NAME), + FakeServicerContext(), + ) + ) + + assert info.supports_multimodal is True + assert list(info.extra["trtllm_supported_modalities"]) == ["image"] + + prompt = ( + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + "Describe this image briefly.<|im_end|>\n<|im_start|>assistant\n" + ) + prompt_token_ids = vlm.tokenizer.encode(prompt) + image_path = llm_models_root(check=True) / "multimodals/test_data/seashore.png" + image_bytes = image_path.read_bytes() + data_uri = "data:image/png;base64," + base64.b64encode(image_bytes).decode() + request = generation_pb2.GenerateRequest( + request_id="openengine-image-conformance", + model=MODEL_NAME, + token_ids=generation_pb2.TokenIds(ids=prompt_token_ids), + stopping=generation_pb2.StoppingOptions(max_tokens=4), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + data_uri=data_uri, + uuid="seashore", + ) + ], + ) + + async def collect(): + return [response async for response in inference.Generate(request, FakeServicerContext())] + + responses = asyncio.run(collect()) + assert responses + assert not any(response.HasField("error") for response in responses) + finished = [response for response in responses if response.HasField("finished")] + assert len(finished) == 1 + assert finished[0].usage.prompt_tokens > len(prompt_token_ids) + assert finished[0].usage.total_tokens == ( + finished[0].usage.prompt_tokens + finished[0].usage.completion_tokens + ) diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index 5855605a06ca..997b140f79f7 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -58,6 +58,15 @@ def _make_audio_file(tmp_path: str) -> str: class TestAsyncLoadImage: + @pytest.mark.asyncio + async def test_load_image_from_raw_bytes_pil(self): + encoded = _make_image_data_url().partition(",")[2] + image = await media_io_module.ImageMediaIO(format="pil").async_load_bytes( + base64.b64decode(encoded) + ) + assert isinstance(image, Image.Image) + assert image.mode == "RGB" + @pytest.mark.asyncio async def test_load_image_from_data_url_pil(self): url = _make_image_data_url() From 1748fb752f16482458d9dc095feb463a9c5525bc Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 16:55:11 -0700 Subject: [PATCH 08/41] fix(openengine): harden multimodal media handling Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 74 ++++++--- tensorrt_llm/grpc/openengine/control.py | 4 +- tensorrt_llm/grpc/openengine/errors.py | 5 + .../grpc/openengine/request_mapping.py | 157 ++++++++++++++---- tensorrt_llm/grpc/openengine/servicer.py | 14 +- tensorrt_llm/inputs/media_io.py | 41 +++-- .../unittest/grpc/openengine/test_control.py | 9 +- .../grpc/openengine/test_multimodal.py | 138 ++++++++++++++- tests/unittest/inputs/test_url_validation.py | 21 ++- 9 files changed, 383 insertions(+), 80 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index a1936ebc4c8a..6affcf574747 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -7,17 +7,15 @@ never discover a capability that Generate then rejects. """ -from typing import Any +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Collection, Mapping from tensorrt_llm.inputs.registry import ( MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, ) -# OpenEngine support is intentionally narrower than the LLM API today. Add a -# modality here only once request mapping and conformance coverage exist for it. -_OPENENGINE_MEDIA_MODALITIES = ("image",) - # Guides each grammar backend can build, keyed by the `GuidedDecoding` oneof # field name. Control advertises from this table and Generate enforces it, so # the two cannot disagree. llguidance has no structural-tag matcher. @@ -29,6 +27,27 @@ } +@dataclass(frozen=True) +class MediaCapabilities: + """Immutable media capability snapshot shared by Control and Generate.""" + + placeholder_templates: Mapping[str, str] + + def __post_init__(self) -> None: + object.__setattr__( + self, "placeholder_templates", MappingProxyType(dict(self.placeholder_templates)) + ) + + @property + def modalities(self) -> frozenset[str]: + """Modalities the loaded model and request mapper both support.""" + return frozenset(self.placeholder_templates) + + def placeholder(self, modality: str, item_number: int) -> str: + """Format the model placeholder for a one-based media item.""" + return self.placeholder_templates[modality].format(item_number) + + def supported_guides(backend: str | None) -> frozenset[str]: """Guides `backend` can build; empty when none can be. @@ -60,36 +79,47 @@ def _loaded_model_type(llm: Any) -> str: return model_type if isinstance(model_type, str) else "" -def supported_media_modalities(llm: Any) -> frozenset[str]: - """OpenEngine media modalities supported by the loaded model. +def resolve_media_capabilities( + llm: Any, implemented_modalities: Collection[str] +) -> MediaCapabilities: + """Resolve media support implemented by both the mapper and loaded model. A registry entry alone describes what this TensorRT-LLM build knows about; it does not prove that the loaded engine instantiated a multimodal input processor. Requiring both prevents a text-only model from advertising - media merely because another image model is registered in the process. + media merely because another image model is registered in the process. The + request mapper supplies ``implemented_modalities`` from its handler table, + so discovery cannot advertise a modality without a request implementation. """ input_processor = getattr(llm, "input_processor", None) if not isinstance(input_processor, BaseMultimodalInputProcessor): - return frozenset() + return MediaCapabilities({}) model_type = _loaded_model_type(llm) if not model_type: - return frozenset() - return frozenset( - modality - for modality in _OPENENGINE_MEDIA_MODALITIES - if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) - ) + return MediaCapabilities({}) - -def media_placeholder(llm: Any, modality: str, item_number: int) -> str: - """Return the model-specific placeholder for a one-based media item.""" - model_type = _loaded_model_type(llm) - return MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality).format(item_number) + placeholders = { + modality: MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality) + for modality in implemented_modalities + if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) + } + # Validate templates once at startup. A malformed registry entry is a + # server configuration failure, not a client request error discovered only + # after Control has advertised the modality. + for modality, template in placeholders.items(): + try: + template.format(1) + except (IndexError, KeyError, ValueError) as error: + raise RuntimeError( + f"invalid multimodal placeholder for model type '{model_type}' " + f"and modality '{modality}': {error}" + ) from error + return MediaCapabilities(placeholders) __all__ = [ "GUIDE_SUPPORT_BY_BACKEND", - "media_placeholder", + "MediaCapabilities", + "resolve_media_capabilities", "supported_guides", - "supported_media_modalities", ] diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index 1495524d4130..8b4292dd33ae 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -35,7 +35,7 @@ openengine_pb2_grpc, server_pb2, ) -from .capabilities import supported_guides, supported_media_modalities +from .capabilities import supported_guides from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError from .errors import AbortFailedError from .kv_events import data_parallel_size, events_config @@ -284,7 +284,7 @@ async def GetModelInfo( ) -> model_pb2.ModelInfo: # Single-model server: like Generate, any non-empty name resolves to the # loaded model. - media_modalities = supported_media_modalities(self._llm) + media_modalities = self._inference.supported_media_modalities info = model_pb2.ModelInfo( model_id=self._model, served_model_name=self._model, diff --git a/tensorrt_llm/grpc/openengine/errors.py b/tensorrt_llm/grpc/openengine/errors.py index 267d2a508c8a..d639794e0c15 100644 --- a/tensorrt_llm/grpc/openengine/errors.py +++ b/tensorrt_llm/grpc/openengine/errors.py @@ -8,6 +8,10 @@ class UnsupportedFeatureError(ValueError): """Raised when OpenEngine requests a feature this adapter cannot map.""" +class MediaUnavailableError(RuntimeError): + """Raised when fetching remote request media fails transiently.""" + + class AbortFailedError(RuntimeError): """Raised when an active request could not be aborted on the engine.""" @@ -18,5 +22,6 @@ def __init__(self, request_id: str, cause: BaseException) -> None: __all__ = [ "AbortFailedError", + "MediaUnavailableError", "UnsupportedFeatureError", ] diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 7e7035e0a1a8..28c948ba89d4 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -4,20 +4,24 @@ """Translate an OpenEngine GenerateRequest into TensorRT-LLM request inputs.""" import asyncio -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass from typing import Any from urllib.parse import urlparse import grpc +import numpy as np +import torch +from PIL import Image from tensorrt_llm.conversation_params import ConversationParams -from tensorrt_llm.inputs.media_io import ImageMediaIO +from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams from .bindings import generation_pb2 -from .capabilities import media_placeholder, supported_guides, supported_media_modalities -from .errors import UnsupportedFeatureError +from .capabilities import MediaCapabilities, resolve_media_capabilities, supported_guides +from .errors import MediaUnavailableError, UnsupportedFeatureError # Cap on the number of stop conditions. Each one costs a KMP automaton per # output index and a scan of every generated token, so an unbounded list lets a @@ -28,10 +32,26 @@ # an unbounded amount of image data while preserving practical multi-image use. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 +# Encoded bodies are separately bounded by the 64 MiB gRPC message limit and +# the media loader's 200 MiB per-URL response limit. This cap covers expansion +# during decode across the entire request. +_MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" +_MediaLoader = Callable[[generation_pb2.MediaItem, int], Awaitable[Any]] + + +@dataclass(frozen=True) +class _MediaHandler: + """One fully implemented OpenEngine media mapping path.""" + + name: str + loader_factory: Callable[[Any], _MediaLoader] + decoded_bytes: Callable[[Any], int] + + def conversation_affinity_enabled(llm: Any) -> bool: """Whether TensorRT-LLM owns attention-DP placement for conversations.""" args = getattr(llm, "args", None) @@ -301,49 +321,113 @@ async def _load_image_item( try: return await image_io.async_load(value) + except MediaFetchError as error: + raise MediaUnavailableError( + f"media[{index}].{source} is temporarily unavailable: {error}" + ) from error except Exception as error: raise ValueError(f"media[{index}].{source} could not be loaded: {error}") from error +def _image_loader_factory(llm: Any) -> _MediaLoader: + image_io = _image_io_for_model(llm) + + async def load(item: generation_pb2.MediaItem, index: int) -> Any: + return await _load_image_item(item, image_io, index) + + return load + + +def _decoded_image_bytes(image: Any) -> int: + """Return retained bytes for a decoded image representation.""" + if isinstance(image, torch.Tensor): + return image.numel() * image.element_size() + if isinstance(image, np.ndarray): + return image.nbytes + if isinstance(image, Image.Image): + return image.width * image.height * len(image.getbands()) + # Custom model media types cannot be sized generically. The built-in image + # handler returns one of the representations above. + return 0 + + +_MEDIA_HANDLERS: dict[int, _MediaHandler] = { + generation_pb2.MODALITY_IMAGE: _MediaHandler( + "image", _image_loader_factory, _decoded_image_bytes + ), +} + + +def media_capabilities_for_model(llm: Any) -> MediaCapabilities: + """Resolve the loaded model against the request handlers in this module.""" + return resolve_media_capabilities(llm, {handler.name for handler in _MEDIA_HANDLERS.values()}) + + +def _modality_name(modality: int) -> str: + try: + return generation_pb2.Modality.Name(modality).removeprefix("MODALITY_").lower() + except ValueError: + return f"unknown ({modality})" + + async def _input_from_request( - request: generation_pb2.GenerateRequest, llm: Any + request: generation_pb2.GenerateRequest, + llm: Any, + media_capabilities: MediaCapabilities | None = None, ) -> str | dict[str, Any]: - """Build LLM API inputs, resolving image media off the event loop.""" + """Build LLM API inputs, resolving media off the event loop.""" inputs = _base_input_from_request(request) if not request.media: return inputs - supported = supported_media_modalities(llm) + if len(request.media) > _MAX_MEDIA_ITEMS: + raise ValueError( + f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " + f"got {len(request.media)}" + ) + + capabilities = media_capabilities or media_capabilities_for_model(llm) + if not capabilities.modalities: + raise UnsupportedFeatureError("media input is not supported by the loaded model") + + resolved_handlers: list[_MediaHandler] = [] for index, item in enumerate(request.media): if item.modality == generation_pb2.MODALITY_UNSPECIFIED: raise ValueError(f"media[{index}].modality must be set") - if item.modality != generation_pb2.MODALITY_IMAGE: - modality = generation_pb2.Modality.Name(item.modality).removeprefix("MODALITY_").lower() - supported_text = ", ".join(sorted(supported)) or "none" + handler = _MEDIA_HANDLERS.get(item.modality) + if handler is None or handler.name not in capabilities.modalities: + modality = _modality_name(item.modality) + supported_text = ", ".join(sorted(capabilities.modalities)) or "none" raise UnsupportedFeatureError( f"multimodal modality '{modality}' is not supported by the loaded model; " f"supported modalities: {supported_text}" ) - if "image" not in supported: - raise UnsupportedFeatureError("image input is not supported by the loaded model") - if len(request.media) > _MAX_MEDIA_ITEMS: - raise ValueError( - f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " - f"got {len(request.media)}" - ) + resolved_handlers.append(handler) - image_io = _image_io_for_model(llm) + loaders = {handler.name: handler.loader_factory(llm) for handler in set(resolved_handlers)} semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) + decoded_bytes = 0 - async def load_bounded(item: generation_pb2.MediaItem, index: int) -> Any: + async def load_bounded( + item: generation_pb2.MediaItem, index: int, handler: _MediaHandler + ) -> Any: + nonlocal decoded_bytes async with semaphore: - return await _load_image_item(item, image_io, index) + decoded = await loaders[handler.name](item, index) + decoded_bytes += handler.decoded_bytes(decoded) + if decoded_bytes > _MAX_TOTAL_DECODED_MEDIA_BYTES: + raise ValueError( + "decoded media exceeds the per-request limit of " + f"{_MAX_TOTAL_DECODED_MEDIA_BYTES} bytes" + ) + return decoded tasks = [ - asyncio.create_task(load_bounded(item, index)) for index, item in enumerate(request.media) + asyncio.create_task(load_bounded(item, index, handler)) + for index, (item, handler) in enumerate(zip(request.media, resolved_handlers)) ] try: - images = await asyncio.gather(*tasks) + decoded_media = await asyncio.gather(*tasks) except (Exception, asyncio.CancelledError): # `gather` propagates the first failure without cancelling siblings. # Explicitly stop any in-flight URL fetches/decodes before returning @@ -352,18 +436,27 @@ async def load_bounded(item: generation_pb2.MediaItem, index: int) -> Any: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) raise + grouped_media: dict[str, list[Any]] = {} + grouped_uuids: dict[str, list[str | None]] = {} + item_order: list[dict[str, Any]] = [] + for item, handler, decoded in zip(request.media, resolved_handlers, decoded_media): + modality_items = grouped_media.setdefault(handler.name, []) + modality_index = len(modality_items) + modality_items.append(decoded) + grouped_uuids.setdefault(handler.name, []).append(item.uuid or None) + item_order.append( + { + "modality": handler.name, + "index": modality_index, + "placeholder": capabilities.placeholder(handler.name, modality_index + 1), + } + ) + mapped: dict[str, Any] = {"prompt": inputs} if isinstance(inputs, str) else dict(inputs) - mapped["multi_modal_data"] = {"image": images} + mapped["multi_modal_data"] = grouped_media if any(item.uuid for item in request.media): - mapped["multi_modal_uuids"] = {"image": [item.uuid or None for item in request.media]} - mapped["mm_item_order"] = [ - { - "modality": "image", - "index": index, - "placeholder": media_placeholder(llm, "image", index + 1), - } - for index in range(len(request.media)) - ] + mapped["multi_modal_uuids"] = grouped_uuids + mapped["mm_item_order"] = item_order return mapped diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index 6bc18ffeb6a4..f1547b65a112 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -25,13 +25,14 @@ Reservation, ) from .disagg import disaggregated_params_from_request -from .errors import AbortFailedError, UnsupportedFeatureError +from .errors import AbortFailedError, MediaUnavailableError, UnsupportedFeatureError from .formatting import _engine_error_response, _stop_texts from .request_mapping import ( _input_from_request, _metadata_from_context, conversation_affinity_enabled, conversation_params_from_request, + media_capabilities_for_model, sampling_params_from_request, ) from .streaming import ( @@ -62,9 +63,15 @@ def __init__( # context worker uses. The actual transfer is driven by opaque_state. self._kv_transfer_backend = kv_transfer_backend self._guided_backend = llm.args.guided_decoding_backend + self._media_capabilities = media_capabilities_for_model(llm) self.subagent_affinity_auth_key = subagent_affinity_auth_key self._active_requests: dict[str, Any] = {} + @property + def supported_media_modalities(self) -> frozenset[str]: + """Media capabilities cached from the loaded model and mapper.""" + return self._media_capabilities.modalities + def active_request_count(self) -> int: """Requests accepted and not yet completed.""" return len(self._active_requests) @@ -175,7 +182,7 @@ async def _generate( if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: raise UnsupportedFeatureError("prefix cache bypass is not supported") - inputs = await _input_from_request(request, self._llm) + inputs = await _input_from_request(request, self._llm, self._media_capabilities) sampling_params = sampling_params_from_request(request, self._guided_backend) trace_headers, target_dp_rank, affinity_headers = _metadata_from_context(context) conversation_params = conversation_params_from_request(request) @@ -230,6 +237,9 @@ async def _generate( except UnsupportedFeatureError as error: await context.abort(grpc.StatusCode.UNIMPLEMENTED, str(error)) return + except MediaUnavailableError as error: + await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) + return except ValueError as error: await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(error)) return diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index e8ab41d4dccf..484fdf5b6f70 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -45,6 +45,10 @@ from tensorrt_llm.logger import logger +class MediaFetchError(RuntimeError): + """A transient failure while fetching remote media.""" + + def rgba_to_rgb( image: Image.Image, background_color: Union[tuple[int, int, int], list[int]] = (255, 255, 255) ) -> Image.Image: @@ -161,7 +165,7 @@ def _validate_url(url: str) -> None: try: infos = socket.getaddrinfo(hostname, None, proto=socket.IPPROTO_TCP) except socket.gaierror as exc: - raise RuntimeError(f"Could not resolve hostname {hostname!r}") from exc + raise MediaFetchError(f"Could not resolve hostname {hostname!r}") from exc allow_private = _allow_private_urls() for _family, _type, _proto, _canon, sockaddr in infos: @@ -257,18 +261,29 @@ async def _safe_aiohttp_get( async def _fetch(fetch_session: aiohttp.ClientSession) -> bytes: current_url = url for _ in range(_MAX_REDIRECTS + 1): - async with fetch_session.get( - current_url, timeout=timeout, allow_redirects=False - ) as response: - if response.status in _REDIRECT_STATUSES: - redirect_url = response.headers.get("Location", "") - current_url = urljoin(current_url, redirect_url) - await asyncio.to_thread(_validate_url, current_url) - continue - maybe_coro = response.raise_for_status() - if asyncio.iscoroutine(maybe_coro): - await maybe_coro - return await _read_aiohttp_content(response) + try: + async with fetch_session.get( + current_url, timeout=timeout, allow_redirects=False + ) as response: + if response.status in _REDIRECT_STATUSES: + redirect_url = response.headers.get("Location", "") + current_url = urljoin(current_url, redirect_url) + await asyncio.to_thread(_validate_url, current_url) + continue + maybe_coro = response.raise_for_status() + if asyncio.iscoroutine(maybe_coro): + await maybe_coro + return await _read_aiohttp_content(response) + except MediaFetchError: + raise + except aiohttp.ClientResponseError as error: + if error.status in (408, 429) or error.status >= 500: + raise MediaFetchError( + f"remote media server returned HTTP {error.status}" + ) from error + raise + except (aiohttp.ClientError, asyncio.TimeoutError) as error: + raise MediaFetchError(f"remote media fetch failed: {error}") from error raise RuntimeError("Too many redirects") if session is not None: diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index f7de8842a222..79b56c77f9d5 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -19,6 +19,7 @@ model_pb2, server_pb2, ) +from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.llmapi import KVEventsConfig # noqa: E402 @@ -172,12 +173,14 @@ async def test_model_info_reports_the_context_window_and_capabilities(): @pytest.mark.asyncio async def test_model_info_advertises_loaded_image_support(monkeypatch): + inference = _inference() monkeypatch.setattr( - "tensorrt_llm.grpc.openengine.control.supported_media_modalities", - lambda llm: frozenset({"image"}), + inference, + "_media_capabilities", + MediaCapabilities({"image": ""}), ) - info = await _servicer().GetModelInfo( + info = await _servicer(inference=inference).GetModelInfo( model_pb2.GetModelInfoRequest(model=MODEL), FakeServicerContext() ) diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index a4bd537e29f5..43ff2297861b 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -8,6 +8,7 @@ from types import SimpleNamespace from typing import Any +import numpy as np import pytest grpc = pytest.importorskip( @@ -17,9 +18,9 @@ import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 -from tensorrt_llm.grpc.openengine.capabilities import supported_media_modalities # noqa: E402 from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, @@ -160,16 +161,75 @@ def is_valid(model_type, modality): def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): llm, seen_model_types = image_llm - assert supported_media_modalities(llm) == {"image"} + assert request_mapping.media_capabilities_for_model(llm).modalities == {"image"} assert seen_model_types == [_ImageConfig.model_type] assert ( - supported_media_modalities( + request_mapping.media_capabilities_for_model( SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) - ) + ).modalities == set() ) +def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): + llm, _ = image_llm + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: "", + ) + + with pytest.raises(RuntimeError, match="invalid multimodal placeholder"): + request_mapping.media_capabilities_for_model(llm) + + +@pytest.mark.asyncio +async def test_handler_table_drives_capability_and_mapping(image_llm, monkeypatch): + llm, _ = image_llm + + def video_loader_factory(_llm): + async def load(item, index): + return f"video-{index}:{bytes(item.raw_bytes).decode()}" + + return load + + monkeypatch.setitem( + request_mapping._MEDIA_HANDLERS, + generation_pb2.MODALITY_VIDEO, + request_mapping._MediaHandler("video", video_loader_factory, lambda decoded: len(decoded)), + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: model_type == _ImageConfig.model_type + and modality in {"image", "video"}, + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + + capabilities = request_mapping.media_capabilities_for_model(llm) + assert capabilities.modalities == {"image", "video"} + + request = generation_pb2.GenerateRequest( + prompt="video", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"clip", + ) + ], + ) + mapped = await request_mapping._input_from_request(request, llm, capabilities) + + assert mapped["multi_modal_data"] == {"video": ["video-0:clip"]} + assert mapped["mm_item_order"] == [ + {"modality": "video", "index": 0, "placeholder": ""} + ] + + @pytest.mark.asyncio async def test_text_only_input_mapping_is_unchanged(): llm = SimpleNamespace() @@ -338,7 +398,9 @@ def create(cls, defaults, runtime): prompt="too many images", media=[ generation_pb2.MediaItem( - modality=generation_pb2.MODALITY_IMAGE, + # Unsupported on purpose: the count guard must win before the + # per-item modality loop. + modality=generation_pb2.MODALITY_VIDEO, raw_bytes=b"image", ) for _ in range(request_mapping._MAX_MEDIA_ITEMS + 1) @@ -349,6 +411,35 @@ def create(cls, defaults, runtime): await request_mapping._input_from_request(request, llm) +@pytest.mark.asyncio +async def test_total_decoded_image_bytes_are_bounded(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return np.zeros((2, 2, 3), dtype=np.uint8) + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(request_mapping, "_MAX_TOTAL_DECODED_MEDIA_BYTES", 20) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + for _ in range(2) + ], + ) + + with pytest.raises(ValueError, match="decoded media exceeds"): + await request_mapping._input_from_request(request, llm) + + @pytest.mark.asyncio async def test_image_mapping_supports_token_ids(image_llm, monkeypatch): llm, _ = image_llm @@ -460,6 +551,43 @@ async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): assert malformed_context.abort_code == grpc.StatusCode.INVALID_ARGUMENT +@pytest.mark.asyncio +async def test_generate_maps_transient_media_fetch_to_unavailable(image_llm, monkeypatch): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load(self, source): + raise MediaFetchError("connection timed out") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + context = FakeServicerContext() + request = generation_pb2.GenerateRequest( + request_id="remote-image", + model="image-model", + prompt="image", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="https://example.com/image.png", + ) + ], + ) + + with pytest.raises(AbortError): + async for _ in servicer.Generate(request, context): + pass + assert context.abort_code == grpc.StatusCode.UNAVAILABLE + + @pytest.mark.asyncio @pytest.mark.parametrize( ("modality", "name"), diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index 9bff8139754c..c2d0421ee5bc 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -26,6 +26,7 @@ from tensorrt_llm.inputs.media_io import ( _MAX_REDIRECTS, + MediaFetchError, _safe_aiohttp_get, _safe_request_get, _validate_url, @@ -129,7 +130,7 @@ def test_rejects_multicast(self, _): side_effect=socket.gaierror("Name or service not known"), ) def test_rejects_unresolvable_hostname(self, _): - with pytest.raises(RuntimeError, match="Could not resolve"): + with pytest.raises(MediaFetchError, match="Could not resolve"): _validate_url("http://this.does.not.exist.invalid/") @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) @@ -346,6 +347,24 @@ class TestSafeAiohttpGet: def _run(self, coro): return asyncio.run(coro) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) + def test_maps_timeout_to_fetch_error(self, _): + class _FakeSession: + def get(self, url, **kwargs): + raise asyncio.TimeoutError("timed out") + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + pass + + with patch( + "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + ): + with pytest.raises(MediaFetchError, match="fetch failed"): + self._run(_safe_aiohttp_get("http://example.com/file.bin")) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_reads_response_chunks_to_eof(self, _): class _FakeResponse: From 856630a6bdc581a6fc0a80fb2cbbc6c316debcb5 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 16:59:18 -0700 Subject: [PATCH 09/41] style(openengine): tighten media comments Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 28 ++++++------------- .../grpc/openengine/request_mapping.py | 23 ++++----------- tensorrt_llm/grpc/openengine/servicer.py | 2 +- 3 files changed, 14 insertions(+), 39 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 6affcf574747..1a7f76bbd90b 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -1,11 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""What this engine can actually do, defined once for Generate and Control. - -Control advertises from this table and Generate enforces it, so a client can -never discover a capability that Generate then rejects. -""" +"""Shared capability resolution for OpenEngine Control and Generate.""" from dataclasses import dataclass from types import MappingProxyType @@ -29,7 +25,7 @@ @dataclass(frozen=True) class MediaCapabilities: - """Immutable media capability snapshot shared by Control and Generate.""" + """Media support shared by Control and Generate.""" placeholder_templates: Mapping[str, str] @@ -40,7 +36,7 @@ def __post_init__(self) -> None: @property def modalities(self) -> frozenset[str]: - """Modalities the loaded model and request mapper both support.""" + """Modalities supported by both the model and request mapper.""" return frozenset(self.placeholder_templates) def placeholder(self, modality: str, item_number: int) -> str: @@ -73,8 +69,7 @@ def _loaded_model_type(llm: Any) -> str: if config is None: return "" - # Composite Hugging Face configs can delegate the instance attribute to - # their text config. The class attribute retains the top-level registry key. + # The class attribute preserves the top-level type for composite configs. model_type = getattr(type(config), "model_type", None) or getattr(config, "model_type", "") return model_type if isinstance(model_type, str) else "" @@ -82,14 +77,9 @@ def _loaded_model_type(llm: Any) -> str: def resolve_media_capabilities( llm: Any, implemented_modalities: Collection[str] ) -> MediaCapabilities: - """Resolve media support implemented by both the mapper and loaded model. - - A registry entry alone describes what this TensorRT-LLM build knows about; - it does not prove that the loaded engine instantiated a multimodal input - processor. Requiring both prevents a text-only model from advertising - media merely because another image model is registered in the process. The - request mapper supplies ``implemented_modalities`` from its handler table, - so discovery cannot advertise a modality without a request implementation. + """Intersect model support with implemented request handlers. + + A live multimodal processor prevents registry-only false positives. """ input_processor = getattr(llm, "input_processor", None) if not isinstance(input_processor, BaseMultimodalInputProcessor): @@ -103,9 +93,7 @@ def resolve_media_capabilities( for modality in implemented_modalities if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) } - # Validate templates once at startup. A malformed registry entry is a - # server configuration failure, not a client request error discovered only - # after Control has advertised the modality. + # Fail startup before Control advertises a malformed template. for modality, template in placeholders.items(): try: template.format(1) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 28c948ba89d4..0f4ab8603c74 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -27,14 +27,10 @@ # output index and a scan of every generated token, so an unbounded list lets a # single request tax the event loop for the whole stream. _MAX_STOP_CONDITIONS = 64 -# Media decoding can retain both encoded and decoded image buffers. Keep one -# request from scheduling an unbounded number of fetch/decode tasks or holding -# an unbounded amount of image data while preserving practical multi-image use. +# Bound per-request decode work and retained image data. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 -# Encoded bodies are separately bounded by the 64 MiB gRPC message limit and -# the media loader's 200 MiB per-URL response limit. This cap covers expansion -# during decode across the entire request. +# Encoded inputs have separate gRPC and remote-loader caps. _MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -45,8 +41,6 @@ @dataclass(frozen=True) class _MediaHandler: - """One fully implemented OpenEngine media mapping path.""" - name: str loader_factory: Callable[[Any], _MediaLoader] decoded_bytes: Callable[[Any], int] @@ -285,7 +279,6 @@ def _base_input_from_request( def _image_io_for_model(llm: Any) -> ImageMediaIO: - """Create image I/O using the loaded model's preferred decode format.""" input_processor = getattr(llm, "input_processor", None) defaults = None if isinstance(input_processor, BaseMultimodalInputProcessor): @@ -297,7 +290,6 @@ def _image_io_for_model(llm: Any) -> ImageMediaIO: async def _load_image_item( item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int ) -> Any: - """Resolve and decode one OpenEngine image item.""" source = item.WhichOneof("source") if source is None: raise ValueError(f"media[{index}] must set exactly one source") @@ -339,15 +331,13 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: def _decoded_image_bytes(image: Any) -> int: - """Return retained bytes for a decoded image representation.""" if isinstance(image, torch.Tensor): return image.numel() * image.element_size() if isinstance(image, np.ndarray): return image.nbytes if isinstance(image, Image.Image): return image.width * image.height * len(image.getbands()) - # Custom model media types cannot be sized generically. The built-in image - # handler returns one of the representations above. + # ImageMediaIO currently returns only the representations above. return 0 @@ -359,7 +349,6 @@ def _decoded_image_bytes(image: Any) -> int: def media_capabilities_for_model(llm: Any) -> MediaCapabilities: - """Resolve the loaded model against the request handlers in this module.""" return resolve_media_capabilities(llm, {handler.name for handler in _MEDIA_HANDLERS.values()}) @@ -375,7 +364,7 @@ async def _input_from_request( llm: Any, media_capabilities: MediaCapabilities | None = None, ) -> str | dict[str, Any]: - """Build LLM API inputs, resolving media off the event loop.""" + """Build LLM API inputs without blocking the event loop on media.""" inputs = _base_input_from_request(request) if not request.media: return inputs @@ -429,9 +418,7 @@ async def load_bounded( try: decoded_media = await asyncio.gather(*tasks) except (Exception, asyncio.CancelledError): - # `gather` propagates the first failure without cancelling siblings. - # Explicitly stop any in-flight URL fetches/decodes before returning - # the request error so their memory and connections do not linger. + # gather leaves siblings running after the first failure. for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index f1547b65a112..f2dabe368a32 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -69,7 +69,7 @@ def __init__( @property def supported_media_modalities(self) -> frozenset[str]: - """Media capabilities cached from the loaded model and mapper.""" + """Cached media modalities.""" return self._media_capabilities.modalities def active_request_count(self) -> int: From 423d2b1704efdf8d1f836e1500c0bdbe5adf4499 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 17:25:55 -0700 Subject: [PATCH 10/41] fix(openengine): harden decoded media accounting Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 3 +++ .../grpc/openengine/request_mapping.py | 25 ++++++++++------- .../grpc/openengine/test_multimodal.py | 27 +++++++++++++++++++ 3 files changed, 45 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 1a7f76bbd90b..4f553a4b8bdc 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -34,6 +34,9 @@ def __post_init__(self) -> None: self, "placeholder_templates", MappingProxyType(dict(self.placeholder_templates)) ) + def __hash__(self) -> int: + return hash(frozenset(self.placeholder_templates.items())) + @property def modalities(self) -> frozenset[str]: """Modalities supported by both the model and request mapper.""" diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 0f4ab8603c74..cfd7c17bf2ee 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -10,9 +10,6 @@ from urllib.parse import urlparse import grpc -import numpy as np -import torch -from PIL import Image from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError @@ -30,7 +27,8 @@ # Bound per-request decode work and retained image data. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 -# Encoded inputs have separate gRPC and remote-loader caps. +# Checked after each decode; concurrency bounds peak overshoot. Encoded inputs +# have separate gRPC and remote-loader caps. _MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -331,14 +329,17 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: def _decoded_image_bytes(image: Any) -> int: + import numpy as np + import torch + from PIL import Image + if isinstance(image, torch.Tensor): return image.numel() * image.element_size() if isinstance(image, np.ndarray): return image.nbytes if isinstance(image, Image.Image): return image.width * image.height * len(image.getbands()) - # ImageMediaIO currently returns only the representations above. - return 0 + raise TypeError(f"unsupported decoded image type: {type(image).__name__}") _MEDIA_HANDLERS: dict[int, _MediaHandler] = { @@ -376,7 +377,8 @@ async def _input_from_request( ) capabilities = media_capabilities or media_capabilities_for_model(llm) - if not capabilities.modalities: + modalities = capabilities.modalities + if not modalities: raise UnsupportedFeatureError("media input is not supported by the loaded model") resolved_handlers: list[_MediaHandler] = [] @@ -384,16 +386,19 @@ async def _input_from_request( if item.modality == generation_pb2.MODALITY_UNSPECIFIED: raise ValueError(f"media[{index}].modality must be set") handler = _MEDIA_HANDLERS.get(item.modality) - if handler is None or handler.name not in capabilities.modalities: + if handler is None or handler.name not in modalities: modality = _modality_name(item.modality) - supported_text = ", ".join(sorted(capabilities.modalities)) or "none" + supported_text = ", ".join(sorted(modalities)) or "none" raise UnsupportedFeatureError( f"multimodal modality '{modality}' is not supported by the loaded model; " f"supported modalities: {supported_text}" ) resolved_handlers.append(handler) - loaders = {handler.name: handler.loader_factory(llm) for handler in set(resolved_handlers)} + loaders: dict[str, _MediaLoader] = {} + for handler in resolved_handlers: + if handler.name not in loaders: + loaders[handler.name] = handler.loader_factory(llm) semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) decoded_bytes = 0 diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 43ff2297861b..f9aa7985f4c6 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -18,6 +18,7 @@ import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 @@ -29,6 +30,15 @@ pytestmark = pytest.mark.cpu_only +def _ignore_synthetic_decoded_size(monkeypatch) -> None: + handler = request_mapping._MEDIA_HANDLERS[generation_pb2.MODALITY_IMAGE] + monkeypatch.setitem( + request_mapping._MEDIA_HANDLERS, + generation_pb2.MODALITY_IMAGE, + request_mapping._MediaHandler(handler.name, handler.loader_factory, lambda _decoded: 0), + ) + + class _ImageConfig: # Composite configs can expose their text model type on the instance; the # top-level class attribute is the key used by the multimodal registry. @@ -171,6 +181,19 @@ def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): ) +def test_media_capabilities_are_hashable(): + first = MediaCapabilities({"image": ""}) + second = MediaCapabilities({"image": ""}) + + assert first == second + assert hash(first) == hash(second) + + +def test_unknown_decoded_image_type_is_rejected(): + with pytest.raises(TypeError, match="unsupported decoded image type"): + request_mapping._decoded_image_bytes(object()) + + def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): llm, _ = image_llm monkeypatch.setattr( @@ -264,6 +287,7 @@ async def async_load_bytes(self, data): return f"decoded-bytes:{data.decode()}" monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) request = generation_pb2.GenerateRequest( prompt=" describe", media=[ @@ -326,6 +350,7 @@ async def async_load_bytes(self, data): type(self).active -= 1 monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) request = generation_pb2.GenerateRequest( prompt="images", media=[ @@ -453,6 +478,7 @@ async def async_load_bytes(self, data): return "decoded" monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) request = generation_pb2.GenerateRequest( token_ids=generation_pb2.TokenIds(ids=[10, 20]), media=[ @@ -487,6 +513,7 @@ async def async_load_bytes(self, data): return "decoded" monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) servicer = OpenEngineInferenceServicer(llm, model="image-model") request = generation_pb2.GenerateRequest( request_id="image-request", From 5d2e3b2c758b4ccc0fa0a3eb0e01cdef934e8089 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Fri, 16 Oct 2026 11:09:34 -0700 Subject: [PATCH 11/41] fix(media): prevent DNS rebinding during async fetches Signed-off-by: krishung5 --- tensorrt_llm/inputs/media_io.py | 68 ++++++++++----- .../inputs/test_async_media_loading.py | 1 + tests/unittest/inputs/test_url_validation.py | 85 +++++++++++++++++-- 3 files changed, 124 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 484fdf5b6f70..dc8cb2de387e 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -97,7 +97,8 @@ async def _get_aiohttp_session() -> aiohttp.ClientSession: """Return the shared aiohttp.ClientSession, creating it on first call.""" global _global_aiohttp_session if _global_aiohttp_session is None or _global_aiohttp_session.closed: - _global_aiohttp_session = aiohttp.ClientSession() + connector = aiohttp.TCPConnector(resolver=_ValidatingResolver()) + _global_aiohttp_session = aiohttp.ClientSession(connector=connector) return _global_aiohttp_session @@ -138,6 +139,42 @@ def _allow_private_urls() -> bool: return True +def _validate_resolved_address(address: str) -> None: + try: + ip = ipaddress.ip_address(address) + except ValueError as exc: + raise MediaFetchError(f"DNS resolver returned an invalid address: {address!r}") from exc + if ip.is_multicast: + raise RuntimeError(f"URL resolves to a multicast address ({ip})") + if not ip.is_global and not _allow_private_urls(): + raise RuntimeError( + f"URL resolves to a non-public address ({ip}); set " + f"{_ALLOW_PRIVATE_URLS_ENV}=1 to allow private media URLs " + f"in trusted network environments" + ) + + +class _ValidatingResolver(aiohttp.abc.AbstractResolver): + """Reject unsafe addresses in the DNS result used by the TCP connector.""" + + def __init__(self) -> None: + self._resolver = aiohttp.DefaultResolver() + + async def resolve( + self, host: str, port: int = 0, family: int = socket.AF_INET + ) -> List[Dict[str, Any]]: + try: + addresses = await self._resolver.resolve(host, port, family) + except OSError as exc: + raise MediaFetchError(f"Could not resolve hostname {host!r}") from exc + for address in addresses: + _validate_resolved_address(address["host"]) + return addresses + + async def close(self) -> None: + await self._resolver.close() + + def _validate_url(url: str) -> None: """Validate that *url* points to a public, non-internal HTTP(S) resource. @@ -147,11 +184,9 @@ def _validate_url(url: str) -> None: rejection (scheme, hostname, and multicast checks still apply); see ``_ALLOW_PRIVATE_URLS_ENV`` above for the security tradeoff. - Note: validation is performed at DNS-resolution time. A DNS-rebinding - attack (TTL=0, resolves to a public IP during validation then a private IP - during the actual TCP connect) could bypass this check. For strict - isolation, supplement with network-level egress filtering that blocks - non-global address ranges at the host firewall. + Async fetches repeat the address check in the resolver used by their TCP + connector, so the address admitted here cannot be replaced with an unsafe + one before connection. """ parsed = urlparse(url) if parsed.scheme not in ("http", "https"): @@ -167,17 +202,8 @@ def _validate_url(url: str) -> None: except socket.gaierror as exc: raise MediaFetchError(f"Could not resolve hostname {hostname!r}") from exc - allow_private = _allow_private_urls() for _family, _type, _proto, _canon, sockaddr in infos: - ip = ipaddress.ip_address(sockaddr[0]) - if ip.is_multicast: - raise RuntimeError(f"URL resolves to a multicast address ({ip})") - if not ip.is_global and not allow_private: - raise RuntimeError( - f"URL resolves to a non-public address ({ip}); set " - f"{_ALLOW_PRIVATE_URLS_ENV}=1 to allow private media URLs " - f"in trusted network environments" - ) + _validate_resolved_address(sockaddr[0]) def _buffer_requests_response(resp: "requests.Response") -> "requests.Response": @@ -255,6 +281,7 @@ async def _safe_aiohttp_get( url: str, timeout_sec: int = 30, session: Optional[aiohttp.ClientSession] = None ) -> bytes: """Aiohttp GET wrapper that validates every redirect hop before following.""" + del session # Kept for API compatibility; fetches must use the validating connector. await asyncio.to_thread(_validate_url, url) timeout = aiohttp.ClientTimeout(total=timeout_sec) @@ -286,11 +313,7 @@ async def _fetch(fetch_session: aiohttp.ClientSession) -> bytes: raise MediaFetchError(f"remote media fetch failed: {error}") from error raise RuntimeError("Too many redirects") - if session is not None: - return await _fetch(session) - - async with aiohttp.ClientSession(timeout=timeout) as owned_session: - return await _fetch(owned_session) + return await _fetch(await _get_aiohttp_session()) def _load_and_convert_image(image): @@ -995,8 +1018,7 @@ async def async_load(self, url: str) -> _MediaT: """ parsed = urlparse(url) if parsed.scheme in ("http", "https"): - session = await _get_aiohttp_session() - data = await _safe_aiohttp_get(url, session=session) + data = await _safe_aiohttp_get(url) return await self._run_in_executor(self.load_bytes, data) elif parsed.scheme == "data": data_spec, b64_data = parsed.path.split(",", 1) diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index 997b140f79f7..90a1b8731fb4 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -170,6 +170,7 @@ async def test_same_session_returned_on_repeated_calls(self): assert first_session is second_session, ( "Expected the same ClientSession object to be returned" ) + assert isinstance(first_session.connector._resolver, media_io_module._ValidatingResolver) @pytest.mark.asyncio async def test_new_session_created_after_close(self): diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index c2d0421ee5bc..374e3128a460 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -20,8 +20,9 @@ import asyncio import socket -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch +import aiohttp import pytest from tensorrt_llm.inputs.media_io import ( @@ -30,6 +31,7 @@ _safe_aiohttp_get, _safe_request_get, _validate_url, + _ValidatingResolver, ) pytestmark = pytest.mark.cpu_only @@ -66,6 +68,35 @@ async def read(self, _): return self._chunks.pop(0) +class _AiohttpErrorResponse: + headers = {} + + def __init__(self, status): + self.status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + pass + + def raise_for_status(self): + raise aiohttp.ClientResponseError( + request_info=MagicMock(), + history=(), + status=self.status, + message="test status", + ) + + +class _AiohttpSession: + def __init__(self, response): + self._response = response + + def get(self, url, **kwargs): + return self._response + + # --------------------------------------------------------------------------- # _validate_url # --------------------------------------------------------------------------- @@ -360,11 +391,31 @@ async def __aexit__(self, *_): pass with patch( - "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_FakeSession()), ): with pytest.raises(MediaFetchError, match="fetch failed"): self._run(_safe_aiohttp_get("http://example.com/file.bin")) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) + def test_preserves_non_retryable_http_error(self, _): + with patch( + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_AiohttpSession(_AiohttpErrorResponse(404))), + ): + with pytest.raises(aiohttp.ClientResponseError) as error: + self._run(_safe_aiohttp_get("http://example.com/missing.jpg")) + assert error.value.status == 404 + + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) + def test_maps_retryable_http_error_to_fetch_error(self, _): + with patch( + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_AiohttpSession(_AiohttpErrorResponse(503))), + ): + with pytest.raises(MediaFetchError, match="HTTP 503"): + self._run(_safe_aiohttp_get("http://example.com/unavailable.jpg")) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_reads_response_chunks_to_eof(self, _): class _FakeResponse: @@ -392,9 +443,15 @@ async def __aexit__(self, *_): pass with patch( - "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_FakeSession()), ): - assert self._run(_safe_aiohttp_get("http://example.com/file.bin")) == b"abcd" + unsafe_session = MagicMock() + unsafe_session.get.side_effect = AssertionError("caller session must not be used") + assert ( + self._run(_safe_aiohttp_get("http://example.com/file.bin", session=unsafe_session)) + == b"abcd" + ) @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_raises_after_too_many_redirects(self, _): @@ -422,7 +479,8 @@ async def __aexit__(self, *_): pass with patch( - "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_FakeSession()), ): with pytest.raises(RuntimeError, match="Too many redirects"): self._run(_safe_aiohttp_get("http://example.com/image.jpg")) @@ -455,7 +513,8 @@ async def __aexit__(self, *_): with ( patch( - "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_FakeSession()), ), patch("tensorrt_llm.inputs.media_io._MAX_RESPONSE_BYTES", 3), ): @@ -469,8 +528,20 @@ def test_rejects_private_url_before_request(self): "tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=_dns("169.254.169.254"), ), - patch("tensorrt_llm.inputs.media_io.aiohttp.ClientSession") as mock_session, + patch("tensorrt_llm.inputs.media_io._get_aiohttp_session") as mock_session, ): with pytest.raises(RuntimeError, match="non-public"): self._run(_safe_aiohttp_get("http://169.254.169.254/latest/")) mock_session.assert_not_called() + + def test_rejects_private_address_from_connection_resolver(self): + async def resolve(): + resolver = _ValidatingResolver() + resolver._resolver.resolve = AsyncMock(return_value=[{"host": "169.254.169.254"}]) + try: + await resolver.resolve("example.com", 80) + finally: + await resolver.close() + + with pytest.raises(RuntimeError, match="non-public"): + self._run(resolve()) From 87f270c033627fcd1c54bb8bfc5c43daf0779f6d Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 11:16:15 -0700 Subject: [PATCH 12/41] feat(openengine): support image inputs Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 57 ++ tensorrt_llm/grpc/openengine/control.py | 6 +- .../grpc/openengine/request_mapping.py | 118 +++- tensorrt_llm/grpc/openengine/servicer.py | 4 +- tensorrt_llm/inputs/media_io.py | 4 + .../integration/test_lists/test-db/l0_a10.yml | 8 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../unittest/grpc/openengine/test_control.py | 17 + .../grpc/openengine/test_multimodal.py | 531 ++++++++++++++++++ .../openengine/test_multimodal_conformance.py | 93 +++ .../inputs/test_async_media_loading.py | 9 + 11 files changed, 837 insertions(+), 11 deletions(-) create mode 100644 tests/unittest/grpc/openengine/test_multimodal.py create mode 100644 tests/unittest/grpc/openengine/test_multimodal_conformance.py diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 6272d1363855..a1936ebc4c8a 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -7,6 +7,17 @@ never discover a capability that Generate then rejects. """ +from typing import Any + +from tensorrt_llm.inputs.registry import ( + MULTIMODAL_PLACEHOLDER_REGISTRY, + BaseMultimodalInputProcessor, +) + +# OpenEngine support is intentionally narrower than the LLM API today. Add a +# modality here only once request mapping and conformance coverage exist for it. +_OPENENGINE_MEDIA_MODALITIES = ("image",) + # Guides each grammar backend can build, keyed by the `GuidedDecoding` oneof # field name. Control advertises from this table and Generate enforces it, so # the two cannot disagree. llguidance has no structural-tag matcher. @@ -32,7 +43,53 @@ def supported_guides(backend: str | None) -> frozenset[str]: return GUIDE_SUPPORT_BY_BACKEND.get(backend.lower(), frozenset()) +def _loaded_model_type(llm: Any) -> str: + """Return the top-level model type for the loaded input processor.""" + input_processor = getattr(llm, "input_processor", None) + config = ( + input_processor.config + if isinstance(input_processor, BaseMultimodalInputProcessor) + else getattr(llm, "_hf_model_config", None) + ) + if config is None: + return "" + + # Composite Hugging Face configs can delegate the instance attribute to + # their text config. The class attribute retains the top-level registry key. + model_type = getattr(type(config), "model_type", None) or getattr(config, "model_type", "") + return model_type if isinstance(model_type, str) else "" + + +def supported_media_modalities(llm: Any) -> frozenset[str]: + """OpenEngine media modalities supported by the loaded model. + + A registry entry alone describes what this TensorRT-LLM build knows about; + it does not prove that the loaded engine instantiated a multimodal input + processor. Requiring both prevents a text-only model from advertising + media merely because another image model is registered in the process. + """ + input_processor = getattr(llm, "input_processor", None) + if not isinstance(input_processor, BaseMultimodalInputProcessor): + return frozenset() + model_type = _loaded_model_type(llm) + if not model_type: + return frozenset() + return frozenset( + modality + for modality in _OPENENGINE_MEDIA_MODALITIES + if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) + ) + + +def media_placeholder(llm: Any, modality: str, item_number: int) -> str: + """Return the model-specific placeholder for a one-based media item.""" + model_type = _loaded_model_type(llm) + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality).format(item_number) + + __all__ = [ "GUIDE_SUPPORT_BY_BACKEND", + "media_placeholder", "supported_guides", + "supported_media_modalities", ] diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index f020d3ae82a1..1495524d4130 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -35,7 +35,7 @@ openengine_pb2_grpc, server_pb2, ) -from .capabilities import supported_guides +from .capabilities import supported_guides, supported_media_modalities from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError from .errors import AbortFailedError from .kv_events import data_parallel_size, events_config @@ -284,6 +284,7 @@ async def GetModelInfo( ) -> model_pb2.ModelInfo: # Single-model server: like Generate, any non-empty name resolves to the # loaded model. + media_modalities = supported_media_modalities(self._llm) info = model_pb2.ModelInfo( model_id=self._model, served_model_name=self._model, @@ -292,9 +293,10 @@ async def GetModelInfo( supports_token_ids_input=True, # Generate rejects `lora_name` even on an enable_lora build. supports_lora=False, - supports_multimodal=False, + supports_multimodal=bool(media_modalities), reasoning_parser=str(self._llm.args.reasoning_parser or ""), ) + info.extra.update({"trtllm_supported_modalities": sorted(media_modalities)}) max_seq_len = _positive_int(self._llm.args.max_seq_len) if max_seq_len is not None: diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index db6ecf866c01..7e7035e0a1a8 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -3,22 +3,31 @@ """Translate an OpenEngine GenerateRequest into TensorRT-LLM request inputs.""" +import asyncio from collections.abc import Mapping from typing import Any +from urllib.parse import urlparse import grpc from tensorrt_llm.conversation_params import ConversationParams +from tensorrt_llm.inputs.media_io import ImageMediaIO +from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams from .bindings import generation_pb2 -from .capabilities import supported_guides +from .capabilities import media_placeholder, supported_guides, supported_media_modalities from .errors import UnsupportedFeatureError # Cap on the number of stop conditions. Each one costs a KMP automaton per # output index and a scan of every generated token, so an unbounded list lets a # single request tax the event loop for the whole stream. _MAX_STOP_CONDITIONS = 64 +# Media decoding can retain both encoded and decoded image buffers. Keep one +# request from scheduling an unbounded number of fetch/decode tasks or holding +# an unbounded amount of image data while preserving practical multi-image use. +_MAX_MEDIA_ITEMS = 64 +_MAX_CONCURRENT_MEDIA_LOADS = 8 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -244,7 +253,9 @@ def _metadata_from_context( return headers or None, target_dp_rank, affinity_headers -def _input_from_request(request: generation_pb2.GenerateRequest) -> str | dict[str, list[int]]: +def _base_input_from_request( + request: generation_pb2.GenerateRequest, +) -> str | dict[str, list[int]]: input_kind = request.WhichOneof("input") if input_kind == "prompt": return request.prompt @@ -253,6 +264,109 @@ def _input_from_request(request: generation_pb2.GenerateRequest) -> str | dict[s raise ValueError("exactly one of prompt or token_ids must be set") +def _image_io_for_model(llm: Any) -> ImageMediaIO: + """Create image I/O using the loaded model's preferred decode format.""" + input_processor = getattr(llm, "input_processor", None) + defaults = None + if isinstance(input_processor, BaseMultimodalInputProcessor): + preferences = input_processor.get_preferred_media_io_kwargs() or {} + defaults = preferences.get("image") + return ImageMediaIO.create(defaults, None) + + +async def _load_image_item( + item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int +) -> Any: + """Resolve and decode one OpenEngine image item.""" + source = item.WhichOneof("source") + if source is None: + raise ValueError(f"media[{index}] must set exactly one source") + if source == "url": + if urlparse(item.url).scheme.lower() not in ("http", "https"): + raise ValueError(f"media[{index}].url must use http or https") + value = item.url + elif source == "data_uri": + if urlparse(item.data_uri).scheme.lower() != "data": + raise ValueError(f"media[{index}].data_uri must use the data scheme") + value = item.data_uri + elif source == "raw_bytes": + if item.mime_type and not item.mime_type.lower().startswith("image/"): + raise ValueError(f"media[{index}].mime_type must be an image type for raw image bytes") + try: + return await image_io.async_load_bytes(bytes(item.raw_bytes)) + except Exception as error: + raise ValueError(f"media[{index}].raw_bytes could not be decoded: {error}") from error + else: + raise ValueError(f"media[{index}] has an unknown source '{source}'") + + try: + return await image_io.async_load(value) + except Exception as error: + raise ValueError(f"media[{index}].{source} could not be loaded: {error}") from error + + +async def _input_from_request( + request: generation_pb2.GenerateRequest, llm: Any +) -> str | dict[str, Any]: + """Build LLM API inputs, resolving image media off the event loop.""" + inputs = _base_input_from_request(request) + if not request.media: + return inputs + + supported = supported_media_modalities(llm) + for index, item in enumerate(request.media): + if item.modality == generation_pb2.MODALITY_UNSPECIFIED: + raise ValueError(f"media[{index}].modality must be set") + if item.modality != generation_pb2.MODALITY_IMAGE: + modality = generation_pb2.Modality.Name(item.modality).removeprefix("MODALITY_").lower() + supported_text = ", ".join(sorted(supported)) or "none" + raise UnsupportedFeatureError( + f"multimodal modality '{modality}' is not supported by the loaded model; " + f"supported modalities: {supported_text}" + ) + if "image" not in supported: + raise UnsupportedFeatureError("image input is not supported by the loaded model") + if len(request.media) > _MAX_MEDIA_ITEMS: + raise ValueError( + f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " + f"got {len(request.media)}" + ) + + image_io = _image_io_for_model(llm) + semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) + + async def load_bounded(item: generation_pb2.MediaItem, index: int) -> Any: + async with semaphore: + return await _load_image_item(item, image_io, index) + + tasks = [ + asyncio.create_task(load_bounded(item, index)) for index, item in enumerate(request.media) + ] + try: + images = await asyncio.gather(*tasks) + except (Exception, asyncio.CancelledError): + # `gather` propagates the first failure without cancelling siblings. + # Explicitly stop any in-flight URL fetches/decodes before returning + # the request error so their memory and connections do not linger. + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + mapped: dict[str, Any] = {"prompt": inputs} if isinstance(inputs, str) else dict(inputs) + mapped["multi_modal_data"] = {"image": images} + if any(item.uuid for item in request.media): + mapped["multi_modal_uuids"] = {"image": [item.uuid or None for item in request.media]} + mapped["mm_item_order"] = [ + { + "modality": "image", + "index": index, + "placeholder": media_placeholder(llm, "image", index + 1), + } + for index in range(len(request.media)) + ] + return mapped + + __all__ = [ "conversation_params_from_request", "sampling_params_from_request", diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index 6a9d6e9018b9..6bc18ffeb6a4 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -169,15 +169,13 @@ async def _generate( f"request_id '{request_id}' is already active", ) return - if request.media: - raise UnsupportedFeatureError("multimodal media is not supported") if request.lora_name: raise UnsupportedFeatureError("LoRA selection is not supported") if request.HasField("kv"): if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: raise UnsupportedFeatureError("prefix cache bypass is not supported") - inputs = _input_from_request(request) + inputs = await _input_from_request(request, self._llm) sampling_params = sampling_params_from_request(request, self._guided_backend) trace_headers, target_dp_rank, affinity_headers = _metadata_from_context(context) conversation_params = conversation_params_from_request(request) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 86f52f71da3c..e8ab41d4dccf 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -1002,6 +1002,10 @@ async def async_load(self, url: str) -> _MediaT: else: raise ValueError(f"Unsupported URL scheme: {parsed.scheme!r}") + async def async_load_bytes(self, data: bytes) -> _MediaT: + """Decode in-memory media without blocking the event loop.""" + return await self._run_in_executor(self.load_bytes, data) + @staticmethod async def _run_in_executor(fn, *args, **kwargs): """Run a blocking decode callable on the configured executor.""" diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 2b5b6f887860..6ff220686f73 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -63,11 +63,11 @@ l0_a10: - unittest/others/test_tracing.py - unittest/grpc/test_grpc_optional.py - unittest/grpc/smg/test_smg.py - # The one OpenEngine test needing a GPU: it builds a real LLM, because a - # stubbed engine would accept whatever its fake generate_async accepts and so - # would agree with an advertisement the real engine rejects. Any GPU will do - # -- the model is TinyLlama-1.1B -- so it sits with the other gateway tests. + # OpenEngine capability conformance needs real loaded models: a stub would + # accept whatever its fake generate_async tolerates and could agree with a + # false advertisement by construction. The image case uses Qwen3-VL-2B. - unittest/grpc/openengine/test_capability_conformance.py + - unittest/grpc/openengine/test_multimodal_conformance.py - unittest/disaggregated/test_coordinator_e2e.py - unittest/disaggregated/test_coordinator_worker.py - unittest/disaggregated/test_agent_multi_backends.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 58ebb7ca4eba..3fca6014da5c 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -160,6 +160,7 @@ l0_cpu: - unittest/grpc/openengine/test_coordinator.py - unittest/grpc/openengine/test_frontend_lifecycle.py - unittest/grpc/openengine/test_control.py + - unittest/grpc/openengine/test_multimodal.py - unittest/grpc/openengine/test_node.py - unittest/grpc/openengine/test_openengine.py - unittest/grpc/openengine/test_servicer.py diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 10420b6df6bd..f7de8842a222 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -166,6 +166,23 @@ async def test_model_info_reports_the_context_window_and_capabilities(): assert model_pb2.GUIDED_DECODING_MODE_JSON_SCHEMA in info.generation.guided_decoding.modes assert info.generation.output_logprobs.supported is True assert info.supports_token_ids_input is True + assert info.supports_multimodal is False + assert list(info.extra["trtllm_supported_modalities"]) == [] + + +@pytest.mark.asyncio +async def test_model_info_advertises_loaded_image_support(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm.grpc.openengine.control.supported_media_modalities", + lambda llm: frozenset({"image"}), + ) + + info = await _servicer().GetModelInfo( + model_pb2.GetModelInfoRequest(model=MODEL), FakeServicerContext() + ) + + assert info.supports_multimodal is True + assert list(info.extra["trtllm_supported_modalities"]) == ["image"] @pytest.mark.asyncio diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py new file mode 100644 index 000000000000..a4bd537e29f5 --- /dev/null +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -0,0 +1,531 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for OpenEngine image input mapping and discovery.""" + +import asyncio +from collections.abc import AsyncIterator, Sequence +from types import SimpleNamespace +from typing import Any + +import pytest + +grpc = pytest.importorskip( + "grpc", reason='gRPC runtime not installed (pip install "grpcio>=1.67.1,<2")' +) +from conftest import AbortError, FakeServicerContext # noqa: E402 + +import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 +from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.capabilities import supported_media_modalities # noqa: E402 +from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 +from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.inputs.registry import ( # noqa: E402 + MULTIMODAL_PLACEHOLDER_REGISTRY, + BaseMultimodalInputProcessor, +) + +pytestmark = pytest.mark.cpu_only + + +class _ImageConfig: + # Composite configs can expose their text model type on the instance; the + # top-level class attribute is the key used by the multimodal registry. + model_type = "openengine_test_image" + + def __init__(self) -> None: + self.model_type = "openengine_test_text" + + +class _ImageInputProcessor(BaseMultimodalInputProcessor): + """Minimal concrete processor used only for capability detection.""" + + def __init__(self, config: _ImageConfig) -> None: + self._config = config + + @property + def processor(self): + return None + + @property + def tokenizer(self): + return None + + @property + def config(self): + return self._config + + @property + def dtype(self): + return None + + def call_with_text_prompt(self, inputs, sampling_params): + raise NotImplementedError + + def get_preferred_media_io_kwargs(self): + return {"image": {"format": "np"}} + + +class _Tokenizer: + def convert_ids_to_tokens( + self, token_ids: Sequence[int], skip_special_tokens: bool = False + ) -> list[str]: + del skip_special_tokens + return [f"token-{token_id}" for token_id in token_ids] + + def decode(self, token_ids: Sequence[int], skip_special_tokens: bool = False) -> str: + del skip_special_tokens + return "".join(f"token-{token_id}" for token_id in token_ids) + + +class _ResultHandle: + def __init__(self, result: Any) -> None: + self._result = result + self.finished = False + self.prompt_token_ids = [] + self.outputs = [] + self.cached_tokens = 0 + self.error = None + self.aborted = False + + async def __aiter__(self) -> AsyncIterator[Any]: + self.finished = self._result.finished + self.prompt_token_ids = self._result.prompt_token_ids + self.outputs = self._result.outputs + self.cached_tokens = self._result.cached_tokens + self.error = self._result.error + yield self._result + + def abort(self) -> None: + self.aborted = True + + +class _ImageLlm: + def __init__(self, input_processor, config) -> None: + self.input_processor = input_processor + self._hf_model_config = config + self.args = SimpleNamespace(guided_decoding_backend=None) + self.tokenizer = _Tokenizer() + output = SimpleNamespace( + index=0, + token_ids=[99], + text="caption", + logprobs=[], + prompt_logprobs=[], + finish_reason="length", + stop_reason=None, + ) + self.result_handle = _ResultHandle( + SimpleNamespace( + # Represents expansion of one placeholder into five image tokens. + prompt_token_ids=[1, 2, 3, 4, 5, 6], + outputs=[output], + cached_tokens=0, + error=None, + finished=True, + ) + ) + self.generate_kwargs = None + + def generate_async(self, **kwargs): + self.generate_kwargs = kwargs + return self.result_handle + + +@pytest.fixture +def image_llm(monkeypatch): + config = _ImageConfig() + input_processor = _ImageInputProcessor(config) + seen_model_types = [] + + def is_valid(model_type, modality): + seen_model_types.append(model_type) + return model_type == _ImageConfig.model_type and modality == "image" + + monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", is_valid) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: "", + ) + return ( + SimpleNamespace( + _hf_model_config=config, + input_processor=input_processor, + ), + seen_model_types, + ) + + +def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): + llm, seen_model_types = image_llm + + assert supported_media_modalities(llm) == {"image"} + assert seen_model_types == [_ImageConfig.model_type] + assert ( + supported_media_modalities( + SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) + ) + == set() + ) + + +@pytest.mark.asyncio +async def test_text_only_input_mapping_is_unchanged(): + llm = SimpleNamespace() + + prompt = generation_pb2.GenerateRequest(prompt="hello") + token_ids = generation_pb2.GenerateRequest(token_ids=generation_pb2.TokenIds(ids=[1, 2, 3])) + + assert await request_mapping._input_from_request(prompt, llm) == "hello" + assert await request_mapping._input_from_request(token_ids, llm) == { + "prompt_token_ids": [1, 2, 3] + } + + +@pytest.mark.asyncio +async def test_image_sources_preserve_order_uuid_and_model_decode_defaults(image_llm, monkeypatch): + llm, _ = image_llm + created_with = [] + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append((defaults, runtime)) + return cls() + + async def async_load(self, source): + # Finish in the opposite order to prove gather preserves request order. + if source.startswith("https"): + await asyncio.sleep(0.01) + return f"decoded:{source}" + + async def async_load_bytes(self, data): + return f"decoded-bytes:{data.decode()}" + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt=" describe", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="https://example.com/first.png", + uuid="first", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + data_uri="data:image/png;base64,c2Vjb25k", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"third", + mime_type="image/png", + uuid="third", + ), + ], + ) + + mapped = await request_mapping._input_from_request(request, llm) + + assert created_with == [({"format": "np"}, None)] + assert mapped["prompt"] == request.prompt + assert mapped["multi_modal_data"] == { + "image": [ + "decoded:https://example.com/first.png", + "decoded:data:image/png;base64,c2Vjb25k", + "decoded-bytes:third", + ] + } + assert mapped["multi_modal_uuids"] == {"image": ["first", None, "third"]} + assert mapped["mm_item_order"] == [ + {"modality": "image", "index": 0, "placeholder": ""}, + {"modality": "image", "index": 1, "placeholder": ""}, + {"modality": "image", "index": 2, "placeholder": ""}, + ] + + +@pytest.mark.asyncio +async def test_image_loading_has_a_bounded_concurrency(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + active = 0 + maximum_active = 0 + + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + type(self).active += 1 + type(self).maximum_active = max(type(self).maximum_active, type(self).active) + try: + await asyncio.sleep(0.01) + return data + finally: + type(self).active -= 1 + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=str(index).encode(), + ) + for index in range(request_mapping._MAX_CONCURRENT_MEDIA_LOADS * 2 + 1) + ], + ) + + mapped = await request_mapping._input_from_request(request, llm) + + assert _ImageIo.maximum_active == request_mapping._MAX_CONCURRENT_MEDIA_LOADS + assert mapped["multi_modal_data"]["image"] == [ + str(index).encode() for index in range(request_mapping._MAX_CONCURRENT_MEDIA_LOADS * 2 + 1) + ] + + +@pytest.mark.asyncio +async def test_image_loading_cancels_siblings_after_a_failure(image_llm, monkeypatch): + llm, _ = image_llm + slow_started = asyncio.Event() + slow_cancelled = asyncio.Event() + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + if data == b"slow": + slow_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + slow_cancelled.set() + raise + await slow_started.wait() + raise ValueError("decode failed") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=value, + ) + for value in (b"slow", b"broken") + ], + ) + + with pytest.raises(ValueError, match="decode failed"): + await request_mapping._input_from_request(request, llm) + assert slow_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_image_count_is_bounded_before_decoding(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + raise AssertionError("the media-count guard must run before decoder creation") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="too many images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + for _ in range(request_mapping._MAX_MEDIA_ITEMS + 1) + ], + ) + + with pytest.raises(ValueError, match="at most 64 media items"): + await request_mapping._input_from_request(request, llm) + + +@pytest.mark.asyncio +async def test_image_mapping_supports_token_ids(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return "decoded" + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + token_ids=generation_pb2.TokenIds(ids=[10, 20]), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + mapped = await request_mapping._input_from_request(request, llm) + + assert mapped["prompt_token_ids"] == [10, 20] + assert mapped["multi_modal_data"] == {"image": ["decoded"]} + assert "multi_modal_uuids" not in mapped + + +@pytest.mark.asyncio +async def test_generate_submits_image_and_reports_expanded_prompt_usage(image_llm, monkeypatch): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return "decoded" + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + request = generation_pb2.GenerateRequest( + request_id="image-request", + model="image-model", + token_ids=generation_pb2.TokenIds(ids=[1, 2]), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + responses = [response async for response in servicer.Generate(request, FakeServicerContext())] + + assert llm.generate_kwargs["inputs"]["multi_modal_data"] == {"image": ["decoded"]} + assert responses[-1].usage.prompt_tokens == 6 + assert responses[-1].usage.total_tokens == 7 + + +@pytest.mark.asyncio +async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + + unsupported_context = FakeServicerContext() + video_request = generation_pb2.GenerateRequest( + request_id="video-request", + model="image-model", + prompt="video", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"video", + ) + ], + ) + with pytest.raises(AbortError): + async for _ in servicer.Generate(video_request, unsupported_context): + pass + assert unsupported_context.abort_code == grpc.StatusCode.UNIMPLEMENTED + + malformed_context = FakeServicerContext() + file_request = generation_pb2.GenerateRequest( + request_id="file-request", + model="image-model", + prompt="image", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="file:///tmp/private.png", + ) + ], + ) + with pytest.raises(AbortError): + async for _ in servicer.Generate(file_request, malformed_context): + pass + assert malformed_context.abort_code == grpc.StatusCode.INVALID_ARGUMENT + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("modality", "name"), + [ + (generation_pb2.MODALITY_VIDEO, "video"), + (generation_pb2.MODALITY_AUDIO, "audio"), + ], +) +async def test_video_and_audio_are_explicitly_unsupported(image_llm, modality, name): + llm, _ = image_llm + request = generation_pb2.GenerateRequest( + prompt="media", + media=[generation_pb2.MediaItem(modality=modality, raw_bytes=b"data")], + ) + + with pytest.raises(UnsupportedFeatureError, match=name): + await request_mapping._input_from_request(request, llm) + + +@pytest.mark.asyncio +async def test_image_is_rejected_for_a_text_only_loaded_model(): + request = generation_pb2.GenerateRequest( + prompt="image", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + with pytest.raises(UnsupportedFeatureError, match="loaded model"): + await request_mapping._input_from_request( + request, SimpleNamespace(input_processor=object()) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("item", "message"), + [ + (generation_pb2.MediaItem(raw_bytes=b"image"), "modality must be set"), + ( + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE), + "must set exactly one source", + ), + ( + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="file:///tmp/private.png", + ), + "must use http or https", + ), + ( + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + mime_type="video/mp4", + ), + "must be an image type", + ), + ], +) +async def test_malformed_image_items_are_rejected(image_llm, item, message): + llm, _ = image_llm + request = generation_pb2.GenerateRequest(prompt="image", media=[item]) + + with pytest.raises(ValueError, match=message): + await request_mapping._input_from_request(request, llm) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py new file mode 100644 index 000000000000..e97c9a708aac --- /dev/null +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Image conformance tests against a real TensorRT-LLM VLM engine.""" + +import asyncio +import base64 + +import pytest + +pytest.importorskip("grpc", reason='gRPC runtime not installed (pip install "grpcio>=1.67.1,<2")') + +import torch # noqa: E402 +from conftest import FakeServicerContext # noqa: E402 +from utils.llm_data import llm_models_root # noqa: E402 + +from tensorrt_llm import LLM # noqa: E402 +from tensorrt_llm.grpc.openengine.bindings import generation_pb2, model_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 +from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.llmapi import KvCacheConfig # noqa: E402 + +MODEL_NAME = "Qwen3/Qwen3-VL-2B-Instruct" + +# The engine owns background executor threads for the lifetime of this module. +pytestmark = pytest.mark.threadleak(enabled=False) + + +@pytest.fixture(scope="module") +def vlm(): + if not torch.cuda.is_available(): + pytest.skip("requires a GPU to run a real VLM engine") + model_path = llm_models_root(check=True) / MODEL_NAME + engine = LLM( + model=str(model_path), + load_format="dummy", + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.5), + max_batch_size=1, + max_seq_len=2048, + ) + try: + yield engine + finally: + engine.shutdown() + + +def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm): + inference = OpenEngineInferenceServicer(vlm, MODEL_NAME) + control = OpenEngineControlServicer(vlm, MODEL_NAME, inference) + info = asyncio.run( + control.GetModelInfo( + model_pb2.GetModelInfoRequest(model=MODEL_NAME), + FakeServicerContext(), + ) + ) + + assert info.supports_multimodal is True + assert list(info.extra["trtllm_supported_modalities"]) == ["image"] + + prompt = ( + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + "Describe this image briefly.<|im_end|>\n<|im_start|>assistant\n" + ) + prompt_token_ids = vlm.tokenizer.encode(prompt) + image_path = llm_models_root(check=True) / "multimodals/test_data/seashore.png" + image_bytes = image_path.read_bytes() + data_uri = "data:image/png;base64," + base64.b64encode(image_bytes).decode() + request = generation_pb2.GenerateRequest( + request_id="openengine-image-conformance", + model=MODEL_NAME, + token_ids=generation_pb2.TokenIds(ids=prompt_token_ids), + stopping=generation_pb2.StoppingOptions(max_tokens=4), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + data_uri=data_uri, + uuid="seashore", + ) + ], + ) + + async def collect(): + return [response async for response in inference.Generate(request, FakeServicerContext())] + + responses = asyncio.run(collect()) + assert responses + assert not any(response.HasField("error") for response in responses) + finished = [response for response in responses if response.HasField("finished")] + assert len(finished) == 1 + assert finished[0].usage.prompt_tokens > len(prompt_token_ids) + assert finished[0].usage.total_tokens == ( + finished[0].usage.prompt_tokens + finished[0].usage.completion_tokens + ) diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index 5855605a06ca..997b140f79f7 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -58,6 +58,15 @@ def _make_audio_file(tmp_path: str) -> str: class TestAsyncLoadImage: + @pytest.mark.asyncio + async def test_load_image_from_raw_bytes_pil(self): + encoded = _make_image_data_url().partition(",")[2] + image = await media_io_module.ImageMediaIO(format="pil").async_load_bytes( + base64.b64decode(encoded) + ) + assert isinstance(image, Image.Image) + assert image.mode == "RGB" + @pytest.mark.asyncio async def test_load_image_from_data_url_pil(self): url = _make_image_data_url() From d2cf1146c0e6dbe26719fc18acfba6a642ef3abf Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 16:55:11 -0700 Subject: [PATCH 13/41] fix(openengine): harden multimodal media handling Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 74 ++++++--- tensorrt_llm/grpc/openengine/control.py | 4 +- tensorrt_llm/grpc/openengine/errors.py | 5 + .../grpc/openengine/request_mapping.py | 157 ++++++++++++++---- tensorrt_llm/grpc/openengine/servicer.py | 14 +- tensorrt_llm/inputs/media_io.py | 41 +++-- .../unittest/grpc/openengine/test_control.py | 9 +- .../grpc/openengine/test_multimodal.py | 138 ++++++++++++++- tests/unittest/inputs/test_url_validation.py | 21 ++- 9 files changed, 383 insertions(+), 80 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index a1936ebc4c8a..6affcf574747 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -7,17 +7,15 @@ never discover a capability that Generate then rejects. """ -from typing import Any +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Collection, Mapping from tensorrt_llm.inputs.registry import ( MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, ) -# OpenEngine support is intentionally narrower than the LLM API today. Add a -# modality here only once request mapping and conformance coverage exist for it. -_OPENENGINE_MEDIA_MODALITIES = ("image",) - # Guides each grammar backend can build, keyed by the `GuidedDecoding` oneof # field name. Control advertises from this table and Generate enforces it, so # the two cannot disagree. llguidance has no structural-tag matcher. @@ -29,6 +27,27 @@ } +@dataclass(frozen=True) +class MediaCapabilities: + """Immutable media capability snapshot shared by Control and Generate.""" + + placeholder_templates: Mapping[str, str] + + def __post_init__(self) -> None: + object.__setattr__( + self, "placeholder_templates", MappingProxyType(dict(self.placeholder_templates)) + ) + + @property + def modalities(self) -> frozenset[str]: + """Modalities the loaded model and request mapper both support.""" + return frozenset(self.placeholder_templates) + + def placeholder(self, modality: str, item_number: int) -> str: + """Format the model placeholder for a one-based media item.""" + return self.placeholder_templates[modality].format(item_number) + + def supported_guides(backend: str | None) -> frozenset[str]: """Guides `backend` can build; empty when none can be. @@ -60,36 +79,47 @@ def _loaded_model_type(llm: Any) -> str: return model_type if isinstance(model_type, str) else "" -def supported_media_modalities(llm: Any) -> frozenset[str]: - """OpenEngine media modalities supported by the loaded model. +def resolve_media_capabilities( + llm: Any, implemented_modalities: Collection[str] +) -> MediaCapabilities: + """Resolve media support implemented by both the mapper and loaded model. A registry entry alone describes what this TensorRT-LLM build knows about; it does not prove that the loaded engine instantiated a multimodal input processor. Requiring both prevents a text-only model from advertising - media merely because another image model is registered in the process. + media merely because another image model is registered in the process. The + request mapper supplies ``implemented_modalities`` from its handler table, + so discovery cannot advertise a modality without a request implementation. """ input_processor = getattr(llm, "input_processor", None) if not isinstance(input_processor, BaseMultimodalInputProcessor): - return frozenset() + return MediaCapabilities({}) model_type = _loaded_model_type(llm) if not model_type: - return frozenset() - return frozenset( - modality - for modality in _OPENENGINE_MEDIA_MODALITIES - if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) - ) + return MediaCapabilities({}) - -def media_placeholder(llm: Any, modality: str, item_number: int) -> str: - """Return the model-specific placeholder for a one-based media item.""" - model_type = _loaded_model_type(llm) - return MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality).format(item_number) + placeholders = { + modality: MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality) + for modality in implemented_modalities + if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) + } + # Validate templates once at startup. A malformed registry entry is a + # server configuration failure, not a client request error discovered only + # after Control has advertised the modality. + for modality, template in placeholders.items(): + try: + template.format(1) + except (IndexError, KeyError, ValueError) as error: + raise RuntimeError( + f"invalid multimodal placeholder for model type '{model_type}' " + f"and modality '{modality}': {error}" + ) from error + return MediaCapabilities(placeholders) __all__ = [ "GUIDE_SUPPORT_BY_BACKEND", - "media_placeholder", + "MediaCapabilities", + "resolve_media_capabilities", "supported_guides", - "supported_media_modalities", ] diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index 1495524d4130..8b4292dd33ae 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -35,7 +35,7 @@ openengine_pb2_grpc, server_pb2, ) -from .capabilities import supported_guides, supported_media_modalities +from .capabilities import supported_guides from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError from .errors import AbortFailedError from .kv_events import data_parallel_size, events_config @@ -284,7 +284,7 @@ async def GetModelInfo( ) -> model_pb2.ModelInfo: # Single-model server: like Generate, any non-empty name resolves to the # loaded model. - media_modalities = supported_media_modalities(self._llm) + media_modalities = self._inference.supported_media_modalities info = model_pb2.ModelInfo( model_id=self._model, served_model_name=self._model, diff --git a/tensorrt_llm/grpc/openengine/errors.py b/tensorrt_llm/grpc/openengine/errors.py index 267d2a508c8a..d639794e0c15 100644 --- a/tensorrt_llm/grpc/openengine/errors.py +++ b/tensorrt_llm/grpc/openengine/errors.py @@ -8,6 +8,10 @@ class UnsupportedFeatureError(ValueError): """Raised when OpenEngine requests a feature this adapter cannot map.""" +class MediaUnavailableError(RuntimeError): + """Raised when fetching remote request media fails transiently.""" + + class AbortFailedError(RuntimeError): """Raised when an active request could not be aborted on the engine.""" @@ -18,5 +22,6 @@ def __init__(self, request_id: str, cause: BaseException) -> None: __all__ = [ "AbortFailedError", + "MediaUnavailableError", "UnsupportedFeatureError", ] diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 7e7035e0a1a8..28c948ba89d4 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -4,20 +4,24 @@ """Translate an OpenEngine GenerateRequest into TensorRT-LLM request inputs.""" import asyncio -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass from typing import Any from urllib.parse import urlparse import grpc +import numpy as np +import torch +from PIL import Image from tensorrt_llm.conversation_params import ConversationParams -from tensorrt_llm.inputs.media_io import ImageMediaIO +from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams from .bindings import generation_pb2 -from .capabilities import media_placeholder, supported_guides, supported_media_modalities -from .errors import UnsupportedFeatureError +from .capabilities import MediaCapabilities, resolve_media_capabilities, supported_guides +from .errors import MediaUnavailableError, UnsupportedFeatureError # Cap on the number of stop conditions. Each one costs a KMP automaton per # output index and a scan of every generated token, so an unbounded list lets a @@ -28,10 +32,26 @@ # an unbounded amount of image data while preserving practical multi-image use. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 +# Encoded bodies are separately bounded by the 64 MiB gRPC message limit and +# the media loader's 200 MiB per-URL response limit. This cap covers expansion +# during decode across the entire request. +_MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" +_MediaLoader = Callable[[generation_pb2.MediaItem, int], Awaitable[Any]] + + +@dataclass(frozen=True) +class _MediaHandler: + """One fully implemented OpenEngine media mapping path.""" + + name: str + loader_factory: Callable[[Any], _MediaLoader] + decoded_bytes: Callable[[Any], int] + + def conversation_affinity_enabled(llm: Any) -> bool: """Whether TensorRT-LLM owns attention-DP placement for conversations.""" args = getattr(llm, "args", None) @@ -301,49 +321,113 @@ async def _load_image_item( try: return await image_io.async_load(value) + except MediaFetchError as error: + raise MediaUnavailableError( + f"media[{index}].{source} is temporarily unavailable: {error}" + ) from error except Exception as error: raise ValueError(f"media[{index}].{source} could not be loaded: {error}") from error +def _image_loader_factory(llm: Any) -> _MediaLoader: + image_io = _image_io_for_model(llm) + + async def load(item: generation_pb2.MediaItem, index: int) -> Any: + return await _load_image_item(item, image_io, index) + + return load + + +def _decoded_image_bytes(image: Any) -> int: + """Return retained bytes for a decoded image representation.""" + if isinstance(image, torch.Tensor): + return image.numel() * image.element_size() + if isinstance(image, np.ndarray): + return image.nbytes + if isinstance(image, Image.Image): + return image.width * image.height * len(image.getbands()) + # Custom model media types cannot be sized generically. The built-in image + # handler returns one of the representations above. + return 0 + + +_MEDIA_HANDLERS: dict[int, _MediaHandler] = { + generation_pb2.MODALITY_IMAGE: _MediaHandler( + "image", _image_loader_factory, _decoded_image_bytes + ), +} + + +def media_capabilities_for_model(llm: Any) -> MediaCapabilities: + """Resolve the loaded model against the request handlers in this module.""" + return resolve_media_capabilities(llm, {handler.name for handler in _MEDIA_HANDLERS.values()}) + + +def _modality_name(modality: int) -> str: + try: + return generation_pb2.Modality.Name(modality).removeprefix("MODALITY_").lower() + except ValueError: + return f"unknown ({modality})" + + async def _input_from_request( - request: generation_pb2.GenerateRequest, llm: Any + request: generation_pb2.GenerateRequest, + llm: Any, + media_capabilities: MediaCapabilities | None = None, ) -> str | dict[str, Any]: - """Build LLM API inputs, resolving image media off the event loop.""" + """Build LLM API inputs, resolving media off the event loop.""" inputs = _base_input_from_request(request) if not request.media: return inputs - supported = supported_media_modalities(llm) + if len(request.media) > _MAX_MEDIA_ITEMS: + raise ValueError( + f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " + f"got {len(request.media)}" + ) + + capabilities = media_capabilities or media_capabilities_for_model(llm) + if not capabilities.modalities: + raise UnsupportedFeatureError("media input is not supported by the loaded model") + + resolved_handlers: list[_MediaHandler] = [] for index, item in enumerate(request.media): if item.modality == generation_pb2.MODALITY_UNSPECIFIED: raise ValueError(f"media[{index}].modality must be set") - if item.modality != generation_pb2.MODALITY_IMAGE: - modality = generation_pb2.Modality.Name(item.modality).removeprefix("MODALITY_").lower() - supported_text = ", ".join(sorted(supported)) or "none" + handler = _MEDIA_HANDLERS.get(item.modality) + if handler is None or handler.name not in capabilities.modalities: + modality = _modality_name(item.modality) + supported_text = ", ".join(sorted(capabilities.modalities)) or "none" raise UnsupportedFeatureError( f"multimodal modality '{modality}' is not supported by the loaded model; " f"supported modalities: {supported_text}" ) - if "image" not in supported: - raise UnsupportedFeatureError("image input is not supported by the loaded model") - if len(request.media) > _MAX_MEDIA_ITEMS: - raise ValueError( - f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " - f"got {len(request.media)}" - ) + resolved_handlers.append(handler) - image_io = _image_io_for_model(llm) + loaders = {handler.name: handler.loader_factory(llm) for handler in set(resolved_handlers)} semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) + decoded_bytes = 0 - async def load_bounded(item: generation_pb2.MediaItem, index: int) -> Any: + async def load_bounded( + item: generation_pb2.MediaItem, index: int, handler: _MediaHandler + ) -> Any: + nonlocal decoded_bytes async with semaphore: - return await _load_image_item(item, image_io, index) + decoded = await loaders[handler.name](item, index) + decoded_bytes += handler.decoded_bytes(decoded) + if decoded_bytes > _MAX_TOTAL_DECODED_MEDIA_BYTES: + raise ValueError( + "decoded media exceeds the per-request limit of " + f"{_MAX_TOTAL_DECODED_MEDIA_BYTES} bytes" + ) + return decoded tasks = [ - asyncio.create_task(load_bounded(item, index)) for index, item in enumerate(request.media) + asyncio.create_task(load_bounded(item, index, handler)) + for index, (item, handler) in enumerate(zip(request.media, resolved_handlers)) ] try: - images = await asyncio.gather(*tasks) + decoded_media = await asyncio.gather(*tasks) except (Exception, asyncio.CancelledError): # `gather` propagates the first failure without cancelling siblings. # Explicitly stop any in-flight URL fetches/decodes before returning @@ -352,18 +436,27 @@ async def load_bounded(item: generation_pb2.MediaItem, index: int) -> Any: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) raise + grouped_media: dict[str, list[Any]] = {} + grouped_uuids: dict[str, list[str | None]] = {} + item_order: list[dict[str, Any]] = [] + for item, handler, decoded in zip(request.media, resolved_handlers, decoded_media): + modality_items = grouped_media.setdefault(handler.name, []) + modality_index = len(modality_items) + modality_items.append(decoded) + grouped_uuids.setdefault(handler.name, []).append(item.uuid or None) + item_order.append( + { + "modality": handler.name, + "index": modality_index, + "placeholder": capabilities.placeholder(handler.name, modality_index + 1), + } + ) + mapped: dict[str, Any] = {"prompt": inputs} if isinstance(inputs, str) else dict(inputs) - mapped["multi_modal_data"] = {"image": images} + mapped["multi_modal_data"] = grouped_media if any(item.uuid for item in request.media): - mapped["multi_modal_uuids"] = {"image": [item.uuid or None for item in request.media]} - mapped["mm_item_order"] = [ - { - "modality": "image", - "index": index, - "placeholder": media_placeholder(llm, "image", index + 1), - } - for index in range(len(request.media)) - ] + mapped["multi_modal_uuids"] = grouped_uuids + mapped["mm_item_order"] = item_order return mapped diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index 6bc18ffeb6a4..f1547b65a112 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -25,13 +25,14 @@ Reservation, ) from .disagg import disaggregated_params_from_request -from .errors import AbortFailedError, UnsupportedFeatureError +from .errors import AbortFailedError, MediaUnavailableError, UnsupportedFeatureError from .formatting import _engine_error_response, _stop_texts from .request_mapping import ( _input_from_request, _metadata_from_context, conversation_affinity_enabled, conversation_params_from_request, + media_capabilities_for_model, sampling_params_from_request, ) from .streaming import ( @@ -62,9 +63,15 @@ def __init__( # context worker uses. The actual transfer is driven by opaque_state. self._kv_transfer_backend = kv_transfer_backend self._guided_backend = llm.args.guided_decoding_backend + self._media_capabilities = media_capabilities_for_model(llm) self.subagent_affinity_auth_key = subagent_affinity_auth_key self._active_requests: dict[str, Any] = {} + @property + def supported_media_modalities(self) -> frozenset[str]: + """Media capabilities cached from the loaded model and mapper.""" + return self._media_capabilities.modalities + def active_request_count(self) -> int: """Requests accepted and not yet completed.""" return len(self._active_requests) @@ -175,7 +182,7 @@ async def _generate( if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: raise UnsupportedFeatureError("prefix cache bypass is not supported") - inputs = await _input_from_request(request, self._llm) + inputs = await _input_from_request(request, self._llm, self._media_capabilities) sampling_params = sampling_params_from_request(request, self._guided_backend) trace_headers, target_dp_rank, affinity_headers = _metadata_from_context(context) conversation_params = conversation_params_from_request(request) @@ -230,6 +237,9 @@ async def _generate( except UnsupportedFeatureError as error: await context.abort(grpc.StatusCode.UNIMPLEMENTED, str(error)) return + except MediaUnavailableError as error: + await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) + return except ValueError as error: await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(error)) return diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index e8ab41d4dccf..484fdf5b6f70 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -45,6 +45,10 @@ from tensorrt_llm.logger import logger +class MediaFetchError(RuntimeError): + """A transient failure while fetching remote media.""" + + def rgba_to_rgb( image: Image.Image, background_color: Union[tuple[int, int, int], list[int]] = (255, 255, 255) ) -> Image.Image: @@ -161,7 +165,7 @@ def _validate_url(url: str) -> None: try: infos = socket.getaddrinfo(hostname, None, proto=socket.IPPROTO_TCP) except socket.gaierror as exc: - raise RuntimeError(f"Could not resolve hostname {hostname!r}") from exc + raise MediaFetchError(f"Could not resolve hostname {hostname!r}") from exc allow_private = _allow_private_urls() for _family, _type, _proto, _canon, sockaddr in infos: @@ -257,18 +261,29 @@ async def _safe_aiohttp_get( async def _fetch(fetch_session: aiohttp.ClientSession) -> bytes: current_url = url for _ in range(_MAX_REDIRECTS + 1): - async with fetch_session.get( - current_url, timeout=timeout, allow_redirects=False - ) as response: - if response.status in _REDIRECT_STATUSES: - redirect_url = response.headers.get("Location", "") - current_url = urljoin(current_url, redirect_url) - await asyncio.to_thread(_validate_url, current_url) - continue - maybe_coro = response.raise_for_status() - if asyncio.iscoroutine(maybe_coro): - await maybe_coro - return await _read_aiohttp_content(response) + try: + async with fetch_session.get( + current_url, timeout=timeout, allow_redirects=False + ) as response: + if response.status in _REDIRECT_STATUSES: + redirect_url = response.headers.get("Location", "") + current_url = urljoin(current_url, redirect_url) + await asyncio.to_thread(_validate_url, current_url) + continue + maybe_coro = response.raise_for_status() + if asyncio.iscoroutine(maybe_coro): + await maybe_coro + return await _read_aiohttp_content(response) + except MediaFetchError: + raise + except aiohttp.ClientResponseError as error: + if error.status in (408, 429) or error.status >= 500: + raise MediaFetchError( + f"remote media server returned HTTP {error.status}" + ) from error + raise + except (aiohttp.ClientError, asyncio.TimeoutError) as error: + raise MediaFetchError(f"remote media fetch failed: {error}") from error raise RuntimeError("Too many redirects") if session is not None: diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index f7de8842a222..79b56c77f9d5 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -19,6 +19,7 @@ model_pb2, server_pb2, ) +from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.llmapi import KVEventsConfig # noqa: E402 @@ -172,12 +173,14 @@ async def test_model_info_reports_the_context_window_and_capabilities(): @pytest.mark.asyncio async def test_model_info_advertises_loaded_image_support(monkeypatch): + inference = _inference() monkeypatch.setattr( - "tensorrt_llm.grpc.openengine.control.supported_media_modalities", - lambda llm: frozenset({"image"}), + inference, + "_media_capabilities", + MediaCapabilities({"image": ""}), ) - info = await _servicer().GetModelInfo( + info = await _servicer(inference=inference).GetModelInfo( model_pb2.GetModelInfoRequest(model=MODEL), FakeServicerContext() ) diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index a4bd537e29f5..43ff2297861b 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -8,6 +8,7 @@ from types import SimpleNamespace from typing import Any +import numpy as np import pytest grpc = pytest.importorskip( @@ -17,9 +18,9 @@ import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 -from tensorrt_llm.grpc.openengine.capabilities import supported_media_modalities # noqa: E402 from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, @@ -160,16 +161,75 @@ def is_valid(model_type, modality): def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): llm, seen_model_types = image_llm - assert supported_media_modalities(llm) == {"image"} + assert request_mapping.media_capabilities_for_model(llm).modalities == {"image"} assert seen_model_types == [_ImageConfig.model_type] assert ( - supported_media_modalities( + request_mapping.media_capabilities_for_model( SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) - ) + ).modalities == set() ) +def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): + llm, _ = image_llm + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: "", + ) + + with pytest.raises(RuntimeError, match="invalid multimodal placeholder"): + request_mapping.media_capabilities_for_model(llm) + + +@pytest.mark.asyncio +async def test_handler_table_drives_capability_and_mapping(image_llm, monkeypatch): + llm, _ = image_llm + + def video_loader_factory(_llm): + async def load(item, index): + return f"video-{index}:{bytes(item.raw_bytes).decode()}" + + return load + + monkeypatch.setitem( + request_mapping._MEDIA_HANDLERS, + generation_pb2.MODALITY_VIDEO, + request_mapping._MediaHandler("video", video_loader_factory, lambda decoded: len(decoded)), + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: model_type == _ImageConfig.model_type + and modality in {"image", "video"}, + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + + capabilities = request_mapping.media_capabilities_for_model(llm) + assert capabilities.modalities == {"image", "video"} + + request = generation_pb2.GenerateRequest( + prompt="video", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"clip", + ) + ], + ) + mapped = await request_mapping._input_from_request(request, llm, capabilities) + + assert mapped["multi_modal_data"] == {"video": ["video-0:clip"]} + assert mapped["mm_item_order"] == [ + {"modality": "video", "index": 0, "placeholder": ""} + ] + + @pytest.mark.asyncio async def test_text_only_input_mapping_is_unchanged(): llm = SimpleNamespace() @@ -338,7 +398,9 @@ def create(cls, defaults, runtime): prompt="too many images", media=[ generation_pb2.MediaItem( - modality=generation_pb2.MODALITY_IMAGE, + # Unsupported on purpose: the count guard must win before the + # per-item modality loop. + modality=generation_pb2.MODALITY_VIDEO, raw_bytes=b"image", ) for _ in range(request_mapping._MAX_MEDIA_ITEMS + 1) @@ -349,6 +411,35 @@ def create(cls, defaults, runtime): await request_mapping._input_from_request(request, llm) +@pytest.mark.asyncio +async def test_total_decoded_image_bytes_are_bounded(image_llm, monkeypatch): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return np.zeros((2, 2, 3), dtype=np.uint8) + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(request_mapping, "_MAX_TOTAL_DECODED_MEDIA_BYTES", 20) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + for _ in range(2) + ], + ) + + with pytest.raises(ValueError, match="decoded media exceeds"): + await request_mapping._input_from_request(request, llm) + + @pytest.mark.asyncio async def test_image_mapping_supports_token_ids(image_llm, monkeypatch): llm, _ = image_llm @@ -460,6 +551,43 @@ async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): assert malformed_context.abort_code == grpc.StatusCode.INVALID_ARGUMENT +@pytest.mark.asyncio +async def test_generate_maps_transient_media_fetch_to_unavailable(image_llm, monkeypatch): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load(self, source): + raise MediaFetchError("connection timed out") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + context = FakeServicerContext() + request = generation_pb2.GenerateRequest( + request_id="remote-image", + model="image-model", + prompt="image", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + url="https://example.com/image.png", + ) + ], + ) + + with pytest.raises(AbortError): + async for _ in servicer.Generate(request, context): + pass + assert context.abort_code == grpc.StatusCode.UNAVAILABLE + + @pytest.mark.asyncio @pytest.mark.parametrize( ("modality", "name"), diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index 9bff8139754c..c2d0421ee5bc 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -26,6 +26,7 @@ from tensorrt_llm.inputs.media_io import ( _MAX_REDIRECTS, + MediaFetchError, _safe_aiohttp_get, _safe_request_get, _validate_url, @@ -129,7 +130,7 @@ def test_rejects_multicast(self, _): side_effect=socket.gaierror("Name or service not known"), ) def test_rejects_unresolvable_hostname(self, _): - with pytest.raises(RuntimeError, match="Could not resolve"): + with pytest.raises(MediaFetchError, match="Could not resolve"): _validate_url("http://this.does.not.exist.invalid/") @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) @@ -346,6 +347,24 @@ class TestSafeAiohttpGet: def _run(self, coro): return asyncio.run(coro) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) + def test_maps_timeout_to_fetch_error(self, _): + class _FakeSession: + def get(self, url, **kwargs): + raise asyncio.TimeoutError("timed out") + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + pass + + with patch( + "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + ): + with pytest.raises(MediaFetchError, match="fetch failed"): + self._run(_safe_aiohttp_get("http://example.com/file.bin")) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_reads_response_chunks_to_eof(self, _): class _FakeResponse: From 1220a29985799427637dd68ee6b09b0f83a00ced Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 16:59:18 -0700 Subject: [PATCH 14/41] style(openengine): tighten media comments Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 28 ++++++------------- .../grpc/openengine/request_mapping.py | 23 ++++----------- tensorrt_llm/grpc/openengine/servicer.py | 2 +- 3 files changed, 14 insertions(+), 39 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 6affcf574747..1a7f76bbd90b 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -1,11 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""What this engine can actually do, defined once for Generate and Control. - -Control advertises from this table and Generate enforces it, so a client can -never discover a capability that Generate then rejects. -""" +"""Shared capability resolution for OpenEngine Control and Generate.""" from dataclasses import dataclass from types import MappingProxyType @@ -29,7 +25,7 @@ @dataclass(frozen=True) class MediaCapabilities: - """Immutable media capability snapshot shared by Control and Generate.""" + """Media support shared by Control and Generate.""" placeholder_templates: Mapping[str, str] @@ -40,7 +36,7 @@ def __post_init__(self) -> None: @property def modalities(self) -> frozenset[str]: - """Modalities the loaded model and request mapper both support.""" + """Modalities supported by both the model and request mapper.""" return frozenset(self.placeholder_templates) def placeholder(self, modality: str, item_number: int) -> str: @@ -73,8 +69,7 @@ def _loaded_model_type(llm: Any) -> str: if config is None: return "" - # Composite Hugging Face configs can delegate the instance attribute to - # their text config. The class attribute retains the top-level registry key. + # The class attribute preserves the top-level type for composite configs. model_type = getattr(type(config), "model_type", None) or getattr(config, "model_type", "") return model_type if isinstance(model_type, str) else "" @@ -82,14 +77,9 @@ def _loaded_model_type(llm: Any) -> str: def resolve_media_capabilities( llm: Any, implemented_modalities: Collection[str] ) -> MediaCapabilities: - """Resolve media support implemented by both the mapper and loaded model. - - A registry entry alone describes what this TensorRT-LLM build knows about; - it does not prove that the loaded engine instantiated a multimodal input - processor. Requiring both prevents a text-only model from advertising - media merely because another image model is registered in the process. The - request mapper supplies ``implemented_modalities`` from its handler table, - so discovery cannot advertise a modality without a request implementation. + """Intersect model support with implemented request handlers. + + A live multimodal processor prevents registry-only false positives. """ input_processor = getattr(llm, "input_processor", None) if not isinstance(input_processor, BaseMultimodalInputProcessor): @@ -103,9 +93,7 @@ def resolve_media_capabilities( for modality in implemented_modalities if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) } - # Validate templates once at startup. A malformed registry entry is a - # server configuration failure, not a client request error discovered only - # after Control has advertised the modality. + # Fail startup before Control advertises a malformed template. for modality, template in placeholders.items(): try: template.format(1) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 28c948ba89d4..0f4ab8603c74 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -27,14 +27,10 @@ # output index and a scan of every generated token, so an unbounded list lets a # single request tax the event loop for the whole stream. _MAX_STOP_CONDITIONS = 64 -# Media decoding can retain both encoded and decoded image buffers. Keep one -# request from scheduling an unbounded number of fetch/decode tasks or holding -# an unbounded amount of image data while preserving practical multi-image use. +# Bound per-request decode work and retained image data. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 -# Encoded bodies are separately bounded by the 64 MiB gRPC message limit and -# the media loader's 200 MiB per-URL response limit. This cap covers expansion -# during decode across the entire request. +# Encoded inputs have separate gRPC and remote-loader caps. _MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -45,8 +41,6 @@ @dataclass(frozen=True) class _MediaHandler: - """One fully implemented OpenEngine media mapping path.""" - name: str loader_factory: Callable[[Any], _MediaLoader] decoded_bytes: Callable[[Any], int] @@ -285,7 +279,6 @@ def _base_input_from_request( def _image_io_for_model(llm: Any) -> ImageMediaIO: - """Create image I/O using the loaded model's preferred decode format.""" input_processor = getattr(llm, "input_processor", None) defaults = None if isinstance(input_processor, BaseMultimodalInputProcessor): @@ -297,7 +290,6 @@ def _image_io_for_model(llm: Any) -> ImageMediaIO: async def _load_image_item( item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int ) -> Any: - """Resolve and decode one OpenEngine image item.""" source = item.WhichOneof("source") if source is None: raise ValueError(f"media[{index}] must set exactly one source") @@ -339,15 +331,13 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: def _decoded_image_bytes(image: Any) -> int: - """Return retained bytes for a decoded image representation.""" if isinstance(image, torch.Tensor): return image.numel() * image.element_size() if isinstance(image, np.ndarray): return image.nbytes if isinstance(image, Image.Image): return image.width * image.height * len(image.getbands()) - # Custom model media types cannot be sized generically. The built-in image - # handler returns one of the representations above. + # ImageMediaIO currently returns only the representations above. return 0 @@ -359,7 +349,6 @@ def _decoded_image_bytes(image: Any) -> int: def media_capabilities_for_model(llm: Any) -> MediaCapabilities: - """Resolve the loaded model against the request handlers in this module.""" return resolve_media_capabilities(llm, {handler.name for handler in _MEDIA_HANDLERS.values()}) @@ -375,7 +364,7 @@ async def _input_from_request( llm: Any, media_capabilities: MediaCapabilities | None = None, ) -> str | dict[str, Any]: - """Build LLM API inputs, resolving media off the event loop.""" + """Build LLM API inputs without blocking the event loop on media.""" inputs = _base_input_from_request(request) if not request.media: return inputs @@ -429,9 +418,7 @@ async def load_bounded( try: decoded_media = await asyncio.gather(*tasks) except (Exception, asyncio.CancelledError): - # `gather` propagates the first failure without cancelling siblings. - # Explicitly stop any in-flight URL fetches/decodes before returning - # the request error so their memory and connections do not linger. + # gather leaves siblings running after the first failure. for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index f1547b65a112..f2dabe368a32 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -69,7 +69,7 @@ def __init__( @property def supported_media_modalities(self) -> frozenset[str]: - """Media capabilities cached from the loaded model and mapper.""" + """Cached media modalities.""" return self._media_capabilities.modalities def active_request_count(self) -> int: From 42d8a1c29234a6265cdfd8797c082f9c2832f343 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 17:25:55 -0700 Subject: [PATCH 15/41] fix(openengine): harden decoded media accounting Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 3 +++ .../grpc/openengine/request_mapping.py | 25 ++++++++++------- .../grpc/openengine/test_multimodal.py | 27 +++++++++++++++++++ 3 files changed, 45 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 1a7f76bbd90b..4f553a4b8bdc 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -34,6 +34,9 @@ def __post_init__(self) -> None: self, "placeholder_templates", MappingProxyType(dict(self.placeholder_templates)) ) + def __hash__(self) -> int: + return hash(frozenset(self.placeholder_templates.items())) + @property def modalities(self) -> frozenset[str]: """Modalities supported by both the model and request mapper.""" diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 0f4ab8603c74..cfd7c17bf2ee 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -10,9 +10,6 @@ from urllib.parse import urlparse import grpc -import numpy as np -import torch -from PIL import Image from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError @@ -30,7 +27,8 @@ # Bound per-request decode work and retained image data. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 -# Encoded inputs have separate gRPC and remote-loader caps. +# Checked after each decode; concurrency bounds peak overshoot. Encoded inputs +# have separate gRPC and remote-loader caps. _MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -331,14 +329,17 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: def _decoded_image_bytes(image: Any) -> int: + import numpy as np + import torch + from PIL import Image + if isinstance(image, torch.Tensor): return image.numel() * image.element_size() if isinstance(image, np.ndarray): return image.nbytes if isinstance(image, Image.Image): return image.width * image.height * len(image.getbands()) - # ImageMediaIO currently returns only the representations above. - return 0 + raise TypeError(f"unsupported decoded image type: {type(image).__name__}") _MEDIA_HANDLERS: dict[int, _MediaHandler] = { @@ -376,7 +377,8 @@ async def _input_from_request( ) capabilities = media_capabilities or media_capabilities_for_model(llm) - if not capabilities.modalities: + modalities = capabilities.modalities + if not modalities: raise UnsupportedFeatureError("media input is not supported by the loaded model") resolved_handlers: list[_MediaHandler] = [] @@ -384,16 +386,19 @@ async def _input_from_request( if item.modality == generation_pb2.MODALITY_UNSPECIFIED: raise ValueError(f"media[{index}].modality must be set") handler = _MEDIA_HANDLERS.get(item.modality) - if handler is None or handler.name not in capabilities.modalities: + if handler is None or handler.name not in modalities: modality = _modality_name(item.modality) - supported_text = ", ".join(sorted(capabilities.modalities)) or "none" + supported_text = ", ".join(sorted(modalities)) or "none" raise UnsupportedFeatureError( f"multimodal modality '{modality}' is not supported by the loaded model; " f"supported modalities: {supported_text}" ) resolved_handlers.append(handler) - loaders = {handler.name: handler.loader_factory(llm) for handler in set(resolved_handlers)} + loaders: dict[str, _MediaLoader] = {} + for handler in resolved_handlers: + if handler.name not in loaders: + loaders[handler.name] = handler.loader_factory(llm) semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) decoded_bytes = 0 diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 43ff2297861b..f9aa7985f4c6 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -18,6 +18,7 @@ import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 @@ -29,6 +30,15 @@ pytestmark = pytest.mark.cpu_only +def _ignore_synthetic_decoded_size(monkeypatch) -> None: + handler = request_mapping._MEDIA_HANDLERS[generation_pb2.MODALITY_IMAGE] + monkeypatch.setitem( + request_mapping._MEDIA_HANDLERS, + generation_pb2.MODALITY_IMAGE, + request_mapping._MediaHandler(handler.name, handler.loader_factory, lambda _decoded: 0), + ) + + class _ImageConfig: # Composite configs can expose their text model type on the instance; the # top-level class attribute is the key used by the multimodal registry. @@ -171,6 +181,19 @@ def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): ) +def test_media_capabilities_are_hashable(): + first = MediaCapabilities({"image": ""}) + second = MediaCapabilities({"image": ""}) + + assert first == second + assert hash(first) == hash(second) + + +def test_unknown_decoded_image_type_is_rejected(): + with pytest.raises(TypeError, match="unsupported decoded image type"): + request_mapping._decoded_image_bytes(object()) + + def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): llm, _ = image_llm monkeypatch.setattr( @@ -264,6 +287,7 @@ async def async_load_bytes(self, data): return f"decoded-bytes:{data.decode()}" monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) request = generation_pb2.GenerateRequest( prompt=" describe", media=[ @@ -326,6 +350,7 @@ async def async_load_bytes(self, data): type(self).active -= 1 monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) request = generation_pb2.GenerateRequest( prompt="images", media=[ @@ -453,6 +478,7 @@ async def async_load_bytes(self, data): return "decoded" monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) request = generation_pb2.GenerateRequest( token_ids=generation_pb2.TokenIds(ids=[10, 20]), media=[ @@ -487,6 +513,7 @@ async def async_load_bytes(self, data): return "decoded" monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + _ignore_synthetic_decoded_size(monkeypatch) servicer = OpenEngineInferenceServicer(llm, model="image-model") request = generation_pb2.GenerateRequest( request_id="image-request", From 14cb645aecbc024ab8c30fee84a4ff4efadb440d Mon Sep 17 00:00:00 2001 From: krishung5 Date: Fri, 16 Oct 2026 11:09:34 -0700 Subject: [PATCH 16/41] fix(media): prevent DNS rebinding during async fetches Signed-off-by: krishung5 --- tensorrt_llm/inputs/media_io.py | 68 ++++++++++----- .../inputs/test_async_media_loading.py | 1 + tests/unittest/inputs/test_url_validation.py | 85 +++++++++++++++++-- 3 files changed, 124 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 484fdf5b6f70..dc8cb2de387e 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -97,7 +97,8 @@ async def _get_aiohttp_session() -> aiohttp.ClientSession: """Return the shared aiohttp.ClientSession, creating it on first call.""" global _global_aiohttp_session if _global_aiohttp_session is None or _global_aiohttp_session.closed: - _global_aiohttp_session = aiohttp.ClientSession() + connector = aiohttp.TCPConnector(resolver=_ValidatingResolver()) + _global_aiohttp_session = aiohttp.ClientSession(connector=connector) return _global_aiohttp_session @@ -138,6 +139,42 @@ def _allow_private_urls() -> bool: return True +def _validate_resolved_address(address: str) -> None: + try: + ip = ipaddress.ip_address(address) + except ValueError as exc: + raise MediaFetchError(f"DNS resolver returned an invalid address: {address!r}") from exc + if ip.is_multicast: + raise RuntimeError(f"URL resolves to a multicast address ({ip})") + if not ip.is_global and not _allow_private_urls(): + raise RuntimeError( + f"URL resolves to a non-public address ({ip}); set " + f"{_ALLOW_PRIVATE_URLS_ENV}=1 to allow private media URLs " + f"in trusted network environments" + ) + + +class _ValidatingResolver(aiohttp.abc.AbstractResolver): + """Reject unsafe addresses in the DNS result used by the TCP connector.""" + + def __init__(self) -> None: + self._resolver = aiohttp.DefaultResolver() + + async def resolve( + self, host: str, port: int = 0, family: int = socket.AF_INET + ) -> List[Dict[str, Any]]: + try: + addresses = await self._resolver.resolve(host, port, family) + except OSError as exc: + raise MediaFetchError(f"Could not resolve hostname {host!r}") from exc + for address in addresses: + _validate_resolved_address(address["host"]) + return addresses + + async def close(self) -> None: + await self._resolver.close() + + def _validate_url(url: str) -> None: """Validate that *url* points to a public, non-internal HTTP(S) resource. @@ -147,11 +184,9 @@ def _validate_url(url: str) -> None: rejection (scheme, hostname, and multicast checks still apply); see ``_ALLOW_PRIVATE_URLS_ENV`` above for the security tradeoff. - Note: validation is performed at DNS-resolution time. A DNS-rebinding - attack (TTL=0, resolves to a public IP during validation then a private IP - during the actual TCP connect) could bypass this check. For strict - isolation, supplement with network-level egress filtering that blocks - non-global address ranges at the host firewall. + Async fetches repeat the address check in the resolver used by their TCP + connector, so the address admitted here cannot be replaced with an unsafe + one before connection. """ parsed = urlparse(url) if parsed.scheme not in ("http", "https"): @@ -167,17 +202,8 @@ def _validate_url(url: str) -> None: except socket.gaierror as exc: raise MediaFetchError(f"Could not resolve hostname {hostname!r}") from exc - allow_private = _allow_private_urls() for _family, _type, _proto, _canon, sockaddr in infos: - ip = ipaddress.ip_address(sockaddr[0]) - if ip.is_multicast: - raise RuntimeError(f"URL resolves to a multicast address ({ip})") - if not ip.is_global and not allow_private: - raise RuntimeError( - f"URL resolves to a non-public address ({ip}); set " - f"{_ALLOW_PRIVATE_URLS_ENV}=1 to allow private media URLs " - f"in trusted network environments" - ) + _validate_resolved_address(sockaddr[0]) def _buffer_requests_response(resp: "requests.Response") -> "requests.Response": @@ -255,6 +281,7 @@ async def _safe_aiohttp_get( url: str, timeout_sec: int = 30, session: Optional[aiohttp.ClientSession] = None ) -> bytes: """Aiohttp GET wrapper that validates every redirect hop before following.""" + del session # Kept for API compatibility; fetches must use the validating connector. await asyncio.to_thread(_validate_url, url) timeout = aiohttp.ClientTimeout(total=timeout_sec) @@ -286,11 +313,7 @@ async def _fetch(fetch_session: aiohttp.ClientSession) -> bytes: raise MediaFetchError(f"remote media fetch failed: {error}") from error raise RuntimeError("Too many redirects") - if session is not None: - return await _fetch(session) - - async with aiohttp.ClientSession(timeout=timeout) as owned_session: - return await _fetch(owned_session) + return await _fetch(await _get_aiohttp_session()) def _load_and_convert_image(image): @@ -995,8 +1018,7 @@ async def async_load(self, url: str) -> _MediaT: """ parsed = urlparse(url) if parsed.scheme in ("http", "https"): - session = await _get_aiohttp_session() - data = await _safe_aiohttp_get(url, session=session) + data = await _safe_aiohttp_get(url) return await self._run_in_executor(self.load_bytes, data) elif parsed.scheme == "data": data_spec, b64_data = parsed.path.split(",", 1) diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index 997b140f79f7..90a1b8731fb4 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -170,6 +170,7 @@ async def test_same_session_returned_on_repeated_calls(self): assert first_session is second_session, ( "Expected the same ClientSession object to be returned" ) + assert isinstance(first_session.connector._resolver, media_io_module._ValidatingResolver) @pytest.mark.asyncio async def test_new_session_created_after_close(self): diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index c2d0421ee5bc..374e3128a460 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -20,8 +20,9 @@ import asyncio import socket -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch +import aiohttp import pytest from tensorrt_llm.inputs.media_io import ( @@ -30,6 +31,7 @@ _safe_aiohttp_get, _safe_request_get, _validate_url, + _ValidatingResolver, ) pytestmark = pytest.mark.cpu_only @@ -66,6 +68,35 @@ async def read(self, _): return self._chunks.pop(0) +class _AiohttpErrorResponse: + headers = {} + + def __init__(self, status): + self.status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + pass + + def raise_for_status(self): + raise aiohttp.ClientResponseError( + request_info=MagicMock(), + history=(), + status=self.status, + message="test status", + ) + + +class _AiohttpSession: + def __init__(self, response): + self._response = response + + def get(self, url, **kwargs): + return self._response + + # --------------------------------------------------------------------------- # _validate_url # --------------------------------------------------------------------------- @@ -360,11 +391,31 @@ async def __aexit__(self, *_): pass with patch( - "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_FakeSession()), ): with pytest.raises(MediaFetchError, match="fetch failed"): self._run(_safe_aiohttp_get("http://example.com/file.bin")) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) + def test_preserves_non_retryable_http_error(self, _): + with patch( + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_AiohttpSession(_AiohttpErrorResponse(404))), + ): + with pytest.raises(aiohttp.ClientResponseError) as error: + self._run(_safe_aiohttp_get("http://example.com/missing.jpg")) + assert error.value.status == 404 + + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) + def test_maps_retryable_http_error_to_fetch_error(self, _): + with patch( + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_AiohttpSession(_AiohttpErrorResponse(503))), + ): + with pytest.raises(MediaFetchError, match="HTTP 503"): + self._run(_safe_aiohttp_get("http://example.com/unavailable.jpg")) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_reads_response_chunks_to_eof(self, _): class _FakeResponse: @@ -392,9 +443,15 @@ async def __aexit__(self, *_): pass with patch( - "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_FakeSession()), ): - assert self._run(_safe_aiohttp_get("http://example.com/file.bin")) == b"abcd" + unsafe_session = MagicMock() + unsafe_session.get.side_effect = AssertionError("caller session must not be used") + assert ( + self._run(_safe_aiohttp_get("http://example.com/file.bin", session=unsafe_session)) + == b"abcd" + ) @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_raises_after_too_many_redirects(self, _): @@ -422,7 +479,8 @@ async def __aexit__(self, *_): pass with patch( - "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_FakeSession()), ): with pytest.raises(RuntimeError, match="Too many redirects"): self._run(_safe_aiohttp_get("http://example.com/image.jpg")) @@ -455,7 +513,8 @@ async def __aexit__(self, *_): with ( patch( - "tensorrt_llm.inputs.media_io.aiohttp.ClientSession", return_value=_FakeSession() + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_FakeSession()), ), patch("tensorrt_llm.inputs.media_io._MAX_RESPONSE_BYTES", 3), ): @@ -469,8 +528,20 @@ def test_rejects_private_url_before_request(self): "tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=_dns("169.254.169.254"), ), - patch("tensorrt_llm.inputs.media_io.aiohttp.ClientSession") as mock_session, + patch("tensorrt_llm.inputs.media_io._get_aiohttp_session") as mock_session, ): with pytest.raises(RuntimeError, match="non-public"): self._run(_safe_aiohttp_get("http://169.254.169.254/latest/")) mock_session.assert_not_called() + + def test_rejects_private_address_from_connection_resolver(self): + async def resolve(): + resolver = _ValidatingResolver() + resolver._resolver.resolve = AsyncMock(return_value=[{"host": "169.254.169.254"}]) + try: + await resolver.resolve("example.com", 80) + finally: + await resolver.close() + + with pytest.raises(RuntimeError, match="non-public"): + self._run(resolve()) From 9db976c9c2fb7f6a03fc5db96bf58c71d3b0d689 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 20:37:00 -0700 Subject: [PATCH 17/41] feat(openengine): support video inputs Signed-off-by: krishung5 --- .../grpc/openengine/request_mapping.py | 87 ++++++-- .../grpc/openengine/test_multimodal.py | 190 ++++++++++++++++-- .../openengine/test_multimodal_conformance.py | 61 +++++- 3 files changed, 303 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index cfd7c17bf2ee..14e9da6f84c8 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -4,6 +4,7 @@ """Translate an OpenEngine GenerateRequest into TensorRT-LLM request inputs.""" import asyncio +import importlib from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from typing import Any @@ -12,7 +13,8 @@ import grpc from tensorrt_llm.conversation_params import ConversationParams -from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError +from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError, VideoMediaIO +from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams @@ -24,7 +26,7 @@ # output index and a scan of every generated token, so an unbounded list lets a # single request tax the event loop for the whole stream. _MAX_STOP_CONDITIONS = 64 -# Bound per-request decode work and retained image data. +# Bound per-request decode work and retained decoded media. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 # Checked after each decode; concurrency bounds peak overshoot. Encoded inputs @@ -37,11 +39,16 @@ _MediaLoader = Callable[[generation_pb2.MediaItem, int], Awaitable[Any]] +def _always_available() -> bool: + return True + + @dataclass(frozen=True) class _MediaHandler: name: str loader_factory: Callable[[Any], _MediaLoader] decoded_bytes: Callable[[Any], int] + runtime_available: Callable[[], bool] = _always_available def conversation_affinity_enabled(llm: Any) -> bool: @@ -276,17 +283,17 @@ def _base_input_from_request( raise ValueError("exactly one of prompt or token_ids must be set") -def _image_io_for_model(llm: Any) -> ImageMediaIO: +def _media_io_for_model(llm: Any, modality: str, media_io_type: type[Any]) -> Any: input_processor = getattr(llm, "input_processor", None) defaults = None if isinstance(input_processor, BaseMultimodalInputProcessor): preferences = input_processor.get_preferred_media_io_kwargs() or {} - defaults = preferences.get("image") - return ImageMediaIO.create(defaults, None) + defaults = preferences.get(modality) + return media_io_type.create(defaults, None) -async def _load_image_item( - item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int +async def _load_media_item( + item: generation_pb2.MediaItem, media_io: Any, index: int, modality: str ) -> Any: source = item.WhichOneof("source") if source is None: @@ -300,17 +307,19 @@ async def _load_image_item( raise ValueError(f"media[{index}].data_uri must use the data scheme") value = item.data_uri elif source == "raw_bytes": - if item.mime_type and not item.mime_type.lower().startswith("image/"): - raise ValueError(f"media[{index}].mime_type must be an image type for raw image bytes") + if item.mime_type and not item.mime_type.lower().startswith(f"{modality}/"): + raise ValueError( + f"media[{index}].mime_type must start with {modality}/ for raw {modality} bytes" + ) try: - return await image_io.async_load_bytes(bytes(item.raw_bytes)) + return await media_io.async_load_bytes(bytes(item.raw_bytes)) except Exception as error: raise ValueError(f"media[{index}].raw_bytes could not be decoded: {error}") from error else: raise ValueError(f"media[{index}] has an unknown source '{source}'") try: - return await image_io.async_load(value) + return await media_io.async_load(value) except MediaFetchError as error: raise MediaUnavailableError( f"media[{index}].{source} is temporarily unavailable: {error}" @@ -320,10 +329,19 @@ async def _load_image_item( def _image_loader_factory(llm: Any) -> _MediaLoader: - image_io = _image_io_for_model(llm) + image_io = _media_io_for_model(llm, "image", ImageMediaIO) async def load(item: generation_pb2.MediaItem, index: int) -> Any: - return await _load_image_item(item, image_io, index) + return await _load_media_item(item, image_io, index, "image") + + return load + + +def _video_loader_factory(llm: Any) -> _MediaLoader: + video_io = _media_io_for_model(llm, "video", VideoMediaIO) + + async def load(item: generation_pb2.MediaItem, index: int) -> Any: + return await _load_media_item(item, video_io, index, "video") return load @@ -342,15 +360,56 @@ def _decoded_image_bytes(image: Any) -> int: raise TypeError(f"unsupported decoded image type: {type(image).__name__}") +def _decoded_audio_data_bytes(audio: AudioData) -> int: + import numpy as np + import torch + + if isinstance(audio.samples, torch.Tensor): + return audio.samples.numel() * audio.samples.element_size() + if isinstance(audio.samples, np.ndarray): + return audio.samples.nbytes + raise TypeError(f"unsupported decoded audio type: {type(audio.samples).__name__}") + + +def _decoded_video_bytes(video: Any) -> int: + import numpy as np + + if not isinstance(video, VideoData): + raise TypeError(f"unsupported decoded video type: {type(video).__name__}") + if isinstance(video.frames, np.ndarray): + total = video.frames.nbytes + elif isinstance(video.frames, list): + total = sum(_decoded_image_bytes(frame) for frame in video.frames) + else: + raise TypeError(f"unsupported decoded video frames type: {type(video.frames).__name__}") + if video.audio is not None: + total += _decoded_audio_data_bytes(video.audio) + return total + + +def _video_runtime_available() -> bool: + try: + importlib.import_module("cv2") + except ImportError: + return False + return True + + _MEDIA_HANDLERS: dict[int, _MediaHandler] = { generation_pb2.MODALITY_IMAGE: _MediaHandler( "image", _image_loader_factory, _decoded_image_bytes ), + generation_pb2.MODALITY_VIDEO: _MediaHandler( + "video", _video_loader_factory, _decoded_video_bytes, _video_runtime_available + ), } def media_capabilities_for_model(llm: Any) -> MediaCapabilities: - return resolve_media_capabilities(llm, {handler.name for handler in _MEDIA_HANDLERS.values()}) + return resolve_media_capabilities( + llm, + {handler.name for handler in _MEDIA_HANDLERS.values() if handler.runtime_available()}, + ) def _modality_name(modality: int) -> str: diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index f9aa7985f4c6..3772ebbe87bc 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for OpenEngine image input mapping and discovery.""" +"""Unit tests for OpenEngine multimodal input mapping and discovery.""" import asyncio from collections.abc import AsyncIterator, Sequence @@ -22,6 +22,7 @@ from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 +from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData # noqa: E402 from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, @@ -35,7 +36,26 @@ def _ignore_synthetic_decoded_size(monkeypatch) -> None: monkeypatch.setitem( request_mapping._MEDIA_HANDLERS, generation_pb2.MODALITY_IMAGE, - request_mapping._MediaHandler(handler.name, handler.loader_factory, lambda _decoded: 0), + request_mapping._MediaHandler( + handler.name, + handler.loader_factory, + lambda _decoded: 0, + handler.runtime_available, + ), + ) + + +def _set_video_runtime_available(monkeypatch, available: bool) -> None: + handler = request_mapping._MEDIA_HANDLERS[generation_pb2.MODALITY_VIDEO] + monkeypatch.setitem( + request_mapping._MEDIA_HANDLERS, + generation_pb2.MODALITY_VIDEO, + request_mapping._MediaHandler( + handler.name, + handler.loader_factory, + handler.decoded_bytes, + lambda: available, + ), ) @@ -74,7 +94,10 @@ def call_with_text_prompt(self, inputs, sampling_params): raise NotImplementedError def get_preferred_media_io_kwargs(self): - return {"image": {"format": "np"}} + return { + "image": {"format": "np"}, + "video": {"format": "np", "num_frames": 4}, + } class _Tokenizer: @@ -145,6 +168,7 @@ def generate_async(self, **kwargs): @pytest.fixture def image_llm(monkeypatch): + _set_video_runtime_available(monkeypatch, True) config = _ImageConfig() input_processor = _ImageInputProcessor(config) seen_model_types = [] @@ -168,11 +192,34 @@ def is_valid(model_type, modality): ) +@pytest.fixture +def image_video_llm(monkeypatch): + _set_video_runtime_available(monkeypatch, True) + config = _ImageConfig() + input_processor = _ImageInputProcessor(config) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: ( + model_type == _ImageConfig.model_type and modality in {"image", "video"} + ), + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + return SimpleNamespace( + _hf_model_config=config, + input_processor=input_processor, + ) + + def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): llm, seen_model_types = image_llm assert request_mapping.media_capabilities_for_model(llm).modalities == {"image"} - assert seen_model_types == [_ImageConfig.model_type] + assert seen_model_types == [_ImageConfig.model_type] * 2 assert ( request_mapping.media_capabilities_for_model( SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) @@ -189,11 +236,40 @@ def test_media_capabilities_are_hashable(): assert hash(first) == hash(second) +def test_media_capability_requires_the_runtime_decoder(image_video_llm, monkeypatch): + _set_video_runtime_available(monkeypatch, False) + + assert request_mapping.media_capabilities_for_model(image_video_llm).modalities == {"image"} + + +def test_video_runtime_availability_requires_importable_opencv(monkeypatch): + monkeypatch.setattr(request_mapping.importlib, "import_module", lambda _name: object()) + assert request_mapping._video_runtime_available() is True + + def missing(_name): + raise ImportError + + monkeypatch.setattr(request_mapping.importlib, "import_module", missing) + assert request_mapping._video_runtime_available() is False + + def test_unknown_decoded_image_type_is_rejected(): with pytest.raises(TypeError, match="unsupported decoded image type"): request_mapping._decoded_image_bytes(object()) +def test_decoded_video_bytes_include_frames_and_audio(): + video = VideoData( + frames=np.zeros((2, 3, 4, 3), dtype=np.uint8), + metadata={}, + audio=AudioData(samples=np.zeros(8, dtype=np.float32), sample_rate=16000), + ) + + assert request_mapping._decoded_video_bytes(video) == 2 * 3 * 4 * 3 + 8 * 4 + with pytest.raises(TypeError, match="unsupported decoded video type"): + request_mapping._decoded_video_bytes(object()) + + def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): llm, _ = image_llm monkeypatch.setattr( @@ -224,8 +300,9 @@ async def load(item, index): monkeypatch.setattr( MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", - lambda model_type, modality: model_type == _ImageConfig.model_type - and modality in {"image", "video"}, + lambda model_type, modality: ( + model_type == _ImageConfig.model_type and modality in {"image", "video"} + ), ) monkeypatch.setattr( MULTIMODAL_PLACEHOLDER_REGISTRY, @@ -328,6 +405,75 @@ async def async_load_bytes(self, data): ] +@pytest.mark.asyncio +async def test_image_and_video_sources_preserve_global_order(image_video_llm, monkeypatch): + created_with = [] + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append(("image", defaults, runtime)) + return cls() + + async def async_load_bytes(self, data): + return np.full((1, 1, 3), data[0], dtype=np.uint8) + + class _VideoIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append(("video", defaults, runtime)) + return cls() + + async def async_load(self, source): + return VideoData( + frames=np.zeros((1, 1, 1, 3), dtype=np.uint8), + metadata={"source": source}, + ) + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(request_mapping, "VideoMediaIO", _VideoIo) + request = generation_pb2.GenerateRequest( + prompt="mixed media", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + url="https://example.com/first.mp4", + uuid="video-id", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"i", + mime_type="image/png", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + data_uri="data:video/mp4;base64,dmlkZW8=", + ), + ], + ) + + mapped = await request_mapping._input_from_request(request, image_video_llm) + + assert created_with == [ + ("video", {"format": "np", "num_frames": 4}, None), + ("image", {"format": "np"}, None), + ] + assert [video.metadata["source"] for video in mapped["multi_modal_data"]["video"]] == [ + "https://example.com/first.mp4", + "data:video/mp4;base64,dmlkZW8=", + ] + assert mapped["multi_modal_data"]["image"][0].shape == (1, 1, 3) + assert mapped["multi_modal_uuids"] == { + "video": ["video-id", None], + "image": [None], + } + assert mapped["mm_item_order"] == [ + {"modality": "video", "index": 0, "placeholder": ""}, + {"modality": "image", "index": 0, "placeholder": ""}, + {"modality": "video", "index": 1, "placeholder": ""}, + ] + + @pytest.mark.asyncio async def test_image_loading_has_a_bounded_concurrency(image_llm, monkeypatch): llm, _ = image_llm @@ -616,21 +762,14 @@ async def async_load(self, source): @pytest.mark.asyncio -@pytest.mark.parametrize( - ("modality", "name"), - [ - (generation_pb2.MODALITY_VIDEO, "video"), - (generation_pb2.MODALITY_AUDIO, "audio"), - ], -) -async def test_video_and_audio_are_explicitly_unsupported(image_llm, modality, name): +async def test_audio_is_explicitly_unsupported(image_llm): llm, _ = image_llm request = generation_pb2.GenerateRequest( prompt="media", - media=[generation_pb2.MediaItem(modality=modality, raw_bytes=b"data")], + media=[generation_pb2.MediaItem(modality=generation_pb2.MODALITY_AUDIO, raw_bytes=b"data")], ) - with pytest.raises(UnsupportedFeatureError, match=name): + with pytest.raises(UnsupportedFeatureError, match="audio"): await request_mapping._input_from_request(request, llm) @@ -674,7 +813,7 @@ async def test_image_is_rejected_for_a_text_only_loaded_model(): raw_bytes=b"image", mime_type="video/mp4", ), - "must be an image type", + "must start with image/", ), ], ) @@ -684,3 +823,20 @@ async def test_malformed_image_items_are_rejected(image_llm, item, message): with pytest.raises(ValueError, match=message): await request_mapping._input_from_request(request, llm) + + +@pytest.mark.asyncio +async def test_raw_video_mime_type_is_validated(image_video_llm): + request = generation_pb2.GenerateRequest( + prompt="video", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"video", + mime_type="image/png", + ) + ], + ) + + with pytest.raises(ValueError, match="must start with video/"): + await request_mapping._input_from_request(request, image_video_llm) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index e97c9a708aac..19c5d089f3e9 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -1,10 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Image conformance tests against a real TensorRT-LLM VLM engine.""" +"""Multimodal conformance tests against a real TensorRT-LLM VLM engine.""" import asyncio import base64 +from pathlib import Path import pytest @@ -44,18 +45,27 @@ def vlm(): engine.shutdown() -def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm): +def _servicers(vlm): inference = OpenEngineInferenceServicer(vlm, MODEL_NAME) control = OpenEngineControlServicer(vlm, MODEL_NAME, inference) - info = asyncio.run( + return inference, control + + +def _model_info(control): + return asyncio.run( control.GetModelInfo( model_pb2.GetModelInfoRequest(model=MODEL_NAME), FakeServicerContext(), ) ) + +def test_image_capability_agrees_with_generate(vlm): + inference, control = _servicers(vlm) + info = _model_info(control) + assert info.supports_multimodal is True - assert list(info.extra["trtllm_supported_modalities"]) == ["image"] + assert "image" in info.extra["trtllm_supported_modalities"] prompt = ( "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" @@ -91,3 +101,46 @@ async def collect(): assert finished[0].usage.total_tokens == ( finished[0].usage.prompt_tokens + finished[0].usage.completion_tokens ) + + +def test_video_capability_agrees_with_generate(vlm): + pytest.importorskip("cv2", reason="video decoding requires optional OpenCV") + inference, control = _servicers(vlm) + info = _model_info(control) + assert info.supports_multimodal is True + assert "video" in info.extra["trtllm_supported_modalities"] + + video_path = ( + Path(__file__).parents[2] / "_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.mp4" + ) + + video_prompt = ( + "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" + "Describe this video briefly.<|im_end|>\n<|im_start|>assistant\n" + ) + video_request = generation_pb2.GenerateRequest( + request_id="openengine-video-conformance", + model=MODEL_NAME, + token_ids=generation_pb2.TokenIds(ids=vlm.tokenizer.encode(video_prompt)), + stopping=generation_pb2.StoppingOptions(max_tokens=4), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + data_uri="data:video/mp4;base64," + + base64.b64encode(video_path.read_bytes()).decode(), + uuid="colors", + ) + ], + ) + + async def collect_video(): + return [ + response async for response in inference.Generate(video_request, FakeServicerContext()) + ] + + video_responses = asyncio.run(collect_video()) + assert video_responses + assert not any(response.HasField("error") for response in video_responses) + video_finished = [response for response in video_responses if response.HasField("finished")] + assert len(video_finished) == 1 + assert video_finished[0].usage.prompt_tokens > len(video_request.token_ids.ids) From a3cff30236e375afce75266badf3eaaebc24e71c Mon Sep 17 00:00:00 2001 From: krishung5 Date: Mon, 19 Oct 2026 11:21:48 -0700 Subject: [PATCH 18/41] fix(openengine): address image review feedback Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 14 +----- .../grpc/openengine/request_mapping.py | 10 ++-- tensorrt_llm/inputs/media_io.py | 15 +++--- tensorrt_llm/inputs/model_config.py | 18 +++++++ tensorrt_llm/inputs/utils.py | 10 ++-- tensorrt_llm/serve/chat_utils.py | 13 +---- .../grpc/openengine/test_multimodal.py | 26 ++++++++-- tests/unittest/inputs/test_model_config.py | 21 +++++++++ tests/unittest/inputs/test_url_validation.py | 47 +++++++++++++++---- 9 files changed, 119 insertions(+), 55 deletions(-) create mode 100644 tensorrt_llm/inputs/model_config.py create mode 100644 tests/unittest/inputs/test_model_config.py diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 4f553a4b8bdc..1c65755be0af 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -7,6 +7,7 @@ from types import MappingProxyType from typing import Any, Collection, Mapping +from tensorrt_llm.inputs.model_config import resolve_top_level_model_type from tensorrt_llm.inputs.registry import ( MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, @@ -63,18 +64,7 @@ def supported_guides(backend: str | None) -> frozenset[str]: def _loaded_model_type(llm: Any) -> str: """Return the top-level model type for the loaded input processor.""" - input_processor = getattr(llm, "input_processor", None) - config = ( - input_processor.config - if isinstance(input_processor, BaseMultimodalInputProcessor) - else getattr(llm, "_hf_model_config", None) - ) - if config is None: - return "" - - # The class attribute preserves the top-level type for composite configs. - model_type = getattr(type(config), "model_type", None) or getattr(config, "model_type", "") - return model_type if isinstance(model_type, str) else "" + return resolve_top_level_model_type(llm.input_processor.config) def resolve_media_capabilities( diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index cfd7c17bf2ee..d58bd45e62bf 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -370,17 +370,17 @@ async def _input_from_request( if not request.media: return inputs + capabilities = media_capabilities or media_capabilities_for_model(llm) + modalities = capabilities.modalities + if not modalities: + raise UnsupportedFeatureError("media input is not supported by the loaded model") + if len(request.media) > _MAX_MEDIA_ITEMS: raise ValueError( f"at most {_MAX_MEDIA_ITEMS} media items are supported per request, " f"got {len(request.media)}" ) - capabilities = media_capabilities or media_capabilities_for_model(llm) - modalities = capabilities.modalities - if not modalities: - raise UnsupportedFeatureError("media input is not supported by the loaded model") - resolved_handlers: list[_MediaHandler] = [] for index, item in enumerate(request.media): if item.modality == generation_pb2.MODALITY_UNSPECIFIED: diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index dc8cb2de387e..f35ab1399099 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -154,6 +154,12 @@ def _validate_resolved_address(address: str) -> None: ) +def _dns_resolution_error(host: str, error: OSError) -> RuntimeError: + if isinstance(error, socket.gaierror) and error.errno == socket.EAI_AGAIN: + return MediaFetchError(f"Could not resolve hostname {host!r}") + return RuntimeError(f"Could not resolve hostname {host!r}") + + class _ValidatingResolver(aiohttp.abc.AbstractResolver): """Reject unsafe addresses in the DNS result used by the TCP connector.""" @@ -166,7 +172,7 @@ async def resolve( try: addresses = await self._resolver.resolve(host, port, family) except OSError as exc: - raise MediaFetchError(f"Could not resolve hostname {host!r}") from exc + raise _dns_resolution_error(host, exc) from exc for address in addresses: _validate_resolved_address(address["host"]) return addresses @@ -200,7 +206,7 @@ def _validate_url(url: str) -> None: try: infos = socket.getaddrinfo(hostname, None, proto=socket.IPPROTO_TCP) except socket.gaierror as exc: - raise MediaFetchError(f"Could not resolve hostname {hostname!r}") from exc + raise _dns_resolution_error(hostname, exc) from exc for _family, _type, _proto, _canon, sockaddr in infos: _validate_resolved_address(sockaddr[0]) @@ -277,11 +283,8 @@ async def _read_aiohttp_content(response: aiohttp.ClientResponse) -> bytes: content.write(chunk) -async def _safe_aiohttp_get( - url: str, timeout_sec: int = 30, session: Optional[aiohttp.ClientSession] = None -) -> bytes: +async def _safe_aiohttp_get(url: str, timeout_sec: int = 30) -> bytes: """Aiohttp GET wrapper that validates every redirect hop before following.""" - del session # Kept for API compatibility; fetches must use the validating connector. await asyncio.to_thread(_validate_url, url) timeout = aiohttp.ClientTimeout(total=timeout_sec) diff --git a/tensorrt_llm/inputs/model_config.py b/tensorrt_llm/inputs/model_config.py new file mode 100644 index 000000000000..e82fe8925eb5 --- /dev/null +++ b/tensorrt_llm/inputs/model_config.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Helpers for resolving model configuration metadata.""" + +from typing import Any + + +def resolve_top_level_model_type(model_config: Any) -> str: + """Return the top-level registry model type for a loaded config. + + Composite configs can delegate the instance attribute to ``text_config``. + The class attribute retains the top-level AutoConfig/registry key. + """ + model_type = getattr(type(model_config), "model_type", None) or getattr( + model_config, "model_type", "" + ) + return model_type if isinstance(model_type, str) else "" diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 4d81366928e0..5b79d238edc4 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -25,7 +25,6 @@ detect_content_format) from tensorrt_llm.inputs.data import prompt_inputs from tensorrt_llm.inputs.media_io import (MEDIA_IO_REGISTRY, - _get_aiohttp_session, _load_and_convert_image, _load_video_by_cv2, _normalize_file_uri, @@ -108,8 +107,7 @@ async def async_load_image( parsed_url = urlparse(image) if parsed_url.scheme in ["http", "https"]: - session = await _get_aiohttp_session() - content = await _safe_aiohttp_get(image, session=session) + content = await _safe_aiohttp_get(image) image = await asyncio.to_thread(_load_and_convert_image, BytesIO(content)) elif parsed_url.scheme == "data": @@ -204,8 +202,7 @@ def _load_from_bytes(data: bytes) -> VideoData: extract_audio=extract_audio) if parsed_url.scheme in ["http", "https"]: - session = await _get_aiohttp_session() - video_data = await _safe_aiohttp_get(video, session=session) + video_data = await _safe_aiohttp_get(video) return await asyncio.to_thread(_load_from_bytes, video_data) elif parsed_url.scheme == "data": decoded_video = await asyncio.to_thread(load_base64_video, video) @@ -254,8 +251,7 @@ async def async_load_audio( parsed_url = urlparse(audio) if parsed_url.scheme in ["http", "https"]: - session = await _get_aiohttp_session() - audio_data = await _safe_aiohttp_get(audio, session=session) + audio_data = await _safe_aiohttp_get(audio) audio = BytesIO(audio_data) elif parsed_url.scheme in ("", "file"): audio = _normalize_file_uri( diff --git a/tensorrt_llm/serve/chat_utils.py b/tensorrt_llm/serve/chat_utils.py index 27c270823972..6da15458b93e 100644 --- a/tensorrt_llm/serve/chat_utils.py +++ b/tensorrt_llm/serve/chat_utils.py @@ -19,6 +19,7 @@ MultimodalData, MultimodalDataTracker, load_base64_image_embeds) from tensorrt_llm.inputs.media_io import MEDIA_IO_REGISTRY, BaseMediaIO +from tensorrt_llm.inputs.model_config import resolve_top_level_model_type from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY from tensorrt_llm.inputs.utils import apply_mm_placeholders @@ -429,18 +430,6 @@ def _parse_tool_message_content(message: Dict[str, Any]) -> Dict[str, Any]: return result -def resolve_top_level_model_type(model_config: AutoConfig) -> str: - """Return the top-level HF model_type for a loaded config. - - Newer composite configs (e.g. Qwen2_5_VLConfig) delegate the instance - attribute to `text_config`, returning e.g. "qwen2_5_vl_text" instead of the - top-level "qwen2_5_vl" used as the AutoConfig/registry key. The class-level - attribute is unaffected by this delegation, so prefer it. - """ - return getattr(type(model_config), "model_type", None) or getattr( - model_config, "model_type", "") - - def parse_chat_messages_coroutines( messages: List[ChatCompletionMessageParam], model_config: AutoConfig, diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index f9aa7985f4c6..60977978bc4d 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -224,8 +224,9 @@ async def load(item, index): monkeypatch.setattr( MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", - lambda model_type, modality: model_type == _ImageConfig.model_type - and modality in {"image", "video"}, + lambda model_type, modality: ( + model_type == _ImageConfig.model_type and modality in {"image", "video"} + ), ) monkeypatch.setattr( MULTIMODAL_PLACEHOLDER_REGISTRY, @@ -423,9 +424,7 @@ def create(cls, defaults, runtime): prompt="too many images", media=[ generation_pb2.MediaItem( - # Unsupported on purpose: the count guard must win before the - # per-item modality loop. - modality=generation_pb2.MODALITY_VIDEO, + modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image", ) for _ in range(request_mapping._MAX_MEDIA_ITEMS + 1) @@ -436,6 +435,23 @@ def create(cls, defaults, runtime): await request_mapping._input_from_request(request, llm) +@pytest.mark.asyncio +async def test_model_capability_is_checked_before_media_count_limit(): + request = generation_pb2.GenerateRequest( + prompt="too many images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + for _ in range(request_mapping._MAX_MEDIA_ITEMS + 1) + ], + ) + + with pytest.raises(UnsupportedFeatureError, match="media input is not supported"): + await request_mapping._input_from_request(request, SimpleNamespace()) + + @pytest.mark.asyncio async def test_total_decoded_image_bytes_are_bounded(image_llm, monkeypatch): llm, _ = image_llm diff --git a/tests/unittest/inputs/test_model_config.py b/tests/unittest/inputs/test_model_config.py new file mode 100644 index 000000000000..3ed6c3032474 --- /dev/null +++ b/tests/unittest/inputs/test_model_config.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +from tensorrt_llm.inputs.model_config import resolve_top_level_model_type + + +def test_resolve_top_level_model_type_prefers_composite_class_attribute(): + class CompositeConfig: + model_type = "top_level" + + def __init__(self): + self.model_type = "text_model" + + assert resolve_top_level_model_type(CompositeConfig()) == "top_level" + + +def test_resolve_top_level_model_type_uses_instance_fallback(): + assert resolve_top_level_model_type(SimpleNamespace(model_type="model")) == "model" + assert resolve_top_level_model_type(SimpleNamespace(model_type=None)) == "" diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index 374e3128a460..b549c402cf57 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -158,11 +158,20 @@ def test_rejects_multicast(self, _): @patch( "tensorrt_llm.inputs.media_io.socket.getaddrinfo", - side_effect=socket.gaierror("Name or service not known"), + side_effect=socket.gaierror(socket.EAI_NONAME, "Name or service not known"), ) def test_rejects_unresolvable_hostname(self, _): - with pytest.raises(MediaFetchError, match="Could not resolve"): + with pytest.raises(RuntimeError, match="Could not resolve") as error: _validate_url("http://this.does.not.exist.invalid/") + assert not isinstance(error.value, MediaFetchError) + + @patch( + "tensorrt_llm.inputs.media_io.socket.getaddrinfo", + side_effect=socket.gaierror(socket.EAI_AGAIN, "Temporary failure in name resolution"), + ) + def test_maps_temporary_resolution_failure_to_fetch_error(self, _): + with pytest.raises(MediaFetchError, match="Could not resolve"): + _validate_url("http://temporarily-unavailable.example/") @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_accepts_public_hostname(self, _): @@ -446,12 +455,7 @@ async def __aexit__(self, *_): "tensorrt_llm.inputs.media_io._get_aiohttp_session", new=AsyncMock(return_value=_FakeSession()), ): - unsafe_session = MagicMock() - unsafe_session.get.side_effect = AssertionError("caller session must not be used") - assert ( - self._run(_safe_aiohttp_get("http://example.com/file.bin", session=unsafe_session)) - == b"abcd" - ) + assert self._run(_safe_aiohttp_get("http://example.com/file.bin")) == b"abcd" @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_raises_after_too_many_redirects(self, _): @@ -545,3 +549,30 @@ async def resolve(): with pytest.raises(RuntimeError, match="non-public"): self._run(resolve()) + + @pytest.mark.parametrize( + ("error", "expected_type"), + [ + ( + socket.gaierror(socket.EAI_AGAIN, "Temporary failure in name resolution"), + MediaFetchError, + ), + ( + socket.gaierror(socket.EAI_NONAME, "Name or service not known"), + RuntimeError, + ), + ], + ) + def test_connection_resolver_classifies_dns_failures(self, error, expected_type): + async def resolve(): + resolver = _ValidatingResolver() + resolver._resolver.resolve = AsyncMock(side_effect=error) + try: + await resolver.resolve("example.com", 80) + finally: + await resolver.close() + + with pytest.raises(expected_type, match="Could not resolve") as raised: + self._run(resolve()) + if error.errno != socket.EAI_AGAIN: + assert not isinstance(raised.value, MediaFetchError) From b777c5c41c4d077efb03de2325c986d62f383e1c Mon Sep 17 00:00:00 2001 From: krishung5 Date: Mon, 19 Oct 2026 15:04:35 -0700 Subject: [PATCH 19/41] fix(openengine): fail closed on video decoder load errors Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/request_mapping.py | 2 +- tests/unittest/grpc/openengine/test_multimodal.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index f74742151263..a2b90958f357 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -390,7 +390,7 @@ def _decoded_video_bytes(video: Any) -> int: def _video_runtime_available() -> bool: try: importlib.import_module("cv2") - except ImportError: + except (ImportError, OSError): return False return True diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 2b45ca022255..755274085e8d 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -246,8 +246,11 @@ def test_video_runtime_availability_requires_importable_opencv(monkeypatch): monkeypatch.setattr(request_mapping.importlib, "import_module", lambda _name: object()) assert request_mapping._video_runtime_available() is True + +@pytest.mark.parametrize("error", [ImportError(), OSError()]) +def test_video_runtime_unavailable_when_opencv_fails_to_load(monkeypatch, error): def missing(_name): - raise ImportError + raise error monkeypatch.setattr(request_mapping.importlib, "import_module", missing) assert request_mapping._video_runtime_available() is False From 6844e5a9b026f52fee7e9d18113aeaa05c56267f Mon Sep 17 00:00:00 2001 From: krishung5 Date: Mon, 19 Oct 2026 15:10:07 -0700 Subject: [PATCH 20/41] feat(openengine): support audio inputs Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 14 ++- .../grpc/openengine/request_mapping.py | 20 +++- .../grpc/openengine/test_multimodal.py | 106 ++++++++++++++++- .../openengine/test_multimodal_conformance.py | 107 +++++++++++++++++- 4 files changed, 235 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 1c65755be0af..edd3575a06f8 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -64,7 +64,11 @@ def supported_guides(backend: str | None) -> frozenset[str]: def _loaded_model_type(llm: Any) -> str: """Return the top-level model type for the loaded input processor.""" - return resolve_top_level_model_type(llm.input_processor.config) + input_processor = getattr(llm, "input_processor", None) + config = getattr(input_processor, "config", None) or getattr(llm, "_hf_model_config", None) + if config is None: + return "" + return resolve_top_level_model_type(config) def resolve_media_capabilities( @@ -72,14 +76,16 @@ def resolve_media_capabilities( ) -> MediaCapabilities: """Intersect model support with implemented request handlers. - A live multimodal processor prevents registry-only false positives. + A live registered processor prevents registry-only false positives. """ input_processor = getattr(llm, "input_processor", None) - if not isinstance(input_processor, BaseMultimodalInputProcessor): - return MediaCapabilities({}) model_type = _loaded_model_type(llm) if not model_type: return MediaCapabilities({}) + if not isinstance(input_processor, BaseMultimodalInputProcessor): + registered_model_type = getattr(type(input_processor), "_registered_model_type", None) + if registered_model_type != model_type or not callable(input_processor): + return MediaCapabilities({}) placeholders = { modality: MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index a2b90958f357..ba3abd720043 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -13,7 +13,7 @@ import grpc from tensorrt_llm.conversation_params import ConversationParams -from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError, VideoMediaIO +from tensorrt_llm.inputs.media_io import AudioMediaIO, ImageMediaIO, MediaFetchError, VideoMediaIO from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams @@ -346,6 +346,15 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: return load +def _audio_loader_factory(llm: Any) -> _MediaLoader: + audio_io = _media_io_for_model(llm, "audio", AudioMediaIO) + + async def load(item: generation_pb2.MediaItem, index: int) -> Any: + return await _load_media_item(item, audio_io, index, "audio") + + return load + + def _decoded_image_bytes(image: Any) -> int: import numpy as np import torch @@ -371,6 +380,12 @@ def _decoded_audio_data_bytes(audio: AudioData) -> int: raise TypeError(f"unsupported decoded audio type: {type(audio.samples).__name__}") +def _decoded_audio_bytes(audio: Any) -> int: + if not isinstance(audio, tuple) or len(audio) != 2: + raise TypeError(f"unsupported decoded audio type: {type(audio).__name__}") + return _decoded_audio_data_bytes(AudioData(samples=audio[0], sample_rate=audio[1])) + + def _decoded_video_bytes(video: Any) -> int: import numpy as np @@ -402,6 +417,9 @@ def _video_runtime_available() -> bool: generation_pb2.MODALITY_VIDEO: _MediaHandler( "video", _video_loader_factory, _decoded_video_bytes, _video_runtime_available ), + generation_pb2.MODALITY_AUDIO: _MediaHandler( + "audio", _audio_loader_factory, _decoded_audio_bytes + ), } diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 755274085e8d..bfcd49063d05 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -100,6 +100,18 @@ def get_preferred_media_io_kwargs(self): } +class _AudioInputProcessor: + """Registered processor that does not use BaseMultimodalInputProcessor.""" + + _registered_model_type = _ImageConfig.model_type + + def __init__(self, config: _ImageConfig) -> None: + self.config = config + + def __call__(self, inputs, sampling_params): + raise NotImplementedError + + class _Tokenizer: def convert_ids_to_tokens( self, token_ids: Sequence[int], skip_special_tokens: bool = False @@ -215,11 +227,34 @@ def image_video_llm(monkeypatch): ) +@pytest.fixture +def image_video_audio_llm(monkeypatch): + _set_video_runtime_available(monkeypatch, True) + config = _ImageConfig() + input_processor = _ImageInputProcessor(config) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: ( + model_type == _ImageConfig.model_type and modality in {"image", "video", "audio"} + ), + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + return SimpleNamespace( + _hf_model_config=config, + input_processor=input_processor, + ) + + def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): llm, seen_model_types = image_llm assert request_mapping.media_capabilities_for_model(llm).modalities == {"image"} - assert seen_model_types == [_ImageConfig.model_type] * 2 + assert seen_model_types == [_ImageConfig.model_type] * 3 assert ( request_mapping.media_capabilities_for_model( SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) @@ -228,6 +263,26 @@ def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): ) +def test_media_capability_accepts_a_registered_non_base_processor(monkeypatch): + config = _ImageConfig() + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: model_type == _ImageConfig.model_type and modality == "audio", + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: "", + ) + llm = SimpleNamespace( + _hf_model_config=config, + input_processor=_AudioInputProcessor(config), + ) + + assert request_mapping.media_capabilities_for_model(llm).modalities == {"audio"} + + def test_media_capabilities_are_hashable(): first = MediaCapabilities({"image": ""}) second = MediaCapabilities({"image": ""}) @@ -273,6 +328,12 @@ def test_decoded_video_bytes_include_frames_and_audio(): request_mapping._decoded_video_bytes(object()) +def test_decoded_audio_bytes_count_samples_and_reject_unknown_types(): + assert request_mapping._decoded_audio_bytes((np.zeros(8, dtype=np.float32), 16000)) == 8 * 4 + with pytest.raises(TypeError, match="unsupported decoded audio type"): + request_mapping._decoded_audio_bytes(object()) + + def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): llm, _ = image_llm monkeypatch.setattr( @@ -409,8 +470,9 @@ async def async_load_bytes(self, data): @pytest.mark.asyncio -async def test_image_and_video_sources_preserve_global_order(image_video_llm, monkeypatch): +async def test_media_sources_preserve_global_order(image_video_audio_llm, monkeypatch): created_with = [] + audio_sources = [] class _ImageIo: @classmethod @@ -433,8 +495,19 @@ async def async_load(self, source): metadata={"source": source}, ) + class _AudioIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append(("audio", defaults, runtime)) + return cls() + + async def async_load(self, source): + audio_sources.append(source) + return np.zeros(8, dtype=np.float32), 16000 + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) monkeypatch.setattr(request_mapping, "VideoMediaIO", _VideoIo) + monkeypatch.setattr(request_mapping, "AudioMediaIO", _AudioIo) request = generation_pb2.GenerateRequest( prompt="mixed media", media=[ @@ -448,6 +521,11 @@ async def async_load(self, source): raw_bytes=b"i", mime_type="image/png", ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + url="https://example.com/audio.wav", + uuid="audio-id", + ), generation_pb2.MediaItem( modality=generation_pb2.MODALITY_VIDEO, data_uri="data:video/mp4;base64,dmlkZW8=", @@ -455,12 +533,15 @@ async def async_load(self, source): ], ) - mapped = await request_mapping._input_from_request(request, image_video_llm) + mapped = await request_mapping._input_from_request(request, image_video_audio_llm) assert created_with == [ ("video", {"format": "np", "num_frames": 4}, None), ("image", {"format": "np"}, None), + ("audio", None, None), ] + assert audio_sources == ["https://example.com/audio.wav"] + assert mapped["multi_modal_data"]["audio"][0][1] == 16000 assert [video.metadata["source"] for video in mapped["multi_modal_data"]["video"]] == [ "https://example.com/first.mp4", "data:video/mp4;base64,dmlkZW8=", @@ -469,10 +550,12 @@ async def async_load(self, source): assert mapped["multi_modal_uuids"] == { "video": ["video-id", None], "image": [None], + "audio": ["audio-id"], } assert mapped["mm_item_order"] == [ {"modality": "video", "index": 0, "placeholder": ""}, {"modality": "image", "index": 0, "placeholder": ""}, + {"modality": "audio", "index": 0, "placeholder": ""}, {"modality": "video", "index": 1, "placeholder": ""}, ] @@ -858,3 +941,20 @@ async def test_raw_video_mime_type_is_validated(image_video_llm): with pytest.raises(ValueError, match="must start with video/"): await request_mapping._input_from_request(request, image_video_llm) + + +@pytest.mark.asyncio +async def test_raw_audio_mime_type_is_validated(image_video_audio_llm): + request = generation_pb2.GenerateRequest( + prompt="audio", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + raw_bytes=b"audio", + mime_type="video/mp4", + ) + ], + ) + + with pytest.raises(ValueError, match="must start with audio/"): + await request_mapping._input_from_request(request, image_video_audio_llm) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index 19c5d089f3e9..dbe5813f74f7 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Multimodal conformance tests against a real TensorRT-LLM VLM engine.""" +"""Multimodal conformance tests against real TensorRT-LLM engines.""" import asyncio import base64 @@ -19,15 +19,17 @@ from tensorrt_llm.grpc.openengine.bindings import generation_pb2, model_pb2 # noqa: E402 from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 -from tensorrt_llm.llmapi import KvCacheConfig # noqa: E402 +from tensorrt_llm.llmapi import KvCacheConfig, SchedulerConfig # noqa: E402 MODEL_NAME = "Qwen3/Qwen3-VL-2B-Instruct" +WHISPER_MODEL_NAME = "openai/whisper-tiny" +_EXPECTED_TRANSCRIPT_FRAGMENT = "thoughts affected hester" -# The engine owns background executor threads for the lifetime of this module. +# The engines own background executor threads while each test runs. pytestmark = pytest.mark.threadleak(enabled=False) -@pytest.fixture(scope="module") +@pytest.fixture def vlm(): if not torch.cuda.is_available(): pytest.skip("requires a GPU to run a real VLM engine") @@ -45,6 +47,53 @@ def vlm(): engine.shutdown() +def _whisper_model_path() -> Path: + root = llm_models_root(check=True) + candidates = [ + root / "whisper-models/whisper-tiny", + root / "whisper/whisper-tiny", + root / "whisper-tiny", + ] + for candidate in candidates: + if candidate.exists(): + return candidate + pytest.skip(f"whisper-tiny not found under {root}") + + +def _whisper_audio_path() -> Path: + root = llm_models_root(check=True) + for checkpoint in ("large-v3", "large-v2"): + candidate = root / f"whisper-models/{checkpoint}/1221-135766-0002.wav" + if candidate.exists(): + return candidate + pytest.skip(f"Whisper test audio not found under {root}/whisper-models") + + +@pytest.fixture +def whisper(): + if not torch.cuda.is_available(): + pytest.skip("requires a GPU to run a real Whisper engine") + engine = LLM( + model=str(_whisper_model_path()), + disable_overlap_scheduler=True, + enable_chunked_prefill=False, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + free_gpu_memory_fraction=0.2, + cross_kv_cache_fraction=0.5, + use_kv_cache_manager_v2=True, + ), + max_batch_size=1, + max_input_len=1500, + max_num_tokens=1500, + scheduler_config=SchedulerConfig(use_python_scheduler=True), + ) + try: + yield engine + finally: + engine.shutdown() + + def _servicers(vlm): inference = OpenEngineInferenceServicer(vlm, MODEL_NAME) control = OpenEngineControlServicer(vlm, MODEL_NAME, inference) @@ -144,3 +193,53 @@ async def collect_video(): video_finished = [response for response in video_responses if response.HasField("finished")] assert len(video_finished) == 1 assert video_finished[0].usage.prompt_tokens > len(video_request.token_ids.ids) + + +def test_audio_capability_agrees_with_generate(whisper): + inference = OpenEngineInferenceServicer(whisper, WHISPER_MODEL_NAME) + control = OpenEngineControlServicer(whisper, WHISPER_MODEL_NAME, inference) + info = asyncio.run( + control.GetModelInfo( + model_pb2.GetModelInfoRequest(model=WHISPER_MODEL_NAME), + FakeServicerContext(), + ) + ) + + assert info.supports_multimodal is True + assert list(info.extra["trtllm_supported_modalities"]) == ["audio"] + + audio_uri = ( + "data:audio/wav;base64," + base64.b64encode(_whisper_audio_path().read_bytes()).decode() + ) + request = generation_pb2.GenerateRequest( + request_id="openengine-audio-conformance", + model=WHISPER_MODEL_NAME, + prompt="", + sampling=generation_pb2.SamplingParams(temperature=0.0), + stopping=generation_pb2.StoppingOptions(max_tokens=96), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + data_uri=audio_uri, + uuid="librispeech-1221-135766-0002", + ) + ], + ) + + async def collect_audio(): + return [response async for response in inference.Generate(request, FakeServicerContext())] + + responses = asyncio.run(collect_audio()) + assert responses + assert not any(response.HasField("error") for response in responses) + token_ids = [ + token.token_id + for response in responses + if response.HasField("token") + for token in response.token.tokens + ] + transcript = whisper.tokenizer.decode(token_ids, skip_special_tokens=True).lower() + assert _EXPECTED_TRANSCRIPT_FRAGMENT in transcript + finished = [response for response in responses if response.HasField("finished")] + assert len(finished) == 1 + assert finished[0].usage.completion_tokens == len(token_ids) From 1e3f7f85497657ae1fc7e1edf125f6dc26ccecbf Mon Sep 17 00:00:00 2001 From: krishung5 Date: Mon, 19 Oct 2026 20:53:49 -0700 Subject: [PATCH 21/41] fix: preserve permanent media client errors Signed-off-by: krishung5 --- tensorrt_llm/inputs/media_io.py | 2 ++ tests/unittest/inputs/test_url_validation.py | 21 ++++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index f35ab1399099..cc32b5ac6e7c 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -312,6 +312,8 @@ async def _fetch(fetch_session: aiohttp.ClientSession) -> bytes: f"remote media server returned HTTP {error.status}" ) from error raise + except (aiohttp.InvalidURL, aiohttp.ClientSSLError): + raise except (aiohttp.ClientError, asyncio.TimeoutError) as error: raise MediaFetchError(f"remote media fetch failed: {error}") from error raise RuntimeError("Too many redirects") diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index b549c402cf57..601609343f13 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -406,6 +406,27 @@ async def __aexit__(self, *_): with pytest.raises(MediaFetchError, match="fetch failed"): self._run(_safe_aiohttp_get("http://example.com/file.bin")) + @pytest.mark.parametrize( + "error", + [ + aiohttp.InvalidURL("not a URL"), + aiohttp.ClientSSLError(MagicMock(), OSError("TLS failed")), + ], + ) + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) + def test_preserves_permanent_client_error(self, _, error): + class _FakeSession: + def get(self, url, **kwargs): + raise error + + with patch( + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=_FakeSession()), + ): + with pytest.raises(type(error)) as raised: + self._run(_safe_aiohttp_get("http://example.com/file.bin")) + assert raised.value is error + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_preserves_non_retryable_http_error(self, _): with patch( From 36e555272e66b1d2db9b61da80dea5c129fa539d Mon Sep 17 00:00:00 2001 From: krishung5 Date: Mon, 19 Oct 2026 20:54:04 -0700 Subject: [PATCH 22/41] test: fail video conformance when decoder is missing Signed-off-by: krishung5 --- .../unittest/grpc/openengine/test_multimodal_conformance.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index 19c5d089f3e9..b5970c285b75 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -5,6 +5,7 @@ import asyncio import base64 +import importlib from pathlib import Path import pytest @@ -104,7 +105,10 @@ async def collect(): def test_video_capability_agrees_with_generate(vlm): - pytest.importorskip("cv2", reason="video decoding requires optional OpenCV") + try: + importlib.import_module("cv2") + except (ImportError, OSError) as error: + pytest.fail(f"video conformance requires a working OpenCV install: {error}", pytrace=False) inference, control = _servicers(vlm) info = _model_info(control) assert info.supports_multimodal is True From a9b3038f9ec23792835d55b8d56dc31d9c5d4ebf Mon Sep 17 00:00:00 2001 From: krishung5 Date: Mon, 19 Oct 2026 20:54:48 -0700 Subject: [PATCH 23/41] test: make multimodal conformance failures explicit Signed-off-by: krishung5 --- .../openengine/test_multimodal_conformance.py | 171 +++++++++--------- 1 file changed, 88 insertions(+), 83 deletions(-) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index 4c3870a4e612..18356c837de8 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -30,7 +30,7 @@ pytestmark = pytest.mark.threadleak(enabled=False) -@pytest.fixture +@pytest.fixture(scope="class") def vlm(): if not torch.cuda.is_available(): pytest.skip("requires a GPU to run a real VLM engine") @@ -58,7 +58,7 @@ def _whisper_model_path() -> Path: for candidate in candidates: if candidate.exists(): return candidate - pytest.skip(f"whisper-tiny not found under {root}") + raise FileNotFoundError(f"whisper-tiny not found under {root}") def _whisper_audio_path() -> Path: @@ -67,7 +67,7 @@ def _whisper_audio_path() -> Path: candidate = root / f"whisper-models/{checkpoint}/1221-135766-0002.wav" if candidate.exists(): return candidate - pytest.skip(f"Whisper test audio not found under {root}/whisper-models") + raise FileNotFoundError(f"Whisper test audio not found under {root}/whisper-models") @pytest.fixture @@ -110,93 +110,98 @@ def _model_info(control): ) -def test_image_capability_agrees_with_generate(vlm): - inference, control = _servicers(vlm) - info = _model_info(control) +class TestVlmMediaConformance: + def test_image_capability_agrees_with_generate(self, vlm): + inference, control = _servicers(vlm) + info = _model_info(control) - assert info.supports_multimodal is True - assert "image" in info.extra["trtllm_supported_modalities"] + assert info.supports_multimodal is True + assert "image" in info.extra["trtllm_supported_modalities"] - prompt = ( - "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" - "Describe this image briefly.<|im_end|>\n<|im_start|>assistant\n" - ) - prompt_token_ids = vlm.tokenizer.encode(prompt) - image_path = llm_models_root(check=True) / "multimodals/test_data/seashore.png" - image_bytes = image_path.read_bytes() - data_uri = "data:image/png;base64," + base64.b64encode(image_bytes).decode() - request = generation_pb2.GenerateRequest( - request_id="openengine-image-conformance", - model=MODEL_NAME, - token_ids=generation_pb2.TokenIds(ids=prompt_token_ids), - stopping=generation_pb2.StoppingOptions(max_tokens=4), - media=[ - generation_pb2.MediaItem( - modality=generation_pb2.MODALITY_IMAGE, - data_uri=data_uri, - uuid="seashore", - ) - ], - ) - - async def collect(): - return [response async for response in inference.Generate(request, FakeServicerContext())] - - responses = asyncio.run(collect()) - assert responses - assert not any(response.HasField("error") for response in responses) - finished = [response for response in responses if response.HasField("finished")] - assert len(finished) == 1 - assert finished[0].usage.prompt_tokens > len(prompt_token_ids) - assert finished[0].usage.total_tokens == ( - finished[0].usage.prompt_tokens + finished[0].usage.completion_tokens - ) + prompt = ( + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + "Describe this image briefly.<|im_end|>\n<|im_start|>assistant\n" + ) + prompt_token_ids = vlm.tokenizer.encode(prompt) + image_path = llm_models_root(check=True) / "multimodals/test_data/seashore.png" + image_bytes = image_path.read_bytes() + data_uri = "data:image/png;base64," + base64.b64encode(image_bytes).decode() + request = generation_pb2.GenerateRequest( + request_id="openengine-image-conformance", + model=MODEL_NAME, + token_ids=generation_pb2.TokenIds(ids=prompt_token_ids), + stopping=generation_pb2.StoppingOptions(max_tokens=4), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + data_uri=data_uri, + uuid="seashore", + ) + ], + ) + async def collect(): + return [ + response async for response in inference.Generate(request, FakeServicerContext()) + ] + + responses = asyncio.run(collect()) + assert responses + assert not any(response.HasField("error") for response in responses) + finished = [response for response in responses if response.HasField("finished")] + assert len(finished) == 1 + assert finished[0].usage.prompt_tokens > len(prompt_token_ids) + assert finished[0].usage.total_tokens == ( + finished[0].usage.prompt_tokens + finished[0].usage.completion_tokens + ) -def test_video_capability_agrees_with_generate(vlm): - try: - importlib.import_module("cv2") - except (ImportError, OSError) as error: - pytest.fail(f"video conformance requires a working OpenCV install: {error}", pytrace=False) - inference, control = _servicers(vlm) - info = _model_info(control) - assert info.supports_multimodal is True - assert "video" in info.extra["trtllm_supported_modalities"] + def test_video_capability_agrees_with_generate(self, vlm): + try: + importlib.import_module("cv2") + except (ImportError, OSError) as error: + pytest.fail( + f"video conformance requires a working OpenCV install: {error}", pytrace=False + ) + inference, control = _servicers(vlm) + info = _model_info(control) + assert info.supports_multimodal is True + assert "video" in info.extra["trtllm_supported_modalities"] - video_path = ( - Path(__file__).parents[2] / "_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.mp4" - ) + video_path = ( + Path(__file__).parents[2] / "_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.mp4" + ) - video_prompt = ( - "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" - "Describe this video briefly.<|im_end|>\n<|im_start|>assistant\n" - ) - video_request = generation_pb2.GenerateRequest( - request_id="openengine-video-conformance", - model=MODEL_NAME, - token_ids=generation_pb2.TokenIds(ids=vlm.tokenizer.encode(video_prompt)), - stopping=generation_pb2.StoppingOptions(max_tokens=4), - media=[ - generation_pb2.MediaItem( - modality=generation_pb2.MODALITY_VIDEO, - data_uri="data:video/mp4;base64," - + base64.b64encode(video_path.read_bytes()).decode(), - uuid="colors", - ) - ], - ) + video_prompt = ( + "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" + "Describe this video briefly.<|im_end|>\n<|im_start|>assistant\n" + ) + video_request = generation_pb2.GenerateRequest( + request_id="openengine-video-conformance", + model=MODEL_NAME, + token_ids=generation_pb2.TokenIds(ids=vlm.tokenizer.encode(video_prompt)), + stopping=generation_pb2.StoppingOptions(max_tokens=4), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + data_uri="data:video/mp4;base64," + + base64.b64encode(video_path.read_bytes()).decode(), + uuid="colors", + ) + ], + ) - async def collect_video(): - return [ - response async for response in inference.Generate(video_request, FakeServicerContext()) - ] - - video_responses = asyncio.run(collect_video()) - assert video_responses - assert not any(response.HasField("error") for response in video_responses) - video_finished = [response for response in video_responses if response.HasField("finished")] - assert len(video_finished) == 1 - assert video_finished[0].usage.prompt_tokens > len(video_request.token_ids.ids) + async def collect_video(): + return [ + response + async for response in inference.Generate(video_request, FakeServicerContext()) + ] + + video_responses = asyncio.run(collect_video()) + assert video_responses + assert not any(response.HasField("error") for response in video_responses) + video_finished = [response for response in video_responses if response.HasField("finished")] + assert len(video_finished) == 1 + assert video_finished[0].usage.prompt_tokens > len(video_request.token_ids.ids) def test_audio_capability_agrees_with_generate(whisper): From 220478b677426ceda1113ea2642eadc1d0ef81b4 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Tue, 20 Oct 2026 10:47:28 -0700 Subject: [PATCH 24/41] fix(openengine): reserve request IDs before media loading Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/servicer.py | 62 ++++++++++++------- .../grpc/openengine/test_multimodal.py | 62 +++++++++++++++++++ 2 files changed, 102 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index f2dabe368a32..dd1df20e2a3a 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -44,6 +44,24 @@ ) +class _LocalRequest: + """Abort-aware local request state spanning mapping and generation.""" + + def __init__(self) -> None: + self.handle: Any = None + self.aborted = False + + def bind(self, handle: Any) -> None: + self.handle = handle + if self.aborted: + handle.abort() + + def abort(self) -> None: + if self.handle is not None: + self.handle.abort() + self.aborted = True + + class OpenEngineInferenceServicer(openengine_pb2_grpc.InferenceServicer): """Translate OpenEngine generation streams to TensorRT-LLM requests.""" @@ -127,15 +145,26 @@ async def Generate( request: generation_pb2.GenerateRequest, context: grpc.aio.ServicerContext, ) -> AsyncIterator[generation_pb2.GenerateResponse]: + local_request = None + registration = None reservation = None try: + if not request.request_id: + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, "request_id must be non-empty" + ) + local_request = _LocalRequest() + registration = ActiveRequest(self._active_requests, request.request_id, local_request) + if not registration.register(): + await context.abort( + grpc.StatusCode.ALREADY_EXISTS, + f"request_id '{request.request_id}' is already active", + ) if self._frontend is not None: - if not request.request_id: - await context.abort( - grpc.StatusCode.INVALID_ARGUMENT, "request_id must be non-empty" - ) reservation = await self._frontend.reserve(request.request_id) - async with aclosing(self._generate(request, context, reservation)) as responses: + async with aclosing( + self._generate(request, context, local_request, registration, reservation) + ) as responses: async for response in responses: yield response except DuplicateRequestError as error: @@ -145,6 +174,8 @@ async def Generate( except CoordinationError as error: await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) finally: + if registration is not None: + registration.release() if reservation is not None: try: await asyncio.shield(reservation.release()) @@ -158,24 +189,18 @@ async def _generate( self, request: generation_pb2.GenerateRequest, context: grpc.aio.ServicerContext, + local_request: _LocalRequest, + registration: ActiveRequest, reservation: Reservation | None = None, ) -> AsyncIterator[generation_pb2.GenerateResponse]: """Run a server-streaming OpenEngine generation request.""" request_id = request.request_id try: - if not request_id: - raise ValueError("request_id must be non-empty") if not request.model: raise ValueError("model must be non-empty") # Like the HTTP /v1/completions single-model server, any non-empty # model name is accepted and served by the loaded model (no NOT_FOUND # on a mismatch) so the two transports behave consistently. - if request_id in self._active_requests: - await context.abort( - grpc.StatusCode.ALREADY_EXISTS, - f"request_id '{request_id}' is already active", - ) - return if request.lora_name: raise UnsupportedFeatureError("LoRA selection is not supported") if request.HasField("kv"): @@ -257,7 +282,7 @@ async def _generate( tokenizer=self._llm.tokenizer, ) - if reservation is not None and reservation.aborted: + if local_request.aborted or (reservation is not None and reservation.aborted): await context.abort(grpc.StatusCode.CANCELLED, "Request aborted before submission") try: @@ -291,12 +316,12 @@ async def _generate( await context.abort(grpc.StatusCode.INTERNAL, str(error)) return + local_request.bind(result_handle) if reservation is not None: reservation.bind(result_handle) engine_terminal = False abort_requested = False - registration = ActiveRequest(self._active_requests, request_id, result_handle) def abort_request(reason: str) -> None: nonlocal abort_requested @@ -341,13 +366,6 @@ def on_rpc_done(_: asyncio.Task) -> None: rpc_task.add_done_callback(on_rpc_done) try: - # Register inside the try so the finally below always releases the - # id. There is no await between the duplicate-id check above and - # here, so the asyncio single-thread atomicity that makes that check - # correct still holds, while an exception before streaming can no - # longer leak the id (which would make it permanently un-reusable, - # since re-submits are rejected with ALREADY_EXISTS). - registration.register() watchdog.start() async for result in result_handle: if context.cancelled(): diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 60977978bc4d..95cbb72071f4 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -550,6 +550,67 @@ async def async_load_bytes(self, data): assert responses[-1].usage.total_tokens == 7 +@pytest.mark.asyncio +async def test_duplicate_request_id_is_reserved_while_media_loads(image_llm, monkeypatch): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + load_started = asyncio.Event() + release_load = asyncio.Event() + load_count = 0 + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + nonlocal load_count + load_count += 1 + load_started.set() + await release_load.wait() + return np.zeros((1, 1, 3), dtype=np.uint8) + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + request = generation_pb2.GenerateRequest( + request_id="duplicate-image-request", + model="image-model", + prompt="", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + async def collect(context): + return [response async for response in servicer.Generate(request, context)] + + first_context = FakeServicerContext() + first = asyncio.create_task(collect(first_context)) + await load_started.wait() + duplicate_context = FakeServicerContext() + try: + with pytest.raises(AbortError): + await asyncio.wait_for(collect(duplicate_context), timeout=1) + assert duplicate_context.abort_code == grpc.StatusCode.ALREADY_EXISTS + assert load_count == 1 + assert servicer.active_request_count() == 1 + assert servicer.abort_request_by_id(request.request_id) + finally: + release_load.set() + first_outcome = await asyncio.gather(first, return_exceptions=True) + + assert isinstance(first_outcome[0], AbortError) + assert first_context.abort_code == grpc.StatusCode.CANCELLED + assert llm.generate_kwargs is None + assert servicer.active_request_count() == 0 + + @pytest.mark.asyncio async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): capability_llm, _ = image_llm @@ -592,6 +653,7 @@ async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): async for _ in servicer.Generate(file_request, malformed_context): pass assert malformed_context.abort_code == grpc.StatusCode.INVALID_ARGUMENT + assert servicer.active_request_count() == 0 @pytest.mark.asyncio From db9a6872f748210abe768ee0a18fc78a4357faaa Mon Sep 17 00:00:00 2001 From: krishung5 Date: Tue, 20 Oct 2026 11:01:27 -0700 Subject: [PATCH 25/41] fix(inputs): scope shared media session to event loop Signed-off-by: krishung5 --- tensorrt_llm/inputs/media_io.py | 16 +++++++++-- tests/unittest/inputs/test_url_validation.py | 30 ++++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index cc32b5ac6e7c..797008910229 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -91,14 +91,24 @@ def convert_image_mode(image: Image.Image, to_mode: str) -> Image.Image: # Created lazily on first use inside the async event loop, then reused so that # TCP connections are kept alive and not re-established per request. _global_aiohttp_session: aiohttp.ClientSession | None = None +_global_aiohttp_session_loop: asyncio.AbstractEventLoop | None = None async def _get_aiohttp_session() -> aiohttp.ClientSession: - """Return the shared aiohttp.ClientSession, creating it on first call.""" - global _global_aiohttp_session - if _global_aiohttp_session is None or _global_aiohttp_session.closed: + """Return the aiohttp session owned by the current event loop.""" + global _global_aiohttp_session, _global_aiohttp_session_loop + current_loop = asyncio.get_running_loop() + if ( + _global_aiohttp_session is None + or _global_aiohttp_session.closed + or _global_aiohttp_session_loop is not current_loop + ): + previous_session = _global_aiohttp_session connector = aiohttp.TCPConnector(resolver=_ValidatingResolver()) _global_aiohttp_session = aiohttp.ClientSession(connector=connector) + _global_aiohttp_session_loop = current_loop + if previous_session is not None and not previous_session.closed: + await previous_session.close() return _global_aiohttp_session diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index 601609343f13..5cf2391104a1 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -25,6 +25,7 @@ import aiohttp import pytest +from tensorrt_llm.inputs import media_io as media_io_module from tensorrt_llm.inputs.media_io import ( _MAX_REDIRECTS, MediaFetchError, @@ -387,6 +388,35 @@ class TestSafeAiohttpGet: def _run(self, coro): return asyncio.run(coro) + def test_shared_session_uses_validating_resolver(self): + async def check(): + session = await media_io_module._get_aiohttp_session() + try: + assert isinstance(session.connector._resolver, _ValidatingResolver) + finally: + await session.close() + media_io_module._global_aiohttp_session = None + media_io_module._global_aiohttp_session_loop = None + + self._run(check()) + + def test_shared_session_is_recreated_for_a_new_event_loop(self): + async def get_session(): + return await media_io_module._get_aiohttp_session() + + first = self._run(get_session()) + second = None + try: + second = self._run(get_session()) + assert second is not first + assert first.closed + finally: + session = second or first + if not session.closed: + self._run(session.close()) + media_io_module._global_aiohttp_session = None + media_io_module._global_aiohttp_session_loop = None + @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_maps_timeout_to_fetch_error(self, _): class _FakeSession: From 1a27ad47fdc31aaa8319e4f0b5600dde032bc69c Mon Sep 17 00:00:00 2001 From: krishung5 Date: Tue, 20 Oct 2026 11:16:47 -0700 Subject: [PATCH 26/41] fix(openengine): reject oversized video before allocation Signed-off-by: krishung5 --- .../grpc/openengine/request_mapping.py | 16 ++++++++++--- tensorrt_llm/inputs/media_io.py | 22 ++++++++++++++++++ .../grpc/openengine/test_multimodal.py | 6 ++++- .../inputs/test_async_media_loading.py | 13 +++++++++++ tests/unittest/inputs/test_video_decode.py | 23 +++++++++++++++++++ 5 files changed, 76 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index a2b90958f357..418487352b7d 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -283,13 +283,18 @@ def _base_input_from_request( raise ValueError("exactly one of prompt or token_ids must be set") -def _media_io_for_model(llm: Any, modality: str, media_io_type: type[Any]) -> Any: +def _media_io_for_model( + llm: Any, + modality: str, + media_io_type: type[Any], + runtime_kwargs: dict[str, Any] | None = None, +) -> Any: input_processor = getattr(llm, "input_processor", None) defaults = None if isinstance(input_processor, BaseMultimodalInputProcessor): preferences = input_processor.get_preferred_media_io_kwargs() or {} defaults = preferences.get(modality) - return media_io_type.create(defaults, None) + return media_io_type.create(defaults, runtime_kwargs) async def _load_media_item( @@ -338,7 +343,12 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: def _video_loader_factory(llm: Any) -> _MediaLoader: - video_io = _media_io_for_model(llm, "video", VideoMediaIO) + video_io = _media_io_for_model( + llm, + "video", + VideoMediaIO, + {"max_decoded_frame_bytes": _MAX_TOTAL_DECODED_MEDIA_BYTES}, + ) async def load(item: generation_pb2.MediaItem, index: int) -> Any: return await _load_media_item(item, video_io, index, "video") diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 797008910229..b29c892f5b70 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -739,6 +739,7 @@ def _load_video_by_cv2( extract_audio: bool = False, cv2_backend: Optional[int] = None, raw_bytes_hash: Optional[str] = None, + max_decoded_frame_bytes: Optional[int] = None, ) -> VideoData: """Decode a video and return sampled frames as a list. @@ -756,6 +757,10 @@ def _load_video_by_cv2( contiguous 4D buffer that HF video processors pass through without an extra frame-by-frame copy. `"pil"` - list[PIL.Image], one per sampled frame. + + When `max_decoded_frame_bytes` is set, the selected output representation + is sized from the first decoded frame and rejected before the frame stack + is allocated. """ assert format in ("pt", "np", "pil"), "format must be one of 'pt', 'np', 'pil'" @@ -829,6 +834,17 @@ def _load_video_by_cv2( fh, fw = bgr_scratch.shape[:2] if stacked_rgb is None: H, W = fh, fw + output_dtype = np.float32 if format == "pt" else np.uint8 + output_item_size = np.dtype(output_dtype).itemsize + decoded_frame_bytes = num_frames_to_sample * H * W * 3 * output_item_size + if ( + max_decoded_frame_bytes is not None + and decoded_frame_bytes > max_decoded_frame_bytes + ): + raise ValueError( + f"Decoded video frames require {decoded_frame_bytes} bytes, " + f"exceeding the limit of {max_decoded_frame_bytes} bytes." + ) stacked_rgb = np.empty((num_frames_to_sample, H, W, 3), dtype=np.uint8) if (fh, fw) == (H, W): cv2.cvtColor( @@ -1161,17 +1177,21 @@ def __init__( format: str = "pt", device: str = "cpu", extract_audio: bool = False, + max_decoded_frame_bytes: Optional[int] = None, ) -> None: # `"pt"` is the safe default every video model handles; models tuned # for the uint8-HWC path opt into `"np"` via # InputProcessor.get_preferred_media_io_kwargs(). if format not in _SUPPORTED_VIDEO_FORMATS: raise ValueError(f"format must be one of {_SUPPORTED_VIDEO_FORMATS}, got {format!r}") + if max_decoded_frame_bytes is not None and max_decoded_frame_bytes <= 0: + raise ValueError("max_decoded_frame_bytes must be positive") self._num_frames = num_frames self._fps = fps self._format = format self._device = device self._extract_audio = extract_audio + self._max_decoded_frame_bytes = max_decoded_frame_bytes @classmethod def merge_kwargs( @@ -1210,6 +1230,7 @@ def load_bytes(self, data: bytes) -> VideoData: extract_audio=self._extract_audio, cv2_backend=cv2_backend, raw_bytes_hash=raw_bytes_hash, + max_decoded_frame_bytes=self._max_decoded_frame_bytes, ) with tempfile.NamedTemporaryFile(suffix=".mp4", dir=_VIDEO_TEMPFILE_DIR) as tmp: tmp.write(data) @@ -1222,6 +1243,7 @@ def load_bytes(self, data: bytes) -> VideoData: self._device, extract_audio=self._extract_audio, raw_bytes_hash=raw_bytes_hash, + max_decoded_frame_bytes=self._max_decoded_frame_bytes, ) def load_base64(self, media_type: str, data: str) -> VideoData: diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 6ff2d160eee4..40056d56ef14 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -458,7 +458,11 @@ async def async_load(self, source): mapped = await request_mapping._input_from_request(request, image_video_llm) assert created_with == [ - ("video", {"format": "np", "num_frames": 4}, None), + ( + "video", + {"format": "np", "num_frames": 4}, + {"max_decoded_frame_bytes": request_mapping._MAX_TOTAL_DECODED_MEDIA_BYTES}, + ), ("image", {"format": "np"}, None), ] assert [video.metadata["source"] for video in mapped["multi_modal_data"]["video"]] == [ diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index 90a1b8731fb4..1675b7357f7c 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -322,3 +322,16 @@ async def test_video_accepts_video_mime(self): with patch.object(media_io_module.VideoMediaIO, "load_base64", return_value="decoded"): result = await io.async_load("data:video/mp4;base64,AAAA") assert result == "decoded" + + +def test_video_media_io_forwards_decoded_frame_limit(): + decoded = object() + io = media_io_module.VideoMediaIO(max_decoded_frame_bytes=123) + + with ( + patch.object(media_io_module, "_select_cv2_stream_buffered_backend", return_value=7), + patch.object(media_io_module, "_load_video_by_cv2", return_value=decoded) as decode, + ): + assert io.load_bytes(b"video") is decoded + + assert decode.call_args.kwargs["max_decoded_frame_bytes"] == 123 diff --git a/tests/unittest/inputs/test_video_decode.py b/tests/unittest/inputs/test_video_decode.py index d2d113645c7a..c9818d296a36 100644 --- a/tests/unittest/inputs/test_video_decode.py +++ b/tests/unittest/inputs/test_video_decode.py @@ -5,6 +5,8 @@ from __future__ import annotations +from unittest.mock import patch + import numpy as np import pytest import torch @@ -66,3 +68,24 @@ def test_np_format_hits_hf_video_processor_fast_path(sample_video_path: str) -> assert len(batched) == 1 assert np.shares_memory(video.frames, batched[0]) + + +@pytest.mark.parametrize( + ("format", "dtype"), [("np", np.uint8), ("pil", np.uint8), ("pt", np.float32)] +) +def test_decode_limit_is_checked_before_frame_stack_allocation( + sample_video_path: str, format: str, dtype: type[np.generic] +) -> None: + decoded_bytes = 10 * 64 * 64 * 3 * np.dtype(dtype).itemsize + + with patch("tensorrt_llm.inputs.media_io.np.empty") as allocate: + with pytest.raises(ValueError, match=f"require {decoded_bytes} bytes"): + _load_video_by_cv2( + sample_video_path, + num_frames=10, + fps=-1, + format=format, + max_decoded_frame_bytes=decoded_bytes - 1, + ) + + allocate.assert_not_called() From 7e630cc7d303070afa748572adb26fb7a6fd6ea8 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Tue, 20 Oct 2026 12:10:55 -0700 Subject: [PATCH 27/41] fix(openengine): share decoded video memory budget Signed-off-by: krishung5 --- .../grpc/openengine/request_mapping.py | 130 ++++++++++------ tensorrt_llm/inputs/media_io.py | 144 ++++++++++++++---- .../grpc/openengine/test_multimodal.py | 67 ++++++-- .../inputs/test_async_media_loading.py | 9 +- tests/unittest/inputs/test_video_decode.py | 68 ++++++++- 5 files changed, 328 insertions(+), 90 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 418487352b7d..99a516025e9e 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -5,6 +5,7 @@ import asyncio import importlib +import threading from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from typing import Any @@ -13,8 +14,12 @@ import grpc from tensorrt_llm.conversation_params import ConversationParams -from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError, VideoMediaIO -from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData +from tensorrt_llm.inputs.media_io import ( + ImageMediaIO, + MediaFetchError, + VideoMediaIO, + _decoded_video_bytes, +) from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams @@ -29,8 +34,9 @@ # Bound per-request decode work and retained decoded media. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 -# Checked after each decode; concurrency bounds peak overshoot. Encoded inputs -# have separate gRPC and remote-loader caps. +# Images are charged after decoding. Video decoders reserve their estimated +# peak before allocating frame buffers, then settle to retained decoded bytes. +# Encoded inputs have separate gRPC and remote-loader caps. _MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -39,6 +45,67 @@ _MediaLoader = Callable[[generation_pb2.MediaItem, int], Awaitable[Any]] +class _DecodedMediaReservation: + def __init__(self, budget: "_DecodedMediaBudget", reserved_bytes: int) -> None: + self._budget = budget + self._reserved_bytes = reserved_bytes + self._active = True + + def commit(self, retained_bytes: int) -> None: + if not self._active: + raise RuntimeError("decoded-media reservation is no longer active") + try: + self._budget._commit(self._reserved_bytes, retained_bytes) + finally: + self._active = False + + def release(self) -> None: + if self._active: + self._budget._release(self._reserved_bytes) + self._active = False + + +class _DecodedMediaBudget: + """Thread-safe shared budget for all decoded media in one request.""" + + def __init__(self, limit: int) -> None: + self._limit = limit + self._used = 0 + self._lock = threading.Lock() + + def _limit_error(self) -> ValueError: + return ValueError(f"decoded media exceeds the per-request limit of {self._limit} bytes") + + def charge(self, retained_bytes: int) -> None: + if retained_bytes < 0: + raise ValueError("decoded media size cannot be negative") + with self._lock: + if self._used + retained_bytes > self._limit: + raise self._limit_error() + self._used += retained_bytes + + def reserve(self, peak_bytes: int) -> _DecodedMediaReservation: + if peak_bytes < 0: + raise ValueError("decoded media size cannot be negative") + with self._lock: + if self._used + peak_bytes > self._limit: + raise self._limit_error() + self._used += peak_bytes + return _DecodedMediaReservation(self, peak_bytes) + + def _commit(self, reserved_bytes: int, retained_bytes: int) -> None: + if retained_bytes < 0: + raise ValueError("decoded media size cannot be negative") + with self._lock: + self._used += retained_bytes - reserved_bytes + if self._used > self._limit: + raise self._limit_error() + + def _release(self, reserved_bytes: int) -> None: + with self._lock: + self._used -= reserved_bytes + + def _always_available() -> bool: return True @@ -46,9 +113,10 @@ def _always_available() -> bool: @dataclass(frozen=True) class _MediaHandler: name: str - loader_factory: Callable[[Any], _MediaLoader] + loader_factory: Callable[[Any, _DecodedMediaBudget], _MediaLoader] decoded_bytes: Callable[[Any], int] runtime_available: Callable[[], bool] = _always_available + reserves_decoded_bytes: bool = False def conversation_affinity_enabled(llm: Any) -> bool: @@ -333,7 +401,7 @@ async def _load_media_item( raise ValueError(f"media[{index}].{source} could not be loaded: {error}") from error -def _image_loader_factory(llm: Any) -> _MediaLoader: +def _image_loader_factory(llm: Any, _budget: _DecodedMediaBudget) -> _MediaLoader: image_io = _media_io_for_model(llm, "image", ImageMediaIO) async def load(item: generation_pb2.MediaItem, index: int) -> Any: @@ -342,12 +410,12 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: return load -def _video_loader_factory(llm: Any) -> _MediaLoader: +def _video_loader_factory(llm: Any, budget: _DecodedMediaBudget) -> _MediaLoader: video_io = _media_io_for_model( llm, "video", VideoMediaIO, - {"max_decoded_frame_bytes": _MAX_TOTAL_DECODED_MEDIA_BYTES}, + {"decoded_memory_reserver": budget.reserve}, ) async def load(item: generation_pb2.MediaItem, index: int) -> Any: @@ -370,33 +438,6 @@ def _decoded_image_bytes(image: Any) -> int: raise TypeError(f"unsupported decoded image type: {type(image).__name__}") -def _decoded_audio_data_bytes(audio: AudioData) -> int: - import numpy as np - import torch - - if isinstance(audio.samples, torch.Tensor): - return audio.samples.numel() * audio.samples.element_size() - if isinstance(audio.samples, np.ndarray): - return audio.samples.nbytes - raise TypeError(f"unsupported decoded audio type: {type(audio.samples).__name__}") - - -def _decoded_video_bytes(video: Any) -> int: - import numpy as np - - if not isinstance(video, VideoData): - raise TypeError(f"unsupported decoded video type: {type(video).__name__}") - if isinstance(video.frames, np.ndarray): - total = video.frames.nbytes - elif isinstance(video.frames, list): - total = sum(_decoded_image_bytes(frame) for frame in video.frames) - else: - raise TypeError(f"unsupported decoded video frames type: {type(video.frames).__name__}") - if video.audio is not None: - total += _decoded_audio_data_bytes(video.audio) - return total - - def _video_runtime_available() -> bool: try: importlib.import_module("cv2") @@ -410,7 +451,11 @@ def _video_runtime_available() -> bool: "image", _image_loader_factory, _decoded_image_bytes ), generation_pb2.MODALITY_VIDEO: _MediaHandler( - "video", _video_loader_factory, _decoded_video_bytes, _video_runtime_available + "video", + _video_loader_factory, + _decoded_video_bytes, + _video_runtime_available, + reserves_decoded_bytes=True, ), } @@ -464,25 +509,20 @@ async def _input_from_request( ) resolved_handlers.append(handler) + budget = _DecodedMediaBudget(_MAX_TOTAL_DECODED_MEDIA_BYTES) loaders: dict[str, _MediaLoader] = {} for handler in resolved_handlers: if handler.name not in loaders: - loaders[handler.name] = handler.loader_factory(llm) + loaders[handler.name] = handler.loader_factory(llm, budget) semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) - decoded_bytes = 0 async def load_bounded( item: generation_pb2.MediaItem, index: int, handler: _MediaHandler ) -> Any: - nonlocal decoded_bytes async with semaphore: decoded = await loaders[handler.name](item, index) - decoded_bytes += handler.decoded_bytes(decoded) - if decoded_bytes > _MAX_TOTAL_DECODED_MEDIA_BYTES: - raise ValueError( - "decoded media exceeds the per-request limit of " - f"{_MAX_TOTAL_DECODED_MEDIA_BYTES} bytes" - ) + if not handler.reserves_decoded_bytes: + budget.charge(handler.decoded_bytes(decoded)) return decoded tasks = [ diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index b29c892f5b70..2aeb5c7065d8 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -18,6 +18,7 @@ from types import MappingProxyType from typing import ( Any, + Callable, ClassVar, Dict, Generic, @@ -25,6 +26,7 @@ Literal, Mapping, Optional, + Protocol, Tuple, Type, TypeVar, @@ -49,6 +51,47 @@ class MediaFetchError(RuntimeError): """A transient failure while fetching remote media.""" +class _DecodedMemoryReservation(Protocol): + def commit(self, retained_bytes: int) -> None: ... + + def release(self) -> None: ... + + +_DecodedMemoryReserver = Callable[[int], _DecodedMemoryReservation] + + +def _decoded_video_bytes(video: Any) -> int: + """Return bytes retained by decoded video frames and optional audio.""" + if not isinstance(video, VideoData): + raise TypeError(f"unsupported decoded video type: {type(video).__name__}") + + if isinstance(video.frames, np.ndarray): + total = video.frames.nbytes + elif isinstance(video.frames, list): + total = 0 + for frame in video.frames: + if isinstance(frame, torch.Tensor): + total += frame.numel() * frame.element_size() + elif isinstance(frame, np.ndarray): + total += frame.nbytes + elif isinstance(frame, Image.Image): + total += frame.width * frame.height * len(frame.getbands()) + else: + raise TypeError(f"unsupported decoded video frame type: {type(frame).__name__}") + else: + raise TypeError(f"unsupported decoded video frames type: {type(video.frames).__name__}") + + if video.audio is not None: + samples = video.audio.samples + if isinstance(samples, torch.Tensor): + total += samples.numel() * samples.element_size() + elif isinstance(samples, np.ndarray): + total += samples.nbytes + else: + raise TypeError(f"unsupported decoded audio type: {type(samples).__name__}") + return total + + def rgba_to_rgb( image: Image.Image, background_color: Union[tuple[int, int, int], list[int]] = (255, 255, 255) ) -> Image.Image: @@ -740,6 +783,7 @@ def _load_video_by_cv2( cv2_backend: Optional[int] = None, raw_bytes_hash: Optional[str] = None, max_decoded_frame_bytes: Optional[int] = None, + decoded_memory_reserver: Optional[_DecodedMemoryReserver] = None, ) -> VideoData: """Decode a video and return sampled frames as a list. @@ -758,9 +802,11 @@ def _load_video_by_cv2( without an extra frame-by-frame copy. `"pil"` - list[PIL.Image], one per sampled frame. - When `max_decoded_frame_bytes` is set, the selected output representation - is sized from the first decoded frame and rejected before the frame stack - is allocated. + When `max_decoded_frame_bytes` or `decoded_memory_reserver` is set, the + counted peak frame buffers are sized from the first decoded frame and + checked before the frame stack is allocated. This covers buffers allocated + here, not codec state, compressed input, Python overhead, or the process's + complete memory peak. """ assert format in ("pt", "np", "pil"), "format must be one of 'pt', 'np', 'pil'" @@ -784,6 +830,7 @@ def _load_video_by_cv2( else: vidcap = cv2.VideoCapture(video) + decoded_memory_reservation: Optional[_DecodedMemoryReservation] = None try: if not vidcap.isOpened(): src_repr = ( @@ -834,17 +881,35 @@ def _load_video_by_cv2( fh, fw = bgr_scratch.shape[:2] if stacked_rgb is None: H, W = fh, fw - output_dtype = np.float32 if format == "pt" else np.uint8 - output_item_size = np.dtype(output_dtype).itemsize - decoded_frame_bytes = num_frames_to_sample * H * W * 3 * output_item_size + frame_values = num_frames_to_sample * H * W * 3 + if format == "pt": + # uint8 RGB stack + float32 NumPy normalization + # buffer + contiguous float32 CHW tensor. A + # non-CPU destination briefly adds another tensor. + peak_item_size = ( + np.dtype(np.uint8).itemsize + 2 * np.dtype(np.float32).itemsize + ) + peak_item_size += ( + np.dtype(np.float32).itemsize + if torch.device(device).type != "cpu" + else 0 + ) + else: + peak_item_size = np.dtype(np.uint8).itemsize + peak_decoded_frame_bytes = frame_values * peak_item_size if ( max_decoded_frame_bytes is not None - and decoded_frame_bytes > max_decoded_frame_bytes + and peak_decoded_frame_bytes > max_decoded_frame_bytes ): raise ValueError( - f"Decoded video frames require {decoded_frame_bytes} bytes, " + "Decoded video frame buffers require " + f"{peak_decoded_frame_bytes} bytes, " f"exceeding the limit of {max_decoded_frame_bytes} bytes." ) + if decoded_memory_reserver is not None: + decoded_memory_reservation = decoded_memory_reserver( + peak_decoded_frame_bytes + ) stacked_rgb = np.empty((num_frames_to_sample, H, W, 3), dtype=np.uint8) if (fh, fw) == (H, W): cv2.cvtColor( @@ -878,6 +943,9 @@ def _load_video_by_cv2( if device != "cpu": tensor_nchw = tensor_nchw.to(device) loaded_frames = list(torch.unbind(tensor_nchw, dim=0)) + del tensor_nchw + del stacked_f32 + del stacked_rgb elif format == "np": loaded_frames = stacked_rgb else: # "pil" @@ -889,36 +957,54 @@ def _load_video_by_cv2( "duration": duration, "frames_indices": valid_indices, } + except BaseException: + if decoded_memory_reservation is not None: + decoded_memory_reservation.release() + raise finally: # Release the OpenCV handle before any downstream re-open (e.g. PyAV # audio extraction on Windows) and to avoid leaking descriptors. vidcap.release() - audio = None - if extract_audio: - # extract_audio_from_video accepts Union[str, BytesIO]. When the - # in-memory cv2 path was taken, `video` is still the original - # bytes parameter — wrap in a fresh BytesIO (the one consumed by - # cv2 is at end-of-stream). When the tempfile fallback was used, - # `video` is the file path string and is passed through as-is. - audio_source = ( - BytesIO(bytes(video)) if isinstance(video, (bytes, bytearray, memoryview)) else video - ) - try: - audio_samples, audio_sample_rate = extract_audio_from_video(audio_source) - audio = AudioData(samples=audio_samples, sample_rate=audio_sample_rate) - except ValueError as e: - if "No audio stream found" in str(e): - logger.warning("Video has no audio track, skipping audio extraction.") - else: - raise + try: + audio = None + if extract_audio: + # extract_audio_from_video accepts Union[str, BytesIO]. When the + # in-memory cv2 path was taken, `video` is still the original + # bytes parameter — wrap in a fresh BytesIO (the one consumed by + # cv2 is at end-of-stream). When the tempfile fallback was used, + # `video` is the file path string and is passed through as-is. + audio_source = ( + BytesIO(bytes(video)) + if isinstance(video, (bytes, bytearray, memoryview)) + else video + ) + try: + audio_samples, audio_sample_rate = extract_audio_from_video(audio_source) + audio = AudioData(samples=audio_samples, sample_rate=audio_sample_rate) + except ValueError as e: + if "No audio stream found" in str(e): + logger.warning("Video has no audio track, skipping audio extraction.") + else: + raise + except BaseException: + if decoded_memory_reservation is not None: + decoded_memory_reservation.release() + raise - return VideoData( + video_data = VideoData( frames=loaded_frames, metadata=metadata, audio=audio, raw_bytes_hash=raw_bytes_hash, ) + if decoded_memory_reservation is not None: + try: + decoded_memory_reservation.commit(_decoded_video_bytes(video_data)) + except Exception: + decoded_memory_reservation.release() + raise + return video_data def _normalize_file_uri(uri: str) -> str: @@ -1178,6 +1264,7 @@ def __init__( device: str = "cpu", extract_audio: bool = False, max_decoded_frame_bytes: Optional[int] = None, + decoded_memory_reserver: Optional[_DecodedMemoryReserver] = None, ) -> None: # `"pt"` is the safe default every video model handles; models tuned # for the uint8-HWC path opt into `"np"` via @@ -1192,6 +1279,7 @@ def __init__( self._device = device self._extract_audio = extract_audio self._max_decoded_frame_bytes = max_decoded_frame_bytes + self._decoded_memory_reserver = decoded_memory_reserver @classmethod def merge_kwargs( @@ -1231,6 +1319,7 @@ def load_bytes(self, data: bytes) -> VideoData: cv2_backend=cv2_backend, raw_bytes_hash=raw_bytes_hash, max_decoded_frame_bytes=self._max_decoded_frame_bytes, + decoded_memory_reserver=self._decoded_memory_reserver, ) with tempfile.NamedTemporaryFile(suffix=".mp4", dir=_VIDEO_TEMPFILE_DIR) as tmp: tmp.write(data) @@ -1244,6 +1333,7 @@ def load_bytes(self, data: bytes) -> VideoData: extract_audio=self._extract_audio, raw_bytes_hash=raw_bytes_hash, max_decoded_frame_bytes=self._max_decoded_frame_bytes, + decoded_memory_reserver=self._decoded_memory_reserver, ) def load_base64(self, media_type: str, data: str) -> VideoData: diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 40056d56ef14..12a955c04633 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -41,6 +41,7 @@ def _ignore_synthetic_decoded_size(monkeypatch) -> None: handler.loader_factory, lambda _decoded: 0, handler.runtime_available, + handler.reserves_decoded_bytes, ), ) @@ -55,6 +56,7 @@ def _set_video_runtime_available(monkeypatch, available: bool) -> None: handler.loader_factory, handler.decoded_bytes, lambda: available, + handler.reserves_decoded_bytes, ), ) @@ -273,6 +275,15 @@ def test_decoded_video_bytes_include_frames_and_audio(): request_mapping._decoded_video_bytes(object()) +def test_decoded_media_budget_settles_peak_to_retained_bytes(): + budget = request_mapping._DecodedMediaBudget(20) + + first = budget.reserve(15) + first.commit(5) + second = budget.reserve(15) + second.release() + + def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): llm, _ = image_llm monkeypatch.setattr( @@ -289,7 +300,7 @@ def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, async def test_handler_table_drives_capability_and_mapping(image_llm, monkeypatch): llm, _ = image_llm - def video_loader_factory(_llm): + def video_loader_factory(_llm, _budget): async def load(item, index): return f"video-{index}:{bytes(item.raw_bytes).decode()}" @@ -457,14 +468,13 @@ async def async_load(self, source): mapped = await request_mapping._input_from_request(request, image_video_llm) - assert created_with == [ - ( - "video", - {"format": "np", "num_frames": 4}, - {"max_decoded_frame_bytes": request_mapping._MAX_TOTAL_DECODED_MEDIA_BYTES}, - ), - ("image", {"format": "np"}, None), - ] + assert len(created_with) == 2 + video_modality, video_defaults, video_runtime = created_with[0] + assert video_modality == "video" + assert video_defaults == {"format": "np", "num_frames": 4} + assert set(video_runtime) == {"decoded_memory_reserver"} + assert callable(video_runtime["decoded_memory_reserver"]) + assert created_with[1] == ("image", {"format": "np"}, None) assert [video.metadata["source"] for video in mapped["multi_modal_data"]["video"]] == [ "https://example.com/first.mp4", "data:video/mp4;base64,dmlkZW8=", @@ -633,6 +643,45 @@ async def async_load_bytes(self, data): await request_mapping._input_from_request(request, llm) +@pytest.mark.asyncio +async def test_concurrent_video_loads_share_budget_before_allocation(image_video_llm, monkeypatch): + allocations = [] + + class _VideoIo: + def __init__(self, reserver): + self._reserver = reserver + + @classmethod + def create(cls, defaults, runtime): + return cls(runtime["decoded_memory_reserver"]) + + async def async_load_bytes(self, data): + reservation = self._reserver(12) + allocations.append(bytes(data)) + try: + await asyncio.Future() + finally: + reservation.release() + + monkeypatch.setattr(request_mapping, "VideoMediaIO", _VideoIo) + monkeypatch.setattr(request_mapping, "_MAX_TOTAL_DECODED_MEDIA_BYTES", 20) + request = generation_pb2.GenerateRequest( + prompt="videos", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=value, + ) + for value in (b"first", b"second") + ], + ) + + with pytest.raises(ValueError, match="decoded media exceeds"): + await request_mapping._input_from_request(request, image_video_llm) + + assert allocations == [b"first"] + + @pytest.mark.asyncio async def test_image_mapping_supports_token_ids(image_llm, monkeypatch): llm, _ = image_llm diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index 1675b7357f7c..ba36fffd5bee 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -13,7 +13,7 @@ import tempfile import threading from io import BytesIO -from unittest.mock import patch +from unittest.mock import Mock, patch import numpy as np import pytest @@ -326,7 +326,11 @@ async def test_video_accepts_video_mime(self): def test_video_media_io_forwards_decoded_frame_limit(): decoded = object() - io = media_io_module.VideoMediaIO(max_decoded_frame_bytes=123) + reserver = Mock() + io = media_io_module.VideoMediaIO( + max_decoded_frame_bytes=123, + decoded_memory_reserver=reserver, + ) with ( patch.object(media_io_module, "_select_cv2_stream_buffered_backend", return_value=7), @@ -335,3 +339,4 @@ def test_video_media_io_forwards_decoded_frame_limit(): assert io.load_bytes(b"video") is decoded assert decode.call_args.kwargs["max_decoded_frame_bytes"] == 123 + assert decode.call_args.kwargs["decoded_memory_reserver"] is reserver diff --git a/tests/unittest/inputs/test_video_decode.py b/tests/unittest/inputs/test_video_decode.py index c9818d296a36..2d01aeae28af 100644 --- a/tests/unittest/inputs/test_video_decode.py +++ b/tests/unittest/inputs/test_video_decode.py @@ -70,22 +70,76 @@ def test_np_format_hits_hf_video_processor_fast_path(sample_video_path: str) -> assert np.shares_memory(video.frames, batched[0]) -@pytest.mark.parametrize( - ("format", "dtype"), [("np", np.uint8), ("pil", np.uint8), ("pt", np.float32)] -) +@pytest.mark.parametrize(("format", "peak_bytes_per_value"), [("np", 1), ("pil", 1), ("pt", 9)]) def test_decode_limit_is_checked_before_frame_stack_allocation( - sample_video_path: str, format: str, dtype: type[np.generic] + sample_video_path: str, format: str, peak_bytes_per_value: int ) -> None: - decoded_bytes = 10 * 64 * 64 * 3 * np.dtype(dtype).itemsize + peak_bytes = 10 * 64 * 64 * 3 * peak_bytes_per_value with patch("tensorrt_llm.inputs.media_io.np.empty") as allocate: - with pytest.raises(ValueError, match=f"require {decoded_bytes} bytes"): + with pytest.raises(ValueError, match=f"require {peak_bytes} bytes"): _load_video_by_cv2( sample_video_path, num_frames=10, fps=-1, format=format, - max_decoded_frame_bytes=decoded_bytes - 1, + max_decoded_frame_bytes=peak_bytes - 1, ) allocate.assert_not_called() + + +def test_pt_decode_reserves_peak_then_commits_retained_bytes(sample_video_path: str) -> None: + class _Reservation: + committed = None + released = False + + def commit(self, retained_bytes: int) -> None: + self.committed = retained_bytes + + def release(self) -> None: + self.released = True + + reservation = _Reservation() + reserved = [] + + def reserve(peak_bytes: int) -> _Reservation: + reserved.append(peak_bytes) + return reservation + + _load_video_by_cv2( + sample_video_path, + num_frames=10, + fps=-1, + format="pt", + decoded_memory_reserver=reserve, + ) + + assert reserved == [10 * 64 * 64 * 3 * 9] + assert reservation.committed == 10 * 64 * 64 * 3 * 4 + assert reservation.released is False + + +def test_decode_releases_reservation_after_allocation_failure(sample_video_path: str) -> None: + class _Reservation: + released = False + + def commit(self, retained_bytes: int) -> None: + raise AssertionError("failed decode must not commit") + + def release(self) -> None: + self.released = True + + reservation = _Reservation() + + with patch("tensorrt_llm.inputs.media_io.np.empty", side_effect=MemoryError("allocation")): + with pytest.raises(MemoryError, match="allocation"): + _load_video_by_cv2( + sample_video_path, + num_frames=10, + fps=-1, + format="np", + decoded_memory_reserver=lambda _peak_bytes: reservation, + ) + + assert reservation.released is True From 7138d1e94f4e27165f166d53e5e3b56034218325 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Tue, 20 Oct 2026 13:49:46 -0700 Subject: [PATCH 28/41] test(openengine): tighten audio conformance coverage Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 4 ++- .../grpc/openengine/test_multimodal.py | 20 +++++++++++++++ .../openengine/test_multimodal_conformance.py | 25 +++++++++++-------- 3 files changed, 38 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index edd3575a06f8..532096d019ed 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -84,7 +84,9 @@ def resolve_media_capabilities( return MediaCapabilities({}) if not isinstance(input_processor, BaseMultimodalInputProcessor): registered_model_type = getattr(type(input_processor), "_registered_model_type", None) - if registered_model_type != model_type or not callable(input_processor): + if registered_model_type != model_type: + return MediaCapabilities({}) + if not callable(input_processor): return MediaCapabilities({}) placeholders = { diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 7c7f698969b5..05e7b69b8320 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -114,6 +114,15 @@ def __call__(self, inputs, sampling_params): raise NotImplementedError +class _NonCallableRegisteredInputProcessor: + """Malformed registration used to verify capability discovery fails closed.""" + + _registered_model_type = _ImageConfig.model_type + + def __init__(self, config: _ImageConfig) -> None: + self.config = config + + class _Tokenizer: def convert_ids_to_tokens( self, token_ids: Sequence[int], skip_special_tokens: bool = False @@ -285,6 +294,17 @@ def test_media_capability_accepts_a_registered_non_base_processor(monkeypatch): assert request_mapping.media_capabilities_for_model(llm).modalities == {"audio"} +def test_media_capability_rejects_a_registered_non_callable_processor(monkeypatch): + config = _ImageConfig() + monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", lambda *_args: True) + llm = SimpleNamespace( + _hf_model_config=config, + input_processor=_NonCallableRegisteredInputProcessor(config), + ) + + assert request_mapping.media_capabilities_for_model(llm).modalities == set() + + def test_media_capabilities_are_hashable(): first = MediaCapabilities({"image": ""}) second = MediaCapabilities({"image": ""}) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index 770e3cde295d..ec748021649a 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -61,21 +61,28 @@ def _whisper_model_path() -> Path: raise FileNotFoundError(f"whisper-tiny not found under {root}") -def _whisper_audio_path() -> Path: +def _whisper_audio_path(model_path: Path) -> Path: root = llm_models_root(check=True) - for checkpoint in ("large-v3", "large-v2"): - candidate = root / f"whisper-models/{checkpoint}/1221-135766-0002.wav" + filename = "1221-135766-0002.wav" + candidates = [ + model_path / filename, + root / f"whisper-models/large-v3/{filename}", + root / f"whisper-models/large-v2/{filename}", + ] + for candidate in candidates: if candidate.exists(): return candidate - raise FileNotFoundError(f"Whisper test audio not found under {root}/whisper-models") + raise FileNotFoundError(f"Whisper test audio not found; tried {candidates}") @pytest.fixture def whisper(): if not torch.cuda.is_available(): pytest.skip("requires a GPU to run a real Whisper engine") + model_path = _whisper_model_path() + audio_path = _whisper_audio_path(model_path) engine = LLM( - model=str(_whisper_model_path()), + model=str(model_path), disable_overlap_scheduler=True, enable_chunked_prefill=False, kv_cache_config=KvCacheConfig( @@ -90,7 +97,7 @@ def whisper(): scheduler_config=SchedulerConfig(use_python_scheduler=True), ) try: - yield engine + yield engine, audio_path finally: engine.shutdown() @@ -205,6 +212,7 @@ async def collect_video(): def test_audio_capability_agrees_with_generate(whisper): + whisper, audio_path = whisper inference = OpenEngineInferenceServicer(whisper, WHISPER_MODEL_NAME) control = OpenEngineControlServicer(whisper, WHISPER_MODEL_NAME, inference) info = asyncio.run( @@ -219,9 +227,7 @@ def test_audio_capability_agrees_with_generate(whisper): assert info.supports_token_ids_input is False assert list(info.extra["trtllm_supported_modalities"]) == ["audio"] - audio_uri = ( - "data:audio/wav;base64," + base64.b64encode(_whisper_audio_path().read_bytes()).decode() - ) + audio_uri = "data:audio/wav;base64," + base64.b64encode(audio_path.read_bytes()).decode() request = generation_pb2.GenerateRequest( request_id="openengine-audio-conformance", model=WHISPER_MODEL_NAME, @@ -232,7 +238,6 @@ def test_audio_capability_agrees_with_generate(whisper): generation_pb2.MediaItem( modality=generation_pb2.MODALITY_AUDIO, data_uri=audio_uri, - uuid="librispeech-1221-135766-0002", ) ], ) From 4941f545a9c00c7082ccf009af85c73dc2ed96f2 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Tue, 20 Oct 2026 14:02:16 -0700 Subject: [PATCH 29/41] fix(inputs): isolate media sessions by event loop Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/server.py | 6 ++- tensorrt_llm/inputs/media_io.py | 40 ++++++++++--------- .../grpc/openengine/test_openengine.py | 3 ++ .../inputs/test_async_media_loading.py | 33 ++++++++++++--- tests/unittest/inputs/test_url_validation.py | 26 +++++------- 5 files changed, 67 insertions(+), 41 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/server.py b/tensorrt_llm/grpc/openengine/server.py index 546adc9b7bb4..8e0a63f82556 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -20,6 +20,7 @@ import uvloop from tensorrt_llm import LLM as PyTorchLLM +from tensorrt_llm.inputs.media_io import _close_aiohttp_session from tensorrt_llm.logger import logger from tensorrt_llm.serve._frontend_processes import ( FrontendStartupCancelled, @@ -170,7 +171,10 @@ async def stop(self, grace: float = 5.0) -> None: Args: grace: Maximum time in seconds to allow active RPCs to finish. """ - await self._server.stop(grace=grace) + try: + await self._server.stop(grace=grace) + finally: + await _close_aiohttp_session() logger.info("OpenEngine server stopped") async def wait_for_termination(self) -> None: diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 2aeb5c7065d8..c7c4922beabf 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -10,6 +10,7 @@ import os import socket import tempfile +import threading from abc import ABC, abstractmethod from concurrent.futures import Executor from io import BytesIO @@ -130,29 +131,32 @@ def convert_image_mode(image: Image.Image, to_mode: str) -> Image.Image: # `_load_video_by_cv2` for the per-format contract. _SUPPORTED_VIDEO_FORMATS = ("pt", "np", "pil") -# Module-level aiohttp session shared across all media fetch calls. -# Created lazily on first use inside the async event loop, then reused so that -# TCP connections are kept alive and not re-established per request. -_global_aiohttp_session: aiohttp.ClientSession | None = None -_global_aiohttp_session_loop: asyncio.AbstractEventLoop | None = None +# One aiohttp session per event loop. Media loading can run on multiple live +# loops in one process, and aiohttp sessions must only be used and closed from +# their owning loop. +_aiohttp_sessions: dict[asyncio.AbstractEventLoop, aiohttp.ClientSession] = {} +_aiohttp_sessions_lock = threading.Lock() async def _get_aiohttp_session() -> aiohttp.ClientSession: """Return the aiohttp session owned by the current event loop.""" - global _global_aiohttp_session, _global_aiohttp_session_loop current_loop = asyncio.get_running_loop() - if ( - _global_aiohttp_session is None - or _global_aiohttp_session.closed - or _global_aiohttp_session_loop is not current_loop - ): - previous_session = _global_aiohttp_session - connector = aiohttp.TCPConnector(resolver=_ValidatingResolver()) - _global_aiohttp_session = aiohttp.ClientSession(connector=connector) - _global_aiohttp_session_loop = current_loop - if previous_session is not None and not previous_session.closed: - await previous_session.close() - return _global_aiohttp_session + with _aiohttp_sessions_lock: + session = _aiohttp_sessions.get(current_loop) + if session is None or session.closed: + connector = aiohttp.TCPConnector(resolver=_ValidatingResolver()) + session = aiohttp.ClientSession(connector=connector) + _aiohttp_sessions[current_loop] = session + return session + + +async def _close_aiohttp_session() -> None: + """Close and forget the media session owned by the current event loop.""" + current_loop = asyncio.get_running_loop() + with _aiohttp_sessions_lock: + session = _aiohttp_sessions.pop(current_loop, None) + if session is not None and not session.closed: + await session.close() # Maximum allowed response size for remote fetches (200 MB). diff --git a/tests/unittest/grpc/openengine/test_openengine.py b/tests/unittest/grpc/openengine/test_openengine.py index fdd464dbe70a..0efb4c0cc886 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -29,6 +29,7 @@ server_pb2, ) from tensorrt_llm.grpc.openengine.server import OpenEngineServer # noqa: E402 +from tensorrt_llm.inputs.media_io import _get_aiohttp_session # noqa: E402 # grpc.aio starts a `_poll_wrapper` daemon thread on server start and tears it # down only once its internal state is released, which happens after the event @@ -78,6 +79,7 @@ async def exercise_server() -> None: assert server.port != 0 await server.start() + media_session = await _get_aiohttp_session() with pytest.raises(RuntimeError, match="Failed to bind"): OpenEngineServer(host="127.0.0.1", port=server.port, llm=llm, model="test-model") channel = grpc.aio.insecure_channel(f"127.0.0.1:{server.port}") @@ -99,6 +101,7 @@ async def exercise_server() -> None: finally: await channel.close() await server.stop(grace=0) + assert media_session.closed asyncio.run(exercise_server()) diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index ba36fffd5bee..06111607c016 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -9,6 +9,7 @@ - MultimodalDataTracker.retrieve_all_async gathers all modalities concurrently """ +import asyncio import base64 import tempfile import threading @@ -24,7 +25,7 @@ import tensorrt_llm.inputs.media_io as media_io_module import tensorrt_llm.inputs.utils as utils_module -from tensorrt_llm.inputs.media_io import _get_aiohttp_session +from tensorrt_llm.inputs.media_io import _close_aiohttp_session, _get_aiohttp_session from tensorrt_llm.inputs.utils import MultimodalDataTracker, async_load_audio, async_load_image pytestmark = pytest.mark.cpu_only @@ -156,12 +157,10 @@ def tracking_read(*args, **kwargs): class TestSessionReuse: @pytest_asyncio.fixture(autouse=True) - async def reset_global_session(self): - media_io_module._global_aiohttp_session = None + async def reset_loop_session(self): + await _close_aiohttp_session() yield - if media_io_module._global_aiohttp_session is not None: - await media_io_module._global_aiohttp_session.close() - media_io_module._global_aiohttp_session = None + await _close_aiohttp_session() @pytest.mark.asyncio async def test_same_session_returned_on_repeated_calls(self): @@ -182,6 +181,28 @@ async def test_new_session_created_after_close(self): ) assert not second_session.closed + @pytest.mark.asyncio + async def test_live_event_loops_keep_independent_sessions(self): + first_session = await _get_aiohttp_session() + worker_ready = threading.Event() + release_worker = threading.Event() + worker_sessions = [] + + async def use_session_on_worker_loop() -> None: + worker_sessions.append(await _get_aiohttp_session()) + worker_ready.set() + await asyncio.to_thread(release_worker.wait, 5) + await _close_aiohttp_session() + + worker = asyncio.create_task(asyncio.to_thread(asyncio.run, use_session_on_worker_loop())) + try: + assert await asyncio.to_thread(worker_ready.wait, 5) + assert worker_sessions[0] is not first_session + assert not first_session.closed + finally: + release_worker.set() + await worker + # ────────────────────────────────────────────────────────────── # MultimodalDataTracker.retrieve_all_async — concurrency diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index 5cf2391104a1..3530dd559127 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -394,28 +394,22 @@ async def check(): try: assert isinstance(session.connector._resolver, _ValidatingResolver) finally: - await session.close() - media_io_module._global_aiohttp_session = None - media_io_module._global_aiohttp_session_loop = None + await media_io_module._close_aiohttp_session() self._run(check()) - def test_shared_session_is_recreated_for_a_new_event_loop(self): + def test_each_event_loop_closes_its_own_session(self): async def get_session(): - return await media_io_module._get_aiohttp_session() + session = await media_io_module._get_aiohttp_session() + await media_io_module._close_aiohttp_session() + return session first = self._run(get_session()) - second = None - try: - second = self._run(get_session()) - assert second is not first - assert first.closed - finally: - session = second or first - if not session.closed: - self._run(session.close()) - media_io_module._global_aiohttp_session = None - media_io_module._global_aiohttp_session_loop = None + second = self._run(get_session()) + + assert second is not first + assert first.closed + assert second.closed @patch("tensorrt_llm.inputs.media_io.socket.getaddrinfo", return_value=PUBLIC_DNS) def test_maps_timeout_to_fetch_error(self, _): From c0e8499a981b28defb27e0d3e5a44faff17ae509 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Wed, 21 Oct 2026 12:03:40 -0700 Subject: [PATCH 30/41] fix(openengine): validate and cancel multimodal request preparation Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 2 + .../grpc/openengine/request_mapping.py | 146 +++- tensorrt_llm/grpc/openengine/servicer.py | 74 +- tensorrt_llm/inputs/media_io.py | 59 +- .../grpc/openengine/test_multimodal.py | 687 +++++++++++++++++- .../openengine/test_multimodal_conformance.py | 40 +- tests/unittest/inputs/test_url_validation.py | 39 + 7 files changed, 958 insertions(+), 89 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 1c65755be0af..83a6f9979b51 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -77,6 +77,8 @@ def resolve_media_capabilities( input_processor = getattr(llm, "input_processor", None) if not isinstance(input_processor, BaseMultimodalInputProcessor): return MediaCapabilities({}) + if getattr(getattr(llm, "args", None), "disable_mm_encoder", False): + return MediaCapabilities({}) model_type = _loaded_model_type(llm) if not model_type: return MediaCapabilities({}) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index d58bd45e62bf..c55fbd075495 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -4,6 +4,7 @@ """Translate an OpenEngine GenerateRequest into TensorRT-LLM request inputs.""" import asyncio +from collections import Counter from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from typing import Any @@ -35,15 +36,24 @@ _MediaLoader = Callable[[generation_pb2.MediaItem, int], Awaitable[Any]] +_MediaValidator = Callable[[generation_pb2.MediaItem, int], None] @dataclass(frozen=True) class _MediaHandler: name: str + validate: _MediaValidator loader_factory: Callable[[Any], _MediaLoader] decoded_bytes: Callable[[Any], int] +@dataclass(frozen=True) +class _PreparedInput: + inputs: str | dict[str, list[int]] + media: tuple[tuple[generation_pb2.MediaItem, _MediaHandler], ...] + capabilities: MediaCapabilities | None + + def conversation_affinity_enabled(llm: Any) -> bool: """Whether TensorRT-LLM owns attention-DP placement for conversations.""" args = getattr(llm, "args", None) @@ -285,30 +295,51 @@ def _image_io_for_model(llm: Any) -> ImageMediaIO: return ImageMediaIO.create(defaults, None) -async def _load_image_item( - item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int -) -> Any: +def _validate_image_item(item: generation_pb2.MediaItem, index: int) -> None: source = item.WhichOneof("source") if source is None: raise ValueError(f"media[{index}] must set exactly one source") if source == "url": - if urlparse(item.url).scheme.lower() not in ("http", "https"): + try: + parsed = urlparse(item.url) + hostname = parsed.hostname + _ = parsed.port # Access validates the port syntax and range. + except ValueError as error: + raise ValueError(f"media[{index}].url is malformed: {error}") from error + if parsed.scheme.lower() not in ("http", "https"): raise ValueError(f"media[{index}].url must use http or https") - value = item.url + if not hostname: + raise ValueError(f"media[{index}].url must include a hostname") elif source == "data_uri": - if urlparse(item.data_uri).scheme.lower() != "data": + parsed = urlparse(item.data_uri) + if parsed.scheme.lower() != "data": raise ValueError(f"media[{index}].data_uri must use the data scheme") - value = item.data_uri + data_spec, separator, payload = parsed.path.partition(",") + if not separator or not payload: + raise ValueError(f"media[{index}].data_uri must include a comma and a payload") + media_type, _, encoding = data_spec.partition(";") + if encoding != "base64": + raise ValueError(f"media[{index}].data_uri must use base64 encoding") + if media_type and not media_type.lower().startswith("image/"): + raise ValueError(f"media[{index}].data_uri must declare an image type") elif source == "raw_bytes": if item.mime_type and not item.mime_type.lower().startswith("image/"): raise ValueError(f"media[{index}].mime_type must be an image type for raw image bytes") + else: + raise ValueError(f"media[{index}] has an unknown source '{source}'") + + +async def _load_image_item( + item: generation_pb2.MediaItem, image_io: ImageMediaIO, index: int +) -> Any: + source = item.WhichOneof("source") + if source == "raw_bytes": try: return await image_io.async_load_bytes(bytes(item.raw_bytes)) except Exception as error: raise ValueError(f"media[{index}].raw_bytes could not be decoded: {error}") from error - else: - raise ValueError(f"media[{index}] has an unknown source '{source}'") + value = item.url if source == "url" else item.data_uri try: return await image_io.async_load(value) except MediaFetchError as error: @@ -344,7 +375,7 @@ def _decoded_image_bytes(image: Any) -> int: _MEDIA_HANDLERS: dict[int, _MediaHandler] = { generation_pb2.MODALITY_IMAGE: _MediaHandler( - "image", _image_loader_factory, _decoded_image_bytes + "image", _validate_image_item, _image_loader_factory, _decoded_image_bytes ), } @@ -360,15 +391,15 @@ def _modality_name(modality: int) -> str: return f"unknown ({modality})" -async def _input_from_request( +def _prepare_input_from_request( request: generation_pb2.GenerateRequest, llm: Any, media_capabilities: MediaCapabilities | None = None, -) -> str | dict[str, Any]: - """Build LLM API inputs without blocking the event loop on media.""" +) -> _PreparedInput: + """Validate request media without starting any fetch or decode work.""" inputs = _base_input_from_request(request) if not request.media: - return inputs + return _PreparedInput(inputs, (), None) capabilities = media_capabilities or media_capabilities_for_model(llm) modalities = capabilities.modalities @@ -395,8 +426,73 @@ async def _input_from_request( ) resolved_handlers.append(handler) + media = tuple(zip(request.media, resolved_handlers)) + for index, (item, handler) in enumerate(media): + handler.validate(item, index) + + return _PreparedInput(inputs, media, capabilities) + + +def _validate_media_placeholders(prepared: _PreparedInput, llm: Any) -> None: + """Reject placeholder mismatches before probing the stateful MM processor.""" + if not prepared.media or prepared.capabilities is None: + return + tokenizer = llm.tokenizer + prompt_ids = ( + tokenizer.encode(prepared.inputs, add_special_tokens=False) + if isinstance(prepared.inputs, str) + else prepared.inputs["prompt_token_ids"] + ) + item_counts: Counter[str] = Counter() + expected_sequences: Counter[tuple[int, ...]] = Counter() + for _, handler in prepared.media: + item_counts[handler.name] += 1 + placeholder = prepared.capabilities.placeholder(handler.name, item_counts[handler.name]) + sequence = tuple(tokenizer.encode(placeholder, add_special_tokens=False)) + if not sequence: + raise ValueError(f"{handler.name} placeholder must encode to at least one token") + expected_sequences[sequence] += 1 + + for sequence, expected in expected_sequences.items(): + width = len(sequence) + actual = sum( + tuple(prompt_ids[index : index + width]) == sequence + for index in range(len(prompt_ids) - width + 1) + ) + if actual != expected: + raise ValueError( + f"media placeholder count does not match media items: expected {expected}, " + f"got {actual} for placeholder token IDs {list(sequence)}" + ) + + processor = llm.input_processor + mm_ids = processor.get_mm_token_ids() + if mm_ids is not None: + special_ids = processor.get_mm_special_token_ids() + content_ids = set(mm_ids.tolist()) - ( + set(special_ids.tolist()) if special_ids is not None else set() + ) + expected_tokens: Counter[int] = Counter() + for sequence, count in expected_sequences.items(): + for token_id in sequence: + if token_id in content_ids: + expected_tokens[token_id] += count + actual_tokens = Counter(token_id for token_id in prompt_ids if token_id in content_ids) + if actual_tokens != expected_tokens: + raise ValueError("multimodal placeholder token counts do not match media items") + + +async def _load_prepared_input(prepared: _PreparedInput, llm: Any) -> str | dict[str, Any]: + """Fetch and decode media from an already validated request plan.""" + if not prepared.media: + return prepared.inputs + + capabilities = prepared.capabilities + if capabilities is None: + raise RuntimeError("media input plan is missing capabilities") + loaders: dict[str, _MediaLoader] = {} - for handler in resolved_handlers: + for _, handler in prepared.media: if handler.name not in loaders: loaders[handler.name] = handler.loader_factory(llm) semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) @@ -418,7 +514,7 @@ async def load_bounded( tasks = [ asyncio.create_task(load_bounded(item, index, handler)) - for index, (item, handler) in enumerate(zip(request.media, resolved_handlers)) + for index, (item, handler) in enumerate(prepared.media) ] try: decoded_media = await asyncio.gather(*tasks) @@ -431,7 +527,7 @@ async def load_bounded( grouped_media: dict[str, list[Any]] = {} grouped_uuids: dict[str, list[str | None]] = {} item_order: list[dict[str, Any]] = [] - for item, handler, decoded in zip(request.media, resolved_handlers, decoded_media): + for (item, handler), decoded in zip(prepared.media, decoded_media): modality_items = grouped_media.setdefault(handler.name, []) modality_index = len(modality_items) modality_items.append(decoded) @@ -444,14 +540,26 @@ async def load_bounded( } ) - mapped: dict[str, Any] = {"prompt": inputs} if isinstance(inputs, str) else dict(inputs) + mapped: dict[str, Any] = ( + {"prompt": prepared.inputs} if isinstance(prepared.inputs, str) else dict(prepared.inputs) + ) mapped["multi_modal_data"] = grouped_media - if any(item.uuid for item in request.media): + if any(item.uuid for item, _ in prepared.media): mapped["multi_modal_uuids"] = grouped_uuids mapped["mm_item_order"] = item_order return mapped +async def _input_from_request( + request: generation_pb2.GenerateRequest, + llm: Any, + media_capabilities: MediaCapabilities | None = None, +) -> str | dict[str, Any]: + """Build LLM API inputs without blocking the event loop on media.""" + prepared = _prepare_input_from_request(request, llm, media_capabilities) + return await _load_prepared_input(prepared, llm) + + __all__ = [ "conversation_params_from_request", "sampling_params_from_request", diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index dd1df20e2a3a..6f04a77cf5a6 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -28,8 +28,10 @@ from .errors import AbortFailedError, MediaUnavailableError, UnsupportedFeatureError from .formatting import _engine_error_response, _stop_texts from .request_mapping import ( - _input_from_request, + _load_prepared_input, _metadata_from_context, + _prepare_input_from_request, + _validate_media_placeholders, conversation_affinity_enabled, conversation_params_from_request, media_capabilities_for_model, @@ -48,17 +50,29 @@ class _LocalRequest: """Abort-aware local request state spanning mapping and generation.""" def __init__(self) -> None: - self.handle: Any = None + self.mapping_task: asyncio.Task[Any] | None = None + self.engine_handle: Any = None self.aborted = False - def bind(self, handle: Any) -> None: - self.handle = handle + def bind_mapping(self, task: asyncio.Task[Any]) -> None: + self.mapping_task = task + if self.aborted: + task.cancel() + + def clear_mapping(self, task: asyncio.Task[Any]) -> None: + if self.mapping_task is task: + self.mapping_task = None + + def bind_engine(self, handle: Any) -> None: + self.engine_handle = handle if self.aborted: handle.abort() def abort(self) -> None: - if self.handle is not None: - self.handle.abort() + if self.engine_handle is not None: + self.engine_handle.abort() + if self.mapping_task is not None: + self.mapping_task.cancel() self.aborted = True @@ -162,6 +176,7 @@ async def Generate( ) if self._frontend is not None: reservation = await self._frontend.reserve(request.request_id) + reservation.bind(local_request) async with aclosing( self._generate(request, context, local_request, registration, reservation) ) as responses: @@ -207,7 +222,9 @@ async def _generate( if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: raise UnsupportedFeatureError("prefix cache bypass is not supported") - inputs = await _input_from_request(request, self._llm, self._media_capabilities) + prepared_input = _prepare_input_from_request( + request, self._llm, self._media_capabilities + ) sampling_params = sampling_params_from_request(request, self._guided_backend) trace_headers, target_dp_rank, affinity_headers = _metadata_from_context(context) conversation_params = conversation_params_from_request(request) @@ -259,11 +276,48 @@ async def _generate( if request.HasField("kv") and request.kv.HasField("cache_salt") else None ) + + async def prepare_inputs() -> Any: + if prepared_input.media: + await asyncio.to_thread(_validate_media_placeholders, prepared_input, self._llm) + inputs = await _load_prepared_input(prepared_input, self._llm) + if prepared_input.media: + return await asyncio.to_thread( + self._llm.preprocess, inputs, sampling_params, disaggregated_params + ) + return inputs + + mapping_task = asyncio.create_task(prepare_inputs()) + local_request.bind_mapping(mapping_task) + try: + # wait() reports child cancellation as completion; RPC cancellation + # raises here instead, so the two paths remain distinguishable. + await asyncio.wait((mapping_task,)) + if mapping_task.cancelled() and local_request.aborted: + await context.abort( + grpc.StatusCode.CANCELLED, "Request aborted before submission" + ) + return + try: + inputs = mapping_task.result() + except (UnsupportedFeatureError, MediaUnavailableError, ValueError): + raise + except Exception as error: + # Unexpected preparation failures are worker errors, not bad inputs. + logger.error(f"OpenEngine request {request_id} could not be prepared: {error}") + await context.abort(grpc.StatusCode.INTERNAL, str(error)) + return + except asyncio.CancelledError: + mapping_task.cancel() + await asyncio.gather(mapping_task, return_exceptions=True) + raise + finally: + local_request.clear_mapping(mapping_task) except UnsupportedFeatureError as error: await context.abort(grpc.StatusCode.UNIMPLEMENTED, str(error)) return except MediaUnavailableError as error: - await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) + await context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(error)) return except ValueError as error: await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(error)) @@ -316,9 +370,7 @@ async def _generate( await context.abort(grpc.StatusCode.INTERNAL, str(error)) return - local_request.bind(result_handle) - if reservation is not None: - reservation.bind(result_handle) + local_request.bind_engine(result_handle) engine_terminal = False abort_requested = False diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 00b7594b95be..2bfd19e45fed 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -297,38 +297,33 @@ async def _safe_aiohttp_get(url: str, timeout_sec: int = 30) -> bytes: """Aiohttp GET wrapper that validates every redirect hop before following.""" await asyncio.to_thread(_validate_url, url) timeout = aiohttp.ClientTimeout(total=timeout_sec) - - async def _fetch(fetch_session: aiohttp.ClientSession) -> bytes: - current_url = url - for _ in range(_MAX_REDIRECTS + 1): - try: - async with fetch_session.get( - current_url, timeout=timeout, allow_redirects=False - ) as response: - if response.status in _REDIRECT_STATUSES: - redirect_url = response.headers.get("Location", "") - current_url = urljoin(current_url, redirect_url) - await asyncio.to_thread(_validate_url, current_url) - continue - maybe_coro = response.raise_for_status() - if asyncio.iscoroutine(maybe_coro): - await maybe_coro - return await _read_aiohttp_content(response) - except MediaFetchError: - raise - except aiohttp.ClientResponseError as error: - if error.status in (408, 429) or error.status >= 500: - raise MediaFetchError( - f"remote media server returned HTTP {error.status}" - ) from error - raise - except (aiohttp.InvalidURL, aiohttp.ClientSSLError): - raise - except (aiohttp.ClientError, asyncio.TimeoutError) as error: - raise MediaFetchError(f"remote media fetch failed: {error}") from error - raise RuntimeError("Too many redirects") - - return await _fetch(await _get_aiohttp_session()) + session = await _get_aiohttp_session() + current_url = url + for _ in range(_MAX_REDIRECTS + 1): + try: + async with session.get(current_url, timeout=timeout, allow_redirects=False) as response: + if response.status in _REDIRECT_STATUSES: + redirect_url = response.headers.get("Location", "") + current_url = urljoin(current_url, redirect_url) + await asyncio.to_thread(_validate_url, current_url) + continue + maybe_coro = response.raise_for_status() + if asyncio.iscoroutine(maybe_coro): + await maybe_coro + return await _read_aiohttp_content(response) + except MediaFetchError: + raise + except aiohttp.ClientResponseError as error: + if error.status in (408, 429) or error.status >= 500: + raise MediaFetchError( + f"remote media server returned HTTP {error.status}" + ) from error + raise + except (aiohttp.InvalidURL, aiohttp.ClientSSLError): + raise + except (aiohttp.ClientError, asyncio.TimeoutError) as error: + raise MediaFetchError(f"remote media fetch failed: {error}") from error + raise RuntimeError("Too many redirects") def _load_and_convert_image(image): diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 95cbb72071f4..25be2fe73946 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -4,6 +4,8 @@ """Unit tests for OpenEngine image input mapping and discovery.""" import asyncio +import re +import threading from collections.abc import AsyncIterator, Sequence from types import SimpleNamespace from typing import Any @@ -17,8 +19,9 @@ from conftest import AbortError, FakeServicerContext # noqa: E402 import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 -from tensorrt_llm.grpc.openengine.bindings import generation_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.bindings import generation_pb2, model_pb2 # noqa: E402 from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 +from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 @@ -26,6 +29,7 @@ MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, ) +from tensorrt_llm.llmapi.llm import PreprocessedInputs # noqa: E402 pytestmark = pytest.mark.cpu_only @@ -35,7 +39,12 @@ def _ignore_synthetic_decoded_size(monkeypatch) -> None: monkeypatch.setitem( request_mapping._MEDIA_HANDLERS, generation_pb2.MODALITY_IMAGE, - request_mapping._MediaHandler(handler.name, handler.loader_factory, lambda _decoded: 0), + request_mapping._MediaHandler( + handler.name, + handler.validate, + handler.loader_factory, + lambda _decoded: 0, + ), ) @@ -78,6 +87,13 @@ def get_preferred_media_io_kwargs(self): class _Tokenizer: + def encode(self, text: str, add_special_tokens: bool = True) -> list[int]: + del add_special_tokens + return [ + 1000 + int(token[7:-1]) if token.startswith("|.", text, re.DOTALL) + ] + def convert_ids_to_tokens( self, token_ids: Sequence[int], skip_special_tokens: bool = False ) -> list[str]: @@ -142,6 +158,50 @@ def generate_async(self, **kwargs): self.generate_kwargs = kwargs return self.result_handle + def preprocess(self, inputs, sampling_params, disaggregated_params): + return inputs + + +class _Reservation: + def __init__(self) -> None: + self.handle: Any = None + self.aborted = False + self.released = False + self._release_task: asyncio.Task[None] | None = None + + def bind(self, handle: Any) -> None: + self.handle = handle + if self.aborted: + handle.abort() + + def abort(self) -> bool: + if self.handle is not None: + self.handle.abort() + self.aborted = True + return True + + def release(self) -> asyncio.Task[None]: + async def mark_released() -> None: + self.released = True + + if self._release_task is None: + self._release_task = asyncio.create_task(mark_released()) + return self._release_task + + +class _Frontend: + def __init__(self, abort_before_bind: bool = False) -> None: + self.abort_before_bind = abort_before_bind + self.reservations: list[_Reservation] = [] + + async def reserve(self, _request_id: str) -> _Reservation: + reservation = _Reservation() + self.reservations.append(reservation) + if self.abort_before_bind: + self.abort_before_bind = False + reservation.abort() + return reservation + @pytest.fixture def image_llm(monkeypatch): @@ -181,6 +241,41 @@ def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): ) +@pytest.mark.asyncio +@pytest.mark.parametrize("disabled", [False, True]) +async def test_disabled_mm_encoder_agrees_between_discovery_and_generate(image_llm, disabled): + capability_llm, _ = image_llm + llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + llm.args.disable_mm_encoder = disabled + llm.args.reasoning_parser = None + llm.args.max_seq_len = 2048 + inference = OpenEngineInferenceServicer(llm, model="image-model") + control = OpenEngineControlServicer(llm, "image-model", inference) + + info = await control.GetModelInfo( + model_pb2.GetModelInfoRequest(model="image-model"), FakeServicerContext() + ) + + assert info.supports_multimodal is (not disabled) + assert list(info.extra["trtllm_supported_modalities"]) == ([] if disabled else ["image"]) + if disabled: + context = FakeServicerContext() + request = generation_pb2.GenerateRequest( + request_id="disabled-encoder", + model="image-model", + prompt="", + media=[ + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image") + ], + ) + with pytest.raises(AbortError): + async for _ in inference.Generate(request, context): + pass + assert context.abort_code == grpc.StatusCode.UNIMPLEMENTED + assert llm.generate_kwargs is None + assert inference.active_request_count() == 0 + + def test_media_capabilities_are_hashable(): first = MediaCapabilities({"image": ""}) second = MediaCapabilities({"image": ""}) @@ -219,7 +314,12 @@ async def load(item, index): monkeypatch.setitem( request_mapping._MEDIA_HANDLERS, generation_pb2.MODALITY_VIDEO, - request_mapping._MediaHandler("video", video_loader_factory, lambda decoded: len(decoded)), + request_mapping._MediaHandler( + "video", + lambda _item, _index: None, + video_loader_factory, + lambda decoded: len(decoded), + ), ) monkeypatch.setattr( MULTIMODAL_PLACEHOLDER_REGISTRY, @@ -329,6 +429,151 @@ async def async_load_bytes(self, data): ] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("invalid_case", "metadata", "expected_code"), + [ + ("sampling", (), grpc.StatusCode.INVALID_ARGUMENT), + ("metadata", (("openengine-priority", "1"),), grpc.StatusCode.UNIMPLEMENTED), + ( + "affinity", + (("x-trtllm-subagent-affinity-id", "parent"),), + grpc.StatusCode.INVALID_ARGUMENT, + ), + ("rank", (("openengine-target-dp-rank", "1"),), grpc.StatusCode.INVALID_ARGUMENT), + ], +) +async def test_request_validation_precedes_media_loading( + image_llm, monkeypatch, invalid_case, metadata, expected_code +): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + raise AssertionError("media loading must not start for an invalid request") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + request_id=f"invalid-{invalid_case}", + model="image-model", + prompt="", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + if invalid_case == "sampling": + request.sampling.num_sequences = 0 + context = FakeServicerContext(metadata=metadata) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + + with pytest.raises(AbortError): + async for _ in servicer.Generate(request, context): + pass + + assert context.abort_code == expected_code + assert llm.generate_kwargs is None + + +@pytest.mark.asyncio +async def test_capability_error_precedes_other_request_validation(image_llm): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + request = generation_pb2.GenerateRequest( + request_id="unsupported-video", + model="image-model", + prompt="video", + sampling=generation_pb2.SamplingParams(num_sequences=0), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"video", + ) + ], + ) + context = FakeServicerContext() + servicer = OpenEngineInferenceServicer(llm, model="image-model") + + with pytest.raises(AbortError): + async for _ in servicer.Generate(request, context): + pass + + assert context.abort_code == grpc.StatusCode.UNIMPLEMENTED + assert "video" in context.abort_details + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "source", + [ + {}, + {"url": "https://"}, + {"url": "https://example.com:bad/image"}, + {"url": "https://example.com:65536/image"}, + {"url": "https://[invalid/image"}, + {"data_uri": "data:image/png;base64"}, + {"data_uri": "data:image/png;base64,"}, + {"data_uri": "data:image/png,AAAA"}, + {"data_uri": "data:audio/wav;base64,AAAA"}, + ], +) +async def test_all_media_descriptions_are_validated_before_loading(image_llm, monkeypatch, source): + llm, _ = image_llm + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + raise AssertionError("media loading must not start before every item is valid") + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + request = generation_pb2.GenerateRequest( + prompt="images", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"valid", + ), + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, **source), + ], + ) + + with pytest.raises(ValueError, match=r"media\[1\]"): + await request_mapping._input_from_request(request, llm) + + +@pytest.mark.parametrize( + "source", + [ + {"url": "https://example.com/image.png"}, + {"url": "https://example.com:8443/image.png"}, + {"data_uri": "data:image/png;base64,AAAA"}, + {"data_uri": "data:;base64,AAAA"}, + {"raw_bytes": b"image", "mime_type": "image/png"}, + ], +) +def test_media_description_validation_accepts_supported_sources(image_llm, source): + llm, _ = image_llm + request = generation_pb2.GenerateRequest( + prompt="image", + media=[generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, **source)], + ) + + prepared = request_mapping._prepare_input_from_request(request, llm) + + assert len(prepared.media) == 1 + assert prepared.media[0][0] == request.media[0] + + @pytest.mark.asyncio async def test_image_loading_has_a_bounded_concurrency(image_llm, monkeypatch): llm, _ = image_llm @@ -534,7 +779,7 @@ async def async_load_bytes(self, data): request = generation_pb2.GenerateRequest( request_id="image-request", model="image-model", - token_ids=generation_pb2.TokenIds(ids=[1, 2]), + token_ids=generation_pb2.TokenIds(ids=[1001, 2]), media=[ generation_pb2.MediaItem( modality=generation_pb2.MODALITY_IMAGE, @@ -550,6 +795,254 @@ async def async_load_bytes(self, data): assert responses[-1].usage.total_tokens == 7 +@pytest.mark.asyncio +@pytest.mark.parametrize("input_kind", ["prompt", "token_ids"]) +@pytest.mark.parametrize( + "shape", ["missing", "duplicate", "extra_bare_pad", "unexpected_video_pad"] +) +async def test_placeholder_mismatch_does_not_poison_hashing( + image_llm, monkeypatch, input_kind, shape +): + import torch + + capability_llm, _ = image_llm + llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + processor = llm.input_processor + processor.multimodal_hashing_supported = None + monkeypatch.setattr(processor, "get_mm_token_ids", lambda: torch.tensor([21, 23])) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda _model_type, _modality: "<|vision_start|><|image_pad|><|vision_end|>", + ) + + class _QwenTokenizer(_Tokenizer): + def encode(self, text, add_special_tokens=True): + ids = { + "<|vision_start|>": 20, + "<|image_pad|>": 21, + "<|vision_end|>": 22, + "<|video_pad|>": 23, + } + return [ids.get(token, 0) for token in re.findall(r"<\|[^>]+\|>|.", text)] + + llm.tokenizer = _QwenTokenizer() + placeholder = "<|vision_start|><|image_pad|><|vision_end|>" + prompts = { + "missing": "describe", + "duplicate": placeholder * 2, + "extra_bare_pad": placeholder + "<|image_pad|>", + "unexpected_video_pad": placeholder + "<|video_pad|>", + } + decoder_creations = 0 + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + nonlocal decoder_creations + decoder_creations += 1 + return cls() + + async def async_load_bytes(self, data): + return np.zeros((1, 1, 3), dtype=np.uint8) + + def preprocess(inputs, sampling_params, disaggregated_params): + processor.multimodal_hashing_supported = True + return inputs + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(llm, "preprocess", preprocess) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + request = generation_pb2.GenerateRequest( + request_id="placeholder-validation", + model="image-model", + media=[ + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image") + ], + ) + + def set_input(prompt): + if input_kind == "prompt": + request.prompt = prompt + else: + request.token_ids.CopyFrom(generation_pb2.TokenIds(ids=llm.tokenizer.encode(prompt))) + + set_input(prompts[shape]) + context = FakeServicerContext() + with pytest.raises(AbortError): + async for _ in servicer.Generate(request, context): + pass + + assert context.abort_code == grpc.StatusCode.INVALID_ARGUMENT + assert "placeholder" in context.abort_details + assert processor.multimodal_hashing_supported is None + assert decoder_creations == 0 + assert llm.generate_kwargs is None + assert servicer.active_request_count() == 0 + + set_input(placeholder) + responses = [response async for response in servicer.Generate(request, FakeServicerContext())] + assert responses + assert decoder_creations == 1 + assert processor.multimodal_hashing_supported is True + assert servicer.active_request_count() == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["complete", "local_abort", "coordinator_abort", "rpc_cancel"]) +async def test_media_preprocessing_is_offloaded_and_abort_aware(image_llm, monkeypatch, outcome): + capability_llm, _ = image_llm + llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + loop = asyncio.get_running_loop() + serving_thread = threading.get_ident() + started = asyncio.Event() + finished = asyncio.Event() + release = threading.Event() + preprocessed = PreprocessedInputs(prompt_token_ids=[1, 2, 3]) + preprocess_calls = [] + submitted = [] + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return np.zeros((1, 1, 3), dtype=np.uint8) + + def preprocess(inputs, sampling_params, disaggregated_params): + assert threading.get_ident() != serving_thread + preprocess_calls.append(inputs) + loop.call_soon_threadsafe(started.set) + try: + assert release.wait(timeout=5), "test must release the preprocessing thread" + return preprocessed + finally: + loop.call_soon_threadsafe(finished.set) + + def generate_async(**kwargs): + assert threading.get_ident() == serving_thread + submitted.append(kwargs["inputs"]) + return llm.result_handle + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(llm, "preprocess", preprocess) + monkeypatch.setattr(llm, "generate_async", generate_async) + frontend = _Frontend() + servicer = OpenEngineInferenceServicer(llm, model="image-model", frontend=frontend) + request = generation_pb2.GenerateRequest( + request_id="preprocessing", + model="image-model", + prompt="", + media=[ + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image") + ], + ) + context = FakeServicerContext() + + async def collect(): + return [response async for response in servicer.Generate(request, context)] + + task = asyncio.create_task(collect()) + try: + await asyncio.wait_for(started.wait(), timeout=2) + text_request = generation_pb2.GenerateRequest( + request_id="concurrent-text", model="image-model", prompt="hello" + ) + + async def collect_text(): + return [ + response + async for response in servicer.Generate(text_request, FakeServicerContext()) + ] + + assert await asyncio.wait_for(collect_text(), timeout=1) + assert submitted == ["hello"] + if outcome == "complete": + release.set() + assert await asyncio.wait_for(task, timeout=2) + assert submitted[1] is preprocessed + elif outcome == "rpc_cancel": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + assert context.abort_code is None + else: + if outcome == "local_abort": + assert servicer.abort_request_by_id(request.request_id) + else: + assert frontend.reservations[0].abort() + with pytest.raises(AbortError): + await asyncio.wait_for(task, timeout=1) + assert context.abort_code == grpc.StatusCode.CANCELLED + finally: + release.set() + await asyncio.wait_for(finished.wait(), timeout=2) + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert len(preprocess_calls) == 1 + assert len(submitted) == (2 if outcome == "complete" else 1) + assert servicer.active_request_count() == 0 + if outcome != "complete": + monkeypatch.setattr(llm, "preprocess", lambda *args: preprocessed) + responses = [ + response async for response in servicer.Generate(request, FakeServicerContext()) + ] + assert responses + assert submitted[-1] is preprocessed + assert len(submitted) == 2 + assert servicer.active_request_count() == 0 + assert all(reservation.released for reservation in frontend.reservations) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error", "expected_code"), + [ + (ValueError("invalid processor input"), grpc.StatusCode.INVALID_ARGUMENT), + (TypeError("processor contract changed"), grpc.StatusCode.INTERNAL), + (RuntimeError("processor failed"), grpc.StatusCode.INTERNAL), + ], +) +async def test_preprocessing_errors_are_classified(image_llm, monkeypatch, error, expected_code): + capability_llm, _ = image_llm + llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + return np.zeros((1, 1, 3), dtype=np.uint8) + + def preprocess(*args): + raise error + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(llm, "preprocess", preprocess) + servicer = OpenEngineInferenceServicer(llm, model="image-model") + request = generation_pb2.GenerateRequest( + request_id="preprocess-error", + model="image-model", + prompt="", + media=[ + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image") + ], + ) + context = FakeServicerContext() + + with pytest.raises(AbortError): + async for _ in servicer.Generate(request, context): + pass + + assert context.abort_code == expected_code + assert llm.generate_kwargs is None + assert servicer.active_request_count() == 0 + + @pytest.mark.asyncio async def test_duplicate_request_id_is_reserved_while_media_loads(image_llm, monkeypatch): capability_llm, _ = image_llm @@ -558,8 +1051,9 @@ async def test_duplicate_request_id_is_reserved_while_media_loads(image_llm, mon capability_llm._hf_model_config, ) load_started = asyncio.Event() - release_load = asyncio.Event() + load_cancelled = asyncio.Event() load_count = 0 + block_load = True class _ImageIo: @classmethod @@ -567,11 +1061,16 @@ def create(cls, defaults, runtime): return cls() async def async_load_bytes(self, data): - nonlocal load_count + nonlocal block_load, load_count load_count += 1 + if not block_load: + return np.zeros((1, 1, 3), dtype=np.uint8) load_started.set() - await release_load.wait() - return np.zeros((1, 1, 3), dtype=np.uint8) + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + load_cancelled.set() + raise monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) servicer = OpenEngineInferenceServicer(llm, model="image-model") @@ -594,21 +1093,157 @@ async def collect(context): first = asyncio.create_task(collect(first_context)) await load_started.wait() duplicate_context = FakeServicerContext() - try: - with pytest.raises(AbortError): - await asyncio.wait_for(collect(duplicate_context), timeout=1) - assert duplicate_context.abort_code == grpc.StatusCode.ALREADY_EXISTS - assert load_count == 1 - assert servicer.active_request_count() == 1 - assert servicer.abort_request_by_id(request.request_id) - finally: - release_load.set() - first_outcome = await asyncio.gather(first, return_exceptions=True) + with pytest.raises(AbortError): + await asyncio.wait_for(collect(duplicate_context), timeout=1) + assert duplicate_context.abort_code == grpc.StatusCode.ALREADY_EXISTS + assert load_count == 1 + assert servicer.active_request_count() == 1 + assert servicer.abort_request_by_id(request.request_id) + first_outcome = await asyncio.wait_for(asyncio.gather(first, return_exceptions=True), timeout=1) assert isinstance(first_outcome[0], AbortError) assert first_context.abort_code == grpc.StatusCode.CANCELLED + assert load_cancelled.is_set() + assert llm.generate_kwargs is None + assert servicer.active_request_count() == 0 + + block_load = False + responses = await collect(FakeServicerContext()) + assert responses + assert servicer.active_request_count() == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel_path", ["coordinator", "rpc", "coordinator_and_rpc"]) +async def test_media_loading_stops_for_coordinator_abort_or_rpc_cancellation( + image_llm, monkeypatch, cancel_path +): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + load_started = asyncio.Event() + load_cancelled = asyncio.Event() + block_load = True + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + return cls() + + async def async_load_bytes(self, data): + if not block_load: + return np.zeros((1, 1, 3), dtype=np.uint8) + load_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + load_cancelled.set() + raise + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + frontend = _Frontend() + servicer = OpenEngineInferenceServicer(llm, model="image-model", frontend=frontend) + request = generation_pb2.GenerateRequest( + request_id=f"{cancel_path}-during-media", + model="image-model", + prompt="", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + async def collect(context): + return [response async for response in servicer.Generate(request, context)] + + context = FakeServicerContext() + task = asyncio.create_task(collect(context)) + await load_started.wait() + if cancel_path == "coordinator": + assert frontend.reservations[0].abort() + with pytest.raises(AbortError): + await asyncio.wait_for(task, timeout=1) + assert context.abort_code == grpc.StatusCode.CANCELLED + else: + if cancel_path == "coordinator_and_rpc": + assert frontend.reservations[0].abort() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + assert context.abort_code is None + + assert load_cancelled.is_set() assert llm.generate_kwargs is None assert servicer.active_request_count() == 0 + assert frontend.reservations[0].released + + block_load = False + responses = await collect(FakeServicerContext()) + assert responses + assert servicer.active_request_count() == 0 + assert len(frontend.reservations) == 2 + assert all(reservation.released for reservation in frontend.reservations) + + +@pytest.mark.asyncio +async def test_coordinator_abort_before_reservation_binding_prevents_media_loading( + image_llm, monkeypatch +): + capability_llm, _ = image_llm + llm = _ImageLlm( + capability_llm.input_processor, + capability_llm._hf_model_config, + ) + decoder_creations = 0 + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + nonlocal decoder_creations + decoder_creations += 1 + return cls() + + async def async_load_bytes(self, data): + return np.zeros((1, 1, 3), dtype=np.uint8) + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + frontend = _Frontend(abort_before_bind=True) + servicer = OpenEngineInferenceServicer(llm, model="image-model", frontend=frontend) + request = generation_pb2.GenerateRequest( + request_id="abort-before-bind", + model="image-model", + prompt="", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"image", + ) + ], + ) + + async def collect(context): + return [response async for response in servicer.Generate(request, context)] + + context = FakeServicerContext() + with pytest.raises(AbortError): + await collect(context) + + assert context.abort_code == grpc.StatusCode.CANCELLED + assert decoder_creations == 0 + assert llm.generate_kwargs is None + assert servicer.active_request_count() == 0 + assert frontend.reservations[0].released + + responses = await collect(FakeServicerContext()) + assert responses + assert decoder_creations == 1 + assert servicer.active_request_count() == 0 + assert len(frontend.reservations) == 2 + assert all(reservation.released for reservation in frontend.reservations) @pytest.mark.asyncio @@ -657,7 +1292,12 @@ async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): @pytest.mark.asyncio -async def test_generate_maps_transient_media_fetch_to_unavailable(image_llm, monkeypatch): +@pytest.mark.parametrize( + "message", ["connection timed out", "remote media server returned HTTP 503"] +) +async def test_generate_maps_transient_media_fetch_to_failed_precondition( + image_llm, monkeypatch, message +): capability_llm, _ = image_llm llm = _ImageLlm( capability_llm.input_processor, @@ -670,7 +1310,7 @@ def create(cls, defaults, runtime): return cls() async def async_load(self, source): - raise MediaFetchError("connection timed out") + raise MediaFetchError(message) monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) servicer = OpenEngineInferenceServicer(llm, model="image-model") @@ -678,7 +1318,7 @@ async def async_load(self, source): request = generation_pb2.GenerateRequest( request_id="remote-image", model="image-model", - prompt="image", + prompt="", media=[ generation_pb2.MediaItem( modality=generation_pb2.MODALITY_IMAGE, @@ -690,7 +1330,10 @@ async def async_load(self, source): with pytest.raises(AbortError): async for _ in servicer.Generate(request, context): pass - assert context.abort_code == grpc.StatusCode.UNAVAILABLE + assert context.abort_code == grpc.StatusCode.FAILED_PRECONDITION + assert message in context.abort_details + assert llm.generate_kwargs is None + assert servicer.active_request_count() == 0 @pytest.mark.asyncio diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index e97c9a708aac..8e0b83fe3b71 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -8,10 +8,12 @@ import pytest -pytest.importorskip("grpc", reason='gRPC runtime not installed (pip install "grpcio>=1.67.1,<2")') +grpc = pytest.importorskip( + "grpc", reason='gRPC runtime not installed (pip install "grpcio>=1.67.1,<2")' +) import torch # noqa: E402 -from conftest import FakeServicerContext # noqa: E402 +from conftest import AbortError, FakeServicerContext # noqa: E402 from utils.llm_data import llm_models_root # noqa: E402 from tensorrt_llm import LLM # noqa: E402 @@ -44,7 +46,8 @@ def vlm(): engine.shutdown() -def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm): +@pytest.mark.parametrize("input_kind", ["prompt", "token_ids"]) +def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm, input_kind): inference = OpenEngineInferenceServicer(vlm, MODEL_NAME) control = OpenEngineControlServicer(vlm, MODEL_NAME, inference) info = asyncio.run( @@ -65,10 +68,14 @@ def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm): image_path = llm_models_root(check=True) / "multimodals/test_data/seashore.png" image_bytes = image_path.read_bytes() data_uri = "data:image/png;base64," + base64.b64encode(image_bytes).decode() + request_input = ( + {"prompt": prompt} + if input_kind == "prompt" + else {"token_ids": generation_pb2.TokenIds(ids=prompt_token_ids)} + ) request = generation_pb2.GenerateRequest( - request_id="openengine-image-conformance", + request_id=f"openengine-image-conformance-{input_kind}", model=MODEL_NAME, - token_ids=generation_pb2.TokenIds(ids=prompt_token_ids), stopping=generation_pb2.StoppingOptions(max_tokens=4), media=[ generation_pb2.MediaItem( @@ -77,8 +84,31 @@ def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm): uuid="seashore", ) ], + **request_input, ) + malformed = generation_pb2.GenerateRequest() + malformed.CopyFrom(request) + malformed.request_id += "-extra-pad" + malformed_prompt = prompt.replace("Describe", "<|image_pad|>Describe", 1) + if input_kind == "prompt": + malformed.prompt = malformed_prompt + else: + malformed.token_ids.CopyFrom( + generation_pb2.TokenIds(ids=vlm.tokenizer.encode(malformed_prompt)) + ) + hashing_supported = vlm.input_processor.multimodal_hashing_supported + malformed_context = FakeServicerContext() + + async def reject_malformed(): + return [response async for response in inference.Generate(malformed, malformed_context)] + + with pytest.raises(AbortError): + asyncio.run(reject_malformed()) + assert malformed_context.abort_code == grpc.StatusCode.INVALID_ARGUMENT + assert vlm.input_processor.multimodal_hashing_supported is hashing_supported + assert inference.active_request_count() == 0 + async def collect(): return [response async for response in inference.Generate(request, FakeServicerContext())] diff --git a/tests/unittest/inputs/test_url_validation.py b/tests/unittest/inputs/test_url_validation.py index 5cf2391104a1..9eafe5d1f182 100644 --- a/tests/unittest/inputs/test_url_validation.py +++ b/tests/unittest/inputs/test_url_validation.py @@ -589,6 +589,45 @@ def test_rejects_private_url_before_request(self): self._run(_safe_aiohttp_get("http://169.254.169.254/latest/")) mock_session.assert_not_called() + @pytest.mark.asyncio + async def test_rejects_redirect_to_private_literal_before_second_request(self): + class _RedirectResponse: + status = 302 + headers = {"Location": "http://127.0.0.1/private"} + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + pass + + class _FakeSession: + def __init__(self): + self.urls = [] + + def get(self, url, **kwargs): + self.urls.append(url) + return _RedirectResponse() + + def resolve(host, *_args, **_kwargs): + return _dns("127.0.0.1") if host == "127.0.0.1" else PUBLIC_DNS + + session = _FakeSession() + with ( + patch( + "tensorrt_llm.inputs.media_io.socket.getaddrinfo", + side_effect=resolve, + ), + patch( + "tensorrt_llm.inputs.media_io._get_aiohttp_session", + new=AsyncMock(return_value=session), + ), + ): + with pytest.raises(RuntimeError, match="non-public"): + await _safe_aiohttp_get("http://example.com/image.jpg") + + assert session.urls == ["http://example.com/image.jpg"] + def test_rejects_private_address_from_connection_resolver(self): async def resolve(): resolver = _ValidatingResolver() From 2d3a2a2d57395529e4e7e485598314705e2c12c4 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Wed, 21 Oct 2026 16:15:44 -0700 Subject: [PATCH 31/41] fix(openengine): budget embedded video audio before allocation Signed-off-by: krishung5 --- .../grpc/openengine/request_mapping.py | 12 ++ tensorrt_llm/inputs/media_io.py | 154 +++++++++++------- .../grpc/openengine/test_multimodal.py | 117 ++++++++++++- 3 files changed, 224 insertions(+), 59 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 5302fc68912c..c8abf1ed27ee 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -52,6 +52,12 @@ def __init__(self, budget: "_DecodedMediaBudget", peak_bytes: int, retained_byte self._retained_bytes = retained_bytes self._active = True + def shrink(self, retained_bytes: int) -> None: + if not self._active or not 0 <= retained_bytes <= self._retained_bytes: + raise ValueError("invalid decoded-media reservation shrink") + self._budget._shrink(self._peak_bytes, self._retained_bytes, retained_bytes) + self._peak_bytes = self._retained_bytes = retained_bytes + def commit(self, retained_bytes: int) -> None: if not self._active: raise RuntimeError("decoded-media reservation is no longer active") @@ -111,6 +117,12 @@ def reserve(self, peak_bytes: int, retained_bytes: int) -> _DecodedMediaReservat self._reserved_retained += retained_bytes return _DecodedMediaReservation(self, peak_bytes, retained_bytes) + def _shrink(self, peak_bytes: int, reserved_retained: int, retained_bytes: int) -> None: + with self._condition: + self._used -= peak_bytes - retained_bytes + self._reserved_retained -= reserved_retained - retained_bytes + self._condition.notify_all() + def _settle(self, peak_bytes: int, reserved_retained: int, retained_bytes: int) -> None: with self._condition: self._used -= peak_bytes diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index f191c9751278..1fb79d49ec00 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -52,6 +52,8 @@ class MediaFetchError(RuntimeError): class _DecodedMemoryReservation(Protocol): + def shrink(self, retained_bytes: int) -> None: ... + def commit(self, retained_bytes: int) -> None: ... def release(self) -> None: ... @@ -463,12 +465,12 @@ def _audio_frame_to_array(frame, mono: bool) -> np.ndarray: return chunk -# TODO(TRTLLM-12001): Add unit tests for this. def extract_audio_from_video( source: Union[str, BytesIO], *, sr: Optional[float] = None, mono: bool = True, + decoded_memory_reserver: Optional[_DecodedMemoryReserver] = None, ) -> Tuple[np.ndarray, int]: """Extract the audio track from a video file using PyAV. @@ -476,6 +478,8 @@ def extract_audio_from_video( source: File path, URL, or BytesIO containing the video. sr: Target sample rate. If `None`, the native sample rate is kept. mono: If `True` (default), average channels to produce a mono waveform. + decoded_memory_reserver: Shared budget for retained chunks and the final + concatenation buffer. Codec and resampler working memory is excluded. Returns: `(waveform, sample_rate)` where *waveform* is a 1-D float32 @@ -494,51 +498,85 @@ def extract_audio_from_video( except ImportError: raise ImportError("PyAV is required for audio extraction from video but is not installed.") + chunks: List[np.ndarray] = [] + reservations: List[Tuple[_DecodedMemoryReservation, int]] = [] + concatenation_reservation: Optional[_DecodedMemoryReservation] = None + + def append_frame(frame) -> None: + if decoded_memory_reserver is not None: + channels = 1 if mono else len(frame.layout.channels) + retained_bytes = frame.samples * channels * np.dtype(np.float32).itemsize + reservation = decoded_memory_reserver(retained_bytes, retained_bytes) + reservations.append((reservation, retained_bytes)) + chunk = _audio_frame_to_array(frame, mono) + if decoded_memory_reserver is not None and ( + chunk.dtype != np.float32 or chunk.nbytes != retained_bytes + ): + raise ValueError("Decoded audio chunk does not match its float32 reservation") + chunks.append(chunk) + try: - with av.open(source) as container: - if not container.streams.audio: - raise ValueError("No audio stream found in the video.") - stream = container.streams.audio[0] - stream.thread_type = "AUTO" - native_sr = stream.rate - target_sr = int(sr) if sr is not None else native_sr - - chunks: List[np.ndarray] = [] - target_layout = "mono" if mono else stream.layout.name - needs_resampling = ( - not math.isclose(float(target_sr), float(native_sr), rel_tol=0.0, abs_tol=1e-6) - or stream.format.name != "fltp" - or stream.layout.name != target_layout - ) - resampler = ( - av.AudioResampler( - format="fltp", - layout=target_layout, - rate=target_sr, + try: + with av.open(source) as container: + if not container.streams.audio: + raise ValueError("No audio stream found in the video.") + stream = container.streams.audio[0] + stream.thread_type = "AUTO" + native_sr = stream.rate + target_sr = int(sr) if sr is not None else native_sr + + target_layout = "mono" if mono else stream.layout.name + needs_resampling = ( + not math.isclose(float(target_sr), float(native_sr), rel_tol=0.0, abs_tol=1e-6) + or stream.format.name != "fltp" + or stream.layout.name != target_layout ) - if needs_resampling - else None - ) - for frame in container.decode(stream): - if needs_resampling: - for out_frame in resampler.resample(frame): - chunks.append(_audio_frame_to_array(out_frame, mono)) - else: - chunks.append(_audio_frame_to_array(frame, mono)) - except ValueError: - raise - except Exception as e: - raise ValueError( - "Invalid or corrupted video data when extracting audio. " - "Ensure the input is a valid video file." - ) from e - - if not chunks: - raise ValueError("No audio frames decoded from the video.") + resampler = ( + av.AudioResampler( + format="fltp", + layout=target_layout, + rate=target_sr, + ) + if needs_resampling + else None + ) + for frame in container.decode(stream): + if needs_resampling: + for out_frame in resampler.resample(frame): + append_frame(out_frame) + else: + append_frame(frame) + except ValueError: + raise + except Exception as e: + raise ValueError( + "Invalid or corrupted video data when extracting audio. " + "Ensure the input is a valid video file." + ) from e - audio = np.concatenate(chunks, axis=-1).astype(np.float32, copy=False) + if not chunks: + raise ValueError("No audio frames decoded from the video.") - return audio, target_sr + if decoded_memory_reserver is not None: + concatenation_reservation = decoded_memory_reserver( + sum(chunk.nbytes for chunk in chunks), 0 + ) + audio = np.concatenate(chunks, axis=-1).astype(np.float32, copy=False) + chunks.clear() + if concatenation_reservation is not None: + concatenation_reservation.release() + concatenation_reservation = None + # The waveform replaces the chunks without changing retained-byte ownership. + for reservation, retained_bytes in reservations: + reservation.commit(retained_bytes) + + return audio, target_sr + finally: + chunks.clear() + if concatenation_reservation is not None: + concatenation_reservation.release() + for reservation, _ in reservations: + reservation.release() def _get_cv2(): @@ -973,9 +1011,18 @@ def _load_video_by_cv2( # audio extraction on Windows) and to avoid leaking descriptors. vidcap.release() + video_data = VideoData( + frames=loaded_frames, + metadata=metadata, + raw_bytes_hash=raw_bytes_hash, + ) try: - audio = None + frame_bytes = ( + _decoded_video_bytes(video_data) if decoded_memory_reservation is not None else 0 + ) if extract_audio: + if decoded_memory_reservation is not None: + decoded_memory_reservation.shrink(frame_bytes) # extract_audio_from_video accepts Union[str, BytesIO]. When the # in-memory cv2 path was taken, `video` is still the original # bytes parameter — wrap in a fresh BytesIO (the one consumed by @@ -987,30 +1034,21 @@ def _load_video_by_cv2( else video ) try: - audio_samples, audio_sample_rate = extract_audio_from_video(audio_source) - audio = AudioData(samples=audio_samples, sample_rate=audio_sample_rate) + audio_samples, audio_sample_rate = extract_audio_from_video( + audio_source, decoded_memory_reserver=decoded_memory_reserver + ) + video_data.audio = AudioData(samples=audio_samples, sample_rate=audio_sample_rate) except ValueError as e: if "No audio stream found" in str(e): logger.warning("Video has no audio track, skipping audio extraction.") else: raise + if decoded_memory_reservation is not None: + decoded_memory_reservation.commit(frame_bytes) except BaseException: if decoded_memory_reservation is not None: decoded_memory_reservation.release() raise - - video_data = VideoData( - frames=loaded_frames, - metadata=metadata, - audio=audio, - raw_bytes_hash=raw_bytes_hash, - ) - if decoded_memory_reservation is not None: - try: - decoded_memory_reservation.commit(_decoded_video_bytes(video_data)) - except Exception: - decoded_memory_reservation.release() - raise return video_data diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index ba30fc8d2de8..79e6f41bc5b6 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -4,9 +4,11 @@ """Unit tests for OpenEngine multimodal input mapping and discovery.""" import asyncio +import sys import threading from collections.abc import AsyncIterator, Sequence from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext from types import SimpleNamespace from typing import Any @@ -23,6 +25,7 @@ from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 +from tensorrt_llm.inputs import media_io # noqa: E402 from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData # noqa: E402 from tensorrt_llm.inputs.registry import ( # noqa: E402 @@ -278,12 +281,124 @@ def test_decoded_video_bytes_include_frames_and_audio(): def test_decoded_media_budget_settles_peak_to_retained_bytes(): - budget = request_mapping._DecodedMediaBudget(20) + budget = request_mapping._DecodedMediaBudget(20, peak_limit=20) first = budget.reserve(15, 5) first.commit(5) + assert (budget._used, budget._retained, budget._reserved_retained) == (5, 5, 0) second = budget.reserve(15, 5) second.release() + assert (budget._used, budget._retained, budget._reserved_retained) == (5, 5, 0) + + +@pytest.mark.parametrize( + "outcome", + [ + "success", + "resampled", + "oversized", + "decode_failure", + "concat_failure", + "no_audio", + "cancelled", + ], +) +def test_embedded_audio_reserves_chunks_and_concatenation(outcome, monkeypatch): + budget = request_mapping._DecodedMediaBudget(27 if outcome == "oversized" else 35) + converted = [] + concatenated = [] + grabs = iter([True]) + capture = SimpleNamespace( + isOpened=lambda: True, + get=lambda _key: 1, + grab=lambda: next(grabs, False), + retrieve=lambda _scratch: (True, np.zeros((1, 1, 3), dtype=np.uint8)), + release=lambda: None, + ) + monkeypatch.setattr( + media_io, + "_get_cv2", + lambda: SimpleNamespace( + VideoCapture=lambda _source: capture, + CAP_PROP_FRAME_COUNT=0, + CAP_PROP_FPS=1, + COLOR_BGR2RGB=2, + cvtColor=lambda _source, _code, dst: dst.fill(0), + ), + ) + + def audio_frame(index): + def to_ndarray(): + converted.append(index) + if outcome == "decode_failure" and index == 1: + raise RuntimeError("audio decode failed") + if outcome == "cancelled" and index == 0: + budget.cancel() + return np.full((1, 4), index, dtype=np.float32) + + return SimpleNamespace( + samples=4, + layout=SimpleNamespace(channels=("mono",)), + to_ndarray=to_ndarray, + ) + + stream = SimpleNamespace( + rate=16000, + format=SimpleNamespace(name="s16" if outcome == "resampled" else "fltp"), + layout=SimpleNamespace(name="mono"), + ) + container = SimpleNamespace( + streams=SimpleNamespace(audio=[] if outcome == "no_audio" else [stream]), + decode=lambda _stream: [audio_frame(0), audio_frame(1)], + ) + monkeypatch.setitem( + sys.modules, + "av", + SimpleNamespace( + open=lambda _source: nullcontext(container), + AudioResampler=lambda **_kwargs: SimpleNamespace(resample=lambda frame: [frame]), + ), + ) + monkeypatch.setenv("TRTLLM_ENABLE_PYAV", "1") + concatenate = np.concatenate + + def checked_concatenate(chunks, **kwargs): + concatenated.append(True) + assert (budget._used, budget._retained, budget._reserved_retained) == (67, 0, 35) + if outcome == "concat_failure": + raise MemoryError("audio concatenation failed") + return concatenate(chunks, **kwargs) + + monkeypatch.setattr(media_io.np, "concatenate", checked_concatenate) + + def load(): + return media_io._load_video_by_cv2( + "video.mp4", + fps=-1, + format="np", + extract_audio=True, + decoded_memory_reserver=budget.reserve, + ) + + if outcome in ("success", "resampled", "no_audio"): + video = load() + retained = media_io._decoded_video_bytes(video) + assert retained == (3 if outcome == "no_audio" else 35) + assert (budget._used, budget._retained, budget._reserved_retained) == ( + retained, + retained, + 0, + ) + if outcome != "no_audio": + np.testing.assert_array_equal(video.audio.samples, [0] * 4 + [1] * 4) + assert video.audio.sample_rate == 16000 + else: + with pytest.raises(MemoryError if outcome == "concat_failure" else ValueError): + load() + assert (budget._used, budget._retained, budget._reserved_retained) == (0, 0, 0) + if outcome in ("oversized", "cancelled"): + assert converted == [0] + assert not concatenated def test_video_retained_limit_does_not_reject_1080p_conversion_peak(): From 8646ea88f85148c1a1bc6116c147b7941ee43af8 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Wed, 21 Oct 2026 16:53:16 -0700 Subject: [PATCH 32/41] fix(openengine): bound audio decoding and checkpoint capabilities Signed-off-by: krishung5 --- tensorrt_llm/grpc/openengine/capabilities.py | 3 + tensorrt_llm/grpc/openengine/control.py | 6 + .../grpc/openengine/request_mapping.py | 8 +- tensorrt_llm/inputs/media_io.py | 44 ++++-- .../unittest/grpc/openengine/test_control.py | 18 +++ .../grpc/openengine/test_multimodal.py | 138 +++++++++++++++++- .../grpc/openengine/test_openengine.py | 14 ++ 7 files changed, 218 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 532096d019ed..4dc90bd77690 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -94,6 +94,9 @@ def resolve_media_capabilities( for modality in implemented_modalities if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) } + # Gemma4 checkpoints may omit the audio encoder despite sharing a model type. + if model_type == "gemma4" and getattr(input_processor.config, "audio_config", None) is None: + placeholders.pop("audio", None) # Fail startup before Control advertises a malformed template. for modality, template in placeholders.items(): try: diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index 85740f505f45..873e9714492c 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -477,6 +477,12 @@ async def _inference_probe(self, model: str) -> lifecycle_pb2.HealthCheck: async def _run_inference_probe(self, model: str) -> lifecycle_pb2.HealthCheck: from tensorrt_llm.sampling_params import SamplingParams + if not token_ids_input_supported(self._llm): + return lifecycle_pb2.HealthCheck( + name="inference_probe", + state=lifecycle_pb2.HEALTH_STATE_READY, + message="probe not applicable: model does not accept token-ID input", + ) handle = None reservation = None # Counted by GetLoad and reachable by Abort(all_requests): a probe diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 306680dbd5a9..14fd66157b7b 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -475,8 +475,10 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: return load -def _audio_loader_factory(llm: Any, _budget: _DecodedMediaBudget) -> _MediaLoader: - audio_io = _media_io_for_model(llm, "audio", AudioMediaIO) +def _audio_loader_factory(llm: Any, budget: _DecodedMediaBudget) -> _MediaLoader: + audio_io = _media_io_for_model( + llm, "audio", AudioMediaIO, {"decoded_memory_reserver": budget.reserve} + ) async def load(item: generation_pb2.MediaItem, index: int) -> Any: return await _load_media_item(item, audio_io, index, "audio") @@ -535,7 +537,7 @@ def _video_runtime_available() -> bool: reserves_decoded_bytes=True, ), generation_pb2.MODALITY_AUDIO: _MediaHandler( - "audio", _audio_loader_factory, _decoded_audio_bytes + "audio", _audio_loader_factory, _decoded_audio_bytes, reserves_decoded_bytes=True ), } diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 9d7bf1baec9b..c67bf89eb7bb 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -1300,29 +1300,55 @@ class AudioMediaIO(BaseMediaIO[Tuple[np.ndarray, int]]): _data_uri_mime_prefix: ClassVar[Optional[str]] = "audio/" - def __init__(self, **kwargs) -> None: - # AudioMediaIO has no configurable parameters. An explicit __init__ - # is needed so that BaseMediaIO.create() does not crash when kwargs - # are present (e.g. via --media_io_kwargs or the per-request API). + def __init__( + self, decoded_memory_reserver: Optional[_DecodedMemoryReserver] = None, **kwargs + ) -> None: + self._decoded_memory_reserver = decoded_memory_reserver if kwargs: logger.warning( - "AudioMediaIO received unexpected kwargs %s — audio loading " - "has no configurable parameters; kwargs will be ignored.", + "AudioMediaIO received unexpected kwargs %s; kwargs will be ignored.", sorted(kwargs), ) def load_bytes(self, data: bytes) -> Tuple[np.ndarray, int]: - return soundfile.read(BytesIO(data)) + return self._read(BytesIO(data)) def load_base64(self, media_type: str, data: str) -> Tuple[np.ndarray, int]: - return soundfile.read(BytesIO(base64.b64decode(data))) + return self.load_bytes(base64.b64decode(data)) def load_file(self, url: str) -> Tuple[np.ndarray, int]: # Strip `file://` and unquote the path for `file:` URIs; pass # empty-scheme URLs through verbatim. parsed = urlparse(url) path = _normalize_file_uri(url) if parsed.scheme == "file" else url - return soundfile.read(path) + return self._read(path) + + def _read(self, source: Union[BytesIO, str]) -> Tuple[np.ndarray, int]: + if self._decoded_memory_reserver is None: + return soundfile.read(source) + + reservation = None + try: + # Read metadata and samples from the same handle. Reserve the float64 + # output before allocation; libsndfile's internal buffers are excluded. + with soundfile.SoundFile(source) as audio: + frames, channels = audio.frames, audio.channels + if frames < 0 or channels < 1: + raise ValueError("invalid audio frame or channel count") + decoded_bytes = frames * channels * np.dtype(np.float64).itemsize + reservation = self._decoded_memory_reserver(decoded_bytes, decoded_bytes) + shape = (frames,) if channels == 1 else (frames, channels) + samples = np.empty(shape, dtype=np.float64) + decoded = audio.read(frames=frames, out=samples) + if decoded.shape != samples.shape: + raise ValueError("audio decoded fewer frames than its header declares") + sample_rate = audio.samplerate + reservation.commit(samples.nbytes) + return samples, sample_rate + finally: + samples = decoded = None + if reservation is not None: + reservation.release() class VideoMediaIO(BaseMediaIO[VideoData]): diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 82c592c6709a..754130cf0719 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -426,6 +426,24 @@ async def test_abort_all_fails_the_rpc_when_any_abort_is_refused(): assert "still running" in context.abort_details +@pytest.mark.asyncio +async def test_health_probe_is_not_applicable_to_text_only_input_models(): + llm = _llm(_supports_token_ids_input=False) + + def unexpected_generate(*args, **kwargs): + pytest.fail("a token-only probe must not reach a prompt-only model") + + llm.generate_async = unexpected_generate + servicer = OpenEngineControlServicer(llm, MODEL, _inference(), kv_transfer_backend="") + response = await servicer.Health( + lifecycle_pb2.HealthRequest(include_inference_probe=True), FakeServicerContext() + ) + probe = next(check for check in response.checks if check.name == "inference_probe") + assert probe.state == lifecycle_pb2.HEALTH_STATE_READY + assert "not applicable" in probe.message + assert response.state == lifecycle_pb2.HEALTH_STATE_READY + + @pytest.mark.asyncio async def test_health_inference_probe_is_bounded_and_aborts_a_stuck_request(): """A wedged engine must yield DEGRADED, not a hung RPC or a leaked request.""" diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 5f3b2e7d3b39..1891ed4772ad 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -9,6 +9,7 @@ from collections.abc import AsyncIterator, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext +from io import BytesIO from types import SimpleNamespace from typing import Any @@ -361,6 +362,137 @@ def test_decoded_audio_bytes_count_samples_and_reject_unknown_types(): request_mapping._decoded_audio_bytes(object()) +@pytest.mark.parametrize("channels", [1, 2]) +@pytest.mark.parametrize("source", ["bytes", "base64", "file"]) +def test_audio_reserves_before_decoding_and_preserves_samples( + channels, source, tmp_path, monkeypatch +): + import base64 + + import soundfile + + samples = np.arange(32 * channels, dtype=np.float64).reshape(32, channels) / 128 + encoded = BytesIO() + soundfile.write(encoded, samples, 16000, format="FLAC") + data = encoded.getvalue() + expected, _ = soundfile.read(BytesIO(data)) + budget = request_mapping._DecodedMediaBudget(samples.nbytes) + original_read = soundfile.SoundFile.read + + def checked_read(audio, *args, **kwargs): + assert (budget._used, budget._reserved_retained) == (samples.nbytes, samples.nbytes) + assert kwargs["out"].nbytes == samples.nbytes + return original_read(audio, *args, **kwargs) + + monkeypatch.setattr(soundfile.SoundFile, "read", checked_read) + loader = media_io.AudioMediaIO(decoded_memory_reserver=budget.reserve) + if source == "bytes": + decoded, rate = loader.load_bytes(data) + elif source == "base64": + decoded, rate = loader.load_base64("audio/flac", base64.b64encode(data).decode()) + else: + path = tmp_path / "audio.flac" + path.write_bytes(data) + decoded, rate = loader.load_file(path.as_uri()) + np.testing.assert_array_equal(decoded, expected) + assert rate == 16000 + assert (budget._used, budget._retained, budget._reserved_retained) == ( + samples.nbytes, + samples.nbytes, + 0, + ) + + +@pytest.mark.parametrize("outcome", ["oversized", "failure", "truncated", "cancelled"]) +def test_audio_budget_rejects_before_allocation_and_cleans_up(outcome, monkeypatch): + import soundfile + + encoded = BytesIO() + soundfile.write(encoded, np.zeros((32, 2)), 16000, format="FLAC") + budget = request_mapping._DecodedMediaBudget(511 if outcome == "oversized" else 512) + loader = media_io.AudioMediaIO(decoded_memory_reserver=budget.reserve) + + def read(audio, *args, **kwargs): + if outcome == "oversized": + pytest.fail("oversized audio must be rejected before decoding") + if outcome == "failure": + raise RuntimeError("decode failed") + if outcome == "cancelled": + raise asyncio.CancelledError() + return kwargs["out"][:0] + + monkeypatch.setattr(soundfile.SoundFile, "read", read) + if outcome == "oversized": + monkeypatch.setattr( + np, "empty", lambda *args, **kwargs: pytest.fail("allocated oversized audio") + ) + error = { + "oversized": ValueError, + "failure": RuntimeError, + "truncated": ValueError, + "cancelled": asyncio.CancelledError, + }[outcome] + with pytest.raises(error): + loader.load_bytes(encoded.getvalue()) + assert (budget._used, budget._retained, budget._reserved_retained) == (0, 0, 0) + + +@pytest.mark.asyncio +async def test_audio_request_uses_one_shared_budget_without_double_charging( + image_video_audio_llm, monkeypatch +): + import soundfile + + encoded = BytesIO() + soundfile.write(encoded, np.zeros(32), 16000, format="FLAC") + monkeypatch.setattr(request_mapping, "_MAX_TOTAL_DECODED_MEDIA_BYTES", 256) + item = generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, raw_bytes=encoded.getvalue(), mime_type="audio/flac" + ) + request = generation_pb2.GenerateRequest(prompt="transcribe", media=[item]) + mapped = await request_mapping._input_from_request(request, image_video_audio_llm) + assert mapped["multi_modal_data"]["audio"][0][0].nbytes == 256 + request.media.append(item) + with pytest.raises(ValueError, match="decoded media exceeds"): + await request_mapping._input_from_request(request, image_video_audio_llm) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("audio_encoder", [False, True]) +async def test_gemma4_audio_capability_requires_checkpoint_encoder(audio_encoder, monkeypatch): + class GemmaConfig: + model_type = "gemma4" + audio_config = object() if audio_encoder else None + + llm = SimpleNamespace(input_processor=_ImageInputProcessor(GemmaConfig())) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: modality in {"image", "audio"}, + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + capabilities = request_mapping.media_capabilities_for_model(llm) + assert "image" in capabilities.modalities + assert ("audio" in capabilities.modalities) == audio_encoder + if not audio_encoder: + request = generation_pb2.GenerateRequest( + prompt="transcribe", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + raw_bytes=b"unused", + mime_type="audio/wav", + ) + ], + ) + with pytest.raises(UnsupportedFeatureError, match="audio"): + await request_mapping._input_from_request(request, llm, capabilities) + + def test_decoded_media_budget_settles_peak_to_retained_bytes(): budget = request_mapping._DecodedMediaBudget(20, peak_limit=20) @@ -748,7 +880,11 @@ async def async_load(self, source): assert set(video_runtime) == {"decoded_memory_reserver"} assert callable(video_runtime["decoded_memory_reserver"]) assert created_with[1] == ("image", {"format": "np"}, None) - assert created_with[2] == ("audio", None, None) + audio_modality, audio_defaults, audio_runtime = created_with[2] + assert audio_modality == "audio" + assert audio_defaults is None + assert set(audio_runtime) == {"decoded_memory_reserver"} + assert callable(audio_runtime["decoded_memory_reserver"]) assert audio_sources == ["https://example.com/audio.wav"] assert mapped["multi_modal_data"]["audio"][0][1] == 16000 assert [video.metadata["source"] for video in mapped["multi_modal_data"]["video"]] == [ diff --git a/tests/unittest/grpc/openengine/test_openengine.py b/tests/unittest/grpc/openengine/test_openengine.py index 0efb4c0cc886..d19928aa31f0 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -17,6 +17,7 @@ import asyncio import gc from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest @@ -48,6 +49,19 @@ def test_format_bind_address_brackets_ipv6() -> None: assert _format_bind_address("[::1]", 8000) == "[::1]:8000" +def test_stop_closes_media_session_when_grpc_stop_fails() -> None: + async def exercise_stop() -> None: + server = OpenEngineServer.__new__(OpenEngineServer) + server._server = SimpleNamespace(stop=AsyncMock(side_effect=RuntimeError("stop failed"))) + media_session = await _get_aiohttp_session() + with pytest.raises(RuntimeError, match="stop failed"): + await server.stop(grace=0) + server._server.stop.assert_awaited_once_with(grace=0) + assert media_session.closed + + asyncio.run(exercise_stop()) + + def test_openengine_server_serves_the_control_contract() -> None: """The server binds, attaches Control, answers, and shuts down cleanly. From c63d619da8f40094d365fa86d8c0f15a5028ae62 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 22 Oct 2026 10:12:15 -0700 Subject: [PATCH 33/41] fix(shared-tensor): avoid global strategy races and reclaim exports Signed-off-by: krishung5 --- .../_torch/shared_tensor/shared_tensor.py | 50 +++++---- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../unittest/_torch/test_shared_tensor_cpu.py | 102 ++++++++++++++++++ 3 files changed, 132 insertions(+), 21 deletions(-) create mode 100644 tests/unittest/_torch/test_shared_tensor_cpu.py diff --git a/tensorrt_llm/_torch/shared_tensor/shared_tensor.py b/tensorrt_llm/_torch/shared_tensor/shared_tensor.py index 24666511249d..b6423c482669 100644 --- a/tensorrt_llm/_torch/shared_tensor/shared_tensor.py +++ b/tensorrt_llm/_torch/shared_tensor/shared_tensor.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import base64 import itertools import logging @@ -6,12 +9,7 @@ from typing import Any, Callable, Dict, Optional, Tuple import torch -from torch.multiprocessing import get_sharing_strategy, set_sharing_strategy -from torch.multiprocessing.reductions import (rebuild_cuda_tensor, - rebuild_storage_filename, - rebuild_tensor, - rebuild_typed_storage, - reduce_storage, reduce_tensor) +from torch.multiprocessing import reductions logger = logging.getLogger(__name__) @@ -88,8 +86,8 @@ class _SharedTensorRebuildMethodRegistry: def initialize(cls): """Initialize the registry with common rebuild methods.""" # Register common methods with fixed keys - cls._registry[cls.REBUILD_CUDA] = rebuild_cuda_tensor - cls._registry[cls.REBUILD_CPU] = rebuild_tensor + cls._registry[cls.REBUILD_CUDA] = reductions.rebuild_cuda_tensor + cls._registry[cls.REBUILD_CPU] = reductions.rebuild_tensor @classmethod def register(cls, method: Callable) -> int: @@ -101,9 +99,9 @@ def register(cls, method: Callable) -> int: Returns: The numeric key assigned to the method """ - if method == rebuild_cuda_tensor: + if method == reductions.rebuild_cuda_tensor: return cls.REBUILD_CUDA - if method == rebuild_tensor: + if method == reductions.rebuild_tensor: return cls.REBUILD_CPU raise NotImplementedError("Other rebuild methods are not supported yet") @@ -326,9 +324,9 @@ def dict_to_cpu_handle(tensor_info: Dict[str, Any]) -> Tuple: storage_handle, tensor_info['storage_size'], str_to_torch_dtype( tensor_info['storage_dtype'])) - storage = rebuild_storage_filename(*storage_metadata) + storage = reductions.rebuild_storage_filename(*storage_metadata) if not isinstance(storage, torch.storage.TypedStorage): - storage = rebuild_typed_storage( + storage = reductions.rebuild_typed_storage( storage, str_to_torch_dtype(tensor_info['storage_dtype'])) meta_data = (tensor_info['tensor_storage_offset'], @@ -371,7 +369,7 @@ def from_tensor(cls, "local_token": token, "producer_pid": os.getpid() }) - rebuild_method, tensor_handle = reduce_tensor(tensor) + rebuild_method, tensor_handle = reductions.reduce_tensor(tensor) method_key = _SharedTensorRebuildMethodRegistry.register(rebuild_method) return cls(method_key, tensor_handle) @@ -433,6 +431,18 @@ def get_local_view(self) -> torch.Tensor: self.method_key) return rebuild_method(*self.tensor_handle) + @classmethod + def release_dict(cls, tensor_info: Dict[str, Any]) -> None: + """Consume an exported handle that will not be delivered to a consumer.""" + if tensor_info[ + 'method_key'] == _SharedTensorRebuildMethodRegistry.REBUILD_CUDA: + torch.UntypedStorage._release_ipc_counter( + base64.b64decode(tensor_info['ref_counter_handle']), + tensor_info['ref_counter_offset'], + device=tensor_info['storage_device']) + else: + cls.from_dict(tensor_info).get_local_view() + def dump_to_dict(self) -> Dict[str, Any]: """Convert this container to a dictionary for direct IPC (Producer side). @@ -452,18 +462,16 @@ def dump_to_dict(self) -> Dict[str, Any]: tensor_dict = SharedTensorContainer.cuda_handle_to_dict( self.tensor_handle) elif self.method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CPU: - sharing_strategy = get_sharing_strategy() - # Here we use file_system sharing strategy to make it serializable between two non-python independent processes - set_sharing_strategy("file_system") storage = self.tensor_handle[1] meta_data = self.tensor_handle[2] - storage_handle = reduce_storage(storage) - # restore the original sharing strategy - set_sharing_strategy(sharing_strategy) - # exclude the first element which is the type of the storage - storage_metadata = storage_handle[-1][1:] + # Use filename sharing directly without changing PyTorch's process-wide strategy. + storage_metadata = storage._share_filename_cpu_() + ( + storage.dtype, ) tensor_dict = SharedTensorContainer.cpu_handle_to_dict( meta_data, storage_metadata) + reductions.shared_cache[ + storage_metadata[1]] = reductions.StorageWeakRef(storage) + storage._shared_incref() else: raise ValueError(f"Unsupported tensor device: {self.method_key}") diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 0ab9365390bb..4cd741dd2810 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -175,6 +175,7 @@ l0_cpu: - unittest/executor/test_ray_stub.py - unittest/executor/test_generation_request_array_prompt.py - unittest/grpc/openengine/test_coordinator.py + - unittest/_torch/test_shared_tensor_cpu.py - unittest/grpc/openengine/test_frontend_lifecycle.py - unittest/grpc/openengine/test_control.py - unittest/grpc/openengine/test_multimodal.py diff --git a/tests/unittest/_torch/test_shared_tensor_cpu.py b/tests/unittest/_torch/test_shared_tensor_cpu.py new file mode 100644 index 000000000000..b8070b82a4ab --- /dev/null +++ b/tests/unittest/_torch/test_shared_tensor_cpu.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU regressions for independent-process tensor exports.""" + +import base64 +import gc +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest +import torch +from torch.multiprocessing import get_sharing_strategy + +from tensorrt_llm._torch.shared_tensor import SharedTensorContainer + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("shape", [(8,), (0,), (3, 4)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.int64]) +def test_filename_roundtrip_preserves_storage_views(shape, dtype): + source = torch.arange(torch.Size(shape).numel(), dtype=dtype).reshape(shape) + if source.ndim == 2: + source = source.T + strategy = get_sharing_strategy() + exported = SharedTensorContainer.from_tensor(source).dump_to_dict() + restored = SharedTensorContainer.from_dict(exported).get_local_view() + torch.testing.assert_close(restored, source) + assert restored.stride() == source.stride() + assert get_sharing_strategy() == strategy + + +def test_concurrent_filename_exports_do_not_change_global_strategy(monkeypatch): + strategy = get_sharing_strategy() + barrier = threading.Barrier(3, timeout=5) + original = torch.storage.TypedStorage._share_filename_cpu_ + + def share(storage): + assert get_sharing_strategy() == strategy + barrier.wait() + result = original(storage) + barrier.wait() + return result + + monkeypatch.setattr(torch.storage.TypedStorage, "_share_filename_cpu_", share) + + def export(value): + source = torch.full((1024,), value) + return SharedTensorContainer.from_tensor(source).dump_to_dict() + + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(export, value) for value in (1, 2)] + barrier.wait() + assert get_sharing_strategy() == strategy + barrier.wait() + exports = [future.result(timeout=5) for future in futures] + for value, exported in zip((1, 2), exports): + restored = SharedTensorContainer.from_dict(exported).get_local_view() + torch.testing.assert_close(restored, torch.full((1024,), value)) + + +def test_export_failure_does_not_change_global_strategy(monkeypatch): + strategy = get_sharing_strategy() + + def fail(*args): + raise RuntimeError("serialization failed") + + monkeypatch.setattr(SharedTensorContainer, "cpu_handle_to_dict", fail) + with pytest.raises(RuntimeError, match="serialization failed"): + SharedTensorContainer.from_tensor(torch.ones(8)).dump_to_dict() + assert get_sharing_strategy() == strategy + + +def test_release_unconsumed_export_reclaims_shared_memory(): + exported = SharedTensorContainer.from_tensor(torch.ones(1024 * 1024)).dump_to_dict() + path = Path("/dev/shm") / base64.b64decode(exported["storage_handle"]).decode().lstrip("/") + assert path.exists() + gc.collect() + assert path.exists() + SharedTensorContainer.release_dict(exported) + gc.collect() + assert not path.exists() + + +def test_release_cuda_export_does_not_reopen_producer_ipc(monkeypatch): + calls = [] + monkeypatch.setattr( + torch.UntypedStorage, + "_release_ipc_counter", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + SharedTensorContainer.release_dict( + { + "method_key": 1, + "ref_counter_handle": base64.b64encode(b"counter"), + "ref_counter_offset": 2, + "storage_device": 0, + } + ) + assert calls == [((b"counter", 2), {"device": 0})] From 67fb3b729ae5e90e9f8917ce417ff121c676e6ee Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 22 Oct 2026 10:12:15 -0700 Subject: [PATCH 34/41] fix(openengine): own preprocessing exports and isolate media work Signed-off-by: krishung5 --- .../grpc/openengine/request_mapping.py | 38 ++- tensorrt_llm/grpc/openengine/server.py | 24 +- tensorrt_llm/grpc/openengine/servicer.py | 75 +++++- tensorrt_llm/inputs/media_io.py | 19 +- tensorrt_llm/inputs/multimodal.py | 12 +- tensorrt_llm/inputs/registry.py | 11 + .../grpc/openengine/test_multimodal.py | 237 +++++++++++++++++- .../grpc/openengine/test_openengine.py | 27 ++ 8 files changed, 420 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index c55fbd075495..393c144bd6f9 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -5,7 +5,7 @@ import asyncio from collections import Counter -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from typing import Any from urllib.parse import urlparse @@ -433,10 +433,38 @@ def _prepare_input_from_request( return _PreparedInput(inputs, media, capabilities) +def _count_token_sequences( + prompt_ids: Sequence[int], sequences: Iterable[tuple[int, ...]] +) -> Counter[tuple[int, ...]]: + """Count overlapping markers in one traversal, sharing common prefixes.""" + trie: dict = {} + for sequence in sequences: + node = trie + for token_id in sequence: + node = node.setdefault(token_id, {}) + node[None] = sequence + counts: Counter[tuple[int, ...]] = Counter() + for start, token_id in enumerate(prompt_ids): + node = trie.get(token_id) + index = start + 1 + while node is not None: + sequence = node.get(None) + if sequence is not None: + counts[sequence] += 1 + if index == len(prompt_ids): + break + node = node.get(prompt_ids[index]) + index += 1 + return counts + + def _validate_media_placeholders(prepared: _PreparedInput, llm: Any) -> None: """Reject placeholder mismatches before probing the stateful MM processor.""" if not prepared.media or prepared.capabilities is None: return + processor = llm.input_processor + if not processor.requires_input_placeholders(tokenized=not isinstance(prepared.inputs, str)): + return tokenizer = llm.tokenizer prompt_ids = ( tokenizer.encode(prepared.inputs, add_special_tokens=False) @@ -453,19 +481,15 @@ def _validate_media_placeholders(prepared: _PreparedInput, llm: Any) -> None: raise ValueError(f"{handler.name} placeholder must encode to at least one token") expected_sequences[sequence] += 1 + actual_sequences = _count_token_sequences(prompt_ids, expected_sequences) for sequence, expected in expected_sequences.items(): - width = len(sequence) - actual = sum( - tuple(prompt_ids[index : index + width]) == sequence - for index in range(len(prompt_ids) - width + 1) - ) + actual = actual_sequences[sequence] if actual != expected: raise ValueError( f"media placeholder count does not match media items: expected {expected}, " f"got {actual} for placeholder token IDs {list(sequence)}" ) - processor = llm.input_processor mm_ids = processor.get_mm_token_ids() if mm_ids is not None: special_ids = processor.get_mm_special_token_ids() diff --git a/tensorrt_llm/grpc/openengine/server.py b/tensorrt_llm/grpc/openengine/server.py index 546adc9b7bb4..8151f187d9b6 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -13,6 +13,7 @@ import threading import uuid from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from typing import Any import click @@ -112,6 +113,12 @@ def __init__( options = list(_SERVER_OPTIONS) options.append(("grpc.so_reuseport", 1 if frontend is not None else 0)) self._server = grpc.aio.server(options=options) + self._input_executor = ThreadPoolExecutor( + max_workers=4, thread_name_prefix="openengine_input" + ) + self._media_executor = ThreadPoolExecutor( + max_workers=8, thread_name_prefix="openengine_media" + ) kv_transfer_backend = _kv_transfer_backend(llm) inference = OpenEngineInferenceServicer( llm, @@ -119,6 +126,8 @@ def __init__( kv_transfer_backend=kv_transfer_backend, frontend=frontend, subagent_affinity_auth_key=subagent_affinity_auth_key, + input_executor=self._input_executor, + media_executor=self._media_executor, ) openengine_pb2_grpc.add_InferenceServicer_to_server(inference, self._server) # Control shares the inference servicer's in-flight request table so @@ -147,7 +156,12 @@ def __init__( ) # grpc.aio raises RuntimeError itself when the bind fails, naming the # address, so there is no zero return to check for. - bound_port = self._server.add_insecure_port(bind_address) + try: + bound_port = self._server.add_insecure_port(bind_address) + except BaseException: + self._input_executor.shutdown(wait=True, cancel_futures=True) + self._media_executor.shutdown(wait=True, cancel_futures=True) + raise if port == 0: self.port = bound_port @@ -170,7 +184,13 @@ async def stop(self, grace: float = 5.0) -> None: Args: grace: Maximum time in seconds to allow active RPCs to finish. """ - await self._server.stop(grace=grace) + try: + await self._server.stop(grace=grace) + finally: + await asyncio.gather( + asyncio.to_thread(self._input_executor.shutdown, wait=True, cancel_futures=True), + asyncio.to_thread(self._media_executor.shutdown, wait=True, cancel_futures=True), + ) logger.info("OpenEngine server stopped") async def wait_for_termination(self) -> None: diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index 6f04a77cf5a6..71a8e5fd3963 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -4,13 +4,18 @@ """OpenEngine Inference service adapter for the TensorRT-LLM LLM API.""" import asyncio -from collections.abc import AsyncIterator +import threading +from collections.abc import AsyncIterator, Callable +from concurrent.futures import Executor from contextlib import aclosing +from contextvars import copy_context +from functools import partial from typing import Any import grpc from tensorrt_llm.executor.request import DEFAULT_REQUEST_PRIORITY +from tensorrt_llm.inputs.media_io import BaseMediaIO from tensorrt_llm.llmapi.llm import LLM from tensorrt_llm.logger import logger from tensorrt_llm.scheduling_params import SchedulingParams @@ -46,6 +51,43 @@ ) +class _PreprocessingWork: + """Own exports until submission, including results from abandoned worker threads.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._abandoned = False + self._result = None + + @staticmethod + def _release(result: Any) -> None: + params = getattr(result, "multimodal_params", None) + if params is not None: + try: + params.release_handles() + except Exception as error: + logger.error(f"Failed to reclaim abandoned preprocessing exports: {error}") + + def run(self, preprocess: Callable[[], Any]) -> Any: + result = preprocess() + with self._lock: + if not self._abandoned: + self._result = result + return result + self._release(result) + return None + + def transfer(self) -> None: + with self._lock: + self._result = None + + def abandon(self) -> None: + with self._lock: + self._abandoned = True + result, self._result = self._result, None + self._release(result) + + class _LocalRequest: """Abort-aware local request state spanning mapping and generation.""" @@ -53,6 +95,7 @@ def __init__(self) -> None: self.mapping_task: asyncio.Task[Any] | None = None self.engine_handle: Any = None self.aborted = False + self.preprocessing: _PreprocessingWork | None = None def bind_mapping(self, task: asyncio.Task[Any]) -> None: self.mapping_task = task @@ -65,6 +108,8 @@ def clear_mapping(self, task: asyncio.Task[Any]) -> None: def bind_engine(self, handle: Any) -> None: self.engine_handle = handle + if self.preprocessing is not None: + self.preprocessing.transfer() if self.aborted: handle.abort() @@ -86,8 +131,12 @@ def __init__( kv_transfer_backend: str = "", frontend: FrontendClient | None = None, subagent_affinity_auth_key: str | None = None, + input_executor: Executor | None = None, + media_executor: Executor | None = None, ) -> None: self._frontend = frontend + self._input_executor = input_executor + self._media_executor = media_executor self._llm = llm self._model = model # Informational label placed in the KvSessionRef of PrefillReady events so @@ -189,6 +238,8 @@ async def Generate( except CoordinationError as error: await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) finally: + if local_request is not None and local_request.preprocessing is not None: + local_request.preprocessing.abandon() if registration is not None: registration.release() if reservation is not None: @@ -278,12 +329,26 @@ async def _generate( ) async def prepare_inputs() -> Any: + loop = asyncio.get_running_loop() if prepared_input.media: - await asyncio.to_thread(_validate_media_placeholders, prepared_input, self._llm) - inputs = await _load_prepared_input(prepared_input, self._llm) + await loop.run_in_executor( + self._input_executor, + copy_context().run, + _validate_media_placeholders, + prepared_input, + self._llm, + ) + with BaseMediaIO.use_executor(self._media_executor): + inputs = await _load_prepared_input(prepared_input, self._llm) if prepared_input.media: - return await asyncio.to_thread( - self._llm.preprocess, inputs, sampling_params, disaggregated_params + work = local_request.preprocessing = _PreprocessingWork() + return await loop.run_in_executor( + self._input_executor, + copy_context().run, + work.run, + partial( + self._llm.preprocess, inputs, sampling_params, disaggregated_params + ), ) return inputs diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 2bfd19e45fed..50a4567bbc18 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -11,7 +11,10 @@ import socket import tempfile from abc import ABC, abstractmethod +from collections.abc import Iterator from concurrent.futures import Executor +from contextlib import contextmanager +from contextvars import ContextVar from io import BytesIO from pathlib import Path from stat import S_ISREG @@ -983,6 +986,9 @@ class BaseMediaIO(ABC, Generic[_MediaT]): # Servers can publish a dedicated pool via `set_executor` so # media decoding does not contend with unrelated `to_thread` callers. _executor: ClassVar[Optional[Executor]] = None + _request_executor: ClassVar[ContextVar[Optional[Executor]]] = ContextVar( + "media_io_executor", default=None + ) # Required MIME type prefix for `data:` URIs handled by this modality # (e.g. `"image/"`). A data URI whose declared media type falls outside @@ -999,6 +1005,16 @@ def set_executor(cls, executor: Executor) -> None: """ BaseMediaIO._executor = executor + @classmethod + @contextmanager + def use_executor(cls, executor: Optional[Executor]) -> Iterator[None]: + """Scope decode work to a server-owned executor without replacing other servers' pools.""" + token = BaseMediaIO._request_executor.set(executor) + try: + yield + finally: + BaseMediaIO._request_executor.reset(token) + @classmethod def create( cls, @@ -1084,7 +1100,8 @@ async def _run_in_executor(fn, *args, **kwargs): if kwargs: fn = functools.partial(fn, **kwargs) loop = asyncio.get_event_loop() - return await loop.run_in_executor(BaseMediaIO._executor, fn, *args) + executor = BaseMediaIO._request_executor.get() or BaseMediaIO._executor + return await loop.run_in_executor(executor, fn, *args) class ImageMediaIO(BaseMediaIO[Union[Image.Image, torch.Tensor, np.ndarray]]): diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 06bae3975faf..c73ce6e29c09 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -600,14 +600,17 @@ def _apply_tensor_operation( # Handle dictionary case if isinstance(input_data, dict): - if operation == "to_tensor" and self._is_shared_tensor_dict( - input_data): + if operation in ("to_tensor", "release_handles" + ) and self._is_shared_tensor_dict(input_data): # Convert shared tensor dict back to tensor try: # Import here to avoid circular imports from tensorrt_llm._torch.shared_tensor import \ SharedTensorContainer + if operation == "release_handles": + SharedTensorContainer.release_dict(input_data) + return None return SharedTensorContainer.from_dict( input_data).get_local_view() except Exception as e: @@ -661,6 +664,11 @@ def _apply_tensor_operation( # For any other type, return as-is return input_data + def release_handles(self) -> None: + """Reclaim exports when preprocessed inputs will not be submitted.""" + self.multimodal_data = self._apply_tensor_operation( + self.multimodal_data, "release_handles") + def to_handle(self, element: str) -> None: """Move specified multimodal data element to shared tensor. diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 9a0fa952eb18..95426434a085 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -34,6 +34,7 @@ from ..sampling_params import SamplingParams from .content_format import ContentFormat from .data import TextPrompt +from .model_config import resolve_top_level_model_type from .multimodal import (MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalInput, _as_cpu_tensor, _compute_mm_masks, _find_mm_embedding_lengths_from_masks, @@ -544,6 +545,16 @@ def get_vocab_size(self) -> Optional[int]: "Please override this method to provide the vocabulary size. ") return None + def requires_input_placeholders(self, *, tokenized: bool = False) -> bool: + """Whether the selected input path needs caller-supplied media markers.""" + if tokenized and self.supports_token_id_mm_expansion: + return True + model_type = getattr(type(self), "_registered_model_type", None) + if model_type is None: + model_type = resolve_top_level_model_type(self.config) + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_content_format( + model_type) != ContentFormat.PASSTHROUGH + def get_mm_token_ids(self) -> Optional[Tensor]: """Token IDs for a logical multimodal unit; include framing tokens for one contiguous span per unit. diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 25be2fe73946..a2358bef0593 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -4,9 +4,14 @@ """Unit tests for OpenEngine image input mapping and discovery.""" import asyncio +import base64 +import gc import re import threading from collections.abc import AsyncIterator, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextvars import ContextVar +from pathlib import Path from types import SimpleNamespace from typing import Any @@ -19,12 +24,17 @@ from conftest import AbortError, FakeServicerContext # noqa: E402 import tensorrt_llm.grpc.openengine.request_mapping as request_mapping # noqa: E402 -from tensorrt_llm.grpc.openengine.bindings import generation_pb2, model_pb2 # noqa: E402 +from tensorrt_llm.grpc.openengine.bindings import ( # noqa: E402 + generation_pb2, + model_pb2, + server_pb2, +) from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 from tensorrt_llm.grpc.openengine.errors import UnsupportedFeatureError # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 +from tensorrt_llm.inputs.multimodal import MultimodalParams # noqa: E402 from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, @@ -203,6 +213,15 @@ async def reserve(self, _request_id: str) -> _Reservation: return reservation +@pytest.fixture +def media_pools(): + with ( + ThreadPoolExecutor(max_workers=2, thread_name_prefix="test_input") as input_pool, + ThreadPoolExecutor(max_workers=2, thread_name_prefix="test_media") as media_pool, + ): + yield input_pool, media_pool + + @pytest.fixture def image_llm(monkeypatch): config = _ImageConfig() @@ -289,6 +308,167 @@ def test_unknown_decoded_image_type_is_rejected(): request_mapping._decoded_image_bytes(object()) +@pytest.mark.parametrize("abandon_first", [False, True]) +def test_abandoned_preprocessing_reclaims_exports_regardless_of_completion_order(abandon_first): + import torch + + from tensorrt_llm.grpc.openengine.servicer import _PreprocessingWork + + params = MultimodalParams(multimodal_data={"image": [torch.ones(1024 * 1024)]}) + params.to_handle("multimodal_data") + handle = params.multimodal_data["image"][0] + path = Path("/dev/shm") / base64.b64decode(handle["storage_handle"]).decode().lstrip("/") + result = PreprocessedInputs(prompt_token_ids=[1], multimodal_params=params) + work = _PreprocessingWork() + assert path.exists() + if abandon_first: + work.abandon() + work.run(lambda: result) + work.abandon() + work.abandon() + gc.collect() + assert not path.exists() + + +@pytest.mark.parametrize("tokenized", [False, True]) +@pytest.mark.parametrize("fast_path", [False, True]) +def test_passthrough_placeholder_policy_follows_processor_input_path( + image_llm, monkeypatch, tokenized, fast_path +): + from tensorrt_llm.inputs.content_format import ContentFormat + + capability_llm, _ = image_llm + llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + monkeypatch.setattr(llm.input_processor, "supports_token_id_mm_expansion", fast_path) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_content_format", + lambda model_type: ContentFormat.PASSTHROUGH, + ) + request = generation_pb2.GenerateRequest( + media=[generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image")] + ) + if tokenized: + request.token_ids.ids.extend(llm.tokenizer.encode("Describe this image")) + else: + request.prompt = "Describe this image" + prepared = request_mapping._prepare_input_from_request(request, llm) + if tokenized and fast_path: + with pytest.raises(ValueError, match="placeholder count"): + request_mapping._validate_media_placeholders(prepared, llm) + else: + request_mapping._validate_media_placeholders(prepared, llm) + assert not llm.input_processor.requires_input_placeholders(tokenized=tokenized) + + +def test_submitted_preprocessing_preserves_exports_for_a_delayed_consumer(): + import torch + + from tensorrt_llm.grpc.openengine.servicer import _PreprocessingWork + + params = MultimodalParams(multimodal_data={"image": torch.ones(8)}) + params.to_handle("multimodal_data") + handle = params.multimodal_data["image"] + path = Path("/dev/shm") / base64.b64decode(handle["storage_handle"]).decode().lstrip("/") + work = _PreprocessingWork() + work.run(lambda: PreprocessedInputs(prompt_token_ids=[1], multimodal_params=params)) + work.transfer() + work.abandon() + assert path.exists() + assert params.multimodal_data["image"] is handle + params.to_tensor("multimodal_data") + torch.testing.assert_close(params.multimodal_data["image"], torch.ones(8)) + params.multimodal_data = {} + gc.collect() + assert not path.exists() + + +def test_native_processor_registration_overrides_hf_config_policy(monkeypatch): + from tensorrt_llm.inputs.content_format import ContentFormat + + class NativeProcessor(_ImageInputProcessor): + _registered_model_type = "mistral_common" + + def content_format(model_type): + assert model_type == "mistral_common" + return ContentFormat.PASSTHROUGH + + monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "get_content_format", content_format) + assert not NativeProcessor(_ImageConfig()).requires_input_placeholders() + + +@pytest.mark.parametrize( + "tokens,sequences", + [ + ([], [(1,)]), + ([1, 1, 1, 2], [(1,), (1, 1), (1, 1, 2), (2,)]), + ([0, 1, 2, 0, 1, 2], [(1, 2), (2,), (9,)]), + ], +) +def test_marker_scan_matches_reference_including_overlaps(tokens, sequences): + counts = request_mapping._count_token_sequences(tokens, sequences) + for sequence in sequences: + expected = sum( + tuple(tokens[index : index + len(sequence)]) == sequence + for index in range(len(tokens) - len(sequence) + 1) + ) + assert counts[sequence] == expected + + +@pytest.mark.asyncio +async def test_media_and_preprocessing_do_not_block_control_executor(image_llm, monkeypatch): + from tensorrt_llm.inputs.media_io import BaseMediaIO + + loop = asyncio.get_running_loop() + capability_llm, _ = image_llm + llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + llm.args = SimpleNamespace( + guided_decoding_backend=None, + tensor_parallel_size=1, + pipeline_parallel_size=1, + context_parallel_size=1, + max_batch_size=8, + max_num_tokens=2048, + kv_cache_config=SimpleNamespace(tokens_per_block=32), + ) + llm._executor = SimpleNamespace( + get_kv_cache_capacity=lambda: {}, get_openengine_discovery=lambda: {} + ) + with ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="test_input") as input_pool, + ThreadPoolExecutor(max_workers=1, thread_name_prefix="test_media") as media_pool, + ): + release = threading.Event() + entered = [asyncio.Event(), asyncio.Event()] + + def block(index): + loop.call_soon_threadsafe(entered[index].set) + assert release.wait(timeout=5) + return threading.current_thread().name + + # Saturate both media pools, then require a real Control GetServerInfo RPC + # (which uses asyncio's default executor) to remain responsive. + tasks = [loop.run_in_executor(input_pool, block, 0)] + with BaseMediaIO.use_executor(media_pool): + tasks.append(asyncio.create_task(BaseMediaIO._run_in_executor(block, 1))) + try: + await asyncio.wait_for(asyncio.gather(*(event.wait() for event in entered)), timeout=2) + inference = OpenEngineInferenceServicer( + llm, "image-model", input_executor=input_pool, media_executor=media_pool + ) + control = OpenEngineControlServicer(llm, "image-model", inference) + info = await asyncio.wait_for( + control.GetServerInfo(server_pb2.GetServerInfoRequest(), FakeServicerContext()), + timeout=1, + ) + assert info.engine_name == "tensorrt_llm" + finally: + release.set() + names = await asyncio.gather(*tasks) + assert names[0].startswith("test_input") + assert names[1].startswith("test_media") + + def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, monkeypatch): llm, _ = image_llm monkeypatch.setattr( @@ -890,15 +1070,28 @@ def set_input(prompt): @pytest.mark.asyncio @pytest.mark.parametrize("outcome", ["complete", "local_abort", "coordinator_abort", "rpc_cancel"]) -async def test_media_preprocessing_is_offloaded_and_abort_aware(image_llm, monkeypatch, outcome): +async def test_media_preprocessing_is_offloaded_and_abort_aware( + image_llm, monkeypatch, outcome, media_pools +): capability_llm, _ = image_llm llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) loop = asyncio.get_running_loop() serving_thread = threading.get_ident() + request_context = ContextVar("preprocess_request_context", default=None) started = asyncio.Event() finished = asyncio.Event() + reclaimed = asyncio.Event() release = threading.Event() - preprocessed = PreprocessedInputs(prompt_token_ids=[1, 2, 3]) + params = MultimodalParams() + preprocessed = PreprocessedInputs(prompt_token_ids=[1, 2, 3], multimodal_params=params) + export_paths = [] + release_handles = params.release_handles + + def reclaim(): + release_handles() + loop.call_soon_threadsafe(reclaimed.set) + + monkeypatch.setattr(params, "release_handles", reclaim) preprocess_calls = [] submitted = [] @@ -908,20 +1101,38 @@ def create(cls, defaults, runtime): return cls() async def async_load_bytes(self, data): - return np.zeros((1, 1, 3), dtype=np.uint8) + from tensorrt_llm.inputs.media_io import BaseMediaIO + + def decode(): + assert threading.current_thread().name.startswith("test_media") + return np.zeros((1, 1, 3), dtype=np.uint8) + + return await BaseMediaIO._run_in_executor(decode) def preprocess(inputs, sampling_params, disaggregated_params): assert threading.get_ident() != serving_thread + assert threading.current_thread().name.startswith("test_input") + assert request_context.get() == "image-request" preprocess_calls.append(inputs) loop.call_soon_threadsafe(started.set) try: assert release.wait(timeout=5), "test must release the preprocessing thread" + import torch + + params.multimodal_data = {"image": {"pixel_values": torch.ones(1024 * 1024)}} + params.to_handle("multimodal_data") + handle = params.multimodal_data["image"]["pixel_values"] + export_paths.append( + Path("/dev/shm") / base64.b64decode(handle["storage_handle"]).decode().lstrip("/") + ) return preprocessed finally: loop.call_soon_threadsafe(finished.set) def generate_async(**kwargs): assert threading.get_ident() == serving_thread + if isinstance(kwargs["inputs"], PreprocessedInputs): + kwargs["inputs"].multimodal_params.to_tensor("multimodal_data") submitted.append(kwargs["inputs"]) return llm.result_handle @@ -929,7 +1140,13 @@ def generate_async(**kwargs): monkeypatch.setattr(llm, "preprocess", preprocess) monkeypatch.setattr(llm, "generate_async", generate_async) frontend = _Frontend() - servicer = OpenEngineInferenceServicer(llm, model="image-model", frontend=frontend) + servicer = OpenEngineInferenceServicer( + llm, + model="image-model", + frontend=frontend, + input_executor=media_pools[0], + media_executor=media_pools[1], + ) request = generation_pb2.GenerateRequest( request_id="preprocessing", model="image-model", @@ -941,7 +1158,11 @@ def generate_async(**kwargs): context = FakeServicerContext() async def collect(): - return [response async for response in servicer.Generate(request, context)] + token = request_context.set("image-request") + try: + return [response async for response in servicer.Generate(request, context)] + finally: + request_context.reset(token) task = asyncio.create_task(collect()) try: @@ -982,6 +1203,10 @@ async def collect_text(): task.cancel() await asyncio.gather(task, return_exceptions=True) + if outcome != "complete": + await asyncio.wait_for(reclaimed.wait(), timeout=2) + gc.collect() + assert export_paths and not any(path.exists() for path in export_paths) assert len(preprocess_calls) == 1 assert len(submitted) == (2 if outcome == "complete" else 1) assert servicer.active_request_count() == 0 diff --git a/tests/unittest/grpc/openengine/test_openengine.py b/tests/unittest/grpc/openengine/test_openengine.py index fdd464dbe70a..3ee59307d1b9 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -16,7 +16,9 @@ import asyncio import gc +from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest @@ -47,6 +49,31 @@ def test_format_bind_address_brackets_ipv6() -> None: assert _format_bind_address("[::1]", 8000) == "[::1]:8000" +@pytest.mark.parametrize("fail_stop", [False, True]) +def test_stop_drains_both_media_pools_even_if_grpc_shutdown_fails(fail_stop) -> None: + async def exercise_stop(): + server = OpenEngineServer.__new__(OpenEngineServer) + server._server = SimpleNamespace( + stop=AsyncMock(side_effect=RuntimeError("stop failed") if fail_stop else None) + ) + server._input_executor = ThreadPoolExecutor(max_workers=1) + server._media_executor = ThreadPoolExecutor(max_workers=1) + futures = [ + pool.submit(lambda: True) for pool in (server._input_executor, server._media_executor) + ] + if fail_stop: + with pytest.raises(RuntimeError, match="stop failed"): + await server.stop(grace=0) + else: + await server.stop(grace=0) + assert all(future.result() for future in futures) + for pool in (server._input_executor, server._media_executor): + with pytest.raises(RuntimeError, match="shutdown"): + pool.submit(lambda: None) + + asyncio.run(exercise_stop()) + + def test_openengine_server_serves_the_control_contract() -> None: """The server binds, attaches Control, answers, and shuts down cleanly. From 2e1b47c206189bcac3817a73a4f30f1bd4a67a0c Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 22 Oct 2026 14:11:49 -0700 Subject: [PATCH 35/41] test(openengine): cover video preprocessing isolation after stack refresh Signed-off-by: krishung5 --- .../grpc/openengine/test_multimodal.py | 27 ++++++++++++------- .../openengine/test_multimodal_conformance.py | 3 ++- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index d5159d8ed6e8..eea56bdb0c8f 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -37,8 +37,8 @@ from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.inputs import media_io # noqa: E402 from tensorrt_llm.inputs.media_io import MediaFetchError # noqa: E402 -from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData # noqa: E402 from tensorrt_llm.inputs.multimodal import MultimodalParams # noqa: E402 +from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData # noqa: E402 from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, @@ -932,7 +932,6 @@ async def async_load(self, source): ] - @pytest.mark.asyncio @pytest.mark.parametrize( ("invalid_case", "metadata", "expected_code"), @@ -1433,10 +1432,11 @@ def set_input(prompt): @pytest.mark.asyncio @pytest.mark.parametrize("outcome", ["complete", "local_abort", "coordinator_abort", "rpc_cancel"]) +@pytest.mark.parametrize("modality", ["image", "video"]) async def test_media_preprocessing_is_offloaded_and_abort_aware( - image_llm, monkeypatch, outcome, media_pools + image_video_llm, monkeypatch, outcome, media_pools, modality ): - capability_llm, _ = image_llm + capability_llm = image_video_llm llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) loop = asyncio.get_running_loop() serving_thread = threading.get_ident() @@ -1468,7 +1468,8 @@ async def async_load_bytes(self, data): def decode(): assert threading.current_thread().name.startswith("test_media") - return np.zeros((1, 1, 3), dtype=np.uint8) + frames = np.zeros((1, 1, 3), dtype=np.uint8) + return frames if modality == "image" else VideoData(frames=[frames], metadata={}) return await BaseMediaIO._run_in_executor(decode) @@ -1482,9 +1483,9 @@ def preprocess(inputs, sampling_params, disaggregated_params): assert release.wait(timeout=5), "test must release the preprocessing thread" import torch - params.multimodal_data = {"image": {"pixel_values": torch.ones(1024 * 1024)}} + params.multimodal_data = {modality: {"pixel_values": torch.ones(1024 * 1024)}} params.to_handle("multimodal_data") - handle = params.multimodal_data["image"]["pixel_values"] + handle = params.multimodal_data[modality]["pixel_values"] export_paths.append( Path("/dev/shm") / base64.b64decode(handle["storage_handle"]).decode().lstrip("/") ) @@ -1500,6 +1501,7 @@ def generate_async(**kwargs): return llm.result_handle monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(request_mapping, "VideoMediaIO", _ImageIo) monkeypatch.setattr(llm, "preprocess", preprocess) monkeypatch.setattr(llm, "generate_async", generate_async) frontend = _Frontend() @@ -1513,9 +1515,16 @@ def generate_async(**kwargs): request = generation_pb2.GenerateRequest( request_id="preprocessing", model="image-model", - prompt="", + prompt=f"<{modality}_1>", media=[ - generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image") + generation_pb2.MediaItem( + modality=( + generation_pb2.MODALITY_IMAGE + if modality == "image" + else generation_pb2.MODALITY_VIDEO + ), + raw_bytes=b"media", + ) ], ) context = FakeServicerContext() diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index 766a5a148c34..80f4a94563e1 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -63,7 +63,8 @@ def _model_info(control): ) -def test_image_capability_agrees_with_generate(vlm): +@pytest.mark.parametrize("input_kind", ["prompt", "token_ids"]) +def test_image_capability_agrees_with_generate(vlm, input_kind): inference, control = _servicers(vlm) info = _model_info(control) From 53e665921f5c7c7f3ce04deb7a6048f03e473d0a Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:49:31 -0700 Subject: [PATCH 36/41] fix: resolve image OpenEngine L0 test failures Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- tensorrt_llm/serve/render/fingerprint.py | 2 +- .../executor/test_forward_error_survival.py | 2 + .../grpc/openengine/test_multimodal.py | 170 +++++++++--------- .../unittest/llmapi/apps/test_render_core.py | 16 ++ 4 files changed, 101 insertions(+), 89 deletions(-) diff --git a/tensorrt_llm/serve/render/fingerprint.py b/tensorrt_llm/serve/render/fingerprint.py index eca1c4201e76..e7becc932f67 100644 --- a/tensorrt_llm/serve/render/fingerprint.py +++ b/tensorrt_llm/serve/render/fingerprint.py @@ -202,7 +202,7 @@ def compute_fingerprint(res: "RenderResources") -> Dict[str, Any]: compared: Dict[str, Any] = { # Harmony renders the same for any model type, so a router that was told to use # Harmony without being able to resolve the type still matches its workers. - "model_type": None if res.use_harmony else res.model_type, + "model_type": None if res.use_harmony else (res.model_type or None), "extension": type(res.extension).__qualname__, # Bumped by an extension whose rendering output changes between revisions. "extension_render_version": getattr(res.extension, "render_version", 1), diff --git a/tests/unittest/executor/test_forward_error_survival.py b/tests/unittest/executor/test_forward_error_survival.py index 267ae48fa9d5..a13c8d9036bc 100644 --- a/tests/unittest/executor/test_forward_error_survival.py +++ b/tests/unittest/executor/test_forward_error_survival.py @@ -53,6 +53,7 @@ def _request(mm_data=None): py_disaggregated_params=None, sampling_config=None, py_multimodal_data=mm_data, + kv_cache_retention_config=None, ) @@ -61,6 +62,7 @@ def _validating_executor(model): executor = object.__new__(PyExecutor) executor.kv_cache_transceiver = None executor.kv_cache_manager = None + executor._is_kv_manager_v2 = False executor.model_engine = SimpleNamespace(model=model) executor.sampler = Mock() executor._validate_token_id_range = Mock() diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index a2358bef0593..0abaae084841 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -213,15 +213,6 @@ async def reserve(self, _request_id: str) -> _Reservation: return reservation -@pytest.fixture -def media_pools(): - with ( - ThreadPoolExecutor(max_workers=2, thread_name_prefix="test_input") as input_pool, - ThreadPoolExecutor(max_workers=2, thread_name_prefix="test_media") as media_pool, - ): - yield input_pool, media_pool - - @pytest.fixture def image_llm(monkeypatch): config = _ImageConfig() @@ -1070,9 +1061,7 @@ def set_input(prompt): @pytest.mark.asyncio @pytest.mark.parametrize("outcome", ["complete", "local_abort", "coordinator_abort", "rpc_cancel"]) -async def test_media_preprocessing_is_offloaded_and_abort_aware( - image_llm, monkeypatch, outcome, media_pools -): +async def test_media_preprocessing_is_offloaded_and_abort_aware(image_llm, monkeypatch, outcome): capability_llm, _ = image_llm llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) loop = asyncio.get_running_loop() @@ -1139,87 +1128,92 @@ def generate_async(**kwargs): monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) monkeypatch.setattr(llm, "preprocess", preprocess) monkeypatch.setattr(llm, "generate_async", generate_async) - frontend = _Frontend() - servicer = OpenEngineInferenceServicer( - llm, - model="image-model", - frontend=frontend, - input_executor=media_pools[0], - media_executor=media_pools[1], - ) - request = generation_pb2.GenerateRequest( - request_id="preprocessing", - model="image-model", - prompt="", - media=[ - generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image") - ], - ) - context = FakeServicerContext() - - async def collect(): - token = request_context.set("image-request") - try: - return [response async for response in servicer.Generate(request, context)] - finally: - request_context.reset(token) - - task = asyncio.create_task(collect()) - try: - await asyncio.wait_for(started.wait(), timeout=2) - text_request = generation_pb2.GenerateRequest( - request_id="concurrent-text", model="image-model", prompt="hello" + # Pools must close during the call phase, before pytest-threadleak's snapshot. + with ( + ThreadPoolExecutor(max_workers=2, thread_name_prefix="test_input") as input_pool, + ThreadPoolExecutor(max_workers=2, thread_name_prefix="test_media") as media_pool, + ): + frontend = _Frontend() + servicer = OpenEngineInferenceServicer( + llm, + model="image-model", + frontend=frontend, + input_executor=input_pool, + media_executor=media_pool, ) + request = generation_pb2.GenerateRequest( + request_id="preprocessing", + model="image-model", + prompt="", + media=[ + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image") + ], + ) + context = FakeServicerContext() - async def collect_text(): - return [ - response - async for response in servicer.Generate(text_request, FakeServicerContext()) - ] + async def collect(): + token = request_context.set("image-request") + try: + return [response async for response in servicer.Generate(request, context)] + finally: + request_context.reset(token) - assert await asyncio.wait_for(collect_text(), timeout=1) - assert submitted == ["hello"] - if outcome == "complete": - release.set() - assert await asyncio.wait_for(task, timeout=2) - assert submitted[1] is preprocessed - elif outcome == "rpc_cancel": - task.cancel() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(task, timeout=1) - assert context.abort_code is None - else: - if outcome == "local_abort": - assert servicer.abort_request_by_id(request.request_id) + task = asyncio.create_task(collect()) + try: + await asyncio.wait_for(started.wait(), timeout=2) + text_request = generation_pb2.GenerateRequest( + request_id="concurrent-text", model="image-model", prompt="hello" + ) + + async def collect_text(): + return [ + response + async for response in servicer.Generate(text_request, FakeServicerContext()) + ] + + assert await asyncio.wait_for(collect_text(), timeout=1) + assert submitted == ["hello"] + if outcome == "complete": + release.set() + assert await asyncio.wait_for(task, timeout=2) + assert submitted[1] is preprocessed + elif outcome == "rpc_cancel": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + assert context.abort_code is None else: - assert frontend.reservations[0].abort() - with pytest.raises(AbortError): - await asyncio.wait_for(task, timeout=1) - assert context.abort_code == grpc.StatusCode.CANCELLED - finally: - release.set() - await asyncio.wait_for(finished.wait(), timeout=2) - if not task.done(): - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - if outcome != "complete": - await asyncio.wait_for(reclaimed.wait(), timeout=2) - gc.collect() - assert export_paths and not any(path.exists() for path in export_paths) - assert len(preprocess_calls) == 1 - assert len(submitted) == (2 if outcome == "complete" else 1) - assert servicer.active_request_count() == 0 - if outcome != "complete": - monkeypatch.setattr(llm, "preprocess", lambda *args: preprocessed) - responses = [ - response async for response in servicer.Generate(request, FakeServicerContext()) - ] - assert responses - assert submitted[-1] is preprocessed - assert len(submitted) == 2 + if outcome == "local_abort": + assert servicer.abort_request_by_id(request.request_id) + else: + assert frontend.reservations[0].abort() + with pytest.raises(AbortError): + await asyncio.wait_for(task, timeout=1) + assert context.abort_code == grpc.StatusCode.CANCELLED + finally: + release.set() + await asyncio.wait_for(finished.wait(), timeout=2) + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + if outcome != "complete": + await asyncio.wait_for(reclaimed.wait(), timeout=2) + gc.collect() + assert export_paths and not any(path.exists() for path in export_paths) + assert len(preprocess_calls) == 1 + assert len(submitted) == (2 if outcome == "complete" else 1) assert servicer.active_request_count() == 0 - assert all(reservation.released for reservation in frontend.reservations) + if outcome != "complete": + monkeypatch.setattr(llm, "preprocess", lambda *args: preprocessed) + responses = [ + response async for response in servicer.Generate(request, FakeServicerContext()) + ] + assert responses + assert submitted[-1] is preprocessed + assert len(submitted) == 2 + assert servicer.active_request_count() == 0 + assert all(reservation.released for reservation in frontend.reservations) @pytest.mark.asyncio diff --git a/tests/unittest/llmapi/apps/test_render_core.py b/tests/unittest/llmapi/apps/test_render_core.py index 438fada4cca0..63bd09c88eaa 100644 --- a/tests/unittest/llmapi/apps/test_render_core.py +++ b/tests/unittest/llmapi/apps/test_render_core.py @@ -471,6 +471,22 @@ async def run(): class TestFingerprint: + @pytest.mark.parametrize("model_type", [None, ""]) + def test_unknown_worker_model_type_matches_the_router(self, tokenizer, model_type) -> None: + worker = self._worker_resources( + tokenizer, model_config=SimpleNamespace(model_type=model_type) + ) + router = RenderResources.from_tokenizer(tokenizer) + assert fingerprints_match(worker.fingerprint(), router.fingerprint()) + + @pytest.mark.parametrize("local,remote", [(None, "qwen3"), ("", "qwen3"), ("qwen2", "qwen3")]) + def test_known_model_types_remain_distinct(self, tokenizer, local, remote) -> None: + base = resources(tokenizer) + assert not fingerprints_match( + dataclasses.replace(base, model_type=local).fingerprint(), + dataclasses.replace(base, model_type=remote).fingerprint(), + ) + def test_a_different_server_template_changes_the_digest(self, tokenizer) -> None: plain = compute_fingerprint(resources(tokenizer)) with_template = compute_fingerprint( From af1d00bf479384ca1dbe38a28cf954bfbc48be50 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:34:56 -0700 Subject: [PATCH 37/41] fix: remove duplicate test request keyword after merge Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- tests/unittest/executor/test_forward_error_survival.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unittest/executor/test_forward_error_survival.py b/tests/unittest/executor/test_forward_error_survival.py index b91aa65f1ab5..2f077085ef0d 100644 --- a/tests/unittest/executor/test_forward_error_survival.py +++ b/tests/unittest/executor/test_forward_error_survival.py @@ -54,7 +54,6 @@ def _request(mm_data=None): py_disaggregated_params=None, sampling_config=None, py_multimodal_data=mm_data, - kv_cache_retention_config=None, ) From 2d69a82634616aab31ba437368ea24c87afbcca0 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:45:53 -0700 Subject: [PATCH 38/41] fix(openengine): stop video decoding between operations on cancellation Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- .../grpc/openengine/request_mapping.py | 5 ++ tensorrt_llm/inputs/media_io.py | 14 ++++- .../grpc/openengine/test_multimodal.py | 58 +++++++++++++++++++ tests/unittest/inputs/test_video_decode.py | 9 +++ 4 files changed, 85 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 9a91e257bf54..c43d0f742ea2 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -54,6 +54,11 @@ def __init__(self, budget: "_DecodedMediaBudget", peak_bytes: int, retained_byte self._retained_bytes = retained_bytes self._active = True + def check_cancelled(self) -> None: + with self._budget._condition: + if self._budget._cancelled: + raise ValueError("decoded media loading was cancelled") + def shrink(self, retained_bytes: int) -> None: if not self._active or not 0 <= retained_bytes <= self._retained_bytes: raise ValueError("invalid decoded-media reservation shrink") diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 4722b173462b..0e12b267b3ee 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -54,6 +54,8 @@ class MediaFetchError(RuntimeError): class _DecodedMemoryReservation(Protocol): + def check_cancelled(self) -> None: ... + def shrink(self, retained_bytes: int) -> None: ... def commit(self, retained_bytes: int) -> None: ... @@ -905,11 +907,19 @@ def _load_video_by_cv2( # affected (e.g. every frame retrieves with a drifted shape). skip_warned = False frame_idx = 0 - while frame_idx <= max_idx and vidcap.grab(): + while frame_idx <= max_idx: + if decoded_memory_reservation is not None: + decoded_memory_reservation.check_cancelled() + if not vidcap.grab(): + break + if decoded_memory_reservation is not None: + decoded_memory_reservation.check_cancelled() if frame_idx in target_set: # Reuse a single BGR buffer across retrieves; cv2 replaces its # contents in place when the argument is shape-compatible. ok, bgr_scratch = vidcap.retrieve(bgr_scratch) + if decoded_memory_reservation is not None: + decoded_memory_reservation.check_cancelled() if ok: fh, fw = bgr_scratch.shape[:2] if stacked_rgb is None: @@ -965,6 +975,8 @@ def _load_video_by_cv2( ) skip_warned = True frame_idx += 1 + if decoded_memory_reservation is not None: + decoded_memory_reservation.check_cancelled() vidcap.release() if stacked_rgb is None or not valid_indices: diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 7b171fa89313..bc2cb9865266 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -521,6 +521,64 @@ def decode(_index): list(pool.map(decode, range(3))) +@pytest.mark.parametrize("cancel_during", ["grab", "retrieve", "convert"]) +def test_video_cancellation_stops_between_frame_operations(cancel_during, monkeypatch): + budget = request_mapping._DecodedMediaBudget(1024) + operations = [] + + def grab(): + operations.append("grab") + if cancel_during == "grab" and operations.count("grab") == 2: + budget.cancel() + return True + + def retrieve(_scratch): + operations.append("retrieve") + if cancel_during == "retrieve" and operations.count("retrieve") == 2: + budget.cancel() + return True, np.zeros((1, 1, 3), dtype=np.uint8) + + def convert(_source, _code, dst): + operations.append("convert") + dst.fill(0) + if cancel_during == "convert": + budget.cancel() + + capture = SimpleNamespace( + isOpened=lambda: True, + get=lambda key: 3 if key == 0 else 1, + grab=grab, + retrieve=retrieve, + release=lambda: operations.append("release"), + ) + monkeypatch.setattr( + media_io, + "_get_cv2", + lambda: SimpleNamespace( + VideoCapture=lambda _source: capture, + CAP_PROP_FRAME_COUNT=0, + CAP_PROP_FPS=1, + COLOR_BGR2RGB=2, + cvtColor=convert, + ), + ) + with pytest.raises(ValueError, match="cancelled"): + media_io._load_video_by_cv2( + "video.mp4", + num_frames=3, + fps=-1, + format="np", + decoded_memory_reserver=budget.reserve, + ) + expected = ["grab", "retrieve", "convert"] + if cancel_during in ("grab", "retrieve"): + expected.append("grab") + if cancel_during == "retrieve": + expected.append("retrieve") + assert operations == expected + ["release"] + assert (budget._used, budget._retained, budget._reserved_retained) == (0, 0, 0) + + @pytest.mark.parametrize("cancel", [False, True]) def test_video_peak_reservation_waits_for_settlement_or_cancellation(cancel, monkeypatch): budget = request_mapping._DecodedMediaBudget(20, peak_limit=20) diff --git a/tests/unittest/inputs/test_video_decode.py b/tests/unittest/inputs/test_video_decode.py index 691903573edb..e39292841b78 100644 --- a/tests/unittest/inputs/test_video_decode.py +++ b/tests/unittest/inputs/test_video_decode.py @@ -102,6 +102,9 @@ class _Reservation: committed = None released = False + def check_cancelled(self) -> None: + pass + def commit(self, retained_bytes: int) -> None: self.committed = retained_bytes @@ -132,6 +135,9 @@ def test_decode_releases_reservation_after_allocation_failure(sample_video_path: class _Reservation: released = False + def check_cancelled(self) -> None: + pass + def commit(self, retained_bytes: int) -> None: raise AssertionError("failed decode must not commit") @@ -159,6 +165,9 @@ def test_truncated_numpy_video_owns_only_the_counted_frames(sample_video_path: s committed = [] class _Reservation: + def check_cancelled(self) -> None: + pass + def commit(self, retained_bytes: int) -> None: committed.append(retained_bytes) From 68b0ef2622b1824b1f3c110387f74c0b93dc20e0 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:40:54 -0700 Subject: [PATCH 39/41] [None][fix] Preserve audio contracts on the shared OpenEngine lifecycle Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- .../grpc/openengine/request_mapping.py | 6 ++- tensorrt_llm/grpc/openengine/server.py | 10 ++-- .../grpc/openengine/test_multimodal.py | 47 ++++++++++++++++++- .../openengine/test_multimodal_conformance.py | 22 +++++++-- .../grpc/openengine/test_openengine.py | 8 ++-- 5 files changed, 79 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index e6e1a3ab6e8b..41334157914e 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -586,7 +586,11 @@ def _video_runtime_available() -> bool: reserves_decoded_bytes=True, ), generation_pb2.MODALITY_AUDIO: _MediaHandler( - "audio", _validate_audio_item, _audio_loader_factory, _decoded_audio_bytes, reserves_decoded_bytes=True + "audio", + _validate_audio_item, + _audio_loader_factory, + _decoded_audio_bytes, + reserves_decoded_bytes=True, ), } diff --git a/tensorrt_llm/grpc/openengine/server.py b/tensorrt_llm/grpc/openengine/server.py index 5fdf610a7398..f5c39e3c27df 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -21,8 +21,8 @@ import uvloop from tensorrt_llm import LLM as PyTorchLLM -from tensorrt_llm.inputs.media_io import _close_aiohttp_session from tensorrt_llm.grpc._protobuf import _warn_if_python_protobuf +from tensorrt_llm.inputs.media_io import _close_aiohttp_session from tensorrt_llm.logger import logger from tensorrt_llm.serve._frontend_processes import ( FrontendStartupCancelled, @@ -193,8 +193,12 @@ async def stop(self, grace: float = 5.0) -> None: await _close_aiohttp_session() finally: await asyncio.gather( - asyncio.to_thread(self._input_executor.shutdown, wait=True, cancel_futures=True), - asyncio.to_thread(self._media_executor.shutdown, wait=True, cancel_futures=True), + asyncio.to_thread( + self._input_executor.shutdown, wait=True, cancel_futures=True + ), + asyncio.to_thread( + self._media_executor.shutdown, wait=True, cancel_futures=True + ), ) logger.info("OpenEngine server stopped") diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 3b79d7112199..751492208d56 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -12,9 +12,8 @@ from collections.abc import AsyncIterator, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext -from io import BytesIO - from contextvars import ContextVar +from io import BytesIO from pathlib import Path from types import SimpleNamespace from typing import Any @@ -366,6 +365,49 @@ def test_media_capability_accepts_a_registered_non_base_processor(monkeypatch): assert request_mapping.media_capabilities_for_model(llm).modalities == {"audio"} +@pytest.mark.asyncio +@pytest.mark.parametrize("prompt", ["", "<|startoftranscript|><|fr|><|translate|><|notimestamps|>"]) +async def test_audio_text_prompt_survives_preparation_without_placeholders(monkeypatch, prompt): + import soundfile + + config = _ImageConfig() + processor = _AudioInputProcessor(config) + processor.supports_token_ids_input = False + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", lambda _, modality: modality == "audio" + ) + monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "get_placeholder", lambda *_: "") + llm = _ImageLlm(processor, config) + llm.args.reasoning_parser = None + llm.args.max_seq_len = 2048 + inference = OpenEngineInferenceServicer(llm, "audio-model") + control = OpenEngineControlServicer(llm, "audio-model", inference) + info = await control.GetModelInfo(model_pb2.GetModelInfoRequest(), FakeServicerContext()) + assert info.supports_text_input is True + assert info.supports_token_ids_input is False + encoded = BytesIO() + soundfile.write(encoded, np.zeros(32), 16000, format="WAV") + request = generation_pb2.GenerateRequest( + request_id="audio-task-prompt", + model="audio-model", + prompt=prompt, + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + raw_bytes=encoded.getvalue(), + mime_type="audio/wav", + ) + ], + ) + responses = [response async for response in inference.Generate(request, FakeServicerContext())] + assert responses[-1].HasField("finished") + inputs = llm.generate_kwargs["inputs"] + assert inputs["prompt"] == prompt + samples, rate = inputs["multi_modal_data"]["audio"][0] + assert samples.shape == (32,) and rate == 16000 + assert inference.active_request_count() == 0 + + def test_media_capability_rejects_a_registered_non_callable_processor(monkeypatch): config = _ImageConfig() monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", lambda *_args: True) @@ -376,6 +418,7 @@ def test_media_capability_rejects_a_registered_non_callable_processor(monkeypatc assert request_mapping.media_capabilities_for_model(llm).modalities == set() + @pytest.mark.asyncio @pytest.mark.parametrize("disabled", [False, True]) async def test_disabled_mm_encoder_agrees_between_discovery_and_generate(image_llm, disabled): diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index cb1189b68819..906d05ce671d 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -236,7 +236,10 @@ def whisper(): engine.shutdown() -def test_audio_capability_agrees_with_generate(whisper): +@pytest.mark.parametrize( + "task_prompt", ["", "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"] +) +def test_audio_capability_agrees_with_generate(whisper, task_prompt): whisper, audio_path = whisper inference = OpenEngineInferenceServicer(whisper, WHISPER_MODEL_NAME) control = OpenEngineControlServicer(whisper, WHISPER_MODEL_NAME, inference) @@ -256,7 +259,7 @@ def test_audio_capability_agrees_with_generate(whisper): request = generation_pb2.GenerateRequest( request_id="openengine-audio-conformance", model=WHISPER_MODEL_NAME, - prompt="", + prompt=task_prompt, sampling=generation_pb2.SamplingParams(temperature=0.0), stopping=generation_pb2.StoppingOptions(max_tokens=96), media=[ @@ -267,6 +270,20 @@ def test_audio_capability_agrees_with_generate(whisper): ], ) + token_request = generation_pb2.GenerateRequest() + token_request.CopyFrom(request) + token_request.request_id += "-token-input" + token_request.token_ids.CopyFrom(generation_pb2.TokenIds(ids=[1])) + token_context = FakeServicerContext() + + async def reject_token_input(): + return [response async for response in inference.Generate(token_request, token_context)] + + with pytest.raises(AbortError): + asyncio.run(reject_token_input()) + assert token_context.abort_code == grpc.StatusCode.UNIMPLEMENTED + assert inference.active_request_count() == 0 + async def collect_audio(): return [response async for response in inference.Generate(request, FakeServicerContext())] @@ -284,4 +301,3 @@ async def collect_audio(): finished = [response for response in responses if response.HasField("finished")] assert len(finished) == 1 assert finished[0].usage.completion_tokens == len(token_ids) - diff --git a/tests/unittest/grpc/openengine/test_openengine.py b/tests/unittest/grpc/openengine/test_openengine.py index 75cd022732e8..cb1dc06e3712 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -30,8 +30,8 @@ openengine_pb2_grpc, server_pb2, ) -from tensorrt_llm.inputs.media_io import _get_aiohttp_session from tensorrt_llm.grpc.openengine.server import OpenEngineServer # noqa: E402 +from tensorrt_llm.inputs.media_io import _get_aiohttp_session # noqa: E402 # grpc.aio starts a `_poll_wrapper` daemon thread on server start and tears it # down only once its internal state is released, which happens after the event @@ -43,8 +43,7 @@ def test_format_bind_address_brackets_ipv6() -> None: """A bare IPv6 host is bracketed; IPv4 and already-bracketed hosts are not.""" - from tensorrt_llm.inputs.media_io import _get_aiohttp_session -from tensorrt_llm.grpc.openengine.server import _format_bind_address + from tensorrt_llm.grpc.openengine.server import _format_bind_address assert _format_bind_address("127.0.0.1", 8000) == "127.0.0.1:8000" assert _format_bind_address("::1", 8000) == "[::1]:8000" @@ -167,8 +166,7 @@ async def exercise_server() -> None: def test_is_loopback_classifies_bind_hosts() -> None: """The listener is unauthenticated, so a non-loopback bind must be flagged.""" - from tensorrt_llm.inputs.media_io import _get_aiohttp_session -from tensorrt_llm.grpc.openengine.server import _is_loopback + from tensorrt_llm.grpc.openengine.server import _is_loopback assert _is_loopback("127.0.0.1") assert _is_loopback("::1") From adc04f7c363ddea7bdf15a24102456ad51e2e21f Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:57:21 -0700 Subject: [PATCH 40/41] [None][fix] Harden OpenEngine audio loading and capability discovery Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- tensorrt_llm/grpc/openengine/README.md | 5 + tensorrt_llm/grpc/openengine/capabilities.py | 5 +- tensorrt_llm/grpc/openengine/control.py | 5 +- tensorrt_llm/inputs/media_io.py | 15 ++- tensorrt_llm/inputs/registry.py | 3 + .../unittest/grpc/openengine/test_control.py | 1 + .../grpc/openengine/test_multimodal.py | 93 ++++++++++++++++--- .../inputs/test_async_media_loading.py | 14 +++ 8 files changed, 123 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index 5ce68cf169c1..ccc11bbf2a4f 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -69,6 +69,11 @@ Set `failureThreshold * periodSeconds` to cover the longest expected model start OpenEngine gRPC does not yet expose sleep or weight updates. When those operations are added, they must publish `inference=NOT_SERVING` for their duration while keeping `liveness=SERVING`; the current executor health predicate alone does not detect those admission states. A sidecar that loses `Watch` or sees a new engine identity must validate health again. +For models that do not accept token-ID input, including Whisper, the optional +`Health(include_inference_probe=True)` check reports the inference probe as not +applicable. It does not submit inference and does not establish inference liveness; +the ordinary lifecycle and executor health checks still apply. + The `Inference.Generate` RPC loads the selected model through TensorRT-LLM's PyTorch `LLM` API and streams incremental token and text events. It supports text and token-ID inputs, native sampling and stopping options, top-N prompt and output log probabilities, TensorRT-LLM guided-decoding modes, cache salt, trace-context propagation, multiple output sequences, finish reasons, and final usage. Clients must continuously consume the response stream. If response delivery remains stalled for 30 seconds, the server aborts the engine request and terminates the stream with a retryable overload error. diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 0c83804921d1..c324a40492e3 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -98,7 +98,10 @@ def resolve_media_capabilities( if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) } # Gemma4 checkpoints may omit the audio encoder despite sharing a model type. - if model_type == "gemma4" and getattr(input_processor.config, "audio_config", None) is None: + if ( + model_type == "gemma4" + and getattr(getattr(input_processor, "config", None), "audio_config", None) is None + ): placeholders.pop("audio", None) # Fail startup before Control advertises a malformed template. for modality, template in placeholders.items(): diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index 68f9864a0f14..b693cdeb885d 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -480,7 +480,10 @@ async def _run_inference_probe(self, model: str) -> lifecycle_pb2.HealthCheck: return lifecycle_pb2.HealthCheck( name="inference_probe", state=lifecycle_pb2.HEALTH_STATE_READY, - message="probe not applicable: model does not accept token-ID input", + message=( + "probe not applicable: model does not accept token-ID input; " + "inference liveness was not tested" + ), ) handle = None reservation = None diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index f960ac233767..4d1317476073 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -135,8 +135,8 @@ def convert_image_mode(image: Image.Image, to_mode: str) -> Image.Image: _SUPPORTED_VIDEO_FORMATS = ("pt", "np", "pil") # One aiohttp session per event loop. Media loading can run on multiple live -# loops in one process, and aiohttp sessions must only be used and closed from -# their owning loop. +# loops in one process. Live sessions must only be used and closed from their +# owning loop; closed-loop entries are disposed on the next lookup. _aiohttp_sessions: dict[asyncio.AbstractEventLoop, aiohttp.ClientSession] = {} _aiohttp_sessions_lock = threading.Lock() @@ -145,11 +145,16 @@ async def _get_aiohttp_session() -> aiohttp.ClientSession: """Return the aiohttp session owned by the current event loop.""" current_loop = asyncio.get_running_loop() with _aiohttp_sessions_lock: + stale_sessions = [ + _aiohttp_sessions.pop(loop) for loop in list(_aiohttp_sessions) if loop.is_closed() + ] session = _aiohttp_sessions.get(current_loop) if session is None or session.closed: connector = aiohttp.TCPConnector(resolver=_ValidatingResolver()) session = aiohttp.ClientSession(connector=connector) _aiohttp_sessions[current_loop] = session + for stale_session in stale_sessions: + await stale_session.close() return session @@ -1361,7 +1366,11 @@ def _read(self, source: Union[BytesIO, str]) -> Tuple[np.ndarray, int]: samples = np.empty(shape, dtype=np.float64) decoded = audio.read(frames=frames, out=samples) if decoded.shape != samples.shape: - raise ValueError("audio decoded fewer frames than its header declares") + # Compressed headers may overestimate frames. Drop the decoder's + # view before shrinking the owning allocation and its budget. + decoded_shape = decoded.shape + decoded = None + samples.resize(decoded_shape) sample_rate = audio.samplerate reservation.commit(samples.nbytes) return samples, sample_rate diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 140c146a7955..ad3465a96157 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -100,6 +100,9 @@ class InputProcessor(Protocol): - Register the inherited class to the model class using @register_input_processor(...) """ + supports_token_ids_input: bool = True + """Whether pre-tokenized prompts are accepted; encoder-driven processors may opt out.""" + model_path: any config: any tokenizer: any diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index a025f69a3a72..ef4a8f9c88c1 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -442,6 +442,7 @@ def unexpected_generate(*args, **kwargs): probe = next(check for check in response.checks if check.name == "inference_probe") assert probe.state == lifecycle_pb2.HEALTH_STATE_READY assert "not applicable" in probe.message + assert "inference liveness was not tested" in probe.message assert response.state == lifecycle_pb2.HEALTH_STATE_READY diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 751492208d56..bc425119bde9 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -43,6 +43,7 @@ from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, + DefaultInputProcessor, ) from tensorrt_llm.llmapi.llm import PreprocessedInputs # noqa: E402 @@ -505,10 +506,11 @@ def test_decoded_audio_bytes_count_samples_and_reject_unknown_types(): request_mapping._decoded_audio_bytes(object()) +@pytest.mark.parametrize("audio_format", ["FLAC", "MP3"]) @pytest.mark.parametrize("channels", [1, 2]) @pytest.mark.parametrize("source", ["bytes", "base64", "file"]) def test_audio_reserves_before_decoding_and_preserves_samples( - channels, source, tmp_path, monkeypatch + audio_format, channels, source, tmp_path, monkeypatch ): import base64 @@ -516,15 +518,17 @@ def test_audio_reserves_before_decoding_and_preserves_samples( samples = np.arange(32 * channels, dtype=np.float64).reshape(32, channels) / 128 encoded = BytesIO() - soundfile.write(encoded, samples, 16000, format="FLAC") + soundfile.write(encoded, samples, 16000, format=audio_format) data = encoded.getvalue() expected, _ = soundfile.read(BytesIO(data)) - budget = request_mapping._DecodedMediaBudget(samples.nbytes) + with soundfile.SoundFile(BytesIO(data)) as audio: + reserved_bytes = audio.frames * audio.channels * np.dtype(np.float64).itemsize + budget = request_mapping._DecodedMediaBudget(reserved_bytes) original_read = soundfile.SoundFile.read def checked_read(audio, *args, **kwargs): - assert (budget._used, budget._reserved_retained) == (samples.nbytes, samples.nbytes) - assert kwargs["out"].nbytes == samples.nbytes + assert (budget._used, budget._reserved_retained) == (reserved_bytes, reserved_bytes) + assert kwargs["out"].nbytes == reserved_bytes return original_read(audio, *args, **kwargs) monkeypatch.setattr(soundfile.SoundFile, "read", checked_read) @@ -532,21 +536,28 @@ def checked_read(audio, *args, **kwargs): if source == "bytes": decoded, rate = loader.load_bytes(data) elif source == "base64": - decoded, rate = loader.load_base64("audio/flac", base64.b64encode(data).decode()) + decoded, rate = loader.load_base64( + "audio/mpeg" if audio_format == "MP3" else "audio/flac", + base64.b64encode(data).decode(), + ) else: - path = tmp_path / "audio.flac" + path = tmp_path / f"audio.{audio_format.lower()}" path.write_bytes(data) decoded, rate = loader.load_file(path.as_uri()) - np.testing.assert_array_equal(decoded, expected) + if audio_format == "MP3": + # Buffered and ordinary MP3 reads can differ by float32 rounding. + np.testing.assert_allclose(decoded, expected, rtol=0, atol=np.finfo(np.float32).eps) + else: + np.testing.assert_array_equal(decoded, expected) assert rate == 16000 assert (budget._used, budget._retained, budget._reserved_retained) == ( - samples.nbytes, - samples.nbytes, + expected.nbytes, + expected.nbytes, 0, ) -@pytest.mark.parametrize("outcome", ["oversized", "failure", "truncated", "cancelled"]) +@pytest.mark.parametrize("outcome", ["oversized", "failure", "cancelled"]) def test_audio_budget_rejects_before_allocation_and_cleans_up(outcome, monkeypatch): import soundfile @@ -562,7 +573,6 @@ def read(audio, *args, **kwargs): raise RuntimeError("decode failed") if outcome == "cancelled": raise asyncio.CancelledError() - return kwargs["out"][:0] monkeypatch.setattr(soundfile.SoundFile, "read", read) if outcome == "oversized": @@ -572,7 +582,6 @@ def read(audio, *args, **kwargs): error = { "oversized": ValueError, "failure": RuntimeError, - "truncated": ValueError, "cancelled": asyncio.CancelledError, }[outcome] with pytest.raises(error): @@ -580,6 +589,39 @@ def read(audio, *args, **kwargs): assert (budget._used, budget._retained, budget._reserved_retained) == (0, 0, 0) +@pytest.mark.parametrize("channels", [1, 2]) +@pytest.mark.parametrize("decoded_frames", [0, 17]) +def test_audio_short_read_preserves_samples_and_accounts_for_owning_storage( + channels, decoded_frames, monkeypatch +): + import soundfile + + encoded = BytesIO() + samples = np.arange(32 * channels, dtype=np.float64).reshape(32, channels) / 128 + soundfile.write(encoded, samples, 16000, format="FLAC") + expected, _ = soundfile.read(BytesIO(encoded.getvalue())) + budget = request_mapping._DecodedMediaBudget(samples.nbytes) + original_read = soundfile.SoundFile.read + + def short_read(audio, *args, **kwargs): + assert budget._used == samples.nbytes + return original_read(audio, *args, **kwargs)[:decoded_frames] + + monkeypatch.setattr(soundfile.SoundFile, "read", short_read) + decoded, rate = media_io.AudioMediaIO(decoded_memory_reserver=budget.reserve).load_bytes( + encoded.getvalue() + ) + np.testing.assert_array_equal(decoded, expected[:decoded_frames]) + assert rate == 16000 + assert decoded.flags.owndata + assert decoded.base is None + assert (budget._used, budget._retained, budget._reserved_retained) == ( + decoded.nbytes, + decoded.nbytes, + 0, + ) + + @pytest.mark.asyncio async def test_audio_request_uses_one_shared_budget_without_double_charging( image_video_audio_llm, monkeypatch @@ -636,6 +678,31 @@ class GemmaConfig: await request_mapping._input_from_request(request, llm, capabilities) +@pytest.mark.parametrize("config_state", ["absent", "none"]) +def test_gemma4_missing_processor_config_fails_closed(config_state, monkeypatch): + config = SimpleNamespace(model_type="gemma4", audio_config=object()) + processor = _ImageInputProcessor(config) + if config_state == "absent": + del processor._config + else: + processor._config = None + llm = SimpleNamespace(input_processor=processor, _hf_model_config=config) + monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", lambda *args: True) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + capabilities = request_mapping.media_capabilities_for_model(llm) + assert "image" in capabilities.modalities + assert "audio" not in capabilities.modalities + + +def test_default_input_processor_declares_token_ids_supported(): + processor = DefaultInputProcessor(model_path="unused", config=None, tokenizer=None) + assert processor.supports_token_ids_input is True + + def test_decoded_media_budget_settles_peak_to_retained_bytes(): budget = request_mapping._DecodedMediaBudget(20, peak_limit=20) diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index 06111607c016..bfc94867dd29 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -181,6 +181,20 @@ async def test_new_session_created_after_close(self): ) assert not second_session.closed + @pytest.mark.asyncio + async def test_closed_event_loop_sessions_are_disposed_on_next_lookup(self): + async def abandoned_session(): + return asyncio.get_running_loop(), await _get_aiohttp_session() + + loop, session = await asyncio.to_thread(asyncio.run, abandoned_session()) + assert loop.is_closed() + assert loop in media_io_module._aiohttp_sessions + assert not session.closed + current_session = await _get_aiohttp_session() + assert loop not in media_io_module._aiohttp_sessions + assert session.closed + assert not current_session.closed + @pytest.mark.asyncio async def test_live_event_loops_keep_independent_sessions(self): first_session = await _get_aiohttp_session() From 8da5d1321f1caae19590fbd32197ef8815ace5c7 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:47:34 -0700 Subject: [PATCH 41/41] Declare media capabilities on loaded input processors Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- .../_torch/models/modeling_gemma4mm.py | 10 ++ .../models/modeling_nemotron_h_multimodal.py | 7 + .../_torch/models/modeling_whisper.py | 4 + tensorrt_llm/grpc/openengine/README.md | 15 +- tensorrt_llm/grpc/openengine/capabilities.py | 10 +- tensorrt_llm/inputs/registry.py | 21 +++ .../grpc/openengine/test_multimodal.py | 134 +++++++++++++----- 7 files changed, 160 insertions(+), 41 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 71132d01513f..dddd934c24c7 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -267,6 +267,16 @@ class Gemma4InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInpu # image, when present) can be split across chunks safely. mm_bidirectional_blocks = False + @property + def supported_modalities(self) -> frozenset[str]: + modalities = super().supported_modalities + config = getattr(self, "config", None) + if getattr(config, "vision_config", None) is None: + modalities -= {"image", "video"} + if getattr(config, "audio_config", None) is None: + modalities -= {"audio"} + return modalities + def __init__( self, model_path: str, diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py index ef0608859a6f..9d5dd7ff74b3 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py @@ -1152,6 +1152,13 @@ class NemotronHMultimodalInputProcessor( ): supports_token_id_mm_expansion: ClassVar[bool] = True + @property + def supported_modalities(self) -> frozenset[str]: + modalities = super().supported_modalities + if getattr(self, "_audio_extractor", None) is None: + return modalities - {"audio"} + return modalities + def get_mm_item_metadata( self, prompt_token_ids: List[int], diff --git a/tensorrt_llm/_torch/models/modeling_whisper.py b/tensorrt_llm/_torch/models/modeling_whisper.py index f88e0e75dbdc..ffc7540b28ce 100644 --- a/tensorrt_llm/_torch/models/modeling_whisper.py +++ b/tensorrt_llm/_torch/models/modeling_whisper.py @@ -602,6 +602,10 @@ class WhisperInputProcessor(InputProcessor): # IDs are not consumed by this processor. supports_token_ids_input = False + @property + def supported_modalities(self) -> frozenset[str]: + return frozenset({"audio"}) + def __init__(self, model_path, config, tokenizer, trust_remote_code: bool = True, **kwargs): self.model_path = model_path self.config = config diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index ccc11bbf2a4f..22ba43f6b297 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -78,7 +78,20 @@ The `Inference.Generate` RPC loads the selected model through TensorRT-LLM's PyT Clients must continuously consume the response stream. If response delivery remains stalled for 30 seconds, the server aborts the engine request and terminates the stream with a retryable overload error. -Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, multimodal media, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`. +Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`. + +### Media capability discovery + +`GetModelInfo.supports_multimodal` and `extra.trtllm_supported_modalities` reflect +the loaded processor's read-only `supported_modalities` property, intersected +with registered placeholders and implemented OpenEngine handlers. Both discovery +and request validation use this result. `InputProcessor` defaults to text-only; +`BaseMultimodalInputProcessor` defaults to its registered modalities. Processors +with checkpoint-dependent encoders override the property to narrow that set. +Gemma4 excludes audio without `audio_config` and images/video without +`vision_config`. Nemotron excludes audio without an +initialized audio extractor, and Whisper declares audio only. Registered custom +processors must declare the property to advertise media support. ### Disaggregated serving diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index c324a40492e3..064ac9178395 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -92,17 +92,13 @@ def resolve_media_capabilities( if not callable(input_processor): return MediaCapabilities({}) + supported_modalities = getattr(input_processor, "supported_modalities", frozenset()) placeholders = { modality: MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality) for modality in implemented_modalities - if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) + if modality in supported_modalities + and MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) } - # Gemma4 checkpoints may omit the audio encoder despite sharing a model type. - if ( - model_type == "gemma4" - and getattr(getattr(input_processor, "config", None), "audio_config", None) is None - ): - placeholders.pop("audio", None) # Fail startup before Control advertises a malformed template. for modality, template in placeholders.items(): try: diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index ad3465a96157..d8af206b8fe7 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -103,6 +103,11 @@ class InputProcessor(Protocol): supports_token_ids_input: bool = True """Whether pre-tokenized prompts are accepted; encoder-driven processors may opt out.""" + @property + def supported_modalities(self) -> frozenset[str]: + """Media accepted by this processor instance; text-only by default.""" + return frozenset() + model_path: any config: any tokenizer: any @@ -159,6 +164,15 @@ class BaseMultimodalInputProcessor(ABC): processes the multimodal inputs. """ + @property + def supported_modalities(self) -> frozenset[str]: + """Registered modalities, narrowed by checkpoint-specific overrides.""" + config = getattr(self, "config", None) + if config is None: + return frozenset() + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_modalities( + resolve_top_level_model_type(config)) + # Whether multimodal soft-token runs need bidirectional attention spanning # the full block. When True, the chunked-prefill scheduler must keep each # MM block intact within a single iteration (snap-up or snap-down at chunk @@ -898,6 +912,13 @@ def remove_placeholder_metadata(self, model_type: str): raise ValueError(f"Model type '{model_type}' is not registered") del self._multimodal_placeholder_by_model_type[model_type] + def get_registered_modalities(self, model_type: str) -> frozenset[str]: + """Return registered media types, or an empty set for an unknown model.""" + self._ensure_provider_imported(model_type) + metadata = self._multimodal_placeholder_by_model_type.get(model_type) + return frozenset( + metadata.placeholder_map) if metadata is not None else frozenset() + def is_valid(self, model_type: str, modality: str) -> bool: self._ensure_provider_imported(model_type) return model_type in self._multimodal_placeholder_by_model_type and \ diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index bc425119bde9..a08f1946c768 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -94,6 +94,10 @@ def __init__(self) -> None: class _ImageInputProcessor(BaseMultimodalInputProcessor): """Minimal concrete processor used only for capability detection.""" + @property + def supported_modalities(self): + return frozenset({"image", "video", "audio"}) + def __init__(self, config: _ImageConfig) -> None: self._config = config @@ -127,6 +131,7 @@ class _AudioInputProcessor: """Registered processor that does not use BaseMultimodalInputProcessor.""" _registered_model_type = _ImageConfig.model_type + supported_modalities = frozenset({"image", "video", "audio"}) def __init__(self, config: _ImageConfig) -> None: self.config = config @@ -643,28 +648,40 @@ async def test_audio_request_uses_one_shared_budget_without_double_charging( @pytest.mark.asyncio +@pytest.mark.parametrize("processor_kind", ["gemma4", "nemotron"]) @pytest.mark.parametrize("audio_encoder", [False, True]) -async def test_gemma4_audio_capability_requires_checkpoint_encoder(audio_encoder, monkeypatch): - class GemmaConfig: - model_type = "gemma4" - audio_config = object() if audio_encoder else None - - llm = SimpleNamespace(input_processor=_ImageInputProcessor(GemmaConfig())) - monkeypatch.setattr( - MULTIMODAL_PLACEHOLDER_REGISTRY, - "is_valid", - lambda model_type, modality: modality in {"image", "audio"}, - ) - monkeypatch.setattr( - MULTIMODAL_PLACEHOLDER_REGISTRY, - "get_placeholder", - lambda model_type, modality: f"<{modality}_{{}}>", - ) - capabilities = request_mapping.media_capabilities_for_model(llm) - assert "image" in capabilities.modalities - assert ("audio" in capabilities.modalities) == audio_encoder +async def test_audio_capability_follows_loaded_processor(processor_kind, audio_encoder): + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor + from tensorrt_llm._torch.models.modeling_nemotron_h_multimodal import ( + NemotronHMultimodalInputProcessor, + ) + + if processor_kind == "gemma4": + processor = Gemma4InputProcessor.__new__(Gemma4InputProcessor) + config = SimpleNamespace( + model_type="gemma4", + vision_config=object(), + audio_config=object() if audio_encoder else None, + ) + else: + processor = NemotronHMultimodalInputProcessor.__new__(NemotronHMultimodalInputProcessor) + config = SimpleNamespace(model_type="nemotron_h_omni") + processor._audio_extractor = object() if audio_encoder else None + processor._config = config + llm = _ImageLlm(processor, config) + llm.args.reasoning_parser = None + llm.args.max_seq_len = 2048 + inference = OpenEngineInferenceServicer(llm, "audio-model") + control = OpenEngineControlServicer(llm, "audio-model", inference) + info = await control.GetModelInfo(model_pb2.GetModelInfoRequest(), FakeServicerContext()) + assert ("audio" in info.extra["trtllm_supported_modalities"]) is audio_encoder + assert "image" in info.extra["trtllm_supported_modalities"] + assert {"image", "video"} <= processor.supported_modalities + assert ("audio" in processor.supported_modalities) is audio_encoder if not audio_encoder: request = generation_pb2.GenerateRequest( + request_id="unsupported-audio", + model="audio-model", prompt="transcribe", media=[ generation_pb2.MediaItem( @@ -675,32 +692,83 @@ class GemmaConfig: ], ) with pytest.raises(UnsupportedFeatureError, match="audio"): - await request_mapping._input_from_request(request, llm, capabilities) + await request_mapping._input_from_request(request, llm) + context = FakeServicerContext() + with pytest.raises(AbortError): + async for _ in inference.Generate(request, context): + pass + assert context.abort_code == grpc.StatusCode.UNIMPLEMENTED + + +@pytest.mark.parametrize("audio_encoder", [False, True]) +def test_gemma4_without_vision_encoder_omits_images_and_video(audio_encoder): + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor + + processor = Gemma4InputProcessor.__new__(Gemma4InputProcessor) + processor._config = SimpleNamespace( + model_type="gemma4", vision_config=None, audio_config=object() if audio_encoder else None + ) + expected = {"audio"} if audio_encoder else frozenset() + assert processor.supported_modalities == expected + assert ( + request_mapping.media_capabilities_for_model( + SimpleNamespace(input_processor=processor) + ).modalities + == expected + ) @pytest.mark.parametrize("config_state", ["absent", "none"]) -def test_gemma4_missing_processor_config_fails_closed(config_state, monkeypatch): +def test_gemma4_missing_processor_config_fails_closed(config_state): + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor + config = SimpleNamespace(model_type="gemma4", audio_config=object()) - processor = _ImageInputProcessor(config) - if config_state == "absent": - del processor._config - else: + processor = Gemma4InputProcessor.__new__(Gemma4InputProcessor) + if config_state == "none": processor._config = None llm = SimpleNamespace(input_processor=processor, _hf_model_config=config) - monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", lambda *args: True) - monkeypatch.setattr( - MULTIMODAL_PLACEHOLDER_REGISTRY, - "get_placeholder", - lambda model_type, modality: f"<{modality}_{{}}>", + assert processor.supported_modalities == frozenset() + assert request_mapping.media_capabilities_for_model(llm).modalities == frozenset() + + +def test_whisper_declares_audio_without_multimodal_base(): + from tensorrt_llm._torch.models.modeling_whisper import WhisperInputProcessor + + processor = WhisperInputProcessor.__new__(WhisperInputProcessor) + processor.config = SimpleNamespace(model_type="whisper") + assert processor.supported_modalities == {"audio"} + llm = SimpleNamespace(input_processor=processor) + assert request_mapping.media_capabilities_for_model(llm).modalities == {"audio"} + + +def test_multimodal_processor_defaults_to_registered_modalities(monkeypatch): + from tensorrt_llm.inputs.registry import MultimodalPlaceholderMetadata + + class RegisteredProcessor(_ImageInputProcessor): + supported_modalities = BaseMultimodalInputProcessor.supported_modalities + + monkeypatch.setitem( + MULTIMODAL_PLACEHOLDER_REGISTRY._multimodal_placeholder_by_model_type, + _ImageConfig.model_type, + MultimodalPlaceholderMetadata(placeholder_map={"image": "", "audio": "