diff --git a/docs/usage/security.md b/docs/usage/security.md index 0d7894a467..428e64258e 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -94,7 +94,9 @@ denial of service: | Environment Variable | Default | Description | | --- | --- | --- | | `VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB` | `256` | Maximum size in MB for a single remote media response. Oversized responses are rejected while streaming before the full body is materialized in memory. | +| `VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB` | `256` | Maximum total HTTP media bytes downloaded for one request, including concurrent downloads. Enforced while streaming. | | `VLLM_MAX_IMAGE_PIXELS` | `178956970` (~179M pixels) | Maximum decoded image size in pixels. Images exceeding this are rejected before raster memory is allocated. Default matches PIL's built-in 2x decompression-bomb threshold (~680 MB for RGB). | +| `VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES` | `715827880` (~683 MiB) | Maximum aggregate decoded image, video-frame, and audio bytes in one request. Images are charged at four bytes per pixel before decoding; audio is limited before decoding; video files are charged after frame extraction. | | `VLLM_MAX_AUDIO_CLIP_FILESIZE_MB` | `25` | Maximum compressed filesize in MB for a single audio file. Enforced on all audio inputs (multimodal chat URLs, speech-to-text uploads, data: URLs, and local file paths) before decoding begins. | | `VLLM_MAX_AUDIO_DECODE_DURATION_S` | `600` | Maximum decoded audio duration in seconds. Prevents compressed audio from expanding into gigabytes of float32 PCM. | | `VLLM_MAX_AUDIO_DECODE_BYTES` | `268435456` (256 MiB) | Maximum float32 PCM bytes that audio decoding may allocate. Guards against sample-rate forgery where an inflated header sample rate bypasses the duration guard while the actual frame count causes a multi-GiB allocation. | diff --git a/tests/multimodal/media/test_audio.py b/tests/multimodal/media/test_audio.py index a75c20979e..5728fa9a46 100644 --- a/tests/multimodal/media/test_audio.py +++ b/tests/multimodal/media/test_audio.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading +from concurrent.futures import ThreadPoolExecutor from io import BytesIO from pathlib import Path from unittest.mock import patch @@ -9,6 +11,7 @@ import pybase64 as base64 import pytest import soundfile as sf +import vllm.envs as envs from vllm.multimodal.media import AudioMediaIO from vllm.multimodal.media import audio as audio_module from vllm.multimodal.media.audio import ( @@ -16,6 +19,7 @@ from vllm.multimodal.media.audio import ( load_audio_soundfile, load_audio_torchcodec, ) +from vllm.multimodal.media.base import MediaDecodedBudget from ...conftest import AudioTestAssets @@ -118,6 +122,102 @@ def _make_flac_bytes(frames: int, channels: int, samplerate: int) -> bytes: return buf.getvalue() +@pytest.mark.skip_global_cleanup +def test_request_budget_rejects_second_decoded_audio_clip(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 384) + clip = _make_flac_bytes(frames=64, channels=1, samplerate=16000) + audio_io = AudioMediaIO(decoded_budget=MediaDecodedBudget()) + + audio_io.load_bytes(clip) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + audio_io.load_bytes(clip) + + +@pytest.mark.skip_global_cleanup +def test_request_budget_allows_parallel_audio_decodes(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 512) + monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_DECODE_BYTES", 256) + budget = MediaDecodedBudget() + audio_io = AudioMediaIO(decoded_budget=budget) + decoding = threading.Barrier(2) + + def decode(*args, **kwargs): + assert kwargs["max_decode_bytes"] == 256 + decoding.wait(timeout=5) + return np.zeros(32, dtype=np.float32), 16000 + + with ( + patch.object(audio_module, "load_audio", side_effect=decode), + ThreadPoolExecutor(max_workers=2) as pool, + ): + results = list(pool.map(audio_io.load_bytes, [b"a", b"b"])) + + assert len(results) == 2 + assert budget.used == 256 + + +@pytest.mark.skip_global_cleanup +def test_request_budget_waits_for_inflight_audio_decode(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 256) + monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_DECODE_BYTES", 256) + budget = MediaDecodedBudget() + audio_io = AudioMediaIO(decoded_budget=budget) + decoding = threading.Event() + finish_decoding = threading.Event() + second_reservation = threading.Event() + reserve = budget.reserve + decode_calls = 0 + + def track_reservation(max_bytes): + if decoding.is_set(): + second_reservation.set() + return reserve(max_bytes) + + def decode(*args, **kwargs): + nonlocal decode_calls + decode_calls += 1 + decoding.set() + assert finish_decoding.wait(timeout=5) + return np.zeros(64, dtype=np.float32), 16000 + + with ( + patch.object(budget, "reserve", side_effect=track_reservation), + patch.object(audio_module, "load_audio", side_effect=decode), + ThreadPoolExecutor(max_workers=2) as pool, + ): + first = pool.submit(audio_io.load_bytes, b"a") + assert decoding.wait(timeout=5) + second = pool.submit(audio_io.load_bytes, b"b") + try: + assert second_reservation.wait(timeout=5) + finally: + finish_decoding.set() + first.result(timeout=5) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + second.result(timeout=5) + + assert decode_calls == 1 + assert budget.used == 256 + + +@pytest.mark.skip_global_cleanup +def test_request_budget_releases_failed_audio_reservation(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 256) + monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_DECODE_BYTES", 256) + budget = MediaDecodedBudget() + audio_io = AudioMediaIO(decoded_budget=budget) + audio = (np.zeros(64, dtype=np.float32), 16000) + + with patch.object( + audio_module, "load_audio", side_effect=[ValueError("invalid audio"), audio] + ): + with pytest.raises(ValueError, match="invalid audio"): + audio_io.load_bytes(b"a") + assert audio_io.load_bytes(b"b")[0].nbytes == 256 + + assert budget.used == 256 + + def test_small_file_passes_memory_guard(): """A small valid file should pass both duration and memory guards.""" payload = _make_flac_bytes(frames=16000, channels=1, samplerate=16000) diff --git a/tests/multimodal/media/test_image.py b/tests/multimodal/media/test_image.py index 94efb312f2..60cfb86cc0 100644 --- a/tests/multimodal/media/test_image.py +++ b/tests/multimodal/media/test_image.py @@ -11,9 +11,9 @@ from vllm.multimodal.image import ( ALLOWED_IMAGE_FORMATS, open_image, ) -from vllm.multimodal.media import ImageMediaIO +from vllm.multimodal.media import ImageMediaIO, MediaConnector, MediaDecodedBudget -pytestmark = pytest.mark.cpu_test +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] ASSETS_DIR = Path(__file__).parent.parent / "assets" assert ASSETS_DIR.exists() @@ -288,6 +288,68 @@ def test_image_pixel_limit_disabled(monkeypatch): assert result.media.size == (1000, 1000) +def _png_bytes(width: int, height: int) -> bytes: + from io import BytesIO + + buf = BytesIO() + Image.new("RGB", (width, height), (1, 2, 3)).save(buf, format="PNG") + return buf.getvalue() + + +def test_request_decoded_budget_rejects_aggregate(monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 600) + budget = MediaDecodedBudget() + image_io = ImageMediaIO(decoded_budget=budget) + data = _png_bytes(10, 10) + + image_io.load_bytes(data) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + image_io.load_bytes(data) + + +def test_request_decoded_budget_accepts_images_within_limit(monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 800) + budget = MediaDecodedBudget() + image_io = ImageMediaIO(decoded_budget=budget) + data = _png_bytes(10, 10) + + image_io.load_bytes(data) + image_io.load_bytes(data) + assert budget.used == 800 + + +def test_request_decoded_budget_disabled(monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 0) + budget = MediaDecodedBudget() + image_io = ImageMediaIO(decoded_budget=budget) + data = _png_bytes(100, 100) + + image_io.load_bytes(data) + image_io.load_bytes(data) + assert budget.used == 0 + + +def test_media_connector_enforces_request_decoded_budget(monkeypatch): + import pybase64 + + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 600) + b64 = pybase64.b64encode(_png_bytes(10, 10)).decode() + url = f"data:image/png;base64,{b64}" + connector = MediaConnector(decoded_budget=MediaDecodedBudget()) + + connector.fetch_image(url) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + connector.fetch_image(url) + + def test_image_media_io_rejects_postscript_payload(): """EPS/PostScript bytes must not reach an external renderer.""" import pybase64 as base64 diff --git a/tests/multimodal/media/test_unprocessable_entity_error.py b/tests/multimodal/media/test_unprocessable_entity_error.py index 6d3c2a0b0a..8d7142d019 100644 --- a/tests/multimodal/media/test_unprocessable_entity_error.py +++ b/tests/multimodal/media/test_unprocessable_entity_error.py @@ -13,11 +13,16 @@ from unittest.mock import AsyncMock, MagicMock, patch import aiohttp import pytest -from vllm.connections import MediaDownloadSizeExceededError +from vllm.connections import ( + MediaDownloadSizeExceededError, + MediaRequestDownloadSizeExceededError, +) from vllm.entrypoints.serve import create_error_response from vllm.exceptions import VLLMClientError, VLLMUnprocessableEntityError from vllm.multimodal.media import MediaConnector +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + class TestVLLMUnprocessableEntityError: """Tests for VLLMUnprocessableEntityError exception.""" @@ -148,6 +153,21 @@ class TestMediaConnectorErrorHandling: assert "VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB=256" in str(exc_info.value) + def test_fetch_image_request_download_size_limit_error(self): + connector = MediaConnector() + + with patch.object( + connector.connection, "get_bytes", new_callable=MagicMock + ) as mock_get: + mock_get.side_effect = MediaRequestDownloadSizeExceededError( + "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB=256" + ) + + with pytest.raises(VLLMUnprocessableEntityError) as exc_info: + connector.fetch_image("https://example.com/image.jpg") + + assert "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB=256" in str(exc_info.value) + class TestErrorResponse: """Tests for error response creation.""" diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index a2fedb1747..e3e2ae6689 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -9,8 +9,10 @@ import numpy as np import numpy.typing as npt import pybase64 import pytest +import torch from PIL import Image +import vllm.envs as envs from vllm.assets.base import get_vllm_public_assets from vllm.assets.video import ( video_get_metadata, @@ -18,6 +20,7 @@ from vllm.assets.video import ( video_to_pil_images_list, ) from vllm.multimodal.media import ImageMediaIO, VideoMediaIO +from vllm.multimodal.media.base import MediaDecodedBudget from vllm.multimodal.video import ( PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, @@ -47,6 +50,18 @@ class Assert10Frames1FPSVideoLoader(VideoLoader): return FAKE_OUTPUT_2 +@pytest.mark.skip_global_cleanup +def test_video_request_budget_counts_torch_frames(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 16) + budget = MediaDecodedBudget() + video_io = VideoMediaIO(ImageMediaIO(decoded_budget=budget)) + frames = torch.zeros((2, 2, 2, 3), dtype=torch.uint8) + video_io.video_loader = SimpleNamespace(load_bytes=lambda *_args, **_kwargs: frames) + + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + video_io.load_bytes(b"video") + + def test_video_media_io_kwargs(monkeypatch: pytest.MonkeyPatch): with monkeypatch.context() as m: m.setenv("VLLM_VIDEO_LOADER_BACKEND", "assert_10_frames_1_fps") diff --git a/tests/test_audio_media_size_precheck.py b/tests/test_audio_media_size_precheck.py index 7635022369..e2f3c57a0f 100644 --- a/tests/test_audio_media_size_precheck.py +++ b/tests/test_audio_media_size_precheck.py @@ -2,15 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import AsyncIterator, Iterator +from io import BytesIO from pathlib import Path from unittest.mock import patch +import numpy as np import pytest +from PIL import Image import vllm.envs as envs -from vllm.connections import HTTPConnection +from vllm.connections import HTTPConnection, MediaRequestDownloadBudget from vllm.exceptions import VLLMValidationError -from vllm.multimodal.media import AudioMediaIO, MediaConnector +from vllm.multimodal.media import AudioMediaIO, ImageMediaIO, MediaConnector +from vllm.multimodal.media.base import MediaDecodedBudget from vllm.utils.mem_constants import KiB_bytes, MiB_bytes pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] @@ -87,6 +91,46 @@ def test_audio_load_bytes_rejects_oversized(monkeypatch: pytest.MonkeyPatch): AudioMediaIO().load_bytes(b"\x00" * (MiB_bytes + 1)) +def test_request_audio_budget_limits_multiple_clips_before_decode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 24) + budget = MediaDecodedBudget() + clip = np.zeros(4, dtype=np.float32) + + def decode(_source, **kwargs): + if kwargs["max_decode_bytes"] < clip.nbytes: + raise ValueError("VLLM_MAX_AUDIO_DECODE_BYTES") + return clip, 16000 + + with patch("vllm.multimodal.media.audio.load_audio", side_effect=decode) as load: + audio_io = AudioMediaIO(decoded_budget=budget) + audio_io.load_bytes(b"first") + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + audio_io.load_bytes(b"second") + + assert budget.used == clip.nbytes + assert load.call_count == 2 + + +def test_request_decoded_budget_is_shared_by_image_and_audio( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 24) + budget = MediaDecodedBudget() + image_bytes = BytesIO() + Image.new("RGB", (2, 3)).save(image_bytes, format="PNG") + ImageMediaIO(decoded_budget=budget).load_bytes(image_bytes.getvalue()) + + with ( + patch("vllm.multimodal.media.audio.load_audio") as decode, + pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"), + ): + AudioMediaIO(decoded_budget=budget).load_bytes(b"audio") + + decode.assert_not_called() + + def test_audio_load_file_rejects_oversized( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ): @@ -138,12 +182,16 @@ async def test_audio_connector_threads_byte_limit_to_http_reader( monkeypatch: pytest.MonkeyPatch, ): monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", 1) + budget = MediaRequestDownloadBudget() class _Connection: async def async_get_bytes(self, _url: str, **kwargs): assert kwargs["max_bytes"] == MiB_bytes + assert kwargs["request_budget"] is budget raise VLLMValidationError("Maximum file size exceeded") - connector = MediaConnector(connection=_Connection()) + connector = MediaConnector( # type: ignore[arg-type] + connection=_Connection(), download_budget=budget + ) with pytest.raises(VLLMValidationError, match="Maximum file size exceeded"): await connector.fetch_audio_async("https://example.test/audio") diff --git a/tests/test_connections.py b/tests/test_connections.py index 983e20f6c1..6385e372e4 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -9,7 +9,13 @@ from unittest.mock import AsyncMock, patch import pytest import vllm.envs as envs -from vllm.connections import HTTPConnection, HTTPResponseSizeExceededError +from vllm.connections import ( + HTTPConnection, + HTTPResponseSizeExceededError, + MediaRequestDownloadBudget, +) + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] _ONE_MIB = 1024 * 1024 _SMALL_BODY = b"a" * (_ONE_MIB // 2) @@ -150,6 +156,49 @@ async def test_async_get_bytes_with_limit_streams_instead_of_materializing() -> assert response.content.iterated == 2 +def test_request_download_budget_counts_across_responses( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", 1) + budget = MediaRequestDownloadBudget() + connection = HTTPConnection() + first = _SyncResponse([b"a" * (3 * _ONE_MIB // 4)]) + second = _SyncResponse([b"b" * (_ONE_MIB // 2)]) + + with patch.object(connection, "get_response", side_effect=[first, second]): + connection.get_bytes("http://example.com/one", request_budget=budget) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB"): + connection.get_bytes("http://example.com/two", request_budget=budget) + + assert budget.used == 3 * _ONE_MIB // 4 + assert second.iterated == 1 + + +@pytest.mark.asyncio +async def test_async_request_download_budget_counts_across_responses( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", 1) + budget = MediaRequestDownloadBudget() + connection = HTTPConnection() + first = _AsyncResponse([b"a" * (3 * _ONE_MIB // 4)]) + second = _AsyncResponse([b"b" * (_ONE_MIB // 2)]) + + with patch.object( + connection, "get_async_response", new=AsyncMock(side_effect=[first, second]) + ): + await connection.async_get_bytes( + "http://example.com/one", request_budget=budget + ) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB"): + await connection.async_get_bytes( + "http://example.com/two", request_budget=budget + ) + + assert budget.used == 3 * _ONE_MIB // 4 + assert second.content.iterated == 1 + + def test_get_bytes_rejects_content_length_over_limit( monkeypatch: pytest.MonkeyPatch, local_media_server: str ) -> None: diff --git a/vllm/connections.py b/vllm/connections.py index b06d2d563f..ded8da430f 100644 --- a/vllm/connections.py +++ b/vllm/connections.py @@ -3,6 +3,7 @@ import asyncio import functools +import threading import time from collections.abc import Callable, Coroutine, Mapping, MutableMapping from pathlib import Path @@ -60,6 +61,33 @@ class MediaDownloadSizeExceededError(ValueError): super().__init__(message) +class MediaRequestDownloadSizeExceededError(ValueError): + """Raised when all HTTP media downloads in one request exceed the limit.""" + + +class MediaRequestDownloadBudget: + """Thread-safe total of HTTP media bytes received for one request.""" + + def __init__(self) -> None: + self.used = 0 + self._lock = threading.Lock() + + def charge(self, size: int) -> None: + limit_mb = envs.VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB + if limit_mb <= 0: + return + limit = limit_mb * MiB_bytes + with self._lock: + received = self.used + size + if received > limit: + raise MediaRequestDownloadSizeExceededError( + "Request media downloads exceed " + f"VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB={limit_mb:g} " + f"({received} bytes received)." + ) + self.used = received + + def _get_media_download_limit_bytes() -> int | None: max_size_mb = envs.VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB if max_size_mb <= 0: @@ -105,19 +133,26 @@ def _append_chunk( received_bytes: int, max_bytes: int | None, media_download_limit: bool, + request_budget: MediaRequestDownloadBudget | None = None, ) -> int: if not chunk: return received_bytes received_bytes += len(chunk) if max_bytes is not None and received_bytes > max_bytes: _raise_size_exceeded(max_bytes, received_bytes, media_download_limit) + if request_budget is not None: + request_budget.charge(len(chunk)) chunks.append(chunk) return received_bytes -def _read_response_bytes(response: requests.Response, max_bytes: int | None) -> bytes: +def _read_response_bytes( + response: requests.Response, + max_bytes: int | None, + request_budget: MediaRequestDownloadBudget | None = None, +) -> bytes: max_bytes, media_download_limit = _get_response_size_limit(max_bytes) - if max_bytes is None: + if max_bytes is None and request_budget is None: return response.content _check_content_length( @@ -129,16 +164,23 @@ def _read_response_bytes(response: requests.Response, max_bytes: int | None) -> received_bytes = 0 for chunk in response.iter_content(_RESPONSE_READ_CHUNK_SIZE): received_bytes = _append_chunk( - chunks, chunk, received_bytes, max_bytes, media_download_limit + chunks, + chunk, + received_bytes, + max_bytes, + media_download_limit, + request_budget, ) return b"".join(chunks) async def _async_read_response_bytes( - response: aiohttp.ClientResponse, max_bytes: int | None + response: aiohttp.ClientResponse, + max_bytes: int | None, + request_budget: MediaRequestDownloadBudget | None = None, ) -> bytes: max_bytes, media_download_limit = _get_response_size_limit(max_bytes) - if max_bytes is None: + if max_bytes is None and request_budget is None: return await response.read() _check_content_length(response.content_length, max_bytes, media_download_limit) @@ -146,7 +188,12 @@ async def _async_read_response_bytes( received_bytes = 0 async for chunk in response.content.iter_chunked(_RESPONSE_READ_CHUNK_SIZE): received_bytes = _append_chunk( - chunks, chunk, received_bytes, max_bytes, media_download_limit + chunks, + chunk, + received_bytes, + max_bytes, + media_download_limit, + request_budget, ) return b"".join(chunks) @@ -407,6 +454,7 @@ class HTTPConnection: timeout: float | None = None, allow_redirects: bool = True, max_bytes: int | None = None, + request_budget: MediaRequestDownloadBudget | None = None, ) -> bytes: with self.get_response( url, @@ -416,7 +464,7 @@ class HTTPConnection: ) as r: r.raise_for_status() - return _read_response_bytes(r, max_bytes) + return _read_response_bytes(r, max_bytes, request_budget) @_async_retry async def async_get_bytes( @@ -426,6 +474,7 @@ class HTTPConnection: timeout: float | None = None, allow_redirects: bool = True, max_bytes: int | None = None, + request_budget: MediaRequestDownloadBudget | None = None, ) -> bytes: async with await self.get_async_response( url, @@ -434,7 +483,7 @@ class HTTPConnection: ) as r: r.raise_for_status() - return await _async_read_response_bytes(r, max_bytes) + return await _async_read_response_bytes(r, max_bytes, request_budget) def get_text(self, url: str, *, timeout: float | None = None) -> str: with self.get_response(url, timeout=timeout) as r: diff --git a/vllm/envs.py b/vllm/envs.py index bb1c03a9da..dbbc49ff59 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -78,6 +78,7 @@ if TYPE_CHECKING: VLLM_MEDIA_CACHE_TTL_HOURS: float = 24 VLLM_MEDIA_FETCH_MAX_RETRIES: int = 3 VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB: int = 256 + VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB: int = 256 VLLM_MEDIA_URL_ALLOW_REDIRECTS: bool = True VLLM_MEDIA_LOADING_THREAD_COUNT: int = 8 VLLM_MAX_AUDIO_CLIP_FILESIZE_MB: int = 25 @@ -86,6 +87,7 @@ if TYPE_CHECKING: VLLM_MAX_AUDIO_PREPROCESS_WORKERS: int = max(1, min(os.cpu_count() or 1, 2)) VLLM_MAX_EMBED_DECODE_BYTES: int = 2_147_483_648 VLLM_MAX_IMAGE_PIXELS: int = 178_956_970 + VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES: int = 715_827_880 VLLM_VIDEO_LOADER_BACKEND: str = "opencv" VLLM_MEDIA_CONNECTOR: str = "http" VLLM_TARGET_DEVICE: str = "cuda" @@ -1063,6 +1065,15 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_MAX_IMAGE_PIXELS": lambda: int( os.getenv("VLLM_MAX_IMAGE_PIXELS", "178956970") ), + # Maximum aggregate decoded bytes across media in one request. + # The default allows one image at the per-image pixel limit at 4 bytes + # per pixel. Set to 0 to disable. + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES": lambda: int( + os.getenv("VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", "715827880") + ), + "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB": lambda: int( + os.getenv("VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", "256") + ), # Backend for Video IO — selects the frame-sampling algorithm. # - "opencv": uniform sampling. # - "opencv_dynamic": duration-aware dynamic sampling. @@ -2412,6 +2423,8 @@ def compile_factors() -> dict[str, object]: "VLLM_MAX_AUDIO_PREPROCESS_WORKERS", "VLLM_MAX_EMBED_DECODE_BYTES", "VLLM_MAX_IMAGE_PIXELS", + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", + "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", "VLLM_VIDEO_LOADER_BACKEND", "VLLM_MEDIA_CONNECTOR", "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME", diff --git a/vllm/multimodal/media/__init__.py b/vllm/multimodal/media/__init__.py index 94b08c4848..84fc7d6ed8 100644 --- a/vllm/multimodal/media/__init__.py +++ b/vllm/multimodal/media/__init__.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from .audio import AudioEmbeddingMediaIO, AudioMediaIO -from .base import MediaIO, MediaWithBytes +from .base import MediaDecodedBudget, MediaIO, MediaWithBytes from .connector import MEDIA_CONNECTOR_REGISTRY, MediaConnector from .image import ImageEmbeddingMediaIO, ImageMediaIO from .video import VIDEO_LOADER_REGISTRY, VideoMediaIO __all__ = [ "MediaIO", + "MediaDecodedBudget", "MediaWithBytes", "AudioEmbeddingMediaIO", "AudioMediaIO", diff --git a/vllm/multimodal/media/audio.py b/vllm/multimodal/media/audio.py index c05f4e4dcc..d36b4a984e 100644 --- a/vllm/multimodal/media/audio.py +++ b/vllm/multimodal/media/audio.py @@ -23,7 +23,7 @@ from vllm.utils.sparse_utils import ( safe_to_dense, ) -from .base import MediaIO +from .base import MediaDecodedBudget, MediaIO logger = init_logger(__name__) @@ -527,8 +527,11 @@ class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]): error handling. """ - def __init__(self, **kwargs) -> None: + def __init__( + self, *, decoded_budget: MediaDecodedBudget | None = None, **kwargs + ) -> None: super().__init__() + self.decoded_budget = decoded_budget # Allow per-server override of the audio decoding backend, e.g.: # --media-io-kwargs '{"audio": {"audio_backend": "torchcodec"}}' @@ -564,13 +567,44 @@ class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]): def load_bytes(self, data: bytes) -> tuple[npt.NDArray, float]: self._validate_encoded_size(len(data)) - return load_audio( - BytesIO(data), - sr=None, - max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, - max_decode_bytes=envs.VLLM_MAX_AUDIO_DECODE_BYTES, - backend=self.audio_backend, - ) + return self._load_with_budget(BytesIO(data)) + + def _load_with_budget(self, source: BytesIO | Path) -> tuple[npt.NDArray, float]: + budget = self.decoded_budget + if budget is None: + return load_audio( + source, + sr=None, + max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, + max_decode_bytes=envs.VLLM_MAX_AUDIO_DECODE_BYTES, + backend=self.audio_backend, + ) + + max_decode_bytes = envs.VLLM_MAX_AUDIO_DECODE_BYTES + reserved = budget.reserve(max_decode_bytes) + decoded_bytes = 0 + try: + try: + audio = load_audio( + source, + sr=None, + max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, + max_decode_bytes=reserved, + backend=self.audio_backend, + ) + except ValueError as exc: + if reserved < max_decode_bytes and "VLLM_MAX_AUDIO_DECODE_BYTES" in str( + exc + ): + raise ValueError( + "Request decoded media exceeds " + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES." + ) from exc + raise + decoded_bytes = audio[0].nbytes + return audio + finally: + budget.settle(reserved, decoded_bytes) def load_base64( self, @@ -588,13 +622,7 @@ class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]): def load_file(self, filepath: Path) -> tuple[npt.NDArray, float]: self._validate_encoded_size(filepath.stat().st_size) - return load_audio( - filepath, - sr=None, - max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, - max_decode_bytes=envs.VLLM_MAX_AUDIO_DECODE_BYTES, - backend=self.audio_backend, - ) + return self._load_with_budget(filepath) def encode_base64( self, diff --git a/vllm/multimodal/media/base.py b/vllm/multimodal/media/base.py index d3e622156f..952b92e969 100644 --- a/vllm/multimodal/media/base.py +++ b/vllm/multimodal/media/base.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading from abc import ABC, abstractmethod from collections.abc import Iterable, Iterator from dataclasses import dataclass, field @@ -9,9 +10,79 @@ from typing import Any, Generic, TypeVar, cast import numpy as np +import vllm.envs as envs + _T = TypeVar("_T") +class MediaDecodedBudget: + """Thread-safe decoded-media byte accumulator for one request.""" + + def __init__(self) -> None: + self.used = 0 + self._reserved = 0 + self._condition = threading.Condition() + + def reserve(self, max_bytes: int) -> int: + """Reserve decoder capacity, waiting for in-flight decodes if needed.""" + limit = envs.VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES + if limit <= 0: + return max_bytes + max_bytes = max(0, max_bytes) + with self._condition: + while True: + available = limit - self.used - self._reserved + if available >= max_bytes or self._reserved == 0: + if available <= 0: + raise ValueError( + "Request decoded media exceeds " + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES." + ) + reserved = min(max_bytes, available) + self._reserved += reserved + return reserved + self._condition.wait() + + def settle(self, reserved: int, decoded_bytes: int) -> None: + """Replace a decoder reservation with its actual decoded size.""" + limit = envs.VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES + if limit <= 0: + return + with self._condition: + self._reserved -= reserved + try: + used = self.used + decoded_bytes + if decoded_bytes > reserved or used > limit: + raise ValueError( + "Request decoded media exceeds " + f"VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES={limit} " + f"({used} bytes)." + ) + self.used = used + finally: + self._condition.notify_all() + + def charge(self, decoded_bytes: int) -> None: + limit = envs.VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES + if limit <= 0 or decoded_bytes <= 0: + return + with self._condition: + while ( + self.used + self._reserved + decoded_bytes > limit + and self._reserved > 0 + and self.used + decoded_bytes <= limit + ): + self._condition.wait() + used = self.used + decoded_bytes + if used > limit: + raise ValueError( + "Request decoded media exceeds " + f"VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES={limit} " + f"({used} bytes)." + ) + self.used = used + + @dataclass class MediaWithBytes(Generic[_T]): """Wrapper that couples a media object with its original encoded bytes. diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index 297dc91af2..e5a52417fe 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -25,6 +25,8 @@ import vllm.envs as envs from vllm.connections import ( HTTPConnection, MediaDownloadSizeExceededError, + MediaRequestDownloadBudget, + MediaRequestDownloadSizeExceededError, global_http_connection, ) from vllm.exceptions import VLLMUnprocessableEntityError, VLLMValidationError @@ -33,7 +35,7 @@ from vllm.multimodal.video import get_video_loader_backend_for_processor from vllm.utils.registry import ExtensionManager from .audio import AudioEmbeddingMediaIO, AudioMediaIO -from .base import MediaIO, MediaWithBytes +from .base import MediaDecodedBudget, MediaIO, MediaWithBytes from .image import ImageEmbeddingMediaIO, ImageMediaIO from .video import VideoEmbeddingMediaIO, VideoMediaIO @@ -109,7 +111,9 @@ def _wrap_media_fetch_error( value=url, ) - if isinstance(exc, MediaDownloadSizeExceededError): + if isinstance( + exc, (MediaDownloadSizeExceededError, MediaRequestDownloadSizeExceededError) + ): return VLLMUnprocessableEntityError( f"Failed to fetch media from URL: {exc}", parameter="image_url", @@ -162,6 +166,8 @@ class MediaConnector: *, allowed_local_media_path: str = "", allowed_media_domains: list[str] | None = None, + decoded_budget: MediaDecodedBudget | None = None, + download_budget: MediaRequestDownloadBudget | None = None, ) -> None: """Args: media_io_kwargs: Additional args passed to process media @@ -201,6 +207,8 @@ class MediaConnector: if allowed_media_domains is None: allowed_media_domains = [] self.allowed_media_domains = allowed_media_domains + self.decoded_budget = decoded_budget + self.download_budget = download_budget # Media download cache (opt-in via VLLM_MEDIA_CACHE) self._media_cache_dir: str | None = None @@ -230,6 +238,12 @@ class MediaConnector: media_cache, ) + def _image_io(self, image_mode: str | None = "RGB") -> ImageMediaIO: + return ImageMediaIO( + **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})), + decoded_budget=self.decoded_budget, + ) + def _get_cached_bytes(self, url: str) -> bytes | None: """Return cached bytes for a URL, or None if not cached/expired.""" if not self._media_cache_dir: @@ -385,12 +399,18 @@ class MediaConnector: return media_io.load_bytes(cached) connection = self.connection + budget_kwargs = ( + {"request_budget": self.download_budget} + if self.download_budget is not None + else {} + ) try: data = connection.get_bytes( url_spec.url, timeout=fetch_timeout, allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, max_bytes=max_bytes, + **budget_kwargs, ) except Exception as e: wrapped = _wrap_media_fetch_error(url, e) @@ -438,12 +458,18 @@ class MediaConnector: return await future connection = self.connection + budget_kwargs = ( + {"request_budget": self.download_budget} + if self.download_budget is not None + else {} + ) try: data = await connection.async_get_bytes( url_spec.url, timeout=fetch_timeout, allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, max_bytes=max_bytes, + **budget_kwargs, ) except Exception as e: wrapped = _wrap_media_fetch_error(url, e) @@ -470,7 +496,9 @@ class MediaConnector: audio_url: str, ) -> tuple[np.ndarray, int | float]: """Load audio from a URL.""" - audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {})) + audio_io = AudioMediaIO( + decoded_budget=self.decoded_budget, **self.media_io_kwargs.get("audio", {}) + ) return self.load_from_url( audio_url, @@ -483,7 +511,9 @@ class MediaConnector: audio_url: str, ) -> tuple[np.ndarray, int | float]: """Asynchronously fetch audio from a URL.""" - audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {})) + audio_io = AudioMediaIO( + decoded_budget=self.decoded_budget, **self.media_io_kwargs.get("audio", {}) + ) return await self.load_from_url_async( audio_url, @@ -503,9 +533,7 @@ class MediaConnector: `media_io_kwargs={"image": {"image_mode": None}}` to keep the original image mode (e.g. preserving the alpha channel). """ - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) try: return self.load_from_url( @@ -529,9 +557,7 @@ class MediaConnector: `media_io_kwargs={"image": {"image_mode": None}}` to keep the original image mode (e.g. preserving the alpha channel). """ - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) try: return await self.load_from_url_async( @@ -551,9 +577,7 @@ class MediaConnector: video_processor: str | None = None, ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: """Load video from an HTTP or base64 data URL.""" - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) if "video_backend" not in video_io_kwargs and ( video_backend := get_video_loader_backend_for_processor(video_processor) @@ -580,9 +604,7 @@ class MediaConnector: `media_io_kwargs={"image": {"image_mode": None}}` to keep the original image mode (e.g. preserving the alpha channel). """ - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) if "video_backend" not in video_io_kwargs and ( video_backend := get_video_loader_backend_for_processor(video_processor) diff --git a/vllm/multimodal/media/image.py b/vllm/multimodal/media/image.py index fb05f4ed74..1e2e106189 100644 --- a/vllm/multimodal/media/image.py +++ b/vllm/multimodal/media/image.py @@ -17,7 +17,7 @@ from vllm.utils.sparse_utils import ( ) from ..image import convert_image_mode, normalize_image, open_image, rgba_to_rgb -from .base import MediaIO, MediaWithBytes +from .base import MediaDecodedBudget, MediaIO, MediaWithBytes MAGIC_NUMPY_PREFIX = b"\x93NUMPY" # https://numpy.org/devdocs/reference/generated/numpy.lib.format.html#format-version-1-0 @@ -28,12 +28,19 @@ class ImageMediaIO(MediaIO[Image.Image]): error handling. """ - def __init__(self, image_mode: str | None = "RGB", **kwargs) -> None: + def __init__( + self, + image_mode: str | None = "RGB", + *, + decoded_budget: MediaDecodedBudget | None = None, + **kwargs, + ) -> None: super().__init__() # Target mode for loaded images; `None` keeps the original mode # (i.e. no conversion, alpha channel is preserved as-is). self.image_mode = image_mode + self.decoded_budget = decoded_budget # `kwargs` contains custom arguments from # --media-io-kwargs for this modality, merged with # per-request runtime media_io_kwargs via merge_kwargs(). @@ -87,6 +94,8 @@ class ImageMediaIO(MediaIO[Image.Image]): f"the maximum of {max_pixels} pixels. Set " f"VLLM_MAX_IMAGE_PIXELS to increase this limit." ) + if self.decoded_budget is not None: + self.decoded_budget.charge(w * h * 4) image = normalize_image(image) image.load() converted = self._convert_image_mode(image) diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index 7528331cdf..569ef3dc93 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -122,6 +122,16 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]] video = self.video_loader.load_bytes( data, num_frames=self.num_frames, **self.kwargs ) + frames = video[0] if isinstance(video, tuple) else video + budget = getattr(self.image_io, "decoded_budget", None) + shape = getattr(frames, "shape", None) + if budget is not None and shape is not None and len(shape) >= 3: + decoded_bytes = ( + frames.numel() * frames.element_size() + if isinstance(frames, torch.Tensor) + else frames.nbytes + ) + budget.charge(decoded_bytes) return MediaWithBytes(video, data) def load_base64( diff --git a/vllm/renderers/chat_utils.py b/vllm/renderers/chat_utils.py index f18407e535..540564baab 100644 --- a/vllm/renderers/chat_utils.py +++ b/vllm/renderers/chat_utils.py @@ -53,6 +53,7 @@ from typing_extensions import TypedDict, override from vllm import envs from vllm.config import ModelConfig +from vllm.connections import MediaRequestDownloadBudget from vllm.exceptions import VLLMValidationError from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict from vllm.logger import init_logger @@ -67,6 +68,7 @@ from vllm.multimodal.inputs import ( VisionChunkVideo, ) from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY, MediaConnector +from vllm.multimodal.media.base import MediaDecodedBudget from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.renderers.embed_utils import ( safe_load_prompt_embeds, @@ -609,6 +611,8 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]): self._items_by_modality = defaultdict[str, list[_T]](list) # Track original modality for each vision_chunk item (image or video) self._modality_order = defaultdict[str, list[str]](list) + self.decoded_budget = MediaDecodedBudget() + self.download_budget = MediaRequestDownloadBudget() @cached_property def use_unified_vision_chunk_modality(self) -> bool: @@ -1052,6 +1056,8 @@ class MultiModalContentParser(BaseMultiModalContentParser): media_io_kwargs=self._tracker.media_io_kwargs, allowed_local_media_path=self._tracker.allowed_local_media_path, allowed_media_domains=self._tracker.allowed_media_domains, + decoded_budget=self._tracker.decoded_budget, + download_budget=self._tracker.download_budget, ) @property @@ -1138,6 +1144,10 @@ class MultiModalContentParser(BaseMultiModalContentParser): def parse_image_pil( self, image_pil: Image.Image | None, uuid: str | None = None ) -> None: + if image_pil is not None: + self._tracker.decoded_budget.charge( + image_pil.size[0] * image_pil.size[1] * 4 + ) placeholder = self._tracker.add("image", (image_pil, uuid)) self._add_placeholder("image", placeholder) @@ -1233,6 +1243,8 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser): media_io_kwargs=self._tracker.media_io_kwargs, allowed_local_media_path=self._tracker.allowed_local_media_path, allowed_media_domains=self._tracker.allowed_media_domains, + decoded_budget=self._tracker.decoded_budget, + download_budget=self._tracker.download_budget, ) @property @@ -1351,6 +1363,10 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser): image_pil: Image.Image | None, uuid: str | None = None, ) -> None: + if image_pil is not None: + self._tracker.decoded_budget.charge( + image_pil.size[0] * image_pil.size[1] * 4 + ) placeholder = self._tracker.add( "image", partial(self._item_with_uuid_async, image_pil, uuid) )