diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index 99e58fce2..2b9ed6ec0 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -398,6 +398,12 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout async def fake_omni_init_app_state(engine_client, state, args): state.engine_client = engine_client state.initialized_by_omni = True + state.openai_serving_video = SimpleNamespace( + shutdown=lambda: captured.__setitem__("video_shutdown", True), + ) + state.openai_serving_speech = SimpleNamespace( + shutdown=lambda: captured.__setitem__("speech_shutdown", True), + ) async def fake_storage_start(): captured["storage_started"] = True @@ -433,6 +439,8 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout assert captured["supported_tasks"] == ("generate",) assert captured["storage_started"] is True assert captured["restrictions"] == {} + assert captured["video_shutdown"] is True + assert captured["speech_shutdown"] is True assert sock.closed is True assert served_app.state.initialized_by_omni is True diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index a7334c177..80deb445b 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -3,6 +3,7 @@ """Unit tests for OpenAI-compatible video API encoding helpers.""" import base64 +import threading from io import BytesIO import av @@ -249,6 +250,95 @@ def test_channel_first_video_tensor_uses_direct_planar_mux(monkeypatch): assert options["fps"] == 12.0 +def test_planar_converter_serial_and_parallel_outputs_match(): + frames = [np.full((2, 3, 3), fill_value=value / 10, dtype=np.float32) for value in (1, 4, 7)] + + serial_converter = video_api_utils._PlanarFrameConverter(max_workers=1) + parallel_converter = video_api_utils._PlanarFrameConverter(max_workers=2) + try: + serial = [ + frame.to_ndarray(format="rgb24") for frame in serial_converter.iter_frames(frames, np.dtype(np.float32)) + ] + parallel = [ + frame.to_ndarray(format="rgb24") for frame in parallel_converter.iter_frames(frames, np.dtype(np.float32)) + ] + finally: + serial_converter.shutdown() + parallel_converter.shutdown() + + np.testing.assert_array_equal(np.stack(serial), np.stack(parallel)) + + +def test_planar_converter_preserves_fifo_order_and_bounds_pending_futures(monkeypatch): + converter = video_api_utils._PlanarFrameConverter(max_workers=2) + started = threading.Event() + release = threading.Event() + four_submitted = threading.Event() + outstanding = 0 + peak_outstanding = 0 + lock = threading.Lock() + real_submit = converter._executor.submit + + def fake_build(frame, _common_dtype): + started.set() + release.wait(timeout=2) + return int(frame[0, 0, 0]) + + def submit(*args, **kwargs): + nonlocal outstanding, peak_outstanding + with lock: + outstanding += 1 + peak_outstanding = max(peak_outstanding, outstanding) + if peak_outstanding >= 4: + four_submitted.set() + future = real_submit(*args, **kwargs) + + def done(_future): + nonlocal outstanding + with lock: + outstanding -= 1 + + future.add_done_callback(done) + return future + + monkeypatch.setattr(converter, "_build_frame", fake_build) + monkeypatch.setattr(converter._executor, "submit", submit) + frames = [np.full((1, 1, 3), value, dtype=np.uint8) for value in range(8)] + result: list[int] = [] + iterator = converter.iter_frames(frames, np.dtype(np.uint8)) + consumer = threading.Thread(target=lambda: result.extend(iterator)) + consumer.start() + try: + assert started.wait(timeout=2) + assert four_submitted.wait(timeout=2) + release.set() + consumer.join(timeout=2) + assert not consumer.is_alive() + finally: + release.set() + iterator.close() + converter.shutdown() + + assert result == list(range(8)) + assert peak_outstanding <= 4 + + +def test_planar_converter_reuses_executor_and_shutdowns_after_cancel(): + converter = video_api_utils._PlanarFrameConverter(max_workers=2) + executor = converter._executor + frames = [np.zeros((2, 2, 3), dtype=np.uint8) for _ in range(4)] + try: + list(converter.iter_frames(frames, np.dtype(np.uint8))) + iterator = converter.iter_frames(frames, np.dtype(np.uint8)) + next(iterator) + iterator.close() + assert converter._executor is executor + finally: + converter.shutdown() + + assert executor._shutdown + + @pytest.mark.parametrize( ("video", "expected_path", "expected_reason", "expected_mux"), [ @@ -347,6 +437,35 @@ def test_video_only_does_not_log_default_audio_sample_rate(monkeypatch): assert log_messages == [] +def test_planar_mux_accepts_none_rate_and_defaults_audio_to_44100(): + def make_frame(): + return next( + video_api_utils._iter_planar_video_frames( + [np.zeros((2, 2, 3), dtype=np.uint8)], + np.dtype(np.uint8), + ) + ) + + video_only = media_utils.mux_av_video_audio_bytes( + [make_frame()], + width=2, + height=2, + audio_sample_rate=None, + ) + with av.open(BytesIO(video_only), mode="r", format="mp4") as container: + assert not container.streams.audio + + with_audio = media_utils.mux_av_video_audio_bytes( + [make_frame()], + width=2, + height=2, + audio_waveform=np.zeros((1, 1024), dtype=np.float32), + audio_sample_rate=None, + ) + with av.open(BytesIO(with_audio), mode="r", format="mp4") as container: + assert container.streams.audio[0].rate == 44100 + + def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch): calls = [] path_logs = [] @@ -383,6 +502,7 @@ def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch): assert "reason=non_contiguous_rgb_planes" in path_logs[0] assert "audio_present=False" in path_logs[0] assert "effective_audio_sample_rate=None" in path_logs[0] + assert "effective_frame_conversion_workers=0" in path_logs[0] assert len(coerce_calls) == 1 diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index b48029c30..6d8977c55 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -95,12 +95,13 @@ class FakeAsyncOmni: yield MockVideoResult(videos) -def test_raw_and_base64_encoders_receive_no_policy_config(mocker: MockerFixture): +def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixture): engine = FakeAsyncOmni() handler = OmniOpenAIServingVideo.for_diffusion( engine, model_name="test-model", ) + assert handler._video_frame_converter.max_workers == 8 raw_encoder = mocker.patch( "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", return_value=b"encoded-video", @@ -117,8 +118,9 @@ def test_raw_and_base64_encoders_receive_no_policy_config(mocker: MockerFixture) asyncio.run(_generate_both_response_types()) - assert "encoding_config" not in raw_encoder.call_args.kwargs - assert "encoding_config" not in base64_encoder.call_args.kwargs + assert raw_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter + assert base64_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter + handler.shutdown() def test_resolve_diffusion_od_config_falls_back_to_attribute(): @@ -1231,8 +1233,9 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi audio=None, audio_sample_rate=None, video_codec_options=None, + frame_converter=None, ): - del video, fps, audio, video_codec_options + del video, fps, audio, video_codec_options, frame_converter audio_sample_rates.append(audio_sample_rate) return b"fake-video" diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index 5c210e996..df67f9d99 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -210,7 +210,7 @@ def mux_av_video_audio_bytes( audio_waveform: np.ndarray | None = None, *, fps: float = 25.0, - audio_sample_rate: int = 44100, + audio_sample_rate: int | None = None, video_codec: str = "h264", audio_codec: str = "aac", crf: str = "18", @@ -236,6 +236,7 @@ def mux_av_video_audio_bytes( samples: np.ndarray | None = None layout: str | None = None if audio_waveform is not None: + effective_audio_sample_rate = 44100 if audio_sample_rate is None else audio_sample_rate samples = audio_waveform.astype(np.float32) if samples.ndim == 1: samples = samples.reshape(1, -1) @@ -243,7 +244,7 @@ def mux_av_video_audio_bytes( samples = np.ascontiguousarray(samples.T) num_channels = samples.shape[0] layout = "stereo" if num_channels >= 2 else "mono" - a_stream = cast(av.AudioStream, container.add_stream(audio_codec, rate=audio_sample_rate)) + a_stream = cast(av.AudioStream, container.add_stream(audio_codec, rate=effective_audio_sample_rate)) a_stream.layout = layout for frame in video_frames: @@ -256,12 +257,12 @@ def mux_av_video_audio_bytes( if samples is None or layout is None: raise ValueError("Audio samples were not prepared for muxing.") audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout) - audio_frame.sample_rate = audio_sample_rate + audio_frame.sample_rate = effective_audio_sample_rate # AAC has a one-frame encoder delay. Mark the input waveform as # starting at t=0 so the MP4 muxer writes the corresponding negative # priming timestamp instead of exposing the delay as leading silence. audio_frame.pts = 0 - audio_frame.time_base = Fraction(1, audio_sample_rate) + audio_frame.time_base = Fraction(1, effective_audio_sample_rate) for packet in a_stream.encode(audio_frame): container.mux(packet) for packet in a_stream.encode(): diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index cdc36e588..cee98a3ac 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -615,6 +615,9 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, await shutdown_task finally: state = getattr(app, "state", None) + serving_video = getattr(state, "openai_serving_video", None) if state is not None else None + if serving_video is not None: + serving_video.shutdown() serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None if serving_speech is not None: serving_speech.shutdown() diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index e8cbcbb48..398c7b81f 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -32,6 +32,7 @@ from vllm_omni.entrypoints.openai.stage_params import ( from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline, parse_lora_request from vllm_omni.entrypoints.openai.video_api_utils import ( _encode_video_bytes, + _PlanarFrameConverter, encode_video_base64, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt @@ -44,6 +45,8 @@ from vllm_omni.outputs.output_metadata import ( logger = init_logger(__name__) +_VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS = 8 + if TYPE_CHECKING: from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -96,6 +99,11 @@ class OmniOpenAIServingVideo: self._engine_client = engine_client self._model_name = model_name self._stage_configs = stage_configs + self._video_frame_converter = _PlanarFrameConverter(max_workers=_VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS) + logger.info( + "Video response frame conversion pool configured: workers=%d", + self._video_frame_converter.max_workers, + ) def _resolve_diffusion_od_config(self) -> OmniDiffusionConfig | SimpleNamespace | None: get_od_config = getattr(self._engine_client, "get_diffusion_od_config", None) @@ -170,6 +178,9 @@ class OmniOpenAIServingVideo: stage_configs=stage_configs, ) + def shutdown(self) -> None: + self._video_frame_converter.shutdown() + async def _run_and_extract( self, request: VideoGenerationRequest, @@ -361,6 +372,7 @@ class OmniOpenAIServingVideo: video, fps=artifacts.output_fps, video_codec_options=video_codec_options, + frame_converter=self._video_frame_converter, ) if artifacts.audios[idx] is None else encode_video_base64( @@ -369,6 +381,7 @@ class OmniOpenAIServingVideo: audio=artifacts.audios[idx], audio_sample_rate=artifacts.audio_sample_rate, video_codec_options=video_codec_options, + frame_converter=self._video_frame_converter, ) ), action=artifacts.actions[idx], @@ -425,6 +438,7 @@ class OmniOpenAIServingVideo: fps=artifacts.output_fps, **({"audio": audio, "audio_sample_rate": artifacts.audio_sample_rate} if audio is not None else {}), video_codec_options=video_codec_options, + frame_converter=self._video_frame_converter, ) _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000 logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms) diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 53d087af4..81c7540ba 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -10,9 +10,12 @@ import base64 import binascii import os import tempfile -from collections.abc import Iterator +import threading +from collections import deque +from collections.abc import Generator +from concurrent.futures import Future, ThreadPoolExecutor from io import BytesIO -from typing import TYPE_CHECKING, Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal import httpx import numpy as np @@ -576,6 +579,112 @@ def _direct_planar_fallback_reason( return None +class _PlanarFrameConverter: + """Convert response frames with one reusable, bounded executor.""" + + def __init__(self, max_workers: int) -> None: + self._max_workers = max_workers + if self._max_workers < 1: + raise ValueError("max_workers must be positive") + self._executor = ( + ThreadPoolExecutor(max_workers=self._max_workers, thread_name_prefix="video-planar") + if self._max_workers > 1 + else None + ) + self._scratch_local = threading.local() + self._shutdown = False + + @property + def max_workers(self) -> int: + return self._max_workers + + def _channel_scratch(self, height: int, width: int, common_dtype: np.dtype) -> np.ndarray: + scratch_dtype = np.dtype(common_dtype) + if not np.issubdtype(scratch_dtype, np.floating): + scratch_dtype = np.dtype(np.float64) + scratch = getattr(self._scratch_local, "channel", None) + if scratch is None or scratch.shape != (height, width) or scratch.dtype != scratch_dtype: + scratch = np.empty((height, width), dtype=scratch_dtype) + self._scratch_local.channel = scratch + return scratch + + def _build_frame(self, frame: np.ndarray, common_dtype: np.dtype) -> av.VideoFrame: + """Build one quantized GBR PyAV frame.""" + import av + + height, width = frame.shape[:2] + scratch = None if frame.dtype == np.uint8 else self._channel_scratch(height, width, common_dtype) + av_frame = av.VideoFrame(width, height, format="gbrp") + for plane, channel in zip(av_frame.planes, (1, 2, 0)): + if plane.height < height or plane.line_size < width: + raise ValueError("PyAV video plane is smaller than the requested frame dimensions.") + plane_view = np.frombuffer( + memoryview(plane), + dtype=np.uint8, + count=plane.height * plane.line_size, + ).reshape(plane.height, plane.line_size) + plane_view.fill(0) + if frame.dtype == np.uint8: + plane_view[:height, :width] = frame[..., channel] + else: + np.copyto(scratch, frame[..., channel], casting="unsafe") + np.clip(scratch, 0.0, 1.0, out=scratch) + scratch *= 255.0 + np.rint(scratch, out=scratch) + plane_view[:height, :width] = scratch + return av_frame + + def iter_frames( + self, + frames: list[np.ndarray], + common_dtype: np.dtype, + ) -> Generator[av.VideoFrame, None, None]: + """Yield converted frames in order with at most two batches in flight.""" + effective_workers = min(len(frames), self._max_workers) + if effective_workers <= 1: + for frame in frames: + yield self._build_frame(frame, common_dtype) + return + + executor = self._executor + assert executor is not None + pending: deque[Future[av.VideoFrame]] = deque() + next_index = 0 + max_pending = 2 * effective_workers + try: + while next_index < len(frames) and len(pending) < max_pending: + pending.append(executor.submit(self._build_frame, frames[next_index], common_dtype)) + next_index += 1 + + while pending: + converted_frame = pending.popleft().result() + if next_index < len(frames): + pending.append(executor.submit(self._build_frame, frames[next_index], common_dtype)) + next_index += 1 + yield converted_frame + finally: + for future in pending: + future.cancel() + + def shutdown(self) -> None: + if not self._shutdown: + self._shutdown = True + if self._executor is not None: + self._executor.shutdown(wait=True, cancel_futures=True) + + +def _iter_planar_video_frames( + frames: list[np.ndarray], + common_dtype: np.dtype, +) -> Generator[av.VideoFrame, None, None]: + """Yield serial planar frames for direct utility callers.""" + converter = _PlanarFrameConverter(max_workers=1) + try: + yield from converter.iter_frames(frames, common_dtype) + finally: + converter.shutdown() + + def _log_video_encoding_path( *, selected_path: str, @@ -585,12 +694,13 @@ def _log_video_encoding_path( fps: int, audio: AudioInput | None, audio_sample_rate: int | None, + effective_frame_conversion_workers: int, reason: str | None = None, ) -> None: reason_field = "" if reason is None else f" reason={reason}" logger.info( "Video response encoding route selected: selected_path=%s%s frames=%d frame_shape=%s dtype=%s fps=%s " - "audio_present=%s effective_audio_sample_rate=%s", + "audio_present=%s effective_audio_sample_rate=%s effective_frame_conversion_workers=%s", selected_path, reason_field, len(frames), @@ -599,6 +709,7 @@ def _log_video_encoding_path( fps, audio is not None, audio_sample_rate, + effective_frame_conversion_workers, ) @@ -611,40 +722,6 @@ def _resolve_audio_sample_rate(audio: AudioInput | None, audio_sample_rate: int return audio_sample_rate or DEFAULT_AUDIO_SAMPLE_RATE -def _iter_planar_video_frames( - frames: list[np.ndarray], - common_dtype: np.dtype, -) -> Iterator[av.VideoFrame]: - """Yield planar PyAV frames while retaining only one channel scratch buffer.""" - import av - - height, width = frames[0].shape[:2] - scratch_dtype = np.float64 if np.issubdtype(common_dtype, np.bool_) else common_dtype - scratch = None if common_dtype == np.uint8 else np.empty((height, width), dtype=scratch_dtype) - - for frame in frames: - av_frame = av.VideoFrame(width, height, format="gbrp") - for plane, channel in zip(av_frame.planes, (1, 2, 0)): - if plane.height < height or plane.line_size < width: - raise ValueError("PyAV video plane is smaller than the requested frame dimensions.") - plane_view = np.frombuffer( - plane, - dtype=np.uint8, - count=plane.height * plane.line_size, - ).reshape(plane.height, plane.line_size) - plane_view.fill(0) - if frame.dtype == np.uint8: - plane_view[:height, :width] = frame[..., channel] - else: - scratch_buffer = cast(np.ndarray, scratch) - np.copyto(scratch_buffer, frame[..., channel], casting="unsafe") - np.clip(scratch_buffer, 0.0, 1.0, out=scratch_buffer) - scratch_buffer *= 255.0 - np.rint(scratch_buffer, out=scratch_buffer) - plane_view[:height, :width] = scratch_buffer - yield av_frame - - def _encode_prepared_video_bytes_legacy( frames: list[np.ndarray], frame_shape: tuple[int, ...], @@ -693,8 +770,9 @@ def _encode_video_bytes( audio: Any | None = None, audio_sample_rate: int | None = None, video_codec_options: dict[str, str] | None = None, + frame_converter: _PlanarFrameConverter | None = None, ) -> bytes: - """Encode a video payload through the automatic capability dispatcher.""" + """Encode a video payload through the direct planar or legacy path.""" from vllm_omni.diffusion.utils.media_utils import mux_av_video_audio_bytes # Prepare once so validation is shared by both paths and malformed common @@ -712,6 +790,7 @@ def _encode_video_bytes( fps=fps, audio=audio, audio_sample_rate=effective_audio_sample_rate, + effective_frame_conversion_workers=0, ) return _encode_prepared_video_bytes_legacy( frames, @@ -723,6 +802,8 @@ def _encode_video_bytes( video_codec_options=video_codec_options, ) + owns_frame_converter = frame_converter is None + converter = frame_converter or _PlanarFrameConverter(max_workers=1) _log_video_encoding_path( selected_path="direct_planar", frames=frames, @@ -731,17 +812,24 @@ def _encode_video_bytes( fps=fps, audio=audio, audio_sample_rate=effective_audio_sample_rate, + effective_frame_conversion_workers=min(len(frames), converter.max_workers), ) audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None - return mux_av_video_audio_bytes( - _iter_planar_video_frames(frames, common_dtype), - width=frame_shape[1], - height=frame_shape[0], - audio_waveform=audio_np, - fps=float(fps), - audio_sample_rate=effective_audio_sample_rate, - video_codec_options=video_codec_options, - ) + video_frames = converter.iter_frames(frames, common_dtype) + try: + return mux_av_video_audio_bytes( + video_frames, + width=frame_shape[1], + height=frame_shape[0], + audio_waveform=audio_np, + fps=float(fps), + audio_sample_rate=effective_audio_sample_rate, + video_codec_options=video_codec_options, + ) + finally: + video_frames.close() + if owns_frame_converter: + converter.shutdown() class FragmentedMP4VideoEncoder: @@ -799,6 +887,7 @@ def encode_video_base64( audio: Any | None = None, audio_sample_rate: int | None = None, video_codec_options: dict[str, str] | None = None, + frame_converter: _PlanarFrameConverter | None = None, ) -> str: """Encode a video (frames/array/tensor) to base64 MP4.""" video_bytes = _encode_video_bytes( @@ -807,5 +896,6 @@ def encode_video_base64( audio=audio, audio_sample_rate=audio_sample_rate, video_codec_options=video_codec_options, + frame_converter=frame_converter, ) return base64.b64encode(video_bytes).decode("utf-8")