diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 393c144bd6f9..c43d0f742ea2 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -4,6 +4,8 @@ """Translate an OpenEngine GenerateRequest into TensorRT-LLM request inputs.""" import asyncio +import importlib +import threading from collections import Counter from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass @@ -13,7 +15,12 @@ import grpc from tensorrt_llm.conversation_params import ConversationParams -from tensorrt_llm.inputs.media_io import ImageMediaIO, MediaFetchError +from tensorrt_llm.inputs.media_io import ( + ImageMediaIO, + MediaFetchError, + VideoMediaIO, + _decoded_video_bytes, +) from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams @@ -25,11 +32,12 @@ # output index and a scan of every generated token, so an unbounded list lets a # single request tax the event loop for the whole stream. _MAX_STOP_CONDITIONS = 64 -# Bound per-request decode work and retained image data. +# Bound per-request decode work and retained decoded media. _MAX_MEDIA_ITEMS = 64 _MAX_CONCURRENT_MEDIA_LOADS = 8 -# Checked after each decode; concurrency bounds peak overshoot. Encoded inputs -# have separate gRPC and remote-loader caps. +# Images are charged after decoding. Video decoders reserve retained bytes and +# counted peak buffers separately before allocating frames. +# Encoded inputs have separate gRPC and remote-loader caps. _MAX_TOTAL_DECODED_MEDIA_BYTES = 512 * 1024 * 1024 _CONVERSATION_ID_EXTRA_KEY = "conversation_id" _DETOKENIZE_EXTRA_KEY = "detokenize" @@ -39,12 +47,123 @@ _MediaValidator = Callable[[generation_pb2.MediaItem, int], None] +class _DecodedMediaReservation: + def __init__(self, budget: "_DecodedMediaBudget", peak_bytes: int, retained_bytes: int) -> None: + self._budget = budget + self._peak_bytes = peak_bytes + self._retained_bytes = retained_bytes + self._active = True + + def check_cancelled(self) -> None: + with self._budget._condition: + if self._budget._cancelled: + raise ValueError("decoded media loading was cancelled") + + def shrink(self, retained_bytes: int) -> None: + if not self._active or not 0 <= retained_bytes <= self._retained_bytes: + raise ValueError("invalid decoded-media reservation shrink") + self._budget._shrink(self._peak_bytes, self._retained_bytes, retained_bytes) + self._peak_bytes = self._retained_bytes = retained_bytes + + def commit(self, retained_bytes: int) -> None: + if not self._active: + raise RuntimeError("decoded-media reservation is no longer active") + try: + self._budget._settle(self._peak_bytes, self._retained_bytes, retained_bytes) + finally: + self._active = False + + def release(self) -> None: + if self._active: + self._budget._settle(self._peak_bytes, self._retained_bytes, 0) + self._active = False + + +class _DecodedMediaBudget: + """Thread-safe shared budget for all decoded media in one request.""" + + def __init__(self, limit: int, peak_limit: int | None = None) -> None: + self._limit = limit + # Conversion buffers have a separate allowance from retained output. + self._peak_limit = peak_limit if peak_limit is not None else 2 * limit + self._retained = 0 + self._reserved_retained = 0 + self._used = 0 + self._condition = threading.Condition() + self._cancelled = False + + def _limit_error(self) -> ValueError: + return ValueError(f"decoded media exceeds the per-request limit of {self._limit} bytes") + + def charge(self, retained_bytes: int) -> None: + if retained_bytes < 0: + raise ValueError("decoded media size cannot be negative") + with self._condition: + if self._retained + self._reserved_retained + retained_bytes > self._limit: + raise self._limit_error() + self._retained += retained_bytes + self._used += retained_bytes + + def reserve(self, peak_bytes: int, retained_bytes: int) -> _DecodedMediaReservation: + if not 0 <= retained_bytes <= peak_bytes: + raise ValueError("decoded media sizes must satisfy 0 <= retained <= peak") + with self._condition: + while True: + if self._cancelled: + raise ValueError("decoded media loading was cancelled") + if self._retained + self._reserved_retained + retained_bytes > self._limit: + raise self._limit_error() + if self._retained + peak_bytes > self._peak_limit: + raise ValueError( + f"decoded media exceeds the peak buffer limit of {self._peak_limit} bytes" + ) + if self._used + peak_bytes <= self._peak_limit: + break + self._condition.wait() + self._used += peak_bytes + self._reserved_retained += retained_bytes + return _DecodedMediaReservation(self, peak_bytes, retained_bytes) + + def _shrink(self, peak_bytes: int, reserved_retained: int, retained_bytes: int) -> None: + with self._condition: + self._used -= peak_bytes - retained_bytes + self._reserved_retained -= reserved_retained - retained_bytes + self._condition.notify_all() + + def _settle(self, peak_bytes: int, reserved_retained: int, retained_bytes: int) -> None: + with self._condition: + self._used -= peak_bytes + self._reserved_retained -= reserved_retained + try: + if retained_bytes < 0: + raise ValueError("decoded media size cannot be negative") + if self._retained + self._reserved_retained + retained_bytes > self._limit: + raise self._limit_error() + if self._used + retained_bytes > self._peak_limit: + raise self._limit_error() + self._used += retained_bytes + self._retained += retained_bytes + finally: + self._condition.notify_all() + + def cancel(self) -> None: + with self._condition: + self._cancelled = True + self._condition.notify_all() + + +def _always_available() -> bool: + return True + + @dataclass(frozen=True) class _MediaHandler: name: str validate: _MediaValidator - loader_factory: Callable[[Any], _MediaLoader] + loader_factory: Callable[[Any, _DecodedMediaBudget], _MediaLoader] decoded_bytes: Callable[[Any], int] + runtime_available: Callable[[], bool] = _always_available + reserves_decoded_bytes: bool = False @dataclass(frozen=True) @@ -286,16 +405,21 @@ def _base_input_from_request( raise ValueError("exactly one of prompt or token_ids must be set") -def _image_io_for_model(llm: Any) -> ImageMediaIO: +def _media_io_for_model( + llm: Any, + modality: str, + media_io_type: type[Any], + runtime_kwargs: dict[str, Any] | None = None, +) -> Any: input_processor = getattr(llm, "input_processor", None) defaults = None if isinstance(input_processor, BaseMultimodalInputProcessor): preferences = input_processor.get_preferred_media_io_kwargs() or {} - defaults = preferences.get("image") - return ImageMediaIO.create(defaults, None) + defaults = preferences.get(modality) + return media_io_type.create(defaults, runtime_kwargs) -def _validate_image_item(item: generation_pb2.MediaItem, index: int) -> None: +def _validate_media_item(item: generation_pb2.MediaItem, index: int, modality: str) -> None: source = item.WhichOneof("source") if source is None: raise ValueError(f"media[{index}] must set exactly one source") @@ -320,28 +444,38 @@ def _validate_image_item(item: generation_pb2.MediaItem, index: int) -> None: 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") + if media_type and not media_type.lower().startswith(f"{modality}/"): + raise ValueError(f"media[{index}].data_uri must declare an {modality} 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") + if item.mime_type and not item.mime_type.lower().startswith(f"{modality}/"): + raise ValueError( + f"media[{index}].mime_type must start with {modality}/ for raw {modality} bytes" + ) 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 +def _validate_image_item(item: generation_pb2.MediaItem, index: int) -> None: + _validate_media_item(item, index, "image") + + +def _validate_video_item(item: generation_pb2.MediaItem, index: int) -> None: + _validate_media_item(item, index, "video") + + +async def _load_media_item( + item: generation_pb2.MediaItem, media_io: Any, index: int, modality: str ) -> Any: source = item.WhichOneof("source") if source == "raw_bytes": try: - return await image_io.async_load_bytes(bytes(item.raw_bytes)) + return await media_io.async_load_bytes(bytes(item.raw_bytes)) except Exception as error: raise ValueError(f"media[{index}].raw_bytes could not be decoded: {error}") from error value = item.url if source == "url" else item.data_uri try: - return await image_io.async_load(value) + return await media_io.async_load(value) except MediaFetchError as error: raise MediaUnavailableError( f"media[{index}].{source} is temporarily unavailable: {error}" @@ -350,11 +484,25 @@ async def _load_image_item( 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) +def _image_loader_factory(llm: Any, _budget: _DecodedMediaBudget) -> _MediaLoader: + image_io = _media_io_for_model(llm, "image", ImageMediaIO) + + async def load(item: generation_pb2.MediaItem, index: int) -> Any: + return await _load_media_item(item, image_io, index, "image") + + return load + + +def _video_loader_factory(llm: Any, budget: _DecodedMediaBudget) -> _MediaLoader: + video_io = _media_io_for_model( + llm, + "video", + VideoMediaIO, + {"decoded_memory_reserver": budget.reserve}, + ) async def load(item: generation_pb2.MediaItem, index: int) -> Any: - return await _load_image_item(item, image_io, index) + return await _load_media_item(item, video_io, index, "video") return load @@ -373,15 +521,34 @@ def _decoded_image_bytes(image: Any) -> int: raise TypeError(f"unsupported decoded image type: {type(image).__name__}") +def _video_runtime_available() -> bool: + try: + importlib.import_module("cv2") + except (ImportError, OSError): + return False + return True + + _MEDIA_HANDLERS: dict[int, _MediaHandler] = { generation_pb2.MODALITY_IMAGE: _MediaHandler( "image", _validate_image_item, _image_loader_factory, _decoded_image_bytes ), + generation_pb2.MODALITY_VIDEO: _MediaHandler( + "video", + _validate_video_item, + _video_loader_factory, + _decoded_video_bytes, + _video_runtime_available, + reserves_decoded_bytes=True, + ), } def media_capabilities_for_model(llm: Any) -> MediaCapabilities: - return resolve_media_capabilities(llm, {handler.name for handler in _MEDIA_HANDLERS.values()}) + return resolve_media_capabilities( + llm, + {handler.name for handler in _MEDIA_HANDLERS.values() if handler.runtime_available()}, + ) def _modality_name(modality: int) -> str: @@ -515,25 +682,20 @@ async def _load_prepared_input(prepared: _PreparedInput, llm: Any) -> str | dict if capabilities is None: raise RuntimeError("media input plan is missing capabilities") + budget = _DecodedMediaBudget(_MAX_TOTAL_DECODED_MEDIA_BYTES) loaders: dict[str, _MediaLoader] = {} for _, handler in prepared.media: if handler.name not in loaders: - loaders[handler.name] = handler.loader_factory(llm) + loaders[handler.name] = handler.loader_factory(llm, budget) semaphore = asyncio.Semaphore(_MAX_CONCURRENT_MEDIA_LOADS) - decoded_bytes = 0 async def load_bounded( item: generation_pb2.MediaItem, index: int, handler: _MediaHandler ) -> Any: - nonlocal decoded_bytes async with semaphore: decoded = await loaders[handler.name](item, index) - decoded_bytes += handler.decoded_bytes(decoded) - if decoded_bytes > _MAX_TOTAL_DECODED_MEDIA_BYTES: - raise ValueError( - "decoded media exceeds the per-request limit of " - f"{_MAX_TOTAL_DECODED_MEDIA_BYTES} bytes" - ) + if not handler.reserves_decoded_bytes: + budget.charge(handler.decoded_bytes(decoded)) return decoded tasks = [ @@ -544,6 +706,7 @@ async def load_bounded( decoded_media = await asyncio.gather(*tasks) except (Exception, asyncio.CancelledError): # gather leaves siblings running after the first failure. + budget.cancel() for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 4c58f2d4fc16..0e12b267b3ee 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -21,12 +21,14 @@ from types import MappingProxyType from typing import ( Any, + Callable, ClassVar, Dict, Generic, List, Mapping, Optional, + Protocol, Tuple, Type, TypeVar, @@ -51,6 +53,51 @@ class MediaFetchError(RuntimeError): """A transient failure while fetching remote media.""" +class _DecodedMemoryReservation(Protocol): + def check_cancelled(self) -> None: ... + + def shrink(self, retained_bytes: int) -> None: ... + + def commit(self, retained_bytes: int) -> None: ... + + def release(self) -> None: ... + + +_DecodedMemoryReserver = Callable[[int, int], _DecodedMemoryReservation] + + +def _decoded_video_bytes(video: Any) -> int: + """Return bytes retained by decoded video frames and optional audio.""" + if not isinstance(video, VideoData): + raise TypeError(f"unsupported decoded video type: {type(video).__name__}") + + if isinstance(video.frames, np.ndarray): + total = video.frames.nbytes + elif isinstance(video.frames, list): + total = 0 + for frame in video.frames: + if isinstance(frame, torch.Tensor): + total += frame.numel() * frame.element_size() + elif isinstance(frame, np.ndarray): + total += frame.nbytes + elif isinstance(frame, Image.Image): + total += frame.width * frame.height * len(frame.getbands()) + else: + raise TypeError(f"unsupported decoded video frame type: {type(frame).__name__}") + else: + raise TypeError(f"unsupported decoded video frames type: {type(video.frames).__name__}") + + if video.audio is not None: + samples = video.audio.samples + if isinstance(samples, torch.Tensor): + total += samples.numel() * samples.element_size() + elif isinstance(samples, np.ndarray): + total += samples.nbytes + else: + raise TypeError(f"unsupported decoded audio type: {type(samples).__name__}") + return total + + def rgba_to_rgb( image: Image.Image, background_color: Union[tuple[int, int, int], list[int]] = (255, 255, 255) ) -> Image.Image: @@ -414,12 +461,12 @@ def _audio_frame_to_array(frame, mono: bool) -> np.ndarray: return chunk -# TODO(TRTLLM-12001): Add unit tests for this. def extract_audio_from_video( source: Union[str, BytesIO], *, sr: Optional[float] = None, mono: bool = True, + decoded_memory_reserver: Optional[_DecodedMemoryReserver] = None, ) -> Tuple[np.ndarray, int]: """Extract the audio track from a video file using PyAV. @@ -427,6 +474,8 @@ def extract_audio_from_video( source: File path, URL, or BytesIO containing the video. sr: Target sample rate. If `None`, the native sample rate is kept. mono: If `True` (default), average channels to produce a mono waveform. + decoded_memory_reserver: Shared budget for retained chunks and the final + concatenation buffer. Codec and resampler working memory is excluded. Returns: `(waveform, sample_rate)` where *waveform* is a 1-D float32 @@ -445,51 +494,85 @@ def extract_audio_from_video( except ImportError: raise ImportError("PyAV is required for audio extraction from video but is not installed.") + chunks: List[np.ndarray] = [] + reservations: List[Tuple[_DecodedMemoryReservation, int]] = [] + concatenation_reservation: Optional[_DecodedMemoryReservation] = None + + def append_frame(frame) -> None: + if decoded_memory_reserver is not None: + channels = 1 if mono else len(frame.layout.channels) + retained_bytes = frame.samples * channels * np.dtype(np.float32).itemsize + reservation = decoded_memory_reserver(retained_bytes, retained_bytes) + reservations.append((reservation, retained_bytes)) + chunk = _audio_frame_to_array(frame, mono) + if decoded_memory_reserver is not None and ( + chunk.dtype != np.float32 or chunk.nbytes != retained_bytes + ): + raise ValueError("Decoded audio chunk does not match its float32 reservation") + chunks.append(chunk) + try: - with av.open(source) as container: - if not container.streams.audio: - raise ValueError("No audio stream found in the video.") - stream = container.streams.audio[0] - stream.thread_type = "AUTO" - native_sr = stream.rate - target_sr = int(sr) if sr is not None else native_sr - - chunks: List[np.ndarray] = [] - target_layout = "mono" if mono else stream.layout.name - needs_resampling = ( - not math.isclose(float(target_sr), float(native_sr), rel_tol=0.0, abs_tol=1e-6) - or stream.format.name != "fltp" - or stream.layout.name != target_layout - ) - resampler = ( - av.AudioResampler( - format="fltp", - layout=target_layout, - rate=target_sr, + try: + with av.open(source) as container: + if not container.streams.audio: + raise ValueError("No audio stream found in the video.") + stream = container.streams.audio[0] + stream.thread_type = "AUTO" + native_sr = stream.rate + target_sr = int(sr) if sr is not None else native_sr + + target_layout = "mono" if mono else stream.layout.name + needs_resampling = ( + not math.isclose(float(target_sr), float(native_sr), rel_tol=0.0, abs_tol=1e-6) + or stream.format.name != "fltp" + or stream.layout.name != target_layout ) - if needs_resampling - else None - ) - for frame in container.decode(stream): - if needs_resampling: - for out_frame in resampler.resample(frame): - chunks.append(_audio_frame_to_array(out_frame, mono)) - else: - chunks.append(_audio_frame_to_array(frame, mono)) - except ValueError: - raise - except Exception as e: - raise ValueError( - "Invalid or corrupted video data when extracting audio. " - "Ensure the input is a valid video file." - ) from e - - if not chunks: - raise ValueError("No audio frames decoded from the video.") + resampler = ( + av.AudioResampler( + format="fltp", + layout=target_layout, + rate=target_sr, + ) + if needs_resampling + else None + ) + for frame in container.decode(stream): + if needs_resampling: + for out_frame in resampler.resample(frame): + append_frame(out_frame) + else: + append_frame(frame) + except ValueError: + raise + except Exception as e: + raise ValueError( + "Invalid or corrupted video data when extracting audio. " + "Ensure the input is a valid video file." + ) from e - audio = np.concatenate(chunks, axis=-1).astype(np.float32, copy=False) + if not chunks: + raise ValueError("No audio frames decoded from the video.") - return audio, target_sr + if decoded_memory_reserver is not None: + concatenation_reservation = decoded_memory_reserver( + sum(chunk.nbytes for chunk in chunks), 0 + ) + audio = np.concatenate(chunks, axis=-1).astype(np.float32, copy=False) + chunks.clear() + if concatenation_reservation is not None: + concatenation_reservation.release() + concatenation_reservation = None + # The waveform replaces the chunks without changing retained-byte ownership. + for reservation, retained_bytes in reservations: + reservation.commit(retained_bytes) + + return audio, target_sr + finally: + chunks.clear() + if concatenation_reservation is not None: + concatenation_reservation.release() + for reservation, _ in reservations: + reservation.release() def _get_cv2(): @@ -733,6 +816,8 @@ def _load_video_by_cv2( extract_audio: bool = False, cv2_backend: Optional[int] = None, raw_bytes_hash: Optional[str] = None, + max_decoded_frame_bytes: Optional[int] = None, + decoded_memory_reserver: Optional[_DecodedMemoryReserver] = None, ) -> VideoData: """Decode a video and return sampled frames as a list. @@ -750,6 +835,13 @@ def _load_video_by_cv2( contiguous 4D buffer that HF video processors pass through without an extra frame-by-frame copy. `"pil"` - list[PIL.Image], one per sampled frame. + + When `max_decoded_frame_bytes` or `decoded_memory_reserver` is set, the + counted peak frame buffers are sized from the first decoded frame and + checked before the frame stack is allocated. This covers buffers allocated + here, not codec state, compressed input, Python overhead, or the process's + complete memory peak. The reserver receives peak and retained byte estimates + and may wait for sibling conversions to release their temporary buffers. """ assert format in ("pt", "np", "pil"), "format must be one of 'pt', 'np', 'pil'" @@ -773,6 +865,7 @@ def _load_video_by_cv2( else: vidcap = cv2.VideoCapture(video) + decoded_memory_reservation: Optional[_DecodedMemoryReservation] = None try: if not vidcap.isOpened(): src_repr = ( @@ -814,15 +907,54 @@ def _load_video_by_cv2( # affected (e.g. every frame retrieves with a drifted shape). skip_warned = False frame_idx = 0 - while frame_idx <= max_idx and vidcap.grab(): + while frame_idx <= max_idx: + if decoded_memory_reservation is not None: + decoded_memory_reservation.check_cancelled() + if not vidcap.grab(): + break + if decoded_memory_reservation is not None: + decoded_memory_reservation.check_cancelled() if frame_idx in target_set: # Reuse a single BGR buffer across retrieves; cv2 replaces its # contents in place when the argument is shape-compatible. ok, bgr_scratch = vidcap.retrieve(bgr_scratch) + if decoded_memory_reservation is not None: + decoded_memory_reservation.check_cancelled() if ok: fh, fw = bgr_scratch.shape[:2] if stacked_rgb is None: H, W = fh, fw + frame_values = num_frames_to_sample * H * W * 3 + retained_frame_bytes = frame_values * (4 if format == "pt" else 1) + if format == "pt": + # uint8 RGB stack + float32 NumPy normalization + # buffer + contiguous float32 CHW tensor. A + # non-CPU destination briefly adds another tensor. + peak_item_size = ( + np.dtype(np.uint8).itemsize + 2 * np.dtype(np.float32).itemsize + ) + peak_item_size += ( + np.dtype(np.float32).itemsize + if torch.device(device).type != "cpu" + else 0 + ) + else: + # RGB stack plus PIL copies or a compact truncated ndarray. + peak_item_size = 2 * np.dtype(np.uint8).itemsize + peak_decoded_frame_bytes = frame_values * peak_item_size + if ( + max_decoded_frame_bytes is not None + and peak_decoded_frame_bytes > max_decoded_frame_bytes + ): + raise ValueError( + "Decoded video frame buffers require " + f"{peak_decoded_frame_bytes} bytes, " + f"exceeding the limit of {max_decoded_frame_bytes} bytes." + ) + if decoded_memory_reserver is not None: + decoded_memory_reservation = decoded_memory_reserver( + peak_decoded_frame_bytes, retained_frame_bytes + ) stacked_rgb = np.empty((num_frames_to_sample, H, W, 3), dtype=np.uint8) if (fh, fw) == (H, W): cv2.cvtColor( @@ -843,11 +975,16 @@ def _load_video_by_cv2( ) skip_warned = True frame_idx += 1 + if decoded_memory_reservation is not None: + decoded_memory_reservation.check_cancelled() vidcap.release() if stacked_rgb is None or not valid_indices: raise ValueError("Video has no readable frames.") - stacked_rgb = stacked_rgb[: len(valid_indices)] + if format == "np" and len(valid_indices) < num_frames_to_sample: + stacked_rgb = stacked_rgb[: len(valid_indices)].copy() + else: + stacked_rgb = stacked_rgb[: len(valid_indices)] if format == "pt": stacked_f32 = stacked_rgb.astype(np.float32) @@ -856,10 +993,14 @@ def _load_video_by_cv2( if device != "cpu": tensor_nchw = tensor_nchw.to(device) loaded_frames = list(torch.unbind(tensor_nchw, dim=0)) + del tensor_nchw + del stacked_f32 + del stacked_rgb elif format == "np": loaded_frames = stacked_rgb else: # "pil" loaded_frames = [Image.fromarray(frame) for frame in stacked_rgb] + del stacked_rgb metadata = { "total_num_frames": frame_count, @@ -867,36 +1008,54 @@ def _load_video_by_cv2( "duration": duration, "frames_indices": valid_indices, } + except BaseException: + if decoded_memory_reservation is not None: + decoded_memory_reservation.release() + raise finally: # Release the OpenCV handle before any downstream re-open (e.g. PyAV # audio extraction on Windows) and to avoid leaking descriptors. vidcap.release() - audio = None - if extract_audio: - # extract_audio_from_video accepts Union[str, BytesIO]. When the - # in-memory cv2 path was taken, `video` is still the original - # bytes parameter — wrap in a fresh BytesIO (the one consumed by - # cv2 is at end-of-stream). When the tempfile fallback was used, - # `video` is the file path string and is passed through as-is. - audio_source = ( - BytesIO(bytes(video)) if isinstance(video, (bytes, bytearray, memoryview)) else video - ) - try: - audio_samples, audio_sample_rate = extract_audio_from_video(audio_source) - audio = AudioData(samples=audio_samples, sample_rate=audio_sample_rate) - except ValueError as e: - if "No audio stream found" in str(e): - logger.warning("Video has no audio track, skipping audio extraction.") - else: - raise - - return VideoData( + video_data = VideoData( frames=loaded_frames, metadata=metadata, - audio=audio, raw_bytes_hash=raw_bytes_hash, ) + try: + frame_bytes = ( + _decoded_video_bytes(video_data) if decoded_memory_reservation is not None else 0 + ) + if extract_audio: + if decoded_memory_reservation is not None: + decoded_memory_reservation.shrink(frame_bytes) + # extract_audio_from_video accepts Union[str, BytesIO]. When the + # in-memory cv2 path was taken, `video` is still the original + # bytes parameter — wrap in a fresh BytesIO (the one consumed by + # cv2 is at end-of-stream). When the tempfile fallback was used, + # `video` is the file path string and is passed through as-is. + audio_source = ( + BytesIO(bytes(video)) + if isinstance(video, (bytes, bytearray, memoryview)) + else video + ) + try: + audio_samples, audio_sample_rate = extract_audio_from_video( + audio_source, decoded_memory_reserver=decoded_memory_reserver + ) + video_data.audio = AudioData(samples=audio_samples, sample_rate=audio_sample_rate) + except ValueError as e: + if "No audio stream found" in str(e): + logger.warning("Video has no audio track, skipping audio extraction.") + else: + raise + if decoded_memory_reservation is not None: + decoded_memory_reservation.commit(frame_bytes) + except BaseException: + if decoded_memory_reservation is not None: + decoded_memory_reservation.release() + raise + return video_data def _normalize_file_uri(uri: str) -> str: @@ -1194,17 +1353,23 @@ def __init__( format: str = "pt", device: str = "cpu", extract_audio: bool = False, + max_decoded_frame_bytes: Optional[int] = None, + decoded_memory_reserver: Optional[_DecodedMemoryReserver] = None, ) -> None: # `"pt"` is the safe default every video model handles; models tuned # for the uint8-HWC path opt into `"np"` via # InputProcessor.get_preferred_media_io_kwargs(). if format not in _SUPPORTED_VIDEO_FORMATS: raise ValueError(f"format must be one of {_SUPPORTED_VIDEO_FORMATS}, got {format!r}") + if max_decoded_frame_bytes is not None and max_decoded_frame_bytes <= 0: + raise ValueError("max_decoded_frame_bytes must be positive") self._num_frames = num_frames self._fps = fps self._format = format self._device = device self._extract_audio = extract_audio + self._max_decoded_frame_bytes = max_decoded_frame_bytes + self._decoded_memory_reserver = decoded_memory_reserver @classmethod def merge_kwargs( @@ -1244,6 +1409,8 @@ def load_bytes(self, data: bytes) -> VideoData: extract_audio=self._extract_audio, cv2_backend=cv2_backend, raw_bytes_hash=raw_bytes_hash, + max_decoded_frame_bytes=self._max_decoded_frame_bytes, + decoded_memory_reserver=self._decoded_memory_reserver, ) with tempfile.NamedTemporaryFile(suffix=".mp4", dir=_VIDEO_TEMPFILE_DIR) as tmp: tmp.write(data) @@ -1256,6 +1423,8 @@ def load_bytes(self, data: bytes) -> VideoData: self._device, extract_audio=self._extract_audio, raw_bytes_hash=raw_bytes_hash, + max_decoded_frame_bytes=self._max_decoded_frame_bytes, + decoded_memory_reserver=self._decoded_memory_reserver, ) def load_base64(self, media_type: str, data: str) -> VideoData: diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 0abaae084841..bc2cb9865266 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -1,15 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for OpenEngine image input mapping and discovery.""" +"""Unit tests for OpenEngine multimodal input mapping and discovery.""" import asyncio import base64 import gc import re +import sys import threading from collections.abc import AsyncIterator, Sequence from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext from contextvars import ContextVar from pathlib import Path from types import SimpleNamespace @@ -33,8 +35,10 @@ 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 import media_io # 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.multimodal_data import AudioData, VideoData # noqa: E402 from tensorrt_llm.inputs.registry import ( # noqa: E402 MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, @@ -54,6 +58,24 @@ def _ignore_synthetic_decoded_size(monkeypatch) -> None: handler.validate, handler.loader_factory, lambda _decoded: 0, + handler.runtime_available, + handler.reserves_decoded_bytes, + ), + ) + + +def _set_video_runtime_available(monkeypatch, available: bool) -> None: + handler = request_mapping._MEDIA_HANDLERS[generation_pb2.MODALITY_VIDEO] + monkeypatch.setitem( + request_mapping._MEDIA_HANDLERS, + generation_pb2.MODALITY_VIDEO, + request_mapping._MediaHandler( + handler.name, + handler.validate, + handler.loader_factory, + handler.decoded_bytes, + lambda: available, + handler.reserves_decoded_bytes, ), ) @@ -93,7 +115,10 @@ def call_with_text_prompt(self, inputs, sampling_params): raise NotImplementedError def get_preferred_media_io_kwargs(self): - return {"image": {"format": "np"}} + return { + "image": {"format": "np"}, + "video": {"format": "np", "num_frames": 4}, + } class _Tokenizer: @@ -215,6 +240,7 @@ async def reserve(self, _request_id: str) -> _Reservation: @pytest.fixture def image_llm(monkeypatch): + _set_video_runtime_available(monkeypatch, True) config = _ImageConfig() input_processor = _ImageInputProcessor(config) seen_model_types = [] @@ -238,11 +264,34 @@ def is_valid(model_type, modality): ) +@pytest.fixture +def image_video_llm(monkeypatch): + _set_video_runtime_available(monkeypatch, True) + config = _ImageConfig() + input_processor = _ImageInputProcessor(config) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "is_valid", + lambda model_type, modality: ( + model_type == _ImageConfig.model_type and modality in {"image", "video"} + ), + ) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_placeholder", + lambda model_type, modality: f"<{modality}_{{}}>", + ) + return SimpleNamespace( + _hf_model_config=config, + input_processor=input_processor, + ) + + def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): llm, seen_model_types = image_llm assert request_mapping.media_capabilities_for_model(llm).modalities == {"image"} - assert seen_model_types == [_ImageConfig.model_type] + assert seen_model_types == [_ImageConfig.model_type] * 2 assert ( request_mapping.media_capabilities_for_model( SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) @@ -294,11 +343,271 @@ def test_media_capabilities_are_hashable(): assert hash(first) == hash(second) +def test_media_capability_requires_the_runtime_decoder(image_video_llm, monkeypatch): + _set_video_runtime_available(monkeypatch, False) + + assert request_mapping.media_capabilities_for_model(image_video_llm).modalities == {"image"} + + +def test_video_runtime_availability_requires_importable_opencv(monkeypatch): + monkeypatch.setattr(request_mapping.importlib, "import_module", lambda _name: object()) + assert request_mapping._video_runtime_available() is True + + +@pytest.mark.parametrize("error", [ImportError(), OSError()]) +def test_video_runtime_unavailable_when_opencv_fails_to_load(monkeypatch, error): + def missing(_name): + raise error + + monkeypatch.setattr(request_mapping.importlib, "import_module", missing) + assert request_mapping._video_runtime_available() is False + + def test_unknown_decoded_image_type_is_rejected(): with pytest.raises(TypeError, match="unsupported decoded image type"): request_mapping._decoded_image_bytes(object()) +def test_decoded_video_bytes_include_frames_and_audio(): + video = VideoData( + frames=np.zeros((2, 3, 4, 3), dtype=np.uint8), + metadata={}, + audio=AudioData(samples=np.zeros(8, dtype=np.float32), sample_rate=16000), + ) + + assert request_mapping._decoded_video_bytes(video) == 2 * 3 * 4 * 3 + 8 * 4 + with pytest.raises(TypeError, match="unsupported decoded video type"): + request_mapping._decoded_video_bytes(object()) + + +def test_decoded_media_budget_settles_peak_to_retained_bytes(): + budget = request_mapping._DecodedMediaBudget(20, peak_limit=20) + + first = budget.reserve(15, 5) + first.commit(5) + assert (budget._used, budget._retained, budget._reserved_retained) == (5, 5, 0) + second = budget.reserve(15, 5) + second.release() + assert (budget._used, budget._retained, budget._reserved_retained) == (5, 5, 0) + + +@pytest.mark.parametrize( + "outcome", + [ + "success", + "resampled", + "oversized", + "decode_failure", + "concat_failure", + "no_audio", + "cancelled", + ], +) +def test_embedded_audio_reserves_chunks_and_concatenation(outcome, monkeypatch): + budget = request_mapping._DecodedMediaBudget(27 if outcome == "oversized" else 35) + converted = [] + concatenated = [] + grabs = iter([True]) + capture = SimpleNamespace( + isOpened=lambda: True, + get=lambda _key: 1, + grab=lambda: next(grabs, False), + retrieve=lambda _scratch: (True, np.zeros((1, 1, 3), dtype=np.uint8)), + release=lambda: None, + ) + monkeypatch.setattr( + media_io, + "_get_cv2", + lambda: SimpleNamespace( + VideoCapture=lambda _source: capture, + CAP_PROP_FRAME_COUNT=0, + CAP_PROP_FPS=1, + COLOR_BGR2RGB=2, + cvtColor=lambda _source, _code, dst: dst.fill(0), + ), + ) + + def audio_frame(index): + def to_ndarray(): + converted.append(index) + if outcome == "decode_failure" and index == 1: + raise RuntimeError("audio decode failed") + if outcome == "cancelled" and index == 0: + budget.cancel() + return np.full((1, 4), index, dtype=np.float32) + + return SimpleNamespace( + samples=4, + layout=SimpleNamespace(channels=("mono",)), + to_ndarray=to_ndarray, + ) + + stream = SimpleNamespace( + rate=16000, + format=SimpleNamespace(name="s16" if outcome == "resampled" else "fltp"), + layout=SimpleNamespace(name="mono"), + ) + container = SimpleNamespace( + streams=SimpleNamespace(audio=[] if outcome == "no_audio" else [stream]), + decode=lambda _stream: [audio_frame(0), audio_frame(1)], + ) + monkeypatch.setitem( + sys.modules, + "av", + SimpleNamespace( + open=lambda _source: nullcontext(container), + AudioResampler=lambda **_kwargs: SimpleNamespace(resample=lambda frame: [frame]), + ), + ) + monkeypatch.setenv("TRTLLM_ENABLE_PYAV", "1") + concatenate = np.concatenate + + def checked_concatenate(chunks, **kwargs): + concatenated.append(True) + assert (budget._used, budget._retained, budget._reserved_retained) == (67, 0, 35) + if outcome == "concat_failure": + raise MemoryError("audio concatenation failed") + return concatenate(chunks, **kwargs) + + monkeypatch.setattr(media_io.np, "concatenate", checked_concatenate) + + def load(): + return media_io._load_video_by_cv2( + "video.mp4", + fps=-1, + format="np", + extract_audio=True, + decoded_memory_reserver=budget.reserve, + ) + + if outcome in ("success", "resampled", "no_audio"): + video = load() + retained = media_io._decoded_video_bytes(video) + assert retained == (3 if outcome == "no_audio" else 35) + assert (budget._used, budget._retained, budget._reserved_retained) == ( + retained, + retained, + 0, + ) + if outcome != "no_audio": + np.testing.assert_array_equal(video.audio.samples, [0] * 4 + [1] * 4) + assert video.audio.sample_rate == 16000 + else: + with pytest.raises(MemoryError if outcome == "concat_failure" else ValueError): + load() + assert (budget._used, budget._retained, budget._reserved_retained) == (0, 0, 0) + if outcome in ("oversized", "cancelled"): + assert converted == [0] + assert not concatenated + + +def test_video_retained_limit_does_not_reject_1080p_conversion_peak(): + budget = request_mapping._DecodedMediaBudget(512 * 1024 * 1024) + values = 10 * 1080 * 1920 * 3 + reservation = budget.reserve(values * 9, values * 4) + reservation.commit(values * 4) + + +@pytest.mark.parametrize("workers", [1, 8]) +def test_three_720p_videos_fit_independently_of_load_concurrency(workers): + budget = request_mapping._DecodedMediaBudget(512 * 1024 * 1024) + values = 10 * 720 * 1280 * 3 + + def decode(_index): + reservation = budget.reserve(values * 9, values * 4) + reservation.commit(values * 4) + + with ThreadPoolExecutor(max_workers=workers) as pool: + list(pool.map(decode, range(3))) + + +@pytest.mark.parametrize("cancel_during", ["grab", "retrieve", "convert"]) +def test_video_cancellation_stops_between_frame_operations(cancel_during, monkeypatch): + budget = request_mapping._DecodedMediaBudget(1024) + operations = [] + + def grab(): + operations.append("grab") + if cancel_during == "grab" and operations.count("grab") == 2: + budget.cancel() + return True + + def retrieve(_scratch): + operations.append("retrieve") + if cancel_during == "retrieve" and operations.count("retrieve") == 2: + budget.cancel() + return True, np.zeros((1, 1, 3), dtype=np.uint8) + + def convert(_source, _code, dst): + operations.append("convert") + dst.fill(0) + if cancel_during == "convert": + budget.cancel() + + capture = SimpleNamespace( + isOpened=lambda: True, + get=lambda key: 3 if key == 0 else 1, + grab=grab, + retrieve=retrieve, + release=lambda: operations.append("release"), + ) + monkeypatch.setattr( + media_io, + "_get_cv2", + lambda: SimpleNamespace( + VideoCapture=lambda _source: capture, + CAP_PROP_FRAME_COUNT=0, + CAP_PROP_FPS=1, + COLOR_BGR2RGB=2, + cvtColor=convert, + ), + ) + with pytest.raises(ValueError, match="cancelled"): + media_io._load_video_by_cv2( + "video.mp4", + num_frames=3, + fps=-1, + format="np", + decoded_memory_reserver=budget.reserve, + ) + expected = ["grab", "retrieve", "convert"] + if cancel_during in ("grab", "retrieve"): + expected.append("grab") + if cancel_during == "retrieve": + expected.append("retrieve") + assert operations == expected + ["release"] + assert (budget._used, budget._retained, budget._reserved_retained) == (0, 0, 0) + + +@pytest.mark.parametrize("cancel", [False, True]) +def test_video_peak_reservation_waits_for_settlement_or_cancellation(cancel, monkeypatch): + budget = request_mapping._DecodedMediaBudget(20, peak_limit=20) + first = budget.reserve(15, 5) + waiting = threading.Event() + wait = budget._condition.wait + + def observed_wait(): + waiting.set() + return wait() + + monkeypatch.setattr(budget._condition, "wait", observed_wait) + with ThreadPoolExecutor(max_workers=1) as pool: + second = pool.submit(budget.reserve, 15, 5) + try: + assert waiting.wait(timeout=5) + assert not second.done() + if cancel: + budget.cancel() + with pytest.raises(ValueError, match="cancelled"): + second.result(timeout=5) + else: + first.commit(5) + second.result(timeout=5).commit(5) + finally: + budget.cancel() + first.release() + + @pytest.mark.parametrize("abandon_first", [False, True]) def test_abandoned_preprocessing_reclaims_exports_regardless_of_completion_order(abandon_first): import torch @@ -476,7 +785,7 @@ def test_media_capability_rejects_a_malformed_placeholder_at_startup(image_llm, async def test_handler_table_drives_capability_and_mapping(image_llm, monkeypatch): llm, _ = image_llm - def video_loader_factory(_llm): + def video_loader_factory(_llm, _budget): async def load(item, index): return f"video-{index}:{bytes(item.raw_bytes).decode()}" @@ -600,6 +909,78 @@ async def async_load_bytes(self, data): ] +@pytest.mark.asyncio +async def test_image_and_video_sources_preserve_global_order(image_video_llm, monkeypatch): + created_with = [] + + class _ImageIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append(("image", defaults, runtime)) + return cls() + + async def async_load_bytes(self, data): + return np.full((1, 1, 3), data[0], dtype=np.uint8) + + class _VideoIo: + @classmethod + def create(cls, defaults, runtime): + created_with.append(("video", defaults, runtime)) + return cls() + + async def async_load(self, source): + return VideoData( + frames=np.zeros((1, 1, 1, 3), dtype=np.uint8), + metadata={"source": source}, + ) + + monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(request_mapping, "VideoMediaIO", _VideoIo) + request = generation_pb2.GenerateRequest( + prompt="mixed media", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + url="https://example.com/first.mp4", + uuid="video-id", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_IMAGE, + raw_bytes=b"i", + mime_type="image/png", + ), + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + data_uri="data:video/mp4;base64,dmlkZW8=", + ), + ], + ) + + mapped = await request_mapping._input_from_request(request, image_video_llm) + + assert len(created_with) == 2 + video_modality, video_defaults, video_runtime = created_with[0] + assert video_modality == "video" + assert video_defaults == {"format": "np", "num_frames": 4} + assert set(video_runtime) == {"decoded_memory_reserver"} + assert callable(video_runtime["decoded_memory_reserver"]) + assert created_with[1] == ("image", {"format": "np"}, None) + assert [video.metadata["source"] for video in mapped["multi_modal_data"]["video"]] == [ + "https://example.com/first.mp4", + "data:video/mp4;base64,dmlkZW8=", + ] + assert mapped["multi_modal_data"]["image"][0].shape == (1, 1, 3) + assert mapped["multi_modal_uuids"] == { + "video": ["video-id", None], + "image": [None], + } + assert mapped["mm_item_order"] == [ + {"modality": "video", "index": 0, "placeholder": ""}, + {"modality": "image", "index": 0, "placeholder": ""}, + {"modality": "video", "index": 1, "placeholder": ""}, + ] + + @pytest.mark.asyncio @pytest.mark.parametrize( ("invalid_case", "metadata", "expected_code"), @@ -897,6 +1278,45 @@ async def async_load_bytes(self, data): await request_mapping._input_from_request(request, llm) +@pytest.mark.asyncio +async def test_concurrent_video_loads_share_budget_before_allocation(image_video_llm, monkeypatch): + allocations = [] + + class _VideoIo: + def __init__(self, reserver): + self._reserver = reserver + + @classmethod + def create(cls, defaults, runtime): + return cls(runtime["decoded_memory_reserver"]) + + async def async_load_bytes(self, data): + reservation = self._reserver(12, 12) + allocations.append(bytes(data)) + try: + await asyncio.Future() + finally: + reservation.release() + + monkeypatch.setattr(request_mapping, "VideoMediaIO", _VideoIo) + monkeypatch.setattr(request_mapping, "_MAX_TOTAL_DECODED_MEDIA_BYTES", 20) + request = generation_pb2.GenerateRequest( + prompt="videos", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=value, + ) + for value in (b"first", b"second") + ], + ) + + with pytest.raises(ValueError, match="decoded media exceeds"): + await request_mapping._input_from_request(request, image_video_llm) + + assert allocations == [b"first"] + + @pytest.mark.asyncio async def test_image_mapping_supports_token_ids(image_llm, monkeypatch): llm, _ = image_llm @@ -1061,8 +1481,11 @@ 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): - capability_llm, _ = image_llm +@pytest.mark.parametrize("modality", ["image", "video"]) +async def test_media_preprocessing_is_offloaded_and_abort_aware( + image_video_llm, monkeypatch, outcome, modality +): + capability_llm = image_video_llm llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) loop = asyncio.get_running_loop() serving_thread = threading.get_ident() @@ -1094,7 +1517,8 @@ async def async_load_bytes(self, data): def decode(): assert threading.current_thread().name.startswith("test_media") - return np.zeros((1, 1, 3), dtype=np.uint8) + frames = np.zeros((1, 1, 3), dtype=np.uint8) + return frames if modality == "image" else VideoData(frames=[frames], metadata={}) return await BaseMediaIO._run_in_executor(decode) @@ -1108,9 +1532,9 @@ def preprocess(inputs, sampling_params, disaggregated_params): assert release.wait(timeout=5), "test must release the preprocessing thread" import torch - params.multimodal_data = {"image": {"pixel_values": torch.ones(1024 * 1024)}} + params.multimodal_data = {modality: {"pixel_values": torch.ones(1024 * 1024)}} params.to_handle("multimodal_data") - handle = params.multimodal_data["image"]["pixel_values"] + handle = params.multimodal_data[modality]["pixel_values"] export_paths.append( Path("/dev/shm") / base64.b64decode(handle["storage_handle"]).decode().lstrip("/") ) @@ -1126,6 +1550,7 @@ def generate_async(**kwargs): return llm.result_handle monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) + monkeypatch.setattr(request_mapping, "VideoMediaIO", _ImageIo) monkeypatch.setattr(llm, "preprocess", preprocess) monkeypatch.setattr(llm, "generate_async", generate_async) # Pools must close during the call phase, before pytest-threadleak's snapshot. @@ -1144,9 +1569,16 @@ def generate_async(**kwargs): request = generation_pb2.GenerateRequest( request_id="preprocessing", model="image-model", - prompt="", + prompt=f"<{modality}_1>", media=[ - generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image") + generation_pb2.MediaItem( + modality=( + generation_pb2.MODALITY_IMAGE + if modality == "image" + else generation_pb2.MODALITY_VIDEO + ), + raw_bytes=b"media", + ) ], ) context = FakeServicerContext() @@ -1556,21 +1988,14 @@ async def async_load(self, source): @pytest.mark.asyncio -@pytest.mark.parametrize( - ("modality", "name"), - [ - (generation_pb2.MODALITY_VIDEO, "video"), - (generation_pb2.MODALITY_AUDIO, "audio"), - ], -) -async def test_video_and_audio_are_explicitly_unsupported(image_llm, modality, name): +async def test_audio_is_explicitly_unsupported(image_llm): llm, _ = image_llm request = generation_pb2.GenerateRequest( prompt="media", - media=[generation_pb2.MediaItem(modality=modality, raw_bytes=b"data")], + media=[generation_pb2.MediaItem(modality=generation_pb2.MODALITY_AUDIO, raw_bytes=b"data")], ) - with pytest.raises(UnsupportedFeatureError, match=name): + with pytest.raises(UnsupportedFeatureError, match="audio"): await request_mapping._input_from_request(request, llm) @@ -1614,7 +2039,7 @@ async def test_image_is_rejected_for_a_text_only_loaded_model(): raw_bytes=b"image", mime_type="video/mp4", ), - "must be an image type", + "must start with image/", ), ], ) @@ -1624,3 +2049,20 @@ async def test_malformed_image_items_are_rejected(image_llm, item, message): with pytest.raises(ValueError, match=message): await request_mapping._input_from_request(request, llm) + + +@pytest.mark.asyncio +async def test_raw_video_mime_type_is_validated(image_video_llm): + request = generation_pb2.GenerateRequest( + prompt="video", + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + raw_bytes=b"video", + mime_type="image/png", + ) + ], + ) + + with pytest.raises(ValueError, match="must start with video/"): + await request_mapping._input_from_request(request, image_video_llm) diff --git a/tests/unittest/grpc/openengine/test_multimodal_conformance.py b/tests/unittest/grpc/openengine/test_multimodal_conformance.py index 8e0b83fe3b71..80f4a94563e1 100644 --- a/tests/unittest/grpc/openengine/test_multimodal_conformance.py +++ b/tests/unittest/grpc/openengine/test_multimodal_conformance.py @@ -1,10 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Image conformance tests against a real TensorRT-LLM VLM engine.""" +"""Multimodal conformance tests against a real TensorRT-LLM VLM engine.""" import asyncio import base64 +import importlib +from pathlib import Path import pytest @@ -46,19 +48,28 @@ def vlm(): engine.shutdown() -@pytest.mark.parametrize("input_kind", ["prompt", "token_ids"]) -def test_image_capability_and_generate_agree_on_a_loaded_vlm(vlm, input_kind): +def _servicers(vlm): inference = OpenEngineInferenceServicer(vlm, MODEL_NAME) control = OpenEngineControlServicer(vlm, MODEL_NAME, inference) - info = asyncio.run( + return inference, control + + +def _model_info(control): + return asyncio.run( control.GetModelInfo( model_pb2.GetModelInfoRequest(model=MODEL_NAME), FakeServicerContext(), ) ) + +@pytest.mark.parametrize("input_kind", ["prompt", "token_ids"]) +def test_image_capability_agrees_with_generate(vlm, input_kind): + inference, control = _servicers(vlm) + info = _model_info(control) + assert info.supports_multimodal is True - assert list(info.extra["trtllm_supported_modalities"]) == ["image"] + assert "image" in info.extra["trtllm_supported_modalities"] prompt = ( "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" @@ -121,3 +132,49 @@ async def collect(): assert finished[0].usage.total_tokens == ( finished[0].usage.prompt_tokens + finished[0].usage.completion_tokens ) + + +def test_video_capability_agrees_with_generate(vlm): + try: + importlib.import_module("cv2") + except (ImportError, OSError) as error: + pytest.fail(f"video conformance requires a working OpenCV install: {error}", pytrace=False) + inference, control = _servicers(vlm) + info = _model_info(control) + assert info.supports_multimodal is True + assert "video" in info.extra["trtllm_supported_modalities"] + + video_path = ( + Path(__file__).parents[2] / "_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.mp4" + ) + + video_prompt = ( + "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" + "Describe this video briefly.<|im_end|>\n<|im_start|>assistant\n" + ) + video_request = generation_pb2.GenerateRequest( + request_id="openengine-video-conformance", + model=MODEL_NAME, + token_ids=generation_pb2.TokenIds(ids=vlm.tokenizer.encode(video_prompt)), + stopping=generation_pb2.StoppingOptions(max_tokens=4), + media=[ + generation_pb2.MediaItem( + modality=generation_pb2.MODALITY_VIDEO, + data_uri="data:video/mp4;base64," + + base64.b64encode(video_path.read_bytes()).decode(), + uuid="colors", + ) + ], + ) + + async def collect_video(): + return [ + response async for response in inference.Generate(video_request, FakeServicerContext()) + ] + + video_responses = asyncio.run(collect_video()) + assert video_responses + assert not any(response.HasField("error") for response in video_responses) + video_finished = [response for response in video_responses if response.HasField("finished")] + assert len(video_finished) == 1 + assert video_finished[0].usage.prompt_tokens > len(video_request.token_ids.ids) diff --git a/tests/unittest/inputs/test_async_media_loading.py b/tests/unittest/inputs/test_async_media_loading.py index 90a1b8731fb4..ba36fffd5bee 100644 --- a/tests/unittest/inputs/test_async_media_loading.py +++ b/tests/unittest/inputs/test_async_media_loading.py @@ -13,7 +13,7 @@ import tempfile import threading from io import BytesIO -from unittest.mock import patch +from unittest.mock import Mock, patch import numpy as np import pytest @@ -322,3 +322,21 @@ async def test_video_accepts_video_mime(self): with patch.object(media_io_module.VideoMediaIO, "load_base64", return_value="decoded"): result = await io.async_load("data:video/mp4;base64,AAAA") assert result == "decoded" + + +def test_video_media_io_forwards_decoded_frame_limit(): + decoded = object() + reserver = Mock() + io = media_io_module.VideoMediaIO( + max_decoded_frame_bytes=123, + decoded_memory_reserver=reserver, + ) + + with ( + patch.object(media_io_module, "_select_cv2_stream_buffered_backend", return_value=7), + patch.object(media_io_module, "_load_video_by_cv2", return_value=decoded) as decode, + ): + assert io.load_bytes(b"video") is decoded + + assert decode.call_args.kwargs["max_decoded_frame_bytes"] == 123 + assert decode.call_args.kwargs["decoded_memory_reserver"] is reserver diff --git a/tests/unittest/inputs/test_video_decode.py b/tests/unittest/inputs/test_video_decode.py index 58289159efa8..e39292841b78 100644 --- a/tests/unittest/inputs/test_video_decode.py +++ b/tests/unittest/inputs/test_video_decode.py @@ -9,7 +9,7 @@ import base64 from io import BytesIO from pathlib import Path -from unittest.mock import patch +from unittest.mock import Mock, patch import numpy as np import pytest @@ -78,6 +78,119 @@ def test_np_format_hits_hf_video_processor_fast_path(sample_video_path: str) -> assert np.shares_memory(video.frames, batched[0]) +@pytest.mark.parametrize(("format", "peak_bytes_per_value"), [("np", 2), ("pil", 2), ("pt", 9)]) +def test_decode_limit_is_checked_before_frame_stack_allocation( + sample_video_path: str, format: str, peak_bytes_per_value: int +) -> None: + peak_bytes = 10 * 64 * 64 * 3 * peak_bytes_per_value + + with patch("tensorrt_llm.inputs.media_io.np.empty") as allocate: + with pytest.raises(ValueError, match=f"require {peak_bytes} bytes"): + _load_video_by_cv2( + sample_video_path, + num_frames=10, + fps=-1, + format=format, + max_decoded_frame_bytes=peak_bytes - 1, + ) + + allocate.assert_not_called() + + +def test_pt_decode_reserves_peak_then_commits_retained_bytes(sample_video_path: str) -> None: + class _Reservation: + committed = None + released = False + + def check_cancelled(self) -> None: + pass + + def commit(self, retained_bytes: int) -> None: + self.committed = retained_bytes + + def release(self) -> None: + self.released = True + + reservation = _Reservation() + reserved = [] + + def reserve(peak_bytes: int, retained_bytes: int) -> _Reservation: + reserved.append((peak_bytes, retained_bytes)) + return reservation + + _load_video_by_cv2( + sample_video_path, + num_frames=10, + fps=-1, + format="pt", + decoded_memory_reserver=reserve, + ) + + assert reserved == [(10 * 64 * 64 * 3 * 9, 10 * 64 * 64 * 3 * 4)] + assert reservation.committed == 10 * 64 * 64 * 3 * 4 + assert reservation.released is False + + +def test_decode_releases_reservation_after_allocation_failure(sample_video_path: str) -> None: + class _Reservation: + released = False + + def check_cancelled(self) -> None: + pass + + def commit(self, retained_bytes: int) -> None: + raise AssertionError("failed decode must not commit") + + def release(self) -> None: + self.released = True + + reservation = _Reservation() + + with patch("tensorrt_llm.inputs.media_io.np.empty", side_effect=MemoryError("allocation")): + with pytest.raises(MemoryError, match="allocation"): + _load_video_by_cv2( + sample_video_path, + num_frames=10, + fps=-1, + format="np", + decoded_memory_reserver=lambda _peak_bytes, _retained_bytes: reservation, + ) + + assert reservation.released is True + + +def test_truncated_numpy_video_owns_only_the_counted_frames(sample_video_path: str) -> None: + capture = cv2.VideoCapture(sample_video_path) + get = capture.get + committed = [] + + class _Reservation: + def check_cancelled(self) -> None: + pass + + def commit(self, retained_bytes: int) -> None: + committed.append(retained_bytes) + + def release(self) -> None: + pass + + # Container metadata overstates the number of frames that actually decode. + truncated = Mock(wraps=capture) + truncated.get.side_effect = lambda key: 50 if key == cv2.CAP_PROP_FRAME_COUNT else get(key) + with patch.object(cv2, "VideoCapture", return_value=truncated): + video = _load_video_by_cv2( + sample_video_path, + num_frames=10, + fps=-1, + format="np", + decoded_memory_reserver=lambda _peak, _retained: _Reservation(), + ) + + assert 0 < len(video.frames) < 10 + assert video.frames.base is None + assert committed == [video.frames.nbytes] + + @pytest.mark.parametrize("input_kind", ["path", "file_uri", "data_uri"]) @pytest.mark.parametrize("force_tempfile", [False, True], ids=["auto_backend", "tempfile"]) def test_offline_loaders_match_serve_video_hash(