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