From 13e0e056d2375406e677f617ebaf486d9500da63 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 11:16:15 -0700 Subject: [PATCH 01/27] 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 02/27] 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 03/27] 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 04/27] 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 05/27] 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 06/27] 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 07/27] 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 08/27] 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 09/27] 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 10/27] 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 11/27] 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 12/27] 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 13/27] 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 1e3f7f85497657ae1fc7e1edf125f6dc26ccecbf Mon Sep 17 00:00:00 2001 From: krishung5 Date: Mon, 19 Oct 2026 20:53:49 -0700 Subject: [PATCH 14/27] 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 15/27] 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 220478b677426ceda1113ea2642eadc1d0ef81b4 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Tue, 20 Oct 2026 10:47:28 -0700 Subject: [PATCH 16/27] 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 17/27] 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 18/27] 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 19/27] 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 c0e8499a981b28defb27e0d3e5a44faff17ae509 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Wed, 21 Oct 2026 12:03:40 -0700 Subject: [PATCH 20/27] 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 21/27] 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 c63d619da8f40094d365fa86d8c0f15a5028ae62 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 22 Oct 2026 10:12:15 -0700 Subject: [PATCH 22/27] 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 23/27] 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 24/27] 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 25/27] 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 26/27] 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 27/27] 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)