diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 71132d01513f..dddd934c24c7 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -267,6 +267,16 @@ class Gemma4InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInpu # image, when present) can be split across chunks safely. mm_bidirectional_blocks = False + @property + def supported_modalities(self) -> frozenset[str]: + modalities = super().supported_modalities + config = getattr(self, "config", None) + if getattr(config, "vision_config", None) is None: + modalities -= {"image", "video"} + if getattr(config, "audio_config", None) is None: + modalities -= {"audio"} + return modalities + def __init__( self, model_path: str, diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py index ef0608859a6f..9d5dd7ff74b3 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py @@ -1152,6 +1152,13 @@ class NemotronHMultimodalInputProcessor( ): supports_token_id_mm_expansion: ClassVar[bool] = True + @property + def supported_modalities(self) -> frozenset[str]: + modalities = super().supported_modalities + if getattr(self, "_audio_extractor", None) is None: + return modalities - {"audio"} + return modalities + def get_mm_item_metadata( self, prompt_token_ids: List[int], diff --git a/tensorrt_llm/_torch/models/modeling_whisper.py b/tensorrt_llm/_torch/models/modeling_whisper.py index 5a1e37c8b19a..ffc7540b28ce 100644 --- a/tensorrt_llm/_torch/models/modeling_whisper.py +++ b/tensorrt_llm/_torch/models/modeling_whisper.py @@ -598,6 +598,13 @@ class WhisperInputProcessor(InputProcessor): # Marks this model as feature-driven for the encoder side: prompts # without audio cannot be served and are rejected at submission. requires_encoder_features = True + # The text prompt selects the decoder language/task; pre-tokenized prompt + # IDs are not consumed by this processor. + supports_token_ids_input = False + + @property + def supported_modalities(self) -> frozenset[str]: + return frozenset({"audio"}) def __init__(self, model_path, config, tokenizer, trust_remote_code: bool = True, **kwargs): self.model_path = model_path diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index 95c7d193db64..5e6fcaf34b57 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -69,11 +69,29 @@ Set `failureThreshold * periodSeconds` to cover the longest expected model start OpenEngine gRPC does not yet expose sleep or weight updates. When those operations are added, they must publish `inference=NOT_SERVING` for their duration while keeping `liveness=SERVING`; the current executor health predicate alone does not detect those admission states. A sidecar that loses `Watch` or sees a new engine identity must validate health again. +For models that do not accept token-ID input, including Whisper, the optional +`Health(include_inference_probe=True)` check reports the inference probe as not +applicable. It does not submit inference and does not establish inference liveness; +the ordinary lifecycle and executor health checks still apply. + The `Inference.Generate` RPC loads the selected model through TensorRT-LLM's PyTorch `LLM` API and streams incremental token and text events. It supports text and token-ID inputs, native sampling and stopping options, top-N prompt and output log probabilities, TensorRT-LLM guided-decoding modes, cache salt, trace-context propagation, multiple output sequences, finish reasons, and final usage. Clients must continuously consume the response stream. If response delivery remains stalled for 30 seconds, the server aborts the engine request and terminates the stream with a retryable overload error. -Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, multimodal media, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`. +Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`. + +### Media capability discovery + +`GetModelInfo.supports_multimodal` and `extra.trtllm_supported_modalities` reflect +the loaded processor's read-only `supported_modalities` property, intersected +with registered placeholders and implemented OpenEngine handlers. Both discovery +and request validation use this result. `InputProcessor` defaults to text-only; +`BaseMultimodalInputProcessor` defaults to its registered modalities. Processors +with checkpoint-dependent encoders override the property to narrow that set. +Gemma4 excludes audio without `audio_config` and images/video without +`vision_config`. Nemotron excludes audio without an +initialized audio extractor, and Whisper declares audio only. Registered custom +processors must declare the property to advertise media support. ### Disaggregated serving diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 83a6f9979b51..064ac9178395 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -64,7 +64,11 @@ def supported_guides(backend: str | None) -> frozenset[str]: def _loaded_model_type(llm: Any) -> str: """Return the top-level model type for the loaded input processor.""" - return resolve_top_level_model_type(llm.input_processor.config) + input_processor = getattr(llm, "input_processor", None) + config = getattr(input_processor, "config", None) or getattr(llm, "_hf_model_config", None) + if config is None: + return "" + return resolve_top_level_model_type(config) def resolve_media_capabilities( @@ -72,21 +76,28 @@ def resolve_media_capabilities( ) -> MediaCapabilities: """Intersect model support with implemented request handlers. - A live multimodal processor prevents registry-only false positives. + A live registered processor prevents registry-only false positives. """ input_processor = getattr(llm, "input_processor", None) - if not isinstance(input_processor, BaseMultimodalInputProcessor): - return MediaCapabilities({}) if getattr(getattr(llm, "args", None), "disable_mm_encoder", False): return MediaCapabilities({}) + model_type = _loaded_model_type(llm) if not model_type: return MediaCapabilities({}) + if not isinstance(input_processor, BaseMultimodalInputProcessor): + registered_model_type = getattr(type(input_processor), "_registered_model_type", None) + if registered_model_type != model_type: + return MediaCapabilities({}) + if not callable(input_processor): + return MediaCapabilities({}) + supported_modalities = getattr(input_processor, "supported_modalities", frozenset()) placeholders = { modality: MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality) for modality in implemented_modalities - if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) + if modality in supported_modalities + and MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality) } # Fail startup before Control advertises a malformed template. for modality, template in placeholders.items(): diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index bbd5a9e824ae..b693cdeb885d 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -39,7 +39,7 @@ from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError from .errors import AbortFailedError from .kv_events import data_parallel_size, events_config -from .request_mapping import conversation_affinity_enabled +from .request_mapping import conversation_affinity_enabled, token_ids_input_supported __all__ = ["OpenEngineControlServicer"] @@ -289,7 +289,7 @@ async def GetModelInfo( served_model_name=self._model, # Without a tokenizer the input processor rejects every string prompt. supports_text_input=getattr(self._llm, "tokenizer", None) is not None, - supports_token_ids_input=True, + supports_token_ids_input=token_ids_input_supported(self._llm), # Generate rejects `lora_name` even on an enable_lora build. supports_lora=False, supports_multimodal=bool(media_modalities), @@ -476,6 +476,15 @@ async def _inference_probe(self, model: str) -> lifecycle_pb2.HealthCheck: async def _run_inference_probe(self, model: str) -> lifecycle_pb2.HealthCheck: from tensorrt_llm.sampling_params import SamplingParams + if not token_ids_input_supported(self._llm): + return lifecycle_pb2.HealthCheck( + name="inference_probe", + state=lifecycle_pb2.HEALTH_STATE_READY, + message=( + "probe not applicable: model does not accept token-ID input; " + "inference liveness was not tested" + ), + ) handle = None reservation = None # Counted by GetLoad and reachable by Abort(all_requests): a probe diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index c43d0f742ea2..41334157914e 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -16,11 +16,13 @@ from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.inputs.media_io import ( + AudioMediaIO, ImageMediaIO, MediaFetchError, VideoMediaIO, _decoded_video_bytes, ) +from tensorrt_llm.inputs.multimodal_data import AudioData from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams @@ -182,6 +184,12 @@ def conversation_affinity_enabled(llm: Any) -> bool: ) +def token_ids_input_supported(llm: Any) -> bool: + """Whether the loaded input processor consumes pre-tokenized prompts.""" + input_processor = getattr(llm, "input_processor", None) + return getattr(input_processor, "supports_token_ids_input", True) is not False + + def _top_n_candidates(selection: Any, name: str) -> int: kind = selection.WhichOneof("selection") if kind is None: @@ -395,12 +403,16 @@ def _metadata_from_context( def _base_input_from_request( - request: generation_pb2.GenerateRequest, + request: generation_pb2.GenerateRequest, llm: Any ) -> str | dict[str, list[int]]: input_kind = request.WhichOneof("input") if input_kind == "prompt": return request.prompt if input_kind == "token_ids": + if not token_ids_input_supported(llm): + raise UnsupportedFeatureError( + "token-ID input is not supported by the loaded model; send a text prompt" + ) return {"prompt_token_ids": request.token_ids.ids[:]} raise ValueError("exactly one of prompt or token_ids must be set") @@ -459,6 +471,10 @@ def _validate_image_item(item: generation_pb2.MediaItem, index: int) -> None: _validate_media_item(item, index, "image") +def _validate_audio_item(item: generation_pb2.MediaItem, index: int) -> None: + _validate_media_item(item, index, "audio") + + def _validate_video_item(item: generation_pb2.MediaItem, index: int) -> None: _validate_media_item(item, index, "video") @@ -507,6 +523,17 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any: return load +def _audio_loader_factory(llm: Any, budget: _DecodedMediaBudget) -> _MediaLoader: + audio_io = _media_io_for_model( + llm, "audio", AudioMediaIO, {"decoded_memory_reserver": budget.reserve} + ) + + async def load(item: generation_pb2.MediaItem, index: int) -> Any: + return await _load_media_item(item, audio_io, index, "audio") + + return load + + def _decoded_image_bytes(image: Any) -> int: import numpy as np import torch @@ -521,6 +548,23 @@ def _decoded_image_bytes(image: Any) -> int: raise TypeError(f"unsupported decoded image type: {type(image).__name__}") +def _decoded_audio_data_bytes(audio: AudioData) -> int: + import numpy as np + import torch + + if isinstance(audio.samples, torch.Tensor): + return audio.samples.numel() * audio.samples.element_size() + if isinstance(audio.samples, np.ndarray): + return audio.samples.nbytes + raise TypeError(f"unsupported decoded audio type: {type(audio.samples).__name__}") + + +def _decoded_audio_bytes(audio: Any) -> int: + if not isinstance(audio, tuple) or len(audio) != 2: + raise TypeError(f"unsupported decoded audio type: {type(audio).__name__}") + return _decoded_audio_data_bytes(AudioData(samples=audio[0], sample_rate=audio[1])) + + def _video_runtime_available() -> bool: try: importlib.import_module("cv2") @@ -541,6 +585,13 @@ def _video_runtime_available() -> bool: _video_runtime_available, reserves_decoded_bytes=True, ), + generation_pb2.MODALITY_AUDIO: _MediaHandler( + "audio", + _validate_audio_item, + _audio_loader_factory, + _decoded_audio_bytes, + reserves_decoded_bytes=True, + ), } @@ -564,7 +615,7 @@ def _prepare_input_from_request( media_capabilities: MediaCapabilities | None = None, ) -> _PreparedInput: """Validate request media without starting any fetch or decode work.""" - inputs = _base_input_from_request(request) + inputs = _base_input_from_request(request, llm) if not request.media: return _PreparedInput(inputs, (), None) @@ -630,6 +681,8 @@ def _validate_media_placeholders(prepared: _PreparedInput, llm: Any) -> None: if not prepared.media or prepared.capabilities is None: return processor = llm.input_processor + if not isinstance(processor, BaseMultimodalInputProcessor): + return if not processor.requires_input_placeholders(tokenized=not isinstance(prepared.inputs, str)): return tokenizer = llm.tokenizer @@ -750,4 +803,5 @@ async def _input_from_request( __all__ = [ "conversation_params_from_request", "sampling_params_from_request", + "token_ids_input_supported", ] diff --git a/tensorrt_llm/grpc/openengine/server.py b/tensorrt_llm/grpc/openengine/server.py index c9d810f65d01..97ccb3210637 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -23,6 +23,7 @@ from tensorrt_llm import LLM as PyTorchLLM from tensorrt_llm.grpc._protobuf import _warn_if_python_protobuf +from tensorrt_llm.inputs.media_io import _close_aiohttp_session from tensorrt_llm.logger import logger from tensorrt_llm.serve._frontend_processes import ( FrontendStartupCancelled, @@ -287,14 +288,17 @@ async def stop(self, grace: float = 5.0) -> None: if self._metrics is not None: await self._metrics.stop() 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 - ), - ) + try: + await _close_aiohttp_session() + finally: + await asyncio.gather( + asyncio.to_thread( + self._input_executor.shutdown, wait=True, cancel_futures=True + ), + asyncio.to_thread( + self._media_executor.shutdown, wait=True, cancel_futures=True + ), + ) logger.info("OpenEngine server stopped") async def wait_for_termination(self) -> None: diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 0e12b267b3ee..4d1317476073 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -10,6 +10,7 @@ import os import socket import tempfile +import threading from abc import ABC, abstractmethod from collections.abc import Iterator from concurrent.futures import Executor @@ -133,29 +134,37 @@ def convert_image_mode(image: Image.Image, to_mode: str) -> Image.Image: # `_load_video_by_cv2` for the per-format contract. _SUPPORTED_VIDEO_FORMATS = ("pt", "np", "pil") -# Module-level aiohttp session shared across all media fetch calls. -# Created lazily on first use inside the async event loop, then reused so that -# TCP connections are kept alive and not re-established per request. -_global_aiohttp_session: aiohttp.ClientSession | None = None -_global_aiohttp_session_loop: asyncio.AbstractEventLoop | None = None +# One aiohttp session per event loop. Media loading can run on multiple live +# loops in one process. Live sessions must only be used and closed from their +# owning loop; closed-loop entries are disposed on the next lookup. +_aiohttp_sessions: dict[asyncio.AbstractEventLoop, aiohttp.ClientSession] = {} +_aiohttp_sessions_lock = threading.Lock() async def _get_aiohttp_session() -> aiohttp.ClientSession: """Return the aiohttp session owned by the current event loop.""" - global _global_aiohttp_session, _global_aiohttp_session_loop current_loop = asyncio.get_running_loop() - if ( - _global_aiohttp_session is None - or _global_aiohttp_session.closed - or _global_aiohttp_session_loop is not current_loop - ): - previous_session = _global_aiohttp_session - connector = aiohttp.TCPConnector(resolver=_ValidatingResolver()) - _global_aiohttp_session = aiohttp.ClientSession(connector=connector) - _global_aiohttp_session_loop = current_loop - if previous_session is not None and not previous_session.closed: - await previous_session.close() - return _global_aiohttp_session + with _aiohttp_sessions_lock: + stale_sessions = [ + _aiohttp_sessions.pop(loop) for loop in list(_aiohttp_sessions) if loop.is_closed() + ] + session = _aiohttp_sessions.get(current_loop) + if session is None or session.closed: + connector = aiohttp.TCPConnector(resolver=_ValidatingResolver()) + session = aiohttp.ClientSession(connector=connector) + _aiohttp_sessions[current_loop] = session + for stale_session in stale_sessions: + await stale_session.close() + return session + + +async def _close_aiohttp_session() -> None: + """Close and forget the media session owned by the current event loop.""" + current_loop = asyncio.get_running_loop() + with _aiohttp_sessions_lock: + session = _aiohttp_sessions.pop(current_loop, None) + if session is not None and not session.closed: + await session.close() # Maximum allowed response size and encoded video input size (200 MiB). @@ -1316,29 +1325,59 @@ class AudioMediaIO(BaseMediaIO[Tuple[np.ndarray, int]]): _data_uri_mime_prefix: ClassVar[Optional[str]] = "audio/" - def __init__(self, **kwargs) -> None: - # AudioMediaIO has no configurable parameters. An explicit __init__ - # is needed so that BaseMediaIO.create() does not crash when kwargs - # are present (e.g. via --media_io_kwargs or the per-request API). + def __init__( + self, decoded_memory_reserver: Optional[_DecodedMemoryReserver] = None, **kwargs + ) -> None: + self._decoded_memory_reserver = decoded_memory_reserver if kwargs: logger.warning( - "AudioMediaIO received unexpected kwargs %s — audio loading " - "has no configurable parameters; kwargs will be ignored.", + "AudioMediaIO received unexpected kwargs %s; kwargs will be ignored.", sorted(kwargs), ) def load_bytes(self, data: bytes) -> Tuple[np.ndarray, int]: - return soundfile.read(BytesIO(data)) + return self._read(BytesIO(data)) def load_base64(self, media_type: str, data: str) -> Tuple[np.ndarray, int]: - return soundfile.read(BytesIO(base64.b64decode(data))) + return self.load_bytes(base64.b64decode(data)) def load_file(self, url: str) -> Tuple[np.ndarray, int]: # Strip `file://` and unquote the path for `file:` URIs; pass # empty-scheme URLs through verbatim. parsed = urlparse(url) path = _normalize_file_uri(url) if parsed.scheme == "file" else url - return soundfile.read(path) + return self._read(path) + + def _read(self, source: Union[BytesIO, str]) -> Tuple[np.ndarray, int]: + if self._decoded_memory_reserver is None: + return soundfile.read(source) + + reservation = None + try: + # Read metadata and samples from the same handle. Reserve the float64 + # output before allocation; libsndfile's internal buffers are excluded. + with soundfile.SoundFile(source) as audio: + frames, channels = audio.frames, audio.channels + if frames < 0 or channels < 1: + raise ValueError("invalid audio frame or channel count") + decoded_bytes = frames * channels * np.dtype(np.float64).itemsize + reservation = self._decoded_memory_reserver(decoded_bytes, decoded_bytes) + shape = (frames,) if channels == 1 else (frames, channels) + samples = np.empty(shape, dtype=np.float64) + decoded = audio.read(frames=frames, out=samples) + if decoded.shape != samples.shape: + # Compressed headers may overestimate frames. Drop the decoder's + # view before shrinking the owning allocation and its budget. + decoded_shape = decoded.shape + decoded = None + samples.resize(decoded_shape) + sample_rate = audio.samplerate + reservation.commit(samples.nbytes) + return samples, sample_rate + finally: + samples = decoded = None + if reservation is not None: + reservation.release() class VideoMediaIO(BaseMediaIO[VideoData]): diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 140c146a7955..d8af206b8fe7 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -100,6 +100,14 @@ class InputProcessor(Protocol): - Register the inherited class to the model class using @register_input_processor(...) """ + supports_token_ids_input: bool = True + """Whether pre-tokenized prompts are accepted; encoder-driven processors may opt out.""" + + @property + def supported_modalities(self) -> frozenset[str]: + """Media accepted by this processor instance; text-only by default.""" + return frozenset() + model_path: any config: any tokenizer: any @@ -156,6 +164,15 @@ class BaseMultimodalInputProcessor(ABC): processes the multimodal inputs. """ + @property + def supported_modalities(self) -> frozenset[str]: + """Registered modalities, narrowed by checkpoint-specific overrides.""" + config = getattr(self, "config", None) + if config is None: + return frozenset() + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_modalities( + resolve_top_level_model_type(config)) + # Whether multimodal soft-token runs need bidirectional attention spanning # the full block. When True, the chunked-prefill scheduler must keep each # MM block intact within a single iteration (snap-up or snap-down at chunk @@ -895,6 +912,13 @@ def remove_placeholder_metadata(self, model_type: str): raise ValueError(f"Model type '{model_type}' is not registered") del self._multimodal_placeholder_by_model_type[model_type] + def get_registered_modalities(self, model_type: str) -> frozenset[str]: + """Return registered media types, or an empty set for an unknown model.""" + self._ensure_provider_imported(model_type) + metadata = self._multimodal_placeholder_by_model_type.get(model_type) + return frozenset( + metadata.placeholder_map) if metadata is not None else frozenset() + def is_valid(self, model_type: str, modality: str) -> bool: self._ensure_provider_imported(model_type) return model_type in self._multimodal_placeholder_by_model_type and \ diff --git a/tests/unittest/_torch/modeling/test_modeling_whisper.py b/tests/unittest/_torch/modeling/test_modeling_whisper.py index 90b13a7daf7c..c26d1e982e78 100644 --- a/tests/unittest/_torch/modeling/test_modeling_whisper.py +++ b/tests/unittest/_torch/modeling/test_modeling_whisper.py @@ -28,8 +28,9 @@ from tensorrt_llm._torch.models.modeling_whisper import WhisperInputProcessor, WhisperLogMelFrontend -def test_whisper_input_processor_requires_encoder_features(): +def test_whisper_input_processor_declares_input_capabilities(): assert WhisperInputProcessor.requires_encoder_features + assert WhisperInputProcessor.supports_token_ids_input is False def _synthetic_waveform_batch(n_samples: int, seed: int = 1234) -> np.ndarray: diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 23a7b27d7b15..ef4a8f9c88c1 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -87,6 +87,9 @@ def _llm(**overrides): args=args, llm_id="instance-1", tokenizer=overrides.get("_tokenizer", object()), + input_processor=SimpleNamespace( + supports_token_ids_input=overrides.get("_supports_token_ids_input", True) + ), _executor=executor, _check_health=lambda: executor.check_health(), ) @@ -424,6 +427,25 @@ async def test_abort_all_fails_the_rpc_when_any_abort_is_refused(): assert "still running" in context.abort_details +@pytest.mark.asyncio +async def test_health_probe_is_not_applicable_to_text_only_input_models(): + llm = _llm(_supports_token_ids_input=False) + + def unexpected_generate(*args, **kwargs): + pytest.fail("a token-only probe must not reach a prompt-only model") + + llm.generate_async = unexpected_generate + servicer = OpenEngineControlServicer(llm, MODEL, _inference(), kv_transfer_backend="") + response = await servicer.Health( + lifecycle_pb2.HealthRequest(include_inference_probe=True), FakeServicerContext() + ) + probe = next(check for check in response.checks if check.name == "inference_probe") + assert probe.state == lifecycle_pb2.HEALTH_STATE_READY + assert "not applicable" in probe.message + assert "inference liveness was not tested" in probe.message + assert response.state == lifecycle_pb2.HEALTH_STATE_READY + + @pytest.mark.asyncio async def test_health_inference_probe_is_bounded_and_aborts_a_stuck_request(): """A wedged engine must yield DEGRADED, not a hung RPC or a leaked request.""" @@ -549,6 +571,15 @@ async def test_text_input_is_advertised_when_a_tokenizer_exists(): assert info.supports_text_input is True +@pytest.mark.asyncio +async def test_token_ids_are_not_advertised_when_the_processor_ignores_them(): + info = await _servicer(_supports_token_ids_input=False).GetModelInfo( + model_pb2.GetModelInfoRequest(), FakeServicerContext() + ) + assert info.supports_text_input is True + assert info.supports_token_ids_input is False + + @pytest.mark.asyncio async def test_health_probe_is_counted_by_get_load_and_released(): """A probe occupies a scheduler slot, so GetLoad must report it. diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index d6daa3338491..b4f9288ad78c 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -13,6 +13,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from contextvars import ContextVar +from io import BytesIO from pathlib import Path from types import SimpleNamespace from typing import Any @@ -42,6 +43,7 @@ from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, + DefaultInputProcessor, ) from tensorrt_llm.llmapi.llm import PreprocessedInputs # noqa: E402 @@ -92,6 +94,10 @@ def __init__(self) -> None: class _ImageInputProcessor(BaseMultimodalInputProcessor): """Minimal concrete processor used only for capability detection.""" + @property + def supported_modalities(self): + return frozenset({"image", "video", "audio"}) + def __init__(self, config: _ImageConfig) -> None: self._config = config @@ -121,6 +127,28 @@ def get_preferred_media_io_kwargs(self): } +class _AudioInputProcessor: + """Registered processor that does not use BaseMultimodalInputProcessor.""" + + _registered_model_type = _ImageConfig.model_type + supported_modalities = frozenset({"image", "video", "audio"}) + + def __init__(self, config: _ImageConfig) -> None: + self.config = config + + def __call__(self, inputs, sampling_params): + raise NotImplementedError + + +class _NonCallableRegisteredInputProcessor: + """Malformed registration used to verify capability discovery fails closed.""" + + _registered_model_type = _ImageConfig.model_type + + def __init__(self, config: _ImageConfig) -> None: + self.config = config + + class _Tokenizer: def encode(self, text: str, add_special_tokens: bool = True) -> list[int]: del add_special_tokens @@ -287,11 +315,34 @@ def image_video_llm(monkeypatch): ) +@pytest.fixture +def image_video_audio_llm(monkeypatch): + _set_video_runtime_available(monkeypatch, True) + config = _ImageConfig() + input_processor = _ImageInputProcessor(config) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: ( + model_type == _ImageConfig.model_type and modality in {"image", "video", "audio"} + ), + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + return SimpleNamespace( + _hf_model_config=config, + input_processor=input_processor, + ) + + def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): llm, seen_model_types = image_llm assert request_mapping.media_capabilities_for_model(llm).modalities == {"image"} - assert seen_model_types == [_ImageConfig.model_type] * 2 + assert seen_model_types == [_ImageConfig.model_type] * 3 assert ( request_mapping.media_capabilities_for_model( SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) @@ -300,6 +351,80 @@ def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): ) +def test_media_capability_accepts_a_registered_non_base_processor(monkeypatch): + config = _ImageConfig() + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: model_type == _ImageConfig.model_type and modality == "audio", + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: "", + ) + llm = SimpleNamespace( + _hf_model_config=config, + input_processor=_AudioInputProcessor(config), + ) + + assert request_mapping.media_capabilities_for_model(llm).modalities == {"audio"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prompt", ["", "<|startoftranscript|><|fr|><|translate|><|notimestamps|>"]) +async def test_audio_text_prompt_survives_preparation_without_placeholders(monkeypatch, prompt): + import soundfile + + config = _ImageConfig() + processor = _AudioInputProcessor(config) + processor.supports_token_ids_input = False + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", lambda _, modality: modality == "audio" + ) + monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "get_placeholder", lambda *_: "") + llm = _ImageLlm(processor, config) + llm.args.reasoning_parser = None + llm.args.max_seq_len = 2048 + inference = OpenEngineInferenceServicer(llm, "audio-model") + control = OpenEngineControlServicer(llm, "audio-model", inference) + info = await control.GetModelInfo(model_pb2.GetModelInfoRequest(), FakeServicerContext()) + assert info.supports_text_input is True + assert info.supports_token_ids_input is False + encoded = BytesIO() + soundfile.write(encoded, np.zeros(32), 16000, format="WAV") + request = generation_pb2.GenerateRequest( + request_id="audio-task-prompt", + model="audio-model", + prompt=prompt, + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + raw_bytes=encoded.getvalue(), + mime_type="audio/wav", + ) + ], + ) + responses = [response async for response in inference.Generate(request, FakeServicerContext())] + assert responses[-1].HasField("finished") + inputs = llm.generate_kwargs["inputs"] + assert inputs["prompt"] == prompt + samples, rate = inputs["multi_modal_data"]["audio"][0] + assert samples.shape == (32,) and rate == 16000 + assert inference.active_request_count() == 0 + + +def test_media_capability_rejects_a_registered_non_callable_processor(monkeypatch): + config = _ImageConfig() + monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", lambda *_args: True) + llm = SimpleNamespace( + _hf_model_config=config, + input_processor=_NonCallableRegisteredInputProcessor(config), + ) + + assert request_mapping.media_capabilities_for_model(llm).modalities == set() + + @pytest.mark.asyncio @pytest.mark.parametrize("disabled", [False, True]) async def test_disabled_mm_encoder_agrees_between_discovery_and_generate(image_llm, disabled): @@ -380,6 +505,272 @@ def test_decoded_video_bytes_include_frames_and_audio(): request_mapping._decoded_video_bytes(object()) +def test_decoded_audio_bytes_count_samples_and_reject_unknown_types(): + assert request_mapping._decoded_audio_bytes((np.zeros(8, dtype=np.float32), 16000)) == 8 * 4 + with pytest.raises(TypeError, match="unsupported decoded audio type"): + request_mapping._decoded_audio_bytes(object()) + + +@pytest.mark.parametrize("audio_format", ["FLAC", "MP3"]) +@pytest.mark.parametrize("channels", [1, 2]) +@pytest.mark.parametrize("source", ["bytes", "base64", "file"]) +def test_audio_reserves_before_decoding_and_preserves_samples( + audio_format, channels, source, tmp_path, monkeypatch +): + import base64 + + import soundfile + + samples = np.arange(32 * channels, dtype=np.float64).reshape(32, channels) / 128 + encoded = BytesIO() + soundfile.write(encoded, samples, 16000, format=audio_format) + data = encoded.getvalue() + expected, _ = soundfile.read(BytesIO(data)) + with soundfile.SoundFile(BytesIO(data)) as audio: + reserved_bytes = audio.frames * audio.channels * np.dtype(np.float64).itemsize + budget = request_mapping._DecodedMediaBudget(reserved_bytes) + original_read = soundfile.SoundFile.read + + def checked_read(audio, *args, **kwargs): + assert (budget._used, budget._reserved_retained) == (reserved_bytes, reserved_bytes) + assert kwargs["out"].nbytes == reserved_bytes + return original_read(audio, *args, **kwargs) + + monkeypatch.setattr(soundfile.SoundFile, "read", checked_read) + loader = media_io.AudioMediaIO(decoded_memory_reserver=budget.reserve) + if source == "bytes": + decoded, rate = loader.load_bytes(data) + elif source == "base64": + decoded, rate = loader.load_base64( + "audio/mpeg" if audio_format == "MP3" else "audio/flac", + base64.b64encode(data).decode(), + ) + else: + path = tmp_path / f"audio.{audio_format.lower()}" + path.write_bytes(data) + decoded, rate = loader.load_file(path.as_uri()) + if audio_format == "MP3": + # Buffered and ordinary MP3 reads can differ by float32 rounding. + np.testing.assert_allclose(decoded, expected, rtol=0, atol=np.finfo(np.float32).eps) + else: + np.testing.assert_array_equal(decoded, expected) + assert rate == 16000 + assert (budget._used, budget._retained, budget._reserved_retained) == ( + expected.nbytes, + expected.nbytes, + 0, + ) + + +@pytest.mark.parametrize("outcome", ["oversized", "failure", "cancelled"]) +def test_audio_budget_rejects_before_allocation_and_cleans_up(outcome, monkeypatch): + import soundfile + + encoded = BytesIO() + soundfile.write(encoded, np.zeros((32, 2)), 16000, format="FLAC") + budget = request_mapping._DecodedMediaBudget(511 if outcome == "oversized" else 512) + loader = media_io.AudioMediaIO(decoded_memory_reserver=budget.reserve) + + def read(audio, *args, **kwargs): + if outcome == "oversized": + pytest.fail("oversized audio must be rejected before decoding") + if outcome == "failure": + raise RuntimeError("decode failed") + if outcome == "cancelled": + raise asyncio.CancelledError() + + monkeypatch.setattr(soundfile.SoundFile, "read", read) + if outcome == "oversized": + monkeypatch.setattr( + np, "empty", lambda *args, **kwargs: pytest.fail("allocated oversized audio") + ) + error = { + "oversized": ValueError, + "failure": RuntimeError, + "cancelled": asyncio.CancelledError, + }[outcome] + with pytest.raises(error): + loader.load_bytes(encoded.getvalue()) + assert (budget._used, budget._retained, budget._reserved_retained) == (0, 0, 0) + + +@pytest.mark.parametrize("channels", [1, 2]) +@pytest.mark.parametrize("decoded_frames", [0, 17]) +def test_audio_short_read_preserves_samples_and_accounts_for_owning_storage( + channels, decoded_frames, monkeypatch +): + import soundfile + + encoded = BytesIO() + samples = np.arange(32 * channels, dtype=np.float64).reshape(32, channels) / 128 + soundfile.write(encoded, samples, 16000, format="FLAC") + expected, _ = soundfile.read(BytesIO(encoded.getvalue())) + budget = request_mapping._DecodedMediaBudget(samples.nbytes) + original_read = soundfile.SoundFile.read + + def short_read(audio, *args, **kwargs): + assert budget._used == samples.nbytes + return original_read(audio, *args, **kwargs)[:decoded_frames] + + monkeypatch.setattr(soundfile.SoundFile, "read", short_read) + decoded, rate = media_io.AudioMediaIO(decoded_memory_reserver=budget.reserve).load_bytes( + encoded.getvalue() + ) + np.testing.assert_array_equal(decoded, expected[:decoded_frames]) + assert rate == 16000 + assert decoded.flags.owndata + assert decoded.base is None + assert (budget._used, budget._retained, budget._reserved_retained) == ( + decoded.nbytes, + decoded.nbytes, + 0, + ) + + +@pytest.mark.asyncio +async def test_audio_request_uses_one_shared_budget_without_double_charging( + image_video_audio_llm, monkeypatch +): + import soundfile + + encoded = BytesIO() + soundfile.write(encoded, np.zeros(32), 16000, format="FLAC") + monkeypatch.setattr(request_mapping, "_MAX_TOTAL_DECODED_MEDIA_BYTES", 256) + item = generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, raw_bytes=encoded.getvalue(), mime_type="audio/flac" + ) + request = generation_pb2.GenerateRequest(prompt="transcribe", media=[item]) + mapped = await request_mapping._input_from_request(request, image_video_audio_llm) + assert mapped["multi_modal_data"]["audio"][0][0].nbytes == 256 + request.media.append(item) + with pytest.raises(ValueError, match="decoded media exceeds"): + await request_mapping._input_from_request(request, image_video_audio_llm) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("processor_kind", ["gemma4", "nemotron"]) +@pytest.mark.parametrize("audio_encoder", [False, True]) +async def test_audio_capability_follows_loaded_processor(processor_kind, audio_encoder): + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor + from tensorrt_llm._torch.models.modeling_nemotron_h_multimodal import ( + NemotronHMultimodalInputProcessor, + ) + + if processor_kind == "gemma4": + processor = Gemma4InputProcessor.__new__(Gemma4InputProcessor) + config = SimpleNamespace( + model_type="gemma4", + vision_config=object(), + audio_config=object() if audio_encoder else None, + ) + else: + processor = NemotronHMultimodalInputProcessor.__new__(NemotronHMultimodalInputProcessor) + config = SimpleNamespace(model_type="nemotron_h_omni") + processor._audio_extractor = object() if audio_encoder else None + processor._config = config + llm = _ImageLlm(processor, config) + llm.args.reasoning_parser = None + llm.args.max_seq_len = 2048 + inference = OpenEngineInferenceServicer(llm, "audio-model") + control = OpenEngineControlServicer(llm, "audio-model", inference) + info = await control.GetModelInfo(model_pb2.GetModelInfoRequest(), FakeServicerContext()) + assert ("audio" in info.extra["trtllm_supported_modalities"]) is audio_encoder + assert "image" in info.extra["trtllm_supported_modalities"] + assert {"image", "video"} <= processor.supported_modalities + assert ("audio" in processor.supported_modalities) is audio_encoder + if not audio_encoder: + request = generation_pb2.GenerateRequest( + request_id="unsupported-audio", + model="audio-model", + prompt="transcribe", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_AUDIO, + raw_bytes=b"unused", + mime_type="audio/wav", + ) + ], + ) + with pytest.raises(UnsupportedFeatureError, match="audio"): + await request_mapping._input_from_request(request, llm) + context = FakeServicerContext() + with pytest.raises(AbortError): + async for _ in inference.Generate(request, context): + pass + assert context.abort_code == grpc.StatusCode.UNIMPLEMENTED + + +@pytest.mark.parametrize("audio_encoder", [False, True]) +def test_gemma4_without_vision_encoder_omits_images_and_video(audio_encoder): + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor + + processor = Gemma4InputProcessor.__new__(Gemma4InputProcessor) + processor._config = SimpleNamespace( + model_type="gemma4", vision_config=None, audio_config=object() if audio_encoder else None + ) + expected = {"audio"} if audio_encoder else frozenset() + assert processor.supported_modalities == expected + assert ( + request_mapping.media_capabilities_for_model( + SimpleNamespace(input_processor=processor) + ).modalities + == expected + ) + + +@pytest.mark.parametrize("config_state", ["absent", "none"]) +def test_gemma4_missing_processor_config_fails_closed(config_state): + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor + + config = SimpleNamespace(model_type="gemma4", audio_config=object()) + processor = Gemma4InputProcessor.__new__(Gemma4InputProcessor) + if config_state == "none": + processor._config = None + llm = SimpleNamespace(input_processor=processor, _hf_model_config=config) + assert processor.supported_modalities == frozenset() + assert request_mapping.media_capabilities_for_model(llm).modalities == frozenset() + + +def test_whisper_declares_audio_without_multimodal_base(): + from tensorrt_llm._torch.models.modeling_whisper import WhisperInputProcessor + + processor = WhisperInputProcessor.__new__(WhisperInputProcessor) + processor.config = SimpleNamespace(model_type="whisper") + assert processor.supported_modalities == {"audio"} + llm = SimpleNamespace(input_processor=processor) + assert request_mapping.media_capabilities_for_model(llm).modalities == {"audio"} + + +def test_multimodal_processor_defaults_to_registered_modalities(monkeypatch): + from tensorrt_llm.inputs.registry import MultimodalPlaceholderMetadata + + class RegisteredProcessor(_ImageInputProcessor): + supported_modalities = BaseMultimodalInputProcessor.supported_modalities + + monkeypatch.setitem( + MULTIMODAL_PLACEHOLDER_REGISTRY._multimodal_placeholder_by_model_type, + _ImageConfig.model_type, + MultimodalPlaceholderMetadata(placeholder_map={"image": "", "audio": "