From 856d19f42cfccb1ae97aea1131b89662368baadb Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Mon, 21 Sep 2026 21:12:00 +0800 Subject: [PATCH] [Benchmark] Add LingBot-World realtime streaming benchmark (#7645) Signed-off-by: tzhouam Co-authored-by: Claude Opus 5 (1M context) --- benchmarks/README.md | 11 + benchmarks/lingbot_world/README.md | 224 +++++ benchmarks/lingbot_world/__init__.py | 2 + .../benchmark_lingbot_world_realtime.py | 642 ++++++++++++++ .../configs/single_gpu_eager.yaml | 45 + .../lingbot_world/configs/usp4_compiled.yaml | 60 ++ benchmarks/lingbot_world/workload.py | 631 ++++++++++++++ .../benchmarks/test_lingbot_world_realtime.py | 798 ++++++++++++++++++ 8 files changed, 2413 insertions(+) create mode 100644 benchmarks/lingbot_world/README.md create mode 100644 benchmarks/lingbot_world/__init__.py create mode 100644 benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py create mode 100644 benchmarks/lingbot_world/configs/single_gpu_eager.yaml create mode 100644 benchmarks/lingbot_world/configs/usp4_compiled.yaml create mode 100644 benchmarks/lingbot_world/workload.py create mode 100644 tests/benchmarks/test_lingbot_world_realtime.py diff --git a/benchmarks/README.md b/benchmarks/README.md index 182bba02d..1a23b1d48 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -41,6 +41,17 @@ LingBot-Video MoE transformer. - **Dense pipeline**: decoded-video MAE, MSE, PSNR, latency, and optional steady-state timings - **MoE transformer**: bitwise router, sparse-block, shared-expert, and full-transformer parity +### [LingBot-World](lingbot_world/README.md) — Realtime Streaming Cadence + +Client-side benchmark for a `WS /v1/realtime/video` world-model session, where one +request is a whole rollout and one AR block streams out as one video chunk. It +measures a cadence rather than a throughput: there is exactly one session, and what +matters is whether the next chunk arrives before the viewer finishes the last one. + +- **Layout**: `lingbot_world/benchmark_lingbot_world_realtime.py` (WS client), `lingbot_world/workload.py` (rollout and metric math), `lingbot_world/configs/` (single-GPU eager and Ulysses-4 compiled deploy configs) +- **Workload**: built-in image-conditioned rollout with a generated camera script, or a JSON spec with per-chunk actions and mid-rollout prompt updates +- **Key metrics**: TTFC (time to first chunk), steady-state chunk inter-arrival percentiles, VIDEO_RTF (wall seconds per video second; lower is better), chunk-deadline attainment, and simulated playback underruns + ### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks. diff --git a/benchmarks/lingbot_world/README.md b/benchmarks/lingbot_world/README.md new file mode 100644 index 000000000..1dce2d7ce --- /dev/null +++ b/benchmarks/lingbot_world/README.md @@ -0,0 +1,224 @@ +# LingBot-World Benchmarks — Realtime Streaming Cadence + +`benchmark_lingbot_world_realtime.py` measures LingBot-World 2.0 as it is +actually deployed: one `WS /v1/realtime/video` session that streams one video +chunk per AR block, for as long as the world keeps running. + +## Why this is not the diffusion serving benchmark + +[`benchmarks/diffusion`](../diffusion/README.md) measures "one request, one +finished video" and reports request throughput. LingBot-World has neither shape: + +- **One request is the whole rollout.** `prepare_encode` runs once, then every AR + block is four DMD steps plus one clean-KV commit forward, and each block + streams out as a chunk. Throughput in requests per second is meaningless. +- **There is exactly one session.** The stepwise deploy config pins + `max_num_seqs: 1`, and the AR-Diffusion consumer rejects a stage with more than + one replica, because session-affine routing does not exist yet. There is no + concurrency to sweep. +- **What matters is cadence.** A deployment is healthy when chunk *N + 1* arrives + before a viewer finishes watching chunk *N*. That is a latency distribution + against a wall clock, not a throughput number. + +## Quick start + +Serve the model, then point the benchmark at it. The benchmark never starts a +server. + +```bash +# Terminal 1 — one GPU, eager (portable smoke configuration) +vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \ + --omni \ + --deploy-config benchmarks/lingbot_world/configs/single_gpu_eager.yaml \ + --port 8000 + +# Terminal 2 +pip install websockets +python benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py \ + --host 127.0.0.1 --port 8000 \ + --num-chunks 16 \ + --output-json /tmp/lingbot_world_realtime.json +``` + +With no `--image`, the first frame comes from the shared vLLM image asset, so the +default workload needs no setup and this repository carries no benchmark binary. +Pass `--image /path/to/frame.png` to condition the rollout on your own frame. + +### Realtime topology + +The single-GPU config is a smoke configuration, not a realtime one: one card runs +all five forwards of every block serially and nothing is compiled. Use the +four-GPU config for numbers that mean anything about realtime behavior. + +```bash +vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \ + --omni \ + --deploy-config benchmarks/lingbot_world/configs/usp4_compiled.yaml \ + --port 8000 + +python benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py \ + --port 8000 --num-chunks 40 --warmup-sessions 1 --sessions 3 \ + --target-fps 12 --output-json /tmp/lingbot_world_repeated.json +``` + +`usp4_compiled.yaml` sets pure Ulysses sequence parallelism over four ranks and +`enforce_eager: false`. The pipeline validates the parallel shape itself: +`sequence_parallel_size` must equal `ulysses_degree`, `ring_degree` and +`allgather_degree` must be 1, `ulysses_mode` must be `strict`, and pipeline, +CFG, and VAE parallel sizes above 1 are rejected. This benchmark config sets +`tensor_parallel_size: 1`; the pipeline validator does not reject larger TP sizes. + +Compiled mode can pay compilation and capture costs on the first rollout. +Use `--warmup-sessions 1` before measured repetitions. Excluding the first six +chunks addresses the attention-window ramp; it does not guarantee that +compilation and capture have finished. + +## Metrics + +| Metric | Meaning | +| --- | --- | +| `ttfc_ms` | **Time to first chunk**: `session.start` sent until the first chunk's bytes arrive. Pays for prompt encode, first-frame VAE encode, and any compile or capture. | +| `interval_all` | Inter-arrival latency across every chunk: mean, median, std, min, p90, p95, p99, max. | +| `interval_steady` | The same summary over steady-state chunks only. | +| `video_rtf` | Wall seconds spent per second of video, matching vLLM-Omni's own RTF (`generation_s / video_duration`). **Lower is better; below 1.0 is real time.** | +| `video_rtfx` | The reciprocal, under the name this repo uses for it: video seconds per wall second, where higher is better. | +| `steady_video_rtf` | `video_rtf` restricted to steady-state chunks. | +| `chunk_deadline_ms` | Wall-clock budget one steady chunk must meet to hold real time: 12 frames ÷ `--target-fps`. Override directly with `--chunk-slo-ms`. | +| `slo_attainment` | Fraction of steady chunks that met the deadline. | +| `playback` | Underrun count, total stall, and stall ratio from a simulated viewer. | + +### Which frame rate counts as real time + +**The checkpoint declares no frame rate.** There is none in `model_index.json`, the +scheduler config, or the transformer config, and the pipeline never reads one. A +chunk is twelve pixel frames; how long those frames *last* is a property of how the +output is played, not of the model. + +That choice moves the verdict without moving anything measured: + +| basis | video per chunk | deadline | RTF at a measured 1357 ms/chunk | +| --- | --- | --- | --- | +| 16 fps | 0.750 s | 750 ms | 1.809 | +| 12 fps | 1.000 s | 1000 ms | 1.357 | + +So `--fps` (sent to the server, which uses it to label the muxed MP4) and +`--target-fps` (the real-time basis for RTF, the deadline, and the playback +simulation) are separate flags. `--target-fps` defaults to `--fps`; set it when the +two differ, and the report prints both. State the basis whenever you quote a number +— an RTF without its fps is not a measurement. (RTF is a cost here, so a bigger number is a +worse result; a steady chunk's RTF is simply its interval divided by its video duration.) +An RTF above 1 means the measured deployment is slower than real time at the +stated target frame rate. The benchmark measures this failure as well as success. + +### Why warmup chunks are not optional + +The DiT attends over a sliding window of `sliding_window_num_frames = 18` latent +frames with `sink_size = 9`, so the key count a block attends over keeps growing +until eighteen latent frames of history exist — six chunks at three frames per +block. **Chunk cost is not stationary before chunk six.** A mean over a short +rollout is therefore lower than anything a running session will sustain, which is +why `--warmup-chunks` defaults to 6 and steady-state metrics are reported +separately. Keep `--num-chunks` comfortably above the warmup, or the benchmark +tells you no chunk reached steady state. The terminal chunk skips next-chunk +preparation, so it is excluded from steady metrics too; it remains in overall +RTF, all-chunk intervals, and playback simulation. + +The default 16-chunk rollout provides only nine steady intervals. Percentiles +are interpolated sample summaries; the report warns when fewer than 100 +intervals contribute, and even 100 samples do not guarantee a stable p99. +Use longer rollouts for tail analysis and at least `--sessions 3` after +`--warmup-sessions 1` for reported comparisons. Session warmup pays compilation +cost; chunk warmup separately excludes the attention-window ramp. The command +above yields 33 steady intervals per session (indices 6 through 38), or 99 across +three sessions. Report each session's mean and the spread across sessions; this +sample still does not establish a reliable p99. + +### Why playback is simulated + +A p99 interval alone cannot say whether a viewer saw a stall: the same p99 is +invisible behind a three-chunk buffer and a visible freeze behind none. The +benchmark replays chunk arrivals against a wall clock — playback starts once +`--playback-buffer-chunks` chunks have landed and then consumes video at `--target-fps` +— and reports every moment the player ran dry. + +### What the client-side number is worth + +On a four-GPU Ulysses run the client's mean inter-arrival over all fifteen +post-first-chunk intervals was **1326.743 ms**, against the server's own +`StageRequestStats.inter_output_latency_ms` of **1327.479 ms** for the same +request — a difference of **0.74 ms**. This cross-check shows agreement in +average output cadence for that run, including warmup and the terminal chunk. +It does not measure encoding or transport latency: a fixed delivery delay can +cancel out of inter-arrival differences. Attributing that latency requires +per-chunk timestamps at both boundaries or server-side profiling. + +Compare matching interval populations when investigating discrepancies; both +measurement bugs and delivery buffering can change the observed cadence. + +### What this benchmark cannot see + +Chunk metadata carries no server-side timestamps, so the benchmark cannot split +DiT time from VAE decode time; use `--enable-diffusion-pipeline-profiler` on the +server for that. It also reports no GPU memory, which is a server-side quantity. + +## Custom rollouts + +`--workload rollout.json` replaces the built-in workload, including mid-rollout +prompt changes, which the CLI flags cannot express: + +```json +{ + "prompt": "The camera moves slowly forward through the scene.", + "image": "/path/to/first_frame.png", + "num_chunks": 24, + "camera_pattern": "orbit", + "fps": 16, + "seed": 42, + "prompt_updates": [ + {"after_chunk": 8, "prompt": "Rain begins to fall", "transition_chunks": 2} + ] +} +``` + +Supply `camera_action_script` directly instead of `num_chunks` to control every +frame's action. It must hold exactly one three-entry action list per chunk; one +AR block is three latent frames, and the server rejects any other shape. +Explicit `--image`, `--prompt`, `--negative-prompt`, `--width`, `--height`, +`--fps`, `--seed`, and `--flow-shift` flags override the matching workload-file +values. `--num-chunks` and `--camera-pattern` apply only to the built-in +workload; when using `--workload`, set those fields in the file instead. +An explicit `camera_action_script` takes precedence over the file's +`num_chunks` and `camera_pattern`. Each prompt update must use a distinct +`after_chunk` boundary so its interaction event ID is unique. + +Camera actions change what the world does, not what it costs, so `--camera-pattern` +(`forward`, `orbit`, `hold`) keeps a rollout representative rather than sweeping a +cost dimension. + +## Options worth knowing + +| Flag | Purpose | +| --- | --- | +| `--target-fps` | Playback rate the real-time verdict is measured against, independent of the mux label. | +| `--sessions N` | N sequential measured rollouts; the report adds spread across them. | +| `--warmup-sessions N` | Unmeasured rollouts first, to pay compile and capture cost. | +| `--print-chunks` | One line per chunk as it arrives. | +| `--save-video PATH` | Remux the last measured session for eyeballing (needs vLLM-Omni). | +| `--ping-interval` | `session.ping` cadence keeping the server's stall clock fresh during a slow first chunk. `0` disables. | +| `--first-chunk-timeout` | Raise it when the server is still compiling or capturing graphs. | + +## Output + +`--output-json` writes the run's configuration, every chunk's arrival time and +size, per-session metrics, and the aggregate, along with the git revision and +hostname so a number can be traced back to what produced it. + +## Tests + +`tests/benchmarks/test_lingbot_world_realtime.py` covers the chunk and frame +arithmetic against the server's own formula, the metric math, and the protocol +handling against a scripted WebSocket server. It needs no GPU: + +```bash +pytest tests/benchmarks/test_lingbot_world_realtime.py +``` diff --git a/benchmarks/lingbot_world/__init__.py b/benchmarks/lingbot_world/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/benchmarks/lingbot_world/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py b/benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py new file mode 100644 index 000000000..186bc1393 --- /dev/null +++ b/benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py @@ -0,0 +1,642 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Realtime streaming benchmark for LingBot-World 2.0 over ``WS /v1/realtime/video``. + +Every other diffusion benchmark in this repository measures "one request, one +finished video". LingBot-World is not that: one request is a whole rollout, and +the thing a deployment lives or dies by is whether chunk *N + 1* keeps arriving +before a viewer has finished watching chunk *N*. This benchmark therefore +measures a **cadence**, not a throughput: + +* ``TTFC`` -- session.start to the first chunk's bytes, which pays for prompt + encode, the first-frame VAE encode, and any compile or capture the deploy + config asks for; +* **inter-arrival latency** over all chunks and, separately, over the steady + state after the attention window saturates; +* ``VIDEO_RTF`` -- wall seconds spent per second of video generated, matching how + vLLM-Omni computes RTF everywhere else (``generation_s / video_duration``). + **Lower is better and below 1.0 is real time**; ``VIDEO_RTFX`` is reported + alongside it as the reciprocal, under the repo's own name for that; +* **playback continuity** -- a simulated viewer with a start buffer, because a + p99 that is harmless behind a three-chunk buffer is a visible stall behind + none. + +The server is not started here. Bring one up first (see the README), then point +this script at it. Requires ``websockets``; ``--save-video`` additionally needs +a vLLM-Omni installation for the remux helper. + +Example: + + python benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py \\ + --host 127.0.0.1 --port 8000 \\ + --num-chunks 16 \\ + --output-json /tmp/lingbot_world_realtime.json +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import json +import mimetypes +import os +import platform +import subprocess +import sys +import time +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from benchmarks.lingbot_world.workload import ( # noqa: E402 + CAMERA_PATTERNS, + DEFAULT_FLOW_SHIFT, + DEFAULT_FPS, + DEFAULT_HEIGHT, + DEFAULT_MODEL, + DEFAULT_PROMPT, + DEFAULT_SEED, + DEFAULT_WARMUP_CHUNKS, + DEFAULT_WIDTH, + ChunkRecord, + Workload, + aggregate_metrics, + compute_metrics, + default_workload, + load_workload, + steady_chunk_deadline_ms, +) + +# The shared vLLM asset keeps the default workload runnable with no setup and +# adds no binary to this repository. Any real photograph conditions a rollout. +_DEFAULT_IMAGE_ASSET = "2560px-Gfp-wisconsin-madison-the-nature-boardwalk" +_WS_PATH = "/v1/realtime/video" + + +class BenchmarkError(RuntimeError): + """A session failed in a way that invalidates its numbers.""" + + +# --------------------------------------------------------------------------- # +# Input resolution +# --------------------------------------------------------------------------- # + + +def resolve_image_reference(value: str | None) -> str: + """Return an ``http(s)`` or ``data:`` URL for the session's first frame.""" + if value is None: + from vllm.assets.image import ImageAsset + + value = str(ImageAsset(_DEFAULT_IMAGE_ASSET).get_path("jpg")) + if value.startswith(("http://", "https://", "data:")): + return value + path = Path(value).expanduser() + if not path.is_file(): + raise BenchmarkError(f"--image is neither a URL nor a readable file: {value}") + media_type = mimetypes.guess_type(path.name)[0] or "image/png" + return f"data:{media_type};base64,{base64.b64encode(path.read_bytes()).decode()}" + + +def build_workload(args: argparse.Namespace) -> Workload: + if args.workload is not None: + # The file names its own first frame, so resolving happens inside + # load_workload -- never eagerly, which would fetch the default asset a + # --workload run has no use for. + return load_workload( + args.workload, + image_resolver=lambda value: resolve_image_reference(value if args.image is None else args.image), + overrides={ + "width": args.width_override, + "height": args.height_override, + "fps": args.fps_override, + "seed": args.seed_override, + "prompt": args.prompt_override, + "flow_shift": args.flow_shift, + "negative_prompt": args.negative_prompt, + }, + ) + return default_workload( + image_reference=resolve_image_reference(args.image), + num_chunks=args.num_chunks, + camera_pattern=args.camera_pattern, + prompt=args.prompt_override if args.prompt_override is not None else DEFAULT_PROMPT, + width=args.width_override if args.width_override is not None else DEFAULT_WIDTH, + height=args.height_override if args.height_override is not None else DEFAULT_HEIGHT, + fps=args.fps_override if args.fps_override is not None else DEFAULT_FPS, + seed=args.seed_override if args.seed_override is not None else DEFAULT_SEED, + flow_shift=args.flow_shift if args.flow_shift is not None else DEFAULT_FLOW_SHIFT, + negative_prompt=args.negative_prompt, + ) + + +# --------------------------------------------------------------------------- # +# One streaming session +# --------------------------------------------------------------------------- # + + +@dataclass +class SessionResult: + request_id: str | None + records: list[ChunkRecord] + media: bytes + ttfc_s: float | None + video_start_s: float | None + wall_s: float + done_event: dict[str, Any] + + +async def _keepalive(websocket: Any, send_lock: asyncio.Lock, interval: float) -> None: + """Refresh the server's stall clock while a slow first chunk is generating. + + A compiled deploy config records CUDA graphs on the first block, which can + outlast the server's ~60 s stall timeout. ``session.ping`` is the documented + way to say the client is still there. + """ + while True: + await asyncio.sleep(interval) + async with send_lock: + await websocket.send(json.dumps({"type": "session.ping"})) + + +async def run_session( + url: str, + model: str, + workload: Workload, + *, + first_chunk_timeout: float, + chunk_timeout: float, + ping_interval: float, + collect_media: bool, + print_chunks: bool, +) -> SessionResult: + """Drive one rollout and timestamp every chunk as its bytes land.""" + try: + from websockets.asyncio.client import connect + except ImportError: # pragma: no cover - exercised only without websockets + try: + from websockets import connect # type: ignore[attr-defined] + except ImportError as exc: + raise BenchmarkError("This benchmark needs the websockets package: pip install websockets") from exc + + payload = workload.session_start_payload(model) + records: list[ChunkRecord] = [] + media = bytearray() + pending: dict[str, Any] | None = None + request_id: str | None = None + ttfc_s: float | None = None + video_start_s: float | None = None + last_arrival: float | None = None + done_event: dict[str, Any] = {} + send_lock = asyncio.Lock() + keepalive: asyncio.Task[None] | None = None + # Fire remaining prompt updates in chunk order as their boundary passes. + updates = sorted(workload.prompt_updates, key=lambda update: update.after_chunk) + update_index = 0 + + async with connect(url, max_size=None, ping_interval=None) as websocket: + async with send_lock: + await websocket.send(json.dumps(payload, ensure_ascii=False)) + started = time.perf_counter() + if ping_interval > 0: + keepalive = asyncio.create_task(_keepalive(websocket, send_lock, ping_interval)) + + try: + while True: + timeout = first_chunk_timeout if not records else chunk_timeout + try: + deadline = started + (last_arrival if last_arrival is not None else 0.0) + timeout + remaining = deadline - time.perf_counter() + if remaining <= 0: + raise asyncio.TimeoutError + message = await asyncio.wait_for(websocket.recv(), timeout=remaining) + except asyncio.TimeoutError as exc: + raise BenchmarkError( + f"No media chunk for {timeout:g} s after {len(records)} chunk(s); " + "the server may still be loading, compiling, or capturing graphs." + ) from exc + now = time.perf_counter() + + if isinstance(message, (bytes, bytearray)): + if pending is None: + raise BenchmarkError("Received a binary frame with no preceding video.chunk_metadata event.") + if collect_media: + media.extend(message) + if pending.get("kind") == "media": + index = int(pending["generation_chunk_index"]) + arrival = now - started + if ttfc_s is None: + ttfc_s = arrival + records.append( + ChunkRecord( + index=index, + arrival_s=arrival, + inter_arrival_s=arrival - (last_arrival if last_arrival is not None else 0.0), + byte_length=len(message), + num_frames=int(pending["num_frames"]), + ) + ) + last_arrival = arrival + if print_chunks: + print( + f" chunk {index:3d} t={arrival:8.3f}s " + f"interval={records[-1].inter_arrival_s * 1000:8.1f}ms " + f"frames={records[-1].num_frames:3d} bytes={len(message)}", + flush=True, + ) + while update_index < len(updates) and updates[update_index].after_chunk <= index: + async with send_lock: + await websocket.send(json.dumps(updates[update_index].to_payload(), ensure_ascii=False)) + update_index += 1 + pending = None + continue + + event = json.loads(message) + event_type = event.get("type") + if event_type == "video.chunk_metadata": + pending = event + elif event_type == "video.start": + request_id = event.get("request_id") + video_start_s = now - started + elif event_type == "error": + # Say what was asked for alongside what the server refused: + # some builds cap a rollout's num_frames, and the operator + # needs to know which --num-chunks would fit. + raise BenchmarkError( + f"Server rejected the session after {len(records)} chunk(s): " + f"{event.get('message', event)} " + f"(requested {workload.num_chunks} chunks = num_frames {workload.num_frames})" + ) + elif event_type == "session.done": + done_event = event + break + finally: + if keepalive is not None: + keepalive.cancel() + try: + await keepalive + except asyncio.CancelledError: + pass + + wall_s = time.perf_counter() - started + + if not records: + raise BenchmarkError("The session finished without delivering a single media chunk.") + if done_event.get("stopped"): + raise BenchmarkError(f"The session was stopped before it finished: {done_event}") + observed = [record.index for record in records] + if observed != list(range(len(records))): + raise BenchmarkError(f"Chunk indices are not contiguous from zero: {observed}") + if len(records) != workload.num_chunks: + raise BenchmarkError( + f"Expected {workload.num_chunks} chunks for num_frames={workload.num_frames}, received {len(records)}." + ) + + return SessionResult( + request_id=request_id, + records=records, + media=bytes(media), + ttfc_s=ttfc_s, + video_start_s=video_start_s, + wall_s=wall_s, + done_event=done_event, + ) + + +# --------------------------------------------------------------------------- # +# Reporting +# --------------------------------------------------------------------------- # + + +def _git_revision() -> str | None: + try: + return subprocess.check_output( + ["git", "-C", str(REPO_ROOT), "rev-parse", "HEAD"], + text=True, + stderr=subprocess.DEVNULL, + ).strip() + except (OSError, subprocess.CalledProcessError): + return None + + +def _print_row(label: str, value: str) -> None: + print(f"{label:<44}{value:>16}") + + +def print_report( + metrics: dict[str, Any], + *, + workload: Workload, + endpoint: str, + sessions: int, + target_fps: float, +) -> None: + header = " LingBot-World realtime streaming benchmark " + print("\n{s:{c}^{n}}".format(s=header, n=60, c="=")) + _print_row("Endpoint:", endpoint) + _print_row("Geometry:", f"{workload.width}x{workload.height}") + _print_row("Chunks per session:", str(workload.num_chunks)) + _print_row("Sessions:", str(sessions)) + frames = workload.num_frames + _print_row("Real-time basis:", f"{target_fps:g} fps") + _print_row("Video seconds per session:", f"{frames / target_fps:.2f} ({frames} frames)") + if abs(target_fps - workload.fps) > 1e-9: + _print_row("Mux label sent to server:", f"{workload.fps} fps") + + if sessions == 1: + print("{s:{c}^{n}}".format(s=" Cadence ", n=60, c="-")) + if metrics.get("ttfc_ms") is not None: + _print_row("TTFC (session.start -> first chunk):", f"{metrics['ttfc_ms']:.1f} ms") + interval_all = metrics.get("interval_all") or {} + interval_steady = metrics.get("interval_steady") or {} + for title, summary in (("All chunks", interval_all), ("Steady state", interval_steady)): + if not summary: + continue + print(f" {title} (n={int(summary['count'])}):") + _print_row(" mean / median:", f"{summary['mean_ms']:.1f} / {summary['median_ms']:.1f} ms") + _print_row(" std / max:", f"{summary['std_ms']:.1f} / {summary['max_ms']:.1f} ms") + tail = f"{summary['p90_ms']:.0f} / {summary['p95_ms']:.0f} / {summary['p99_ms']:.0f} ms" + _print_row(" p90 / p95 / p99:", tail) + if summary["count"] < 100: + print( + " NOTE: fewer than 100 intervals; tail percentiles are descriptive, not reliable tail estimates." + ) + + print("{s:{c}^{n}}".format(s=" Real time ", n=60, c="-")) + _print_row("VIDEO_RTF (wall s / video s, <1 real time):", f"{metrics['video_rtf']:.3f}") + _print_row("VIDEO_RTFX (video s / wall s):", f"{metrics['video_rtfx']:.3f}") + if metrics.get("steady_video_rtf") is not None: + _print_row("VIDEO_RTF steady state:", f"{metrics['steady_video_rtf']:.3f}") + _print_row("Frames per second:", f"{metrics['frames_per_second']:.2f}") + _print_row("Chunk deadline:", f"{metrics['chunk_deadline_ms']:.1f} ms") + if metrics.get("slo_attainment") is not None: + _print_row( + "Steady chunks within deadline:", + f"{metrics['slo_attainment'] * 100:.1f}% ({metrics['slo_violations']} late)", + ) + if not metrics.get("steady_state_reached", False): + print( + f" NOTE: no chunk reached the steady state; the attention window needs " + f"{metrics['warmup_chunks']} chunks. Raise --num-chunks." + ) + + playback = metrics.get("playback") or {} + print("{s:{c}^{n}}".format(s=" Playback continuity ", n=60, c="-")) + _print_row("Start buffer:", f"{playback.get('buffer_chunks')} chunk(s)") + _print_row("Underruns:", str(playback.get("underrun_count"))) + _print_row("Total stall:", f"{playback.get('total_stall_ms', 0.0):.1f} ms") + _print_row("Stall ratio:", f"{playback.get('stall_ratio', 0.0) * 100:.2f}%") + if playback.get("first_underrun_chunk") is not None: + _print_row("First underrun at chunk:", str(playback["first_underrun_chunk"])) + else: + print("{s:{c}^{n}}".format(s=" Across sessions ", n=60, c="-")) + if metrics.get("ttfc_ms_mean") is not None: + _print_row("TTFC mean / max:", f"{metrics['ttfc_ms_mean']:.1f} / {metrics['ttfc_ms_max']:.1f} ms") + if metrics.get("steady_interval_ms_mean") is not None: + _print_row("Steady interval mean:", f"{metrics['steady_interval_ms_mean']:.1f} ms") + _print_row("Steady interval spread:", f"{metrics['steady_interval_ms_spread']:.1f} ms") + _print_row("VIDEO_RTF mean (lower is better):", f"{metrics['video_rtf_mean']:.3f}") + _print_row( + "VIDEO_RTF best / worst:", + f"{metrics['video_rtf_best']:.3f} / {metrics['video_rtf_worst']:.3f}", + ) + _print_row("Chunk deadline:", f"{metrics['chunk_deadline_ms']:.1f} ms") + _print_row("Underruns (all sessions):", str(metrics["underrun_count"])) + print("=" * 60) + + +def save_video(media: bytes, path: Path, *, fps: int) -> None: + from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes + + playback_bytes = finalize_streaming_video_bytes(media, input_format="m4s", fps=fps) + path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = path.with_suffix(path.suffix + ".tmp") + tmp_path.write_bytes(playback_bytes) + os.replace(tmp_path, path) + print(f"Saved {len(media)} streamed bytes -> {len(playback_bytes)} playback bytes at {path}") + + +# --------------------------------------------------------------------------- # +# CLI +# --------------------------------------------------------------------------- # + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark LingBot-World 2.0 realtime streaming over WS /v1/realtime/video.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + endpoint = parser.add_argument_group("endpoint") + endpoint.add_argument("--host", default="127.0.0.1") + endpoint.add_argument("--port", type=int, default=8000) + endpoint.add_argument("--base-url", default=None, help="Full ws:// URL; overrides --host/--port.") + endpoint.add_argument("--model", default=DEFAULT_MODEL) + + workload = parser.add_argument_group("workload") + workload.add_argument("--workload", default=None, help="JSON rollout spec; overrides the built-in workload.") + workload.add_argument("--image", default=None, help="First frame: path or URL. Defaults to the shared vLLM asset.") + workload.add_argument("--prompt", dest="prompt_override", default=None, help=f"Default: {DEFAULT_PROMPT!r}") + workload.add_argument("--negative-prompt", default=None) + workload.add_argument("--num-chunks", type=int, default=16, help="AR blocks per session; ignored with --workload.") + workload.add_argument( + "--camera-pattern", + choices=CAMERA_PATTERNS, + default="forward", + help="Camera pattern for the built-in workload only; ignored with --workload.", + ) + workload.add_argument("--width", dest="width_override", type=int, default=None, help=f"Default: {DEFAULT_WIDTH}") + workload.add_argument("--height", dest="height_override", type=int, default=None, help=f"Default: {DEFAULT_HEIGHT}") + workload.add_argument("--fps", dest="fps_override", type=int, default=None, help=f"Default: {DEFAULT_FPS}") + workload.add_argument("--seed", dest="seed_override", type=int, default=None, help=f"Default: {DEFAULT_SEED}") + workload.add_argument("--flow-shift", type=float, default=None, help=f"Default: {DEFAULT_FLOW_SHIFT}") + + run = parser.add_argument_group("run") + run.add_argument("--sessions", type=int, default=1, help="Sequential measured rollouts.") + run.add_argument("--warmup-sessions", type=int, default=0, help="Unmeasured rollouts first (pays compile cost).") + run.add_argument( + "--warmup-chunks", + type=int, + default=DEFAULT_WARMUP_CHUNKS, + help="Chunks excluded from steady state; the attention window saturates after 6.", + ) + run.add_argument( + "--chunk-slo-ms", + type=float, + default=None, + help="Deadline a steady chunk must meet. Default: one chunk of video at --target-fps.", + ) + run.add_argument( + "--target-fps", + type=float, + default=None, + help=( + "Playback rate used for real-time factor, the chunk deadline, and the playback " + "simulation. The checkpoint declares no frame rate, so this is a consumption " + "choice, not a model property. Default: the --fps sent to the server." + ), + ) + run.add_argument("--playback-buffer-chunks", type=int, default=1, help="Chunks buffered before playback starts.") + run.add_argument("--first-chunk-timeout", type=float, default=900.0) + run.add_argument("--chunk-timeout", type=float, default=300.0) + run.add_argument("--ping-interval", type=float, default=20.0, help="session.ping cadence; 0 disables.") + + output = parser.add_argument_group("output") + output.add_argument("--output-json", default=None) + output.add_argument("--save-video", default=None, help="Write the last measured session's video here.") + output.add_argument("--print-chunks", action="store_true") + + args = parser.parse_args(argv) + if args.sessions < 1: + parser.error("--sessions must be at least 1.") + if args.warmup_sessions < 0: + parser.error("--warmup-sessions must be non-negative.") + if args.warmup_chunks < 0: + parser.error("--warmup-chunks must be non-negative.") + if args.playback_buffer_chunks < 1: + parser.error("--playback-buffer-chunks must be at least 1.") + if args.target_fps is not None and args.target_fps <= 0: + parser.error("--target-fps must be positive.") + if args.workload is None and args.num_chunks < 1: + parser.error("--num-chunks must be at least 1.") + return args + + +def endpoint_url(args: argparse.Namespace) -> str: + if args.base_url is not None: + base = args.base_url.rstrip("/") + for prefix, replacement in (("http://", "ws://"), ("https://", "wss://")): + if base.startswith(prefix): + base = replacement + base[len(prefix) :] + break + return base if base.endswith(_WS_PATH) else base + _WS_PATH + return f"ws://{args.host}:{args.port}{_WS_PATH}" + + +async def run_benchmark(args: argparse.Namespace) -> dict[str, Any]: + workload = build_workload(args) + url = endpoint_url(args) + # The mux label and the real-time basis are different questions; only the + # latter decides whether a run passes. + target_fps = args.target_fps if args.target_fps is not None else float(workload.fps) + if max(1, args.warmup_chunks) >= workload.num_chunks - 1: + print( + f"WARNING: --warmup-chunks={args.warmup_chunks} leaves no nonterminal steady-state chunk in a " + f"{workload.num_chunks}-chunk rollout; steady-state metrics will be empty.", + file=sys.stderr, + ) + + print(f"Connecting to {url}") + print( + f"Rollout: {workload.num_chunks} chunks, num_frames={workload.num_frames}, " + f"{workload.video_seconds:.2f}s of video" + ) + print( + f"Real-time basis: {target_fps:g} fps -> a steady 12-frame chunk is " + f"{steady_chunk_deadline_ms(target_fps):.1f} ms of video" + + ("" if args.target_fps is None else f" (mux label stays {workload.fps} fps)") + ) + + for index in range(args.warmup_sessions): + print(f"Warmup session {index + 1}/{args.warmup_sessions} (not measured)...") + await run_session( + url, + args.model, + workload, + first_chunk_timeout=args.first_chunk_timeout, + chunk_timeout=args.chunk_timeout, + ping_interval=args.ping_interval, + collect_media=False, + print_chunks=False, + ) + + session_metrics: list[dict[str, Any]] = [] + session_payloads: list[dict[str, Any]] = [] + last_result: SessionResult | None = None + for index in range(args.sessions): + if args.sessions > 1: + print(f"Session {index + 1}/{args.sessions}...") + result = await run_session( + url, + args.model, + workload, + first_chunk_timeout=args.first_chunk_timeout, + chunk_timeout=args.chunk_timeout, + ping_interval=args.ping_interval, + collect_media=args.save_video is not None, + print_chunks=args.print_chunks, + ) + last_result = result + metrics = compute_metrics( + result.records, + fps=target_fps, + warmup_chunks=args.warmup_chunks, + slo_ms=args.chunk_slo_ms, + playback_buffer_chunks=args.playback_buffer_chunks, + ttfc_s=result.ttfc_s, + session_wall_s=result.wall_s, + ) + session_metrics.append(metrics) + session_payloads.append( + { + "session_index": index, + "request_id": result.request_id, + "video_start_ms": (result.video_start_s * 1000.0) if result.video_start_s is not None else None, + "chunks": [record.to_dict() for record in result.records], + "metrics": metrics, + } + ) + + aggregate = aggregate_metrics(session_metrics, fps=target_fps) + print_report(aggregate, workload=workload, endpoint=url, sessions=args.sessions, target_fps=target_fps) + + if args.save_video is not None and last_result is not None: + save_video(last_result.media, Path(args.save_video), fps=workload.fps) + + document = { + "benchmark": "lingbot_world_realtime", + "timestamp": datetime.now(timezone.utc).isoformat(), + "endpoint": url, + "model": args.model, + "git_revision": _git_revision(), + "hostname": platform.node(), + "workload": workload.describe(), + "config": { + "sessions": args.sessions, + "warmup_sessions": args.warmup_sessions, + "warmup_chunks": args.warmup_chunks, + "camera_pattern": args.camera_pattern, + "chunk_slo_ms": args.chunk_slo_ms, + "target_fps": target_fps, + "mux_fps": workload.fps, + "playback_buffer_chunks": args.playback_buffer_chunks, + "workload_file": args.workload, + }, + "sessions": session_payloads, + "metrics": aggregate, + } + if args.output_json is not None: + path = Path(args.output_json) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(document, indent=2, sort_keys=True) + "\n") + print(f"Wrote {path}") + return document + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + try: + asyncio.run(run_benchmark(args)) + except BenchmarkError as error: + print(f"ERROR: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/lingbot_world/configs/single_gpu_eager.yaml b/benchmarks/lingbot_world/configs/single_gpu_eager.yaml new file mode 100644 index 000000000..58f36e4f0 --- /dev/null +++ b/benchmarks/lingbot_world/configs/single_gpu_eager.yaml @@ -0,0 +1,45 @@ +# Benchmark deploy config: LingBot-World 2.0 on one GPU, eager. +# +# This is the portable smoke configuration -- it matches the shipped +# vllm_omni/deploy/lingbot_world_v2_stepwise.yaml so a number measured here is +# comparable to the default deployment. It is NOT a realtime configuration: one +# GPU runs all five forwards of every AR block serially and nothing is compiled, +# so expect a chunk interval several times the real-time deadline. +# +# Use configs/usp4_compiled.yaml for the realtime topology. +# +# vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \ +# --omni --deploy-config benchmarks/lingbot_world/configs/single_gpu_eager.yaml \ +# --port 8000 + +pipeline: lingbot_world +async_chunk: false +distributed_executor_backend: mp +dtype: bfloat16 + +stages: + - stage_id: 0 + devices: "0" + # AR-Diffusion binds one session per runner invocation; the runner rejects + # anything larger. + max_num_seqs: 1 + enforce_eager: true + model_class_name: LingBotWorldCausalDMDPipeline + engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine + step_execution: true + streaming_output: true + default_sampling_params: + # The causal-fast checkpoint pins a four-step DMD schedule; this sizes the + # step budget before the pipeline is reached. + num_inference_steps: 4 + output_type: pil + extra_args: + flow_shift: 5.0 + model_config: + # The AR cache geometry is fixed at load time, so a benchmark request must + # ask for exactly this resolution. + ar_diffusion_height: 480 + ar_diffusion_width: 832 + ar_diffusion_kv_config: + gpu_memory_fraction: 0.6 + warmup_cudagraph: true diff --git a/benchmarks/lingbot_world/configs/usp4_compiled.yaml b/benchmarks/lingbot_world/configs/usp4_compiled.yaml new file mode 100644 index 000000000..ed71bc13e --- /dev/null +++ b/benchmarks/lingbot_world/configs/usp4_compiled.yaml @@ -0,0 +1,60 @@ +# Benchmark deploy config: LingBot-World 2.0 on four GPUs, Ulysses SP-4, compiled. +# +# This is the realtime topology. Two things differ from the single-GPU config, +# and both matter to the number the benchmark reports: +# +# parallel_config Pure Ulysses sequence parallelism over four ranks. The +# pipeline validates this: LingBot requires +# sequence_parallel_size == ulysses_degree, ring_degree and +# allgather_degree of 1, ulysses_mode "strict", and +# ulysses_a2a_permute disabled. Pipeline, CFG, and VAE +# parallel sizes above 1 are rejected. This config uses +# TP=1; the validator does not reject larger TP sizes. +# enforce_eager false puts LingBotDMDBlockRunner on the compiled path +# (regional torch.compile with cudagraph step marking). The +# first rollout therefore pays compile and capture cost -- +# use --warmup-sessions 1 before measured repetitions. +# Chunk warmup only excludes the attention-window ramp. +# +# Serve it under a GPU reservation, since it needs four cards: +# +# gpu run --gpus 4 -- vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \ +# --omni --deploy-config benchmarks/lingbot_world/configs/usp4_compiled.yaml \ +# --port 8000 +# +# `devices` is the rank->card mapping inside the reservation, so it stays +# "0,1,2,3" no matter which physical cards CUDA_VISIBLE_DEVICES selected. + +pipeline: lingbot_world +async_chunk: false +distributed_executor_backend: mp +dtype: bfloat16 + +stages: + - stage_id: 0 + devices: "0,1,2,3" + max_num_seqs: 1 + enforce_eager: false + model_class_name: LingBotWorldCausalDMDPipeline + engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine + step_execution: true + streaming_output: true + parallel_config: + tensor_parallel_size: 1 + sequence_parallel_size: 4 + ulysses_degree: 4 + ring_degree: 1 + default_sampling_params: + num_inference_steps: 4 + output_type: pil + extra_args: + flow_shift: 5.0 + model_config: + ar_diffusion_height: 480 + ar_diffusion_width: 832 + ar_diffusion_kv_config: + # Each rank holds its own shard of the paged KV cache, so the per-card + # fraction can stay at the single-GPU value; lower it if the four-rank + # start OOMs on smaller cards. + gpu_memory_fraction: 0.6 + warmup_cudagraph: true diff --git a/benchmarks/lingbot_world/workload.py b/benchmarks/lingbot_world/workload.py new file mode 100644 index 000000000..4b59c8899 --- /dev/null +++ b/benchmarks/lingbot_world/workload.py @@ -0,0 +1,631 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Workload construction and metric math for the LingBot-World realtime benchmark. + +Nothing here imports torch, vllm, or websockets, so the arithmetic that decides +what a benchmark number *means* is unit-testable without a GPU or a server. + +Two model properties drive every formula in this module, both read from the +checkpoint's ``transformer/config.json`` rather than assumed: + +``num_frames_per_block = 3`` + One AR block commits three latent frames, so one streamed chunk is three + latent frames and a request's ``camera_action_script`` carries exactly one + three-entry action list per chunk. + +``sliding_window_num_frames = 18`` with ``sink_size = 9`` + Self-attention sees nine sink frames plus a nine-frame recent window, so the + key count a block attends over keeps growing until eighteen latent frames of + history exist -- six chunks. Chunk cost is therefore *not* stationary before + chunk six, and a mean taken over a short rollout reports a number no steady + session will ever produce. :data:`DEFAULT_WARMUP_CHUNKS` exists for that + reason and is not a round-number guess. + +The causal VAE expands the first latent frame to one pixel frame and every later +latent frame to four, so chunk zero carries nine pixel frames and every chunk +after it carries twelve. Video-duration math has to respect that asymmetry or +real-time factor comes out wrong on short rollouts. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +# transformer.config.num_frames_per_block: latent frames committed per AR block. +FRAMES_PER_BLOCK = 3 +# vae_scale_factor_temporal: latent frames -> pixel frames, after the first. +TEMPORAL_COMPRESSION = 4 +# (sink_size 9 + recent window 9) / FRAMES_PER_BLOCK: chunks until the attention +# window saturates and per-chunk cost becomes stationary. +DEFAULT_WARMUP_CHUNKS = 6 + +DEFAULT_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" +DEFAULT_PROMPT = "The camera moves slowly forward through the scene." +# The AR KV cache geometry is fixed at load time, so a request must ask for +# exactly the resolution the deploy config declares. +DEFAULT_WIDTH = 832 +DEFAULT_HEIGHT = 480 +DEFAULT_FPS = 16 +DEFAULT_FLOW_SHIFT = 5.0 +DEFAULT_SEED = 42 + +_HOLD = "hold" +_ACTION_KEYS = ("w", "a", "s", "d") + + +def num_frames_for_chunks(num_chunks: int) -> int: + """Pixel-frame count whose rollout is exactly ``num_chunks`` chunks. + + Inverse of :func:`chunks_for_num_frames`, which is the server's own + ``((num_frames - 1) // 4 + 1) // 3``. + """ + if num_chunks < 1: + raise ValueError("num_chunks must be at least 1.") + return (num_chunks * FRAMES_PER_BLOCK - 1) * TEMPORAL_COMPRESSION + 1 + + +def chunks_for_num_frames(num_frames: int) -> int: + """Chunks the server generates for ``num_frames``, mirroring its own formula.""" + if num_frames < 1: + raise ValueError("num_frames must be at least 1.") + return ((num_frames - 1) // TEMPORAL_COMPRESSION + 1) // FRAMES_PER_BLOCK + + +def pixel_frames_in_chunk(chunk_index: int) -> int: + """Pixel frames carried by one chunk. + + The causal decoder expands the session's opening latent frame once and every + later latent frame four times, so chunk zero is short by three frames. + """ + if chunk_index < 0: + raise ValueError("chunk_index must be non-negative.") + if chunk_index == 0: + return (FRAMES_PER_BLOCK - 1) * TEMPORAL_COMPRESSION + 1 + return FRAMES_PER_BLOCK * TEMPORAL_COMPRESSION + + +def camera_actions_for_chunk(pattern: str, chunk_index: int) -> list[list[str]]: + """One three-entry action list for ``chunk_index``. + + Camera actions select camera embeddings; they change what the world does, + not how much arithmetic it costs. The pattern therefore exists to keep a + benchmark rollout representative rather than to sweep a cost dimension. + """ + if chunk_index < 0: + raise ValueError("chunk_index must be non-negative.") + if pattern == "forward": + return [["w"] for _ in range(FRAMES_PER_BLOCK)] + if pattern == _HOLD: + return [[] for _ in range(FRAMES_PER_BLOCK)] + if pattern == "orbit": + key = _ACTION_KEYS[chunk_index % len(_ACTION_KEYS)] + return [[key] for _ in range(FRAMES_PER_BLOCK)] + raise ValueError(f"Unknown camera pattern {pattern!r}; expected one of {sorted(CAMERA_PATTERNS)}.") + + +CAMERA_PATTERNS = ("forward", "orbit", _HOLD) + + +def build_camera_script(num_chunks: int, pattern: str) -> list[list[list[str]]]: + """Full ``camera_action_script``: one action list per generated chunk.""" + if num_chunks < 1: + raise ValueError("num_chunks must be at least 1.") + return [camera_actions_for_chunk(pattern, index) for index in range(num_chunks)] + + +def _validate_camera_script(script: Sequence[Any]) -> list[list[list[str]]]: + validated: list[list[list[str]]] = [] + for chunk_index, actions in enumerate(script): + if not isinstance(actions, Sequence) or isinstance(actions, (str, bytes)): + raise ValueError(f"camera_action_script[{chunk_index}] must be a list of three action lists.") + if len(actions) != FRAMES_PER_BLOCK: + raise ValueError( + f"camera_action_script[{chunk_index}] must hold exactly {FRAMES_PER_BLOCK} action lists " + f"(got {len(actions)}); one AR block is {FRAMES_PER_BLOCK} latent frames." + ) + frames: list[list[str]] = [] + for frame_index, frame in enumerate(actions): + if not isinstance(frame, Sequence) or isinstance(frame, (str, bytes)): + raise ValueError(f"camera_action_script[{chunk_index}][{frame_index}] must be a list of strings.") + if any(not isinstance(action, str) for action in frame): + raise ValueError(f"camera_action_script[{chunk_index}][{frame_index}] must contain only strings.") + frames.append([str(action) for action in frame]) + validated.append(frames) + if not validated: + raise ValueError("camera_action_script must not be empty.") + return validated + + +@dataclass(frozen=True) +class PromptUpdate: + """A mid-rollout ``session.interaction`` scheduled on a chunk boundary.""" + + after_chunk: int + prompt: str + transition_chunks: int = 3 + + def __post_init__(self) -> None: + if self.after_chunk < 0: + raise ValueError("PromptUpdate.after_chunk must be non-negative.") + if not self.prompt.strip(): + raise ValueError("PromptUpdate.prompt must contain non-whitespace text.") + if self.transition_chunks < 0: + raise ValueError("PromptUpdate.transition_chunks must be non-negative.") + + def to_payload(self) -> dict[str, Any]: + return { + "type": "session.interaction", + "interaction": { + "event_id": f"chunk-{self.after_chunk}", + "event": {"prompt": self.prompt}, + "transition_chunks": self.transition_chunks, + }, + } + + +@dataclass(frozen=True) +class Workload: + """One rollout: what the benchmark asks the server to generate.""" + + prompt: str + image_reference: str + camera_script: list[list[list[str]]] + width: int = DEFAULT_WIDTH + height: int = DEFAULT_HEIGHT + fps: int = DEFAULT_FPS + seed: int = DEFAULT_SEED + flow_shift: float = DEFAULT_FLOW_SHIFT + negative_prompt: str | None = None + prompt_updates: tuple[PromptUpdate, ...] = () + extra_params: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.prompt.strip(): + raise ValueError("Workload.prompt must contain non-whitespace text.") + if not self.image_reference: + raise ValueError("Workload.image_reference is required; LingBot-World is image-conditioned.") + if self.width <= 0 or self.height <= 0 or self.width % 16 or self.height % 16: + raise ValueError("Workload width/height must be positive multiples of 16.") + if self.fps <= 0: + raise ValueError("Workload.fps must be positive.") + object.__setattr__(self, "camera_script", _validate_camera_script(self.camera_script)) + seen_update_chunks: set[int] = set() + for update in self.prompt_updates: + if update.after_chunk in seen_update_chunks: + raise ValueError(f"Duplicate prompt update after chunk {update.after_chunk}; event IDs must be unique.") + seen_update_chunks.add(update.after_chunk) + if update.after_chunk >= self.num_chunks: + raise ValueError( + f"Prompt update after chunk {update.after_chunk} never fires in a {self.num_chunks}-chunk rollout." + ) + + @property + def num_chunks(self) -> int: + return len(self.camera_script) + + @property + def num_frames(self) -> int: + return num_frames_for_chunks(self.num_chunks) + + @property + def video_seconds(self) -> float: + return self.num_frames / self.fps + + def session_start_payload(self, model: str) -> dict[str, Any]: + """The ``session.start`` event for ``WS /v1/realtime/video``.""" + extra_params: dict[str, Any] = { + "flow_shift": self.flow_shift, + "camera_action_script": self.camera_script, + } + extra_params.update(self.extra_params) + payload: dict[str, Any] = { + "type": "session.start", + "model": model, + "prompt": self.prompt, + "image_reference": {"image_url": self.image_reference}, + "width": self.width, + "height": self.height, + "num_frames": self.num_frames, + "fps": self.fps, + "seed": self.seed, + "extra_params": extra_params, + } + if self.negative_prompt is not None: + payload["negative_prompt"] = self.negative_prompt + return payload + + def describe(self) -> dict[str, Any]: + return { + "num_chunks": self.num_chunks, + "num_frames": self.num_frames, + "width": self.width, + "height": self.height, + "fps": self.fps, + "seed": self.seed, + "flow_shift": self.flow_shift, + "video_seconds": self.video_seconds, + "prompt_updates": [ + { + "after_chunk": update.after_chunk, + "transition_chunks": update.transition_chunks, + } + for update in self.prompt_updates + ], + } + + +def default_workload( + *, + image_reference: str, + num_chunks: int, + camera_pattern: str = "forward", + prompt: str = DEFAULT_PROMPT, + width: int = DEFAULT_WIDTH, + height: int = DEFAULT_HEIGHT, + fps: int = DEFAULT_FPS, + seed: int = DEFAULT_SEED, + flow_shift: float = DEFAULT_FLOW_SHIFT, + negative_prompt: str | None = None, +) -> Workload: + """Built-in rollout: a fixed prompt and a generated camera script.""" + return Workload( + prompt=prompt, + image_reference=image_reference, + camera_script=build_camera_script(num_chunks, camera_pattern), + width=width, + height=height, + fps=fps, + seed=seed, + flow_shift=flow_shift, + negative_prompt=negative_prompt, + ) + + +def load_workload( + path: str | Path, + *, + image_resolver: Any = None, + overrides: Mapping[str, Any] | None = None, +) -> Workload: + """Read a rollout spec from JSON. + + ``image_resolver`` turns the spec's ``image`` field into an ``http(s)`` or + ``data:`` URL; the caller owns that because inlining a local file is I/O. + Keys in ``overrides`` that are not ``None`` win over the file, so CLI flags + stay meaningful next to ``--workload``. + """ + document = json.loads(Path(path).read_text()) + if not isinstance(document, dict): + raise ValueError("A workload file must contain a JSON object.") + + merged = dict(document) + for key, value in (overrides or {}).items(): + if value is not None: + merged[key] = value + + camera_script = merged.get("camera_action_script") + if camera_script is None: + num_chunks = merged.get("num_chunks") + if num_chunks is None: + raise ValueError("A workload file needs either camera_action_script or num_chunks.") + camera_script = build_camera_script(int(num_chunks), str(merged.get("camera_pattern", "forward"))) + + image = merged.get("image") or merged.get("image_reference") + if image is None: + raise ValueError("A workload file needs an image or image_reference field.") + image_reference = image_resolver(str(image)) if image_resolver is not None else str(image) + + updates_field = merged.get("prompt_updates") or () + prompt_updates = tuple( + PromptUpdate( + after_chunk=int(entry["after_chunk"]), + prompt=str(entry["prompt"]), + transition_chunks=int(entry.get("transition_chunks", 3)), + ) + for entry in updates_field + ) + + extra_params = merged.get("extra_params") or {} + if not isinstance(extra_params, Mapping): + raise ValueError("Workload extra_params must be a JSON object.") + + return Workload( + prompt=str(merged.get("prompt", DEFAULT_PROMPT)), + image_reference=image_reference, + camera_script=camera_script, + width=int(merged.get("width", DEFAULT_WIDTH)), + height=int(merged.get("height", DEFAULT_HEIGHT)), + fps=int(merged.get("fps", DEFAULT_FPS)), + seed=int(merged.get("seed", DEFAULT_SEED)), + flow_shift=float(merged.get("flow_shift", DEFAULT_FLOW_SHIFT)), + negative_prompt=merged.get("negative_prompt"), + prompt_updates=prompt_updates, + extra_params=dict(extra_params), + ) + + +# --------------------------------------------------------------------------- # +# Metrics +# --------------------------------------------------------------------------- # + + +@dataclass +class ChunkRecord: + """One media chunk as the client observed it.""" + + index: int + arrival_s: float + """Seconds from ``session.start`` being sent to this chunk's bytes arriving.""" + inter_arrival_s: float + """Seconds since the previous chunk; for chunk zero this equals ``arrival_s``.""" + byte_length: int + num_frames: int + + def to_dict(self) -> dict[str, Any]: + return { + "index": self.index, + "arrival_s": self.arrival_s, + "inter_arrival_ms": self.inter_arrival_s * 1000.0, + "byte_length": self.byte_length, + "num_frames": self.num_frames, + } + + +def percentile(values: Sequence[float], fraction: float) -> float: + """Linearly interpolated percentile, so a 10-sample run still reports a p95. + + Implemented here rather than pulled from numpy to keep this module importable + in a bare environment; the benchmark's own numbers must not depend on which + optional package happens to be installed. + """ + if not values: + raise ValueError("percentile() needs at least one value.") + if not 0.0 <= fraction <= 1.0: + raise ValueError("percentile fraction must be within [0, 1].") + ordered = sorted(values) + if len(ordered) == 1: + return ordered[0] + position = fraction * (len(ordered) - 1) + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def _stddev(values: Sequence[float]) -> float: + if len(values) < 2: + return 0.0 + mean = sum(values) / len(values) + return (sum((value - mean) ** 2 for value in values) / (len(values) - 1)) ** 0.5 + + +def _latency_summary(values: Sequence[float]) -> dict[str, float]: + """Millisecond summary of a set of inter-arrival intervals.""" + if not values: + return {} + milliseconds = [value * 1000.0 for value in values] + return { + "count": float(len(milliseconds)), + "mean_ms": sum(milliseconds) / len(milliseconds), + "median_ms": percentile(milliseconds, 0.5), + "std_ms": _stddev(milliseconds), + "min_ms": min(milliseconds), + "p90_ms": percentile(milliseconds, 0.90), + "p95_ms": percentile(milliseconds, 0.95), + "p99_ms": percentile(milliseconds, 0.99), + "max_ms": max(milliseconds), + } + + +@dataclass(frozen=True) +class PlaybackResult: + """What a viewer would have seen, given a start buffer and a playback rate.""" + + underrun_count: int + total_stall_s: float + playback_span_s: float + first_underrun_chunk: int | None + + def to_dict(self) -> dict[str, Any]: + return { + "underrun_count": self.underrun_count, + "total_stall_ms": self.total_stall_s * 1000.0, + "playback_span_s": self.playback_span_s, + "stall_ratio": (self.total_stall_s / self.playback_span_s) if self.playback_span_s > 0 else 0.0, + "first_underrun_chunk": self.first_underrun_chunk, + } + + +def simulate_playback( + records: Sequence[ChunkRecord], + *, + fps: float, + buffer_chunks: int = 1, +) -> PlaybackResult: + """Replay arrivals against a wall clock and count stalls. + + Playback starts when ``buffer_chunks`` chunks have arrived and then consumes + video at ``fps``. A chunk that has not arrived by the moment the player needs + it is an underrun, and the player resumes only when the bytes land -- which + is what a viewer experiences, and what a single p99 number cannot tell you: + the same p99 is harmless behind a three-chunk buffer and fatal behind none. + """ + if fps <= 0: + raise ValueError("fps must be positive.") + if buffer_chunks < 1: + raise ValueError("buffer_chunks must be at least 1.") + if not records: + return PlaybackResult(0, 0.0, 0.0, None) + if buffer_chunks > len(records): + raise ValueError(f"buffer_chunks={buffer_chunks} exceeds the {len(records)} chunks received.") + + ordered = sorted(records, key=lambda record: record.index) + start = ordered[buffer_chunks - 1].arrival_s + cursor = start + underruns = 0 + total_stall = 0.0 + first_underrun: int | None = None + + for record in ordered: + if record.arrival_s > cursor: + stall = record.arrival_s - cursor + underruns += 1 + total_stall += stall + if first_underrun is None: + first_underrun = record.index + cursor = record.arrival_s + cursor += record.num_frames / fps + + return PlaybackResult( + underrun_count=underruns, + total_stall_s=total_stall, + playback_span_s=cursor - start, + first_underrun_chunk=first_underrun, + ) + + +def steady_chunk_deadline_ms(fps: float) -> float: + """Wall-clock budget a steady chunk must meet to hold real time. + + ``fps`` here is the rate the frames are *intended to be played at*, which the + checkpoint does not declare: nothing in ``model_index.json``, the scheduler + config, or the transformer config carries a frame rate, and the pipeline + never reads one. It is a property of how the output is consumed, so both + this deadline and real-time factor scale linearly with the choice and the + caller has to state it rather than inherit it from a muxing default. + """ + if fps <= 0: + raise ValueError("fps must be positive.") + return pixel_frames_in_chunk(1) / fps * 1000.0 + + +def compute_metrics( + records: Sequence[ChunkRecord], + *, + fps: float, + warmup_chunks: int = DEFAULT_WARMUP_CHUNKS, + slo_ms: float | None = None, + playback_buffer_chunks: int = 1, + ttfc_s: float | None = None, + session_wall_s: float | None = None, +) -> dict[str, Any]: + """Summarize one session. + + ``warmup_chunks`` are excluded from the steady-state block because the + attention window has not saturated before then (see the module docstring). + The terminal chunk also stays out of steady metrics because it skips + next-chunk preparation. All chunks stay in the playback simulation, which is + what a session really costs end to end. + """ + if not records: + raise ValueError("compute_metrics() needs at least one chunk record.") + if warmup_chunks < 0: + raise ValueError("warmup_chunks must be non-negative.") + + ordered = sorted(records, key=lambda record: record.index) + deadline_ms = slo_ms if slo_ms is not None else steady_chunk_deadline_ms(fps) + + all_intervals = [record.inter_arrival_s for record in ordered[1:]] + # The terminal chunk skips next-chunk preparation and is not representative + # of sustained generation. Keep it in end-to-end and playback metrics. + steady = [record for record in ordered[:-1] if record.index >= warmup_chunks] + steady_intervals = [record.inter_arrival_s for record in steady if record.index > 0] + + total_frames = sum(record.num_frames for record in ordered) + video_seconds = total_frames / fps + wall_seconds = session_wall_s if session_wall_s is not None else ordered[-1].arrival_s + + metrics: dict[str, Any] = { + "chunks_received": len(ordered), + "warmup_chunks": warmup_chunks, + "steady_chunks": len(steady_intervals), + "steady_state_reached": bool(steady_intervals), + "total_frames": total_frames, + "total_bytes": sum(record.byte_length for record in ordered), + "video_seconds": video_seconds, + "wall_seconds": wall_seconds, + # RTF is processing time over content duration, as vLLM-Omni computes it + # (`patch.py`: generation_s / video_duration; `compute_audio_rtf`: latency + # / audio_duration, "SLO red line < 1"). LOWER IS BETTER and below 1.0 is + # real time. RTFX is the reciprocal the repo reports under its own name, + # kept here because "fraction of real time achieved" reads better on a + # dashboard -- but the two must never be confused for one another. + "video_rtf": (wall_seconds / video_seconds) if video_seconds > 0 else 0.0, + "video_rtfx": (video_seconds / wall_seconds) if wall_seconds > 0 else 0.0, + "frames_per_second": (total_frames / wall_seconds) if wall_seconds > 0 else 0.0, + "chunk_deadline_ms": deadline_ms, + "interval_all": _latency_summary(all_intervals), + "interval_steady": _latency_summary(steady_intervals), + } + if ttfc_s is not None: + metrics["ttfc_ms"] = ttfc_s * 1000.0 + + if steady_intervals: + steady_frames = sum(record.num_frames for record in steady if record.index > 0) + steady_wall = sum(steady_intervals) + steady_video = steady_frames / fps + metrics["steady_video_rtf"] = (steady_wall / steady_video) if steady_video > 0 else 0.0 + metrics["steady_video_rtfx"] = (steady_video / steady_wall) if steady_wall > 0 else 0.0 + within = sum(1 for value in steady_intervals if value * 1000.0 <= deadline_ms) + metrics["slo_attainment"] = within / len(steady_intervals) + metrics["slo_violations"] = len(steady_intervals) - within + + metrics["playback"] = simulate_playback( + ordered, + fps=fps, + buffer_chunks=playback_buffer_chunks, + ).to_dict() + metrics["playback"]["buffer_chunks"] = playback_buffer_chunks + return metrics + + +def aggregate_metrics(sessions: Iterable[Mapping[str, Any]], *, fps: float) -> dict[str, Any]: + """Combine per-session metrics across sequential repeats. + + Steady intervals are pooled rather than averaging per-session means, so a + session that happened to produce fewer steady chunks does not carry the same + weight as a full one. + """ + entries = list(sessions) + if not entries: + raise ValueError("aggregate_metrics() needs at least one session.") + if len(entries) == 1: + single = dict(entries[0]) + single["sessions"] = 1 + return single + + deadline_ms = entries[0]["chunk_deadline_ms"] + if any(entry["chunk_deadline_ms"] != deadline_ms for entry in entries): + raise ValueError("Cannot aggregate sessions with different chunk deadlines.") + ttfcs = [float(entry["ttfc_ms"]) for entry in entries if entry.get("ttfc_ms") is not None] + # Lower is better, so "max" is the worst session, not the best. + rtfs = [float(entry["video_rtf"]) for entry in entries] + steady_means = [ + float(entry["interval_steady"]["mean_ms"]) for entry in entries if entry.get("interval_steady", {}).get("count") + ] + steady_count = sum(entry.get("interval_steady", {}).get("count", 0) for entry in entries) + steady_total_ms = sum( + entry["interval_steady"]["mean_ms"] * entry["interval_steady"]["count"] + for entry in entries + if entry.get("interval_steady", {}).get("count") + ) + return { + "sessions": len(entries), + "chunks_received": sum(int(entry["chunks_received"]) for entry in entries), + "total_frames": sum(int(entry["total_frames"]) for entry in entries), + "video_seconds": sum(float(entry["video_seconds"]) for entry in entries), + "wall_seconds": sum(float(entry["wall_seconds"]) for entry in entries), + "video_rtf_mean": sum(rtfs) / len(rtfs), + "video_rtf_best": min(rtfs), + "video_rtf_worst": max(rtfs), + "ttfc_ms_mean": (sum(ttfcs) / len(ttfcs)) if ttfcs else None, + "ttfc_ms_max": max(ttfcs) if ttfcs else None, + "steady_interval_ms_mean": (steady_total_ms / steady_count) if steady_count else None, + "steady_interval_ms_spread": (max(steady_means) - min(steady_means)) if steady_means else None, + "underrun_count": sum(int(entry["playback"]["underrun_count"]) for entry in entries), + "chunk_deadline_ms": deadline_ms, + } diff --git a/tests/benchmarks/test_lingbot_world_realtime.py b/tests/benchmarks/test_lingbot_world_realtime.py new file mode 100644 index 000000000..9926513b4 --- /dev/null +++ b/tests/benchmarks/test_lingbot_world_realtime.py @@ -0,0 +1,798 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the LingBot-World realtime streaming benchmark. + +The arithmetic is tested against the server's own formulas rather than against +hand-copied expectations, and the protocol handling runs against a real +WebSocket server that replays ``/v1/realtime/video`` with scripted delays -- so +a benchmark number can be wrong only if the server itself is, not because the +client parsed incorrectly an event or timed the wrong edge. +""" + +from __future__ import annotations + +# ruff: noqa: E402, I001 +import asyncio +import json +import sys +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from benchmarks.lingbot_world.benchmark_lingbot_world_realtime import ( + BenchmarkError, + build_workload, + print_report, + endpoint_url, + parse_args, + resolve_image_reference, + run_session, +) +from benchmarks.lingbot_world.workload import ( + DEFAULT_WARMUP_CHUNKS, + FRAMES_PER_BLOCK, + ChunkRecord, + PromptUpdate, + Workload, + aggregate_metrics, + build_camera_script, + chunks_for_num_frames, + compute_metrics, + default_workload, + load_workload, + num_frames_for_chunks, + percentile, + pixel_frames_in_chunk, + simulate_playback, + steady_chunk_deadline_ms, +) + +_IMAGE = "data:image/png;base64,Zm9v" + + +def _workload(num_chunks: int = 8, fps: int = 16) -> Workload: + return default_workload(image_reference=_IMAGE, num_chunks=num_chunks, fps=fps) + + +# --------------------------------------------------------------------------- # +# Chunk and frame arithmetic +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("num_chunks", list(range(1, 33))) +def test_frame_count_round_trips_through_the_server_formula(num_chunks: int) -> None: + """``num_frames`` must land on a rollout of exactly the requested chunks. + + The server derives chunk count as ``((num_frames - 1) // 4 + 1) // 3``; if + the benchmark's inverse is off by one frame the run silently generates a + different number of blocks than it reports. + """ + assert chunks_for_num_frames(num_frames_for_chunks(num_chunks)) == num_chunks + + +def test_recipe_frame_counts_match_the_documented_chunk_counts() -> None: + # The recipe states three chunks for num_frames 33 and seven for 81. + assert chunks_for_num_frames(33) == 3 + assert chunks_for_num_frames(81) == 7 + assert num_frames_for_chunks(3) == 33 + assert num_frames_for_chunks(7) == 81 + + +def test_first_chunk_is_short_because_the_decoder_expands_it_once() -> None: + assert pixel_frames_in_chunk(0) == 9 + assert pixel_frames_in_chunk(1) == 12 + assert pixel_frames_in_chunk(7) == 12 + # A whole rollout's frames are the sum of its chunks' frames. + assert sum(pixel_frames_in_chunk(index) for index in range(5)) == num_frames_for_chunks(5) + + +def test_camera_script_holds_one_three_entry_action_list_per_chunk() -> None: + script = build_camera_script(4, "orbit") + assert len(script) == 4 + assert all(len(actions) == FRAMES_PER_BLOCK for actions in script) + # orbit advances through the action keys so successive chunks differ. + assert script[0] != script[1] + + +def test_unknown_camera_pattern_is_rejected() -> None: + with pytest.raises(ValueError, match="Unknown camera pattern"): + build_camera_script(2, "sideways") + + +# --------------------------------------------------------------------------- # +# Workload +# --------------------------------------------------------------------------- # + + +def test_session_start_payload_matches_the_documented_protocol() -> None: + workload = _workload(num_chunks=3) + payload = workload.session_start_payload("robbyant/lingbot-world-v2-14b-causal-fast-diffusers") + assert payload["type"] == "session.start" + assert payload["num_frames"] == 33 + assert payload["image_reference"] == {"image_url": _IMAGE} + assert payload["extra_params"]["flow_shift"] == 5.0 + assert len(payload["extra_params"]["camera_action_script"]) == 3 + + +def test_workload_rejects_a_camera_script_with_the_wrong_block_width() -> None: + with pytest.raises(ValueError, match="exactly 3 action lists"): + Workload(prompt="p", image_reference=_IMAGE, camera_script=[[["w"], ["w"]]]) + + +def test_workload_rejects_geometry_the_cache_cannot_serve() -> None: + with pytest.raises(ValueError, match="multiples of 16"): + Workload(prompt="p", image_reference=_IMAGE, camera_script=build_camera_script(1, "forward"), width=831) + + +def test_prompt_update_beyond_the_rollout_is_rejected() -> None: + with pytest.raises(ValueError, match="never fires"): + Workload( + prompt="p", + image_reference=_IMAGE, + camera_script=build_camera_script(2, "forward"), + prompt_updates=(PromptUpdate(after_chunk=5, prompt="later"),), + ) + + +@pytest.mark.parametrize("from_file", [False, True]) +def test_duplicate_prompt_update_boundaries_are_rejected(tmp_path: Path, from_file: bool) -> None: + updates = [ + {"after_chunk": 1, "prompt": "Rain begins"}, + {"after_chunk": 0, "prompt": "Clouds gather"}, + {"after_chunk": 1, "prompt": "Rain stops"}, + ] + with pytest.raises(ValueError, match="Duplicate prompt update after chunk 1"): + if from_file: + path = tmp_path / "rollout.json" + path.write_text(json.dumps({"image": _IMAGE, "num_chunks": 3, "prompt_updates": updates})) + build_workload(parse_args(["--workload", str(path)])) + else: + Workload( + prompt="p", + image_reference=_IMAGE, + camera_script=build_camera_script(3, "forward"), + prompt_updates=( + PromptUpdate(after_chunk=1, prompt="Rain begins"), + PromptUpdate(after_chunk=0, prompt="Clouds gather"), + PromptUpdate(after_chunk=1, prompt="Rain stops"), + ), + ) + + +def test_workload_file_is_loaded_and_cli_overrides_win(tmp_path: Path) -> None: + path = tmp_path / "rollout.json" + path.write_text( + json.dumps( + { + "prompt": "A quiet street", + "image": "https://example.invalid/frame.png", + "num_chunks": 4, + "camera_pattern": "hold", + "fps": 12, + "seed": 7, + "prompt_updates": [{"after_chunk": 1, "prompt": "Rain begins", "transition_chunks": 2}], + } + ) + ) + workload = load_workload(path, overrides={"fps": 24, "seed": None}) + assert workload.prompt == "A quiet street" + assert workload.num_chunks == 4 + assert workload.fps == 24, "an explicit CLI value must beat the file" + assert workload.seed == 7, "a None override must leave the file value alone" + assert workload.prompt_updates[0].transition_chunks == 2 + + +def test_workload_file_without_a_rollout_length_is_rejected(tmp_path: Path) -> None: + path = tmp_path / "rollout.json" + path.write_text(json.dumps({"image": "https://example.invalid/frame.png"})) + with pytest.raises(ValueError, match="camera_action_script or num_chunks"): + load_workload(path) + + +def test_local_image_is_inlined_as_a_data_url(tmp_path: Path) -> None: + image = tmp_path / "frame.png" + image.write_bytes(b"\x89PNG\r\n\x1a\n") + reference = resolve_image_reference(str(image)) + assert reference.startswith("data:image/png;base64,") + assert resolve_image_reference("https://example.invalid/a.png") == "https://example.invalid/a.png" + with pytest.raises(BenchmarkError, match="neither a URL nor a readable file"): + resolve_image_reference(str(tmp_path / "missing.png")) + + +# --------------------------------------------------------------------------- # +# Metric math +# --------------------------------------------------------------------------- # + + +def test_rtf_is_a_cost_not_a_speed() -> None: + """RTF must match vLLM-Omni's own formula, where lower is better. + + ``vllm_omni/metrics/definitions.py::compute_audio_rtf`` is + ``latency / duration`` with the docstring "SLO red line < 1", and + ``benchmarks/patch/patch.py`` sets ``video_rtf = generation_s / + video_duration``. Emitting the reciprocal under the same field name would + silently invert every dashboard that aggregates it, so the direction is + pinned here rather than left to a comment. + """ + # Two chunks, 9 + 12 = 21 frames at 12 fps = 1.75 s of video, generated in 3.5 s. + slow = compute_metrics(_records([1.0, 2.5]), fps=12, warmup_chunks=0, session_wall_s=3.5) + assert slow["video_seconds"] == pytest.approx(1.75) + assert slow["video_rtf"] == pytest.approx(2.0), "twice as slow as real time" + assert slow["video_rtfx"] == pytest.approx(0.5) + + fast = compute_metrics(_records([0.4, 0.475]), fps=12, warmup_chunks=0, session_wall_s=0.875) + assert fast["video_rtf"] == pytest.approx(0.5), "twice as fast as real time" + assert fast["video_rtf"] < 1.0 < slow["video_rtf"] + assert fast["video_rtf"] * fast["video_rtfx"] == pytest.approx(1.0) + + +def test_steady_rtf_is_the_chunk_interval_over_its_deadline() -> None: + """A steady chunk's RTF is just its interval divided by its video duration.""" + metrics = compute_metrics(_records([1.0] + [1.357] * 15), fps=12, warmup_chunks=6) + # 12 frames at 12 fps = 1000 ms of video per steady chunk. + assert metrics["chunk_deadline_ms"] == pytest.approx(1000.0) + assert metrics["steady_video_rtf"] == pytest.approx(1.357, abs=1e-3) + + +def test_percentile_interpolates_so_short_runs_still_report_tails() -> None: + values = [10.0, 20.0, 30.0, 40.0] + assert percentile(values, 0.0) == 10.0 + assert percentile(values, 1.0) == 40.0 + assert percentile(values, 0.5) == 25.0 + assert percentile([5.0], 0.99) == 5.0 + + +def _records(intervals: list[float]) -> list[ChunkRecord]: + """One record per entry; ``intervals[0]`` is the TTFC, the rest are gaps.""" + records: list[ChunkRecord] = [] + arrival = 0.0 + for index, interval in enumerate(intervals): + arrival += interval + records.append( + ChunkRecord( + index=index, + arrival_s=arrival, + inter_arrival_s=interval, + byte_length=1000, + num_frames=pixel_frames_in_chunk(index), + ) + ) + return records + + +def test_steady_state_excludes_the_unsaturated_attention_window() -> None: + """Warmup chunks must not drag the steady-state mean down. + + The first six chunks attend over a growing window, so they are genuinely + cheaper. A benchmark that averages them in reports a cadence the session + will never sustain. + """ + intervals = [1.0] + [0.2] * (DEFAULT_WARMUP_CHUNKS - 1) + [0.8] * 4 + metrics = compute_metrics(_records(intervals), fps=16, warmup_chunks=DEFAULT_WARMUP_CHUNKS) + assert metrics["interval_steady"]["count"] == 3 + assert metrics["interval_steady"]["mean_ms"] == pytest.approx(800.0) + assert metrics["interval_all"]["mean_ms"] < metrics["interval_steady"]["mean_ms"] + + +def test_short_rollouts_report_that_steady_state_was_never_reached() -> None: + metrics = compute_metrics(_records([1.0, 0.5, 0.5]), fps=16, warmup_chunks=DEFAULT_WARMUP_CHUNKS) + assert metrics["steady_state_reached"] is False + assert metrics["interval_steady"] == {} + assert "slo_attainment" not in metrics + + +def test_real_time_factor_and_deadline_use_the_true_frame_counts() -> None: + # Eight chunks: 9 + 7*12 = 93 frames at 16 fps = 5.8125 s of video. + metrics = compute_metrics(_records([1.0] + [0.5] * 7), fps=16, warmup_chunks=0, session_wall_s=5.8125) + assert metrics["total_frames"] == 93 + assert metrics["video_seconds"] == pytest.approx(5.8125) + assert metrics["video_rtf"] == pytest.approx(1.0) + assert metrics["video_rtfx"] == pytest.approx(1.0) + # A steady chunk carries 12 frames, so real time allows 750 ms at 16 fps. + assert metrics["chunk_deadline_ms"] == pytest.approx(750.0) + assert steady_chunk_deadline_ms(12) == pytest.approx(1000.0) + + +def test_slo_attainment_counts_only_steady_chunks() -> None: + intervals = [1.0] + [0.1] * 5 + [0.7, 0.9, 0.7, 0.9] + metrics = compute_metrics(_records(intervals), fps=16, warmup_chunks=6, slo_ms=750.0) + assert metrics["slo_attainment"] == pytest.approx(2 / 3) + assert metrics["slo_violations"] == 1 + + +def test_playback_simulation_finds_stalls_a_percentile_would_hide() -> None: + """One late chunk behind a one-chunk buffer is a visible stall.""" + # 12 frames at 16 fps = 750 ms of playback per chunk; chunk 3 takes 2 s. + intervals = [1.0, 0.5, 0.5, 2.0, 0.5] + result = simulate_playback(_records(intervals), fps=16, buffer_chunks=1) + assert result.underrun_count >= 1 + assert result.first_underrun_chunk == 3 + assert result.total_stall_s > 0.0 + + +def test_a_deeper_buffer_absorbs_the_same_jitter() -> None: + intervals = [1.0, 0.5, 0.5, 1.4, 0.2, 0.2] + fast = simulate_playback(_records(intervals), fps=16, buffer_chunks=1) + buffered = simulate_playback(_records(intervals), fps=16, buffer_chunks=3) + assert buffered.underrun_count < fast.underrun_count + + +def test_playback_rejects_a_buffer_deeper_than_the_rollout() -> None: + with pytest.raises(ValueError, match="exceeds"): + simulate_playback(_records([1.0, 0.5]), fps=16, buffer_chunks=9) + + +def test_aggregate_pools_sessions_and_reports_spread() -> None: + fast = compute_metrics(_records([1.0] + [0.1] * 5 + [0.6] * 4), fps=16, warmup_chunks=6) + slow = compute_metrics(_records([1.0] + [0.1] * 5 + [0.9] * 4), fps=16, warmup_chunks=6) + aggregate = aggregate_metrics([fast, slow], fps=16) + assert aggregate["sessions"] == 2 + assert aggregate["steady_interval_ms_mean"] == pytest.approx(750.0) + assert aggregate["steady_interval_ms_spread"] == pytest.approx(300.0) + assert aggregate["video_rtf_best"] <= aggregate["video_rtf_mean"] <= aggregate["video_rtf_worst"] + + +def test_single_session_aggregate_is_the_session_itself() -> None: + metrics = compute_metrics(_records([1.0] + [0.5] * 3), fps=16, warmup_chunks=0) + aggregate = aggregate_metrics([metrics], fps=16) + assert aggregate["sessions"] == 1 + assert aggregate["video_rtf"] == metrics["video_rtf"] + + +# --------------------------------------------------------------------------- # +# CLI +# --------------------------------------------------------------------------- # + + +def test_endpoint_url_normalizes_http_and_appends_the_ws_path() -> None: + assert endpoint_url(parse_args([])) == "ws://127.0.0.1:8000/v1/realtime/video" + assert endpoint_url(parse_args(["--host", "h", "--port", "9"])) == "ws://h:9/v1/realtime/video" + args = parse_args(["--base-url", "http://example.invalid:8099"]) + assert endpoint_url(args) == "ws://example.invalid:8099/v1/realtime/video" + args = parse_args(["--base-url", "ws://example.invalid:8099/v1/realtime/video"]) + assert endpoint_url(args) == "ws://example.invalid:8099/v1/realtime/video" + + +def test_cli_rejects_a_zero_chunk_buffer() -> None: + with pytest.raises(SystemExit): + parse_args(["--playback-buffer-chunks", "0"]) + with pytest.raises(SystemExit): + parse_args(["--target-fps", "0"]) + + +def test_the_real_time_basis_is_independent_of_the_mux_label() -> None: + """The checkpoint declares no frame rate, so the basis must be statable. + + The same measured cadence has RTF 1.809 at 16 fps and 1.357 at 12 fps; + tying that verdict to the fps sent to the muxer would let a playback label + silently decide whether a run passes. + """ + records = _records([1.0] + [1.357] * 15) + at_16 = compute_metrics(records, fps=16, warmup_chunks=6) + at_12 = compute_metrics(records, fps=12, warmup_chunks=6) + assert at_16["chunk_deadline_ms"] == pytest.approx(750.0) + assert at_12["chunk_deadline_ms"] == pytest.approx(1000.0) + # RTF is cost, so the harsher 16 fps basis scores WORSE (higher). + assert at_16["video_rtf"] > at_12["video_rtf"] + # Same wall clock, same bytes: only the basis moved. + assert at_16["video_rtf"] / at_12["video_rtf"] == pytest.approx(16 / 12) + assert parse_args([]).target_fps is None, "default defers to the workload fps" + + +# --------------------------------------------------------------------------- # +# Protocol handling against a scripted server +# --------------------------------------------------------------------------- # + + +def _serve(script): + """Return an async context manager yielding a ws:// URL for ``script``.""" + from websockets.asyncio.server import Server, serve + + class _Server: + def __init__(self) -> None: + self.server: Server | None = None + self.start_payloads: list[dict] = [] + + async def __aenter__(self) -> str: + async def handler(websocket) -> None: + self.start_payloads.append(json.loads(await websocket.recv())) + await script(websocket, self.start_payloads[-1]) + + self.server = await serve(handler, "127.0.0.1", 0) + port = next(iter(self.server.sockets)).getsockname()[1] + return f"ws://127.0.0.1:{port}/v1/realtime/video" + + async def __aexit__(self, *exc) -> None: + assert self.server is not None + self.server.close() + await self.server.wait_closed() + + return _Server() + + +async def _emit_rollout(websocket, payload, *, intervals=None, chunks=None, trailer=True) -> None: + """Replay the documented event order: start, metadata+binary per chunk, done.""" + num_chunks = chunks if chunks is not None else len(payload["extra_params"]["camera_action_script"]) + await websocket.send(json.dumps({"type": "video.start", "request_id": "req-1", "format": "m4s"})) + for index in range(num_chunks): + if intervals is not None: + await asyncio.sleep(intervals[index]) + frames = pixel_frames_in_chunk(index) + body = bytes([index % 256]) * (100 + index) + await websocket.send( + json.dumps( + { + "type": "video.chunk_metadata", + "request_id": "req-1", + "kind": "media", + "num_frames": frames, + "byte_length": len(body), + "generation_chunk_index": index, + } + ) + ) + await websocket.send(body) + if trailer: + tail = b"\x00" * 8 + await websocket.send( + json.dumps( + { + "type": "video.chunk_metadata", + "request_id": "req-1", + "kind": "trailer", + "num_frames": 0, + "byte_length": len(tail), + "generation_chunk_index": None, + } + ) + ) + await websocket.send(tail) + await websocket.send(json.dumps({"type": "session.done", "stopped": False, "chunks": num_chunks})) + + +def _run(coro): + return asyncio.run(asyncio.wait_for(coro, timeout=30)) + + +def test_session_records_one_chunk_per_block_and_ignores_the_trailer() -> None: + workload = _workload(num_chunks=4) + + async def scenario() -> None: + server = _serve(_emit_rollout) + url = await server.__aenter__() + try: + result = await run_session( + url, + "model", + workload, + first_chunk_timeout=10, + chunk_timeout=10, + ping_interval=0, + collect_media=True, + print_chunks=False, + ) + finally: + await server.__aexit__(None, None, None) + + assert [record.index for record in result.records] == [0, 1, 2, 3] + assert [record.num_frames for record in result.records] == [9, 12, 12, 12] + assert result.request_id == "req-1" + assert result.ttfc_s is not None and result.ttfc_s > 0 + # The trailer's bytes are kept for remuxing but carry no chunk record. + assert len(result.media) == sum(record.byte_length for record in result.records) + 8 + assert server.start_payloads[0]["num_frames"] == workload.num_frames + + _run(scenario()) + + +def test_measured_intervals_track_the_server_cadence() -> None: + workload = _workload(num_chunks=3) + + async def script(websocket, payload) -> None: + await _emit_rollout(websocket, payload, intervals=[0.05, 0.30, 0.05]) + + async def scenario() -> None: + server = _serve(script) + url = await server.__aenter__() + try: + result = await run_session( + url, + "model", + workload, + first_chunk_timeout=10, + chunk_timeout=10, + ping_interval=0, + collect_media=False, + print_chunks=False, + ) + finally: + await server.__aexit__(None, None, None) + + intervals = [record.inter_arrival_s for record in result.records] + assert intervals[1] == pytest.approx(0.30, abs=0.15) + assert intervals[2] < intervals[1], "a fast chunk must not inherit the slow one's interval" + assert result.records[0].arrival_s == pytest.approx(result.ttfc_s) + + _run(scenario()) + + +def test_a_short_rollout_fails_loudly_instead_of_reporting_a_partial_cadence() -> None: + workload = _workload(num_chunks=4) + + async def script(websocket, payload) -> None: + await _emit_rollout(websocket, payload, chunks=2) + + async def scenario() -> None: + server = _serve(script) + url = await server.__aenter__() + try: + with pytest.raises(BenchmarkError, match="Expected 4 chunks"): + await run_session( + url, + "model", + workload, + first_chunk_timeout=10, + chunk_timeout=10, + ping_interval=0, + collect_media=False, + print_chunks=False, + ) + finally: + await server.__aexit__(None, None, None) + + _run(scenario()) + + +def test_a_server_error_event_aborts_the_session() -> None: + workload = _workload(num_chunks=2) + + async def script(websocket, payload) -> None: + await websocket.send(json.dumps({"type": "error", "message": "resolution mismatch"})) + + async def scenario() -> None: + server = _serve(script) + url = await server.__aenter__() + try: + with pytest.raises(BenchmarkError, match="resolution mismatch.*requested 2 chunks"): + await run_session( + url, + "model", + workload, + first_chunk_timeout=10, + chunk_timeout=10, + ping_interval=0, + collect_media=False, + print_chunks=False, + ) + finally: + await server.__aexit__(None, None, None) + + _run(scenario()) + + +def test_a_stalled_server_times_out_rather_than_hanging() -> None: + workload = _workload(num_chunks=2) + + async def script(websocket, payload) -> None: + await websocket.send(json.dumps({"type": "video.start", "request_id": "req-1", "format": "m4s"})) + await asyncio.sleep(5) + + async def scenario() -> None: + server = _serve(script) + url = await server.__aenter__() + try: + with pytest.raises(BenchmarkError, match="No media chunk for"): + await run_session( + url, + "model", + workload, + first_chunk_timeout=0.4, + chunk_timeout=0.4, + ping_interval=0, + collect_media=False, + print_chunks=False, + ) + finally: + await server.__aexit__(None, None, None) + + _run(scenario()) + + +def test_keepalive_pings_reach_the_server_during_a_slow_first_chunk() -> None: + """A compiled first block can outlast the server's stall timeout.""" + workload = _workload(num_chunks=1) + seen: list[str] = [] + + async def script(websocket, payload) -> None: + await websocket.send(json.dumps({"type": "video.start", "request_id": "req-1", "format": "m4s"})) + + async def drain() -> None: + async for message in websocket: + event = json.loads(message) + seen.append(event["type"]) + if event["type"] == "session.ping": + await websocket.send(json.dumps({"type": "session.pong"})) + + drainer = asyncio.create_task(drain()) + await asyncio.sleep(0.35) + await _emit_rollout(websocket, payload, trailer=False) + drainer.cancel() + + async def scenario() -> None: + server = _serve(script) + url = await server.__aenter__() + try: + result = await run_session( + url, + "model", + workload, + first_chunk_timeout=10, + chunk_timeout=10, + ping_interval=0.1, + collect_media=False, + print_chunks=False, + ) + finally: + await server.__aexit__(None, None, None) + assert len(result.records) == 1 + assert "session.ping" in seen + + _run(scenario()) + + +def test_prompt_updates_are_sent_on_their_chunk_boundary() -> None: + workload = Workload( + prompt="start", + image_reference=_IMAGE, + camera_script=build_camera_script(3, "forward"), + prompt_updates=(PromptUpdate(after_chunk=0, prompt="Rain begins", transition_chunks=1),), + ) + received: list[dict] = [] + + async def script(websocket, payload) -> None: + async def drain() -> None: + async for message in websocket: + received.append(json.loads(message)) + + drainer = asyncio.create_task(drain()) + await _emit_rollout(websocket, payload, intervals=[0.02, 0.15, 0.02], trailer=False) + await asyncio.sleep(0.05) + drainer.cancel() + + async def scenario() -> None: + server = _serve(script) + url = await server.__aenter__() + try: + await run_session( + url, + "model", + workload, + first_chunk_timeout=10, + chunk_timeout=10, + ping_interval=0, + collect_media=False, + print_chunks=False, + ) + finally: + await server.__aexit__(None, None, None) + interactions = [event for event in received if event.get("type") == "session.interaction"] + assert len(interactions) == 1 + assert interactions[0]["interaction"]["event"]["prompt"] == "Rain begins" + assert interactions[0]["interaction"]["transition_chunks"] == 1 + + _run(scenario()) + + +@pytest.mark.parametrize("delivered", [0, 1]) +def test_pongs_do_not_extend_media_deadlines(delivered) -> None: + seen = [] + + async def script(websocket, payload) -> None: + if delivered: + await websocket.send( + json.dumps( + { + "type": "video.chunk_metadata", + "kind": "media", + "generation_chunk_index": 0, + "num_frames": 9, + } + ) + ) + await websocket.send(b"media") + async for message in websocket: + seen.append(json.loads(message)["type"]) + await websocket.send(json.dumps({"type": "session.pong"})) + + async def scenario() -> None: + async with _serve(script) as url: + with pytest.raises(BenchmarkError, match=f"after {delivered} chunk"): + await asyncio.wait_for( + run_session( + url, + "model", + _workload(num_chunks=2), + first_chunk_timeout=0.25, + chunk_timeout=0.25, + ping_interval=0.03, + collect_media=False, + print_chunks=False, + ), + timeout=2, + ) + assert "session.ping" in seen + + _run(scenario()) + + +def test_terminal_chunk_is_only_excluded_from_steady_metrics() -> None: + records = _records([1.0] + [0.1] * 5 + [1.2, 0.8, 0.1]) + metrics = compute_metrics(records, fps=12, warmup_chunks=6) + assert metrics["steady_chunks"] == 2 + assert metrics["interval_steady"]["mean_ms"] == pytest.approx(1000) + assert metrics["steady_video_rtf"] == pytest.approx(1) + assert metrics["slo_attainment"] == pytest.approx(0.5) + assert metrics["interval_all"]["count"] == 8 + assert metrics["total_frames"] == 105 + assert metrics["wall_seconds"] == pytest.approx(records[-1].arrival_s) + assert metrics["playback"]["total_stall_ms"] == pytest.approx( + simulate_playback(records, fps=12).total_stall_s * 1000 + ) + + +def test_terminal_chunk_alone_does_not_establish_steady_state() -> None: + metrics = compute_metrics(_records([1.0] * 7), fps=12, warmup_chunks=6) + assert not metrics["steady_state_reached"] + assert "steady_video_rtf" not in metrics + + +def test_aggregate_preserves_custom_deadline_and_weights_intervals() -> None: + fast = compute_metrics(_records([1.0, 0.5, 0.1]), fps=12, warmup_chunks=0, slo_ms=700) + slow = compute_metrics(_records([1.0, 0.9, 0.9, 0.9, 0.1]), fps=12, warmup_chunks=0, slo_ms=700) + aggregate = aggregate_metrics([fast, slow], fps=12) + assert aggregate["chunk_deadline_ms"] == 700 + assert aggregate["steady_interval_ms_mean"] == pytest.approx(800) + slow["chunk_deadline_ms"] = 800 + with pytest.raises(ValueError, match="different chunk deadlines"): + aggregate_metrics([fast, slow], fps=12) + + +@pytest.mark.parametrize( + "flags, expected_shift, expected_negative", + [ + ([], 7.0, "file negative"), + (["--flow-shift", "5", "--negative-prompt", "CLI negative"], 5.0, "CLI negative"), + (["--negative-prompt", ""], 7.0, ""), + ], +) +def test_cli_workload_generation_overrides(tmp_path, flags, expected_shift, expected_negative) -> None: + path = tmp_path / "workload.json" + path.write_text( + json.dumps( + { + "image": _IMAGE, + "num_chunks": 3, + "flow_shift": 7.0, + "negative_prompt": "file negative", + } + ) + ) + workload = build_workload(parse_args(["--workload", str(path), *flags])) + assert workload.flow_shift == expected_shift + assert workload.negative_prompt == expected_negative + + +def test_builtin_workload_keeps_default_flow_shift() -> None: + assert build_workload(parse_args(["--image", _IMAGE])).flow_shift == 5.0 + + +def test_short_sample_report_qualifies_tail_percentiles(capsys) -> None: + metrics = compute_metrics(_records([1.0] * 16), fps=12) + print_report(metrics, workload=_workload(num_chunks=16), endpoint="ws://test", sessions=1, target_fps=12) + assert "not reliable tail estimates" in capsys.readouterr().out