From e284d907b5554038bef6ab861bdb78b0b99dd5ea Mon Sep 17 00:00:00 2001 From: Qi Jia Date: Sun, 13 Sep 2026 20:33:10 +0800 Subject: [PATCH] [Perf][Frontend] Chunk in-memory image file responses (#7459) Signed-off-by: QI JIA Co-authored-by: QI JIA Co-authored-by: Cursor --- .../benchmark_image_file_response.py | 202 ++++++++++++++++++ .../openai_api/test_image_protocol.py | 75 +++++++ .../entrypoints/openai/protocol/images.py | 14 +- 3 files changed, 289 insertions(+), 2 deletions(-) create mode 100644 tests/benchmarks/benchmark_image_file_response.py create mode 100644 tests/entrypoints/openai_api/test_image_protocol.py diff --git a/tests/benchmarks/benchmark_image_file_response.py b/tests/benchmarks/benchmark_image_file_response.py new file mode 100644 index 000000000..acd8967a6 --- /dev/null +++ b/tests/benchmarks/benchmark_image_file_response.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Benchmark newline-split and fixed-size chunking of image file responses. + +Both variants stream the exact payload `ImageGenerationResponse.stream_response()` +produces, so the measurement isolates the response layer: PNG encoding and ZIP +building are shared by the variants and excluded from the timings. +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import hashlib +import io +import json +import statistics +import time +from collections.abc import AsyncIterator, Callable, Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import cast + +import numpy as np +from PIL import Image +from starlette.responses import StreamingResponse + +from vllm_omni.entrypoints.openai.protocol.images import ( + _FILE_RESPONSE_CHUNK_SIZE, + ImageData, + ImageGenerationResponse, + _iter_file_chunks, +) + + +@dataclass(frozen=True) +class BenchmarkVariant: + label: str + content: Callable[[bytes], Iterable[bytes] | AsyncIterator[memoryview]] + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("value must be at least 1") + return parsed + + +def _encode_png(*, width: int, height: int, seed: int) -> bytes: + """Encode incompressible noise, so the payload looks like a real photo.""" + rng = np.random.default_rng(seed) + pixels = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8) + buffer = io.BytesIO() + Image.fromarray(pixels, mode="RGB").save(buffer, format="png") + return buffer.getvalue() + + +async def _drain(response: StreamingResponse) -> tuple[bytes, int]: + """Run a response through ASGI, returning its body and non-empty frame count.""" + body = bytearray() + frames = 0 + + async def send(message: dict[str, object]) -> None: + nonlocal frames + if message["type"] == "http.response.body" and message["body"]: + frames += 1 + body.extend(cast(bytes, message["body"])) + + async def receive() -> dict[str, str]: + return {"type": "http.disconnect"} + + scope = {"type": "http", "method": "GET", "asgi": {"spec_version": "2.4"}} + await response(scope, receive, send) + return bytes(body), frames + + +async def _build_payload(*, images: int, width: int, height: int, seed: int) -> bytes: + png = _encode_png(width=width, height=height, seed=seed) + response = ImageGenerationResponse( + created=0, + data=[ImageData(b64_json=base64.b64encode(png).decode()) for _ in range(images)], + output_format="png", + size=f"{width}x{height}", + ) + payload, _ = await _drain(response.stream_response()) + return payload + + +async def _measure(variant: BenchmarkVariant, payload: bytes) -> dict[str, object]: + response = StreamingResponse(variant.content(payload), media_type="application/octet-stream") + cpu_start = time.process_time_ns() + wall_start = time.perf_counter_ns() + body, frames = await _drain(response) + wall_ms = (time.perf_counter_ns() - wall_start) / 1_000_000 + process_cpu_ms = (time.process_time_ns() - cpu_start) / 1_000_000 + return { + "label": variant.label, + "wall_ms": wall_ms, + "process_cpu_ms": process_cpu_ms, + "body_frames": frames, + "output_bytes": len(body), + "output_sha256": hashlib.sha256(body).hexdigest(), + } + + +def _summarize(records: list[dict[str, object]], label: str) -> dict[str, object]: + selected = [record for record in records if record["label"] == label] + wall_values = [cast(float, record["wall_ms"]) for record in selected] + cpu_values = [cast(float, record["process_cpu_ms"]) for record in selected] + return { + "runs": len(selected), + "body_frames": cast(int, selected[0]["body_frames"]), + "wall_ms": { + "median": statistics.median(wall_values), + "min": min(wall_values), + "max": max(wall_values), + }, + "process_cpu_ms": { + "median": statistics.median(cpu_values), + "min": min(cpu_values), + "max": max(cpu_values), + }, + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--images", type=_positive_int, default=1, help="2 or more streams a ZIP archive") + parser.add_argument("--width", type=_positive_int, default=1024) + parser.add_argument("--height", type=_positive_int, default=1024) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--rounds", type=_positive_int, default=3) + parser.add_argument("--seed", type=int, default=20260912) + parser.add_argument("--output-json", type=Path) + args = parser.parse_args() + if args.warmup < 0: + parser.error("--warmup must be non-negative") + return args + + +async def _run(args: argparse.Namespace) -> dict[str, object]: + payload = await _build_payload(images=args.images, width=args.width, height=args.height, seed=args.seed) + baseline = BenchmarkVariant("newline_split", io.BytesIO) + candidate = BenchmarkVariant("fixed_size_chunks", _iter_file_chunks) + + for _ in range(args.warmup): + for variant in (baseline, candidate): + await _measure(variant, payload) + + records: list[dict[str, object]] = [] + for round_index in range(args.rounds): + order = (baseline, candidate) if round_index % 2 == 0 else (candidate, baseline) + for variant in order: + record = await _measure(variant, payload) + record["round"] = round_index + 1 + records.append(record) + + output_hashes = {str(record["output_sha256"]) for record in records} + if len(output_hashes) != 1: + raise RuntimeError(f"benchmark variants produced different outputs: {sorted(output_hashes)}") + + baseline_summary = _summarize(records, baseline.label) + candidate_summary = _summarize(records, candidate.label) + baseline_ms = cast(float, cast(dict[str, object], baseline_summary["wall_ms"])["median"]) + candidate_ms = cast(float, cast(dict[str, object], candidate_summary["wall_ms"])["median"]) + return { + "config": { + "images": args.images, + "width": args.width, + "height": args.height, + "warmup": args.warmup, + "rounds": args.rounds, + "seed": args.seed, + "chunk_size": _FILE_RESPONSE_CHUNK_SIZE, + "payload_bytes": len(payload), + "payload_newline_bytes": payload.count(b"\n"), + }, + "records": records, + "summary": { + baseline.label: baseline_summary, + candidate.label: candidate_summary, + "median_wall_speedup": baseline_ms / candidate_ms, + "body_frame_ratio": ( + cast(int, baseline_summary["body_frames"]) / cast(int, candidate_summary["body_frames"]) + ), + "output_sha256": output_hashes.pop(), + }, + } + + +def main() -> None: + args = _parse_args() + result = asyncio.run(_run(args)) + output_json = json.dumps(result, indent=2, sort_keys=True) + if args.output_json is not None: + args.output_json.write_text(output_json + "\n", encoding="utf-8") + print(output_json) + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoints/openai_api/test_image_protocol.py b/tests/entrypoints/openai_api/test_image_protocol.py new file mode 100644 index 000000000..b77038933 --- /dev/null +++ b/tests/entrypoints/openai_api/test_image_protocol.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import base64 +import io +import zipfile + +import pytest +from fastapi.responses import StreamingResponse + +from vllm_omni.entrypoints.openai.protocol.images import ( + _FILE_RESPONSE_CHUNK_SIZE, + ImageData, + ImageGenerationResponse, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_NEWLINE_HEAVY_PAYLOAD = (b"\x89PNG\n\x1a\n" * 64) + bytes(range(256)) * 8 + + +def _response(*payloads: bytes) -> ImageGenerationResponse: + return ImageGenerationResponse( + created=0, + data=[ImageData(b64_json=base64.b64encode(payload).decode()) for payload in payloads], + output_format="png", + size="1x1", + ) + + +async def _collect(response: StreamingResponse) -> list[bytes]: + return [bytes(chunk) async for chunk in response.body_iterator] + + +@pytest.mark.asyncio +async def test_single_file_response_chunks_ignore_newlines() -> None: + payload = _NEWLINE_HEAVY_PAYLOAD + assert b"\n" in payload + + response = _response(payload).stream_response() + chunks = await _collect(response) + + assert len(chunks) == 1 + assert b"".join(chunks) == payload + assert response.headers["content-length"] == str(len(payload)) + assert response.headers["content-type"] == "image/png" + + +@pytest.mark.asyncio +async def test_single_file_response_chunks_are_bounded() -> None: + payload = _NEWLINE_HEAVY_PAYLOAD * 200 + expected_chunks = -(-len(payload) // _FILE_RESPONSE_CHUNK_SIZE) + assert expected_chunks > 1 + + chunks = await _collect(_response(payload).stream_response()) + + assert len(chunks) == expected_chunks + assert max(len(chunk) for chunk in chunks) <= _FILE_RESPONSE_CHUNK_SIZE + assert b"".join(chunks) == payload + + +@pytest.mark.asyncio +async def test_zip_file_response_streams_whole_archive() -> None: + payloads = (_NEWLINE_HEAVY_PAYLOAD, b"second image\nwith a newline") + + response = _response(*payloads).stream_response() + chunks = await _collect(response) + + zip_bytes = b"".join(chunks) + assert len(chunks) == -(-len(zip_bytes) // _FILE_RESPONSE_CHUNK_SIZE) + assert response.headers["content-length"] == str(len(zip_bytes)) + assert response.headers["content-type"] == "application/zip" + with zipfile.ZipFile(io.BytesIO(zip_bytes)) as archive: + assert archive.read("image_0.png") == payloads[0] + assert archive.read("image_1.png") == payloads[1] diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index b7800c13e..067098f78 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -11,6 +11,7 @@ import base64 import io import uuid import zipfile +from collections.abc import AsyncIterator from enum import Enum from http import HTTPStatus from typing import Any, Literal @@ -25,6 +26,8 @@ from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers _INT64_MIN = -(2**63) _INT64_MAX = 2**63 - 1 +_FILE_RESPONSE_CHUNK_SIZE = 64 * 1024 + _IMAGE_FILE_METADATA = { "jpg": ("jpg", "image/jpeg"), "jpeg": ("jpeg", "image/jpeg"), @@ -33,6 +36,13 @@ _IMAGE_FILE_METADATA = { } +async def _iter_file_chunks(data: bytes) -> AsyncIterator[memoryview]: + """Yield fixed-size chunks of an in-memory payload.""" + view = memoryview(data) + for offset in range(0, len(view), _FILE_RESPONSE_CHUNK_SIZE): + yield view[offset : offset + _FILE_RESPONSE_CHUNK_SIZE] + + class ResponseFormat(str, Enum): """Image response format""" @@ -235,7 +245,7 @@ class ImageGenerationResponse(BaseModel): image_bytes = base64.b64decode(self.data[0].b64_json) filename = f"image_{uuid.uuid4().hex[:8]}.{extension}" return StreamingResponse( - io.BytesIO(image_bytes), + _iter_file_chunks(image_bytes), media_type=media_type, headers={ "Content-Disposition": f'attachment; filename="{filename}"', @@ -251,7 +261,7 @@ class ImageGenerationResponse(BaseModel): zip_bytes = zip_buffer.getvalue() filename = f"images_{uuid.uuid4().hex[:8]}.zip" return StreamingResponse( - io.BytesIO(zip_bytes), + _iter_file_chunks(zip_bytes), media_type="application/zip", headers={ "Content-Disposition": f'attachment; filename="{filename}"',