[Frontend] Encode interleaved video frames in parallel (#6776)

Signed-off-by: mokashliu <mokashliu@tencent.com>
Co-authored-by: mokashliu <mokashliu@tencent.com>

tests/benchmarks/benchmark_video_response_encoding.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Benchmark legacy and automatic MP4 response encoding on synthetic input."""
 
 from __future__ import annotations
@@ -11,13 +11,16 @@ import statistics
 import time
 from collections.abc import Callable
 from dataclasses import dataclass
+from functools import partial
 from pathlib import Path
+from typing import cast
 
 import numpy as np
 
 from vllm_omni.entrypoints.openai.video_api_utils import (
     _encode_video_bytes,
     _encode_video_bytes_legacy,
+    _PlanarFrameConverter,
 )
 
 
@@ -42,12 +45,19 @@ def _build_inputs(
     fps: int,
     audio_sample_rate: int,
     seed: int,
+    input_layout: str = "planar-backed",
 ) -> tuple[np.ndarray, np.ndarray]:
     rng = np.random.default_rng(seed)
     planar = rng.random((3, frames, height, width), dtype=np.float32)
     video = planar.transpose(1, 2, 3, 0)
-    if not all(video[..., channel].flags.c_contiguous for channel in range(3)):
+    if input_layout == "interleaved":
+        video = np.ascontiguousarray(video)
+        if not video.flags.c_contiguous:
+            raise RuntimeError("synthetic video is not C-contiguous")
+    elif input_layout == "planar-backed" and not all(video[..., channel].flags.c_contiguous for channel in range(3)):
         raise RuntimeError("synthetic video does not expose contiguous channel planes")
+    elif input_layout != "planar-backed":
+        raise ValueError(f"unsupported input layout: {input_layout}")
 
     sample_count = max(1, round(frames / fps * audio_sample_rate))
     timeline = np.arange(sample_count, dtype=np.float32) / audio_sample_rate
@@ -86,8 +96,8 @@ def _measure(
 
 def _summarize(records: list[dict[str, object]], label: str) -> dict[str, object]:
     selected = [record for record in records if record["label"] == label]
-    wall_values = [float(record["wall_ms"]) for record in selected]
-    cpu_values = [float(record["process_cpu_ms"]) for record in selected]
+    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),
         "wall_ms": {
@@ -113,6 +123,16 @@ def _parse_args() -> argparse.Namespace:
     parser.add_argument("--warmup", type=int, default=1)
     parser.add_argument("--rounds", type=_positive_int, default=3)
     parser.add_argument("--seed", type=int, default=20260817)
+    parser.add_argument(
+        "--input-layout",
+        choices=("planar-backed", "interleaved"),
+        default="planar-backed",
+    )
+    parser.add_argument(
+        "--frame-conversion-workers",
+        type=_positive_int,
+        default=8,
+    )
     parser.add_argument("--output-json", type=Path)
     args = parser.parse_args()
     if args.warmup < 0:
@@ -129,33 +149,41 @@ def main() -> None:
         fps=args.fps,
         audio_sample_rate=args.audio_sample_rate,
         seed=args.seed,
+        input_layout=args.input_layout,
     )
     baseline = BenchmarkVariant("legacy", _encode_video_bytes_legacy)
-    candidate = BenchmarkVariant("automatic", _encode_video_bytes)
+    converter = _PlanarFrameConverter(max_workers=args.frame_conversion_workers)
+    candidate = BenchmarkVariant(
+        "automatic",
+        partial(_encode_video_bytes, frame_converter=converter),
+    )
 
-    for _ in range(args.warmup):
-        for variant in (baseline, candidate):
-            _measure(
-                variant,
-                video=video,
-                audio=audio,
-                fps=args.fps,
-                audio_sample_rate=args.audio_sample_rate,
-            )
+    try:
+        for _ in range(args.warmup):
+            for variant in (baseline, candidate):
+                _measure(
+                    variant,
+                    video=video,
+                    audio=audio,
+                    fps=args.fps,
+                    audio_sample_rate=args.audio_sample_rate,
+                )
 
-    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 = _measure(
-                variant,
-                video=video,
-                audio=audio,
-                fps=args.fps,
-                audio_sample_rate=args.audio_sample_rate,
-            )
-            record["round"] = round_index + 1
-            records.append(record)
+        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 = _measure(
+                    variant,
+                    video=video,
+                    audio=audio,
+                    fps=args.fps,
+                    audio_sample_rate=args.audio_sample_rate,
+                )
+                record["round"] = round_index + 1
+                records.append(record)
+    finally:
+        converter.shutdown()
 
     output_hashes = {str(record["output_sha256"]) for record in records}
     if len(output_hashes) != 1:
@@ -165,8 +193,8 @@ def main() -> None:
     paired_savings_percent = []
     for round_index in range(args.rounds):
         pair = [record for record in records if record["round"] == round_index + 1]
-        baseline_ms = float(next(record["wall_ms"] for record in pair if record["label"] == baseline.label))
-        candidate_ms = float(next(record["wall_ms"] for record in pair if record["label"] == candidate.label))
+        baseline_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == baseline.label))
+        candidate_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == candidate.label))
         paired_savings_ms.append(baseline_ms - candidate_ms)
         paired_savings_percent.append((baseline_ms - candidate_ms) / baseline_ms * 100)
 
@@ -180,6 +208,8 @@ def main() -> None:
             "warmup": args.warmup,
             "rounds": args.rounds,
             "seed": args.seed,
+            "input_layout": args.input_layout,
+            "frame_conversion_workers": args.frame_conversion_workers,
             "video_shape": list(video.shape),
             "video_strides": list(video.strides),
             "codec_options": {"preset": "ultrafast", "threads": "0"},

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for OpenAI-compatible video API encoding helpers."""
 
 import base64
 import threading
 from io import BytesIO
+from typing import Any
 
 import av
 import httpx
@@ -41,7 +42,7 @@ def _install_http_transport(monkeypatch, handler, client_kwargs):
 @pytest.mark.asyncio
 async def test_decode_image_url_follows_redirects_when_allowed(monkeypatch):
     requested_paths = []
-    client_kwargs = []
+    client_kwargs: list[dict[str, object]] = []
 
     def _handler(request):
         requested_paths.append(request.url.path)
@@ -63,7 +64,7 @@ async def test_decode_image_url_follows_redirects_when_allowed(monkeypatch):
 @pytest.mark.asyncio
 async def test_decode_image_url_rejects_redirects_when_disabled(monkeypatch):
     requested_paths = []
-    client_kwargs = []
+    client_kwargs: list[dict[str, object]] = []
 
     def _handler(request):
         requested_paths.append(request.url.path)
@@ -81,7 +82,7 @@ async def test_decode_image_url_rejects_redirects_when_disabled(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_decode_image_url_reports_http_status(monkeypatch):
-    client_kwargs = []
+    client_kwargs: list[dict[str, object]] = []
 
     def _handler(request):
         return httpx.Response(404)
@@ -95,7 +96,7 @@ async def test_decode_image_url_reports_http_status(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_decode_image_url_reports_connection_failure(monkeypatch):
-    client_kwargs = []
+    client_kwargs: list[dict[str, object]] = []
 
     def _handler(request):
         raise httpx.ConnectError("connection refused", request=request)
@@ -141,7 +142,7 @@ def _install_fake_video_mux(monkeypatch, mux_calls):
 
 
 def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
-    mux_calls = []
+    mux_calls: list[dict[str, Any]] = []
     _install_fake_video_mux(monkeypatch, mux_calls)
 
     frames = [np.full((2, 2, 3), fill_value=i / 5, dtype=np.float32) for i in range(5)]
@@ -549,6 +550,44 @@ def test_prepared_automatic_fallback_preserves_legacy_quantization(monkeypatch,
     np.testing.assert_array_equal(mux_inputs[0], reference)
 
 
+@pytest.mark.parametrize(
+    ("worker_count", "expected_path", "expected_reason"),
+    [
+        (None, "legacy_fallback", "non_contiguous_rgb_planes"),
+        (1, "legacy_fallback", "non_contiguous_rgb_planes"),
+        (8, "direct_planar", None),
+    ],
+)
+def test_interleaved_video_routing(monkeypatch, worker_count, expected_path, expected_reason):
+    routes = []
+
+    def fake_planar_mux(*args, **kwargs):
+        return b"planar-video"
+
+    def fake_compat_mux(frames, audio, **kwargs):
+        return b"legacy-video"
+
+    monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fake_planar_mux)
+    monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux)
+    monkeypatch.setattr(video_api_utils, "_log_video_encoding_path", lambda **kwargs: routes.append(kwargs))
+
+    converter = None if worker_count is None else video_api_utils._PlanarFrameConverter(max_workers=worker_count)
+    try:
+        encoded = video_api_utils._encode_video_bytes(
+            np.zeros((2, 4, 6, 3), dtype=np.float32),
+            fps=12,
+            frame_converter=converter,
+        )
+    finally:
+        if converter is not None:
+            converter.shutdown()
+
+    assert encoded == (b"planar-video" if expected_path == "direct_planar" else b"legacy-video")
+    assert len(routes) == 1
+    assert routes[0]["selected_path"] == expected_path
+    assert routes[0].get("reason") == expected_reason
+
+
 @pytest.mark.parametrize(
     ("video", "reason"),
     [
@@ -690,7 +729,8 @@ def test_planar_bool_frames_match_bounded_compatible_output():
 )
 def test_planar_frame_rejects_undersized_plane(monkeypatch, plane_height, line_size):
     class FakePlane:
-        pass
+        height: int
+        line_size: int
 
     plane = FakePlane()
     plane.height = plane_height
@@ -758,6 +798,32 @@ def test_direct_and_compatible_paths_produce_equivalent_mp4(with_audio):
     np.testing.assert_array_equal(direct_frames, compatible_frames)
 
 
+@pytest.mark.parametrize("with_audio", [False, True])
+def test_parallel_interleaved_and_legacy_paths_produce_equivalent_mp4(with_audio):
+    rng = np.random.default_rng(7)
+    video = rng.random((3, 32, 32, 3), dtype=np.float32)
+    audio = np.zeros((2, 2400), dtype=np.float32) if with_audio else None
+    kwargs = {
+        "fps": 12,
+        "audio": audio,
+        "audio_sample_rate": 24000,
+        "video_codec_options": {"preset": "ultrafast", "threads": "1"},
+    }
+
+    converter = video_api_utils._PlanarFrameConverter(max_workers=8)
+    try:
+        parallel_bytes = video_api_utils._encode_video_bytes(
+            video,
+            **kwargs,
+            frame_converter=converter,
+        )
+    finally:
+        converter.shutdown()
+    legacy_bytes = video_api_utils._encode_video_bytes_legacy(video, **kwargs)
+
+    assert parallel_bytes == legacy_bytes
+
+
 def test_mux_closes_container_and_preserves_generator_error(monkeypatch):
     class GeneratorError(RuntimeError):
         pass