diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index bc80fe0db..b534d6c67 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -1,14 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for the serving-layer streaming video WebSocket handler.""" from __future__ import annotations import asyncio import base64 +import gc import io import json import threading +import weakref from typing import Any import pytest @@ -38,10 +40,10 @@ def _b64(data: bytes) -> str: def _text_result(text: str) -> OmniRequestOutput: class Output: - pass + text: str class RequestOutput: - pass + outputs: list[Output] output = Output() output.text = text @@ -52,10 +54,10 @@ def _text_result(text: str) -> OmniRequestOutput: def _audio_result(audio_data: Any) -> OmniRequestOutput: class Output: - pass + multimodal_output: dict[str, Any] class RequestOutput: - pass + outputs: list[Output] output = Output() output.multimodal_output = {"audio": audio_data} @@ -513,6 +515,88 @@ async def test_new_query_cancels_in_flight_query(): assert "session.done" in ws.sent_types() +@pytest.mark.asyncio +@pytest.mark.parametrize("delay_abort", [False, True], ids=["immediate-ack", "delayed-ack"]) +async def test_interrupted_queries_wait_for_abort_without_fixed_delay(monkeypatch, delay_abort): + started: asyncio.Queue[str] = asyncio.Queue() + abort_started: asyncio.Queue[str] = asyncio.Queue() + allow_abort = asyncio.Event() + if not delay_abort: + allow_abort.set() + request_ids: list[str] = [] + closed: list[str] = [] + aborted: list[str] = [] + delays: list[float] = [] + original_sleep = asyncio.sleep + + async def record_sleep(delay, result=None): + # Observe requested delays without a machine-dependent latency limit. + delays.append(delay) + return await original_sleep(0, result) + + monkeypatch.setattr(video_stream_base.asyncio, "sleep", record_sleep) + + class BlockingEngine: + async def generate(self, *, request_id, **kwargs): + if request_ids: + assert aborted[-1] == request_ids[-1] + request_ids.append(request_id) + started.put_nowait(request_id) + try: + if len(request_ids) <= 3: + yield _text_result("partial") + await asyncio.Event().wait() + else: + yield _text_result("final") + finally: + closed.append(request_id) + + async def abort(self, request_id): + assert request_id in closed + abort_started.put_nowait(request_id) + await allow_abort.wait() + aborted.append(request_id) + + class PreprocessedHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"prompt_token_ids": [1]} + + ws = TimedWebSocket() + handler = PreprocessedHandler(chat_service=object(), engine_client=BlockingEngine(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + try: + ws.put({"type": "session.config", "modalities": ["text"], "enable_frame_filter": False}) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + ws.put({"type": "video.query", "text": "first"}) + previous_id = await asyncio.wait_for(started.get(), timeout=2.0) + + for _ in range(3): + ws.put({"type": "video.query", "text": "interrupt"}) + assert await asyncio.wait_for(abort_started.get(), timeout=2.0) == previous_id + if delay_abort: + assert previous_id not in aborted + assert started.empty() + allow_abort.set() + previous_id = await asyncio.wait_for(started.get(), timeout=2.0) + if delay_abort: + allow_abort.clear() + + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert aborted == request_ids[:-1] + assert len(set(request_ids)) == 4 + assert [msg["text"] for msg in ws.sent if msg["type"] == "response.text.done"] == ["final"] + assert "error" not in ws.sent_types() + assert "session.done" in ws.sent_types() + assert not [delay for delay in delays if delay > 0], "Restart must not add a timed grace period after abort" + finally: + allow_abort.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.asyncio async def test_video_done_waits_for_in_flight_query(): query_started = asyncio.Event() @@ -588,6 +672,68 @@ async def test_frame_prewarm_does_not_block_following_query(monkeypatch): assert "session.done" in ws.sent_types() +@pytest.mark.asyncio +@pytest.mark.parametrize("max_frames", [1, 2], ids=["evicted", "retained"]) +async def test_frame_prewarm_only_keeps_retained_images(monkeypatch, max_frames): + frame_a = _make_jpeg(255, 0, 0) + frame_b = _make_jpeg(0, 255, 0) + decode_started = asyncio.Event() + release_decode = asyncio.Event() + frame_b_accepted = asyncio.Event() + decoded_images: dict[bytes, weakref.ReferenceType[Image.Image]] = {} + prewarm_task: asyncio.Task | None = None + original_to_thread = asyncio.to_thread + + async def controlled_to_thread(function, *args, **kwargs): + nonlocal prewarm_task + if function is video_stream_base._decode_frame_bytes: + if args[0] == frame_a: + prewarm_task = asyncio.current_task() + decode_started.set() + await release_decode.wait() + image = await original_to_thread(function, *args, **kwargs) + decoded_images[args[0]] = weakref.ref(image) + return image + return await original_to_thread(function, *args, **kwargs) + + class AckWebSocket(TimedWebSocket): + async def send_json(self, data): + await super().send_json(data) + if data.get("type") == "video.frame.ack" and data.get("frame_id") == "B": + frame_b_accepted.set() + + monkeypatch.setattr(video_stream_base.asyncio, "to_thread", controlled_to_thread) + ws = AckWebSocket() + handler = QwenOmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0) + session_task = asyncio.create_task(handler.handle_session(ws)) + ws.put({"type": "session.config", "max_frames": max_frames, "enable_frame_filter": False}) + try: + ws.put({"type": "video.frame", "frame_id": "A", "data": _b64(frame_a)}) + await asyncio.wait_for(decode_started.wait(), timeout=5.0) + ws.put({"type": "video.frame", "frame_id": "B", "data": _b64(frame_b)}) + await asyncio.wait_for(frame_b_accepted.wait(), timeout=5.0) + ack = next(message for message in ws.sent if message.get("frame_id") == "B") + assert ack["accepted"] is True + assert ack.get("dropped_frame_id") == ("A" if max_frames == 1 else None) + + # Finish A's real decode only after B has either evicted A or joined it. + release_decode.set() + assert prewarm_task is not None + await asyncio.wait_for(asyncio.shield(prewarm_task), timeout=5.0) + gc.collect() + assert not session_task.done() + # A finished task must not keep an evicted PIL image alive for the session. + assert (decoded_images[frame_a]() is not None) == (max_frames == 2) + finally: + release_decode.set() + ws.put({"type": "video.done"}) + await asyncio.wait_for(session_task, timeout=5.0) + + gc.collect() + assert decoded_images[frame_a]() is None + assert not any(message.get("type") == "error" for message in ws.sent) + + @pytest.mark.asyncio async def test_client_cannot_send_internal_frame_decode_failed_message(): captured_frames: list[list[str]] = [] diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index af0083fcb..e05a40eef 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.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 """Qwen-Omni streaming video WebSocket handler. Accepts video frames incrementally via WebSocket, buffers them, and @@ -26,12 +26,14 @@ Protocol: from __future__ import annotations +from collections.abc import Mapping from typing import Any from vllm_omni.entrypoints.openai.video_stream_base import ( _BAD_FRAME, _DEFAULT_CONFIG_TIMEOUT, _DEFAULT_IDLE_TIMEOUT, + PrewarmedFrame, StreamingVideoSessionConfig, VideoStreamTurnTrigger, ) @@ -59,7 +61,7 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase): audio_buffer: bytearray, message_history: list[dict[str, Any]], query_text: str, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: n_buf = len(frame_buffer) if n_buf <= config.num_frames: diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index f975adbef..ab5d51c33 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.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 """Base WebSocket handler for streaming video input understanding. Shared session loop, frame/audio buffering, EVS pre-filter, prewarm, @@ -37,7 +37,8 @@ import uuid import wave from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Protocol, runtime_checkable +from enum import Enum +from typing import Any, Final, Protocol, TypeAlias, runtime_checkable import torch from fastapi import WebSocket, WebSocketDisconnect @@ -61,7 +62,14 @@ _MAX_BUFFER_FRAMES = 64 _MAX_AUDIO_BUFFER_BYTES = 4 * 1024 * 1024 _MAX_MSG_QUEUE = 200 _CODEC_FRAME_SAMPLES = 1920 # CausalConv leading-edge artifact length -_BAD_FRAME = object() + + +class _FrameStatus(Enum): + BAD = "bad" + + +_BAD_FRAME: Final = _FrameStatus.BAD +PrewarmedFrame: TypeAlias = tuple[Any, str] | _FrameStatus def _decode_frame_bytes(raw_bytes: bytes) -> Any: @@ -83,7 +91,7 @@ class VideoStreamPipelineHooks(Protocol): audio_buffer: bytearray, message_history: list[dict[str, Any]], query_text: str, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: """Build OpenAI-style messages and the current user message.""" ... @@ -169,7 +177,7 @@ class OmniStreamingVideoHandler: audio_buffer: bytearray, message_history: list[dict[str, Any]], query_text: str, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: raise NotImplementedError @@ -219,7 +227,7 @@ class OmniStreamingVideoHandler: frame_buffer: list[str] = [] # base64-encoded JPEG frames frame_metadata: list[dict[str, Any]] = [] # Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index. - frame_pil_cache: dict[str, tuple[Any, str] | object] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME + frame_pil_cache: dict[str, PrewarmedFrame] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME frame_filter = ( FrameSimilarityFilter(threshold=config.frame_filter_threshold) if config.enable_frame_filter else None ) @@ -306,7 +314,6 @@ class OmniStreamingVideoHandler: await self._engine_client.abort(prev_request_id) except Exception: pass - await asyncio.sleep(0.1) prev_was_interrupted = False request_id = f"video-{uuid.uuid4().hex[:12]}" @@ -429,8 +436,12 @@ class OmniStreamingVideoHandler: async def _prewarm(b64: str, b: bytes, u: str) -> None: try: pil = await asyncio.to_thread(_decode_frame_bytes, b) - frame_pil_cache[b64] = (pil, u) + # The frame may have been evicted while decoding. + if b64 in frame_buffer: + frame_pil_cache[b64] = (pil, u) except Exception: + if b64 not in frame_buffer: + return frame_pil_cache[b64] = _BAD_FRAME logger.warning("prewarm decode failed for frame (len=%d)", len(b)) try: @@ -577,7 +588,7 @@ class OmniStreamingVideoHandler: query_text: str, request_id: str, interrupt_event: asyncio.Event, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], frame_metadata: list[dict[str, Any]] | None = None, ) -> None: """Build prompt, run inference, stream text + audio response.""" @@ -616,10 +627,13 @@ class OmniStreamingVideoHandler: query_text: str, request_id: str, interrupt_event: asyncio.Event, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], frame_metadata: list[dict[str, Any]] | None = None, ) -> None: """Direct engine_client.generate() path for async_chunk audio.""" + engine_client = self._engine_client + assert engine_client is not None, "_process_query must validate the engine client" + from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) @@ -686,7 +700,7 @@ class OmniStreamingVideoHandler: audio_tail_tensors: list[Any] = [] try: - result_gen = self._engine_client.generate( + result_gen = engine_client.generate( prompt=engine_prompt, request_id=request_id, output_modalities=config.modalities, @@ -989,6 +1003,8 @@ class OmniStreamingVideoHandler: ) mixin = AudioMixin() resp = mixin.create_audio(audio_obj) + # base64_encode=True selects the string result of create_audio(). + assert isinstance(resp.audio_data, str) return resp.audio_data @staticmethod