From 7bceb1a76b99652b7a7abb315c7cc6f1f952b967 Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:49:07 +0800 Subject: [PATCH] [Bugfix] Align video frame consumption reporting with prompt sampling (#7329) --- docs/serving/video_stream_api.md | 8 + .../openai_api/test_serving_video_stream.py | 194 ++++++++++++++++++ .../openai/serving_video_stream.py | 20 +- .../entrypoints/openai/video_stream_base.py | 32 ++- 4 files changed, 232 insertions(+), 22 deletions(-) diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md index 101289fc0..23aa4cb08 100644 --- a/docs/serving/video_stream_api.md +++ b/docs/serving/video_stream_api.md @@ -50,6 +50,14 @@ WebSocket /v1/video/chat/stream | Server -> Client | `session.done` | none | Session closed. | | Server -> Client | `error` | `message` | Recoverable protocol or generation error. | +### Frame Consumption Reporting + +When buffered frames include client-supplied `frame_id` values, the server emits `video.frames.consumed` after the first engine output. Its `frame_ids`, `frames`, and `latest_pts_ms` describe the images included in that query's prompt. + +Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. Excluded frames are not replaced with other buffered frames. Frames whose background decoding has not finished remain eligible through the image URL path. + +**Bugfix compatibility note:** `video.frames.consumed` now excludes known decode failures that were already excluded from the prompt. Older versions could report those frames and their timestamps as consumed. The event name and fields are unchanged; an empty selection reports empty lists and `latest_pts_ms: null`. + ### `session.config` Fields | Field | Type | Default | Description | diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index f7267a061..c9c4f6a4b 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -375,6 +375,200 @@ async def test_video_frames_consumed_is_emitted_after_engine_uses_frame_prompt() ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "num_frames,bad_index,expected_indices", + [ + pytest.param(2, 4, [0], id="bad-last-no-refill"), + pytest.param(1, 4, [], id="all-selected-frames-bad"), + pytest.param(8, 4, [0, 1, 2, 3], id="below-sampling-limit"), + pytest.param(2, 1, [0, 4], id="bad-frame-not-selected"), + pytest.param(2, 0, [4], id="bad-first-frame"), + ], +) +async def test_video_frames_consumed_excludes_bad_frame_in_query_snapshot( + monkeypatch, num_frames, bad_index, expected_indices +): + """A query queued before bad-frame cleanup must report only its prompt images.""" + query_queued = asyncio.Event() + decode_failed = asyncio.Event() + release_decode = asyncio.Event() + pong_blocked = asyncio.Event() + release_pong = asyncio.Event() + turn_done = asyncio.Event() + original_to_thread = asyncio.to_thread + bad_image = b"not-a-jpeg" + frames = [_make_jpeg(r=index * 40) for index in range(5)] + frames[bad_index] = bad_image + engine_prompts = [] + + class ObservedQueue(asyncio.Queue): + def put_nowait(self, item): + super().put_nowait(item) + if isinstance(item, dict): + if item.get("type") == "video.query": + query_queued.set() + elif item.get("type") == "_internal.frame_decode_failed": + decode_failed.set() + + async def gated_to_thread(function, *args, **kwargs): + if function is video_stream_base._decode_frame_bytes and args[0] == bad_image: + await release_decode.wait() + return await original_to_thread(function, *args, **kwargs) + + class BackpressuredWebSocket(TimedWebSocket): + async def send_json(self, data): + await super().send_json(data) + if data["type"] == "pong": + pong_blocked.set() + await release_pong.wait() + elif data["type"] == "response.text.done": + turn_done.set() + + class OneOutputEngine: + async def generate(self, **kwargs): + engine_prompts.append(kwargs["prompt"]) + yield _text_result("visible") + + class CapturingHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"messages": request.messages} + + # Construct the client queue before observing the server's message queue. + ws = BackpressuredWebSocket() + monkeypatch.setattr(asyncio, "Queue", ObservedQueue) + monkeypatch.setattr(asyncio, "to_thread", gated_to_thread) + handler = CapturingHandler(chat_service=object(), engine_client=OneOutputEngine(), idle_timeout=5.0) + ws.put( + { + "type": "session.config", + "model": "test", + "modalities": ["text"], + "num_frames": num_frames, + "enable_frame_filter": False, + } + ) + for index, frame in enumerate(frames): + ws.put({"type": "video.frame", "data": _b64(frame), "frame_id": f"frame-{index}", "pts_ms": index * 100}) + ws.put({"type": "ping"}) + ws.put({"type": "video.query", "text": "describe"}) + task = asyncio.create_task(handler.handle_session(ws)) + try: + # Hold the processor at a WebSocket send. Queue the query first, then + # finish the real failing decode so its cleanup lands behind the query. + await asyncio.wait_for(pong_blocked.wait(), timeout=5.0) + await asyncio.wait_for(query_queued.wait(), timeout=5.0) + release_decode.set() + await asyncio.wait_for(decode_failed.wait(), timeout=5.0) + release_pong.set() + await asyncio.wait_for(turn_done.wait(), timeout=5.0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=5.0) + finally: + release_decode.set() + release_pong.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert len(engine_prompts) == 1 + images = [part for part in engine_prompts[0]["messages"][-1]["content"] if part["type"].startswith("image_")] + # Expected positions preserve sample-then-filter: do not refill a dropped + # position or shift the remaining metadata when a frame fails to decode. + assert len(images) == len(expected_indices) + for image, index in zip(images, expected_indices): + if image["type"] == "image_pil": + assert image["image_pil"].tobytes() == Image.open(io.BytesIO(frames[index])).convert("RGB").tobytes() + else: + assert image["image_url"]["url"] == f"data:image/jpeg;base64,{_b64(frames[index])}" + consumed_events = [message for message in ws.sent if message.get("type") == "video.frames.consumed"] + assert len(consumed_events) == 1 + consumed = consumed_events[0] + expected_ids = [f"frame-{index}" for index in expected_indices] + assert consumed["frame_ids"] == expected_ids + assert [frame["frame_id"] for frame in consumed["frames"]] == expected_ids + assert [frame["pts_ms"] for frame in consumed["frames"]] == [index * 100 for index in expected_indices] + assert consumed["latest_pts_ms"] == (expected_indices[-1] * 100 if expected_indices else None) + assert {"type": "error", "message": "Frame decode failed"} in ws.sent + assert ws.sent[-1]["type"] == "session.done" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "num_frames,frame_ids,repeat_image,expected_indices", + [ + pytest.param(2, ["f0", "f1", "f2", "f3", "f4"], False, [0, 4], id="stride-and-last"), + pytest.param(1, ["f0", "f1", "f2"], False, [2], id="last-only"), + pytest.param(8, ["f0", "f1", "f2"], False, [0, 1, 2], id="below-sampling-limit"), + pytest.param(2, [None, "f1", None, "f3", None], False, [0, 4], id="unidentified-selected-frames"), + pytest.param(2, [None, None, None], False, [0, 2], id="no-frame-ids"), + pytest.param(2, ["same", "same", "same"], False, [0, 2], id="repeated-frame-id"), + pytest.param(2, ["f0", "f1", "f2"], True, [0, 2], id="repeated-image"), + ], +) +async def test_video_frame_selection_preserves_uncached_frames( + monkeypatch, num_frames, frame_ids, repeat_image, expected_indices +): + """Prewarm misses still use image URLs, with metadata paired by buffer position.""" + engine_prompts = [] + pending_decode = asyncio.Event() + + async def blocked_to_thread(*args, **kwargs): + await pending_decode.wait() + + monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread) + + class OneOutputEngine: + async def generate(self, **kwargs): + engine_prompts.append(kwargs["prompt"]) + yield _text_result("visible") + + class CapturingHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"messages": request.messages} + + frames = [_b64(_make_jpeg(r=0 if repeat_image else index * 40)) for index in range(len(frame_ids))] + messages = [ + { + "type": "session.config", + "model": "test", + "modalities": ["text"], + "num_frames": num_frames, + "enable_frame_filter": False, + }, + *[ + {"type": "video.frame", "data": frame, "frame_id": frame_ids[index], "pts_ms": index * 100} + for index, frame in enumerate(frames) + ], + {"type": "video.query", "text": "describe"}, + {"type": "video.done"}, + ] + ws = MockWebSocket([json.dumps(message) for message in messages]) + handler = CapturingHandler(chat_service=object(), engine_client=OneOutputEngine()) + await asyncio.wait_for(handler.handle_session(ws), timeout=5.0) + + assert not [message for message in ws.sent if message["type"] == "error"] + assert len(engine_prompts) == 1 + content = engine_prompts[0]["messages"][-1]["content"] + assert content == [ + *[ + {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{frames[index]}"}} + for index in expected_indices + ], + {"type": "text", "text": "describe"}, + ] + consumed_events = [message for message in ws.sent if message["type"] == "video.frames.consumed"] + if not any(frame_ids): + assert consumed_events == [] + else: + assert len(consumed_events) == 1 + consumed = consumed_events[0] + assert consumed["frame_ids"] == [frame_ids[index] for index in expected_indices if frame_ids[index] is not None] + assert [frame["frame_id"] for frame in consumed["frames"]] == [frame_ids[index] for index in expected_indices] + assert [frame["pts_ms"] for frame in consumed["frames"]] == [index * 100 for index in expected_indices] + assert consumed["latest_pts_ms"] == expected_indices[-1] * 100 + + @pytest.mark.asyncio async def test_audio_in_video_sets_mm_processor_kwargs(): captured_requests = [] diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index 5b30ff0fa..5f7334bd2 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -30,7 +30,6 @@ 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, @@ -62,22 +61,19 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase): message_history: list[dict[str, Any]], query_text: str, prewarmed_frames: Mapping[str, PrewarmedFrame], + *, + frame_indices: list[int] | None = None, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: - n_buf = len(frame_buffer) - if n_buf <= config.num_frames: - frames = list(frame_buffer) - else: - stride = max(1, n_buf // config.num_frames) - idx = [i * stride for i in range(config.num_frames - 1)] + [n_buf - 1] - frames = [frame_buffer[i] for i in idx] - prewarmed = prewarmed_frames or {} + if frame_indices is None: + frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed) user_content: list[dict] = [] - for frame_b64 in frames: + for index in frame_indices: + frame_b64 = frame_buffer[index] cached = prewarmed.get(frame_b64) - if cached is _BAD_FRAME: - continue if cached is not None: + # The shared selector excludes failed decodes before this point. + assert isinstance(cached, tuple) pil, pil_uuid = cached user_content.append( { diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index 8885e9a33..52ae9d894 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -93,8 +93,10 @@ class VideoStreamPipelineHooks(Protocol): message_history: list[dict[str, Any]], query_text: str, prewarmed_frames: Mapping[str, PrewarmedFrame], + *, + frame_indices: list[int] | None = None, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: - """Build OpenAI-style messages and the current user message.""" + """Build messages using the supplied frame selection, or sample if omitted.""" ... def on_turn_complete( @@ -179,6 +181,8 @@ class OmniStreamingVideoHandler: message_history: list[dict[str, Any]], query_text: str, prewarmed_frames: Mapping[str, PrewarmedFrame], + *, + frame_indices: list[int] | None = None, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: raise NotImplementedError @@ -639,6 +643,9 @@ class OmniStreamingVideoHandler: ChatCompletionRequest, ) + # Select once from this turn's snapshot so prompt images and the + # consumed event refer to exactly the same frames, including decode failures. + frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed_frames) messages, user_message = self.build_engine_prompt( config, frame_buffer, @@ -646,6 +653,7 @@ class OmniStreamingVideoHandler: message_history, query_text, prewarmed_frames, + frame_indices=frame_indices, ) request_kwargs: dict[str, Any] = { @@ -678,7 +686,7 @@ class OmniStreamingVideoHandler: await self._send_error(websocket, f"Preprocess failed: {e}") return decoded_ready_ts_ms = _time.monotonic() * 1000 - selected_metadata = self._sample_frame_metadata(frame_metadata or [], config.num_frames) + selected_metadata = [frame_metadata[index] for index in frame_indices] if frame_metadata else [] model_selected_ts_ms = _time.monotonic() * 1000 await websocket.send_json({"type": "response.start"}) @@ -1071,15 +1079,19 @@ class OmniStreamingVideoHandler: pass @staticmethod - def _sample_frame_metadata( - frame_metadata: list[dict[str, Any]], + def _sample_frame_indices( + frame_buffer: list[str], num_frames: int, - ) -> list[dict[str, Any]]: - if len(frame_metadata) <= num_frames: - return list(frame_metadata) - stride = max(1, len(frame_metadata) // num_frames) - indices = [index * stride for index in range(num_frames - 1)] + [len(frame_metadata) - 1] - return [frame_metadata[index] for index in indices] + prewarmed_frames: Mapping[str, PrewarmedFrame], + ) -> list[int]: + """Stride-sample with the last frame, then drop known bad frames without refilling.""" + n_buf = len(frame_buffer) + if n_buf <= num_frames: + indices = list(range(n_buf)) + else: + stride = max(1, n_buf // num_frames) + indices = [index * stride for index in range(num_frames - 1)] + [n_buf - 1] + return [index for index in indices if prewarmed_frames.get(frame_buffer[index]) is not _BAD_FRAME] @staticmethod async def _send_frame_ack(