From 89d54b7d916aaa61a4e8e80dd918556d12db09da Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Thu, 1 Oct 2026 02:38:45 -0700 Subject: [PATCH 1/2] [None][perf] bound multimodal frontend CPU memory Add opt-in --max_multimodal_cpu_bytes and --max_multimodal_cpu_bytes_per_request options to trtllm-serve. The per-request limit makes the media tracker count unique CPU tensor, NumPy, and PIL storage as items finish decoding; once the limit is exceeded, the remaining loads are cancelled and the request fails with HTTP 413. The total limit reserves one worst-case request slot per frontend process before any media load starts and returns it once the request has been handed to the engine. If only the total limit is set, the per-request limit defaults to it, so one request with media is preprocessed at a time; the server logs the slot count at startup. Text-only requests and servers without the options keep the existing control flow. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- tensorrt_llm/commands/serve.py | 36 +++++- tensorrt_llm/inputs/multimodal.py | 17 +++ tensorrt_llm/inputs/utils.py | 107 +++++++++++++++++- tensorrt_llm/serve/openai_server.py | 40 +++++++ .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../references/trtllm_serve_cli.yaml | 18 +++ ...test_openai_server_multimodal_admission.py | 83 ++++++++++++++ 7 files changed, 298 insertions(+), 4 deletions(-) create mode 100644 tests/unittest/llmapi/apps/test_openai_server_multimodal_admission.py diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 9036efd4a951..45d7e23b30fc 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -58,7 +58,8 @@ parse_disagg_config_file, parse_metadata_server_config_file, validate_config_bool) -from tensorrt_llm.llmapi.llm_args import MultimodalConfig, TorchLlmArgs +from tensorrt_llm.llmapi.llm_args import (MultimodalConfig, TorchLlmArgs, + _parse_binary_byte_string) from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr, split_mpi_env from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory, @@ -1211,6 +1212,27 @@ def launch_visual_gen_server( "\"num_frames\": 16}}' to enable audio extraction from " "video files.", status="prototype") +@stability_option( + "--max_multimodal_cpu_bytes", + type=str, + default=None, + help=("Maximum CPU working set for concurrently preprocessing multimodal " + "requests in each serving frontend process. Each request with media " + "reserves --max_multimodal_cpu_bytes_per_request, so " + "max_multimodal_cpu_bytes // max_multimodal_cpu_bytes_per_request " + "of them run at a time and the rest wait. If " + "--max_multimodal_cpu_bytes_per_request is not set, it defaults to " + "this value and only one runs at a time. Accepts byte values such " + "as '8GiB'. Disabled by default."), + status="prototype") +@stability_option( + "--max_multimodal_cpu_bytes_per_request", + type=str, + default=None, + help=("Maximum decoded CPU storage for one multimodal request. Accepts " + "byte values such as '1GiB'. Defaults to the total multimodal CPU " + "limit when that limit is set."), + status="prototype") @stability_option("--video_pruning_rate", type=float, default=None, @@ -1332,6 +1354,8 @@ def serve( enable_attention_dp: bool, disagg_cluster_uri: Optional[str], media_io_kwargs: Optional[str], + max_multimodal_cpu_bytes: Optional[str], + max_multimodal_cpu_bytes_per_request: Optional[str], agent_percentage: float, agent_types: Optional[str], video_pruning_rate: Optional[float], @@ -1495,7 +1519,11 @@ def _serve_llm(): raise ValueError(f"Invalid JSON for media_io_kwargs: {e}") multimodal_server_config = MultimodalServerConfig( - media_io_kwargs=parsed_media_io_kwargs) + media_io_kwargs=parsed_media_io_kwargs, + max_cpu_bytes=_parse_binary_byte_string(max_multimodal_cpu_bytes), + max_cpu_bytes_per_request=_parse_binary_byte_string( + max_multimodal_cpu_bytes_per_request), + ) if grpc: if num_serve_frontends != 1: @@ -1522,6 +1550,10 @@ def _serve_llm(): server_role, "disagg_cluster_config": disagg_cluster_config, + "max_multimodal_cpu_bytes": + max_multimodal_cpu_bytes, + "max_multimodal_cpu_bytes_per_request": + max_multimodal_cpu_bytes_per_request, } for name, value in unsupported_args.items(): if value is not None: diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 06bae3975faf..7d31e4a16108 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -823,6 +823,23 @@ def has_content(self) -> bool: @dataclass class MultimodalServerConfig(): media_io_kwargs: Optional[dict] = None + max_cpu_bytes: Optional[int] = None + max_cpu_bytes_per_request: Optional[int] = None + + def __post_init__(self) -> None: + if self.max_cpu_bytes is not None and self.max_cpu_bytes <= 0: + raise ValueError("max_cpu_bytes must be positive") + if (self.max_cpu_bytes_per_request is not None + and self.max_cpu_bytes_per_request <= 0): + raise ValueError("max_cpu_bytes_per_request must be positive") + if (self.max_cpu_bytes is not None + and self.max_cpu_bytes_per_request is None): + self.max_cpu_bytes_per_request = self.max_cpu_bytes + if (self.max_cpu_bytes is not None + and self.max_cpu_bytes_per_request is not None + and self.max_cpu_bytes_per_request > self.max_cpu_bytes): + raise ValueError( + "max_cpu_bytes_per_request cannot exceed max_cpu_bytes") def _update_hash(hasher, item: object) -> None: diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index df5e692a30b3..0a34d1434709 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -14,7 +14,7 @@ import numpy as np import soundfile import torch -from PIL import Image +from PIL import Image, ImageMode from torchvision.transforms import ToTensor from transformers import AutoProcessor, PreTrainedTokenizerBase, ProcessorMixin from transformers.utils import logging @@ -47,6 +47,58 @@ logger = logging.get_logger(__name__) +class MultimodalDataTooLargeError(ValueError): + """A request's materialized multimodal inputs exceed its CPU limit.""" + + +def _cpu_storage_bytes( + value: Any, + seen_storages: Optional[set[tuple[str, int, int]]] = None) -> int: + """Count unique CPU tensor, array, and image storage reachable from value.""" + if seen_storages is None: + seen_storages = set() + + if isinstance(value, torch.Tensor): + if value.device.type != "cpu": + return 0 + storage = value.untyped_storage() + key = ("cpu", storage.data_ptr(), storage.nbytes()) + if key in seen_storages: + return 0 + seen_storages.add(key) + return storage.nbytes() + + if isinstance(value, np.ndarray): + owner = value + while isinstance(owner.base, np.ndarray): + owner = owner.base + data_ptr = owner.__array_interface__["data"][0] + key = ("cpu", data_ptr, owner.nbytes) + if key in seen_storages: + return 0 + seen_storages.add(key) + return owner.nbytes + + if isinstance(value, Image.Image): + mode = ImageMode.getmode(value.mode) + size_bytes = (value.width * value.height * len(mode.bands) * + np.dtype(mode.typestr).itemsize) + key = ("pil", id(value), size_bytes) + if key in seen_storages: + return 0 + seen_storages.add(key) + return size_bytes + + if isinstance(value, BaseModalityData): + return _cpu_storage_bytes(vars(value), seen_storages) + if isinstance(value, dict): + return sum( + _cpu_storage_bytes(item, seen_storages) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(_cpu_storage_bytes(item, seen_storages) for item in value) + return 0 + + def load_base64_image(parsed_url: str) -> Image.Image: data_spec, data = parsed_url.path.split(",", 1) media_type, data_type = data_spec.split(";", 1) @@ -426,7 +478,12 @@ async def _retrieve( out[modality].append(result) return dict(out) - # _data and _embeddings also gathered concurrently + max_bytes = self._multimodal_server_config.max_cpu_bytes_per_request + if max_bytes is not None: + return await self._retrieve_with_byte_limit(max_bytes) + + # _data and _embeddings also gathered concurrently. Preserve this + # zero-overhead path when frontend byte admission is not configured. data_result, embed_result = await asyncio.gather( _retrieve(self._data), _retrieve(self._embeddings)) return data_result, embed_result @@ -452,6 +509,52 @@ async def discard_all_async( pending.clear() return None, None + async def _retrieve_with_byte_limit( + self, max_bytes: int + ) -> tuple[Optional[Dict[str, List[Any]]], Optional[Dict[str, List[Any]]]]: + """Collect completed items until their unique CPU storage hits the limit.""" + pairs = [(False, modality, item) + for modality, items in self._data.items() for item in items] + pairs.extend((True, modality, item) + for modality, items in self._embeddings.items() + for item in items) + if not pairs: + return None, None + + async def _indexed_result(index: int, + item: Coroutine) -> tuple[int, Any]: + return index, await item + + tasks = [ + asyncio.create_task(_indexed_result(index, item)) + for index, (_, _, item) in enumerate(pairs) + ] + results: list[Any] = [None] * len(tasks) + seen_storages: set[tuple[str, int, int]] = set() + resident_bytes = 0 + try: + for completed in asyncio.as_completed(tasks): + index, result = await completed + resident_bytes += _cpu_storage_bytes(result, seen_storages) + if resident_bytes > max_bytes: + raise MultimodalDataTooLargeError( + "Multimodal inputs require at least " + f"{resident_bytes} CPU bytes, exceeding the per-request " + f"limit of {max_bytes} bytes") + results[index] = result + except BaseException: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + data_result: dict[str, list] = defaultdict(list) + embed_result: dict[str, list] = defaultdict(list) + for (is_embedding, modality, _), result in zip(pairs, results): + (embed_result + if is_embedding else data_result)[modality].append(result) + return (dict(data_result) or None, dict(embed_result) or None) + def retrieve_all_sync( self ) -> tuple[Optional[Dict[str, List[Any]]], Optional[Dict[str, List[Any]]]]: diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 2fc87d828883..f50d4bca9a5e 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -49,6 +49,7 @@ from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.inputs.utils import (ConversationMessage, + MultimodalDataTooLargeError, async_apply_chat_template) from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams from tensorrt_llm.llmapi import MultimodalEncoder, SchedulingParams, tracing @@ -574,6 +575,17 @@ def __init__( merged.setdefault(modality, {}).update(kw) cfg.media_io_kwargs = merged self.multimodal_server_config = cfg + self._mm_cpu_request_slots: Optional[asyncio.BoundedSemaphore] = None + if (self.multimodal_server_config is not None + and self.multimodal_server_config.max_cpu_bytes is not None): + request_bytes = self.multimodal_server_config.max_cpu_bytes_per_request + assert request_bytes is not None + num_slots = self.multimodal_server_config.max_cpu_bytes // request_bytes + self._mm_cpu_request_slots = asyncio.BoundedSemaphore(num_slots) + logger.info( + f"Multimodal CPU limit: at most {num_slots} request(s) with " + "media are decoded and preprocessed at a time in this " + f"frontend process, each capped at {request_bytes} bytes") self.allow_request_chat_template = allow_request_chat_template self._internal_disagg_auth_key = internal_disagg_auth_key self._enable_rl_control_endpoints = enable_rl_control_endpoints @@ -1102,6 +1114,14 @@ def create_error_response( return JSONResponse(content=error_response.model_dump(), status_code=error_response.code) + async def _acquire_mm_cpu_slot(self) -> bool: + """Reserve one request's worst-case MM CPU working set, if limited.""" + config = self.multimodal_server_config + if config is None or config.max_cpu_bytes is None: + return False + await self._mm_cpu_request_slots.acquire() + return True + def _create_invalid_response_id_error(self, response_id: str) -> Response: return self.create_error_response( err_type="InvalidRequestError", @@ -1842,6 +1862,7 @@ async def chat_stream_generator( logger.error(traceback.format_exc()) raise + mm_cpu_slot_acquired = False try: ensure_request_chat_template_allowed( request, self.allow_request_chat_template) @@ -2098,6 +2119,14 @@ async def chat_stream_generator( resolve_media=resolve_media, ) + # Media loads have not started yet; reserve CPU before they do. + # mm_item_order also lists items whose placeholder is empty (e.g. + # Whisper audio); embedding items appear only in the counts. + has_multimodal_data = resolve_media and (bool(mm_item_order) or + any(mm_placeholder_counts)) + if has_multimodal_data: + mm_cpu_slot_acquired = await self._acquire_mm_cpu_slot() + # Decode base64 int32 prompt_token_ids relayed by the orchestrator. if request.prompt_token_ids is None and request.prompt_token_ids_b64: import numpy as np @@ -2223,6 +2252,9 @@ async def chat_stream_generator( priority=request.priority if request.priority is not None else DEFAULT_REQUEST_PRIORITY, ) + if mm_cpu_slot_acquired: + self._mm_cpu_request_slots.release() + mm_cpu_slot_acquired = False asyncio.create_task(self.await_disconnected(raw_request, promise)) if not self.postproc_worker_enabled: postproc_args.tokenizer = self.tokenizer @@ -2254,11 +2286,19 @@ async def chat_stream_generator( # If internal executor error is raised, shutdown the server _record_generator_termination(self.generator) signal.raise_signal(signal.SIGINT) + except MultimodalDataTooLargeError as e: + return self.create_error_response( + str(e), + err_type="RequestTooLargeError", + status_code=HTTPStatus.REQUEST_ENTITY_TOO_LARGE) except ValueError as e: return self.create_error_response(str(e)) except Exception as e: logger.error(traceback.format_exc()) return self.create_error_response(str(e)) + finally: + if mm_cpu_slot_acquired: + self._mm_cpu_request_slots.release() async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index cc3e467efd54..a49bc5296fac 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -147,6 +147,7 @@ l0_cpu: - unittest/llmapi/apps/test_harmony_parsing.py::TestStripIncompleteMessagesReporting - unittest/llmapi/apps/test_kimi_serve_extensions.py - unittest/llmapi/apps/test_openai_chat_disagg_multimodal.py + - unittest/llmapi/apps/test_openai_server_multimodal_admission.py - unittest/llmapi/apps/test_reasoning_prompt_resolution.py # test_request_metrics.py as a whole is not registered on any list, so only the # two cases added with the KV-transfer fix are selected here. Registering the diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index 875c5bd2a057..28175fa0a4cf 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -238,6 +238,24 @@ commands: is_flag: false flags: - "--max_beam_width" + max_multimodal_cpu_bytes: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--max_multimodal_cpu_bytes" + max_multimodal_cpu_bytes_per_request: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--max_multimodal_cpu_bytes_per_request" max_num_tokens: type: int default: 8192 diff --git a/tests/unittest/llmapi/apps/test_openai_server_multimodal_admission.py b/tests/unittest/llmapi/apps/test_openai_server_multimodal_admission.py new file mode 100644 index 000000000000..ce25b16a831d --- /dev/null +++ b/tests/unittest/llmapi/apps/test_openai_server_multimodal_admission.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU memory admission for multimodal requests in the OpenAI server.""" + +import asyncio + +import numpy as np +import pytest + +from tensorrt_llm.inputs.multimodal import MultimodalServerConfig +from tensorrt_llm.inputs.utils import MultimodalDataTooLargeError, MultimodalDataTracker +from tensorrt_llm.serve.openai_server import OpenAIServer + +pytestmark = pytest.mark.cpu_only + + +def _make_tracker(max_bytes: int) -> MultimodalDataTracker: + return MultimodalDataTracker( + model_type="test_model", + multimodal_server_config=MultimodalServerConfig(max_cpu_bytes_per_request=max_bytes), + ) + + +async def _loaded(value): + return value + + +@pytest.mark.asyncio +async def test_shared_storage_is_counted_once(): + storage = np.zeros((16,), dtype=np.uint8) + tracker = _make_tracker(max_bytes=16) + tracker._data["image"].extend([_loaded(storage[:8]), _loaded(storage[8:])]) + + data, _ = await tracker.retrieve_all_async() + + assert len(data["image"]) == 2 + + +@pytest.mark.asyncio +async def test_limit_cancels_remaining_items(): + cancelled = asyncio.Event() + + async def _pending(): + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + + tracker = _make_tracker(max_bytes=16) + tracker._data["image"].extend([_loaded(np.zeros((17,), dtype=np.uint8)), _pending()]) + + with pytest.raises(MultimodalDataTooLargeError, match="17 CPU bytes"): + await tracker.retrieve_all_async() + + assert cancelled.is_set() + + +def test_multimodal_server_cpu_limits_are_consistent(): + config = MultimodalServerConfig(max_cpu_bytes=1024) + assert config.max_cpu_bytes_per_request == 1024 + + with pytest.raises(ValueError, match="cannot exceed"): + MultimodalServerConfig( + max_cpu_bytes=1024, + max_cpu_bytes_per_request=2048, + ) + + +@pytest.mark.asyncio +async def test_total_limit_serializes_multimodal_preprocessing(): + server = object.__new__(OpenAIServer) + server.multimodal_server_config = MultimodalServerConfig(max_cpu_bytes=1) + server._mm_cpu_request_slots = asyncio.BoundedSemaphore(1) + + assert await server._acquire_mm_cpu_slot() + next_request = asyncio.create_task(server._acquire_mm_cpu_slot()) + await asyncio.sleep(0) + assert not next_request.done() + + server._mm_cpu_request_slots.release() + assert await next_request + server._mm_cpu_request_slots.release() From ee506fbe54a8d40d5d89bc0f8c350cc777c0730d Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Thu, 1 Oct 2026 02:44:33 -0700 Subject: [PATCH 2/2] [None][perf] account multimodal frontend request memory When --max_multimodal_cpu_bytes_per_request is set, a pure ASGI middleware rejects /v1/chat/completions, /v1/messages, and /v1/responses bodies above the limit with HTTP 413 before FastAPI buffers and parses them, for both declared and chunked bodies. The received byte count seeds the per-request decoded-media ledger, and the decoded media plus the processed multimodal CPU tensors and shared CPU handles are checked again before the request is handed to the engine. A request rejected by that check rebuilds and drops the shared CPU handles preprocessing exported, since no worker will consume them; otherwise their shared memory would stay allocated until the process exits. The OpenAI routes report 413 with error type RequestTooLargeError; /v1/messages uses the Anthropic request_too_large type. Once a request is handed off, openai_chat drops its references to the decoded media so that a non-streaming request does not keep them alive until generation completes. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- .../commands/trtllm-serve/trtllm-serve.rst | 30 +++ tensorrt_llm/commands/serve.py | 8 +- tensorrt_llm/inputs/utils.py | 60 +++++- tensorrt_llm/serve/chat_utils.py | 7 +- tensorrt_llm/serve/openai_server.py | 112 ++++++++++++ ...test_openai_server_multimodal_admission.py | 173 +++++++++++++++++- 6 files changed, 378 insertions(+), 12 deletions(-) diff --git a/docs/source/commands/trtllm-serve/trtllm-serve.rst b/docs/source/commands/trtllm-serve/trtllm-serve.rst index 36e154c02408..cc6417978b91 100644 --- a/docs/source/commands/trtllm-serve/trtllm-serve.rst +++ b/docs/source/commands/trtllm-serve/trtllm-serve.rst @@ -214,6 +214,36 @@ model. {"type": "audio_url", "audio_url": {"url": "https://example.com/audio.mp3"}} ]} +Multimodal CPU Memory Limits +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Decoded media can need far more CPU memory than the request that carried it. +Two prototype options bound this memory in each serving frontend process; both +are disabled by default: + +* ``--max_multimodal_cpu_bytes_per_request`` caps one request. Request bodies + above the cap on ``/v1/chat/completions``, ``/v1/messages``, and + ``/v1/responses`` are rejected before they are parsed, including text-only + requests. For chat and Messages requests, the raw body plus the decoded media + and processed multimodal tensors must also stay under the cap; media decoding + stops as soon as it is exceeded. Rejected requests get HTTP 413. +* ``--max_multimodal_cpu_bytes`` caps the requests with media that are decoded + and preprocessed at the same time. Each one reserves the per-request cap, so + at most ``max_multimodal_cpu_bytes // max_multimodal_cpu_bytes_per_request`` + run concurrently and the rest wait. Text-only requests do not wait. If + ``--max_multimodal_cpu_bytes_per_request`` is not set, it defaults to + ``--max_multimodal_cpu_bytes``, so only one request with media is decoded and + preprocessed at a time. + +.. code-block:: bash + + trtllm-serve Qwen/Qwen2-VL-7B-Instruct \ + --max_multimodal_cpu_bytes 8GiB \ + --max_multimodal_cpu_bytes_per_request 1GiB + +The limits count tensor, array, and image storage, not process RSS. Leave +headroom for decoder temporaries and allocator overhead. + Visual Generation Serving diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 45d7e23b30fc..d1362aa933df 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1229,9 +1229,11 @@ def launch_visual_gen_server( "--max_multimodal_cpu_bytes_per_request", type=str, default=None, - help=("Maximum decoded CPU storage for one multimodal request. Accepts " - "byte values such as '1GiB'. Defaults to the total multimodal CPU " - "limit when that limit is set."), + help=("Maximum CPU storage for one multimodal request: its raw HTTP body " + "plus decoded media and processed multimodal tensors. Chat, " + "Messages, and Responses request bodies above this size are " + "rejected with HTTP 413. Accepts byte values such as '1GiB'. " + "Defaults to the total multimodal CPU limit when that limit is set."), status="prototype") @stability_option("--video_pruning_rate", type=float, diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 0a34d1434709..f3439ef692e0 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -5,6 +5,7 @@ import base64 import tempfile from collections import defaultdict +from dataclasses import fields, is_dataclass from io import BytesIO from pathlib import Path from typing import (Any, Collection, Coroutine, Dict, List, Optional, Tuple, @@ -51,9 +52,15 @@ class MultimodalDataTooLargeError(ValueError): """A request's materialized multimodal inputs exceed its CPU limit.""" +def _is_shared_cpu_tensor_handle(value: dict) -> bool: + """Whether value is a CPU handle from `SharedTensorContainer.dump_to_dict`.""" + return ("method_key" in value and "storage_size" in value + and "storage_dtype" in value) + + def _cpu_storage_bytes( value: Any, - seen_storages: Optional[set[tuple[str, int, int]]] = None) -> int: + seen_storages: Optional[set[tuple[Any, ...]]] = None) -> int: """Count unique CPU tensor, array, and image storage reachable from value.""" if seen_storages is None: seen_storages = set() @@ -92,13 +99,56 @@ def _cpu_storage_bytes( if isinstance(value, BaseModalityData): return _cpu_storage_bytes(vars(value), seen_storages) if isinstance(value, dict): + if _is_shared_cpu_tensor_handle(value): + dtype = getattr( + torch, + str(value["storage_dtype"]).removeprefix("torch."), + None, + ) + if isinstance(dtype, torch.dtype): + size_bytes = int(value["storage_size"]) * dtype.itemsize + key = ("shared_cpu", value.get("storage_handle"), size_bytes) + if key in seen_storages: + return 0 + seen_storages.add(key) + return size_bytes return sum( _cpu_storage_bytes(item, seen_storages) for item in value.values()) if isinstance(value, (list, tuple)): return sum(_cpu_storage_bytes(item, seen_storages) for item in value) + if is_dataclass(value) and not isinstance(value, type): + return sum( + _cpu_storage_bytes(getattr(value, field.name), seen_storages) + for field in fields(value)) return 0 +def _release_shared_cpu_tensors(value: Any) -> None: + """Free the shared memory behind CPU tensor handles reachable from value. + + Exporting a CPU tensor as a handle (`MultimodalParams.to_handle`) keeps its + shared-memory segment alive until a consumer rebuilds the tensor. A request + dropped before any worker consumes its handles must rebuild them here, or + the segments stay in /dev/shm until the process exits. CUDA handles are + left alone. + """ + if isinstance(value, dict): + if not _is_shared_cpu_tensor_handle(value): + for item in value.values(): + _release_shared_cpu_tensors(item) + return + from tensorrt_llm._torch.shared_tensor import SharedTensorContainer + try: + # The rebuild takes over the reference added by the export, so + # dropping the rebuilt view frees the segment. + SharedTensorContainer.from_dict(value).get_local_view() + except (KeyError, RuntimeError, ValueError) as e: + logger.warning(f"Failed to release a shared CPU tensor: {e}") + elif isinstance(value, (list, tuple)): + for item in value: + _release_shared_cpu_tensors(item) + + def load_base64_image(parsed_url: str) -> Image.Image: data_spec, data = parsed_url.path.split(",", 1) media_type, data_type = data_spec.split(";", 1) @@ -427,6 +477,7 @@ def __init__( model_type: str, multimodal_server_config: Optional[MultimodalServerConfig] = None, request_media_io_kwargs: Optional[Dict[str, Dict[str, Any]]] = None, + initial_cpu_bytes: int = 0, ): self._model_type = model_type self._data = defaultdict[str, list](list) @@ -446,6 +497,9 @@ def __init__( # Per-request override merged with the server default at media-load # time; see `BaseMediaIO.merge_kwargs` in `inputs/media_io.py`. self._request_media_io_kwargs = request_media_io_kwargs + if initial_cpu_bytes < 0: + raise ValueError("initial_cpu_bytes must be non-negative") + self._initial_cpu_bytes = initial_cpu_bytes @property def request_media_io_kwargs(self) -> Optional[Dict[str, Dict[str, Any]]]: @@ -530,8 +584,8 @@ async def _indexed_result(index: int, for index, (_, _, item) in enumerate(pairs) ] results: list[Any] = [None] * len(tasks) - seen_storages: set[tuple[str, int, int]] = set() - resident_bytes = 0 + seen_storages: set[tuple[Any, ...]] = set() + resident_bytes = self._initial_cpu_bytes try: for completed in asyncio.as_completed(tasks): index, result = await completed diff --git a/tensorrt_llm/serve/chat_utils.py b/tensorrt_llm/serve/chat_utils.py index 27c270823972..27bc345e57d1 100644 --- a/tensorrt_llm/serve/chat_utils.py +++ b/tensorrt_llm/serve/chat_utils.py @@ -448,6 +448,7 @@ def parse_chat_messages_coroutines( request_media_io_kwargs: Optional[Dict[str, Dict[str, Any]]] = None, model_type_override: Optional[str] = None, resolve_media: bool = True, + initial_cpu_bytes: int = 0, ) -> Tuple[List[ConversationMessage], Coroutine[Any, Any, tuple[Optional[Dict[ str, List[Any]]], Optional[Dict[str, List[Any]]]]], list[dict[str, int]]]: @@ -468,6 +469,9 @@ def parse_chat_messages_coroutines( resolve_media: When False the returned coroutine closes the pending loads instead of running them and yields `(None, None)`. + initial_cpu_bytes: CPU bytes already held by the request (e.g. + its raw body), charged against + `multimodal_server_config.max_cpu_bytes_per_request`. Returns: `(conversation, mm_coroutine, mm_placeholder_counts)` where @@ -481,7 +485,8 @@ def parse_chat_messages_coroutines( mm_data_tracker = MultimodalDataTracker( model_type, multimodal_server_config, - request_media_io_kwargs=request_media_io_kwargs) + request_media_io_kwargs=request_media_io_kwargs, + initial_cpu_bytes=initial_cpu_bytes) # Determine content format to decide placeholder strategy. # diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index f50d4bca9a5e..79aa7b2e08ea 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -34,6 +34,7 @@ from fastapi.routing import APIRoute from pydantic import ValidationError from starlette.routing import Mount +from starlette.types import ASGIApp, Message, Receive, Scope, Send from transformers import AutoProcessor from tensorrt_llm._startup import _StartupTimer @@ -50,6 +51,8 @@ from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.inputs.utils import (ConversationMessage, MultimodalDataTooLargeError, + _cpu_storage_bytes, + _release_shared_cpu_tensors, async_apply_chat_template) from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams from tensorrt_llm.llmapi import MultimodalEncoder, SchedulingParams, tracing @@ -161,6 +164,77 @@ def _is_visual_gen_instance(obj) -> bool: # yapf: enable + +class _MultimodalRequestBodyLimitMiddleware: + """Reject oversized MM-capable request bodies before FastAPI parses them.""" + + _PATHS = frozenset( + {"/v1/chat/completions", "/v1/messages", "/v1/responses"}) + + def __init__(self, app: ASGIApp, max_bytes: int) -> None: + self.app = app + self.max_bytes = max_bytes + + async def _reject(self, scope: Scope, receive: Receive, send: Send) -> None: + message = (f"Request body exceeds the {self.max_bytes}-byte " + "multimodal CPU limit") + if scope["path"] == "/v1/messages": + response = anthropic_error_response( + message, "request_too_large", + HTTPStatus.REQUEST_ENTITY_TOO_LARGE) + else: + response = JSONResponse( + status_code=HTTPStatus.REQUEST_ENTITY_TOO_LARGE, + content=ErrorResponse( + message=message, + type="RequestTooLargeError", + code=HTTPStatus.REQUEST_ENTITY_TOO_LARGE).model_dump(), + ) + await response(scope, receive, send) + + async def __call__(self, scope: Scope, receive: Receive, + send: Send) -> None: + if (scope["type"] != "http" or scope.get("method") != "POST" + or scope.get("path") not in self._PATHS): + await self.app(scope, receive, send) + return + + content_length = dict(scope.get("headers", [])).get(b"content-length") + if content_length is not None: + try: + declared_bytes = int(content_length) + except ValueError: + # The HTTP server rejects malformed Content-Length headers. + declared_bytes = 0 + if declared_bytes > self.max_bytes: + await self._reject(scope, receive, send) + return + + received_bytes = 0 + body_too_large = False + + async def receive_with_limit() -> Message: + nonlocal body_too_large, received_bytes + message = await receive() + if message["type"] == "http.request": + received_bytes += len(message.get("body", b"")) + scope.setdefault( + "state", + {})["multimodal_request_body_bytes"] = received_bytes + if received_bytes > self.max_bytes: + body_too_large = True + return {"type": "http.disconnect"} + return message + + async def send_if_within_limit(message: Message) -> None: + if not body_too_large: + await send(message) + + await self.app(scope, receive_with_limit, send_if_within_limit) + if body_too_large: + await self._reject(scope, receive, send) + + _msgpack_decoder = msgspec.msgpack.Decoder() @@ -800,6 +874,14 @@ async def validation_exception_handler(request, exc): adjusted_clock=self._adjusted_steady_clock) self.app.add_middleware(ServerArrivalTimeMiddleware, adjusted_clock=self._adjusted_steady_clock) + if (self.multimodal_server_config is not None + and self.multimodal_server_config.max_cpu_bytes_per_request + is not None): + self.app.add_middleware( + _MultimodalRequestBodyLimitMiddleware, + max_bytes=self.multimodal_server_config. + max_cpu_bytes_per_request, + ) def _init_visual_gen(self): self.processor = None @@ -1863,6 +1945,8 @@ async def chat_stream_generator( raise mm_cpu_slot_acquired = False + request_body_bytes = (0 if raw_request is None else getattr( + raw_request.state, "multimodal_request_body_bytes", 0)) try: ensure_request_chat_template_allowed( request, self.allow_request_chat_template) @@ -2106,6 +2190,7 @@ async def chat_stream_generator( self.multimodal_server_config, request_media_io_kwargs=request.media_io_kwargs, resolve_media=resolve_media, + initial_cpu_bytes=request_body_bytes, ) except ValidationError: # ValidatorIterator rejects extra fields; fall back to raw JSON. @@ -2117,6 +2202,7 @@ async def chat_stream_generator( self.multimodal_server_config, request_media_io_kwargs=request.media_io_kwargs, resolve_media=resolve_media, + initial_cpu_bytes=request_body_bytes, ) # Media loads have not started yet; reserve CPU before they do. @@ -2236,6 +2322,24 @@ async def chat_stream_generator( self._input_proc_executor, functools.partial(preprocess_fn, prompt, sampling_params, disaggregated_params)) + max_cpu_bytes = ( + self.multimodal_server_config.max_cpu_bytes_per_request + if self.multimodal_server_config is not None else None) + if max_cpu_bytes is not None and has_multimodal_data: + processed_mm = getattr(generate_inputs, "multimodal_params", + None) + materialized_bytes = request_body_bytes + _cpu_storage_bytes( + (mm_data, processed_mm)) + if materialized_bytes > max_cpu_bytes: + # No worker will consume the handles preprocess + # exported, so free their shared memory here. + if processed_mm is not None: + _release_shared_cpu_tensors( + processed_mm.multimodal_data) + raise MultimodalDataTooLargeError( + "Multimodal request materialized at least " + f"{materialized_bytes} CPU bytes, exceeding the " + f"per-request limit of {max_cpu_bytes} bytes") promise = self.generator.generate_async( inputs=generate_inputs, @@ -2255,6 +2359,11 @@ async def chat_stream_generator( if mm_cpu_slot_acquired: self._mm_cpu_request_slots.release() mm_cpu_slot_acquired = False + # The request is handed off; drop this frame's references to the + # decoded media so a non-streaming request does not keep it alive + # until generation completes. + prompt = generate_inputs = mm_data = mm_embeddings = None + processed_mm = None asyncio.create_task(self.await_disconnected(raw_request, promise)) if not self.postproc_worker_enabled: postproc_args.tokenizer = self.tokenizer @@ -2340,6 +2449,9 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest, message = "Internal server error" err_type = ("invalid_request_error" if 400 <= status < 500 else "api_error") + if status == HTTPStatus.REQUEST_ENTITY_TOO_LARGE: + # Same type as the request body limit's 413 on this route. + err_type = "request_too_large" return anthropic_error_response(message, err_type, status) try: diff --git a/tests/unittest/llmapi/apps/test_openai_server_multimodal_admission.py b/tests/unittest/llmapi/apps/test_openai_server_multimodal_admission.py index ce25b16a831d..21b322df05b3 100644 --- a/tests/unittest/llmapi/apps/test_openai_server_multimodal_admission.py +++ b/tests/unittest/llmapi/apps/test_openai_server_multimodal_admission.py @@ -3,21 +3,32 @@ """CPU memory admission for multimodal requests in the OpenAI server.""" import asyncio +import base64 +import json +import os import numpy as np import pytest - -from tensorrt_llm.inputs.multimodal import MultimodalServerConfig -from tensorrt_llm.inputs.utils import MultimodalDataTooLargeError, MultimodalDataTracker -from tensorrt_llm.serve.openai_server import OpenAIServer +import torch +from fastapi import FastAPI + +from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalServerConfig +from tensorrt_llm.inputs.utils import ( + MultimodalDataTooLargeError, + MultimodalDataTracker, + _cpu_storage_bytes, + _release_shared_cpu_tensors, +) +from tensorrt_llm.serve.openai_server import OpenAIServer, _MultimodalRequestBodyLimitMiddleware pytestmark = pytest.mark.cpu_only -def _make_tracker(max_bytes: int) -> MultimodalDataTracker: +def _make_tracker(max_bytes: int, initial_cpu_bytes: int = 0) -> MultimodalDataTracker: return MultimodalDataTracker( model_type="test_model", multimodal_server_config=MultimodalServerConfig(max_cpu_bytes_per_request=max_bytes), + initial_cpu_bytes=initial_cpu_bytes, ) @@ -36,6 +47,15 @@ async def test_shared_storage_is_counted_once(): assert len(data["image"]) == 2 +@pytest.mark.asyncio +async def test_raw_body_bytes_reduce_decoded_capacity(): + tracker = _make_tracker(max_bytes=16, initial_cpu_bytes=8) + tracker._data["image"].append(_loaded(np.zeros((9,), dtype=np.uint8))) + + with pytest.raises(MultimodalDataTooLargeError, match="17 CPU bytes"): + await tracker.retrieve_all_async() + + @pytest.mark.asyncio async def test_limit_cancels_remaining_items(): cancelled = asyncio.Event() @@ -67,6 +87,149 @@ def test_multimodal_server_cpu_limits_are_consistent(): ) +def test_shared_cpu_tensor_handles_are_counted_once_and_released(): + storage = torch.ones(12) + params = MultimodalParams(multimodal_data={"image": [storage[:6], storage[6:]]}) + params.to_handle("multimodal_data") + del storage + handles = params.multimodal_data["image"] + shm_path = "/dev/shm" + base64.b64decode(handles[0]["storage_handle"]).decode() + + assert _cpu_storage_bytes(handles) == 48 + # Dropping unconsumed handles would leave the segment until process exit. + assert os.path.exists(shm_path) + _release_shared_cpu_tensors(params.multimodal_data) + assert not os.path.exists(shm_path) + + +async def _run_body_limit( + chunks: list[bytes], + *, + max_bytes: int, + path: str = "/v1/chat/completions", + content_length: int | None = None, +): + received_body = bytearray() + observed_state = {} + + async def app(scope, receive, send): + while True: + message = await receive() + received_body.extend(message.get("body", b"")) + if not message.get("more_body", False): + break + observed_state.update(scope.get("state", {})) + await send({"type": "http.response.start", "status": 204, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + messages = [ + { + "type": "http.request", + "body": chunk, + "more_body": index + 1 < len(chunks), + } + for index, chunk in enumerate(chunks) + ] + + async def receive(): + return messages.pop(0) + + sent = [] + + async def send(message): + sent.append(message) + + headers = [] + if content_length is not None: + headers.append((b"content-length", str(content_length).encode())) + scope = { + "type": "http", + "method": "POST", + "path": path, + "headers": headers, + } + middleware = _MultimodalRequestBodyLimitMiddleware(app, max_bytes) + await middleware(scope, receive, send) + return sent[0]["status"], bytes(received_body), observed_state + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/messages"]) +async def test_content_length_rejects_before_reading_body(path): + status, body, _ = await _run_body_limit([b"not read"], max_bytes=8, path=path, content_length=9) + + assert status == 413 + assert body == b"" + + +@pytest.mark.asyncio +async def test_chunked_body_is_rejected_at_limit(): + status, body, _ = await _run_body_limit([b"12345", b"678901"], max_bytes=10) + + assert status == 413 + assert body == b"12345" + + +@pytest.mark.asyncio +async def test_chunked_body_stays_413_through_fastapi(): + app = FastAPI() + + @app.post("/v1/chat/completions") + async def route(payload: dict): + return payload + + messages = [ + {"type": "http.request", "body": b'{"value":', "more_body": True}, + {"type": "http.request", "body": b'"too long"}', "more_body": False}, + ] + + async def receive(): + return messages.pop(0) + + sent = [] + + async def send(message): + sent.append(message) + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/chat/completions", + "raw_path": b"/v1/chat/completions", + "query_string": b"", + "root_path": "", + "headers": [(b"content-type", b"application/json")], + "client": ("test", 1), + "server": ("test", 80), + "state": {}, + } + middleware = _MultimodalRequestBodyLimitMiddleware(app, max_bytes=10) + await middleware(scope, receive, send) + + assert sent[0]["status"] == 413 + assert json.loads(sent[1]["body"])["type"] == "RequestTooLargeError" + + +@pytest.mark.asyncio +async def test_accepted_body_bytes_are_recorded_for_decoded_limit(): + status, body, state = await _run_body_limit([b"12345", b"67890"], max_bytes=10) + + assert status == 204 + assert body == b"1234567890" + assert state["multimodal_request_body_bytes"] == 10 + + +@pytest.mark.asyncio +async def test_non_multimodal_route_is_unchanged(): + status, body, _ = await _run_body_limit([b"123456789"], max_bytes=8, path="/v1/completions") + + assert status == 204 + assert body == b"123456789" + + @pytest.mark.asyncio async def test_total_limit_serializes_multimodal_preprocessing(): server = object.__new__(OpenAIServer)