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 9036efd4a951..d1362aa933df 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,29 @@ 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 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, default=None, @@ -1332,6 +1356,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 +1521,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 +1552,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..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, @@ -14,7 +15,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 +48,107 @@ logger = logging.get_logger(__name__) +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[Any, ...]]] = 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): + 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) @@ -375,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) @@ -394,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]]]: @@ -426,7 +532,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 +563,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[Any, ...]] = set() + resident_bytes = self._initial_cpu_bytes + 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/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 2fc87d828883..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 @@ -49,6 +50,9 @@ from tensorrt_llm.inputs.multimodal import MultimodalServerConfig 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 @@ -160,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() @@ -574,6 +649,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 @@ -788,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 @@ -1102,6 +1196,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 +1944,9 @@ async def chat_stream_generator( logger.error(traceback.format_exc()) 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) @@ -2085,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. @@ -2096,8 +2202,17 @@ 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. + # 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 @@ -2207,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, @@ -2223,6 +2356,14 @@ 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 + # 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 @@ -2254,11 +2395,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: @@ -2300,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/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..21b322df05b3 --- /dev/null +++ b/tests/unittest/llmapi/apps/test_openai_server_multimodal_admission.py @@ -0,0 +1,246 @@ +# 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 base64 +import json +import os + +import numpy as np +import pytest +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, 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, + ) + + +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_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() + + 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, + ) + + +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) + 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()