[Bugfix] Fix video prewarm cache retention and cancel-restart delay (#7363)

Signed-off-by: psv666 <2693925048@qq.com>

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]] = []

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:

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