diff --git a/requirements.txt b/requirements.txt index a243f9b26ab5..b8b0435f83ee 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,6 +7,7 @@ cuda-python>=13 diffusers>=0.37.1 ftfy lark +lazy_loader~=0.5 mpi4py numpy>=2.0.0,<2.4 # numba 0.63.1 requires numpy<2.4 onnx>=1.21.0 @@ -91,3 +92,4 @@ smg-grpc-proto>=0.4.2 cache-dit>=1.3.5 librosa msgpack +uvloop>=0.19.0 diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py index 0acd057fdaa6..3dd866015dd3 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py @@ -19,7 +19,8 @@ import functools import math import os -from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast +from typing import (Any, Callable, Dict, List, Optional, Tuple, TypedDict, + Union, cast) import torch import torch.nn.functional as F @@ -63,6 +64,34 @@ def has_raw_multimodal_payload(param: MultimodalParams) -> bool: }) +def _get_cached_merged_typed_dict(schema, cache): + """Return a stable copy of ProcessorMixin's ephemeral merged TypedDict. + + `ProcessorMixin._merge_kwargs` creates a fresh + `TypedDict("merged_typed_dict", ...)` for image/video kwargs on every + processor call. `huggingface_hub` caches strict dataclass validators by + schema-object identity, so a fresh type defeats that cache. Reuse an + equivalent TypedDict class keyed by (totality, annotation items) so the + upstream validator hits cache and skips the recursive type validation. + """ + if getattr(schema, "__name__", None) != "merged_typed_dict": + return schema + try: + cache_key = (getattr(schema, "__total__", + True), tuple(schema.__annotations__.items())) + cached_schema = cache.get(cache_key) + except TypeError: + return schema + if cached_schema is None: + cached_schema = TypedDict( + "merged_typed_dict", + dict(schema.__annotations__), + total=getattr(schema, "__total__", True), + ) + cache[cache_key] = cached_schema + return cached_schema + + @functools.lru_cache(maxsize=None) def _install_processor_output_validation_filter(): """Install a process-wide filter over transformers' ``validate_typed_dict``. @@ -107,6 +136,7 @@ def _install_processor_output_validation_filter(): "No transformers module exposes validate_typed_dict; " "cannot patch processor output validation.") base_orig = binders[0].validate_typed_dict + merged_schema_cache: dict = {} def _filtered_validate(schema, data): if isinstance(data, dict): @@ -114,6 +144,7 @@ def _filtered_validate(schema, data): k: v for k, v in data.items() if k not in _PROCESSOR_OUTPUT_KEYS } + schema = _get_cached_merged_typed_dict(schema, merged_schema_cache) return base_orig(schema, data) for b in binders: diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 8bccc3edc980..4fc949c2b2db 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -2,6 +2,7 @@ # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy +import math import re from functools import lru_cache from typing import Any, Dict, List, Optional, Tuple, Union @@ -161,6 +162,61 @@ def _expand_prompt_token_ids_for_mm_handoff( ) +def _decide_do_sample_frames( + video_datas: Optional[List[Any]], + mm_processor_kwargs: Dict[str, Any], +) -> bool: + """Pick a single `do_sample_frames` flag for the HF processor call. + + HF's video processor takes a scalar `do_sample_frames` that applies to + every video in the request. Decide it as follows: + + 1. If `mm_processor_kwargs.do_sample_frames` is explicitly set + (True or False), honor it. + 2. If the caller supplies no frame target (`num_frames` / `fps`), + match HF's class default, which samples frames (returns True). + 3. Otherwise, for each video compute the target frame count from the + kwargs (`num_frames` directly, or `floor(duration * fps)` if + `fps` is given) and compare to `len(vd.frames)`. If any video + needs a different count, the batch is sampled (returns True). + + Per-video targets that match the IO-decoded count don't need HF + sampling; the all-or-nothing reduction over the batch means a single + video needing resampling pulls the rest along through a no-op + identity `np.linspace`. + """ + if "do_sample_frames" in mm_processor_kwargs: + return bool(mm_processor_kwargs["do_sample_frames"]) + + if not video_datas: + return False + + user_num_frames = mm_processor_kwargs.get("num_frames") + user_fps = mm_processor_kwargs.get("fps") + has_num_frames = user_num_frames is not None and user_num_frames != -1 + has_fps = user_fps is not None and user_fps != -1 + + # No explicit frame target from the caller: defer to HF's class-default + # sampling (the stock processor sets `do_sample_frames=True` when neither + # `num_frames` nor `fps` is given). Returning False here would hand the + # IO-decoded frames straight to HF unchanged and diverge from stock HF + # whenever the IO loader decoded a different number of frames than HF's + # default sampler would select. + if not has_num_frames and not has_fps: + return True + + for vd in video_datas: + n_decoded = len(vd.frames) + if has_num_frames: + n_target = user_num_frames + else: # has_fps + duration = (vd.metadata or {}).get("duration") or 0 + n_target = math.floor(duration * user_fps) + if n_target != n_decoded: + return True + return False + + class Qwen3VLInputProcessorBase(Qwen2VLInputProcessorBase): """Qwen3-VL input processor. @@ -253,19 +309,42 @@ def _preprocess( if videos and isinstance(videos[0][0], torch.Tensor): do_rescale = False - # Forward video metadata only when the caller opts into per-request kwargs; - # the default path pre-samples frames in the IO loader, so unconditional - # metadata triggers IndexError in HF's _decode_and_sample_videos. - video_metadata = ( - [vd.metadata for vd in video_datas] if video_datas and mm_processor_kwargs else None - ) - - # num_frames and fps are mutually exclusive in the HF processor's sample_frames. - # If the caller set num_frames without fps, null fps explicitly so the class-level - # default fps=2 does not interfere. - proc_kwargs = dict(mm_processor_kwargs) - if "num_frames" in proc_kwargs and "fps" not in proc_kwargs: - proc_kwargs["fps"] = None + do_sample_frames = _decide_do_sample_frames(video_datas, mm_processor_kwargs) + + # Pass `do_sample_frames` plus, when sampling is needed, the + # caller's `num_frames` / `fps` target. Everything else the caller + # supplied (resize, normalize knobs, etc.) flows through unchanged. + proc_kwargs: Dict[str, Any] = {"do_sample_frames": do_sample_frames} + for k, v in mm_processor_kwargs.items(): + if k in ("num_frames", "fps", "do_sample_frames"): + continue + proc_kwargs[k] = v + if do_sample_frames: + if "num_frames" in mm_processor_kwargs: + proc_kwargs["num_frames"] = mm_processor_kwargs["num_frames"] + if "fps" in mm_processor_kwargs: + proc_kwargs["fps"] = mm_processor_kwargs["fps"] + elif "num_frames" in mm_processor_kwargs: + # HF's `sample_frames` honors `num_frames` only when `fps` is + # not also set; the class-default `fps=2` would otherwise cap + # the returned count below the caller's requested + # `num_frames` for short clips. Null `fps` so `num_frames` is + # respected verbatim. + proc_kwargs["fps"] = None + + # Forward per-video metadata with `total_num_frames` rewritten to the + # actual decoded frame count. HF's `sample_frames` computes indices + # via `np.linspace(0, total_num_frames - 1, num_frames)` and indexes + # the frame tensor with them; the rewrite keeps those indices in + # range and the no-sampling path consistent for downstream qwen3vl + # code that consults the metadata. + video_metadata: Optional[List[Dict[str, Any]]] = None + if video_datas: + video_metadata = [] + for vd in video_datas: + m = dict(vd.metadata or {}) + m["total_num_frames"] = len(vd.frames) + video_metadata.append(m) return self.processor( text=[text], diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index bb5cab5ddd81..440b6ac42d41 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -15,6 +15,7 @@ import click import torch +import uvloop import yaml from strenum import StrEnum from torch.cuda import device_count @@ -350,7 +351,9 @@ def launch_server( disagg_cluster_config: Optional[DisaggClusterConfig] = None, multimodal_server_config: Optional[MultimodalServerConfig] = None, served_model_name: Optional[str] = None, - allow_request_chat_template: bool = False): + allow_request_chat_template: bool = False, + input_processor_workers: int = 8, + media_load_workers: int = 8): backend = llm_args["backend"] model = served_model_name or llm_args["model"] @@ -394,14 +397,16 @@ def launch_server( disagg_cluster_config=disagg_cluster_config, multimodal_server_config=multimodal_server_config, chat_template=chat_template, - allow_request_chat_template=allow_request_chat_template) + allow_request_chat_template=allow_request_chat_template, + input_processor_workers=input_processor_workers, + media_load_workers=media_load_workers) _apply_fastapi_middlewares(server.app, middleware) # Optionally disable GC (default: not disabled) if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": gc.disable() - asyncio.run(server(host, port, sockets=[s])) + uvloop.run(server(host, port, sockets=[s])) def launch_grpc_server(host: str, @@ -527,7 +532,7 @@ def signal_handler(): logger.info("Shutdown complete") - asyncio.run(serve_grpc_async()) + uvloop.run(serve_grpc_async()) def launch_mm_encoder_server( @@ -548,7 +553,7 @@ def launch_mm_encoder_server( metadata_server_cfg=metadata_server_cfg, tool_parser=None, allow_request_chat_template=allow_request_chat_template) - asyncio.run(server(host, port)) + uvloop.run(server(host, port)) def launch_visual_gen_server( @@ -601,7 +606,7 @@ def launch_visual_gen_server( metadata_server_cfg=metadata_server_cfg, tool_parser=None) _apply_fastapi_middlewares(server.app, middleware) - asyncio.run(server(host, port, sockets=[s])) + uvloop.run(server(host, port, sockets=[s])) class ChoiceWithAlias(click.Choice): @@ -769,6 +774,22 @@ def convert(self, value: Any, param: Optional["click.Parameter"], help="Number of workers to postprocess raw responses " "to comply with OpenAI protocol.", status="prototype") +@stability_option("--input-processor-workers", + "input_processor_workers", + type=click.IntRange(min=1), + default=8, + help="Size of the dedicated thread pool that runs the HF " + "input processor (multimodal preprocess) on the chat " + "and completion endpoints.", + status="prototype") +@stability_option("--media-load-workers", + "media_load_workers", + type=click.IntRange(min=1), + default=8, + help="Size of the dedicated thread pool that decodes media " + "payloads (image / video / audio) for multimodal " + "requests.", + status="prototype") @stability_option("--trust_remote_code", is_flag=True, default=False, @@ -945,7 +966,8 @@ def serve( moe_expert_parallel_size: Optional[int], moe_cluster_parallel_size: Optional[int], gpus_per_node: Optional[int], free_gpu_memory_fraction: float, kv_cache_dtype: str, - num_postprocess_workers: int, trust_remote_code: bool, + num_postprocess_workers: int, input_processor_workers: int, + media_load_workers: int, trust_remote_code: bool, revision: Optional[str], extra_llm_api_options: Optional[str], reasoning_parser: Optional[str], tool_parser: Optional[str], metadata_server_config_file: Optional[str], server_role: Optional[str], @@ -1146,7 +1168,9 @@ def _serve_llm(): disagg_cluster_config, multimodal_server_config, served_model_name=served_model_name, - allow_request_chat_template=allow_request_chat_template) + allow_request_chat_template=allow_request_chat_template, + input_processor_workers=input_processor_workers, + media_load_workers=media_load_workers) def _serve_visual_gen(): parsed_visual_gen_args = (VisualGenArgs.from_yaml(visual_gen_args) @@ -1421,7 +1445,7 @@ def disaggregated( if os.getenv("TRTLLM_DISAGG_SERVER_DISABLE_GC", "1") == "1": gc.disable() - asyncio.run(server(disagg_cfg.hostname, disagg_cfg.port, sockets=[s])) + uvloop.run(server(disagg_cfg.hostname, disagg_cfg.port, sockets=[s])) def set_cuda_device(): diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 3e411d550d4d..bd57c93e2d3a 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -4,28 +4,50 @@ import asyncio import base64 +import functools import ipaddress import math import os import socket import tempfile from abc import ABC, abstractmethod +from concurrent.futures import Executor from io import BytesIO from pathlib import Path from types import MappingProxyType -from typing import Any, Dict, Generic, List, Literal, Mapping, Optional, Tuple, Type, TypeVar, Union +from typing import ( + Any, + ClassVar, + Dict, + Generic, + List, + Literal, + Mapping, + Optional, + Tuple, + Type, + TypeVar, + Union, +) from urllib.parse import unquote, urljoin, urlparse import aiohttp +import lazy_loader as lazy import numpy as np import requests import soundfile import torch +from packaging.version import Version from PIL import Image from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData from tensorrt_llm.logger import logger +# Lazy import: OpenCV is large and only needed when the cv2-backed video +# decode path is exercised. The proxy triggers the actual `import cv2` on +# first attribute access (e.g. `cv2.VideoCapture`). +cv2 = lazy.load("cv2") + def rgba_to_rgb( image: Image.Image, background_color: Union[tuple[int, int, int], list[int]] = (255, 255, 255) @@ -56,10 +78,14 @@ def convert_image_mode(image: Image.Image, to_mode: str) -> Image.Image: # Canonical set of supported media modalities for Pydantic field validation. MediaModality = Literal["image", "video", "audio"] -# Output representations supported by `ImageMediaIO` and `VideoMediaIO`: -# `"pt"` → `torch.Tensor`, `"pil"` → `PIL.Image.Image` (per frame for video). +# Output representations supported by `ImageMediaIO`: +# `"pt"` -> `torch.Tensor`, `"pil"` -> `PIL.Image.Image`. _SUPPORTED_IMAGE_FORMATS = ("pt", "pil") +# Output representations supported by `VideoMediaIO`. See +# `_load_video_by_cv2` for the per-format contract. +_SUPPORTED_VIDEO_FORMATS = ("pt", "hwc_uint8", "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. @@ -312,25 +338,118 @@ def extract_audio_from_video( return audio, target_sr +def _select_cv2_stream_buffered_backend() -> Optional[int]: + """Return a VideoCapture backend that can read from a Python `BytesIO`. + + Returns `None` if no such backend is available in this OpenCV build — + the caller falls back to the tempfile path. The `videoio_registry` + stream-buffered API was introduced in OpenCV 4.13.0, so older builds + are gated out by a version check. + + Plugin-provided backends must implement the stream-buffered API at + version 1.2+; older plugins exist that load fine but crash on stream + open. Built-in backends (the FFMPEG path in the PyPI wheels) are always + safe to use. + """ + # The stream-buffered API was introduced in OpenCV 4.13.0. Older builds + # don't have `cv2.videoio_registry.getStreamBufferedBackends`, so signal + # "no usable backend" and let the caller fall back to the tempfile path. + if Version(cv2.__version__) < Version("4.13.0"): + return None + + vr = cv2.videoio_registry + + # `getStreamBufferedBackends()` enumerates every backend in this OpenCV + # build that *claims* to support stream-buffered reads. We filter that + # list down to one that's safe to use: + # 1. `hasBackend(backend)` — the backend's shared library is actually + # loadable in this process (a backend can be enumerated by the + # registry but missing at link time on a stripped build). + # 2. For plugin (non-built-in) backends, check the plugin's announced + # stream-buffered API/ABI version. Plugins reporting ABI<1 or + # (ABI==1, API<2) load fine but segfault on stream open in the + # wild — skip them. Built-in backends (FFMPEG in the PyPI wheels) + # have no plugin-version concept and are always safe. + # The first surviving backend is returned. Order is OpenCV's preference. + for backend in vr.getStreamBufferedBackends(): + if not vr.hasBackend(backend): + continue + if not vr.isBackendBuiltIn(backend): + _, abi, api = vr.getStreamBufferedBackendPluginVersion(backend) + if abi < 1 or (abi == 1 and api < 2): + continue + return backend + return None + + +# Prefer /dev/shm (tmpfs, RAM-backed) over the system tempdir to avoid a disk +# round-trip; falls back to `None` (system default) when unavailable. Only +# consumed as the `dir=` argument to `tempfile.NamedTemporaryFile`, so the +# B108 predictable-path concern does not apply. +_VIDEO_TEMPFILE_DIR: Optional[str] = ( # nosec B108 + "/dev/shm" # nosec B108 + if os.path.isdir("/dev/shm") and os.access("/dev/shm", os.W_OK) # nosec B108 + else None +) + + def _load_video_by_cv2( - video: str, + video: Union[str, bytes], num_frames: int = 10, fps: int = 30, format: str = "pt", device: str = "cpu", extract_audio: bool = False, + cv2_backend: Optional[int] = None, ) -> VideoData: - # Keep this import local to avoid importing cv2 if not needed - import cv2 - - assert format in ["pt", "pil"], "format must be either Pytorch or PIL" - - vidcap = cv2.VideoCapture(video) + """Decode a video and return sampled frames as a list. + + `video` is either a file path / URL (`str`) or raw mp4 bytes. When + bytes are passed the caller must also pass `cv2_backend` (a + stream-buffered backend id from `_select_cv2_stream_buffered_backend`); + cv2.VideoCapture is then opened over a `BytesIO` via that backend, no + tempfile required. Callers that have no stream-buffered backend + available should spill to a tempfile themselves and pass a path. + + `format` controls the per-frame return type: + `"pt"` - list[torch.Tensor], dtype=float32, shape=(C, H, W), + range [0, 1]. Frames are rescaled and permuted to + CHW inside this function. + `"hwc_uint8"` - list[np.ndarray], dtype=uint8, shape=(H, W, 3). + Frames are returned unchanged from the cv2 decode; + rescale and permute are deferred to the downstream + HF processor (its `do_rescale=True` path). + `"pil"` - list[PIL.Image], one per sampled frame. + """ + assert format in ("pt", "hwc_uint8", "pil"), "format must be one of 'pt', 'hwc_uint8', 'pil'" + + # Open the source. Two cases: + # (a) `video` is a file path / URL str -> hand it straight to cv2. + # (b) `video` is mp4 bytes -> feed cv2 from an in-memory BytesIO via + # the caller-supplied stream-buffered backend. The buffer is held + # alive in `video_buf` because cv2 keeps a non-owning view into it. + video_buf: Optional[BytesIO] = None + if isinstance(video, (bytes, bytearray, memoryview)): + if cv2_backend is None: + raise ValueError( + "cv2_backend must be provided when `video` is bytes; " + "callers without a stream-buffered backend should spill " + "the bytes to a tempfile and pass the path instead." + ) + video_buf = BytesIO(bytes(video)) + vidcap = cv2.VideoCapture(video_buf, cv2_backend, []) + else: + vidcap = cv2.VideoCapture(video) try: if not vidcap.isOpened(): + src_repr = ( + f"<{len(video)} bytes>" + if isinstance(video, (bytes, bytearray, memoryview)) + else f"'{video}'" + ) raise ValueError( - f"Video '{video}' could not be opened. Make sure opencv is installed with video support." + f"Video {src_repr} could not be opened. Make sure opencv is installed with video support." ) frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT)) @@ -371,16 +490,20 @@ def _load_video_by_cv2( valid_indices = [i for i in indices if i in raw_frames] if format == "pt": - # Bypass PIL: direct numpy HWC uint8 -> torch CHW float32 - loaded_frames = [ - torch.from_numpy(raw_frames[i]) - .permute(2, 0, 1) - .float() - .div_(255.0) - .to(device=device) - for i in valid_indices - ] - else: + # uint8 -> float32 + /255 rescale done once on the stacked buffer + # so the dtype conversion is a single memory pass and there's one + # Python torch call instead of one per frame. + stacked_uint8 = np.stack([raw_frames[i] for i in valid_indices]) + stacked_f32 = stacked_uint8.astype(np.float32) * (1.0 / 255.0) + tensor_nchw = torch.from_numpy(stacked_f32).permute(0, 3, 1, 2).contiguous() + if device != "cpu": + tensor_nchw = tensor_nchw.to(device) + loaded_frames = list(torch.unbind(tensor_nchw, dim=0)) + elif format == "hwc_uint8": + # Return uint8 HWC frames as-is; let the downstream HF processor + # handle this internally. + loaded_frames = [raw_frames[i] for i in valid_indices] + else: # "pil" loaded_frames = [Image.fromarray(raw_frames[i]) for i in valid_indices] metadata = { @@ -396,8 +519,16 @@ def _load_video_by_cv2( audio = None if extract_audio: + # extract_audio_from_video accepts Union[str, BytesIO]. When the + # in-memory cv2 path was taken, `video` is still the original + # bytes parameter — wrap in a fresh BytesIO (the one consumed by + # cv2 is at end-of-stream). When the tempfile fallback was used, + # `video` is the file path string and is passed through as-is. + audio_source = ( + BytesIO(bytes(video)) if isinstance(video, (bytes, bytearray, memoryview)) else video + ) try: - audio_samples, audio_sample_rate = extract_audio_from_video(video) + audio_samples, audio_sample_rate = extract_audio_from_video(audio_source) audio = AudioData(samples=audio_samples, sample_rate=audio_sample_rate) except ValueError as e: if "No audio stream found" in str(e): @@ -428,6 +559,21 @@ class BaseMediaIO(ABC, Generic[_MediaT]): modalities. """ + # Executor used by `async_load` to run blocking decode work off the + # event loop. `None` selects the asyncio loop's default executor. + # 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 + + @classmethod + def set_executor(cls, executor: Executor) -> None: + """Publish a shared executor for blocking decode work. + + Affects every existing and future subclass instance. Calling more + than once replaces the previously configured executor. + """ + BaseMediaIO._executor = executor + @classmethod def create( cls, @@ -476,13 +622,15 @@ async def async_load(self, url: str) -> _MediaT: """Fetch and decode media from a URL. Dispatches on scheme: http/https (remote fetch), data: (inline - base64), or file:// / bare path (local file). + base64), or file:// / bare path (local file). Blocking decode work + runs on the executor published via `set_executor`, falling + back to the asyncio loop's default executor when none is set. """ parsed = urlparse(url) if parsed.scheme in ("http", "https"): session = await _get_aiohttp_session() data = await _safe_aiohttp_get(url, session=session) - return await asyncio.to_thread(self.load_bytes, data) + return await self._run_in_executor(self.load_bytes, data) elif parsed.scheme == "data": data_spec, b64_data = parsed.path.split(",", 1) parts = data_spec.split(";", 1) @@ -490,12 +638,20 @@ async def async_load(self, url: str) -> _MediaT: encoding = parts[1] if len(parts) > 1 else "" if encoding != "base64": raise NotImplementedError("Only base64 data URLs are supported for now.") - return await asyncio.to_thread(self.load_base64, media_type, b64_data) + return await self._run_in_executor(self.load_base64, media_type, b64_data) elif parsed.scheme in ("", "file"): - return await asyncio.to_thread(self.load_file, url) + return await self._run_in_executor(self.load_file, url) else: raise ValueError(f"Unsupported URL scheme: {parsed.scheme!r}") + @staticmethod + async def _run_in_executor(fn, *args, **kwargs): + """Run a blocking decode callable on the configured executor.""" + if kwargs: + fn = functools.partial(fn, **kwargs) + loop = asyncio.get_event_loop() + return await loop.run_in_executor(BaseMediaIO._executor, fn, *args) + class ImageMediaIO(BaseMediaIO[Union[Image.Image, torch.Tensor]]): """I/O for the image modality.""" @@ -558,12 +714,15 @@ def __init__( self, num_frames: int = 10, fps: int = 30, - format: str = "pt", + format: str = "hwc_uint8", device: str = "cpu", extract_audio: bool = False, ) -> None: - if format not in _SUPPORTED_IMAGE_FORMATS: - raise ValueError(f"format must be one of {_SUPPORTED_IMAGE_FORMATS}, got {format!r}") + # Default format is `"hwc_uint8"`: frames are returned as uint8 HWC + # and the downstream HF processor handles rescale + permute. Pass + # `format="pt"` to get pre-rescaled float32 CHW tensors instead. + if format not in _SUPPORTED_VIDEO_FORMATS: + raise ValueError(f"format must be one of {_SUPPORTED_VIDEO_FORMATS}, got {format!r}") self._num_frames = num_frames self._fps = fps self._format = format @@ -590,11 +749,27 @@ def merge_kwargs( return merged def load_bytes(self, data: bytes) -> VideoData: - with tempfile.NamedTemporaryFile(delete=True, suffix=".mp4") as f: - f.write(data) - f.flush() + # In-memory fast path when cv2 has a stream-buffered backend; spill + # to a tempfile otherwise so cv2 can open it by path. The tempfile + # is unlinked on context exit; the inode survives until cv2 closes + # its own fd (Linux semantics), so the decode inside the `with` + # block reads safely. + cv2_backend = _select_cv2_stream_buffered_backend() + if cv2_backend is not None: + return _load_video_by_cv2( + data, + self._num_frames, + self._fps, + self._format, + self._device, + extract_audio=self._extract_audio, + cv2_backend=cv2_backend, + ) + with tempfile.NamedTemporaryFile(suffix=".mp4", dir=_VIDEO_TEMPFILE_DIR) as tmp: + tmp.write(data) + tmp.flush() return _load_video_by_cv2( - f.name, + tmp.name, self._num_frames, self._fps, self._format, diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 1424e409fc98..030d4bce94bd 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1,6 +1,7 @@ #!/usr/bin/env python import asyncio import base64 +import functools import json import os import re @@ -10,6 +11,7 @@ import traceback import uuid from collections import deque +from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from datetime import datetime from http import HTTPStatus @@ -34,6 +36,7 @@ from tensorrt_llm.executor.postproc_worker import PostprocParams from tensorrt_llm.inputs import prompt_inputs from tensorrt_llm.inputs.data import TokensPrompt +from tensorrt_llm.inputs.media_io import BaseMediaIO from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.utils import ConversationMessage, apply_chat_template from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams @@ -189,7 +192,9 @@ def __init__( disagg_cluster_config: Optional[DisaggClusterConfig] = None, multimodal_server_config: Optional[MultimodalServerConfig] = None, chat_template: Optional[str] = None, - allow_request_chat_template: bool = False): + allow_request_chat_template: bool = False, + input_processor_workers: int = 8, + media_load_workers: int = 8): self.generator = generator self._is_visual_gen = isinstance(generator, VisualGen) self.tool_parser = tool_parser @@ -203,6 +208,20 @@ def __init__( self.host = None self.port = None + # Dedicated thread pools for the chat / completion path. Keeping + # multimodal preprocessing and decode work off the asyncio default + # executor avoids contention with unrelated `to_thread` callers and + # lets the two stages be sized independently. + self._input_proc_executor = ThreadPoolExecutor( + max_workers=input_processor_workers, + thread_name_prefix="trtllm_inputproc", + ) + self._media_load_executor = ThreadPoolExecutor( + max_workers=media_load_workers, + thread_name_prefix="trtllm_media_load", + ) + BaseMediaIO.set_executor(self._media_load_executor) + model_dir = Path(model) if model_dir.exists() and model_dir.is_dir(): self.model = model_dir.name @@ -1318,9 +1337,11 @@ async def chat_stream_generator( generate_inputs = prompt preprocess_fn = getattr(self.generator, "preprocess", None) if preprocess_fn is not None: - generate_inputs = await asyncio.to_thread( - preprocess_fn, prompt, sampling_params, - disaggregated_params) + loop = asyncio.get_event_loop() + generate_inputs = await loop.run_in_executor( + self._input_proc_executor, + functools.partial(preprocess_fn, prompt, sampling_params, + disaggregated_params)) promise = self.generator.generate_async( inputs=generate_inputs, @@ -1619,8 +1640,11 @@ async def generator_wrapper(generator: AsyncIterator[Any]): prompt = prompt_inputs(prompt) if prompt.get("prompt") is not None: - prompt_token_ids, extra_processed_inputs = await asyncio.to_thread( - self.generator.input_processor, prompt, sampling_params) + loop = asyncio.get_event_loop() + prompt_token_ids, extra_processed_inputs = await loop.run_in_executor( + self._input_proc_executor, + functools.partial(self.generator.input_processor, + prompt, sampling_params)) tokens_prompt = TokensPrompt( prompt_token_ids=prompt_token_ids, query_token_ids=extra_processed_inputs.get( diff --git a/tests/unittest/_torch/modeling/test_modeling_multimodal_utils.py b/tests/unittest/_torch/modeling/test_modeling_multimodal_utils.py new file mode 100644 index 000000000000..b32613e39093 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_multimodal_utils.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 + +from typing import Optional, TypedDict + +from tensorrt_llm._torch.models.modeling_multimodal_utils import _get_cached_merged_typed_dict + + +def test_merged_typed_dict_reuses_cached_schema(): + """Two `merged_typed_dict` schemas with the same shape collapse to one cached + class. This is what makes huggingface_hub's per-schema validator cache hit + on subsequent ProcessorMixin._merge_kwargs calls.""" + cache: dict = {} + first = TypedDict("merged_typed_dict", {"do_rescale": Optional[bool]}, total=False) + second = TypedDict("merged_typed_dict", {"do_rescale": Optional[bool]}, total=False) + + cached_first = _get_cached_merged_typed_dict(first, cache) + cached_second = _get_cached_merged_typed_dict(second, cache) + + assert cached_first is cached_second # identity stable across equivalent shapes + + +def test_non_merged_typed_dict_schemas_pass_through(): + """Schemas whose `__name__` is not the ephemeral `merged_typed_dict` are + returned unchanged — only HF's per-call class is the one we want to collapse.""" + other = TypedDict("other_typed_dict", {"do_rescale": Optional[bool]}, total=False) + assert _get_cached_merged_typed_dict(other, {}) is other diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl_preprocess.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl_preprocess.py index 8803f3fd3a4b..df0a4815bd68 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl_preprocess.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl_preprocess.py @@ -12,32 +12,31 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for Qwen3VLInputProcessorBase._preprocess kwargs handling. - -Covers the per-request ``mm_processor_kwargs`` plumbing added to Qwen3-VL: - * Video metadata from the IO loader is forwarded to the HF processor so - sample_frames sees the real source fps (rather than HF's 24 fps default). - * ``num_frames`` and ``fps`` are mutually exclusive in HF's sample_frames; if - the caller sets ``num_frames`` without ``fps``, ``_preprocess`` must - explicitly null ``fps`` so the processor's class-level ``fps=2`` default - does not interfere. - * The caller-supplied ``mm_processor_kwargs`` dict must not be mutated. - -The tests bind ``_preprocess`` to a stand-in object so they avoid the heavy -``__init__`` (which would download an HF processor) and only exercise the new -control-flow. +"""Unit tests for Qwen3-VL preprocess control flow. + +Covers two pieces of logic that decide what reaches HF's video processor: + + 1. ``_decide_do_sample_frames`` — the decision tree that picks the single + ``do_sample_frames`` flag for the whole batch. + 2. ``_preprocess`` — metadata rewriting and kwargs threading around the + processor call. + +The tests bind ``_preprocess`` to a stand-in object so they exercise the +control flow without constructing a real HF processor. """ from types import SimpleNamespace from unittest.mock import MagicMock -from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLInputProcessorBase +from tensorrt_llm._torch.models.modeling_qwen3vl import ( + Qwen3VLInputProcessorBase, + _decide_do_sample_frames, +) -def _fake_video(metadata): - """Minimal stand-in for VideoData: just needs ``.frames`` and ``.metadata``.""" - # A single non-Tensor frame so the do_rescale branch stays in its default. - return SimpleNamespace(frames=[object()], metadata=metadata) +def _fake_video(metadata, *, n_frames=8): + """Stand-in for VideoData: only `.frames` (count) and `.metadata` are read.""" + return SimpleNamespace(frames=[object()] * n_frames, metadata=metadata) def _call_preprocess(mm_data, mm_processor_kwargs): @@ -52,98 +51,103 @@ def _call_preprocess(mm_data, mm_processor_kwargs): return fake_processor -class TestVideoMetadataForwarding: - """``video_metadata`` is forwarded only when the caller passes - ``mm_processor_kwargs``.""" - - def test_metadata_forwarded_when_opted_in(self): - """Opted-in path: list built per-video, in order.""" - md0 = {"total_num_frames": 64, "fps": 25.0, "duration": 2.5} - md1 = {"total_num_frames": 32, "fps": 30.0, "duration": 1.0} - mm_data = {"video": [_fake_video(md0), _fake_video(md1)]} - - processor = _call_preprocess(mm_data, {"fps": 2.0}) - - kwargs = processor.call_args.kwargs - assert kwargs["video_metadata"] == [md0, md1] - assert len(kwargs["videos"]) == 2 - - def test_metadata_not_forwarded_with_empty_kwargs(self): - """Empty kwargs: videos present but metadata is not forwarded.""" - mm_data = {"video": [_fake_video({"fps": 30.0})]} - - processor = _call_preprocess(mm_data, {}) - - assert processor.call_args.kwargs["video_metadata"] is None +class TestDecideDoSampleFrames: + """Decision tree for whether HF should run ``sample_frames``.""" + + def test_explicit_true_wins_even_when_counts_match(self): + # Even with a matching target, caller's explicit True is honored. + vd = _fake_video({"duration": 4.0}, n_frames=8) + assert _decide_do_sample_frames([vd], {"do_sample_frames": True}) is True + + def test_explicit_false_wins_even_when_counts_differ(self): + # Even when num_frames would force sampling, caller's False is honored. + vd = _fake_video({"duration": 4.0}, n_frames=8) + assert ( + _decide_do_sample_frames([vd], {"do_sample_frames": False, "num_frames": 16}) is False + ) + + def test_num_frames_matches_decoded_no_sampling(self): + vd = _fake_video({}, n_frames=8) + assert _decide_do_sample_frames([vd], {"num_frames": 8}) is False + + def test_num_frames_differs_triggers_sampling(self): + vd = _fake_video({}, n_frames=8) + assert _decide_do_sample_frames([vd], {"num_frames": 4}) is True + + def test_fps_target_computed_from_duration(self): + # Target = floor(duration * fps) = floor(4.0 * 2.0) = 8 → no resample. + vd_match = _fake_video({"duration": 4.0}, n_frames=8) + assert _decide_do_sample_frames([vd_match], {"fps": 2.0}) is False + # Target = floor(4.0 * 4.0) = 16 ≠ 8 → resample. + vd_diff = _fake_video({"duration": 4.0}, n_frames=8) + assert _decide_do_sample_frames([vd_diff], {"fps": 4.0}) is True + + def test_silent_caller_io_loaded_all_triggers_fallback_sampling(self): + # No num_frames/fps from caller, IO loaded every source frame + # (decoded count == total_num_frames): defer to HF's class-default + # sampling. + vd = _fake_video({"total_num_frames": 240}, n_frames=240) + assert _decide_do_sample_frames([vd], {}) is True + + def test_silent_caller_io_subsampled_triggers_sampling(self): + # No num_frames/fps from caller: defer to HF's class-default sampling + # even when IO already subsampled (decoded count < total_num_frames). + vd = _fake_video({"total_num_frames": 240}, n_frames=8) + assert _decide_do_sample_frames([vd], {}) is True + + def test_batch_reduction_is_any_video_needs_sampling(self): + # One video matches, one needs resampling → batch is sampled. + match = _fake_video({}, n_frames=8) + differ = _fake_video({}, n_frames=16) + assert _decide_do_sample_frames([match, differ], {"num_frames": 8}) is True + + +class TestPreprocessMetadataForwarding: + """`_preprocess` forwards per-video metadata with controlled rewrites.""" + + def test_total_num_frames_rewritten_to_decoded_count(self): + # Source clip had 240 frames; IO subsampled to 8. + # `total_num_frames` must be rewritten so HF's sample_frames indices + # stay in range of the 8-frame tensor we hand it. + vd = _fake_video({"total_num_frames": 240, "fps": 24.0, "duration": 10.0}, n_frames=8) + processor = _call_preprocess({"video": [vd]}, {"num_frames": 8}) + m = processor.call_args.kwargs["video_metadata"][0] + assert m["total_num_frames"] == 8 + assert m["fps"] == 24.0 # other fields unchanged def test_metadata_is_none_when_no_videos(self): - """No videos: metadata is None regardless of kwargs.""" processor = _call_preprocess({}, {"fps": 2.0}) - kwargs = processor.call_args.kwargs - assert kwargs["video_metadata"] is None - assert kwargs["videos"] is None - - -class TestNumFramesFpsMutualExclusivity: - """Truth table over ("num_frames" present, "fps" present). - - HF's Qwen3VLProcessor.sample_frames treats num_frames and fps as mutually - exclusive; the class-level default ``fps=2`` would otherwise win over the - caller's ``num_frames``. The fix injects ``fps=None`` only in the - (True, False) corner. - """ - - def test_num_frames_without_fps_injects_null_fps(self): - """(True, False): the corner the fix exists for.""" - mm_data = {"video": [_fake_video({"fps": 30.0})]} - - processor = _call_preprocess(mm_data, {"num_frames": 8}) - - kwargs = processor.call_args.kwargs - assert kwargs["num_frames"] == 8 - assert "fps" in kwargs - assert kwargs["fps"] is None + assert processor.call_args.kwargs["video_metadata"] is None - def test_explicit_fps_is_preserved(self): - """(True, True): the fix must not trample a user-supplied fps.""" - mm_data = {"video": [_fake_video({"fps": 30.0})]} - processor = _call_preprocess(mm_data, {"num_frames": 8, "fps": 4.0}) +class TestPreprocessKwargsForwarding: + """`do_sample_frames` is always passed; sampling kwargs are conditional.""" + def test_sampling_kwargs_forwarded_when_sampling(self): + # num_frames mismatches frames count → sampling fires → num_frames forwarded. + vd = _fake_video({}, n_frames=8) + processor = _call_preprocess({"video": [vd]}, {"num_frames": 4}) kwargs = processor.call_args.kwargs - assert kwargs["num_frames"] == 8 - assert kwargs["fps"] == 4.0 - - def test_fps_only_passes_through_unchanged(self): - """(False, True): the condition must not fire for fps alone.""" - processor = _call_preprocess({}, {"fps": 4.0}) - - kwargs = processor.call_args.kwargs - assert kwargs["fps"] == 4.0 - assert "num_frames" not in kwargs - - def test_no_kwargs_means_no_num_frames_no_fps(self): - """(False, False): the empty-kwargs case adds no spurious keys.""" - processor = _call_preprocess({}, {}) - + assert kwargs["do_sample_frames"] is True + assert kwargs["num_frames"] == 4 + + def test_sampling_kwargs_dropped_when_not_sampling(self): + # num_frames matches frames count → no sampling → num_frames NOT forwarded + # (HF's class-default ``fps=2`` cannot collide because do_sample_frames=False). + vd = _fake_video({}, n_frames=8) + processor = _call_preprocess({"video": [vd]}, {"num_frames": 8}) kwargs = processor.call_args.kwargs + assert kwargs["do_sample_frames"] is False assert "num_frames" not in kwargs assert "fps" not in kwargs - -class TestInputDictNotMutated: - """The caller's mm_processor_kwargs dict is copied, not mutated.""" - - def test_num_frames_only_dict_not_mutated(self): - original = {"num_frames": 8} + def test_caller_dict_not_mutated(self): + original = {"num_frames": 8, "max_pixels": 1234} snapshot = dict(original) _call_preprocess({}, original) assert original == snapshot - -class TestUnrelatedKwargsForwarded: - """Keys unrelated to the num_frames/fps fix pass through verbatim.""" - - def test_unrelated_kwargs_reach_processor(self): - processor = _call_preprocess({}, {"num_frames": 8, "max_pixels": 1234}) + def test_unrelated_kwargs_pass_through(self): + # Resize/normalize knobs flow through unchanged. + processor = _call_preprocess({}, {"max_pixels": 1234}) assert processor.call_args.kwargs["max_pixels"] == 1234 diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index 09affe820303..d536fcdeb319 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -303,6 +303,24 @@ commands: is_flag: false flags: - "--num_postprocess_workers" + input_processor_workers: + type: int + default: 8 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--input-processor-workers" + media_load_workers: + type: int + default: 8 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--media-load-workers" otlp_traces_endpoint: type: str default: null