From 13e0e056d2375406e677f617ebaf486d9500da63 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 15 Oct 2026 11:16:15 -0700 Subject: [PATCH 01/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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 a3cff30236e375afce75266badf3eaaebc24e71c Mon Sep 17 00:00:00 2001 From: krishung5 Date: Mon, 19 Oct 2026 11:21:48 -0700 Subject: [PATCH 11/28] 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 1e3f7f85497657ae1fc7e1edf125f6dc26ccecbf Mon Sep 17 00:00:00 2001 From: krishung5 Date: Mon, 19 Oct 2026 20:53:49 -0700 Subject: [PATCH 12/28] 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 220478b677426ceda1113ea2642eadc1d0ef81b4 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Tue, 20 Oct 2026 10:47:28 -0700 Subject: [PATCH 13/28] 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 14/28] 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 c0e8499a981b28defb27e0d3e5a44faff17ae509 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Wed, 21 Oct 2026 12:03:40 -0700 Subject: [PATCH 15/28] 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 c63d619da8f40094d365fa86d8c0f15a5028ae62 Mon Sep 17 00:00:00 2001 From: krishung5 Date: Thu, 22 Oct 2026 10:12:15 -0700 Subject: [PATCH 16/28] 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 17/28] 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 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 18/28] 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 19/28] 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 16048326a524000653cbbd3e4cec1c0a7f7eb46f Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:27:47 -0700 Subject: [PATCH 20/28] [None][feat] Publish image routing hints in OpenEngine KV events Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- .../nanobind/batch_manager/bindings.cpp | 10 ++ .../kv_cache/kv_cache_manager_v2.py | 9 ++ .../_torch/pyexecutor/kv_cache_events.py | 73 ++++++++-- tensorrt_llm/grpc/openengine/capabilities.py | 15 ++ tensorrt_llm/grpc/openengine/control.py | 4 +- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../unittest/grpc/openengine/test_control.py | 22 ++- .../grpc/openengine/test_image_routing.py | 73 ++++++++++ .../test_image_routing_conformance.py | 130 ++++++++++++++++++ .../test_kv_cache_event_manager.py | 41 ++++++ 11 files changed, 366 insertions(+), 13 deletions(-) create mode 100644 tests/unittest/grpc/openengine/test_image_routing.py create mode 100644 tests/unittest/grpc/openengine/test_image_routing_conformance.py diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 087e7ddae60c..8e8a83bd03f2 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -238,6 +238,16 @@ void initBindings(nb::module_& m) } return hashes; }) + .def_prop_ro("multimodal_uuids", + [](GenLlmReq& self) + { + std::optional>> uuids = std::nullopt; + if (self.getMultimodalUuids()) + { + uuids = *self.getMultimodalUuids().value(); + } + return uuids; + }) .def_prop_ro("multimodal_positions", [](GenLlmReq& self) { diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 4930afc9e55e..02ed2f36a81e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -5689,6 +5689,15 @@ def _augment_tokens_for_block_reuse( "reuse." ) + if isinstance(self.event_manager, StreamingKVCacheEventManager): + uuids = req.multimodal_uuids + if uuids is not None: + self.event_manager.remember_multimodal_uuids( + (_hash_to_digest(digest), uuid) + for digest, uuid in zip(req.multimodal_hashes, uuids, strict=True) + if uuid is not None + ) + # Multimodal path: materialize a Python-int list (digest bytes get spliced in below), # which flows through the per-element binding fallback. tokens may be a zero-copy numpy # int32 view (get_tokens_view) — use tolist() so elements are Python ints, not np.int32. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 29533afa3c37..bf8cab5d4f45 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -29,10 +29,10 @@ import time import traceback from abc import ABC, abstractmethod -from collections import deque +from collections import OrderedDict, deque from itertools import count from queue import Queue -from typing import Any, Iterator, Optional +from typing import Any, Iterable, Iterator, Optional import msgspec import zmq @@ -80,6 +80,15 @@ class MultimodalKey( hash: str start_offset: int + uuid: str | None = None + + +class ImageRoutingMetadata(msgspec.Struct, frozen=True, gc=False): + """Versioned digest-first image normalization metadata for KV routers.""" + + version: int + token_id_offset: int + keys: list[list[MultimodalKey]] class BlockStored(KVCacheWireEvent): @@ -99,6 +108,7 @@ class BlockStored(KVCacheWireEvent): locality: str | None = None # Aligned one-for-one with block_hashes when multimodal decoding is enabled. mm_keys: list[list[MultimodalKey]] | None = None + trtllm_image_routing: ImageRoutingMetadata | None = None class BlockRemoved(KVCacheWireEvent): @@ -551,6 +561,10 @@ def __init__( raise ValueError("mm_token_id_offset must be non-negative") self._block_size = block_size self._include_mm_keys = mm_token_id_offset is not None + self._mm_token_id_offset = mm_token_id_offset + self._uuid_capacity = max_entries + self._uuids: OrderedDict[bytes, str] = OrderedDict() + self._uuid_lock = threading.Lock() self._event_sink = kv_cache_manager_v2_runtime.StreamingEventSink( max_entries=max_entries, mm_token_id_offset=mm_token_id_offset, @@ -560,6 +574,28 @@ def __init__( def event_sink(self) -> kv_cache_manager_v2_runtime.StreamingEventSink: return self._event_sink + def remember_multimodal_uuids(self, items: Iterable[tuple[bytes, str]]) -> None: + """Retain bounded routing hints without changing native cache keys.""" + for digest, uuid in items: + if len(uuid.encode("utf-8")) > 256: + continue + with self._uuid_lock: + previous = self._uuids.get(digest) + if previous is not None and previous != uuid: + raise ValueError("Conflicting UUIDs for a multimodal content digest") + self._uuids[digest] = uuid + self._uuids.move_to_end(digest) + while len(self._uuids) > self._uuid_capacity: + self._uuids.popitem(last=False) + + def _multimodal_key(self, key: tuple) -> MultimodalKey: + digest = bytes(key[0]) + with self._uuid_lock: + uuid = key[2] if len(key) > 2 else None + if uuid is None: + uuid = self._uuids.get(digest) + return MultimodalKey(hash=digest.hex(), start_offset=key[1], uuid=uuid) + def set_target_life_cycle(self, life_cycle_id: int) -> None: self._event_sink.set_target_life_cycle(life_cycle_id) @@ -567,6 +603,14 @@ def drain_events(self) -> list[BlockStored | BlockRemoved | AllBlocksCleared]: result: list[BlockStored | BlockRemoved | AllBlocksCleared] = [] for event in self._event_sink.drain_iteration_events(): if isinstance(event, kv_cache_manager_v2_runtime.StreamingBlockStoredData): + mm_keys = ( + [ + [self._multimodal_key(key) for key in block_keys] + for block_keys in event.mm_keys + ] + if self._include_mm_keys + else None + ) result.append( BlockStored( block_hashes=list(event.block_hashes), @@ -576,15 +620,18 @@ def drain_events(self) -> list[BlockStored | BlockRemoved | AllBlocksCleared]: lora_id=event.lora_id, medium="GPU", lora_name=None, - mm_keys=( - [ - [ - MultimodalKey(hash=bytes(key[0]).hex(), start_offset=key[1]) - for key in block_keys - ] - for block_keys in event.mm_keys - ] - if self._include_mm_keys + mm_keys=mm_keys, + trtllm_image_routing=( + ImageRoutingMetadata( + version=1, token_id_offset=self._mm_token_id_offset, keys=mm_keys + ) + if mm_keys is not None + and self._mm_token_id_offset is not None + and any( + key.uuid is not None and key.uuid.startswith("dynamo-image-v1:") + for block_keys in mm_keys + for key in block_keys + ) else None ), ) @@ -646,6 +693,10 @@ def start(self) -> None: """Bind the publisher's sockets and start its background thread.""" self._publisher.start() + def remember_multimodal_uuids(self, items: Iterable[tuple[bytes, str]]) -> None: + """Associate native MM digests with optional request identities.""" + self._event_source.remember_multimodal_uuids(items) + def get_source_descriptor(self) -> dict[str, str | int] | None: """Describe a live ZMQ publisher for OpenEngine node discovery.""" publisher = self._publisher diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 83a6f9979b51..c2e5e879742a 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -67,6 +67,21 @@ def _loaded_model_type(llm: Any) -> str: return resolve_top_level_model_type(llm.input_processor.config) +def supports_image_routing(llm: Any, modalities: Collection[str]) -> bool: + """Whether the loaded model exposes the validated image KV-event contract.""" + processor_config = getattr(getattr(llm, "input_processor", None), "config", None) + if "image" not in modalities or resolve_top_level_model_type(processor_config) != "qwen3_vl": + return False + cache = getattr(llm.args, "kv_cache_config", None) + events = getattr(cache, "kv_events_config", None) + return ( + getattr(cache, "use_kv_cache_manager_v2", False) is True + and getattr(cache, "enable_block_reuse", False) is True + and getattr(events, "enable_kv_cache_events", False) is True + and getattr(events, "publisher", None) == "zmq" + ) + + def resolve_media_capabilities( llm: Any, implemented_modalities: Collection[str] ) -> MediaCapabilities: diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index 8b4292dd33ae..392843377290 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, supports_image_routing from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError from .errors import AbortFailedError from .kv_events import data_parallel_size, events_config @@ -297,6 +297,8 @@ async def GetModelInfo( reasoning_parser=str(self._llm.args.reasoning_parser or ""), ) info.extra.update({"trtllm_supported_modalities": sorted(media_modalities)}) + if supports_image_routing(self._llm, media_modalities): + info.extra.update({"trtllm_image_routing_version": 1}) max_seq_len = _positive_int(self._llm.args.max_seq_len) if max_seq_len is not None: diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index f5935fd9a92b..024810ae5d43 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -69,6 +69,7 @@ l0_a10: # 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/grpc/openengine/test_image_routing_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 bcf16bc75a4f..b0412cbd2835 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -191,6 +191,7 @@ l0_cpu: - unittest/_torch/test_shared_tensor_cpu.py - unittest/grpc/openengine/test_frontend_lifecycle.py - unittest/grpc/openengine/test_control.py + - unittest/grpc/openengine/test_image_routing.py - unittest/grpc/openengine/test_multimodal.py - unittest/grpc/openengine/test_node.py - unittest/grpc/openengine/test_openengine.py diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 79b56c77f9d5..b0a75486ecaf 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -19,7 +19,10 @@ model_pb2, server_pb2, ) -from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 +from tensorrt_llm.grpc.openengine.capabilities import ( # noqa: E402 + MediaCapabilities, + supports_image_routing, +) 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 @@ -30,6 +33,23 @@ MODEL = "test-model" +@pytest.mark.parametrize("model_type", ["qwen3_vl", "qwen2_vl", "text"]) +@pytest.mark.parametrize("enabled", [False, True]) +def test_image_routing_capability_requires_streaming_v2_and_validated_model(model_type, enabled): + llm = SimpleNamespace( + input_processor=SimpleNamespace(config=SimpleNamespace(model_type=model_type)), + args=SimpleNamespace( + kv_cache_config=SimpleNamespace( + use_kv_cache_manager_v2=enabled, + enable_block_reuse=True, + kv_events_config=SimpleNamespace(enable_kv_cache_events=True, publisher="zmq"), + ) + ), + ) + assert supports_image_routing(llm, {"image"}) is (enabled and model_type == "qwen3_vl") + assert not supports_image_routing(llm, set()) + + class _FakeHandle: def __init__(self, fail: bool = False) -> None: self.aborted = False diff --git a/tests/unittest/grpc/openengine/test_image_routing.py b/tests/unittest/grpc/openengine/test_image_routing.py new file mode 100644 index 000000000000..bbd29ba36989 --- /dev/null +++ b/tests/unittest/grpc/openengine/test_image_routing.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Image routing hints preserve engine identities and survive the event wire.""" + +from types import SimpleNamespace + +import msgspec +import pytest +from PIL import Image + +from tensorrt_llm._torch.pyexecutor import kv_cache_events as events +from tensorrt_llm.inputs.multimodal import apply_mm_hashes +from tensorrt_llm.llmapi.llm_args import KVEventsConfig + +pytestmark = pytest.mark.cpu_only + + +def test_streaming_image_hints_are_bounded_and_content_sensitive(monkeypatch): + uuid = "dynamo-image-v1:0123456789abcdef" + red = Image.new("RGB", (32, 32), "red") + blue = Image.new("RGB", (32, 32), "blue") + hashes = [ + apply_mm_hashes({"image": [image]}, {"image": [uuid]}).hashes["image"][0] + for image in (red, blue) + ] + assert hashes[0] != hashes[1], "routing hints must not replace content-sensitive cache keys" + + manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + source = manager._event_source + native_sink = source._event_sink + try: + manager.remember_multimodal_uuids([(bytes.fromhex(hashes[0]), uuid)]) + + class Stored: + block_hashes = [123] + parent_block_hash = None + token_ids = [1, hashes[0], 1001, 7] + lora_id = None + mm_keys = [[(bytes.fromhex(hashes[0]), 0)]] + + monkeypatch.setattr(events.kv_cache_manager_v2_runtime, "StreamingBlockStoredData", Stored) + source._event_sink = SimpleNamespace(drain_iteration_events=lambda: [Stored()]) + wire = msgspec.msgpack.decode(msgspec.msgpack.encode(source.drain_events()[0])) + metadata = wire["trtllm_image_routing"] + assert metadata["version"] == 1 and metadata["token_id_offset"] == 1000 + assert metadata["keys"][0][0]["uuid"] == uuid + assert metadata["keys"][0][0]["hash"] == hashes[0] + assert wire["token_ids"] == Stored.token_ids + + Stored.token_ids = [1, 2, 3, 4] + Stored.mm_keys = [[]] + text_wire = msgspec.msgpack.decode(msgspec.msgpack.encode(source.drain_events()[0])) + assert "trtllm_image_routing" not in text_wire + + manager.remember_multimodal_uuids([(bytes.fromhex(hashes[1]), uuid)]) + assert len(source._uuids) == 1 + assert source._multimodal_key((bytes.fromhex(hashes[0]), 0)).uuid is None + assert source._multimodal_key((bytes.fromhex(hashes[1]), 0)).uuid == uuid + manager.remember_multimodal_uuids([(bytes.fromhex(hashes[0]), "x" * 257)]) + assert len(source._uuids) == 1 + with pytest.raises(ValueError, match="Conflicting UUIDs"): + manager.remember_multimodal_uuids([(bytes.fromhex(hashes[1]), "other")]) + finally: + source._event_sink = native_sink + manager.shutdown() diff --git a/tests/unittest/grpc/openengine/test_image_routing_conformance.py b/tests/unittest/grpc/openengine/test_image_routing_conformance.py new file mode 100644 index 000000000000..abbd93effeb3 --- /dev/null +++ b/tests/unittest/grpc/openengine/test_image_routing_conformance.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Real-weight image cache reuse and streaming routing-hint conformance.""" + +import asyncio +import base64 +from io import BytesIO + +import grpc +import msgspec +import pytest +import torch +import zmq +from PIL import Image +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM +from tensorrt_llm.grpc.openengine.bindings import generation_pb2 as pb +from tensorrt_llm.grpc.openengine.bindings import model_pb2, openengine_pb2_grpc +from tensorrt_llm.grpc.openengine.server import OpenEngineServer +from tensorrt_llm.llmapi import KvCacheConfig, KVEventsConfig + +pytestmark = pytest.mark.threadleak(enabled=False) + + +def test_real_image_routing_keeps_cache_identity_content_sensitive(unused_tcp_port, tmp_path): + if not torch.cuda.is_available(): + pytest.skip("requires a GPU and the real Qwen3-VL-2B checkpoint") + endpoint = f"tcp://127.0.0.1:{unused_tcp_port}" + context = zmq.Context() + subscriber = context.socket(zmq.SUB) + subscriber.setsockopt(zmq.SUBSCRIBE, b"") + subscriber.connect(endpoint) + llm = None + try: + llm = LLM( + model=str(llm_models_root(check=True) / "Qwen3/Qwen3-VL-2B-Instruct"), + max_batch_size=1, + max_seq_len=2048, + max_num_tokens=2048, + kv_cache_config=KvCacheConfig( + max_tokens=4096, + free_gpu_memory_fraction=0.3, + enable_block_reuse=True, + use_kv_cache_manager_v2=True, + kv_events_config=KVEventsConfig( + enable_kv_cache_events=True, publisher="zmq", endpoint=endpoint + ), + ), + ) + + async def run(): + server = OpenEngineServer(host="127.0.0.1", port=0, llm=llm, model="image-routing") + try: + await server.start() + async with grpc.aio.insecure_channel(f"127.0.0.1:{server.port}") as channel: + control = openengine_pb2_grpc.ControlStub(channel) + info = await control.GetModelInfo( + model_pb2.GetModelInfoRequest(model="image-routing") + ) + assert info.extra["trtllm_image_routing_version"] == 1 + inference = openengine_pb2_grpc.InferenceStub(channel) + + async def generate(color, request_id): + image = BytesIO() + Image.new("RGB", (224, 224), color).save(image, format="PNG") + request = pb.GenerateRequest( + model="image-routing", + request_id=request_id, + prompt=( + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + "Name the solid background color. Answer with one color word." + "<|im_end|>\n<|im_start|>assistant\n" + ), + sampling=pb.SamplingParams(temperature=0), + stopping=pb.StoppingOptions(max_tokens=16), + media=[ + pb.MediaItem( + modality=pb.MODALITY_IMAGE, + uuid="dynamo-image-v1:0123456789abcdef", + data_uri="data:image/png;base64," + + base64.b64encode(image.getvalue()).decode(), + ) + ], + ) + ids, usages, finished = [], [], 0 + async for response in inference.Generate(request, timeout=180): + assert not response.HasField("error"), response + if response.HasField("token"): + ids.extend(token.token_id for token in response.token.tokens) + if response.HasField("usage"): + usages.append(response.usage) + finished += int(response.HasField("finished")) + assert finished == 1 and usages + answer = llm.tokenizer.decode(ids).lower() + assert color in answer, answer + return usages[-1].cached_prompt_tokens + + first = await generate("red", "red-first") + repeated = await generate("red", "red-repeat") + changed = await generate("blue", "blue-same-hint") + assert repeated > first + assert changed < repeated, "changed content reused the first image's KV prefix" + finally: + await server.stop() + + asyncio.run(run()) + captured, hashes = [], set() + while subscriber.poll(1000): + frames = subscriber.recv_multipart() + captured.append(frames[-1]) + batch = msgspec.msgpack.decode(frames[-1]) + for event in batch[1]: + metadata = event.get("trtllm_image_routing") + if not metadata: + continue + for keys in metadata["keys"]: + for key in keys: + assert key["uuid"] == "dynamo-image-v1:0123456789abcdef" + hashes.add(key["hash"]) + assert len(hashes) == 2, ( + "both content digests must retain the routing hint on the event wire" + ) + (tmp_path / "image-routing-events.msgpack").write_bytes(msgspec.msgpack.encode(captured)) + finally: + if llm is not None: + llm.shutdown() + subscriber.close(linger=0) + context.term() diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py index 8198529e4069..fee1ec228ac8 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -393,6 +393,12 @@ def test_native_streaming_sink_preserves_multimodal_event_data(real_block_factor make_block = real_block_factory(event_sink, tokens_per_block=4) digest_a = bytes(range(32)) digest_b = bytes(reversed(range(32))) + manager.remember_multimodal_uuids( + [ + (digest_a, "dynamo-image-v1:0123456789abcdef"), + (digest_b, "dynamo-image-v1:fedcba9876543210"), + ] + ) first = make_block([1, digest_a, 1001, 1002], [4]) gap = make_block([2, 3, 4, 5], [4], parent=first) continued = make_block([1003, 7, 1004, 1005], [4], parent=gap) @@ -432,6 +438,41 @@ def test_native_streaming_sink_preserves_multimodal_event_data(real_block_factor [(digest_a.hex(), 6), (digest_b.hex(), 0)], ] assert manager.stored_blocks == 4 + assert stored.trtllm_image_routing.version == 1 + assert stored.trtllm_image_routing.token_id_offset == 1000 + assert [[key.uuid for key in keys] for keys in stored.trtllm_image_routing.keys] == [ + ["dynamo-image-v1:0123456789abcdef"], + [], + ["dynamo-image-v1:0123456789abcdef"] * 2, + ["dynamo-image-v1:0123456789abcdef", "dynamo-image-v1:fedcba9876543210"], + ] + finally: + manager.shutdown() + + +def test_streaming_uuid_hints_are_bounded_and_do_not_replace_content_hashes(): + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + a, b = bytes(range(32)), bytes(reversed(range(32))) + try: + manager.remember_multimodal_uuids([(a, "same-routing-hint"), (b, "same-routing-hint")]) + source = manager._event_source + assert len(source._uuids) == 1 + first = source._multimodal_key((a, 0)) + second = source._multimodal_key((b, 0)) + assert first.uuid is None + assert second.uuid == "same-routing-hint" + assert first.hash != second.hash + manager.remember_multimodal_uuids([(a, "x" * 257)]) + assert len(source._uuids) == 1 + with pytest.raises(ValueError, match="Conflicting UUIDs"): + manager.remember_multimodal_uuids([(b, "different-hint")]) finally: manager.shutdown() From 3c872a49711c7099be132ebe08cfa1e726d7b325 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:42:08 -0700 Subject: [PATCH 21/28] [None][fix] Gate image routing and keep UUID metadata client-neutral Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_events.py | 4 +--- tensorrt_llm/grpc/openengine/README.md | 24 ++++++++++++++++++- tensorrt_llm/grpc/openengine/capabilities.py | 3 ++- .../unittest/grpc/openengine/test_control.py | 2 ++ .../grpc/openengine/test_image_routing.py | 4 ++-- 5 files changed, 30 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index bf8cab5d4f45..d14663ec64a5 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -628,9 +628,7 @@ def drain_events(self) -> list[BlockStored | BlockRemoved | AllBlocksCleared]: if mm_keys is not None and self._mm_token_id_offset is not None and any( - key.uuid is not None and key.uuid.startswith("dynamo-image-v1:") - for block_keys in mm_keys - for key in block_keys + key.uuid is not None for block_keys in mm_keys for key in block_keys ) else None ), diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index 5ce68cf169c1..a801b3cee8a2 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -73,7 +73,7 @@ The `Inference.Generate` RPC loads the selected model through TensorRT-LLM's PyT Clients must continuously consume the response stream. If response delivery remains stalled for 30 seconds, the server aborts the engine request and terminates the stream with a retryable overload error. -Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, multimodal media, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`. +Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`. ### Disaggregated serving @@ -118,6 +118,28 @@ Enable load reporting independently with `--grpc --grpc-protocol openengine --op For attention DP, each publisher binds `base_port + global_rank`. When event streaming is enabled, OpenEngine collects the live publisher descriptors at executor startup and advertises every rank, including remote nodes. Those hostnames and ports must be reachable from the subscriber. A missing publisher fails startup; an unavailable discovery snapshot fails the RPC rather than synthesizing socket addresses. PP and CP streaming remain unsupported. The extra placement collection is OpenEngine-only; native HTTP serving and OpenEngine with events disabled do not enable it. +### Image routing metadata + +`GetModelInfo.extra.trtllm_image_routing_version = 1` advertises image routing +metadata for aggregated Qwen3-VL with KV cache manager V2 explicitly enabled, +block reuse enabled, and ZMQ event publishing. Image passthrough separately +requires `supports_multimodal` and `extra.trtllm_supported_modalities` to advertise +image support. Video, audio, and disaggregated image routing are not advertised. + +For image blocks with caller UUIDs, `BlockStored.trtllm_image_routing` contains +`version`, `token_id_offset` (the vocabulary size), and `keys`, aligned with the +stored blocks. Each key carries its native content `hash`, `start_offset`, and +optional `uuid`. Native image tokens start with a digest and continue with IDs +above `token_id_offset`; a router can normalize those tokens using identities it +sent in `MediaItem.uuid`. Native block hashes, parent chains, and removals retain +their existing meaning. Cache hashing continues to include image content. + +The publisher retains at most `max_entries` digest-to-UUID associations and skips +UUIDs longer than 256 UTF-8 bytes. Eviction can leave UUIDs unavailable. Consumers +must ignore unmapped image stores while continuing to process removals. Routing +identities must refer to stable media content. This extension does not change the +OpenEngine schema. + ### Node-local discovery OpenEngine exposes optional, versioned `ServerInfo.extra.trtllm_node` metadata without changing the protocol schema. It describes a shared engine incarnation (`engine_id`), local hostname (`node_id`), whether this is the leader, node count, KV block size, local publishing DP ranks, and the global rank-to-node ownership map. The engine ID matches `GetServerInfo.instance_id` on the serving ingress and node metadata listeners. Consumers can use this information to partition telemetry collection without introducing routing-framework-specific server options. diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index c2e5e879742a..e407448aebf7 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -75,7 +75,8 @@ def supports_image_routing(llm: Any, modalities: Collection[str]) -> bool: cache = getattr(llm.args, "kv_cache_config", None) events = getattr(cache, "kv_events_config", None) return ( - getattr(cache, "use_kv_cache_manager_v2", False) is True + getattr(llm.args, "cache_transceiver_config", None) is None + and getattr(cache, "use_kv_cache_manager_v2", False) is True and getattr(cache, "enable_block_reuse", False) is True and getattr(events, "enable_kv_cache_events", False) is True and getattr(events, "publisher", None) == "zmq" diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 11820772255d..9ea851ac9812 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -49,6 +49,8 @@ def test_image_routing_capability_requires_streaming_v2_and_validated_model(mode ) assert supports_image_routing(llm, {"image"}) is (enabled and model_type == "qwen3_vl") assert not supports_image_routing(llm, set()) + llm.args.cache_transceiver_config = object() + assert not supports_image_routing(llm, {"image"}) class _FakeHandle: diff --git a/tests/unittest/grpc/openengine/test_image_routing.py b/tests/unittest/grpc/openengine/test_image_routing.py index bbd29ba36989..b787fb83eb73 100644 --- a/tests/unittest/grpc/openengine/test_image_routing.py +++ b/tests/unittest/grpc/openengine/test_image_routing.py @@ -16,8 +16,8 @@ pytestmark = pytest.mark.cpu_only -def test_streaming_image_hints_are_bounded_and_content_sensitive(monkeypatch): - uuid = "dynamo-image-v1:0123456789abcdef" +@pytest.mark.parametrize("uuid", ["dynamo-image-v1:0123456789abcdef", "caller-owned"]) +def test_streaming_image_hints_are_bounded_and_content_sensitive(monkeypatch, uuid): red = Image.new("RGB", (32, 32), "red") blue = Image.new("RGB", (32, 32), "blue") hashes = [ From 1fffd609c9028d048c9918b1eb7a0aaf21c7e561 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Tue, 27 Oct 2026 12:32:19 -0700 Subject: [PATCH 22/28] [None][fix] Default multimodal processor media IO preferences Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- tensorrt_llm/inputs/registry.py | 4 ++++ tests/unittest/grpc/openengine/test_multimodal.py | 9 +++++++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 6a4bf3d6f8f0..412417a0f051 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -174,6 +174,10 @@ class BaseMultimodalInputProcessor(ABC): # uses `_init_tokenization_cache` and `_encode_with_tokenization_cache`. supports_tokenization_cache: ClassVar[bool] = False + def get_preferred_media_io_kwargs(self) -> dict[str, dict[str, Any]]: + """Return per-modality decode defaults, below server and request overrides.""" + return {} + def get_mm_item_metadata( self, prompt_token_ids: List[int], diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 0abaae084841..be38a0fb04f2 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -539,8 +539,13 @@ async def test_text_only_input_mapping_is_unchanged(): @pytest.mark.asyncio -async def test_image_sources_preserve_order_uuid_and_model_decode_defaults(image_llm, monkeypatch): +@pytest.mark.parametrize("model_defaults", [None, {"format": "np"}]) +async def test_image_sources_preserve_order_uuid_and_model_decode_defaults( + image_llm, monkeypatch, model_defaults +): llm, _ = image_llm + if model_defaults is None: + monkeypatch.delattr(_ImageInputProcessor, "get_preferred_media_io_kwargs") created_with = [] class _ImageIo: @@ -583,7 +588,7 @@ async def async_load_bytes(self, data): mapped = await request_mapping._input_from_request(request, llm) - assert created_with == [({"format": "np"}, None)] + assert created_with == [(model_defaults, None)] assert mapped["prompt"] == request.prompt assert mapped["multi_modal_data"] == { "image": [ From d83bef8a2d1f2ece2d403acdcab784208083a05b Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Tue, 27 Oct 2026 20:51:45 -0700 Subject: [PATCH 23/28] [None][test] Fix KV-cache, ADP and mRoPE test fixtures Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- .../kv_cache/test_kv_cache_v2_first_new_block_probe.py | 1 + .../executor/kv_cache/test_kv_cache_v2_multimodal_runs.py | 1 + tests/unittest/_torch/executor/test_pytorch_model_engine.py | 2 ++ tests/unittest/_torch/test_connector_adp.py | 5 ++++- 4 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py index 855289820381..41683b60e3a2 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py @@ -52,6 +52,7 @@ def make_stub_manager( # for a one-model draft, which reads D prompt tokens past the target's end. mgr.reuse_match_backoff = reuse_match_backoff mgr.vocab_size = 32000 + mgr.event_manager = None mgr.conversation_manager = None mgr.kv_connector_manager = None mgr.kv_cache_map = {} diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py index b31abb20a596..e037230de177 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py @@ -23,6 +23,7 @@ def _make_manager(vocab_size: int): manager = KVCacheManagerV2.__new__(KVCacheManagerV2) manager.vocab_size = vocab_size + manager.event_manager = None return manager diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 6bae538cca81..28e4abe907f7 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1904,6 +1904,7 @@ def test_mrope_multi_beam_generation_positions(self) -> None: """Every beam row gets its request's `past + delta` on all 3 axes.""" beam_width = 2 engine, kv_cache_manager = create_model_engine_and_kvcache( + self, llm_args=TorchLlmArgs(model="dummy", max_batch_size=4, max_num_tokens=64, @@ -1973,6 +1974,7 @@ def test_mrope_mixed_batched_and_fallback_positions_match_fallback( self) -> None: """Mixed scalar batching matches the old per-request writes exactly.""" engine, kv_cache_manager = create_model_engine_and_kvcache( + self, llm_args=TorchLlmArgs(model="dummy", max_batch_size=2, max_num_tokens=64), diff --git a/tests/unittest/_torch/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index ad1e6829630e..3e778e0445c0 100644 --- a/tests/unittest/_torch/test_connector_adp.py +++ b/tests/unittest/_torch/test_connector_adp.py @@ -80,10 +80,12 @@ def _manager(rank: int = 1) -> KvCacheConnectorManager: def _configure_recovery_controls(executor: "PyExecutor") -> None: from tensorrt_llm._torch.pyexecutor.executor_request_queue import ExecutorRequestQueue - executor.dist = MagicMock(tp_size=1, cp_size=1, pp_size=1) + executor.dist = MagicMock(tp_size=1, cp_size=1, pp_size=1, tp_rank=0, has_cp_helix=False) + executor.dist.tp_cp_allgather_int64.side_effect = executor.dist.tp_allgather_int64 executor.enable_attention_dp = True executor.enable_early_first_token_response = False executor.kv_connector_manager = None + executor.kv_cache_transceiver = None executor._is_kv_manager_v2 = False executor._kv_load_tracker = KvLoadTracker(None, rank=0, enabled=False) executor._kv_cache_capacity = {} @@ -414,6 +416,7 @@ def test_batch_start_combines_readiness_with_status(batch_sizes: tuple[int, int] _configure_recovery_controls(executor) executor.kv_connector_manager = _manager(rank) executor.dist.tp_size = 2 + executor.dist.tp_rank = rank executor.dist.tp_allgather_int64.side_effect = partial(collective.gather_int64, rank) executors.append(executor) From 756422c354491f904cb4dcc1f907bd5d6ee1a987 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Tue, 27 Oct 2026 21:15:54 -0700 Subject: [PATCH 24/28] [None][fix] Validate image routing UUID pairs before registration Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 5 +- ...st_kv_cache_manager_v2_helix_superblock.py | 10 +++ .../unittest/grpc/openengine/test_control.py | 51 ++++++++++++-- .../grpc/openengine/test_image_routing.py | 67 +++++++++++++++++++ 4 files changed, 127 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index f99e0a7ee546..4408059d688f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6045,11 +6045,12 @@ def _augment_tokens_for_block_reuse( if isinstance(self.event_manager, StreamingKVCacheEventManager): uuids = req.multimodal_uuids if uuids is not None: - self.event_manager.remember_multimodal_uuids( + uuid_pairs = [ (_hash_to_digest(digest), uuid) for digest, uuid in zip(req.multimodal_hashes, uuids, strict=True) if uuid is not None - ) + ] + self.event_manager.remember_multimodal_uuids(uuid_pairs) # Multimodal path: materialize a Python-int list (digest bytes get spliced in below), # which flows through the per-element binding fallback. tokens may be a zero-copy numpy diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py index 7f9d7ef36e06..213e0acc09b9 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py @@ -429,6 +429,16 @@ def test_resize_context_rejects_disagg_helix_request(prompt_len: int) -> None: ) +def test_resize_context_allows_disagg_helix_dummy_request() -> None: + resizes = [] + req = _context_req(prompt_len=200, total_input_len_cp=200) + req.is_dummy_request = True + assert KVCacheManagerV2.resize_context( + _context_resize_mgr(is_disagg=True, resizes=resizes), req, num_tokens=32 + ) + assert resizes == [64 + 32] + + def test_dummy_frozen_fields_sum_invariant() -> None: """The frozen dummy fiction (last rank active, shared synthetic global length) keeps the same books as real requests: per-rank lengths sum to diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 9ea851ac9812..70704911d1d3 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -36,18 +36,33 @@ @pytest.mark.parametrize("model_type", ["qwen3_vl", "qwen2_vl", "text"]) @pytest.mark.parametrize("enabled", [False, True]) -def test_image_routing_capability_requires_streaming_v2_and_validated_model(model_type, enabled): +@pytest.mark.parametrize( + "block_reuse,events_enabled,publisher", + [(True, True, "zmq"), (False, True, "zmq"), (True, False, "zmq"), (True, True, "null")], +) +def test_image_routing_capability_requires_streaming_v2_and_validated_model( + model_type, enabled, block_reuse, events_enabled, publisher +): llm = SimpleNamespace( input_processor=SimpleNamespace(config=SimpleNamespace(model_type=model_type)), args=SimpleNamespace( kv_cache_config=SimpleNamespace( use_kv_cache_manager_v2=enabled, - enable_block_reuse=True, - kv_events_config=SimpleNamespace(enable_kv_cache_events=True, publisher="zmq"), + enable_block_reuse=block_reuse, + kv_events_config=SimpleNamespace( + enable_kv_cache_events=events_enabled, publisher=publisher + ), ) ), ) - assert supports_image_routing(llm, {"image"}) is (enabled and model_type == "qwen3_vl") + expected = ( + enabled + and model_type == "qwen3_vl" + and block_reuse + and events_enabled + and publisher == "zmq" + ) + assert supports_image_routing(llm, {"image"}) is expected assert not supports_image_routing(llm, set()) llm.args.cache_transceiver_config = object() assert not supports_image_routing(llm, {"image"}) @@ -211,6 +226,34 @@ async def test_model_info_advertises_loaded_image_support(monkeypatch): assert list(info.extra["trtllm_supported_modalities"]) == ["image"] +@pytest.mark.asyncio +@pytest.mark.parametrize("events_enabled", [False, True]) +async def test_model_info_advertises_image_routing_only_with_events(monkeypatch, events_enabled): + llm = _llm( + kv_cache_config=SimpleNamespace( + use_kv_cache_manager_v2=True, + enable_block_reuse=True, + tokens_per_block=32, + kv_events_config=SimpleNamespace( + enable_kv_cache_events=events_enabled, publisher="zmq" + ), + ) + ) + llm.input_processor = SimpleNamespace(config=SimpleNamespace(model_type="qwen3_vl")) + inference = OpenEngineInferenceServicer(llm, MODEL) + monkeypatch.setattr( + inference, "_media_capabilities", MediaCapabilities({"image": ""}) + ) + info = await OpenEngineControlServicer(llm, MODEL, inference).GetModelInfo( + model_pb2.GetModelInfoRequest(model=MODEL), FakeServicerContext() + ) + assert info.supports_multimodal + if events_enabled: + assert info.extra["trtllm_image_routing_version"] == 1 + else: + assert "trtllm_image_routing_version" not in info.extra + + @pytest.mark.asyncio async def test_model_info_leaves_the_context_window_unset_when_unknown(): """An unset limit means "unknown"; advertising zero would be a real claim.""" diff --git a/tests/unittest/grpc/openengine/test_image_routing.py b/tests/unittest/grpc/openengine/test_image_routing.py index b787fb83eb73..6e6c0c6d31c4 100644 --- a/tests/unittest/grpc/openengine/test_image_routing.py +++ b/tests/unittest/grpc/openengine/test_image_routing.py @@ -10,6 +10,7 @@ from PIL import Image from tensorrt_llm._torch.pyexecutor import kv_cache_events as events +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm.inputs.multimodal import apply_mm_hashes from tensorrt_llm.llmapi.llm_args import KVEventsConfig @@ -71,3 +72,69 @@ class Stored: finally: source._event_sink = native_sink manager.shutdown() + + +@pytest.mark.parametrize( + "hashes,uuids,error", + [ + ([tuple(range(8)), tuple(range(8, 16))], ["image-a"], "shorter"), + ([tuple(range(8))], ["image-a", "image-b"], "longer"), + ([tuple(range(8)), (1, 2, 3)], ["image-a", "image-b"], "Expected 8 int32"), + ], +) +def test_invalid_request_uuid_pairs_leave_routing_hints_unchanged(hashes, uuids, error): + event_manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + try: + event_manager.remember_multimodal_uuids([(b"x" * 32, "existing")]) + before = event_manager._event_source._uuids.copy() + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + manager.vocab_size = 1000 + manager.event_manager = event_manager + request = SimpleNamespace( + multimodal_hashes=hashes, + multimodal_uuids=uuids, + multimodal_positions=[0] * len(hashes), + multimodal_lengths=[1] * len(hashes), + ) + with pytest.raises(ValueError, match=error): + manager._augment_tokens_for_block_reuse([1, 2], request) + assert event_manager._event_source._uuids == before + finally: + event_manager.shutdown() + + +def test_request_uuid_pairs_preserve_partial_uuid_slots(): + hashes = [tuple(range(8)), tuple(range(8, 16))] + event_manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + try: + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + manager.vocab_size = 1000 + manager.event_manager = event_manager + request = SimpleNamespace( + multimodal_hashes=hashes, + multimodal_uuids=["image-a", None], + multimodal_positions=[0, 1], + multimodal_lengths=[1, 1], + multimodal_item_run_cu_offsets=None, + multimodal_run_positions=None, + multimodal_run_lengths=None, + ) + manager._augment_tokens_for_block_reuse([1, 2], request) + digest = b"".join(value.to_bytes(4, "big", signed=True) for value in hashes[0]) + assert dict(event_manager._event_source._uuids) == {digest: "image-a"} + finally: + event_manager.shutdown() From 7495cfbe164c311e95e96576b781491d24b254de Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Tue, 27 Oct 2026 21:23:55 -0700 Subject: [PATCH 25/28] [None][fix] Preserve inference and metadata for unmapped image hints Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_events.py | 15 ++-- tensorrt_llm/grpc/openengine/README.md | 6 +- .../grpc/openengine/test_image_routing.py | 80 ++++++++++++++++++- .../test_kv_cache_event_manager.py | 7 +- 4 files changed, 94 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index d14663ec64a5..84df77eeb4a4 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -563,7 +563,7 @@ def __init__( self._include_mm_keys = mm_token_id_offset is not None self._mm_token_id_offset = mm_token_id_offset self._uuid_capacity = max_entries - self._uuids: OrderedDict[bytes, str] = OrderedDict() + self._uuids: OrderedDict[bytes, str | None] = OrderedDict() self._uuid_lock = threading.Lock() self._event_sink = kv_cache_manager_v2_runtime.StreamingEventSink( max_entries=max_entries, @@ -580,10 +580,11 @@ def remember_multimodal_uuids(self, items: Iterable[tuple[bytes, str]]) -> None: if len(uuid.encode("utf-8")) > 256: continue with self._uuid_lock: - previous = self._uuids.get(digest) - if previous is not None and previous != uuid: - raise ValueError("Conflicting UUIDs for a multimodal content digest") - self._uuids[digest] = uuid + if digest in self._uuids and self._uuids[digest] != uuid: + # Keep conflicting hints ambiguous until their bounded entry is evicted. + self._uuids[digest] = None + else: + self._uuids[digest] = uuid self._uuids.move_to_end(digest) while len(self._uuids) > self._uuid_capacity: self._uuids.popitem(last=False) @@ -627,9 +628,7 @@ def drain_events(self) -> list[BlockStored | BlockRemoved | AllBlocksCleared]: ) if mm_keys is not None and self._mm_token_id_offset is not None - and any( - key.uuid is not None for block_keys in mm_keys for key in block_keys - ) + and any(mm_keys) else None ), ) diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index a801b3cee8a2..aa831170dc22 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -126,7 +126,7 @@ block reuse enabled, and ZMQ event publishing. Image passthrough separately requires `supports_multimodal` and `extra.trtllm_supported_modalities` to advertise image support. Video, audio, and disaggregated image routing are not advertised. -For image blocks with caller UUIDs, `BlockStored.trtllm_image_routing` contains +For image blocks, `BlockStored.trtllm_image_routing` contains `version`, `token_id_offset` (the vocabulary size), and `keys`, aligned with the stored blocks. Each key carries its native content `hash`, `start_offset`, and optional `uuid`. Native image tokens start with a digest and continue with IDs @@ -135,7 +135,9 @@ sent in `MediaItem.uuid`. Native block hashes, parent chains, and removals retai their existing meaning. Cache hashing continues to include image content. The publisher retains at most `max_entries` digest-to-UUID associations and skips -UUIDs longer than 256 UTF-8 bytes. Eviction can leave UUIDs unavailable. Consumers +UUIDs longer than 256 UTF-8 bytes. Conflicting hints mark a digest ambiguous until +its entry is evicted, without failing inference. Missing, evicted, or ambiguous +hints omit the UUID while retaining the image-routing metadata. Consumers must ignore unmapped image stores while continuing to process removals. Routing identities must refer to stable media content. This extension does not change the OpenEngine schema. diff --git a/tests/unittest/grpc/openengine/test_image_routing.py b/tests/unittest/grpc/openengine/test_image_routing.py index 6e6c0c6d31c4..ff0c8fb755ec 100644 --- a/tests/unittest/grpc/openengine/test_image_routing.py +++ b/tests/unittest/grpc/openengine/test_image_routing.py @@ -67,8 +67,8 @@ class Stored: assert source._multimodal_key((bytes.fromhex(hashes[1]), 0)).uuid == uuid manager.remember_multimodal_uuids([(bytes.fromhex(hashes[0]), "x" * 257)]) assert len(source._uuids) == 1 - with pytest.raises(ValueError, match="Conflicting UUIDs"): - manager.remember_multimodal_uuids([(bytes.fromhex(hashes[1]), "other")]) + manager.remember_multimodal_uuids([(bytes.fromhex(hashes[1]), "other")]) + assert source._multimodal_key((bytes.fromhex(hashes[1]), 0)).uuid is None finally: source._event_sink = native_sink manager.shutdown() @@ -138,3 +138,79 @@ def test_request_uuid_pairs_preserve_partial_uuid_slots(): assert dict(event_manager._event_source._uuids) == {digest: "image-a"} finally: event_manager.shutdown() + + +def test_conflicting_request_hints_preserve_inference_cache_keys(): + manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + cache = KVCacheManagerV2.__new__(KVCacheManagerV2) + cache.vocab_size = 1000 + cache.event_manager = manager + hashes = [tuple(range(8))] + digest = b"".join(value.to_bytes(4, "big", signed=True) for value in hashes[0]) + results = [] + try: + for uuid in ["image-a", "image-b", "image-a"]: + request = SimpleNamespace( + multimodal_hashes=hashes, + multimodal_uuids=[uuid], + multimodal_positions=[0], + multimodal_lengths=[1], + multimodal_item_run_cu_offsets=None, + multimodal_run_positions=None, + multimodal_run_lengths=None, + ) + results.append(cache._augment_tokens_for_block_reuse([1, 2], request)) + hint = manager._event_source._multimodal_key((digest, 0)).uuid + assert hint == ("image-a" if len(results) == 1 else None) + assert results[0] == results[1] == results[2] + finally: + manager.shutdown() + + +@pytest.mark.parametrize("hint_state", ["absent", "evicted", "ambiguous"]) +def test_unmapped_integer_image_continuations_keep_routing_metadata(monkeypatch, hint_state): + manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + source = manager._event_source + native_sink = source._event_sink + digest = bytes(range(32)) + try: + if hint_state != "absent": + manager.remember_multimodal_uuids([(digest, "image-a")]) + if hint_state == "evicted": + manager.remember_multimodal_uuids([(bytes(reversed(digest)), "image-b")]) + else: + manager.remember_multimodal_uuids([(digest, "image-b")]) + + class Stored: + block_hashes = [123] + parent_block_hash = 122 + token_ids = [1003, 1004, 1005, 7] + lora_id = None + mm_keys = [[(digest, 3)]] + + monkeypatch.setattr(events.kv_cache_manager_v2_runtime, "StreamingBlockStoredData", Stored) + source._event_sink = SimpleNamespace(drain_iteration_events=lambda: [Stored()]) + wire = msgspec.msgpack.decode(msgspec.msgpack.encode(source.drain_events()[0])) + metadata = wire["trtllm_image_routing"] + assert metadata["version"] == 1 and metadata["token_id_offset"] == 1000 + key = metadata["keys"][0][0] + assert key["hash"] == digest.hex() and key["start_offset"] == 3 + assert key.get("uuid") is None + assert wire["token_ids"] == Stored.token_ids + finally: + source._event_sink = native_sink + manager.shutdown() diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py index fee1ec228ac8..e5894a412113 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -471,8 +471,11 @@ def test_streaming_uuid_hints_are_bounded_and_do_not_replace_content_hashes(): assert first.hash != second.hash manager.remember_multimodal_uuids([(a, "x" * 257)]) assert len(source._uuids) == 1 - with pytest.raises(ValueError, match="Conflicting UUIDs"): - manager.remember_multimodal_uuids([(b, "different-hint")]) + manager.remember_multimodal_uuids([(b, "different-hint")]) + assert source._multimodal_key((b, 0)).uuid is None + manager.remember_multimodal_uuids([(b, "same-routing-hint")]) + assert source._multimodal_key((b, 0)).uuid is None + assert len(source._uuids) == 1 finally: manager.shutdown() From e6b1db517e06d822391590b861cf4854cc2d1ac2 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:45:40 -0700 Subject: [PATCH 26/28] Avoid duplicate image routing keys in KV events Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/kv_cache_events.py | 3 +-- tensorrt_llm/grpc/openengine/README.md | 9 +++++---- tests/unittest/grpc/openengine/test_image_routing.py | 11 ++++++++--- .../grpc/openengine/test_image_routing_conformance.py | 2 +- .../test_kv_cache_event_manager.py | 2 +- 5 files changed, 16 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 84df77eeb4a4..72750ef4a123 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -88,7 +88,6 @@ class ImageRoutingMetadata(msgspec.Struct, frozen=True, gc=False): version: int token_id_offset: int - keys: list[list[MultimodalKey]] class BlockStored(KVCacheWireEvent): @@ -624,7 +623,7 @@ def drain_events(self) -> list[BlockStored | BlockRemoved | AllBlocksCleared]: mm_keys=mm_keys, trtllm_image_routing=( ImageRoutingMetadata( - version=1, token_id_offset=self._mm_token_id_offset, keys=mm_keys + version=1, token_id_offset=self._mm_token_id_offset ) if mm_keys is not None and self._mm_token_id_offset is not None diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index aa831170dc22..eae2ed74470e 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -127,10 +127,11 @@ requires `supports_multimodal` and `extra.trtllm_supported_modalities` to advert image support. Video, audio, and disaggregated image routing are not advertised. For image blocks, `BlockStored.trtllm_image_routing` contains -`version`, `token_id_offset` (the vocabulary size), and `keys`, aligned with the -stored blocks. Each key carries its native content `hash`, `start_offset`, and -optional `uuid`. Native image tokens start with a digest and continue with IDs -above `token_id_offset`; a router can normalize those tokens using identities it +`version` and `token_id_offset` (the vocabulary size). The existing +`BlockStored.mm_keys` contains the keys, aligned with the stored blocks. Each key +retains its native content `hash` and `start_offset` and gains an optional `uuid`. +Keys are serialized only in `mm_keys`. Native image tokens start with a digest +and continue with IDs above `token_id_offset`; a router can normalize those tokens using identities it sent in `MediaItem.uuid`. Native block hashes, parent chains, and removals retain their existing meaning. Cache hashing continues to include image content. diff --git a/tests/unittest/grpc/openengine/test_image_routing.py b/tests/unittest/grpc/openengine/test_image_routing.py index ff0c8fb755ec..652519966f59 100644 --- a/tests/unittest/grpc/openengine/test_image_routing.py +++ b/tests/unittest/grpc/openengine/test_image_routing.py @@ -52,8 +52,12 @@ class Stored: wire = msgspec.msgpack.decode(msgspec.msgpack.encode(source.drain_events()[0])) metadata = wire["trtllm_image_routing"] assert metadata["version"] == 1 and metadata["token_id_offset"] == 1000 - assert metadata["keys"][0][0]["uuid"] == uuid - assert metadata["keys"][0][0]["hash"] == hashes[0] + assert set(metadata) == {"version", "token_id_offset"} + assert wire["mm_keys"][0][0]["uuid"] == uuid + assert wire["mm_keys"][0][0]["hash"] == hashes[0] + encoded = msgspec.msgpack.encode(source.drain_events()[0]) + assert encoded.count(hashes[0].encode()) == 2 # Token stream and mm_keys. + assert encoded.count(uuid.encode()) == 1 assert wire["token_ids"] == Stored.token_ids Stored.token_ids = [1, 2, 3, 4] @@ -207,7 +211,8 @@ class Stored: wire = msgspec.msgpack.decode(msgspec.msgpack.encode(source.drain_events()[0])) metadata = wire["trtllm_image_routing"] assert metadata["version"] == 1 and metadata["token_id_offset"] == 1000 - key = metadata["keys"][0][0] + assert set(metadata) == {"version", "token_id_offset"} + key = wire["mm_keys"][0][0] assert key["hash"] == digest.hex() and key["start_offset"] == 3 assert key.get("uuid") is None assert wire["token_ids"] == Stored.token_ids diff --git a/tests/unittest/grpc/openengine/test_image_routing_conformance.py b/tests/unittest/grpc/openengine/test_image_routing_conformance.py index abbd93effeb3..d122ba6de045 100644 --- a/tests/unittest/grpc/openengine/test_image_routing_conformance.py +++ b/tests/unittest/grpc/openengine/test_image_routing_conformance.py @@ -115,7 +115,7 @@ async def generate(color, request_id): metadata = event.get("trtllm_image_routing") if not metadata: continue - for keys in metadata["keys"]: + for keys in event["mm_keys"]: for key in keys: assert key["uuid"] == "dynamo-image-v1:0123456789abcdef" hashes.add(key["hash"]) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py index e5894a412113..90d3ec7508bc 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -440,7 +440,7 @@ def test_native_streaming_sink_preserves_multimodal_event_data(real_block_factor assert manager.stored_blocks == 4 assert stored.trtllm_image_routing.version == 1 assert stored.trtllm_image_routing.token_id_offset == 1000 - assert [[key.uuid for key in keys] for keys in stored.trtllm_image_routing.keys] == [ + assert [[key.uuid for key in keys] for keys in stored.mm_keys] == [ ["dynamo-image-v1:0123456789abcdef"], [], ["dynamo-image-v1:0123456789abcdef"] * 2, From bb4e1f3174899e6aee654b699f8b867ce14064d4 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:35:08 -0700 Subject: [PATCH 27/28] Avoid repeated image hint registration during chunked prefill Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 2 +- .../_torch/pyexecutor/kv_cache_events.py | 4 +- tensorrt_llm/grpc/openengine/README.md | 13 ++++-- .../unittest/grpc/openengine/test_control.py | 4 +- .../grpc/openengine/test_image_routing.py | 46 +++++++++++++++++++ .../test_image_routing_conformance.py | 3 +- 6 files changed, 60 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 47a7a58e06cd..0f76181101fa 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6126,7 +6126,7 @@ def _augment_tokens_for_block_reuse( "reuse." ) - if isinstance(self.event_manager, StreamingKVCacheEventManager): + if chunk_start == 0 and isinstance(self.event_manager, StreamingKVCacheEventManager): uuids = req.multimodal_uuids if uuids is not None: uuid_pairs = [ diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 72750ef4a123..d6063b5fe5f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -591,9 +591,7 @@ def remember_multimodal_uuids(self, items: Iterable[tuple[bytes, str]]) -> None: def _multimodal_key(self, key: tuple) -> MultimodalKey: digest = bytes(key[0]) with self._uuid_lock: - uuid = key[2] if len(key) > 2 else None - if uuid is None: - uuid = self._uuids.get(digest) + uuid = self._uuids.get(digest) return MultimodalKey(hash=digest.hex(), start_offset=key[1], uuid=uuid) def set_target_life_cycle(self, life_cycle_id: int) -> None: diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index eae2ed74470e..0c77e0a504f6 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -124,15 +124,20 @@ For attention DP, each publisher binds `base_port + global_rank`. When event str metadata for aggregated Qwen3-VL with KV cache manager V2 explicitly enabled, block reuse enabled, and ZMQ event publishing. Image passthrough separately requires `supports_multimodal` and `extra.trtllm_supported_modalities` to advertise -image support. Video, audio, and disaggregated image routing are not advertised. +image support. Set `use_kv_cache_manager_v2=True` explicitly: `"auto"` is +excluded from capability discovery even if the engine resolves it to V2. +Video, audio, and disaggregated image routing are not advertised. For image blocks, `BlockStored.trtllm_image_routing` contains `version` and `token_id_offset` (the vocabulary size). The existing `BlockStored.mm_keys` contains the keys, aligned with the stored blocks. Each key retains its native content `hash` and `start_offset` and gains an optional `uuid`. -Keys are serialized only in `mm_keys`. Native image tokens start with a digest -and continue with IDs above `token_id_offset`; a router can normalize those tokens using identities it -sent in `MediaItem.uuid`. Native block hashes, parent chains, and removals retain +Keys are serialized only in `mm_keys`. The publisher emits this extension for +all multimodal keys, including video/audio keys; the name identifies the current +image-routing contract, not a modality filter. Only image routing has been +validated and advertised by OpenEngine. Native image tokens start with a digest +and continue with IDs above `token_id_offset`; a router can normalize those +tokens using identities it sent in `MediaItem.uuid`. Native block hashes, parent chains, and removals retain their existing meaning. Cache hashing continues to include image content. The publisher retains at most `max_entries` digest-to-UUID associations and skips diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 70704911d1d3..e101b647be29 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -35,7 +35,7 @@ @pytest.mark.parametrize("model_type", ["qwen3_vl", "qwen2_vl", "text"]) -@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("enabled", [False, True, "auto"]) @pytest.mark.parametrize( "block_reuse,events_enabled,publisher", [(True, True, "zmq"), (False, True, "zmq"), (True, False, "zmq"), (True, True, "null")], @@ -56,7 +56,7 @@ def test_image_routing_capability_requires_streaming_v2_and_validated_model( ), ) expected = ( - enabled + enabled is True and model_type == "qwen3_vl" and block_reuse and events_enabled diff --git a/tests/unittest/grpc/openengine/test_image_routing.py b/tests/unittest/grpc/openengine/test_image_routing.py index 652519966f59..cbbb321b285d 100644 --- a/tests/unittest/grpc/openengine/test_image_routing.py +++ b/tests/unittest/grpc/openengine/test_image_routing.py @@ -144,6 +144,52 @@ def test_request_uuid_pairs_preserve_partial_uuid_slots(): event_manager.shutdown() +def test_chunked_prefill_registers_hints_before_images_without_repeating(monkeypatch): + event_manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=2, + max_window_size=128, + mm_token_id_offset=1000, + ) + cache = KVCacheManagerV2.__new__(KVCacheManagerV2) + cache.vocab_size = 1000 + cache.event_manager = event_manager + request = SimpleNamespace( + multimodal_hashes=[tuple(range(8))], + multimodal_uuids=["image-a"], + multimodal_positions=[2], + multimodal_lengths=[3], + multimodal_item_run_cu_offsets=None, + multimodal_run_positions=None, + multimodal_run_lengths=None, + ) + registrations = [] + original = event_manager.remember_multimodal_uuids + + def remember(pairs): + registrations.append(pairs) + original(pairs) + + monkeypatch.setattr(event_manager, "remember_multimodal_uuids", remember) + digest = b"".join(value.to_bytes(4, "big", signed=True) for value in range(8)) + try: + tokens = [1, 2, 3, 4, 5, 6] + first = cache._augment_tokens_for_block_reuse(tokens, request, start=0, end=2) + assert len(registrations) == 1 + assert event_manager._event_source._multimodal_key((digest, 0)).uuid == "image-a" + # Later chunks cannot read or register the request's UUID vector again. + del request.multimodal_uuids + middle = cache._augment_tokens_for_block_reuse(tokens, request, start=2, end=4) + last = cache._augment_tokens_for_block_reuse(tokens, request, start=4, end=6) + assert len(registrations) == 1 + request.multimodal_uuids = ["image-a"] + expected = cache._augment_tokens_for_block_reuse(tokens, request) + assert first + middle + last == expected + finally: + event_manager.shutdown() + + def test_conflicting_request_hints_preserve_inference_cache_keys(): manager = events.StreamingKVCacheEventManager( KVEventsConfig(enable_kv_cache_events=True, publisher="null"), diff --git a/tests/unittest/grpc/openengine/test_image_routing_conformance.py b/tests/unittest/grpc/openengine/test_image_routing_conformance.py index d122ba6de045..7ffefc99f1cb 100644 --- a/tests/unittest/grpc/openengine/test_image_routing_conformance.py +++ b/tests/unittest/grpc/openengine/test_image_routing_conformance.py @@ -25,8 +25,7 @@ def test_real_image_routing_keeps_cache_identity_content_sensitive(unused_tcp_port, tmp_path): - if not torch.cuda.is_available(): - pytest.skip("requires a GPU and the real Qwen3-VL-2B checkpoint") + assert torch.cuda.is_available(), "GPU conformance requires a CUDA device" endpoint = f"tcp://127.0.0.1:{unused_tcp_port}" context = zmq.Context() subscriber = context.socket(zmq.SUB) From 1daf67ac6e2a47c5bb573d26f43e1e44dd707610 Mon Sep 17 00:00:00 2001 From: krishung5 <43719498+krishung5@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:56:44 -0700 Subject: [PATCH 28/28] Narrow image routing scope and clarify event metadata docs Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> --- docs/source/features/kvcache.md | 10 ++++++++++ tensorrt_llm/grpc/openengine/README.md | 7 +++++-- tensorrt_llm/inputs/registry.py | 4 ---- .../test_kv_cache_manager_v2_helix_superblock.py | 10 ---------- .../_torch/executor/test_pytorch_model_engine.py | 2 -- tests/unittest/_torch/test_connector_adp.py | 5 +---- tests/unittest/grpc/openengine/test_multimodal.py | 9 ++------- 7 files changed, 18 insertions(+), 29 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 3f46d098c863..684a5fd695a8 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -323,6 +323,16 @@ prompt = TextPrompt( V2 derives `mm_keys` directly from the cached token sequence. Each entry identifies a continuous multimodal segment within that block: `hash` is the item's digest, and `start_offset` is the segment's first token offset within the item. An item spanning multiple blocks retains the same digest with increasing offsets. Text may separate segments of the same item. Items are processed in prompt order; one item cannot resume after another item has started. The item digest is distinct from `block_hash`, which also depends on the preceding token sequence. +Streaming V2 events can also include an optional `uuid` in each `mm_keys` entry. +It is a bounded routing hint; it does not replace the content-sensitive digest. +Missing, evicted, or conflicting hints omit `uuid`. +For multimodal blocks, the optional `BlockStored.trtllm_image_routing` extension +provides the routing format version and token-ID offset. The extension can accompany +video and audio keys as well as image keys; the OpenEngine capability flag advertises +only the currently validated image-routing support. See the +[OpenEngine image routing contract](source:tensorrt_llm/grpc/openengine/README.md#image-routing-metadata) +for capability requirements and consumer handling. + ### Enable Offloading to Host Memory diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index 0c77e0a504f6..1ecae66b131a 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -142,8 +142,11 @@ their existing meaning. Cache hashing continues to include image content. The publisher retains at most `max_entries` digest-to-UUID associations and skips UUIDs longer than 256 UTF-8 bytes. Conflicting hints mark a digest ambiguous until -its entry is evicted, without failing inference. Missing, evicted, or ambiguous -hints omit the UUID while retaining the image-routing metadata. Consumers +its entry is evicted or the publisher restarts, without failing inference. +Every registration refreshes the entry's LRU position, including ambiguous entries, +so a frequently registered digest can remain ambiguous for the publisher's lifetime. +An agreeing registration does not restore a hint after a conflict. Missing, evicted, +or ambiguous hints omit the UUID while retaining the image-routing metadata. Consumers must ignore unmapped image stores while continuing to process removals. Routing identities must refer to stable media content. This extension does not change the OpenEngine schema. diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 8f905eec8917..140c146a7955 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -174,10 +174,6 @@ class BaseMultimodalInputProcessor(ABC): # uses `_init_tokenization_cache` and `_encode_with_tokenization_cache`. supports_tokenization_cache: ClassVar[bool] = False - def get_preferred_media_io_kwargs(self) -> dict[str, dict[str, Any]]: - """Return per-modality decode defaults, below server and request overrides.""" - return {} - def get_mm_item_metadata( self, prompt_token_ids: List[int], diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py index 213e0acc09b9..7f9d7ef36e06 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py @@ -429,16 +429,6 @@ def test_resize_context_rejects_disagg_helix_request(prompt_len: int) -> None: ) -def test_resize_context_allows_disagg_helix_dummy_request() -> None: - resizes = [] - req = _context_req(prompt_len=200, total_input_len_cp=200) - req.is_dummy_request = True - assert KVCacheManagerV2.resize_context( - _context_resize_mgr(is_disagg=True, resizes=resizes), req, num_tokens=32 - ) - assert resizes == [64 + 32] - - def test_dummy_frozen_fields_sum_invariant() -> None: """The frozen dummy fiction (last rank active, shared synthetic global length) keeps the same books as real requests: per-rank lengths sum to diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 6929621f6d42..42fe06ecb65e 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1903,7 +1903,6 @@ def test_mrope_multi_beam_generation_positions(self) -> None: """Every beam row gets its request's `past + delta` on all 3 axes.""" beam_width = 2 engine, kv_cache_manager = create_model_engine_and_kvcache( - self, llm_args=TorchLlmArgs(model="dummy", max_batch_size=4, max_num_tokens=64, @@ -1973,7 +1972,6 @@ def test_mrope_mixed_batched_and_fallback_positions_match_fallback( self) -> None: """Mixed scalar batching matches the old per-request writes exactly.""" engine, kv_cache_manager = create_model_engine_and_kvcache( - self, llm_args=TorchLlmArgs(model="dummy", max_batch_size=2, max_num_tokens=64), diff --git a/tests/unittest/_torch/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index 3e778e0445c0..ad1e6829630e 100644 --- a/tests/unittest/_torch/test_connector_adp.py +++ b/tests/unittest/_torch/test_connector_adp.py @@ -80,12 +80,10 @@ def _manager(rank: int = 1) -> KvCacheConnectorManager: def _configure_recovery_controls(executor: "PyExecutor") -> None: from tensorrt_llm._torch.pyexecutor.executor_request_queue import ExecutorRequestQueue - executor.dist = MagicMock(tp_size=1, cp_size=1, pp_size=1, tp_rank=0, has_cp_helix=False) - executor.dist.tp_cp_allgather_int64.side_effect = executor.dist.tp_allgather_int64 + executor.dist = MagicMock(tp_size=1, cp_size=1, pp_size=1) executor.enable_attention_dp = True executor.enable_early_first_token_response = False executor.kv_connector_manager = None - executor.kv_cache_transceiver = None executor._is_kv_manager_v2 = False executor._kv_load_tracker = KvLoadTracker(None, rank=0, enabled=False) executor._kv_cache_capacity = {} @@ -416,7 +414,6 @@ def test_batch_start_combines_readiness_with_status(batch_sizes: tuple[int, int] _configure_recovery_controls(executor) executor.kv_connector_manager = _manager(rank) executor.dist.tp_size = 2 - executor.dist.tp_rank = rank executor.dist.tp_allgather_int64.side_effect = partial(collective.gather_int64, rank) executors.append(executor) diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 9554b7050171..bc2cb9865266 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -848,13 +848,8 @@ async def test_text_only_input_mapping_is_unchanged(): @pytest.mark.asyncio -@pytest.mark.parametrize("model_defaults", [None, {"format": "np"}]) -async def test_image_sources_preserve_order_uuid_and_model_decode_defaults( - image_llm, monkeypatch, model_defaults -): +async def test_image_sources_preserve_order_uuid_and_model_decode_defaults(image_llm, monkeypatch): llm, _ = image_llm - if model_defaults is None: - monkeypatch.delattr(_ImageInputProcessor, "get_preferred_media_io_kwargs") created_with = [] class _ImageIo: @@ -897,7 +892,7 @@ async def async_load_bytes(self, data): mapped = await request_mapping._input_from_request(request, llm) - assert created_with == [(model_defaults, None)] + assert created_with == [({"format": "np"}, None)] assert mapped["prompt"] == request.prompt assert mapped["multi_modal_data"] == { "image": [