[FEAT][Fronted] Support incremental prefill for realtime video streaming (#5894)
Signed-off-by: ZhengWG <zwg0606@gmail.com> Signed-off-by: Zheng Wengang <zwg0606@gmail.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
This commit is contained in:
1 parent
e372aa3a41
commit
eb1699bfca
6 files changed
+1187
-234
No files matched your search
@@ -16,6 +16,10 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
To enable [incremental prefill](#incremental-prefill), set
|
||||
`enable_prefix_caching: true` on stage 0 (the thinker) in the deploy config. Otherwise
|
||||
the server falls back to rebuilding the prompt per query.
|
||||
|
||||
### Run the Example Client
|
||||
|
||||
```bash
|
||||
@@ -54,7 +58,7 @@ WebSocket /v1/video/chat/stream
|
||||
|
||||
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.
|
||||
Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. With incremental prefill active, every buffered frame is included instead of stride-sampling, with the same decode-failure exclusion. Excluded frames are not replaced with other buffered frames. The incremental path waits for the whole window to finish decoding. The legacy path waits only for its sampled positions, so an unselected frame still decoding does not delay the query.
|
||||
|
||||
**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`.
|
||||
|
||||
@@ -64,10 +68,10 @@ Queries stride-sample the buffered frames, keeping the last frame, then exclude
|
||||
| ------- | ------ | --------- | ------------- |
|
||||
| `model` | string or null | null | Optional model name. Usually omitted because the server hosts one model. |
|
||||
| `modalities` | list[string] | `["text", "audio"]` | Output modalities. Use `["text"]`, `["audio"]`, or both. |
|
||||
| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. |
|
||||
| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Oldest frames are evicted first. |
|
||||
| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. Legacy path only: with incremental prefill active, every buffered frame is submitted and frame density is the client's responsibility (clients push discrete frames at their own rate). |
|
||||
| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Below the cap both paths keep every frame. Legacy evicts the oldest frame once full. Incremental prefill drops the oldest half at once, then appends until full again. |
|
||||
| `system_prompt` | string or null | null | Optional custom system prompt. |
|
||||
| `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. |
|
||||
| `use_audio_in_video` | bool | `true` | Incremental prefill: forwarded on warmup and every query when this is true, so multimodal hashes stay aligned. Legacy: forwarded only when this query has input audio. |
|
||||
| `sampling_params_list` | list or null | null | Optional per-stage parameter dictionaries. Each provided entry replaces that stage's deployment sampling settings. |
|
||||
| `enable_frame_filter` | bool | `true` | Enable EVS near-duplicate frame filtering. |
|
||||
| `frame_filter_threshold` | float, 0.0-1.0 | `0.95` | EVS similarity threshold. Higher keeps more frames; lower drops more near-duplicates. |
|
||||
@@ -107,9 +111,42 @@ The server accepts these legacy field names and rewrites them before validation.
|
||||
|
||||
EVS compares downsampled frames and drops near-duplicate frames before they enter the session frame buffer. `frame_filter_threshold` controls retention: higher values are more permissive and keep more frames; lower values are more aggressive and drop more similar frames.
|
||||
|
||||
## Incremental Prefill
|
||||
|
||||
Server-side only (stage-0 `enable_prefix_caching: true` and `modalities: ["text"]`).
|
||||
There is no `session.config` flag. Audio-output sessions stay on the legacy path.
|
||||
|
||||
Input is unchanged: `video.frame` and `audio.chunk` are separate buffers. Warmup
|
||||
prefills `history + frames` as they arrive (`max_tokens=1`). When
|
||||
`use_audio_in_video` is true, warmup and query both pass that processor kwarg.
|
||||
`video.query` is the same prompt plus optional input audio and the question text,
|
||||
so the warmed vision prefix cache is reused and the query pays for the audio/text
|
||||
suffix. All buffered frames are
|
||||
submitted in arrival order (`num_frames` subsample is legacy-only). Once the buffer
|
||||
hits `max_frames`, the oldest half is dropped in one step and a warning is logged;
|
||||
until then the buffer matches the legacy path. After a turn,
|
||||
history compresses to the last two text messages and the next warmup follows the
|
||||
new prefix.
|
||||
|
||||
### Legacy path (incremental prefill inactive)
|
||||
|
||||
Stage-0 prefix caching is off, or `modalities` includes `"audio"`: no warmup,
|
||||
frames are re-sampled up to `num_frames` each query, and the query waits only
|
||||
for those sampled frames to decode. Prompt shape is the same
|
||||
(`image_pil` frames, optional trailing `input_audio`, text). Both paths keep the
|
||||
frame buffer across queries (`max_frames`) and compress history the same way.
|
||||
|
||||
## Known Limitations
|
||||
|
||||
- Session KV reuse and incremental prefill are not implemented in this PR. Each `video.query` rebuilds the model prompt from the retained frame and audio buffers.
|
||||
- Past `max_frames`, incremental prefill keeps a suffix of the newest frames, not a resample of the whole stream. Rebuilding that longer window is left to a follow-up.
|
||||
- Incremental prefill applies only to text-only sessions with stage-0 prefix caching enabled
|
||||
on `WS /v1/video/chat/stream`. It does not accelerate `/v1/realtime` camera input, which
|
||||
uses duplex conversation items rather than `video.frame` / `video.query`. Audio-output
|
||||
sessions on this endpoint rebuild the prompt each query.
|
||||
- Stage-0 prefix caching is safe for audio-output requests: on a hit, the omni prefix cache
|
||||
restores the thinker hidden states the talker needs. A preempted request that resumes on a
|
||||
hit does not get them back, so size the deployment to avoid preemption; see
|
||||
[Prefix Caching](../design/feature/prefix_caching.md).
|
||||
- Back-to-back short replies can still expose an engine-layer scheduler race. The PR notes an observed workaround of at least 200 ms idle between turns when clients repeatedly see idle timeouts.
|
||||
- If the audio buffer exceeds the server limit, the server emits `Audio buffer overflow` and clears the currently buffered audio for the session.
|
||||
- The API is intended for Qwen3-Omni streaming video understanding; other models may not support the same multimodal processor arguments.
|
||||
@@ -7,10 +7,12 @@ 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
|
||||
|
||||
@@ -26,7 +28,9 @@ from vllm_omni.entrypoints.openai.serving_video_stream import (
|
||||
QwenOmniStreamingVideoHandler,
|
||||
StreamingVideoSessionConfig,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.video_stream_base import OmniStreamingVideoHandler
|
||||
from vllm_omni.entrypoints.openai.video_stream_base import (
|
||||
OmniStreamingVideoHandler,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -506,17 +510,9 @@ async def test_video_frames_consumed_excludes_bad_frame_in_query_snapshot(
|
||||
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."""
|
||||
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 = []
|
||||
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):
|
||||
@@ -550,13 +546,13 @@ async def test_video_frame_selection_preserves_uncached_frames(
|
||||
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"},
|
||||
]
|
||||
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 == []
|
||||
@@ -909,6 +905,9 @@ async def test_interrupted_queries_wait_for_abort_without_fixed_delay(monkeypatc
|
||||
|
||||
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)
|
||||
@@ -1004,7 +1003,7 @@ async def test_video_done_waits_for_in_flight_query():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
|
||||
async def test_query_waits_for_frame_prewarm(monkeypatch):
|
||||
decode_started = threading.Event()
|
||||
release_decode = threading.Event()
|
||||
query_started = asyncio.Event()
|
||||
@@ -1035,14 +1034,80 @@ async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
|
||||
assert decode_started.is_set()
|
||||
|
||||
ws.put({"type": "video.query", "text": "describe"})
|
||||
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
||||
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):
|
||||
@@ -1272,3 +1337,452 @@ def test_build_messages_keeps_recent_history_text_only():
|
||||
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}
|
||||
@@ -174,7 +174,9 @@ async def test_real_producer_deltas_reach_video_client(
|
||||
produced.append(OmniRequestOutput.from_stage_output(result, final_output_type="audio"))
|
||||
|
||||
class Engine:
|
||||
async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None):
|
||||
async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None, submitted=None):
|
||||
if submitted is not None and not submitted.done():
|
||||
submitted.set_result(None)
|
||||
if explicit_sampling:
|
||||
assert sampling_params_list is not None
|
||||
assert sampling_params_list[0].output_kind == RequestOutputKind.DELTA
|
||||
|
||||
@@ -128,6 +128,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
|
||||
reasoning_ended: bool | None = None,
|
||||
reasoning_parser_kwargs: dict[str, Any] | None = None,
|
||||
arrival_time: float | None = None,
|
||||
submitted: asyncio.Future[None] | None = None,
|
||||
) -> AsyncGenerator[OmniRequestOutput, None]:
|
||||
"""Generate outputs for the given prompt(s) asynchronously.
|
||||
|
||||
@@ -296,6 +297,9 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
|
||||
if admitting:
|
||||
await self._release_generate_admission()
|
||||
admitting = False
|
||||
# Callers wait on this before abort so Stage0 observes the sender cache.
|
||||
if submitted is not None and not submitted.done():
|
||||
submitted.set_result(None)
|
||||
# Refresh gauges on arrival.
|
||||
self._publish_request_gauges(len(self.request_states))
|
||||
|
||||
|
||||
@@ -66,29 +66,10 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase):
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
||||
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 index in frame_indices:
|
||||
frame_b64 = frame_buffer[index]
|
||||
cached = prewarmed.get(frame_b64)
|
||||
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(
|
||||
{
|
||||
"type": "image_pil",
|
||||
"image_pil": pil,
|
||||
"uuid": pil_uuid,
|
||||
}
|
||||
)
|
||||
else:
|
||||
user_content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/jpeg;base64,{frame_b64}"},
|
||||
}
|
||||
)
|
||||
frame_indices = self._select_frame_indices(config, frame_buffer, prewarmed)
|
||||
frames = [frame_buffer[index] for index in frame_indices]
|
||||
|
||||
user_content: list[dict] = self._build_frame_image_parts(frames, prewarmed)
|
||||
|
||||
if len(audio_buffer) > 0:
|
||||
wav_b64 = self._pcm_to_wav_b64(bytes(audio_buffer))
|
||||
@@ -107,14 +88,7 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase):
|
||||
|
||||
user_message: dict[str, Any] = {"role": "user", "content": user_content}
|
||||
|
||||
messages: list[dict[str, Any]] = []
|
||||
if config.system_prompt:
|
||||
messages.append({"role": "system", "content": config.system_prompt})
|
||||
|
||||
recent_history = message_history[-2:] if len(message_history) > 2 else message_history
|
||||
for hist_msg in recent_history:
|
||||
messages.append(self._text_only_message(hist_msg))
|
||||
|
||||
messages = self._history_prefix_messages(config, message_history)
|
||||
messages.append(user_message)
|
||||
|
||||
return messages, user_message
|
||||
|
||||
@@ -38,7 +38,7 @@ import wave
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any, Final, Protocol, TypeAlias, runtime_checkable
|
||||
from typing import Any, Final, Protocol, TypeAlias, TypeGuard, runtime_checkable
|
||||
|
||||
import torch
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
@@ -77,6 +77,21 @@ def _decode_frame_bytes(raw_bytes: bytes) -> Any:
|
||||
return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
|
||||
|
||||
|
||||
def _usable_prewarmed_frame(cached: PrewarmedFrame | None) -> TypeGuard[tuple[Any, str]]:
|
||||
return cached is not None and cached is not _BAD_FRAME
|
||||
|
||||
|
||||
def _stage0_prefix_caching_enabled(engine_client: Any) -> bool:
|
||||
"""Read stage-0 (thinker) prefix caching from the engine. Called once
|
||||
at handler init: incremental prefill is a server-side capability, not a
|
||||
per-request choice."""
|
||||
engine = getattr(engine_client, "engine", None)
|
||||
stage_vllm_configs = getattr(engine, "stage_vllm_configs", None)
|
||||
stage0 = stage_vllm_configs[0] if stage_vllm_configs else None
|
||||
cache_config = getattr(stage0, "cache_config", None)
|
||||
return bool(getattr(cache_config, "enable_prefix_caching", False))
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class VideoStreamPipelineHooks(Protocol):
|
||||
"""Pipeline-specific hooks for streaming video handlers."""
|
||||
@@ -108,6 +123,20 @@ class VideoStreamPipelineHooks(Protocol):
|
||||
"""Update session state after a successful turn."""
|
||||
...
|
||||
|
||||
def build_engine_prompt_prefix(
|
||||
self,
|
||||
config: "StreamingVideoSessionConfig",
|
||||
frame_buffer: list[str],
|
||||
message_history: list[dict[str, Any]],
|
||||
prewarmed_frames: Mapping[str, PrewarmedFrame],
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""Prefix of the next :meth:`build_engine_prompt` through the last vision token.
|
||||
|
||||
History plus arriving video frames, without query text or input audio.
|
||||
Return None when there are no usable frames (or to disable incremental prefill).
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VideoStreamTurnTrigger:
|
||||
@@ -144,7 +173,10 @@ class StreamingVideoSessionConfig(BaseModel):
|
||||
)
|
||||
use_audio_in_video: bool = Field(
|
||||
default=True,
|
||||
description="Interleave audio chunks with video frames when audio input is present.",
|
||||
description=(
|
||||
"Pass use_audio_in_video to the processor when this is true and "
|
||||
"(incremental prefill is active, or this query has input audio)."
|
||||
),
|
||||
)
|
||||
sampling_params_list: list[dict[str, Any]] | None = Field(
|
||||
default=None,
|
||||
@@ -194,10 +226,77 @@ class OmniStreamingVideoHandler:
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def _build_frame_image_parts(
|
||||
frames: list[str],
|
||||
prewarmed_frames: Mapping[str, PrewarmedFrame] | None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Decoded frame parts. Prefix and query both wait for PIL prewarm so
|
||||
every frame is ``image_pil`` plus the arrival-time uuid."""
|
||||
prewarmed = prewarmed_frames or {}
|
||||
parts: list[dict[str, Any]] = []
|
||||
for frame_b64 in frames:
|
||||
cached = prewarmed.get(frame_b64)
|
||||
if not _usable_prewarmed_frame(cached):
|
||||
continue
|
||||
pil, mm_uuid = cached
|
||||
parts.append(
|
||||
{
|
||||
"type": "image_pil",
|
||||
"image_pil": pil,
|
||||
"uuid": mm_uuid,
|
||||
}
|
||||
)
|
||||
return parts
|
||||
|
||||
def _history_prefix_messages(
|
||||
self,
|
||||
config: StreamingVideoSessionConfig,
|
||||
message_history: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Shared prompt prefix (system + compressed text-only history).
|
||||
|
||||
:meth:`build_engine_prompt_prefix` and :meth:`build_engine_prompt`
|
||||
MUST build this identically: any divergence before the frame parts
|
||||
voids the warmed KV from that token onward.
|
||||
"""
|
||||
messages: list[dict[str, Any]] = []
|
||||
if config.system_prompt:
|
||||
messages.append({"role": "system", "content": config.system_prompt})
|
||||
recent_history = message_history[-2:] if len(message_history) > 2 else message_history
|
||||
for hist_msg in recent_history:
|
||||
messages.append(self._text_only_message(hist_msg))
|
||||
return messages
|
||||
|
||||
def build_engine_prompt_prefix(
|
||||
self,
|
||||
config: StreamingVideoSessionConfig,
|
||||
frame_buffer: list[str],
|
||||
message_history: list[dict[str, Any]],
|
||||
prewarmed_frames: Mapping[str, PrewarmedFrame],
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""Prefix of the next :meth:`build_engine_prompt` through the last vision token.
|
||||
|
||||
History plus arriving video frames, without query text or input audio.
|
||||
A mismatch voids the warmed KV from the first diverging token.
|
||||
Return None when there are no usable frames (or override to disable).
|
||||
"""
|
||||
frame_parts = self._build_frame_image_parts(frame_buffer, prewarmed_frames)
|
||||
if not frame_parts:
|
||||
return None
|
||||
|
||||
messages = self._history_prefix_messages(config, message_history)
|
||||
messages.append({"role": "user", "content": frame_parts})
|
||||
return messages
|
||||
|
||||
def create_message_history(self, config: StreamingVideoSessionConfig) -> Any:
|
||||
"""Per-session conversation state (default: empty OpenAI-style list)."""
|
||||
return []
|
||||
|
||||
def _incremental_prefill_active(self, config: StreamingVideoSessionConfig) -> bool:
|
||||
"""Incremental frame prefill: stage-0 prefix caching on, text-only output."""
|
||||
return self._incremental_prefill_supported and "audio" not in config.modalities
|
||||
|
||||
def on_frame_buffered(
|
||||
self,
|
||||
raw_bytes: bytes,
|
||||
@@ -219,6 +318,7 @@ class OmniStreamingVideoHandler:
|
||||
self._idle_timeout = idle_timeout
|
||||
self._config_timeout = config_timeout
|
||||
self._engine_client = engine_client
|
||||
self._incremental_prefill_supported = _stage0_prefix_caching_enabled(engine_client)
|
||||
|
||||
async def handle_session(self, websocket: WebSocket) -> None:
|
||||
"""Main session loop for a single WebSocket connection."""
|
||||
@@ -244,6 +344,53 @@ class OmniStreamingVideoHandler:
|
||||
interrupt_event = asyncio.Event()
|
||||
prewarm_tasks: set[asyncio.Task[Any]] = set()
|
||||
query_task: asyncio.Task[Any] | None = None
|
||||
session_alive = True
|
||||
# Session-owned decode readiness. Warmup/query wait on these
|
||||
# Events, not on the prewarm tasks, so cancelling an engine
|
||||
# request cannot cascade into JPEG decode.
|
||||
frame_ready: dict[str, asyncio.Event] = {}
|
||||
# In-flight warmup/query snapshots keep their PIL entries if FIFO
|
||||
# evicts the live buffer while they await readiness.
|
||||
pinned_frame_refs: dict[str, int] = {}
|
||||
|
||||
def _drop_frame_cache(frame_b64: str) -> None:
|
||||
frame_pil_cache.pop(frame_b64, None)
|
||||
frame_ready.pop(frame_b64, None)
|
||||
|
||||
def _pin_frames(frames: list[str]) -> None:
|
||||
for frame_b64 in frames:
|
||||
pinned_frame_refs[frame_b64] = pinned_frame_refs.get(frame_b64, 0) + 1
|
||||
|
||||
def _unpin_frames(frames: list[str]) -> None:
|
||||
for frame_b64 in frames:
|
||||
remaining = pinned_frame_refs.get(frame_b64, 0) - 1
|
||||
if remaining <= 0:
|
||||
pinned_frame_refs.pop(frame_b64, None)
|
||||
# Settle the eviction this pin deferred: the frame left
|
||||
# the buffer, so its cache entry is orphaned.
|
||||
if frame_b64 not in frame_buffer:
|
||||
_drop_frame_cache(frame_b64)
|
||||
else:
|
||||
pinned_frame_refs[frame_b64] = remaining
|
||||
|
||||
def _snapshot_prewarmed(frames: list[str]) -> dict[str, Any]:
|
||||
return {frame_b64: frame_pil_cache[frame_b64] for frame_b64 in frames if frame_b64 in frame_pil_cache}
|
||||
|
||||
def _frames_ready(frames: list[str]) -> bool:
|
||||
for frame_b64 in frames:
|
||||
event = frame_ready.get(frame_b64)
|
||||
if event is None or not event.is_set():
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _await_frames_ready(frames: list[str]) -> None:
|
||||
waits: list[Any] = []
|
||||
for frame_b64 in dict.fromkeys(frames):
|
||||
event = frame_ready.get(frame_b64)
|
||||
if event is not None and not event.is_set():
|
||||
waits.append(event.wait())
|
||||
if waits:
|
||||
await asyncio.gather(*waits)
|
||||
|
||||
msg_queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue(maxsize=_MAX_MSG_QUEUE)
|
||||
|
||||
@@ -294,15 +441,95 @@ class OmniStreamingVideoHandler:
|
||||
interrupt_event.set()
|
||||
prev_was_interrupted = True
|
||||
logger.info("Interrupt signaled for %s", active_request_id)
|
||||
if query_task is not None and not query_task.done():
|
||||
query_task.cancel()
|
||||
await asyncio.gather(query_task, return_exceptions=True)
|
||||
query_task = None
|
||||
# Abort after the task reaches Stage0 so a cancel during preprocess still submits.
|
||||
if abort_now and self._engine_client:
|
||||
try:
|
||||
await self._engine_client.abort(active_request_id)
|
||||
except Exception:
|
||||
logger.debug("Abort failed for %s", active_request_id, exc_info=True)
|
||||
if query_task is not None and not query_task.done():
|
||||
query_task.cancel()
|
||||
await asyncio.gather(query_task, return_exceptions=True)
|
||||
query_task = None
|
||||
|
||||
# --- Incremental prefill: keep the prefix cache warmed to the latest
|
||||
# buffered frame so a query only pays for its own text suffix. ---
|
||||
warmup_task: asyncio.Task[Any] | None = None
|
||||
warmup_request_id: str | None = None
|
||||
warmed_signature: tuple[Any, ...] | None = None
|
||||
|
||||
def _context_signature() -> tuple[Any, ...]:
|
||||
history_len = len(message_history) if isinstance(message_history, list) else None
|
||||
return (history_len, tuple(frame_buffer))
|
||||
|
||||
def _maybe_start_warmup() -> None:
|
||||
nonlocal warmup_task, warmup_request_id
|
||||
if not session_alive:
|
||||
return
|
||||
if self._engine_client is None or not self._incremental_prefill_active(config):
|
||||
return
|
||||
if active_request_id is not None:
|
||||
return
|
||||
if warmup_task is not None and not warmup_task.done():
|
||||
return
|
||||
if not frame_buffer or _context_signature() == warmed_signature:
|
||||
return
|
||||
frames = list(frame_buffer)
|
||||
# Decode is session-owned: do not start prefill until PIL is
|
||||
# already in cache, so this task is only engine work.
|
||||
if not _frames_ready(frames):
|
||||
return
|
||||
signature = _context_signature()
|
||||
request_id = f"video-warmup-{uuid.uuid4().hex[:12]}"
|
||||
warmup_request_id = request_id
|
||||
_pin_frames(frames)
|
||||
|
||||
async def _run_warmup() -> None:
|
||||
nonlocal warmup_request_id, warmed_signature
|
||||
cancelled = False
|
||||
try:
|
||||
messages = self.build_engine_prompt_prefix(
|
||||
config, frames, message_history, _snapshot_prewarmed(frames)
|
||||
)
|
||||
if not messages:
|
||||
warmed_signature = signature
|
||||
return
|
||||
await self._prefill_context(config, messages, request_id)
|
||||
warmed_signature = signature
|
||||
except asyncio.CancelledError:
|
||||
cancelled = True
|
||||
raise
|
||||
except Exception:
|
||||
# Mark attempted so a persistently failing state does
|
||||
# not hot-loop; the next frame retriggers naturally.
|
||||
warmed_signature = signature
|
||||
logger.debug("Context warmup failed for %s", request_id, exc_info=True)
|
||||
finally:
|
||||
_unpin_frames(frames)
|
||||
if warmup_request_id == request_id:
|
||||
warmup_request_id = None
|
||||
# Skip on cancel: query already owns the session and will
|
||||
# call_soon after it finishes. Restarting here can spawn a
|
||||
# second warmup with the same frames before active_request_id
|
||||
# is set.
|
||||
if not cancelled:
|
||||
asyncio.get_running_loop().call_soon(_maybe_start_warmup)
|
||||
|
||||
warmup_task = asyncio.create_task(_run_warmup())
|
||||
|
||||
async def _cancel_warmup() -> None:
|
||||
nonlocal warmup_task, warmup_request_id
|
||||
request_id = warmup_request_id
|
||||
if warmup_task is not None and not warmup_task.done():
|
||||
warmup_task.cancel()
|
||||
await asyncio.gather(warmup_task, return_exceptions=True)
|
||||
warmup_task = None
|
||||
warmup_request_id = None
|
||||
if request_id and self._engine_client:
|
||||
try:
|
||||
await self._engine_client.abort(request_id)
|
||||
except Exception:
|
||||
logger.debug("Warmup abort failed for %s", request_id, exc_info=True)
|
||||
|
||||
async def _start_query_turn(*, query_text: str) -> None:
|
||||
"""Schedule a new inference turn from the current buffers."""
|
||||
@@ -311,28 +538,43 @@ class OmniStreamingVideoHandler:
|
||||
await _cancel_active_query()
|
||||
|
||||
if not frame_buffer:
|
||||
await _cancel_warmup()
|
||||
await self._send_error(websocket, "No frames buffered")
|
||||
return
|
||||
|
||||
# Claim the session before cancelling warmup so a leftover
|
||||
# call_soon from a just-finished warmup hits the 405 gate.
|
||||
request_id = f"video-{uuid.uuid4().hex[:12]}"
|
||||
active_request_id = request_id
|
||||
interrupt_event.clear()
|
||||
# Free prefill capacity for the query; completed warmup blocks
|
||||
# stay in the prefix cache and are reused by it.
|
||||
await _cancel_warmup()
|
||||
|
||||
if prev_was_interrupted and prev_request_id and self._engine_client:
|
||||
try:
|
||||
await self._engine_client.abort(prev_request_id)
|
||||
except Exception:
|
||||
pass
|
||||
prev_was_interrupted = False
|
||||
|
||||
request_id = f"video-{uuid.uuid4().hex[:12]}"
|
||||
active_request_id = request_id
|
||||
interrupt_event.clear()
|
||||
query_frames = list(frame_buffer)
|
||||
query_frame_metadata = list(frame_metadata)
|
||||
query_audio_buffer = bytearray(audio_buffer)
|
||||
audio_buffer.clear()
|
||||
query_prewarmed_frames = dict(frame_pil_cache)
|
||||
# Legacy stride positions ignore decode status. Waiting on the
|
||||
# whole buffer would stall the query on a frame it will not use.
|
||||
if self._incremental_prefill_active(config):
|
||||
wait_indices = list(range(len(query_frames)))
|
||||
else:
|
||||
wait_indices = self._candidate_frame_indices(query_frames, config.num_frames)
|
||||
wait_frames = [query_frames[index] for index in wait_indices]
|
||||
_pin_frames(wait_frames)
|
||||
|
||||
async def _run_query() -> None:
|
||||
nonlocal active_request_id, prev_request_id
|
||||
try:
|
||||
await _await_frames_ready(wait_frames)
|
||||
query_prewarmed_frames = _snapshot_prewarmed(query_frames)
|
||||
process_kwargs: dict[str, Any] = {}
|
||||
if any(metadata.get("frame_id") for metadata in query_frame_metadata):
|
||||
process_kwargs["frame_metadata"] = query_frame_metadata
|
||||
@@ -349,19 +591,25 @@ class OmniStreamingVideoHandler:
|
||||
**process_kwargs,
|
||||
)
|
||||
finally:
|
||||
_unpin_frames(wait_frames)
|
||||
if active_request_id == request_id:
|
||||
prev_request_id = request_id
|
||||
active_request_id = None
|
||||
# Warm frames that arrived during the turn (fires after
|
||||
# this task completes; gated on no active query).
|
||||
asyncio.get_running_loop().call_soon(_maybe_start_warmup)
|
||||
|
||||
query_task = asyncio.create_task(_run_query())
|
||||
|
||||
async def _processor() -> None:
|
||||
"""Process enqueued messages."""
|
||||
nonlocal active_request_id, prev_request_id, prev_was_interrupted, query_task
|
||||
nonlocal session_alive
|
||||
|
||||
while True:
|
||||
msg = await msg_queue.get()
|
||||
if msg is None:
|
||||
session_alive = False
|
||||
await _cancel_active_query(abort_now=True)
|
||||
return
|
||||
|
||||
@@ -376,8 +624,8 @@ class OmniStreamingVideoHandler:
|
||||
]
|
||||
frame_buffer[:] = [frame_buffer[index] for index in retained_indices]
|
||||
frame_metadata[:] = [frame_metadata[index] for index in retained_indices]
|
||||
if frame_pil_cache.get(frame_data) is _BAD_FRAME:
|
||||
frame_pil_cache.pop(frame_data, None)
|
||||
if frame_pil_cache.get(frame_data) is _BAD_FRAME and frame_data not in pinned_frame_refs:
|
||||
_drop_frame_cache(frame_data)
|
||||
if removed:
|
||||
await self._send_error(websocket, "Frame decode failed")
|
||||
|
||||
@@ -410,10 +658,24 @@ class OmniStreamingVideoHandler:
|
||||
max_buf = config.max_frames
|
||||
dropped_frame_id: str | None = None
|
||||
if len(frame_buffer) >= max_buf:
|
||||
dropped = frame_buffer.pop(0)
|
||||
dropped_metadata = frame_metadata.pop(0)
|
||||
dropped_frame_id = dropped_metadata.get("frame_id")
|
||||
frame_pil_cache.pop(dropped, None)
|
||||
# Legacy drops one frame. Incremental drops the oldest half so later frames append.
|
||||
drop_count = max(1, max_buf // 2) if self._incremental_prefill_active(config) else 1
|
||||
for _ in range(drop_count):
|
||||
dropped = frame_buffer.pop(0)
|
||||
dropped_metadata = frame_metadata.pop(0)
|
||||
if dropped_frame_id is None:
|
||||
dropped_frame_id = dropped_metadata.get("frame_id")
|
||||
# Duplicate frames share one cache entry keyed by
|
||||
# b64: keep it while any duplicate remains buffered.
|
||||
if dropped not in pinned_frame_refs and dropped not in frame_buffer:
|
||||
_drop_frame_cache(dropped)
|
||||
if drop_count > 1:
|
||||
logger.warning(
|
||||
"incremental prefill dropped %d oldest frames at max_frames=%d; "
|
||||
"buffered_frames is the remaining suffix",
|
||||
drop_count,
|
||||
max_buf,
|
||||
)
|
||||
frame_buffer.append(frame_data)
|
||||
frame_metadata.append(
|
||||
{
|
||||
@@ -435,17 +697,23 @@ class OmniStreamingVideoHandler:
|
||||
)
|
||||
# Prewarm: decode PIL off the event loop so query-time chat_template
|
||||
# can skip base64+Image.open. uuid=md5 lets mm_cache dedupe identical frames.
|
||||
if frame_data not in frame_pil_cache:
|
||||
if frame_data not in frame_pil_cache and frame_data not in frame_ready:
|
||||
mm_uuid = hashlib.md5(raw_bytes, usedforsecurity=False).hexdigest()
|
||||
frame_ready[frame_data] = asyncio.Event()
|
||||
|
||||
async def _prewarm(b64: str, b: bytes, u: str) -> None:
|
||||
cancelled = False
|
||||
try:
|
||||
pil = await asyncio.to_thread(_decode_frame_bytes, b)
|
||||
# The frame may have been evicted while decoding.
|
||||
if b64 in frame_buffer:
|
||||
# The frame may have been evicted (and unpinned)
|
||||
# while decoding; frame_ready tracks liveness.
|
||||
if b64 in frame_ready:
|
||||
frame_pil_cache[b64] = (pil, u)
|
||||
except asyncio.CancelledError:
|
||||
cancelled = True
|
||||
raise
|
||||
except Exception:
|
||||
if b64 not in frame_buffer:
|
||||
if b64 not in frame_ready:
|
||||
return
|
||||
frame_pil_cache[b64] = _BAD_FRAME
|
||||
logger.warning("prewarm decode failed for frame (len=%d)", len(b))
|
||||
@@ -455,10 +723,18 @@ class OmniStreamingVideoHandler:
|
||||
logger.warning(
|
||||
"frame decode failure event dropped because message queue is full"
|
||||
)
|
||||
finally:
|
||||
event = frame_ready.get(b64)
|
||||
if event is not None:
|
||||
event.set()
|
||||
if not cancelled:
|
||||
asyncio.get_running_loop().call_soon(_maybe_start_warmup)
|
||||
|
||||
task = asyncio.create_task(_prewarm(frame_data, raw_bytes, mm_uuid))
|
||||
prewarm_tasks.add(task)
|
||||
task.add_done_callback(prewarm_tasks.discard)
|
||||
elif frame_data in frame_pil_cache:
|
||||
frame_ready.setdefault(frame_data, asyncio.Event()).set()
|
||||
|
||||
is_generating = active_request_id is not None or (
|
||||
query_task is not None and not query_task.done()
|
||||
@@ -471,6 +747,8 @@ class OmniStreamingVideoHandler:
|
||||
)
|
||||
):
|
||||
await _start_query_turn(query_text="")
|
||||
else:
|
||||
_maybe_start_warmup()
|
||||
|
||||
elif msg_type == "audio.chunk":
|
||||
data_b64 = msg.get("data", "")
|
||||
@@ -501,6 +779,7 @@ class OmniStreamingVideoHandler:
|
||||
await _start_query_turn(query_text=query_text)
|
||||
|
||||
elif msg_type == "video.done":
|
||||
session_alive = False
|
||||
if query_task is not None and not query_task.done():
|
||||
await asyncio.gather(query_task, return_exceptions=True)
|
||||
query_task = None
|
||||
@@ -520,6 +799,7 @@ class OmniStreamingVideoHandler:
|
||||
try:
|
||||
await _processor()
|
||||
finally:
|
||||
session_alive = False
|
||||
reader_task.cancel()
|
||||
try:
|
||||
await reader_task
|
||||
@@ -529,6 +809,7 @@ class OmniStreamingVideoHandler:
|
||||
t.cancel()
|
||||
if prewarm_tasks:
|
||||
await asyncio.gather(*prewarm_tasks, return_exceptions=True)
|
||||
await _cancel_warmup()
|
||||
if query_task is not None and not query_task.done():
|
||||
await _cancel_active_query(abort_now=True)
|
||||
|
||||
@@ -643,9 +924,10 @@ class OmniStreamingVideoHandler:
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
|
||||
reuse_active = self._incremental_prefill_active(config)
|
||||
# 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)
|
||||
frame_indices = self._select_frame_indices(config, frame_buffer, prewarmed_frames)
|
||||
messages, user_message = self.build_engine_prompt(
|
||||
config,
|
||||
frame_buffer,
|
||||
@@ -665,7 +947,9 @@ class OmniStreamingVideoHandler:
|
||||
"continue_final_message": False,
|
||||
"add_special_tokens": False,
|
||||
}
|
||||
if config.use_audio_in_video and len(audio_buffer) > 0:
|
||||
# Incremental: same kwargs as warmup so frame mm hashes stay a prefix.
|
||||
# Legacy: only when this query actually has input audio.
|
||||
if config.use_audio_in_video and (reuse_active or len(audio_buffer) > 0):
|
||||
request_kwargs["mm_processor_kwargs"] = {
|
||||
"use_audio_in_video": True,
|
||||
}
|
||||
@@ -680,97 +964,142 @@ class OmniStreamingVideoHandler:
|
||||
await self._send_error(websocket, f"Failed to build request: {e}")
|
||||
return
|
||||
|
||||
try:
|
||||
engine_prompt = await self._preprocess_to_engine_prompt(chat_request)
|
||||
except Exception as e:
|
||||
await self._send_error(websocket, f"Preprocess failed: {e}")
|
||||
return
|
||||
decoded_ready_ts_ms = _time.monotonic() * 1000
|
||||
selected_metadata = [frame_metadata[index] for index in frame_indices] if frame_metadata else []
|
||||
model_selected_ts_ms = _time.monotonic() * 1000
|
||||
submitted = asyncio.get_running_loop().create_future()
|
||||
|
||||
await websocket.send_json({"type": "response.start"})
|
||||
text_parts: list[str] = []
|
||||
text_done_sent = False
|
||||
audio_chunk_count = 0
|
||||
# Count emitted tensors only to identify the first audio emission.
|
||||
# Default and explicit sampling parameters use DELTA audio outputs.
|
||||
audio_chunks_drained = 0
|
||||
previous_text = ""
|
||||
interrupted = False
|
||||
frames_consumed_sent = False
|
||||
t_start = _time.monotonic()
|
||||
t_first_text = None
|
||||
t_first_audio = None
|
||||
async def _drive_query() -> None:
|
||||
decoded_ready_ts_ms = _time.monotonic() * 1000
|
||||
selected_metadata = [frame_metadata[index] for index in frame_indices] if frame_metadata else []
|
||||
model_selected_ts_ms = _time.monotonic() * 1000
|
||||
|
||||
# Wire-level async-chunk switch. "off" means
|
||||
# buffer all deltas server-side and flush once at the end; the engine
|
||||
# pipeline still overlaps internally.
|
||||
async_chunk_mode = video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK
|
||||
streaming = async_chunk_mode == "on"
|
||||
audio_tail_tensors: list[Any] = []
|
||||
if not interrupt_event.is_set():
|
||||
await websocket.send_json({"type": "response.start"})
|
||||
text_parts: list[str] = []
|
||||
text_done_sent = False
|
||||
audio_chunk_count = 0
|
||||
# Count emitted tensors only to identify the first audio emission.
|
||||
# Default and explicit sampling parameters use DELTA audio outputs.
|
||||
audio_chunks_drained = 0
|
||||
previous_text = ""
|
||||
interrupted = False
|
||||
frames_consumed_sent = False
|
||||
t_start = _time.monotonic()
|
||||
t_first_text = None
|
||||
t_first_audio = None
|
||||
|
||||
try:
|
||||
result_gen = engine_client.generate(
|
||||
prompt=engine_prompt,
|
||||
request_id=request_id,
|
||||
output_modalities=config.modalities,
|
||||
**sampling_kwargs,
|
||||
)
|
||||
# Wire-level async-chunk switch. "off" means
|
||||
# buffer all deltas server-side and flush once at the end; the engine
|
||||
# pipeline still overlaps internally.
|
||||
async_chunk_mode = video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK
|
||||
streaming = async_chunk_mode == "on"
|
||||
audio_tail_tensors: list[Any] = []
|
||||
|
||||
async for output in result_gen:
|
||||
# Soft interrupt: drain without sending
|
||||
if interrupt_event.is_set():
|
||||
if not interrupted:
|
||||
logger.info("Generation interrupted — draining")
|
||||
interrupted = True
|
||||
continue
|
||||
try:
|
||||
assert engine_client is not None
|
||||
result_gen = engine_client.generate(
|
||||
prompt=engine_prompt,
|
||||
request_id=request_id,
|
||||
output_modalities=config.modalities,
|
||||
submitted=submitted,
|
||||
**sampling_kwargs,
|
||||
)
|
||||
|
||||
if not isinstance(output, OmniRequestOutput):
|
||||
continue
|
||||
async for output in result_gen:
|
||||
# Soft interrupt: drain without sending
|
||||
if interrupt_event.is_set():
|
||||
if not interrupted:
|
||||
logger.info("Generation interrupted — draining")
|
||||
interrupted = True
|
||||
continue
|
||||
|
||||
if not frames_consumed_sent and frame_metadata:
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "video.frames.consumed",
|
||||
"request_id": request_id,
|
||||
"model_selected_ts_ms": model_selected_ts_ms,
|
||||
"frame_ids": [
|
||||
metadata["frame_id"]
|
||||
for metadata in selected_metadata
|
||||
if isinstance(metadata.get("frame_id"), str)
|
||||
],
|
||||
"frames": [
|
||||
{
|
||||
"frame_id": metadata.get("frame_id"),
|
||||
"pts_ms": metadata.get("pts_ms"),
|
||||
"source_pts_ms": metadata.get("source_pts_ms"),
|
||||
"quality_profile": metadata.get("quality_profile"),
|
||||
"receiver_received_ts_ms": metadata.get("receiver_received_ts_ms"),
|
||||
"decoded_ready_ts_ms": decoded_ready_ts_ms,
|
||||
}
|
||||
for metadata in selected_metadata
|
||||
],
|
||||
"latest_pts_ms": selected_metadata[-1].get("pts_ms") if selected_metadata else None,
|
||||
}
|
||||
)
|
||||
frames_consumed_sent = True
|
||||
if not isinstance(output, OmniRequestOutput):
|
||||
continue
|
||||
|
||||
out_type = getattr(output, "final_output_type", "text")
|
||||
if not frames_consumed_sent and frame_metadata:
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "video.frames.consumed",
|
||||
"request_id": request_id,
|
||||
"model_selected_ts_ms": model_selected_ts_ms,
|
||||
"frame_ids": [
|
||||
metadata["frame_id"]
|
||||
for metadata in selected_metadata
|
||||
if isinstance(metadata.get("frame_id"), str)
|
||||
],
|
||||
"frames": [
|
||||
{
|
||||
"frame_id": metadata.get("frame_id"),
|
||||
"pts_ms": metadata.get("pts_ms"),
|
||||
"source_pts_ms": metadata.get("source_pts_ms"),
|
||||
"quality_profile": metadata.get("quality_profile"),
|
||||
"receiver_received_ts_ms": metadata.get("receiver_received_ts_ms"),
|
||||
"decoded_ready_ts_ms": decoded_ready_ts_ms,
|
||||
}
|
||||
for metadata in selected_metadata
|
||||
],
|
||||
"latest_pts_ms": selected_metadata[-1].get("pts_ms") if selected_metadata else None,
|
||||
}
|
||||
)
|
||||
frames_consumed_sent = True
|
||||
|
||||
if out_type == "audio":
|
||||
if streaming and not text_done_sent:
|
||||
full_text = "".join(text_parts)
|
||||
await websocket.send_json({"type": "response.text.done", "text": full_text})
|
||||
text_done_sent = True
|
||||
out_type = getattr(output, "final_output_type", "text")
|
||||
|
||||
if t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
audio_chunk_count += 1
|
||||
if streaming:
|
||||
b64, audio_chunks_drained = self._extract_audio_delta_b64(
|
||||
if out_type == "audio":
|
||||
if streaming and not text_done_sent:
|
||||
full_text = "".join(text_parts)
|
||||
await websocket.send_json({"type": "response.text.done", "text": full_text})
|
||||
text_done_sent = True
|
||||
|
||||
if t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
audio_chunk_count += 1
|
||||
if streaming:
|
||||
b64, audio_chunks_drained = self._extract_audio_delta_b64(
|
||||
output,
|
||||
audio_chunks_drained,
|
||||
)
|
||||
if b64:
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "response.output_audio.delta",
|
||||
"data": b64,
|
||||
"format": "wav",
|
||||
}
|
||||
)
|
||||
else:
|
||||
audio_data = self._get_audio_data(output)
|
||||
if audio_data is not None:
|
||||
if isinstance(audio_data, list):
|
||||
audio_tail_tensors.extend(audio_data)
|
||||
else:
|
||||
audio_tail_tensors.append(audio_data)
|
||||
else:
|
||||
delta_text, previous_text = self._extract_text_delta(
|
||||
output,
|
||||
audio_chunks_drained,
|
||||
previous_text,
|
||||
)
|
||||
if delta_text:
|
||||
if t_first_text is None:
|
||||
t_first_text = _time.monotonic()
|
||||
text_parts.append(delta_text)
|
||||
if streaming:
|
||||
await websocket.send_json({"type": "response.text.delta", "delta": delta_text})
|
||||
|
||||
if not text_done_sent:
|
||||
full_text = "".join(text_parts)
|
||||
await websocket.send_json({"type": "response.text.done", "text": full_text})
|
||||
text_done_sent = True
|
||||
|
||||
if not streaming and audio_tail_tensors:
|
||||
try:
|
||||
coalesced = audio_tail_tensors[0]
|
||||
if len(audio_tail_tensors) != 1:
|
||||
coalesced = torch.cat(audio_tail_tensors, dim=-1)
|
||||
tail_np = self._tensor_to_1d_np(coalesced)
|
||||
b64, _ = self._encode_tail(
|
||||
tail_np,
|
||||
0,
|
||||
new_drained=len(audio_tail_tensors),
|
||||
is_first=True,
|
||||
)
|
||||
if b64:
|
||||
await websocket.send_json(
|
||||
@@ -780,76 +1109,45 @@ class OmniStreamingVideoHandler:
|
||||
"format": "wav",
|
||||
}
|
||||
)
|
||||
else:
|
||||
audio_data = self._get_audio_data(output)
|
||||
if audio_data is not None:
|
||||
if isinstance(audio_data, list):
|
||||
audio_tail_tensors.extend(audio_data)
|
||||
else:
|
||||
audio_tail_tensors.append(audio_data)
|
||||
else:
|
||||
delta_text, previous_text = self._extract_text_delta(
|
||||
output,
|
||||
previous_text,
|
||||
)
|
||||
if delta_text:
|
||||
if t_first_text is None:
|
||||
t_first_text = _time.monotonic()
|
||||
text_parts.append(delta_text)
|
||||
if streaming:
|
||||
await websocket.send_json({"type": "response.text.delta", "delta": delta_text})
|
||||
except Exception:
|
||||
logger.exception("Failed to coalesce off-path audio")
|
||||
|
||||
if not text_done_sent:
|
||||
if audio_chunk_count > 0:
|
||||
await websocket.send_json({"type": "response.output_audio.done"})
|
||||
|
||||
response_text = "".join(text_parts)
|
||||
self.on_turn_complete(message_history, user_message, response_text)
|
||||
|
||||
t_end = _time.monotonic()
|
||||
logger.info(
|
||||
"[TIMING] mode=%s total=%.2fs first_text=%.2fs first_audio=%.2fs audio_chunks=%d",
|
||||
async_chunk_mode,
|
||||
t_end - t_start,
|
||||
(t_first_text - t_start) if t_first_text else -1,
|
||||
(t_first_audio - t_start) if t_first_audio else -1,
|
||||
audio_chunk_count,
|
||||
)
|
||||
|
||||
except Exception:
|
||||
logger.exception("Engine query failed")
|
||||
await self._send_error(websocket, "Query processing failed")
|
||||
|
||||
if not text_done_sent and not interrupt_event.is_set():
|
||||
full_text = "".join(text_parts)
|
||||
await websocket.send_json({"type": "response.text.done", "text": full_text})
|
||||
text_done_sent = True
|
||||
|
||||
if not streaming and audio_tail_tensors:
|
||||
try:
|
||||
coalesced = (
|
||||
audio_tail_tensors[0] if len(audio_tail_tensors) == 1 else torch.cat(audio_tail_tensors, dim=-1)
|
||||
)
|
||||
tail_np = self._tensor_to_1d_np(coalesced)
|
||||
b64, _ = self._encode_tail(
|
||||
tail_np,
|
||||
0,
|
||||
new_drained=len(audio_tail_tensors),
|
||||
is_first=True,
|
||||
)
|
||||
if b64:
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "response.output_audio.delta",
|
||||
"data": b64,
|
||||
"format": "wav",
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to coalesce off-path audio")
|
||||
|
||||
if audio_chunk_count > 0:
|
||||
await websocket.send_json({"type": "response.output_audio.done"})
|
||||
|
||||
response_text = "".join(text_parts)
|
||||
self.on_turn_complete(message_history, user_message, response_text)
|
||||
|
||||
t_end = _time.monotonic()
|
||||
logger.info(
|
||||
"[TIMING] mode=%s total=%.2fs first_text=%.2fs first_audio=%.2fs audio_chunks=%d",
|
||||
async_chunk_mode,
|
||||
t_end - t_start,
|
||||
(t_first_text - t_start) if t_first_text else -1,
|
||||
(t_first_audio - t_start) if t_first_audio else -1,
|
||||
audio_chunk_count,
|
||||
)
|
||||
|
||||
except Exception:
|
||||
logger.exception("Engine query failed")
|
||||
await self._send_error(websocket, "Query processing failed")
|
||||
|
||||
if not text_done_sent:
|
||||
full_text = "".join(text_parts)
|
||||
await websocket.send_json({"type": "response.text.done", "text": full_text})
|
||||
prompt_task = asyncio.ensure_future(self._preprocess_to_engine_prompt(chat_request))
|
||||
try:
|
||||
engine_prompt = await asyncio.shield(prompt_task)
|
||||
except asyncio.CancelledError:
|
||||
# Sender cache is updated in preprocess; submit before the caller aborts.
|
||||
engine_prompt = await asyncio.shield(prompt_task)
|
||||
await self._stop_after_submit(asyncio.ensure_future(_drive_query()), submitted)
|
||||
raise
|
||||
except Exception as e:
|
||||
await self._send_error(websocket, f"Preprocess failed: {e}")
|
||||
return
|
||||
await self._await_until_submitted(asyncio.ensure_future(_drive_query()), submitted)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Audio helpers
|
||||
@@ -1009,6 +1307,105 @@ class OmniStreamingVideoHandler:
|
||||
return text[len(previous_text) :], text
|
||||
return text, text
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Incremental prefill (context warmup)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _stop_after_submit(self, drive: asyncio.Future[Any], submitted: asyncio.Future[None]) -> None:
|
||||
"""Stage0 must observe the request before the caller aborts it."""
|
||||
if not submitted.done():
|
||||
submit_wait = asyncio.ensure_future(asyncio.shield(submitted))
|
||||
await asyncio.wait({submit_wait, drive}, return_when=asyncio.FIRST_COMPLETED)
|
||||
if not submitted.done():
|
||||
submit_wait.cancel()
|
||||
exc = drive.exception() if drive.done() else None
|
||||
if exc is not None:
|
||||
raise exc
|
||||
if not drive.done():
|
||||
drive.cancel()
|
||||
await asyncio.gather(drive, return_exceptions=True)
|
||||
|
||||
async def _await_until_submitted(self, drive: asyncio.Future[Any], submitted: asyncio.Future[None]) -> None:
|
||||
"""On cancel, wait until Stage0 has accepted the request before aborting."""
|
||||
try:
|
||||
# Shield so cancelling this task does not cancel generate before Stage0 submit.
|
||||
await asyncio.shield(drive)
|
||||
except asyncio.CancelledError:
|
||||
await self._stop_after_submit(drive, submitted)
|
||||
raise
|
||||
|
||||
async def _drive_warmup_generate(
|
||||
self,
|
||||
engine_client: Any,
|
||||
engine_prompt: Any,
|
||||
request_id: str,
|
||||
modalities: list[str],
|
||||
*,
|
||||
stop_after_submit: bool = False,
|
||||
) -> None:
|
||||
"""max_tokens=1 generate; cancellation still submits before abort."""
|
||||
from vllm import SamplingParams
|
||||
|
||||
submitted = asyncio.get_running_loop().create_future()
|
||||
result_gen = engine_client.generate(
|
||||
prompt=engine_prompt,
|
||||
request_id=request_id,
|
||||
sampling_params=SamplingParams(max_tokens=1, temperature=0.0),
|
||||
output_modalities=modalities,
|
||||
submitted=submitted,
|
||||
)
|
||||
|
||||
async def _drive() -> None:
|
||||
async for _ in result_gen:
|
||||
pass
|
||||
|
||||
drive = asyncio.ensure_future(_drive())
|
||||
if stop_after_submit:
|
||||
await self._stop_after_submit(drive, submitted)
|
||||
return
|
||||
await self._await_until_submitted(drive, submitted)
|
||||
|
||||
async def _prefill_context(
|
||||
self,
|
||||
config: StreamingVideoSessionConfig,
|
||||
messages: list[dict[str, Any]],
|
||||
request_id: str,
|
||||
) -> None:
|
||||
"""max_tokens=1 generate so arriving frames land in the prefix cache."""
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
|
||||
engine_client = self._engine_client
|
||||
if engine_client is None:
|
||||
return
|
||||
|
||||
request_kwargs: dict[str, Any] = {
|
||||
"model": config.model or "default",
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"modalities": config.modalities,
|
||||
"add_generation_prompt": True,
|
||||
"continue_final_message": False,
|
||||
"add_special_tokens": False,
|
||||
}
|
||||
if config.use_audio_in_video:
|
||||
request_kwargs["mm_processor_kwargs"] = {
|
||||
"use_audio_in_video": True,
|
||||
}
|
||||
chat_request = ChatCompletionRequest(**request_kwargs)
|
||||
prompt_task = asyncio.ensure_future(self._preprocess_to_engine_prompt(chat_request))
|
||||
try:
|
||||
engine_prompt = await asyncio.shield(prompt_task)
|
||||
except asyncio.CancelledError:
|
||||
# Sender cache is updated in preprocess; submit before the caller aborts.
|
||||
engine_prompt = await asyncio.shield(prompt_task)
|
||||
await self._drive_warmup_generate(
|
||||
engine_client, engine_prompt, request_id, config.modalities, stop_after_submit=True
|
||||
)
|
||||
raise
|
||||
await self._drive_warmup_generate(engine_client, engine_prompt, request_id, config.modalities)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Preprocessing
|
||||
# ------------------------------------------------------------------
|
||||
@@ -1043,6 +1440,18 @@ class OmniStreamingVideoHandler:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _candidate_frame_indices(frame_buffer: list[str], num_frames: int) -> list[int]:
|
||||
"""Stride positions for a legacy query, including the last frame.
|
||||
|
||||
Ignores decode status so a frame still decoding is waited on, not skipped.
|
||||
"""
|
||||
n_buf = len(frame_buffer)
|
||||
if n_buf <= num_frames:
|
||||
return list(range(n_buf))
|
||||
stride = max(1, n_buf // num_frames)
|
||||
return [index * stride for index in range(num_frames - 1)] + [n_buf - 1]
|
||||
|
||||
@staticmethod
|
||||
def _sample_frame_indices(
|
||||
frame_buffer: list[str],
|
||||
@@ -1050,13 +1459,26 @@ class OmniStreamingVideoHandler:
|
||||
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]
|
||||
indices = OmniStreamingVideoHandler._candidate_frame_indices(frame_buffer, num_frames)
|
||||
return [index for index in indices if _usable_prewarmed_frame(prewarmed_frames.get(frame_buffer[index]))]
|
||||
|
||||
def _select_frame_indices(
|
||||
self,
|
||||
config: StreamingVideoSessionConfig,
|
||||
frame_buffer: list[str],
|
||||
prewarmed_frames: Mapping[str, PrewarmedFrame],
|
||||
) -> list[int]:
|
||||
"""This turn's frame selection, shared by the prompt and the consumed event."""
|
||||
if self._incremental_prefill_active(config):
|
||||
# Frames are sampled at arrival time only (EVS + max_frames), so
|
||||
# the buffer already is the warmed sequence. Query-time
|
||||
# subsampling would pick a subset, which is not a prefix.
|
||||
return [
|
||||
index
|
||||
for index, frame_b64 in enumerate(frame_buffer)
|
||||
if _usable_prewarmed_frame(prewarmed_frames.get(frame_b64))
|
||||
]
|
||||
return self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed_frames)
|
||||
|
||||
@staticmethod
|
||||
async def _send_frame_ack(
|
||||
|
||||
Reference in new issue
Block a user