Files

1148 lines
46 KiB
Python

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""OmniInteract session execution and artifacts for serving benchmarks."""
from __future__ import annotations
import asyncio
import binascii
import contextlib
import copy
import fcntl
import hashlib
import json
import logging
import math
import os
import subprocess
import tempfile
import time
from collections import Counter
from collections.abc import Callable, Iterator, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from urllib.parse import urljoin, urlsplit
import pybase64 as base64
from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
DEFAULT_MAX_VIDEO_DURATION_S,
OmniInteractCase,
OmniInteractPreparedInput,
case_manifest,
sampled_case_row,
)
# The duplex client library (vllm_omni.clients) is imported lazily inside the
# functions that need it: this module reaches the CLI import graph through the
# benchmark patch package (`vllm-omni serve` included), and the client package
# must stay out of that graph
# (tests/engine/test_duplex_import_boundary.py enforces the boundary).
if TYPE_CHECKING:
from vllm_omni.clients.duplex import EventCollector, WebSocketTransport
# The duplex input wire format (16 kHz mono PCM16); mirrors
# vllm_omni.clients.duplex without importing it at module scope.
PCM16_SAMPLE_RATE = 16_000
PCM16_BYTES_PER_SAMPLE = 2
OUTPUT_SAMPLE_RATE = 24_000
SUCCESS_ARTIFACTS = (".done", "output.wav", "wav_transcript.json", "events.json", "result.json")
BATCH_ARTIFACTS = ("batch_summary.json", "official_eval_manifest.jsonl", "sampled_cases.jsonl")
ARTIFACT_LOCK_FILE = ".omniinteract.lock"
_INPUT_CHUNK_MS = 200
# Minimum new played audio between two incremental playback.ack sends for one
# response; completion acks always flush regardless of the step.
_ACK_STEP_MS = 200
VIDEO_FPS = 1.0
_COMPLETION_SETTLE_S = 2.0
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class OmniInteractBenchmarkConfig:
model: str
base_url: str = "http://127.0.0.1:8000"
endpoint: str = "/v1/realtime"
output_root: Path = Path("omniinteract-output")
timeout_s: float = 900.0
media_timeout_s: float = 600.0
max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S
ref_audio: str | None = None
require_response: bool = False
extra_headers: dict[str, str] | None = None
extra_body: dict[str, object] | None = None
@dataclass
class OmniInteractCaseResult:
subset: str
video: str
output_dir: str
success: bool = False
error: str = ""
session_id: str = ""
latency_s: float = 0.0
input_audio_chunks: int = 0
input_video_frames: int = 0
pacing_mean_lag_s: float = 0.0
pacing_max_lag_s: float = 0.0
responses: int = 0
audio_bytes: int = 0
audio_clipped_bytes: int = 0
transcript: str = ""
eligible_for_official_eval: bool = False
official_eval_ineligible_reasons: list[str] = field(default_factory=list)
artifact_warnings: list[str] = field(default_factory=list)
output_tokens: int = 0
duplex_request_metrics: list[dict[str, object]] = field(default_factory=list)
duplex_session_metrics: dict[str, object] = field(default_factory=dict)
_artifact_context: _DeferredArtifactContext | None = field(default=None, repr=False, compare=False)
def as_dict(self) -> dict[str, Any]:
return copy.deepcopy({key: value for key, value in vars(self).items() if not key.startswith("_")})
@dataclass(frozen=True)
class _DeferredArtifactContext:
horizon_bytes: int
spans: tuple[tuple[int, bytes], ...]
chunks: list[dict[str, object]]
events: list[dict[str, object]]
def benchmark_summary(results: list[OmniInteractCaseResult]) -> dict[str, Any]:
succeeded = sum(result.success for result in results)
eligible = sum(result.success and result.eligible_for_official_eval for result in results)
return {
"total": len(results),
"success": succeeded,
"failed": len(results) - succeeded,
"eligible_for_official_eval": eligible,
"successful_but_ineligible": succeeded - eligible,
"audio_clipped_bytes": sum(result.audio_clipped_bytes for result in results),
"results": [result.as_dict() for result in results],
}
def _run_media_command(
command: list[str],
*,
timeout_s: float,
text: bool = False,
) -> subprocess.CompletedProcess[Any]:
try:
result = subprocess.run(
command,
capture_output=True,
check=False,
text=text,
timeout=timeout_s,
)
except subprocess.TimeoutExpired as exc:
raise TimeoutError(f"Media command timed out after {timeout_s:g}s: {command[0]}") from exc
except FileNotFoundError as exc:
raise RuntimeError(f"Required media command is unavailable: {command[0]}") from exc
if result.returncode:
error = result.stderr if text else result.stderr.decode("utf-8", "ignore")
raise RuntimeError(f"{command[0]} failed: {error.strip()}")
return result
def prepare_media(
video: Path,
fps: float,
*,
timeout_s: float,
max_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S,
) -> tuple[float, bytes, list[str | None]]:
probe = [
"ffprobe",
*"-v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1".split(),
str(video),
]
duration = float(_run_media_command(probe, text=True, timeout_s=timeout_s).stdout.strip())
if not math.isfinite(duration) or duration <= 0 or duration > max_duration_s:
raise ValueError(f"Invalid video duration for {video}: {duration!r}")
target_bytes = math.ceil(duration) * PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
audio = [
"ffmpeg",
*"-loglevel error -i".split(),
str(video),
"-t",
str(duration),
*f"-vn -f s16le -ac 1 -ar {PCM16_SAMPLE_RATE} -fs".split(),
str(target_bytes),
"pipe:1",
]
pcm = _run_media_command(audio, timeout_s=timeout_s).stdout
pcm = (pcm + bytes(max(0, target_bytes - len(pcm))))[:target_bytes]
frame_count = math.ceil(duration * fps)
frames: list[str | None] = [None] * frame_count
with tempfile.TemporaryDirectory(prefix="vllm-omni-frames-") as temp_dir:
output_pattern = Path(temp_dir) / "frame-%06d.jpg"
frame_filter = f"select=gte(t\\,(selected_n+0.5)/{fps}),scale=640:640:force_original_aspect_ratio=decrease"
frame_command = [
"ffmpeg",
*"-loglevel error -i".split(),
str(video),
"-an",
"-vf",
frame_filter,
*f"-frames:v {frame_count} -vsync vfr -q:v 5".split(),
str(output_pattern),
]
_run_media_command(frame_command, timeout_s=timeout_s)
for index, frame_path in enumerate(sorted(Path(temp_dir).glob("frame-*.jpg"))[:frame_count]):
frames[index] = base64.b64encode(frame_path.read_bytes()).decode("ascii")
return duration, pcm, frames
@dataclass(frozen=True)
class _AudioSegment:
event_index: int
response_id: str
start_s: float
end_s: float
pcm16: bytes
class _Playback:
def __init__(self) -> None:
self.cursor = 0
self.end_s = 0.0
self.segments: list[_AudioSegment] = []
self.completed: set[str] = set()
self.completion_acked: set[str] = set()
self.warnings: list[str] = []
self._total_samples: dict[str, int] = {}
self._response_end_s: dict[str, float] = {}
self._acked_ms: dict[str, int] = {}
def _warn_once(self, warning: str) -> None:
if warning not in self.warnings:
self.warnings.append(warning)
def ingest(self, events: EventCollector) -> None:
while self.cursor < len(events.events):
index, self.cursor = self.cursor, self.cursor + 1
event = events.events[index]
response_id = events.response_id(event)
if event.get("type") == "response.done":
if response_id:
self.completed.add(response_id)
continue
if event.get("type") != "response.output_audio.delta":
continue
encoded = event.get("delta") or event.get("audio")
if not response_id:
raise ValueError("response audio has no response_id")
if not isinstance(encoded, str):
raise ValueError("response audio payload is missing")
if event.get("format") is None:
self._warn_once("response.output_audio.delta omitted format; assumed pcm16")
elif event.get("format") != "pcm16":
raise ValueError("OmniInteract output must be pcm16")
rate = event.get("sample_rate_hz")
if rate is None:
rate = events.output_sample_rate_hz or OUTPUT_SAMPLE_RATE
self._warn_once(f"response.output_audio.delta omitted sample_rate_hz; assumed {rate}")
if rate != OUTPUT_SAMPLE_RATE:
raise ValueError(f"OmniInteract output must use {OUTPUT_SAMPLE_RATE} Hz audio")
try:
raw = base64.b64decode(encoded, validate=True)
except (ValueError, binascii.Error) as exc:
raise ValueError("response audio is not valid base64") from exc
if not raw or len(raw) % PCM16_BYTES_PER_SAMPLE:
raise ValueError("response audio is empty or not PCM16 aligned")
start = max(events.event_received_at_s[index], self.end_s)
samples = len(raw) // PCM16_BYTES_PER_SAMPLE
segment = _AudioSegment(index, response_id, start, start + samples / rate, raw)
self.segments.append(segment)
self.end_s = segment.end_s
self._total_samples[response_id] = self._total_samples.get(response_id, 0) + samples
self._response_end_s[response_id] = segment.end_s
async def acknowledge(self, client: _RealtimeSession, now: float | None = None) -> None:
"""Ack playback progress incrementally along the serialized clock.
A real listener reports cumulative played audio while it plays, not
once after the fact. The first ack for a response goes out as soon as
its audio starts arriving, whatever the played amount: it checkpoints
the response's history position on the server, so a user input
committed later updates that slot in place instead of displacing it.
Later acks follow ``_ACK_STEP_MS``, with the exact total at drain.
"""
events = client.events
self.ingest(events)
now = time.monotonic() if now is None else now
for response_id, total_samples in self._total_samples.items():
if response_id in self.completion_acked:
continue
acked_ms = self._acked_ms.get(response_id)
if response_id in self.completed and now >= self._response_end_s[response_id]:
played_ms = total_samples * 1000 // OUTPUT_SAMPLE_RATE
if acked_ms is None or played_ms > acked_ms:
await client.send_playback_ack(response_id, played_ms)
self._acked_ms[response_id] = played_ms
self.completion_acked.add(response_id)
continue
played_ms = self._played_ms(response_id, now)
if acked_ms is not None and played_ms - acked_ms < _ACK_STEP_MS:
continue
await client.send_playback_ack(response_id, played_ms)
self._acked_ms[response_id] = played_ms
def _played_ms(self, response_id: str, now: float) -> int:
"""Cumulative audio of one response played by ``now`` on the serial clock."""
played_samples = 0
for segment in self.segments:
if segment.response_id != response_id or now <= segment.start_s:
continue
samples = len(segment.pcm16) // PCM16_BYTES_PER_SAMPLE
if now >= segment.end_s:
played_samples += samples
else:
played_samples += min(samples, int((now - segment.start_s) * OUTPUT_SAMPLE_RATE + 0.5))
return played_samples * 1000 // OUTPUT_SAMPLE_RATE
async def stream_inputs(
client: _RealtimeSession,
pcm: bytes,
frames: Sequence[str | None],
playback: _Playback,
) -> tuple[int, int, float, float]:
bytes_per_second = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
chunk_bytes = bytes_per_second * _INPUT_CHUNK_MS // 1000
started_at = time.monotonic()
frame_cursor = sent_frames = 0
lags: list[float] = []
for chunks, offset in enumerate(range(0, len(pcm), chunk_bytes), start=1):
end = min(offset + chunk_bytes, len(pcm))
end_ms = end * 1000 // bytes_per_second
ready: list[str] = []
while frame_cursor < len(frames) and end_ms >= (frame_cursor + 0.5) * 1000 / VIDEO_FPS:
if frames[frame_cursor]:
ready.append(frames[frame_cursor] or "")
frame_cursor += 1
sent_frames += len(ready)
event: dict[str, object] = {
"type": "input_audio_buffer.append",
"audio": base64.b64encode(pcm[offset:end]).decode("ascii"),
"input_audio_format": "pcm16",
"sample_rate_hz": PCM16_SAMPLE_RATE,
"duration_ms": (end - offset) * 1000 // bytes_per_second,
"audio_end_ms": end_ms,
}
if ready:
event["video_frames"] = ready
lags.append(max(0.0, time.monotonic() - started_at - offset / bytes_per_second))
await client.send(event)
await playback.acknowledge(client)
await asyncio.sleep(max(0.0, started_at + end_ms / 1000 - time.monotonic()))
return chunks if pcm else 0, sent_frames, sum(lags) / len(lags) if lags else 0.0, max(lags, default=0.0)
def _response_status(event: dict[str, object]) -> str | None:
response = event.get("response")
if isinstance(response, dict) and response.get("status"):
return str(response["status"])
status = event.get("status")
return str(status) if status else None
def response_ledger(
collector: EventCollector,
*,
end: int | None = None,
) -> tuple[set[str], set[str]]:
created: set[str] = set()
done: set[str] = set()
for event in collector.events[:end]:
event_type = event.get("type")
if event_type not in {"response.created", "response.done"}:
continue
response_id = collector.response_id(event)
if not response_id:
raise ValueError(f"{event_type} has no response_id")
if event_type == "response.created":
if response_id in created:
raise ValueError(f"duplicate response.created for {response_id}")
created.add(response_id)
continue
if response_id not in created:
raise ValueError(f"response.done without response.created for {response_id}")
if response_id in done:
raise ValueError(f"duplicate response.done for {response_id}")
response = event.get("response") if isinstance(event.get("response"), dict) else {}
details = event.get("status_details") if isinstance(event.get("status_details"), dict) else {}
nested = response.get("status_details") if isinstance(response.get("status_details"), dict) else {}
if "failed" in {event.get("status"), response.get("status"), details.get("type"), nested.get("type")}:
raise ValueError(f"response.done reports failure for {response_id}")
done.add(response_id)
return created, done
def _is_model_listen_decision(event: dict[str, object]) -> bool:
response = event.get("response")
metadata = response.get("metadata") if isinstance(response, dict) else event.get("metadata")
return isinstance(metadata, dict) and metadata.get("buffering") is not True and metadata.get("model_listen") is True
def _has_post_commit_decision(
collector: EventCollector,
events: list[dict[str, object]],
*,
prior_response_ids: set[str],
) -> bool:
pending_prior = set(prior_response_ids)
for event in events:
event_type = event.get("type")
if event_type == "response.done":
response_id = collector.response_id(event)
if response_id in pending_prior:
pending_prior.remove(response_id)
continue
if not pending_prior and _response_status(event) != "cancelled":
return True
elif event_type == "response.listen" and not pending_prior and _is_model_listen_decision(event):
return True
return False
# An ack computed just before a commit event was delivered can still lose the
# race server-side; the response it acks was already superseded, so the
# refusal is harmless — record it instead of failing the case.
_TOLERATED_ERROR_CODES = frozenset({"playback_ack_too_late"})
# ``session.close`` carries no reason of its own, so the server stamps the
# terminal event with the reason it inferred for it. Anything else on a
# ``session.closed`` after our own close request means the session ended for a
# reason we did not ask for (a disconnect, a lease timeout, a shutdown), which
# is what this guard is looking for.
_CLIENT_CLOSE_REASONS = frozenset({"client_close"})
def _error_code(event: dict[str, object]) -> str | None:
error = event.get("error")
code = error.get("code") if isinstance(error, dict) else event.get("code")
return code if isinstance(code, str) else None
def _raise_if_session_terminated(
collector: EventCollector,
from_index: int,
*,
explicit_close_from: int | None = None,
warnings: list[str] | None = None,
) -> None:
errors = collector.errors()
fatal = [event for event in errors if _error_code(event) not in _TOLERATED_ERROR_CODES]
if fatal:
raise RuntimeError(str(fatal[-1]))
if warnings is not None:
for event in errors:
warning = f"tolerated in-flight server rejection: {_error_code(event)}"
if warning not in warnings:
warnings.append(warning)
for index, event in enumerate(collector.events[from_index:], start=from_index):
event_type = event.get("type")
if event_type not in {"session.expired", "session.closed"}:
continue
reason = event.get("reason")
nested = event.get("event")
if reason is None and isinstance(nested, dict):
reason = nested.get("reason")
expected = (
event_type == "session.closed"
and explicit_close_from is not None
and index >= explicit_close_from
and (reason is None or reason in _CLIENT_CLOSE_REASONS)
)
if expected:
continue
detail = f": {reason}" if reason else ""
prefix = "Unexpected " if explicit_close_from is not None else ""
raise RuntimeError(f"{prefix}{event_type}{detail}")
def _ensure_final_commit_tail(pcm: bytes, events: list[dict[str, object]]) -> bytes:
"""Reserve one PCM sample so an exact model unit is flushed by commit.
A complete unit is otherwise emitted before commit, so its asynchronous
decision cannot be correlated with the accepted final input. Keeping an
almost-full unit buffered makes the server's final residual flush carry
that correlation; the server pads the missing sample back to a full unit.
"""
from vllm_omni.clients.duplex import chunk_period_ms, has_residual_model_unit
period_ms = chunk_period_ms(events)
if len(pcm) >= PCM16_BYTES_PER_SAMPLE and not has_residual_model_unit(pcm, chunk_period_ms=period_ms):
return pcm[:-PCM16_BYTES_PER_SAMPLE]
return pcm
async def wait_for_session_completion(
client: _RealtimeSession,
playback: _Playback,
*,
commit_from: int,
session_from: int | None = None,
timeout_s: float,
settle_s: float,
) -> int:
deadline = time.monotonic() + timeout_s
committed_index: int | None = None
prior_response_ids: set[str] = set()
last_event_count = len(client.events.events)
stable_since = time.monotonic()
while time.monotonic() < deadline:
client.raise_if_reader_stopped()
_raise_if_session_terminated(
client.events,
commit_from if session_from is None else session_from,
warnings=playback.warnings,
)
await playback.acknowledge(client)
if len(client.events.events) != last_event_count:
last_event_count = len(client.events.events)
stable_since = time.monotonic()
if committed_index is None:
committed_index = next(
(
index
for index in range(commit_from, len(client.events.events))
if client.events.events[index].get("type") == "input_audio_buffer.committed"
),
None,
)
if committed_index is not None:
committed_event = client.events.events[committed_index]
committed = committed_event.get("event")
if isinstance(committed, dict) and committed.get("overlap_deferred") is True:
prior_created, prior_done = response_ledger(client.events, end=committed_index)
prior_response_ids = prior_created - prior_done
stable_since = time.monotonic()
if committed_index is not None:
created, done = response_ledger(client.events)
post_commit = client.events.events[committed_index + 1 :]
decision = _has_post_commit_decision(
client.events,
post_commit,
prior_response_ids=prior_response_ids,
)
if created == done and decision and time.monotonic() - stable_since >= settle_s:
return committed_index
await asyncio.sleep(0.05)
missing: set[str] = set()
with contextlib.suppress(ValueError):
created, done = response_ledger(client.events)
missing = created - done
raise TimeoutError(
"Timed out waiting for committed input and stable responses"
+ (f"; unfinished response_ids={sorted(missing)}" if missing else "")
)
def _output_dir(root: Path, case: OmniInteractCase) -> Path:
relative = case.video_rel.replace("\\", "/")
stem = Path(relative).with_suffix("").as_posix().replace("/", "__")
digest = hashlib.sha256(f"{case.subset}/{relative}".encode()).hexdigest()[:8]
return root / case.subset / f"{stem}--{digest}"
@contextlib.contextmanager
def omniinteract_output_lock(root: Path) -> Iterator[None]:
root.mkdir(parents=True, exist_ok=True)
with (root / ARTIFACT_LOCK_FILE).open("a+b") as lock_file:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
def _atomic_replace(path: Path, writer: Callable[[Path], None]) -> None:
descriptor, temporary_name = tempfile.mkstemp(
dir=path.parent,
prefix=f".{path.name}.",
suffix=".tmp",
)
os.close(descriptor)
temporary = Path(temporary_name)
try:
writer(temporary)
temporary.replace(path)
finally:
temporary.unlink(missing_ok=True)
def _atomic_write_text(path: Path, value: str) -> None:
_atomic_replace(path, lambda temporary: temporary.write_text(value, encoding="utf-8"))
def _atomic_write_json(path: Path, value: object) -> None:
_atomic_write_text(path, json.dumps(value, ensure_ascii=False, indent=2) + "\n")
def _atomic_write_wav(path: Path, pcm: bytes | bytearray, rate: int) -> None:
from vllm_omni.clients.duplex import write_pcm16_wav
_atomic_replace(
path,
lambda temporary: write_pcm16_wav(
temporary,
cast(bytes, pcm), # wave.writeframes accepts bytearray.
sample_rate_hz=rate,
),
)
def _clear_artifacts(directory: Path, names: Sequence[str]) -> None:
directory.mkdir(parents=True, exist_ok=True)
for name in names:
(directory / name).unlink(missing_ok=True)
def _artifact_summary(case: OmniInteractCase, result: OmniInteractCaseResult) -> dict[str, Any]:
return {
**result.as_dict(),
"annotation": str(case.annotation_path),
"scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type,
}
def _publish_success_artifacts(
directory: Path,
result: OmniInteractCaseResult,
context: _DeferredArtifactContext,
summary: dict[str, Any],
) -> None:
pcm = bytearray(context.horizon_bytes)
for offset, chunk in context.spans:
pcm[offset : offset + len(chunk)] = chunk
try:
_clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json"))
_atomic_write_wav(directory / "output.wav", pcm, OUTPUT_SAMPLE_RATE)
_atomic_write_json(
directory / "wav_transcript.json",
{
"text": result.transcript,
"chunks": context.chunks,
"timestamp_semantics": "serialized playback queue time relative to input streaming start",
},
)
_atomic_write_json(directory / "events.json", context.events)
_atomic_write_json(directory / "result.json", result.as_dict())
_atomic_write_json(directory / ".done", summary)
except Exception:
_clear_artifacts(directory, SUCCESS_ARTIFACTS)
result.success = result.eligible_for_official_eval = False
if "artifact_write_failed" not in result.official_eval_ineligible_reasons:
result.official_eval_ineligible_reasons.append("artifact_write_failed")
raise
def _collect_output(
collector: EventCollector,
playback: _Playback,
*,
stream_start: float,
video_duration_s: float,
require_response: bool,
retain_events: bool,
result: OmniInteractCaseResult,
) -> _DeferredArtifactContext:
playback.ingest(collector)
created, done = response_ledger(collector)
if created != done:
raise ValueError(f"unfinished response_ids: {sorted(created - done)}")
done_events = [event for event in collector.events if event.get("type") == "response.done"]
cancelled = any(_response_status(event) == "cancelled" for event in done_events)
terminal_responses = {
response_id
for event in done_events
if _response_status(event) != "cancelled" and (response_id := collector.response_id(event))
}
horizon_s = math.ceil(video_duration_s)
horizon_bytes = horizon_s * OUTPUT_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
spans: list[tuple[int, bytes]] = []
clipped_bytes = 0
response_times: dict[str, list[float]] = {}
responses_with_audio: set[str] = set()
for segment in playback.segments:
response_id = segment.response_id
if response_id not in created:
raise ValueError("response audio has no matching response.created")
start_s = max(0.0, segment.start_s - stream_start)
end_s = max(start_s, segment.end_s - stream_start)
offset = round(start_s * OUTPUT_SAMPLE_RATE) * PCM16_BYTES_PER_SAMPLE
writable = max(0, min(len(segment.pcm16), horizon_bytes - offset))
clipped_bytes += len(segment.pcm16) - writable
if writable:
clipped = segment.pcm16[:writable]
spans.append((offset, clipped))
if any(clipped):
responses_with_audio.add(response_id)
timing = response_times.setdefault(response_id, [start_s, end_s])
timing[0], timing[1] = min(timing[0], start_s), max(timing[1], end_s)
text_event_types = {"response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta"}
text_ids = {collector.response_id(event) for event in collector.events if event.get("type") in text_event_types}
if None in text_ids or not text_ids <= created:
raise ValueError("response transcript has no matching response.created")
chunks: list[dict[str, object]] = []
texts: list[str] = []
for response_id in collector.response_ids:
text = collector.response_text(response_id).strip()
if not text:
continue
response_timing = response_times.get(response_id)
if response_timing is None or response_timing[0] >= horizon_s:
continue
texts.append(text)
chunks.append(
{
"response_id": response_id,
"text": text,
"timestamp": [round(response_timing[0], 6), round(min(response_timing[1], horizon_s), 6)],
}
)
complete_outputs = terminal_responses & responses_with_audio & {str(chunk["response_id"]) for chunk in chunks}
if require_response and not complete_outputs:
raise ValueError("OmniInteract E2E requires a response with audio and transcript")
result.audio_bytes, result.audio_clipped_bytes = (
sum(len(segment.pcm16) for segment in playback.segments),
clipped_bytes,
)
result.transcript, result.responses, result.success = " ".join(texts).strip(), len(created), True
result.artifact_warnings = list(playback.warnings)
result.official_eval_ineligible_reasons = []
if clipped_bytes:
result.official_eval_ineligible_reasons.append("audio_clipped")
if cancelled:
result.official_eval_ineligible_reasons.append("cancelled_response")
result.eligible_for_official_eval = not result.official_eval_ineligible_reasons
audio_sizes = {segment.event_index: len(segment.pcm16) for segment in playback.segments}
events = (
[
{
**{key: value for key, value in event.items() if key not in {"delta", "audio"}},
"audio_bytes": audio_sizes[index],
}
if event.get("type") == "response.output_audio.delta"
else dict(event)
for index, event in enumerate(collector.events)
]
if retain_events
else []
)
return _DeferredArtifactContext(horizon_bytes, tuple(spans), chunks, events)
def prepare_success_artifacts(
root: Path,
case: OmniInteractCase,
collector: EventCollector,
*,
playback: _Playback | None = None,
stream_start: float,
video_duration_s: float,
require_response: bool,
result: OmniInteractCaseResult,
capture_artifacts: bool = True,
) -> dict[str, Any]:
directory = _output_dir(root, case)
context = _collect_output(
collector,
playback or _Playback(),
stream_start=stream_start,
video_duration_s=video_duration_s,
require_response=require_response,
retain_events=capture_artifacts,
result=result,
)
result.output_dir = str(directory.resolve())
summary = _artifact_summary(case, result)
if capture_artifacts:
result._artifact_context = context
return summary
def write_failure_artifacts(root: Path, case: OmniInteractCase, result: OmniInteractCaseResult) -> None:
directory = _output_dir(root, case)
result.output_dir = str(directory.resolve())
result.success = result.eligible_for_official_eval = False
_clear_artifacts(directory, SUCCESS_ARTIFACTS)
_atomic_write_json(directory / ".failed.json", _artifact_summary(case, result))
def publish_deferred_case_artifacts(
root: Path,
case: OmniInteractCase,
result: OmniInteractCaseResult,
) -> None:
context = result._artifact_context
result._artifact_context = None
if not result.success:
write_failure_artifacts(root, case, result)
return
if context is None:
raise RuntimeError("OmniInteract benchmark output lost its deferred artifact context")
directory = _output_dir(root, case)
_publish_success_artifacts(directory, result, context, _artifact_summary(case, result))
def clear_case_artifacts(root: Path, case: OmniInteractCase) -> None:
directory = _output_dir(root, case)
_clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json"))
def clear_batch_artifacts(root: Path) -> None:
_clear_artifacts(root, BATCH_ARTIFACTS)
def write_batch_artifacts(
root: Path,
cases: list[OmniInteractCase],
results: list[OmniInteractCaseResult],
) -> None:
root.mkdir(parents=True, exist_ok=True)
rows = [
case_manifest(case, _output_dir(root, case))
for case, result in zip(cases, results, strict=True)
if result.success and result.eligible_for_official_eval
]
sampled_rows = [sampled_case_row(case) for case in cases]
ineligible = [result for result in results if result.success and not result.eligible_for_official_eval]
if ineligible:
reasons = Counter(
reason for result in ineligible for reason in (result.official_eval_ineligible_reasons or ["unspecified"])
)
logger.warning(
"%d successful OmniInteract cases were excluded from official evaluation: %s",
len(ineligible),
", ".join(f"{reason}={count}" for reason, count in sorted(reasons.items())),
)
_atomic_write_json(root / "batch_summary.json", benchmark_summary(results))
_atomic_write_text(
root / "official_eval_manifest.jsonl",
"".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
)
# Every sampled case (including failed / ineligible) so official MiniCPM-o
# ``--video_list`` can replay the same video set without re-deriving ids.
_atomic_write_text(
root / "sampled_cases.jsonl",
"".join(json.dumps(row, ensure_ascii=False) + "\n" for row in sampled_rows),
)
def _websocket_url(config: OmniInteractBenchmarkConfig) -> str:
endpoint = (
config.endpoint
if urlsplit(config.endpoint).scheme
else urljoin(config.base_url.rstrip("/") + "/", config.endpoint.lstrip("/"))
)
from vllm_omni.clients.duplex import build_realtime_url
return build_realtime_url(endpoint, config.model, autostart=False)
class _RealtimeSession:
"""The public :class:`DuplexClient` behind the benchmark's session surface.
The case runner works in event-cursor style over one growing collector,
so this wrapper pairs a ``DuplexClient`` with an :class:`EventCollector`
fed by a background consumer task and exposes the few calls the runner
needs: raw sends, the final commit, cumulative playback acks, reader
liveness checks, and a confirmed close.
"""
_MAX_FRAME_BYTES = 64 * 1024 * 1024
def __init__(self, config: OmniInteractBenchmarkConfig, reference_audio: str) -> None:
from vllm_omni.clients.duplex import DuplexClient, EventCollector, SessionConfig
self.session_config = SessionConfig(
ref_audio=reference_audio,
overlap_policy="listen_only",
playback_commit_policy="ack_only",
idle_timeout_s=float(config.timeout_s),
extra_body={
"force_listen_count": 0,
**(config.extra_body or {}),
},
)
self.url = _websocket_url(config)
headers = dict(config.extra_headers or {})
async def connect(url: str) -> WebSocketTransport:
import websockets
return await websockets.connect(
url,
max_size=self._MAX_FRAME_BYTES,
additional_headers=headers or None,
)
self.events = EventCollector()
self._client = DuplexClient(
self.url,
model=config.model,
config=self.session_config,
reconnect=None,
heartbeat_interval_s=None,
handshake_timeout_s=min(config.timeout_s, 20.0),
connect=connect,
)
self._consume_task: asyncio.Task[None] | None = None
async def __aenter__(self) -> _RealtimeSession:
await self._client.__aenter__()
self._consume_task = asyncio.create_task(self.events.consume(self._client))
# The handshake consumed session.created before the collector
# subscribed; seed it so capability lookups and cursors still work.
self.events.add({"type": "session.created", "session": dict(self._client.session_info)})
return self
async def __aexit__(self, exc_type, exc, traceback) -> None:
try:
await self._client.__aexit__(exc_type, exc, traceback)
finally:
if self._consume_task is not None:
self._consume_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self._consume_task
async def send(self, event: dict[str, object]) -> None:
await self._client.send(event)
async def commit(self) -> None:
await self._client.commit()
async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
await self._client.ack_playback(
played_ms,
response_id=response_id,
item_id=f"item_{response_id}",
)
def raise_if_reader_stopped(self) -> None:
"""Fail a caller waiting on events after the event consumer exits."""
task = self._consume_task
if task is None or not task.done():
return
if task.cancelled():
raise ConnectionError("Realtime event consumer was cancelled")
error = task.exception()
if error is not None:
raise ConnectionError("Realtime session failed") from error
raise ConnectionError("Realtime session closed before the requested event arrived")
async def close_session(self, *, timeout_s: float = 20.0) -> None:
await self._client.close(timeout_s=timeout_s)
if self._consume_task is not None:
# Let the collector drain session.closed before callers check it.
with contextlib.suppress(asyncio.TimeoutError, asyncio.CancelledError):
await asyncio.wait_for(asyncio.shield(self._consume_task), timeout=5.0)
def _audio_rtf_from_raw_metric(raw_metric: dict[str, object]) -> float | None:
"""Derive the audio RTF from a client-reported raw request metric."""
from vllm_omni.benchmarks.duplex_session_metrics import audio_rtf_from_raw_metric
return audio_rtf_from_raw_metric(raw_metric)
def _populate_response_metrics(
result: OmniInteractCaseResult,
collector: EventCollector,
*,
stream_start: float,
) -> None:
from vllm_omni.benchmarks.duplex_session_metrics import collect_duplex_session_metrics
bundle = collect_duplex_session_metrics(
collector,
stream_start=stream_start,
session_id=result.session_id,
)
result.output_tokens = bundle.output_tokens
result.duplex_request_metrics = bundle.request_metrics
result.duplex_session_metrics = bundle.session_metrics
def _session_capabilities(client: _RealtimeSession) -> dict[str, object]:
info = getattr(getattr(client, "_client", None), "session_info", None)
if isinstance(info, dict) and isinstance(info.get("capabilities"), dict):
return info["capabilities"]
for event in client.events.events:
if event.get("type") != "session.created":
continue
session = event.get("session")
if isinstance(session, dict) and isinstance(session.get("capabilities"), dict):
return session["capabilities"]
return {}
def _wants_annotation_video_clock(client: _RealtimeSession) -> bool:
"""Video-required turn-commit sessions use the AURA annotation clock.
The shared 1 FPS soundtrack path is unchanged unless both capability bits
are set. The AURA module is imported only in that case.
"""
caps = _session_capabilities(client)
required = caps.get("required_input_modalities") or ()
if isinstance(required, str):
required = (required,)
try:
video_required = "video" in {str(item) for item in required}
except TypeError:
return False
if not video_required or not caps.get("supports_turn_commit_only"):
return False
from vllm_omni.model_executor.models.aura_omni.benchmarks.omniinteract_clock import (
wants_annotation_video_clock,
)
return wants_annotation_video_clock(caps)
async def run_omniinteract_case(
case: OmniInteractCase,
config: OmniInteractBenchmarkConfig,
*,
request_index: int | str,
capture_artifacts: bool = True,
prepared_input: OmniInteractPreparedInput | None = None,
) -> OmniInteractCaseResult:
if not config.ref_audio:
raise ValueError("ref_audio is required for MiniCPM-o native-duplex audio output")
for name, value in (
("timeout_s", config.timeout_s),
("media_timeout_s", config.media_timeout_s),
("max_video_duration_s", config.max_video_duration_s),
):
if not math.isfinite(value) or value <= 0:
raise ValueError(f"{name} must be finite and positive")
session_id = f"omniinteract:{case.subset}:{request_index}:{time.monotonic_ns()}"
result = OmniInteractCaseResult(
case.subset,
str(case.video_path),
str(_output_dir(config.output_root, case).resolve()),
session_id=session_id,
)
started_at = time.monotonic()
try:
if prepared_input is None:
from vllm_omni.clients.duplex import reference_audio_data_url
reference_audio = reference_audio_data_url(config.ref_audio)
duration, pcm, frames = await asyncio.to_thread(
prepare_media,
case.video_path,
VIDEO_FPS,
timeout_s=config.media_timeout_s,
max_duration_s=config.max_video_duration_s,
)
else:
reference_audio = prepared_input.ref_audio_data_url
duration = prepared_input.duration_s
pcm = prepared_input.pcm16
frames = prepared_input.video_frames
if not any(frames):
raise ValueError(f"No video frames were decoded from {case.video_path}")
async with _RealtimeSession(config, reference_audio) as client:
session_from = 0 # the collector holds only this session's events
pcm = _ensure_final_commit_tail(pcm, client.events.events)
playback = _Playback()
stream_start = time.monotonic()
try:
input_duration_s = len(pcm) / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
upload_timeout_s = config.timeout_s + input_duration_s
if _wants_annotation_video_clock(client):
from vllm_omni.model_executor.models.aura_omni.benchmarks.omniinteract_clock import (
ANNOTATION_VIDEO_FPS,
stream_annotation_clock,
)
duration, _, frames = await asyncio.to_thread(
prepare_media,
case.video_path,
ANNOTATION_VIDEO_FPS,
timeout_s=config.media_timeout_s,
max_duration_s=config.max_video_duration_s,
)
if not any(frames):
raise ValueError(f"No video frames were decoded from {case.video_path}")
upload_timeout_s = config.timeout_s + duration
stream = stream_annotation_clock(
client,
frames,
playback,
video_path=case.video_path,
fps=ANNOTATION_VIDEO_FPS,
)
logger.info(
"OmniInteract annotation video clock: video=%s fps=%s",
case.video_path.name,
ANNOTATION_VIDEO_FPS,
)
else:
stream = stream_inputs(client, pcm, frames, playback)
try:
chunks, frame_count, mean_lag, max_lag = await asyncio.wait_for(stream, timeout=upload_timeout_s)
except asyncio.TimeoutError as exc:
raise TimeoutError(f"Realtime upload timed out after {upload_timeout_s:g}s") from exc
commit_from = len(client.events.events)
await client.commit()
await wait_for_session_completion(
client,
playback,
commit_from=commit_from,
session_from=session_from,
timeout_s=config.timeout_s,
settle_s=_COMPLETION_SETTLE_S,
)
await playback.acknowledge(client, playback.end_s)
_raise_if_session_terminated(client.events, session_from, warnings=playback.warnings)
close_from = len(client.events.events)
await client.close_session(timeout_s=min(config.timeout_s, 20.0))
except Exception:
with contextlib.suppress(Exception):
await client.close_session(timeout_s=min(config.timeout_s, 20.0))
raise
_raise_if_session_terminated(
client.events,
session_from,
explicit_close_from=close_from,
warnings=playback.warnings,
)
result.latency_s = time.monotonic() - started_at
result.input_audio_chunks, result.input_video_frames = chunks, frame_count
result.pacing_mean_lag_s, result.pacing_max_lag_s = mean_lag, max_lag
_populate_response_metrics(result, client.events, stream_start=stream_start)
await asyncio.to_thread(
prepare_success_artifacts,
config.output_root,
case,
client.events,
playback=playback,
stream_start=stream_start,
video_duration_s=duration,
require_response=config.require_response,
result=result,
capture_artifacts=capture_artifacts,
)
except Exception as exc:
result.success = result.eligible_for_official_eval = False
if "case_failed" not in result.official_eval_ineligible_reasons:
result.official_eval_ineligible_reasons.append("case_failed")
result.error = str(exc)
result.latency_s = max(0.0, time.monotonic() - started_at)
return result