Zheng Wengangandamy-why-3459
eb1699bfca
Signed-off-by: ZhengWG <zwg0606@gmail.com> Signed-off-by: Zheng Wengang <zwg0606@gmail.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
1789 lines
65 KiB
Python
1789 lines
65 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# 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 hashlib
|
|
import io
|
|
import json
|
|
import threading
|
|
import weakref
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
|
|
|
from tests.helpers.serving_chat import build_serving_chat
|
|
from vllm_omni.config.stage_config import StageConfig
|
|
from vllm_omni.entrypoints.openai import video_stream_base, video_stream_envs
|
|
from vllm_omni.entrypoints.openai.serving_video_stream import (
|
|
QwenOmniStreamingVideoHandler,
|
|
StreamingVideoSessionConfig,
|
|
)
|
|
from vllm_omni.entrypoints.openai.video_stream_base import (
|
|
OmniStreamingVideoHandler,
|
|
)
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _make_jpeg(r: int = 128, g: int = 128, b: int = 128) -> bytes:
|
|
img = Image.new("RGB", (64, 64), (r, g, b))
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="JPEG", quality=95)
|
|
return buf.getvalue()
|
|
|
|
|
|
def _b64(data: bytes) -> str:
|
|
return base64.b64encode(data).decode()
|
|
|
|
|
|
def _text_result(text: str) -> OmniRequestOutput:
|
|
class Output:
|
|
text: str
|
|
|
|
class RequestOutput:
|
|
outputs: list[Output]
|
|
|
|
output = Output()
|
|
output.text = text
|
|
request_output = RequestOutput()
|
|
request_output.outputs = [output]
|
|
return OmniRequestOutput.from_stage_output(request_output, final_output_type="text")
|
|
|
|
|
|
def _audio_result(audio_data: Any) -> OmniRequestOutput:
|
|
class Output:
|
|
multimodal_output: dict[str, Any]
|
|
|
|
class RequestOutput:
|
|
outputs: list[Output]
|
|
|
|
output = Output()
|
|
output.multimodal_output = {"audio": audio_data}
|
|
request_output = RequestOutput()
|
|
request_output.outputs = [output]
|
|
return OmniRequestOutput.from_stage_output(request_output, final_output_type="audio")
|
|
|
|
|
|
class MockWebSocket:
|
|
def __init__(self, messages: list[str] | None = None):
|
|
self._messages = list(messages or [])
|
|
self._idx = 0
|
|
self.accepted = False
|
|
self.sent: list[dict[str, Any]] = []
|
|
|
|
async def accept(self):
|
|
self.accepted = True
|
|
|
|
async def receive_text(self) -> str:
|
|
if self._idx >= len(self._messages):
|
|
await asyncio.sleep(999)
|
|
msg = self._messages[self._idx]
|
|
self._idx += 1
|
|
return msg
|
|
|
|
async def send_json(self, data: dict[str, Any]):
|
|
self.sent.append(data)
|
|
|
|
|
|
class TimedWebSocket:
|
|
def __init__(self):
|
|
self._q: asyncio.Queue[str] = asyncio.Queue()
|
|
self.accepted = False
|
|
self.sent: list[dict[str, Any]] = []
|
|
|
|
async def accept(self):
|
|
self.accepted = True
|
|
|
|
async def receive_text(self) -> str:
|
|
return await self._q.get()
|
|
|
|
async def send_json(self, data: dict[str, Any]):
|
|
self.sent.append(data)
|
|
|
|
def put(self, msg: dict[str, Any]):
|
|
self._q.put_nowait(json.dumps(msg))
|
|
|
|
def sent_types(self) -> list[str]:
|
|
return [m.get("type", "") for m in self.sent]
|
|
|
|
|
|
def test_api_server_registers_video_stream_route():
|
|
from vllm_omni.entrypoints.openai.api_server import router
|
|
|
|
assert any(getattr(route, "path", None) == "/v1/video/chat/stream" for route in router.routes)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_receive_config_accepts_client_legacy_aliases():
|
|
ws = MockWebSocket(
|
|
[
|
|
json.dumps(
|
|
{
|
|
"type": "session.config",
|
|
"model": "test",
|
|
"num_sample_frames": 7,
|
|
"evs_enabled": False,
|
|
"evs_threshold": 0.87,
|
|
}
|
|
)
|
|
]
|
|
)
|
|
handler = OmniStreamingVideoHandler(chat_service=object())
|
|
|
|
config = await handler._receive_config(ws)
|
|
|
|
assert config is not None
|
|
assert config.num_frames == 7
|
|
assert config.enable_frame_filter is False
|
|
assert config.frame_filter_threshold == 0.87
|
|
|
|
|
|
@pytest.fixture
|
|
def run_sampling_session(monkeypatch):
|
|
async def run(config_overrides):
|
|
engine = MagicMock()
|
|
engine.stage_configs = [
|
|
StageConfig(stage_id=index, model_stage=name)
|
|
for index, name in enumerate(("thinker", "talker", "code2wav"))
|
|
]
|
|
engine.default_sampling_params_list = [
|
|
SamplingParams(temperature=0.4, max_tokens=64, top_p=0.85),
|
|
SamplingParams(temperature=0.7, max_tokens=96),
|
|
SamplingParams(temperature=0.8, max_tokens=128),
|
|
]
|
|
|
|
async def generate(**kwargs):
|
|
yield _text_result("answer")
|
|
|
|
engine.generate.side_effect = generate
|
|
handler = QwenOmniStreamingVideoHandler(
|
|
chat_service=build_serving_chat(engine_client=engine),
|
|
engine_client=engine,
|
|
)
|
|
monkeypatch.setattr(handler, "_preprocess_to_engine_prompt", AsyncMock(return_value={"prompt": "video"}))
|
|
ws = MockWebSocket(
|
|
[
|
|
json.dumps(
|
|
{
|
|
"type": "session.config",
|
|
"model": "test",
|
|
"modalities": ["text"],
|
|
"enable_frame_filter": False,
|
|
**config_overrides,
|
|
}
|
|
),
|
|
json.dumps({"type": "video.frame", "data": _b64(_make_jpeg())}),
|
|
json.dumps({"type": "video.query", "text": "describe"}),
|
|
json.dumps({"type": "video.done"}),
|
|
]
|
|
)
|
|
await handler.handle_session(ws)
|
|
return ws, engine
|
|
|
|
return run
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("async_chunk_mode", ["on", "off"])
|
|
@pytest.mark.parametrize("num_stage_overrides", [1, 3])
|
|
async def test_video_sampling_params_reach_engine(
|
|
run_sampling_session, monkeypatch, async_chunk_mode, num_stage_overrides
|
|
):
|
|
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", async_chunk_mode)
|
|
overrides = [
|
|
{"temperature": 0.2, "max_tokens": 1, "seed": 42},
|
|
{"temperature": 0.6, "max_tokens": 12},
|
|
{"temperature": 0.9, "max_tokens": 24},
|
|
][:num_stage_overrides]
|
|
|
|
ws, engine = await run_sampling_session({"sampling_params_list": overrides})
|
|
|
|
assert not [msg for msg in ws.sent if msg["type"] == "error"]
|
|
engine.generate.assert_called_once()
|
|
params = engine.generate.call_args.kwargs["sampling_params_list"]
|
|
assert len(params) == 3
|
|
for index, param in enumerate(params):
|
|
assert isinstance(param, SamplingParams)
|
|
default = engine.default_sampling_params_list[index]
|
|
expected = (
|
|
overrides[index]
|
|
if index < len(overrides)
|
|
else {
|
|
"temperature": default.temperature,
|
|
"max_tokens": default.max_tokens,
|
|
}
|
|
)
|
|
assert param.temperature == expected["temperature"]
|
|
assert param.max_tokens == expected["max_tokens"]
|
|
assert param.output_kind == RequestOutputKind.DELTA
|
|
assert param is not default
|
|
assert params[0].seed == 42
|
|
assert all(param.output_kind == RequestOutputKind.CUMULATIVE for param in engine.default_sampling_params_list)
|
|
assert next(msg for msg in ws.sent if msg["type"] == "response.text.done")["text"] == "answer"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_video_sampling_params_provided_stage_uses_constructor_defaults(run_sampling_session):
|
|
ws, engine = await run_sampling_session({"sampling_params_list": [{"temperature": 0.2}]})
|
|
|
|
assert not [msg for msg in ws.sent if msg["type"] == "error"]
|
|
engine.generate.assert_called_once()
|
|
params = engine.generate.call_args.kwargs["sampling_params_list"]
|
|
constructor_defaults = SamplingParams()
|
|
assert params[0].temperature == 0.2
|
|
for field in ("max_tokens", "top_p"):
|
|
assert getattr(params[0], field) == getattr(constructor_defaults, field)
|
|
assert getattr(params[0], field) != getattr(engine.default_sampling_params_list[0], field)
|
|
assert len(params) == 3
|
|
for param, default in zip(params[1:], engine.default_sampling_params_list[1:]):
|
|
assert param.temperature == default.temperature
|
|
assert param.max_tokens == default.max_tokens
|
|
assert param.top_p == default.top_p
|
|
assert param is not default
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("config_overrides", [{}, {"sampling_params_list": None}, {"sampling_params_list": []}])
|
|
async def test_video_sampling_params_omitted_preserve_engine_defaults(run_sampling_session, config_overrides):
|
|
ws, engine = await run_sampling_session(config_overrides)
|
|
|
|
assert not [msg for msg in ws.sent if msg["type"] == "error"]
|
|
engine.generate.assert_called_once()
|
|
assert engine.generate.call_args.kwargs.get("sampling_params_list") is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("invalid_params", [{"temperature": -1}, {"max_tokens": 0}, {"unknown_sampling_option": 1}])
|
|
async def test_video_sampling_params_invalid_do_not_reach_engine(run_sampling_session, invalid_params):
|
|
ws, engine = await run_sampling_session({"sampling_params_list": [invalid_params]})
|
|
|
|
engine.generate.assert_not_called()
|
|
assert any(msg["type"] == "error" for msg in ws.sent)
|
|
assert ws.sent[-1]["type"] == "session.done"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_video_frame_ack_reports_receiver_buffer_state():
|
|
ws = MockWebSocket(
|
|
[
|
|
json.dumps(
|
|
{
|
|
"type": "session.config",
|
|
"model": "test",
|
|
"enable_frame_filter": False,
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
"type": "video.frame",
|
|
"data": _b64(_make_jpeg()),
|
|
"frame_id": "frame-7",
|
|
"pts_ms": 700,
|
|
"capture_ts_ms": 1234.5,
|
|
}
|
|
),
|
|
json.dumps({"type": "video.done"}),
|
|
]
|
|
)
|
|
handler = QwenOmniStreamingVideoHandler(chat_service=object())
|
|
|
|
await handler.handle_session(ws)
|
|
|
|
ack = next(message for message in ws.sent if message.get("type") == "video.frame.ack")
|
|
assert ack["frame_id"] == "frame-7"
|
|
assert ack["pts_ms"] == 700
|
|
assert ack["capture_ts_ms"] == 1234.5
|
|
assert ack["accepted"] is True
|
|
assert ack["buffered_frames"] == 1
|
|
assert ack["server_receive_ts_ms"] > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_video_frames_consumed_is_emitted_after_engine_uses_frame_prompt():
|
|
class OneOutputEngine:
|
|
def generate(self, **_kwargs):
|
|
async def _gen():
|
|
await asyncio.sleep(0.05)
|
|
yield _text_result("visible")
|
|
|
|
return _gen()
|
|
|
|
class CapturingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _preprocess_to_engine_prompt(self, request):
|
|
return {"prompt": "with-video"}
|
|
|
|
ws = MockWebSocket(
|
|
[
|
|
json.dumps(
|
|
{
|
|
"type": "session.config",
|
|
"model": "test",
|
|
"modalities": ["text"],
|
|
"num_frames": 1,
|
|
"enable_frame_filter": False,
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
"type": "video.frame",
|
|
"data": _b64(_make_jpeg()),
|
|
"frame_id": "frame-9",
|
|
"pts_ms": 900,
|
|
"source_pts_ms": 880,
|
|
"quality_profile": "balanced",
|
|
}
|
|
),
|
|
json.dumps({"type": "video.query", "text": "describe"}),
|
|
json.dumps({"type": "video.done"}),
|
|
]
|
|
)
|
|
handler = CapturingHandler(
|
|
chat_service=object(),
|
|
engine_client=OneOutputEngine(),
|
|
idle_timeout=2.0,
|
|
)
|
|
|
|
await handler.handle_session(ws)
|
|
|
|
consumed = next(message for message in ws.sent if message.get("type") == "video.frames.consumed")
|
|
assert consumed["frame_ids"] == ["frame-9"]
|
|
assert consumed["latest_pts_ms"] == 900
|
|
assert consumed["request_id"].startswith("video-")
|
|
assert consumed["model_selected_ts_ms"] > 0
|
|
assert consumed["frames"] == [
|
|
{
|
|
"frame_id": "frame-9",
|
|
"pts_ms": 900,
|
|
"source_pts_ms": 880,
|
|
"quality_profile": "balanced",
|
|
"receiver_received_ts_ms": consumed["frames"][0]["receiver_received_ts_ms"],
|
|
"decoded_ready_ts_ms": consumed["frames"][0]["decoded_ready_ts_ms"],
|
|
}
|
|
]
|
|
assert consumed["frames"][0]["receiver_received_ts_ms"] > 0
|
|
assert consumed["frames"][0]["decoded_ready_ts_ms"] >= consumed["frames"][0]["receiver_received_ts_ms"]
|
|
assert consumed["model_selected_ts_ms"] >= consumed["frames"][0]["decoded_ready_ts_ms"]
|
|
assert ws.sent.index(consumed) < next(
|
|
index for index, message in enumerate(ws.sent) if message.get("type") == "response.text.delta"
|
|
)
|
|
|
|
|
|
@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_pairs_metadata_by_position(num_frames, frame_ids, repeat_image, expected_indices):
|
|
"""Queries wait for prewarmed PIL frames, with metadata paired by buffer position."""
|
|
engine_prompts = []
|
|
|
|
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[-1] == {"type": "text", "text": "describe"}
|
|
image_parts = content[:-1]
|
|
assert [part["type"] for part in image_parts] == ["image_pil"] * len(expected_indices)
|
|
for part, index in zip(image_parts, expected_indices):
|
|
raw = base64.b64decode(frames[index])
|
|
assert part["image_pil"].tobytes() == Image.open(io.BytesIO(raw)).convert("RGB").tobytes()
|
|
assert part["uuid"] == hashlib.md5(raw, usedforsecurity=False).hexdigest()
|
|
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 = []
|
|
|
|
class EmptyEngine:
|
|
def generate(self, **_kwargs):
|
|
async def _gen():
|
|
if False:
|
|
yield None
|
|
|
|
return _gen()
|
|
|
|
class CapturingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _preprocess_to_engine_prompt(self, request):
|
|
captured_requests.append(request)
|
|
return {"prompt": "x"}
|
|
|
|
ws = MockWebSocket()
|
|
handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine())
|
|
config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"], use_audio_in_video=True)
|
|
|
|
await handler._process_query_engine(
|
|
ws,
|
|
config,
|
|
[_b64(_make_jpeg())],
|
|
bytearray(b"\x00\x00"),
|
|
[],
|
|
"what is happening?",
|
|
"req-1",
|
|
asyncio.Event(),
|
|
{},
|
|
)
|
|
|
|
assert captured_requests
|
|
assert captured_requests[0].mm_processor_kwargs == {"use_audio_in_video": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_audio_in_video_disabled_omits_mm_processor_kwargs():
|
|
captured_requests = []
|
|
|
|
class EmptyEngine:
|
|
def generate(self, **_kwargs):
|
|
async def _gen():
|
|
if False:
|
|
yield None
|
|
|
|
return _gen()
|
|
|
|
class CapturingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _preprocess_to_engine_prompt(self, request):
|
|
captured_requests.append(request)
|
|
return {"prompt": "x"}
|
|
|
|
ws = MockWebSocket()
|
|
handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine())
|
|
config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"], use_audio_in_video=False)
|
|
|
|
await handler._process_query_engine(
|
|
ws,
|
|
config,
|
|
[_b64(_make_jpeg())],
|
|
bytearray(b"\x00\x00"),
|
|
[],
|
|
"what is happening?",
|
|
"req-1",
|
|
asyncio.Event(),
|
|
{},
|
|
)
|
|
|
|
assert captured_requests
|
|
assert captured_requests[0].mm_processor_kwargs is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_inline_audio_data_sets_mm_processor_kwargs():
|
|
captured_requests = []
|
|
|
|
class EmptyEngine:
|
|
def generate(self, **_kwargs):
|
|
async def _gen():
|
|
if False:
|
|
yield None
|
|
|
|
return _gen()
|
|
|
|
class CapturingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _preprocess_to_engine_prompt(self, request):
|
|
captured_requests.append(request)
|
|
return {"prompt": "x"}
|
|
|
|
ws = MockWebSocket(
|
|
[
|
|
json.dumps({"type": "session.config", "model": "test"}),
|
|
json.dumps({"type": "video.frame", "data": _b64(_make_jpeg())}),
|
|
json.dumps(
|
|
{
|
|
"type": "video.query",
|
|
"text": "describe",
|
|
"audio_data": _b64(b"\x00\x00"),
|
|
}
|
|
),
|
|
json.dumps({"type": "video.done"}),
|
|
]
|
|
)
|
|
handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=2.0)
|
|
|
|
await handler.handle_session(ws)
|
|
|
|
assert captured_requests
|
|
assert captured_requests[0].mm_processor_kwargs == {"use_audio_in_video": True}
|
|
assert "session.done" in [m.get("type") for m in ws.sent]
|
|
|
|
|
|
def test_audio_delta_mode_is_read_by_serving_code_at_runtime(monkeypatch):
|
|
handler = OmniStreamingVideoHandler(chat_service=object())
|
|
result = _audio_result([object()])
|
|
|
|
monkeypatch.setattr(
|
|
OmniStreamingVideoHandler,
|
|
"_delta_fast",
|
|
classmethod(lambda cls, audio_data, chunks_drained: ("fast-path", chunks_drained)),
|
|
)
|
|
monkeypatch.setattr(
|
|
OmniStreamingVideoHandler,
|
|
"_delta_slow",
|
|
classmethod(lambda cls, audio_data, chunks_drained: ("slow-path", chunks_drained)),
|
|
)
|
|
|
|
monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", "fast")
|
|
assert handler._extract_audio_delta_b64(result, 0)[0] == "fast-path"
|
|
|
|
monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", "slow")
|
|
assert handler._extract_audio_delta_b64(result, 0)[0] == "slow-path"
|
|
|
|
|
|
def test_video_stream_envs_strip_and_warn_once_per_invalid_value(monkeypatch):
|
|
warnings = []
|
|
|
|
video_stream_envs._warned_invalid_envs.clear()
|
|
try:
|
|
monkeypatch.setattr(
|
|
video_stream_envs.logger,
|
|
"warning",
|
|
lambda message, *args, **_kwargs: warnings.append((message, args)),
|
|
)
|
|
|
|
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", " off ")
|
|
assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "off"
|
|
assert not warnings
|
|
|
|
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "bad")
|
|
assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on"
|
|
assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on"
|
|
assert len(warnings) == 1
|
|
|
|
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "still_bad")
|
|
assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on"
|
|
assert len(warnings) == 2
|
|
finally:
|
|
video_stream_envs._warned_invalid_envs.clear()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_chunk_mode_is_read_by_engine_path_at_runtime(monkeypatch):
|
|
class TextEngine:
|
|
def generate(self, **_kwargs):
|
|
async def _gen():
|
|
yield _text_result("hello")
|
|
|
|
return _gen()
|
|
|
|
class CapturingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _preprocess_to_engine_prompt(self, request):
|
|
return {"prompt": "x"}
|
|
|
|
handler = CapturingHandler(chat_service=object(), engine_client=TextEngine())
|
|
config = StreamingVideoSessionConfig(model="test", modalities=["text"])
|
|
|
|
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "on")
|
|
ws_on = MockWebSocket()
|
|
await handler._process_query_engine(
|
|
ws_on,
|
|
config,
|
|
[_b64(_make_jpeg())],
|
|
bytearray(),
|
|
[],
|
|
"describe",
|
|
"req-on",
|
|
asyncio.Event(),
|
|
{},
|
|
)
|
|
assert {"type": "response.text.delta", "delta": "hello"} in ws_on.sent
|
|
|
|
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "off")
|
|
ws_off = MockWebSocket()
|
|
await handler._process_query_engine(
|
|
ws_off,
|
|
config,
|
|
[_b64(_make_jpeg())],
|
|
bytearray(),
|
|
[],
|
|
"describe",
|
|
"req-off",
|
|
asyncio.Event(),
|
|
{},
|
|
)
|
|
assert {"type": "response.text.done", "text": "hello"} in ws_off.sent
|
|
assert not any(m.get("type") == "response.text.delta" for m in ws_off.sent)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("async_chunk_mode", ["on", "off"])
|
|
async def test_audio_events_use_output_audio_names(monkeypatch, async_chunk_mode):
|
|
class AudioEngine:
|
|
def generate(self, **_kwargs):
|
|
async def _gen():
|
|
yield _audio_result([torch.zeros(1)])
|
|
|
|
return _gen()
|
|
|
|
class AudioHandler(QwenOmniStreamingVideoHandler):
|
|
async def _preprocess_to_engine_prompt(self, request):
|
|
return {"prompt": "x"}
|
|
|
|
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", async_chunk_mode)
|
|
monkeypatch.setattr(
|
|
OmniStreamingVideoHandler,
|
|
"_encode_audio_wav_b64",
|
|
staticmethod(lambda _audio: "audio-b64"),
|
|
)
|
|
|
|
ws = MockWebSocket()
|
|
handler = AudioHandler(chat_service=object(), engine_client=AudioEngine())
|
|
config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"])
|
|
|
|
await handler._process_query_engine(
|
|
ws,
|
|
config,
|
|
[_b64(_make_jpeg())],
|
|
bytearray(),
|
|
[],
|
|
"describe",
|
|
"req-audio-events",
|
|
asyncio.Event(),
|
|
{},
|
|
)
|
|
|
|
audio_events = [message for message in ws.sent if "audio" in message.get("type", "")]
|
|
assert audio_events == [
|
|
{
|
|
"type": "response.output_audio.delta",
|
|
"data": "audio-b64",
|
|
"format": "wav",
|
|
},
|
|
{"type": "response.output_audio.done"},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_without_engine_client_sends_error():
|
|
ws = MockWebSocket()
|
|
handler = OmniStreamingVideoHandler(chat_service=object(), engine_client=None)
|
|
|
|
await handler._process_query(
|
|
ws,
|
|
StreamingVideoSessionConfig(model="test"),
|
|
[],
|
|
bytearray(),
|
|
[],
|
|
"describe",
|
|
"req-1",
|
|
asyncio.Event(),
|
|
{},
|
|
)
|
|
|
|
assert {"type": "error", "message": "Streaming video requires an engine client"} in ws.sent
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_new_query_cancels_in_flight_query():
|
|
query_started = asyncio.Event()
|
|
query_cancelled = asyncio.Event()
|
|
calls = 0
|
|
|
|
class BlockingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _process_query(self, *args, **kwargs):
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls > 1:
|
|
return
|
|
query_started.set()
|
|
try:
|
|
await asyncio.sleep(999)
|
|
except asyncio.CancelledError:
|
|
query_cancelled.set()
|
|
raise
|
|
|
|
ws = TimedWebSocket()
|
|
handler = BlockingHandler(chat_service=object(), idle_timeout=5.0)
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
|
|
ws.put({"type": "session.config", "model": "test"})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
|
|
|
ws.put({"type": "video.query", "text": "interrupt"})
|
|
await asyncio.wait_for(query_cancelled.wait(), timeout=2.0)
|
|
ws.put({"type": "video.done"})
|
|
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
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):
|
|
submitted = kwargs.get("submitted")
|
|
if submitted is not None and not submitted.done():
|
|
submitted.set_result(None)
|
|
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()
|
|
allow_finish = asyncio.Event()
|
|
query_finished = asyncio.Event()
|
|
|
|
class BlockingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _process_query(self, *args, **kwargs):
|
|
query_started.set()
|
|
await allow_finish.wait()
|
|
query_finished.set()
|
|
|
|
ws = TimedWebSocket()
|
|
handler = BlockingHandler(chat_service=object(), idle_timeout=5.0)
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
|
|
ws.put({"type": "session.config", "model": "test"})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
|
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.sleep(0.05)
|
|
assert not task.done()
|
|
assert not query_finished.is_set()
|
|
|
|
allow_finish.set()
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
assert query_finished.is_set()
|
|
assert "session.done" in ws.sent_types()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_waits_for_frame_prewarm(monkeypatch):
|
|
decode_started = threading.Event()
|
|
release_decode = threading.Event()
|
|
query_started = asyncio.Event()
|
|
|
|
def blocked_decode(raw_bytes: bytes):
|
|
decode_started.set()
|
|
release_decode.wait(timeout=2.0)
|
|
return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
|
|
|
|
class BlockingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _process_query(self, *args, **kwargs):
|
|
query_started.set()
|
|
|
|
monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", blocked_decode)
|
|
|
|
ws = TimedWebSocket()
|
|
handler = BlockingHandler(chat_service=object(), idle_timeout=5.0)
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
|
|
ws.put({"type": "session.config", "model": "test"})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
|
|
for _ in range(100):
|
|
if decode_started.is_set():
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
assert decode_started.is_set()
|
|
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
await asyncio.sleep(0.05)
|
|
assert not query_started.is_set()
|
|
|
|
release_decode.set()
|
|
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
assert "session.done" in ws.sent_types()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_legacy_query_does_not_wait_for_unselected_frame(monkeypatch):
|
|
"""APC off and num_frames=1: a blocked older frame must not delay the query."""
|
|
frame_a = _make_jpeg(1, 0, 0)
|
|
frame_b = _make_jpeg(0, 1, 0)
|
|
decode_a_started = threading.Event()
|
|
release_a = threading.Event()
|
|
decode_b_done = threading.Event()
|
|
query_started = asyncio.Event()
|
|
|
|
def gated_decode(raw_bytes: bytes):
|
|
if raw_bytes == frame_a:
|
|
decode_a_started.set()
|
|
release_a.wait(timeout=5.0)
|
|
image = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
|
|
if raw_bytes == frame_b:
|
|
decode_b_done.set()
|
|
return image
|
|
|
|
class RecordingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _process_query(self, *args, **kwargs):
|
|
query_started.set()
|
|
|
|
monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
|
|
ws = TimedWebSocket()
|
|
handler = RecordingHandler(chat_service=object(), engine_client=object(), idle_timeout=5.0)
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
try:
|
|
ws.put(
|
|
{
|
|
"type": "session.config",
|
|
"model": "test",
|
|
"modalities": ["text"],
|
|
"num_frames": 1,
|
|
"enable_frame_filter": False,
|
|
}
|
|
)
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(frame_a), "frame_id": "A"})
|
|
for _ in range(100):
|
|
if decode_a_started.is_set():
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
assert decode_a_started.is_set()
|
|
ws.put({"type": "video.frame", "data": _b64(frame_b), "frame_id": "B"})
|
|
for _ in range(200):
|
|
if decode_b_done.is_set():
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
assert decode_b_done.is_set()
|
|
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
|
assert not release_a.is_set()
|
|
release_a.set()
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
finally:
|
|
release_a.set()
|
|
if not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
|
|
|
|
@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]] = []
|
|
frame = _b64(_make_jpeg())
|
|
|
|
class CapturingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _process_query(
|
|
self,
|
|
websocket,
|
|
config,
|
|
frame_buffer,
|
|
audio_buffer,
|
|
message_history,
|
|
query_text,
|
|
request_id,
|
|
interrupt_event,
|
|
prewarmed_frames,
|
|
):
|
|
captured_frames.append(list(frame_buffer))
|
|
|
|
ws = TimedWebSocket()
|
|
handler = CapturingHandler(chat_service=object(), idle_timeout=5.0)
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
|
|
ws.put({"type": "session.config", "model": "test"})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": frame})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "_internal.frame_decode_failed", "b64": frame})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
assert {"type": "error", "message": "Unknown type: _internal.frame_decode_failed"} in ws.sent
|
|
assert captured_frames == [[frame]]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_frame_prewarm_removes_frame_before_query():
|
|
ws = TimedWebSocket()
|
|
handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0)
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
|
|
ws.put({"type": "session.config", "model": "test", "enable_frame_filter": False})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(b"not-a-jpeg")})
|
|
|
|
for _ in range(100):
|
|
if any(m.get("message") == "Frame decode failed" for m in ws.sent):
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
|
|
assert {"type": "error", "message": "Frame decode failed"} in ws.sent
|
|
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
assert {"type": "error", "message": "No frames buffered"} in ws.sent
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_frame_filter_error_sends_invalid_image(monkeypatch):
|
|
def fail_should_retain(self, frame_jpeg):
|
|
raise ValueError("decode failed")
|
|
|
|
monkeypatch.setattr(video_stream_base.FrameSimilarityFilter, "should_retain", fail_should_retain)
|
|
|
|
ws = TimedWebSocket()
|
|
handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0)
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
|
|
ws.put({"type": "session.config", "model": "test"})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
assert {"type": "error", "message": "Invalid image data"} in ws.sent
|
|
assert "session.done" in ws.sent_types()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_audio_buffer_overflow_clears_buffer_before_query(monkeypatch):
|
|
captured_audio_lengths: list[int] = []
|
|
|
|
class EmptyEngine:
|
|
def generate(self, **_kwargs):
|
|
async def _gen():
|
|
if False:
|
|
yield None
|
|
|
|
return _gen()
|
|
|
|
class CapturingHandler(QwenOmniStreamingVideoHandler):
|
|
async def _process_query_engine(
|
|
self,
|
|
websocket,
|
|
config,
|
|
frame_buffer,
|
|
audio_buffer,
|
|
message_history,
|
|
query_text,
|
|
request_id,
|
|
interrupt_event,
|
|
prewarmed_frames,
|
|
):
|
|
captured_audio_lengths.append(len(audio_buffer))
|
|
|
|
monkeypatch.setattr(video_stream_base, "_MAX_AUDIO_BUFFER_BYTES", 4)
|
|
|
|
ws = TimedWebSocket()
|
|
handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=5.0)
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
|
|
ws.put({"type": "session.config", "model": "test"})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "audio.chunk", "data": _b64(b"1234")})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "audio.chunk", "data": _b64(b"5")})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
assert {"type": "error", "message": "Audio buffer overflow"} in ws.sent
|
|
assert captured_audio_lengths == [0]
|
|
|
|
|
|
def test_build_messages_keeps_recent_history_text_only():
|
|
handler = QwenOmniStreamingVideoHandler(chat_service=object())
|
|
old_frame = _b64(_make_jpeg(1, 2, 3))
|
|
current_frame = _b64(_make_jpeg(4, 5, 6))
|
|
history = [
|
|
{"role": "user", "content": [{"type": "text", "text": "old question"}]},
|
|
{"role": "assistant", "content": "old answer"},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{old_frame}"}},
|
|
{"type": "input_audio", "input_audio": {"data": "ignored", "format": "wav"}},
|
|
{"type": "text", "text": "recent question"},
|
|
],
|
|
},
|
|
{"role": "assistant", "content": "recent answer"},
|
|
]
|
|
|
|
messages, user_message = handler._build_messages(
|
|
StreamingVideoSessionConfig(model="test", num_frames=1),
|
|
[current_frame],
|
|
bytearray(),
|
|
history,
|
|
"current question",
|
|
{},
|
|
)
|
|
|
|
assert messages[0] == {"role": "user", "content": "recent question"}
|
|
assert messages[1] == {"role": "assistant", "content": "recent answer"}
|
|
assert messages[2] == user_message
|
|
assert user_message["content"][-1] == {"type": "text", "text": "current question"}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Incremental prefill (warmup) lifecycle
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class RecordingReuseEngine:
|
|
"""Engine stub with stage-0 prefix caching on; records generate() and abort()."""
|
|
|
|
def __init__(self, warmup_delay: float = 0.0, query_delay: float = 0.0):
|
|
self.requests: list[dict[str, Any]] = []
|
|
self.aborted: list[str] = []
|
|
self._warmup_delay = warmup_delay
|
|
self._query_delay = query_delay
|
|
self.engine = SimpleNamespace(
|
|
stage_vllm_configs=[SimpleNamespace(cache_config=SimpleNamespace(enable_prefix_caching=True))]
|
|
)
|
|
|
|
def warmup_ids(self) -> list[str]:
|
|
return [r["request_id"] for r in self.requests if r["request_id"].startswith("video-warmup-")]
|
|
|
|
def query_ids(self) -> list[str]:
|
|
return [r["request_id"] for r in self.requests if not r["request_id"].startswith("video-warmup-")]
|
|
|
|
def generate(self, **kwargs):
|
|
self.requests.append(kwargs)
|
|
submitted = kwargs.get("submitted")
|
|
if submitted is not None and not submitted.done():
|
|
submitted.set_result(None)
|
|
is_warmup = kwargs["request_id"].startswith("video-warmup-")
|
|
delay = self._warmup_delay if is_warmup else self._query_delay
|
|
|
|
async def _gen():
|
|
if delay:
|
|
await asyncio.sleep(delay)
|
|
yield _text_result("ok")
|
|
|
|
return _gen()
|
|
|
|
async def abort(self, request_id: str) -> None:
|
|
self.aborted.append(request_id)
|
|
|
|
|
|
class PromptRecordingHandler(QwenOmniStreamingVideoHandler):
|
|
"""Records each query turn's user-message content and processor kwargs."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.query_contents: list[list[dict[str, Any]]] = []
|
|
self.warmup_contents: list[list[dict[str, Any]]] = []
|
|
self.mm_processor_kwargs: list[Any] = []
|
|
|
|
async def _preprocess_to_engine_prompt(self, request):
|
|
self.mm_processor_kwargs.append(getattr(request, "mm_processor_kwargs", None))
|
|
return {"prompt": "p"}
|
|
|
|
def build_engine_prompt(
|
|
self, config, frame_buffer, audio_buffer, message_history, query_text, prewarmed_frames, *, frame_indices=None
|
|
):
|
|
messages, user_message = super().build_engine_prompt(
|
|
config,
|
|
frame_buffer,
|
|
audio_buffer,
|
|
message_history,
|
|
query_text,
|
|
prewarmed_frames,
|
|
frame_indices=frame_indices,
|
|
)
|
|
self.query_contents.append(list(user_message["content"]))
|
|
return messages, user_message
|
|
|
|
def build_engine_prompt_prefix(self, config, frame_buffer, message_history, prewarmed_frames):
|
|
messages = super().build_engine_prompt_prefix(config, frame_buffer, message_history, prewarmed_frames)
|
|
if messages:
|
|
self.warmup_contents.append(list(messages[-1]["content"]))
|
|
return messages
|
|
|
|
|
|
async def _poll(predicate, timeout: float = 3.0) -> bool:
|
|
for _ in range(int(timeout / 0.01)):
|
|
if predicate():
|
|
return True
|
|
await asyncio.sleep(0.01)
|
|
return bool(predicate())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_warmup_fires_on_frame_and_query_cancels_it():
|
|
engine = RecordingReuseEngine(warmup_delay=10.0, query_delay=0.3)
|
|
handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
|
|
|
|
ws = TimedWebSocket()
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
|
|
assert await _poll(lambda: engine.warmup_ids())
|
|
warmup_id = engine.warmup_ids()[0]
|
|
assert engine.requests[0]["sampling_params"].max_tokens == 1
|
|
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
assert await _poll(lambda: engine.query_ids())
|
|
# in-flight warmup was cancelled + aborted, and not restarted mid-query
|
|
assert warmup_id in engine.aborted
|
|
assert len(engine.warmup_ids()) == 1
|
|
assert handler.mm_processor_kwargs[:2] == [
|
|
{"use_audio_in_video": True},
|
|
{"use_audio_in_video": True},
|
|
]
|
|
|
|
assert await _poll(lambda: "response.text.done" in ws.sent_types())
|
|
# the committed turn changed the context -> a fresh warmup fires
|
|
assert await _poll(lambda: len(engine.warmup_ids()) == 2)
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_warmup_reaches_stage0_before_abort():
|
|
"""Cancel during submit must not abort a warmup the engine never received."""
|
|
|
|
class _HoldSubmitEngine(RecordingReuseEngine):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.started = asyncio.Event()
|
|
self.release_submit = asyncio.Event()
|
|
|
|
def generate(self, **kwargs):
|
|
submitted = kwargs.get("submitted")
|
|
|
|
async def _gen():
|
|
self.started.set()
|
|
await self.release_submit.wait()
|
|
self.requests.append(kwargs)
|
|
if submitted is not None and not submitted.done():
|
|
submitted.set_result(None)
|
|
yield _text_result("ok")
|
|
|
|
return _gen()
|
|
|
|
engine = _HoldSubmitEngine()
|
|
handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
|
|
ws = TimedWebSocket()
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
assert await _poll(engine.started.is_set)
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
await asyncio.sleep(0.05)
|
|
assert engine.requests == []
|
|
assert engine.aborted == []
|
|
engine.release_submit.set()
|
|
assert await _poll(lambda: engine.aborted)
|
|
assert engine.warmup_ids()
|
|
assert engine.aborted[0] == engine.warmup_ids()[0]
|
|
assert await _poll(lambda: engine.query_ids())
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_query_reaches_stage0_before_abort():
|
|
"""Cancel during query submit must not abort a request the engine never received."""
|
|
|
|
class _HoldQueryEngine(RecordingReuseEngine):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.query_started = asyncio.Event()
|
|
self.release_submit = asyncio.Event()
|
|
|
|
def generate(self, **kwargs):
|
|
if str(kwargs.get("request_id", "")).startswith("video-warmup-"):
|
|
return super().generate(**kwargs)
|
|
submitted = kwargs.get("submitted")
|
|
|
|
async def _gen():
|
|
self.query_started.set()
|
|
await self.release_submit.wait()
|
|
self.requests.append(kwargs)
|
|
if submitted is not None and not submitted.done():
|
|
submitted.set_result(None)
|
|
yield _text_result("ok")
|
|
|
|
return _gen()
|
|
|
|
engine = _HoldQueryEngine()
|
|
handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
|
|
ws = TimedWebSocket()
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
assert await _poll(lambda: engine.warmup_ids())
|
|
ws.put({"type": "video.query", "text": "first"})
|
|
assert await _poll(engine.query_started.is_set)
|
|
ws.put({"type": "video.query", "text": "second"})
|
|
await asyncio.sleep(0.05)
|
|
assert engine.query_ids() == []
|
|
assert engine.aborted == []
|
|
engine.release_submit.set()
|
|
assert await _poll(lambda: engine.query_ids())
|
|
first_query = engine.query_ids()[0]
|
|
assert await _poll(lambda: first_query in engine.aborted)
|
|
assert await _poll(lambda: len(engine.query_ids()) == 2)
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_audio_session_never_starts_warmup():
|
|
engine = RecordingReuseEngine()
|
|
handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
|
|
|
|
ws = TimedWebSocket()
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
ws.put({"type": "session.config", "model": "test", "modalities": ["text", "audio"], "enable_frame_filter": False})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg(9, 9, 9))})
|
|
await asyncio.sleep(0.3)
|
|
|
|
assert engine.warmup_ids() == []
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_warmup_starts_only_after_frames_ready(monkeypatch):
|
|
decode_started = threading.Event()
|
|
release_decode = threading.Event()
|
|
|
|
def gated_decode(raw_bytes: bytes):
|
|
decode_started.set()
|
|
release_decode.wait(timeout=5.0)
|
|
return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
|
|
|
|
monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
|
|
engine = RecordingReuseEngine()
|
|
handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
|
|
|
|
ws = TimedWebSocket()
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
|
assert await _poll(decode_started.is_set)
|
|
await asyncio.sleep(0.1)
|
|
assert engine.warmup_ids() == []
|
|
|
|
release_decode.set()
|
|
assert await _poll(lambda: engine.warmup_ids())
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_cancelling_warmup_keeps_prewarm_alive(monkeypatch):
|
|
"""Cancel warmup during prefill must not drop a later frame still decoding."""
|
|
frame1 = _b64(_make_jpeg(1, 1, 1))
|
|
frame2 = _b64(_make_jpeg(2, 2, 2))
|
|
raw2 = base64.b64decode(frame2)
|
|
decode_started = threading.Event()
|
|
release_decode = threading.Event()
|
|
|
|
def gated_decode(raw_bytes: bytes):
|
|
if raw_bytes == raw2:
|
|
decode_started.set()
|
|
release_decode.wait(timeout=5.0)
|
|
return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
|
|
|
|
monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
|
|
engine = RecordingReuseEngine(warmup_delay=10.0)
|
|
handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
|
|
|
|
ws = TimedWebSocket()
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": frame1})
|
|
assert await _poll(lambda: engine.warmup_ids())
|
|
|
|
ws.put({"type": "video.frame", "data": frame2})
|
|
assert await _poll(decode_started.is_set)
|
|
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
await asyncio.sleep(0.1)
|
|
release_decode.set()
|
|
|
|
assert await _poll(lambda: "response.text.done" in ws.sent_types())
|
|
assert len(handler.query_contents) == 1
|
|
image_parts = [p for p in handler.query_contents[0] if p.get("type") == "image_pil"]
|
|
assert len(image_parts) == 2
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_keeps_frames_evicted_while_pinned(monkeypatch):
|
|
"""Regression: frames FIFO-evicted while an in-flight query awaits prewarm
|
|
must keep their PIL entries and stay in the query prompt."""
|
|
frame_a = _b64(_make_jpeg(1, 1, 1))
|
|
frame_a2 = _b64(_make_jpeg(2, 2, 2))
|
|
frame_b = _b64(_make_jpeg(3, 3, 3))
|
|
raw_a = base64.b64decode(frame_a)
|
|
decoded: list[bytes] = []
|
|
release_others = threading.Event()
|
|
|
|
def gated_decode(raw_bytes: bytes):
|
|
if raw_bytes != raw_a:
|
|
release_others.wait(timeout=5.0)
|
|
decoded.append(raw_bytes)
|
|
return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
|
|
|
|
monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
|
|
engine = RecordingReuseEngine()
|
|
handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
|
|
|
|
ws = TimedWebSocket()
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
ws.put(
|
|
{
|
|
"type": "session.config",
|
|
"model": "test",
|
|
"modalities": ["text"],
|
|
"max_frames": 2,
|
|
"enable_frame_filter": False,
|
|
}
|
|
)
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": frame_a})
|
|
assert await _poll(lambda: raw_a in decoded)
|
|
ws.put({"type": "video.frame", "data": frame_a2}) # decode gated -> query will wait
|
|
|
|
ws.put({"type": "video.query", "text": "describe"}) # pins [a, a2]
|
|
await asyncio.sleep(0.1)
|
|
ws.put({"type": "video.frame", "data": frame_b}) # buffer full -> evicts pinned a
|
|
await asyncio.sleep(0.1)
|
|
release_others.set()
|
|
|
|
assert await _poll(lambda: "response.text.done" in ws.sent_types())
|
|
assert len(handler.query_contents) == 1
|
|
uuids = [p["uuid"] for p in handler.query_contents[0] if p.get("type") == "image_pil"]
|
|
expected = [hashlib.md5(base64.b64decode(f), usedforsecurity=False).hexdigest() for f in (frame_a, frame_a2)]
|
|
assert uuids == expected
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
|
|
|
|
async def _query_image_uuids(modalities: list[str], frames: list[str], *, max_frames: int) -> list[str]:
|
|
engine = RecordingReuseEngine()
|
|
handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
|
|
ws = TimedWebSocket()
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
ws.put(
|
|
{
|
|
"type": "session.config",
|
|
"model": "test",
|
|
"modalities": modalities,
|
|
"max_frames": max_frames,
|
|
"num_frames": max_frames,
|
|
"enable_frame_filter": False,
|
|
}
|
|
)
|
|
await asyncio.sleep(0)
|
|
for frame in frames:
|
|
ws.put({"type": "video.frame", "data": frame})
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
assert await _poll(lambda: handler.query_contents)
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
return [part["uuid"] for part in handler.query_contents[0] if part.get("type") == "image_pil"]
|
|
|
|
|
|
def _frame_uuid(frame_b64: str) -> str:
|
|
return hashlib.md5(base64.b64decode(frame_b64), usedforsecurity=False).hexdigest()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_below_max_frames_incremental_matches_legacy_window():
|
|
frames = [_b64(_make_jpeg(i, 0, 0)) for i in range(3)]
|
|
expected = [_frame_uuid(frame) for frame in frames]
|
|
assert await _query_image_uuids(["text"], frames, max_frames=4) == expected
|
|
assert await _query_image_uuids(["text", "audio"], frames, max_frames=4) == expected
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_incremental_full_window_drops_oldest_half():
|
|
frames = [_b64(_make_jpeg(i, 1, 0)) for i in range(5)]
|
|
kept = await _query_image_uuids(["text"], frames, max_frames=4)
|
|
legacy = await _query_image_uuids(["text", "audio"], frames, max_frames=4)
|
|
assert kept == [_frame_uuid(frame) for frame in frames[2:]]
|
|
assert legacy == [_frame_uuid(frame) for frame in frames[1:]]
|
|
|
|
|
|
async def _run_frame_audio_query(modalities: list[str], frame: str, pcm: bytes) -> PromptRecordingHandler:
|
|
engine = RecordingReuseEngine()
|
|
handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
|
|
ws = TimedWebSocket()
|
|
task = asyncio.create_task(handler.handle_session(ws))
|
|
ws.put(
|
|
{
|
|
"type": "session.config",
|
|
"model": "test",
|
|
"modalities": modalities,
|
|
"enable_frame_filter": False,
|
|
}
|
|
)
|
|
await asyncio.sleep(0)
|
|
ws.put({"type": "video.frame", "data": frame})
|
|
if "audio" not in modalities:
|
|
assert await _poll(lambda: engine.warmup_ids())
|
|
else:
|
|
await asyncio.sleep(0.05)
|
|
assert engine.warmup_ids() == []
|
|
ws.put({"type": "audio.chunk", "data": _b64(pcm)})
|
|
ws.put({"type": "video.query", "text": "describe"})
|
|
assert await _poll(lambda: engine.query_ids())
|
|
ws.put({"type": "video.done"})
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
return handler
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reuse_query_with_audio_matches_legacy_prompt():
|
|
"""Incremental and legacy query prompts match; both set the processor kwarg."""
|
|
frame = _b64(_make_jpeg())
|
|
pcm = b"\x01\x00\x02\x00"
|
|
reuse = await _run_frame_audio_query(["text"], frame, pcm)
|
|
legacy = await _run_frame_audio_query(["text", "audio"], frame, pcm)
|
|
|
|
assert reuse.warmup_contents
|
|
assert all(part.get("type") != "input_audio" for part in reuse.warmup_contents[0])
|
|
assert [part.get("type") for part in reuse.warmup_contents[0]] == ["image_pil"]
|
|
|
|
assert reuse.query_contents and legacy.query_contents
|
|
reuse_q = reuse.query_contents[0]
|
|
legacy_q = legacy.query_contents[0]
|
|
assert [part.get("type") for part in reuse_q] == ["image_pil", "input_audio", "text"]
|
|
assert [part.get("type") for part in legacy_q] == ["image_pil", "input_audio", "text"]
|
|
assert reuse_q[0]["uuid"] == legacy_q[0]["uuid"]
|
|
assert reuse_q[1] == legacy_q[1]
|
|
assert reuse_q[2] == legacy_q[2] == {"type": "text", "text": "describe"}
|
|
|
|
assert reuse.mm_processor_kwargs[0] == {"use_audio_in_video": True}
|
|
assert reuse.mm_processor_kwargs[1] == {"use_audio_in_video": True}
|
|
assert legacy.mm_processor_kwargs[0] == {"use_audio_in_video": True}
|