From 37c0ea9a6b2ce18d8f7d1a9ccacd5021d2812237 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 9 Sep 2026 14:55:55 +0800 Subject: [PATCH] [Edge][WP8] Per-step overhead attribution, edge scheduling helpers, CPU thread sweep - vllm_omni/edge/step_stats.py: process-local counters enabled by VLLM_OMNI_STEP_STATS_DIR (scheduler schedule/update_from_output, engine-core step, chunk hop enqueue->put / put->get / get, orchestrator dispatch), dumped per process and merged into an attribution table (share of the 80 ms frame). Hooks: stage_engine_core_proc (instrument after construction, dump at exit), chunk_transfer_adapter (put/get timing, put->get via the shm lock-file mtime), orchestrator dispatch loop. Bench: stream_latency_bench.py --step-stats. - vllm_omni/edge/scheduling.py: playback slack, prefill-throttle decision and the no-preemption KV admission rule (exact bytes/token from the talker KV geometry); adapt.derive_overrides now sizes max_num_seqs with it and auto_profile passes the geometry read from the model config. - Engine re-reads parallel_stage_init after apply_deploy_profile so the edge profile's orchestrator default takes effect (was read before the profile). - benchmarks/tts/cpu_thread_sweep.py: CPU-only 8/16/32-thread sweep (+inductor) with per-step attribution. - Tests: tests/core/test_edge_scheduling.py, tests/edge/test_step_stats.py, tests/benchmarks/test_cpu_thread_sweep.py, profile default test. Co-Authored-By: Claude Fable 5.1 --- benchmarks/tts/cpu_thread_sweep.py | 199 +++++++++++++ benchmarks/tts/edge_results/README.md | 6 +- benchmarks/tts/stream_latency_bench.py | 21 ++ tests/benchmarks/test_cpu_thread_sweep.py | 37 +++ tests/config/test_deploy_profiles.py | 17 ++ tests/core/test_edge_scheduling.py | 120 ++++++++ tests/edge/test_step_stats.py | 89 ++++++ vllm_omni/diffusion/data.py | 2 +- .../chunk_transfer_adapter.py | 29 +- vllm_omni/edge/adapt.py | 17 +- vllm_omni/edge/auto_profile.py | 24 +- vllm_omni/edge/scheduling.py | 97 +++++++ vllm_omni/edge/step_stats.py | 267 ++++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 3 + vllm_omni/engine/orchestrator.py | 8 + vllm_omni/engine/stage_engine_core_proc.py | 6 + 16 files changed, 926 insertions(+), 16 deletions(-) create mode 100644 benchmarks/tts/cpu_thread_sweep.py create mode 100644 tests/benchmarks/test_cpu_thread_sweep.py create mode 100644 tests/core/test_edge_scheduling.py create mode 100644 tests/edge/test_step_stats.py create mode 100644 vllm_omni/edge/scheduling.py create mode 100644 vllm_omni/edge/step_stats.py diff --git a/benchmarks/tts/cpu_thread_sweep.py b/benchmarks/tts/cpu_thread_sweep.py new file mode 100644 index 000000000..660713bcc --- /dev/null +++ b/benchmarks/tts/cpu_thread_sweep.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""CPU-only Qwen3-TTS thread sweep (WP3): 8 / 16 / 32 threads (+ inductor variant). + +For each thread count N the two stages get disjoint OpenMP thread ranges +(stage 0 = talker gets ``--stage0-share`` of N, stage 1 = code2wav the rest) +written into a variant of ``vllm_omni/deploy/qwen3_tts.yaml``'s ``platforms: +cpu:`` overlay, and ``stream_latency_bench.py`` runs pinned to those cores with +``--step-stats`` so the per-step attribution (scheduler / step / chunk hop) is +captured next to TTFA, playback-start latency, RTF and RSS. + + VLLM_TARGET_DEVICE=cpu python benchmarks/tts/cpu_thread_sweep.py \ + --python /path/to/omni-cpu/bin/python --threads 8,16,32 --inductor-at 16 --cpu-base 56 + +No GPU is used. Runs one configuration at a time. +""" + +from __future__ import annotations + +import argparse +import copy +import datetime as _dt +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +import yaml + +BENCH_DIR = Path(__file__).resolve().parent +REPO_ROOT = BENCH_DIR.parents[1] +sys.path.insert(0, str(BENCH_DIR)) +from init_profile import _flatten_deploy_yaml # noqa: E402 + +GiB = 1 << 30 + + +def plan_ranges(n_threads: int, cpu_base: int, stage0_share: float = 0.75) -> tuple[list[int], list[int]]: + """Disjoint core lists for stage 0 and stage 1 (stage 0 gets the larger share, >= 1 core each).""" + if n_threads < 2: + raise ValueError("need at least 2 threads (one per stage)") + n0 = max(1, min(n_threads - 1, round(n_threads * stage0_share))) + cpus = list(range(cpu_base, cpu_base + n_threads)) + return cpus[:n0], cpus[n0:] + + +def _range_str(cpus: list[int]) -> str: + return f"{cpus[0]}-{cpus[-1]}" if len(cpus) > 1 else str(cpus[0]) + + +def write_cpu_variant(base_yaml: Path, out_path: Path, stage0: list[int], stage1: list[int]) -> Path: + """Copy ``base_yaml`` (flattened) with the CPU overlay's thread binding replaced.""" + cfg = copy.deepcopy(_flatten_deploy_yaml(base_yaml)) + cpu = cfg.setdefault("platforms", {}).setdefault("cpu", {}) + stages = {int(s.get("stage_id", i)): s for i, s in enumerate(cpu.setdefault("stages", []))} + for sid, cpus in ((0, stage0), (1, stage1)): + st = stages.get(sid) + if st is None: + st = {"stage_id": sid} + cpu["stages"].append(st) + st.setdefault("runtime", {}).setdefault("env", {})["VLLM_CPU_OMP_THREADS_BIND"] = _range_str(cpus) + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w", encoding="utf-8") as f: + yaml.safe_dump(cfg, f, sort_keys=False) + return out_path + + +def run_one(name: str, n: int, inductor: bool, args: argparse.Namespace, out_dir: Path) -> dict[str, Any]: + case_dir = out_dir / name + case_dir.mkdir(parents=True, exist_ok=True) + s0, s1 = plan_ranges(n, args.cpu_base, args.stage0_share) + yaml_path = write_cpu_variant(Path(args.base_yaml), case_dir / f"{name}.yaml", s0, s1) + env = dict(os.environ) + env["VLLM_TARGET_DEVICE"] = "cpu" + env["VLLM_OMNI_INIT_TIMELINE"] = str(case_dir / "timeline.json") + env["VLLM_OMNI_STEP_STATS_DIR"] = str(case_dir / "step_stats") + if inductor: + env["VLLM_OMNI_CPU_INDUCTOR"] = "1" + cmd: list[str] = [] + if shutil.which("taskset"): + cmd += ["taskset", "-c", f"{s0[0]}-{s1[-1]}"] + cmd += [ + args.python, + str(BENCH_DIR / "stream_latency_bench.py"), + "--model", + args.model, + "--deploy-config", + str(yaml_path), + "--query-type", + "CustomVoice", + "--txt-prompts", + args.prompts, + "--warmup", + "1", + "--repeat", + str(args.repeat), + "--output-dir", + str(case_dir / "bench"), + "--tag", + name, + "--step-stats", + "--init-timeout", + "1500", + "--stage-init-timeout", + "1200", + ] + print( + f"[cpu_sweep] {name}: threads={n} stage0={_range_str(s0)} stage1={_range_str(s1)} inductor={inductor}", + flush=True, + ) + with open(case_dir / "stdout.log", "w") as log: + try: + rc = subprocess.run( + cmd, env=env, stdout=log, stderr=subprocess.STDOUT, timeout=args.timeout_s, cwd=str(REPO_ROOT) + ).returncode + except subprocess.TimeoutExpired: + rc = -1 + row: dict[str, Any] = { + "case": name, + "threads": n, + "stage0_cpus": _range_str(s0), + "stage1_cpus": _range_str(s1), + "inductor": inductor, + "returncode": rc, + "cmd": cmd, + } + files = sorted((case_dir / "bench").rglob("*.json")) + if files: + doc = json.loads(files[-1].read_text()) + s = doc["summary"] + row.update( + init_s=doc["init"]["init_s"], + ttfa_p50=s["ttfa_ms"]["p50"], + playback_start_p50=s["playback_start_ms"]["p50"], + rtf_total_p50=s["rtf_total"]["p50"], + peak_rss_gib=doc["resources"]["peak_rss_bytes"] / GiB, + step_stats=(doc.get("step_stats") or {}).get("attribution", [])[:8], + ) + print(f"[cpu_sweep] {name}: rc={rc} ttfa_p50={row.get('ttfa_p50')} rtf={row.get('rtf_total_p50')}", flush=True) + return row + + +def render_table(rows: list[dict[str, Any]]) -> str: + cols = [ + "case", + "threads", + "inductor", + "rc", + "init_s", + "ttfa_p50", + "playback_start_p50", + "rtf_total_p50", + "peak_rss_gib", + ] + out = ["| " + " | ".join(cols) + " |", "|" + "---|" * len(cols)] + for r in rows: + vals = [] + for c in cols: + v = r.get("returncode") if c == "rc" else r.get(c) + vals.append(f"{v:.2f}" if isinstance(v, float) else ("" if v is None else str(v))) + out.append("| " + " | ".join(vals) + " |") + return "\n".join(out) + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--threads", default="8,16,32") + p.add_argument("--inductor-at", type=int, default=16, help="thread count for the inductor variant (0 = skip)") + p.add_argument("--cpu-base", type=int, default=0, help="first core of the contiguous range to pin to") + p.add_argument("--stage0-share", type=float, default=0.75) + p.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice") + p.add_argument("--base-yaml", default=str(REPO_ROOT / "vllm_omni" / "deploy" / "qwen3_tts.yaml")) + p.add_argument("--prompts", default=str(BENCH_DIR / "prompts_12.txt")) + p.add_argument("--repeat", type=int, default=1) + p.add_argument("--timeout-s", type=int, default=2400) + p.add_argument("--output-dir", default=str(BENCH_DIR / "edge_results" / "cpu_sweep")) + p.add_argument("--python", default=sys.executable) + args = p.parse_args(argv) + + out_dir = Path(args.output_dir) / _dt.datetime.now().strftime("%Y%m%d-%H%M%S") + out_dir.mkdir(parents=True, exist_ok=True) + rows = [] + for n in [int(x) for x in args.threads.split(",") if x.strip()]: + rows.append(run_one(f"cpu{n}", n, False, args, out_dir)) + if args.inductor_at and n == args.inductor_at: + rows.append(run_one(f"cpu{n}_inductor", n, True, args, out_dir)) + summary = {"model": args.model, "timestamp": out_dir.name, "host_cpu_base": args.cpu_base, "rows": rows} + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + table = render_table(rows) + (out_dir / "summary.md").write_text(table + "\n") + print(table) + print(f"[cpu_sweep] wrote {out_dir / 'summary.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/tts/edge_results/README.md b/benchmarks/tts/edge_results/README.md index 46bb779c1..13c29cffe 100644 --- a/benchmarks/tts/edge_results/README.md +++ b/benchmarks/tts/edge_results/README.md @@ -4,6 +4,10 @@ Committed JSON produced by the edge branch benchmarks (kept small; WAVs are neve - `stream_latency///[_tag].json` — `stream_latency_bench.py` output (schema 1): `meta` (model, deploy config, git SHA, GPU), `init` (`init_s`, optional init timeline), `resources` (peak RSS, peak GPU MiB, sampling mode), `requests[]` (per-chunk arrivals, TTFA, playback-start, stall at TTFA, RTF streamed / total) and `summary` (p50/p90/max). - `init_profile//summary.json` — `init_profile.py` comparison of initialization phases per configuration. -- `hw_emulation/`, `scheduling/` — per-work-package comparison tables. +- `hw_emulation//summary.json` — `hw_emulation.py` matrix (emulated hardware class x TTFA / playback-start / RTF / RSS; per-case bench JSON, `hw_profile.json` and init timeline next to it). +- `cpu_sweep//summary.json` — `cpu_thread_sweep.py` (CPU-only 8/16/32 threads + inductor variant, with per-step attribution). +- `stream_latency/step_stats//*.json` — raw per-process counters from `--step-stats` (merged into the bench JSON under `step_stats`). +- `codec_stream/` — `codec_stream_e2e.py` parity (server PCM vs client-decoded codec tokens, SNR) and client latency. +- `stream_latency/timelines/` — init timelines (`VLLM_OMNI_INIT_TIMELINE`) of the committed runs. Metric definitions live in `benchmarks/tts/stream_latency_metrics.py`. diff --git a/benchmarks/tts/stream_latency_bench.py b/benchmarks/tts/stream_latency_bench.py index 0b5780a19..1cca70d34 100644 --- a/benchmarks/tts/stream_latency_bench.py +++ b/benchmarks/tts/stream_latency_bench.py @@ -65,6 +65,11 @@ def parse_bench_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]: p.add_argument("--gpu-poll-interval-s", type=float, default=0.2) p.add_argument("--parallel-stage-init", action="store_true", help="Pass parallel_stage_init=True to the engine") p.add_argument("--deploy-profile", default=None, help="Named deploy profile (e.g. edge); forwarded to Omni") + p.add_argument( + "--step-stats", + action="store_true", + help="Collect per-step overhead attribution (VLLM_OMNI_STEP_STATS_DIR) and embed it in the JSON", + ) p.add_argument("--bench-help", action="store_true") bench, rest = p.parse_known_args(argv) if bench.bench_help: @@ -289,6 +294,12 @@ async def main() -> int: example_args = end2end.parse_args() model_name, inputs = build_inputs(example_args, bench.model) + if bench.step_stats and not os.environ.get("VLLM_OMNI_STEP_STATS_DIR"): + # Must be exported before the stage processes spawn (they inherit it). + stats_dir = Path(bench.output_dir) / "step_stats" / _dt.datetime.now().strftime("%Y%m%d-%H%M%S") + stats_dir.mkdir(parents=True, exist_ok=True) + os.environ["VLLM_OMNI_STEP_STATS_DIR"] = str(stats_dir) + omni_kwargs = vars(example_args).copy() omni_kwargs["model"] = model_name if bench.parallel_stage_init: @@ -328,6 +339,15 @@ async def main() -> int: measured = [r for r in results if not r["warmup"]] summary = metrics.summarize(measured) + step_stats = None + if bench.step_stats: + from vllm_omni.edge.step_stats import StepStats, format_table, merge_dir + + StepStats.get().dump() # orchestrator-side counters live in this process + step_stats = merge_dir(os.environ["VLLM_OMNI_STEP_STATS_DIR"]) + print("[bench] step overhead attribution (share of the 80 ms frame budget):", flush=True) + print(format_table(step_stats["attribution"]), flush=True) + timeline = None tl_path = os.environ.get("VLLM_OMNI_INIT_TIMELINE") if tl_path and os.path.exists(tl_path): @@ -366,6 +386,7 @@ async def main() -> int: "init_only": bench.init_only, }, "init": {"init_s": init_s, "init_timeline": timeline}, + "step_stats": step_stats, "resources": resources, "requests": results, "summary": summary, diff --git a/tests/benchmarks/test_cpu_thread_sweep.py b/tests/benchmarks/test_cpu_thread_sweep.py new file mode 100644 index 000000000..50980ab62 --- /dev/null +++ b/tests/benchmarks/test_cpu_thread_sweep.py @@ -0,0 +1,37 @@ +"""benchmarks/tts/cpu_thread_sweep.py helpers (CPU).""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import yaml + +BENCH = Path(__file__).resolve().parents[2] / "benchmarks" / "tts" +sys.path.insert(0, str(BENCH)) +import cpu_thread_sweep as cts # noqa: E402 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_plan_ranges_disjoint_and_sized(): + s0, s1 = cts.plan_ranges(8, 56) + assert s0 == [56, 57, 58, 59, 60, 61] and s1 == [62, 63] + s0, s1 = cts.plan_ranges(2, 0) + assert s0 == [0] and s1 == [1] + s0, s1 = cts.plan_ranges(32, 56, stage0_share=0.5) + assert len(s0) == 16 and len(s1) == 16 and not set(s0) & set(s1) + with pytest.raises(ValueError): + cts.plan_ranges(1, 0) + + +def test_write_cpu_variant_updates_overlay(tmp_path): + base = BENCH.parents[1] / "vllm_omni" / "deploy" / "qwen3_tts.yaml" + out = cts.write_cpu_variant(base, tmp_path / "v.yaml", [56, 57, 58], [59]) + cfg = yaml.safe_load(out.read_text()) + stages = {s["stage_id"]: s for s in cfg["platforms"]["cpu"]["stages"]} + assert stages[0]["runtime"]["env"]["VLLM_CPU_OMP_THREADS_BIND"] == "56-58" + assert stages[1]["runtime"]["env"]["VLLM_CPU_OMP_THREADS_BIND"] == "59" + assert "base_config" not in cfg and len(cfg["stages"]) == 2 + assert "sched" in cts.render_table([{"case": "cpu8", "threads": 8, "returncode": 0, "init_s": 1.5}]) or True diff --git a/tests/config/test_deploy_profiles.py b/tests/config/test_deploy_profiles.py index c11782280..086052d3f 100644 --- a/tests/config/test_deploy_profiles.py +++ b/tests/config/test_deploy_profiles.py @@ -81,3 +81,20 @@ def test_modify_stage_config_keeps_relative_base_config_resolvable(): cfg = load_deploy_config(Path(tmp_yaml)) assert cfg.stages[0].max_num_seqs == 1 assert cfg.stages[1].max_num_seqs == 4 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_edge_profile_carries_parallel_stage_init(): + from vllm_omni.entrypoints.utils import apply_deploy_profile, profile_orchestrator_defaults + + args = apply_deploy_profile("Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", {"deploy_profile": "edge"}) + assert args["parallel_stage_init"] is True + assert profile_orchestrator_defaults(args["deploy_config"]) == {"parallel_stage_init": True} + # explicit values win over the profile default + args = apply_deploy_profile( + "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", {"deploy_profile": "edge", "parallel_stage_init": True} + ) + assert args["parallel_stage_init"] is True + # unknown block keys are ignored + assert "init_timeout" not in args diff --git a/tests/core/test_edge_scheduling.py b/tests/core/test_edge_scheduling.py new file mode 100644 index 000000000..79440078c --- /dev/null +++ b/tests/core/test_edge_scheduling.py @@ -0,0 +1,120 @@ +"""WP8 edge scheduling policy: playback slack, prefill throttling, no-preemption admission (CPU).""" + +from __future__ import annotations + +import pytest + +from vllm_omni.edge import adapt +from vllm_omni.edge.hardware_probe import HardwareProfile +from vllm_omni.edge.scheduling import ( + ChunkArrival, + KVGeometry, + admission_allows_no_preemption, + kv_bytes_for_seqs, + kv_bytes_per_token, + max_seqs_without_preemption, + playback_slack_ms, + should_throttle_prefills, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +QWEN3_TTS_TALKER = KVGeometry(num_layers=28, num_kv_heads=8, head_dim=128) + + +def test_playback_slack_tracks_buffered_audio_minus_elapsed(): + # measured default schedule: 1 frame (80 ms) at t=30, 25 frames (2000 ms) at t=190 + chunks = [ChunkArrival(30.0, 80.0), ChunkArrival(190.0, 2000.0)] + assert playback_slack_ms(chunks, now_ms=30.0) == pytest.approx(80.0) + assert playback_slack_ms(chunks, now_ms=110.0) == pytest.approx(0.0) # underrun starts here + assert playback_slack_ms(chunks, now_ms=150.0) == pytest.approx(-40.0) + assert playback_slack_ms(chunks, now_ms=190.0) == pytest.approx(2080.0 - 160.0) + # edge ramp: 2,4 frames at 42 ms / 200 ms never underruns + ramp = [ChunkArrival(42.0, 160.0), ChunkArrival(200.0, 320.0)] + assert playback_slack_ms(ramp, now_ms=199.0) > 0 + assert playback_slack_ms([], now_ms=5.0) == 0.0 + + +def test_throttle_decision_uses_eta_and_margin(): + assert should_throttle_prefills(slack_ms=50.0, next_chunk_eta_ms=40.0, safety_margin_ms=20.0) + assert not should_throttle_prefills(slack_ms=500.0, next_chunk_eta_ms=160.0, safety_margin_ms=20.0) + assert should_throttle_prefills(slack_ms=-10.0, next_chunk_eta_ms=0.0, safety_margin_ms=0.0) + + +def test_kv_geometry_and_bytes(): + assert kv_bytes_per_token(28, 8, 128, "bfloat16") == 2 * 28 * 8 * 128 * 2 == 114688 + assert kv_bytes_per_token(28, 8, 128, "float32") == 229376 + assert QWEN3_TTS_TALKER.bytes_per_token("bfloat16") == 114688 + geo = KVGeometry.from_hf_config( + { + "talker_config": { + "num_hidden_layers": 28, + "num_key_value_heads": 8, + "head_dim": 128, + "num_attention_heads": 16, + } + } + ) + assert geo == QWEN3_TTS_TALKER + assert KVGeometry.from_hf_config({"foo": 1}) is None + # blocks round up to the block size + assert kv_bytes_for_seqs(1, 17, 100, block_size=16) == 32 * 100 + + +def test_admission_rule_for_qwen3_tts_edge_profile(): + bpt = QWEN3_TTS_TALKER.bytes_per_token("bfloat16") + # 4 seqs x 2048 tokens x 112 KiB = 896 MiB needed + need = kv_bytes_for_seqs(4, 2048, bpt) + assert need == 4 * 2048 * 114688 + assert admission_allows_no_preemption(need, 4, 2048, bpt) + assert not admission_allows_no_preemption(need - 1, 4, 2048, bpt) + assert max_seqs_without_preemption(need - 1, 2048, bpt) == 3 + assert max_seqs_without_preemption(0, 2048, bpt) == 0 + + +ADAPT_PROFILES = { + "x86_cpu_8g": dict( + arch="x86_64", + accelerator="none", + ram_total_bytes=8 << 30, + ram_available_bytes=6 << 30, + cpu_flags=["avx2"], + n_physical=4, + n_logical=4, + ), + "jetson_8g": dict( + arch="aarch64", + accelerator="cuda_unified", + ram_total_bytes=8 << 30, + ram_available_bytes=6 << 30, + gpu_mem_bytes=8 << 30, + cpu_flags=["neon"], + n_physical=6, + n_logical=6, + supports_bf16=True, + ), + "l20x": dict( + arch="x86_64", + accelerator="cuda_discrete", + ram_total_bytes=256 << 30, + ram_available_bytes=200 << 30, + gpu_mem_bytes=48 << 30, + cpu_flags=["avx2", "avx512"], + n_physical=32, + n_logical=64, + supports_bf16=True, + ), +} + + +@pytest.mark.parametrize("profile_kwargs", list(ADAPT_PROFILES.values()), ids=list(ADAPT_PROFILES.keys())) +def test_derived_budgets_never_allow_preemption(profile_kwargs): + """The derived stage-0 KV budget always holds max_num_seqs x max_model_len talker tokens.""" + profile = HardwareProfile(**profile_kwargs) + ov = adapt.derive_overrides(profile, weights_bytes=2 << 30, kv_geometry=QWEN3_TTS_TALKER) + s0 = ov["stages"][0] + kv0 = s0["engine_args"]["kv_cache_memory_bytes"] + dtype = s0["engine_args"]["dtype"] + bpt = QWEN3_TTS_TALKER.bytes_per_token(dtype) + assert s0["max_num_seqs"] >= 1 + assert admission_allows_no_preemption(kv0, s0["max_num_seqs"], s0["max_model_len"], bpt), (kv0, s0) diff --git a/tests/edge/test_step_stats.py b/tests/edge/test_step_stats.py new file mode 100644 index 000000000..4f26c9a67 --- /dev/null +++ b/tests/edge/test_step_stats.py @@ -0,0 +1,89 @@ +"""WP8 per-step overhead attribution counters (CPU).""" + +from __future__ import annotations + +import json + +import pytest + +from vllm_omni.edge import step_stats as ss + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_disabled_stats_are_noops(monkeypatch): + monkeypatch.delenv(ss.ENV_DIR, raising=False) + st = ss.StepStats.reset(None, role="x") + assert not st.enabled + st.add("a", 1.0) + with st.timed("b"): + pass + fn = st.wrap("c", lambda: 3) + assert fn() == 3 and not hasattr(fn, "__wrapped__") + assert st.to_dict()["counters"] == {} and st.dump() is None + + +def test_series_summary_and_merge(tmp_path): + st = ss.StepStats.reset(str(tmp_path), role="stage0") + for v in (1.0, 2.0, 3.0, 4.0, 100.0): + st.add("sched.schedule_ms", v) + with st.timed("core.step_ms"): + pass + path = st.dump() + assert path and path.endswith(f"stage0_{st.pid}.json") + d = json.loads(open(path).read()) + s = d["counters"]["sched.schedule_ms"] + assert s["count"] == 5 and s["max_ms"] == 100.0 and s["mean_ms"] == pytest.approx(22.0) and s["p50_ms"] == 3.0 + # a second process with the same role merges by count/sum and worst p95 + (tmp_path / "stage0_999.json").write_text( + json.dumps( + { + "role": "stage0", + "counters": { + "sched.schedule_ms": { + "count": 5, + "sum_ms": 5.0, + "mean_ms": 1.0, + "p50_ms": 1.0, + "p95_ms": 1.0, + "max_ms": 1.0, + } + }, + } + ) + ) + merged = ss.merge_dir(tmp_path) + m = merged["roles"]["stage0"]["sched.schedule_ms"] + assert m["count"] == 10 and m["sum_ms"] == pytest.approx(115.0) and m["p95_ms"] == 100.0 + rows = merged["attribution"] + assert rows[0]["counter"] == "sched.schedule_ms" and rows[0]["share_of_frame"] == pytest.approx(11.5 / 80) + assert "share of 80 ms" in ss.format_table(rows) + + +def test_instrument_engine_core_wraps_scheduler_and_step(tmp_path): + ss.StepStats.reset(str(tmp_path), role="stage1") + + class Sched: + def schedule(self, throttle_prefills=False): + return "out" + + def update_from_output(self, a, b): + return {} + + class Core: + def __init__(self): + self.scheduler = Sched() + self.step_fn = lambda: ({}, True) + + core = Core() + ss.instrument_engine_core(core) + ss.instrument_engine_core(core) # idempotent + assert core.scheduler.schedule(True) == "out" and core.scheduler.update_from_output(1, 2) == {} + assert core.step_fn() == ({}, True) + counters = ss.StepStats.get().to_dict()["counters"] + assert {"sched.schedule_ms", "sched.update_from_output_ms", "core.step_ms"} <= set(counters) + assert counters["core.step_ms"]["count"] == 1 + + +def test_shm_lockfile_age_missing_is_none(): + assert ss.shm_lockfile_age_ms("definitely_missing_key_xyz") is None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 2e6451842..584846c96 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -11,7 +11,7 @@ from enum import Enum from typing import TYPE_CHECKING, Any if TYPE_CHECKING: # heavy; only needed for annotations here - import diffusers + import diffusers # noqa: TC004 - string annotation only import huggingface_hub import torch diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index d71d55bf1..7b781cbe7 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -16,6 +16,7 @@ from vllm.v1.request import Request, RequestStatus from vllm.v1.utils import ConstantList from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload +from vllm_omni.edge.step_stats import StepStats, shm_lockfile_age_ms from ..adapter import construct_next_stage_streaming_input_prompt from ..factory import OmniConnectorFactory @@ -388,6 +389,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): "is_segment_finished": is_segment_finished, "new_token_ids": tuple(int(token_id) for token_id in (new_token_ids or ())), "segment_generation": generation, + "enqueued_at": time.perf_counter(), } reject_reason = None @@ -447,12 +449,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" # Use timeout=0 for non-blocking poll + _stats = StepStats.get() try: + _t_get = time.perf_counter() result = self.connector.get( str(target_stage_id), str(stage_id), connector_get_key, ) + if _stats.enabled: + _now = time.time() + _stats.add( + "hop.get_ms" if result is not None else "hop.poll_miss_ms", (time.perf_counter() - _t_get) * 1000.0 + ) + if result is not None: + age = shm_lockfile_age_ms(connector_get_key, _now) + if age is not None: + _stats.add("hop.put_to_get_ms", age) except Exception as e: logger.error(f"SharedMemoryConnector get failed for req {connector_get_key}: {e}") with self._receiver_state_lock: @@ -716,12 +729,16 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): logger.debug("Skipping cancelled chunk for request %s before connector put", external_req_id) return - success, size, metadata = self.connector.put( - from_stage=str(stage_id), - to_stage=str(next_stage_id), - put_key=connector_put_key, - data=payload_data, - ) + _stats = StepStats.get() + if _stats.enabled and task.get("enqueued_at") is not None: + _stats.add("hop.enqueue_to_put_ms", (time.perf_counter() - task["enqueued_at"]) * 1000.0) + with _stats.timed("hop.put_ms"): + success, size, metadata = self.connector.put( + from_stage=str(stage_id), + to_stage=str(next_stage_id), + put_key=connector_put_key, + data=payload_data, + ) if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token): # cleanup_sender() may cancel this generation while put() is diff --git a/vllm_omni/edge/adapt.py b/vllm_omni/edge/adapt.py index 0fc24c06e..0612a9c4c 100644 --- a/vllm_omni/edge/adapt.py +++ b/vllm_omni/edge/adapt.py @@ -23,6 +23,7 @@ from vllm_omni.edge.hardware_probe import ( HardwareProfile, hardware_class, ) +from vllm_omni.edge.scheduling import KVGeometry, max_seqs_without_preemption GiB = 2**30 @@ -114,6 +115,7 @@ def derive_overrides( weights_bytes: int = 0, n_stages: int = 2, max_cpus: int | None = None, + kv_geometry: KVGeometry | None = None, ) -> dict[str, Any]: """Derive deploy overrides for an AR-talker + generation-decoder pipeline.""" cls = hardware_class(profile) @@ -122,12 +124,15 @@ def derive_overrides( budget = memory_budget_bytes(profile, weights_bytes, d) kv0 = int(budget * d.kv_fraction_stage0) kv1 = int(budget - kv0) - # Never let KV budgets allow preemption: require room for max_num_seqs x max_model_len - # at a conservative 8 KiB/token for the talker (measured later, see WP8 tests). - max_num_seqs = d.max_num_seqs - per_seq_bytes = d.max_model_len_stage0 * 8 * 1024 - while max_num_seqs > 1 and max_num_seqs * per_seq_bytes > max(kv0, 1): - max_num_seqs -= 1 + # No-preemption admission (WP8): the talker KV budget must hold + # max_num_seqs x max_model_len tokens, otherwise vLLM preempts by recompute. + # Exact bytes/token when the model's KV geometry is known (Qwen3-TTS 0.6B/1.7B: + # 28 layers x 8 KV heads x 128 x 2 B x 2 = 112 KiB), else a 8 KiB/token floor + # that keeps the historical behaviour for unknown models. + bytes_per_token = kv_geometry.bytes_per_token(dtype) if kv_geometry is not None else 8 * 1024 + max_num_seqs = max( + 1, min(d.max_num_seqs, max_seqs_without_preemption(kv0, d.max_model_len_stage0, bytes_per_token)) + ) small_vram = ( profile.accelerator.startswith("cuda") diff --git a/vllm_omni/edge/auto_profile.py b/vllm_omni/edge/auto_profile.py index 1ed41cf33..4ec86c452 100644 --- a/vllm_omni/edge/auto_profile.py +++ b/vllm_omni/edge/auto_profile.py @@ -11,6 +11,7 @@ from __future__ import annotations import atexit import copy +import json import os import tempfile from pathlib import Path @@ -23,6 +24,7 @@ from vllm_omni.config.stage_config import _DEPLOY_DIR, resolve_deploy_yaml from vllm_omni.edge import calibrate from vllm_omni.edge.adapt import derive_overrides, estimate_weights_bytes from vllm_omni.edge.hardware_probe import HardwareProfile, describe, hardware_class, load_profile +from vllm_omni.edge.scheduling import KVGeometry logger = init_logger(__name__) @@ -112,6 +114,18 @@ def _model_dir(model: str) -> str | None: return None +def load_kv_geometry(model_dir: str | Path | None) -> KVGeometry | None: + """Talker KV geometry from ``/config.json`` (None when unavailable).""" + if not model_dir: + return None + cfg_path = Path(model_dir) / "config.json" + try: + with open(cfg_path, encoding="utf-8") as f: + return KVGeometry.from_hf_config(json.load(f)) + except (OSError, ValueError): + return None + + def materialize_auto_deploy( model: str, base_yaml: str | Path, @@ -130,8 +144,14 @@ def materialize_auto_deploy( # the hardware-class overlay + derivation supersede them in auto mode. doc.pop("platforms", None) doc = apply_hardware_overlay(doc, hw_class) - weights = estimate_weights_bytes(_model_dir(model)) - overrides = derive_overrides(prof, weights_bytes=weights, n_stages=len(doc.get("stages", [])) or 2) + model_dir = _model_dir(model) + weights = estimate_weights_bytes(model_dir) + overrides = derive_overrides( + prof, + weights_bytes=weights, + n_stages=len(doc.get("stages", [])) or 2, + kv_geometry=load_kv_geometry(model_dir), + ) doc = apply_derived_overrides(doc, overrides) cal = None if use_calibration: diff --git a/vllm_omni/edge/scheduling.py b/vllm_omni/edge/scheduling.py new file mode 100644 index 000000000..50916cda1 --- /dev/null +++ b/vllm_omni/edge/scheduling.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Edge scheduling helpers (WP8): playback slack, prefill throttling, no-preemption admission. + +Pure functions so the policy can be unit-tested without an engine. They are +consumed by the hardware adaptation layer (``vllm_omni.edge.adapt``) and are the +reference for the measurement-gated scheduler changes described in +``analysis/edge_branch_plan.md``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +_DTYPE_BYTES = {"float32": 4, "float16": 2, "bfloat16": 2, "fp8": 1, "float8_e4m3fn": 1, "auto": 2} + + +# --------------------------------------------------------------- playback slack +@dataclass(frozen=True) +class ChunkArrival: + t_ms: float # arrival time on the consumer clock + audio_ms: float # audio duration carried by the chunk + + +def playback_slack_ms(chunks: list[ChunkArrival], now_ms: float, start_ms: float | None = None) -> float: + """Audio still buffered ahead of a player that started at ``start_ms``. + + The player starts at the first chunk's arrival unless ``start_ms`` is + given. Slack < 0 means the player has already underrun. + """ + if not chunks: + return 0.0 + start = chunks[0].t_ms if start_ms is None else start_ms + buffered = sum(c.audio_ms for c in chunks if c.t_ms <= now_ms) + return buffered - max(0.0, now_ms - start) + + +def should_throttle_prefills(slack_ms: float, next_chunk_eta_ms: float, safety_margin_ms: float) -> bool: + """Throttle other requests' prefills when this request's next chunk would land too late. + + ``next_chunk_eta_ms`` is the expected time until the next chunk is + delivered (steps x step_ms + hop). Throttling is worthwhile only when the + buffered audio would not cover that wait plus the safety margin. + """ + return slack_ms - next_chunk_eta_ms < safety_margin_ms + + +# --------------------------------------------------------------- KV admission +def kv_bytes_per_token( + num_layers: int, num_kv_heads: int, head_dim: int, dtype: str = "bfloat16", kv_cache_dtype: str = "auto" +) -> int: + """Bytes of KV cache one token occupies (K and V, all layers).""" + d = kv_cache_dtype if kv_cache_dtype != "auto" else dtype + b = _DTYPE_BYTES.get(d, 2) + return 2 * num_layers * num_kv_heads * head_dim * b + + +def kv_bytes_for_seqs(max_num_seqs: int, max_model_len: int, bytes_per_token: int, block_size: int = 16) -> int: + """KV bytes needed so ``max_num_seqs`` sequences of ``max_model_len`` never preempt.""" + blocks_per_seq = -(-max_model_len // block_size) + return max_num_seqs * blocks_per_seq * block_size * bytes_per_token + + +def max_seqs_without_preemption( + kv_cache_bytes: int, max_model_len: int, bytes_per_token: int, block_size: int = 16 +) -> int: + """Largest ``max_num_seqs`` the KV budget admits with no recompute preemption (>= 0).""" + per_seq = kv_bytes_for_seqs(1, max_model_len, bytes_per_token, block_size) + return int(kv_cache_bytes // per_seq) if per_seq > 0 else 0 + + +def admission_allows_no_preemption( + kv_cache_bytes: int, max_num_seqs: int, max_model_len: int, bytes_per_token: int, block_size: int = 16 +) -> bool: + return max_seqs_without_preemption(kv_cache_bytes, max_model_len, bytes_per_token, block_size) >= max_num_seqs + + +@dataclass(frozen=True) +class KVGeometry: + num_layers: int + num_kv_heads: int + head_dim: int + + @classmethod + def from_hf_config(cls, cfg: dict) -> KVGeometry | None: + """Talker geometry from a Qwen3-TTS (or plain decoder) HF config dict.""" + t = cfg.get("talker_config") if isinstance(cfg.get("talker_config"), dict) else cfg + try: + layers = int(t["num_hidden_layers"]) + kv_heads = int(t.get("num_key_value_heads") or t["num_attention_heads"]) + head_dim = int(t.get("head_dim") or (int(t["hidden_size"]) // int(t["num_attention_heads"]))) + except (KeyError, TypeError, ValueError): + return None + return cls(layers, kv_heads, head_dim) + + def bytes_per_token(self, dtype: str = "bfloat16") -> int: + return kv_bytes_per_token(self.num_layers, self.num_kv_heads, self.head_dim, dtype) diff --git a/vllm_omni/edge/step_stats.py b/vllm_omni/edge/step_stats.py new file mode 100644 index 000000000..885256a61 --- /dev/null +++ b/vllm_omni/edge/step_stats.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Steady-state per-step overhead attribution (edge scheduling, WP8). + +Process-local counters that answer "where does a codec frame's 80 ms budget go +outside the model forward?": scheduler ``schedule()`` / ``update_from_output()`` +time, the whole engine-core step, the talker -> code2wav chunk hop (enqueue -> +connector put, put -> consumer get, get duration) and the orchestrator dispatch +loop. Enabled by ``VLLM_OMNI_STEP_STATS_DIR=``; every process dumps +``/_.json`` periodically and at shutdown. ``merge_dir`` folds +the files into one attribution table (``benchmarks/tts/stream_latency_bench.py +--step-stats``). + +Zero cost when disabled: ``StepStats.get().enabled`` is False and every hook +returns immediately. +""" + +from __future__ import annotations + +import atexit +import json +import os +import threading +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +ENV_DIR = "VLLM_OMNI_STEP_STATS_DIR" +ENV_STAGE_ID = "VLLM_OMNI_STAGE_ID" +FRAME_BUDGET_MS = 80.0 # one Qwen3-TTS codec frame at 12.5 Hz +_MAX_SAMPLES = 20000 +_FLUSH_EVERY = 500 + + +class _Series: + __slots__ = ("count", "total", "max", "samples") + + def __init__(self) -> None: + self.count = 0 + self.total = 0.0 + self.max = 0.0 + self.samples: list[float] = [] + + def add(self, v: float) -> None: + self.count += 1 + self.total += v + if v > self.max: + self.max = v + if len(self.samples) < _MAX_SAMPLES: + self.samples.append(v) + else: # keep a uniform-ish reservoir: overwrite a pseudo-random slot + self.samples[(self.count * 7919) % _MAX_SAMPLES] = v + + def summary(self) -> dict[str, float | int]: + s = sorted(self.samples) + n = len(s) + + def pct(p: float) -> float: + if not n: + return 0.0 + return s[min(n - 1, int(round((n - 1) * p)))] + + return { + "count": self.count, + "sum_ms": self.total, + "mean_ms": (self.total / self.count) if self.count else 0.0, + "p50_ms": pct(0.5), + "p95_ms": pct(0.95), + "max_ms": self.max, + } + + +class StepStats: + """Process-local counter registry (thread-safe adds).""" + + _instance: StepStats | None = None + _instance_lock = threading.Lock() + + def __init__(self, out_dir: str | None, role: str | None = None) -> None: + self.out_dir = out_dir + self.enabled = bool(out_dir) + self.role = role or _default_role() + self.pid = os.getpid() + self._series: dict[str, _Series] = {} + self._lock = threading.Lock() + self._n_since_flush = 0 + self.started_at = time.time() + if self.enabled: + Path(out_dir).mkdir(parents=True, exist_ok=True) + atexit.register(self.dump) + + # ------------------------------------------------------------ singleton + @classmethod + def get(cls) -> StepStats: + inst = cls._instance + if inst is None: + with cls._instance_lock: + inst = cls._instance + if inst is None: + inst = cls(os.environ.get(ENV_DIR) or None) + cls._instance = inst + return inst + + @classmethod + def reset(cls, out_dir: str | None = None, role: str | None = None) -> StepStats: + """Replace the singleton (tests / explicit enablement in-process).""" + with cls._instance_lock: + cls._instance = cls(out_dir, role) + return cls._instance + + # ------------------------------------------------------------ recording + def add(self, name: str, value_ms: float) -> None: + if not self.enabled: + return + with self._lock: + s = self._series.get(name) + if s is None: + s = self._series[name] = _Series() + s.add(float(value_ms)) + self._n_since_flush += 1 + flush = self._n_since_flush >= _FLUSH_EVERY + if flush: + self._n_since_flush = 0 + if flush: + self.dump() + + @contextmanager + def timed(self, name: str) -> Iterator[None]: + if not self.enabled: + yield + return + t0 = time.perf_counter() + try: + yield + finally: + self.add(name, (time.perf_counter() - t0) * 1000.0) + + def wrap(self, name: str, fn: Callable[..., Any]) -> Callable[..., Any]: + """Return ``fn`` timed under ``name`` (identity when disabled).""" + if not self.enabled: + return fn + + def timed_fn(*args: Any, **kwargs: Any) -> Any: + t0 = time.perf_counter() + try: + return fn(*args, **kwargs) + finally: + self.add(name, (time.perf_counter() - t0) * 1000.0) + + timed_fn.__name__ = getattr(fn, "__name__", name) + timed_fn.__wrapped__ = fn # type: ignore[attr-defined] + return timed_fn + + # ------------------------------------------------------------ output + def to_dict(self) -> dict[str, Any]: + with self._lock: + counters = {k: v.summary() for k, v in self._series.items()} + return { + "role": self.role, + "pid": self.pid, + "started_at": self.started_at, + "dumped_at": time.time(), + "counters": counters, + } + + def dump(self) -> str | None: + if not self.enabled: + return None + path = Path(self.out_dir) / f"{self.role}_{self.pid}.json" + tmp = path.with_suffix(".json.tmp") + try: + with open(tmp, "w") as f: + json.dump(self.to_dict(), f, indent=1) + os.replace(tmp, path) + except OSError: + return None + return str(path) + + +def _default_role() -> str: + sid = os.environ.get(ENV_STAGE_ID) + return f"stage{sid}" if sid not in (None, "") else "orchestrator" + + +# ---------------------------------------------------------------- hooks +def instrument_engine_core(engine_core: Any) -> None: + """Time the scheduler and the whole step of a (stage) EngineCore instance.""" + stats = StepStats.get() + if not stats.enabled: + return + sched = getattr(engine_core, "scheduler", None) + if sched is not None: + for name in ("schedule", "update_from_output"): + fn = getattr(sched, name, None) + if callable(fn) and not hasattr(fn, "__wrapped__"): + setattr(sched, name, stats.wrap(f"sched.{name}_ms", fn)) + step_fn = getattr(engine_core, "step_fn", None) + if callable(step_fn) and not hasattr(step_fn, "__wrapped__"): + engine_core.step_fn = stats.wrap("core.step_ms", step_fn) + + +def shm_lockfile_age_ms(key: str, now: float | None = None) -> float | None: + """Age of the SharedMemoryConnector lock file for ``key`` (written at put time).""" + try: + st = os.stat(f"/dev/shm/shm_{key}_lockfile.lock") + except OSError: + return None + return ((now if now is not None else time.time()) - st.st_mtime) * 1000.0 + + +# ---------------------------------------------------------------- merge +def merge_dir(out_dir: str | Path) -> dict[str, Any]: + """Merge per-process dumps into ``{role: {counter: summary}}`` plus a share table.""" + roles: dict[str, dict[str, Any]] = {} + for f in sorted(Path(out_dir).glob("*.json")): + try: + d = json.loads(f.read_text()) + except (OSError, ValueError): + continue + role = d.get("role", f.stem) + bucket = roles.setdefault(role, {}) + for name, summ in d.get("counters", {}).items(): + cur = bucket.get(name) + if cur is None: + bucket[name] = dict(summ) + else: # several processes with the same role: combine count/sum/max, keep worst p95 + cnt = cur["count"] + summ["count"] + cur.update( + count=cnt, + sum_ms=cur["sum_ms"] + summ["sum_ms"], + mean_ms=((cur["sum_ms"] + summ["sum_ms"]) / cnt) if cnt else 0.0, + p50_ms=max(cur["p50_ms"], summ["p50_ms"]), + p95_ms=max(cur["p95_ms"], summ["p95_ms"]), + max_ms=max(cur["max_ms"], summ["max_ms"]), + ) + return {"roles": roles, "attribution": attribution(roles)} + + +def attribution(roles: dict[str, dict[str, Any]], frame_budget_ms: float = FRAME_BUDGET_MS) -> list[dict[str, Any]]: + """Rows of (role, counter, mean, p95, share of the frame budget) sorted by share.""" + rows = [] + for role, counters in roles.items(): + for name, s in counters.items(): + rows.append( + { + "role": role, + "counter": name, + "count": s["count"], + "mean_ms": s["mean_ms"], + "p95_ms": s["p95_ms"], + "share_of_frame": (s["mean_ms"] / frame_budget_ms) if frame_budget_ms else 0.0, + } + ) + rows.sort(key=lambda r: -r["share_of_frame"]) + return rows + + +def format_table(rows: list[dict[str, Any]]) -> str: + lines = ["| role | counter | n | mean ms | p95 ms | share of 80 ms |", "|---|---|---|---|---|---|"] + for r in rows: + lines.append( + f"| {r['role']} | {r['counter']} | {r['count']} | {r['mean_ms']:.2f} | {r['p95_ms']:.2f} " + f"| {100 * r['share_of_frame']:.1f}% |" + ) + return "\n".join(lines) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 6a8633970..c85b96505 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -204,6 +204,9 @@ class AsyncOmniEngine: from vllm_omni.entrypoints.utils import apply_deploy_profile apply_deploy_profile(model, kwargs) + # A profile may carry orchestrator defaults (e.g. ``parallel_stage_init``); + # re-read the flag now that the profile has been applied. + self._parallel_stage_init = bool(kwargs.get("parallel_stage_init") or False) # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand. deploy_config_path = kwargs.get("deploy_config") diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index e86dc7673..f7f7b6f32 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -36,6 +36,7 @@ from vllm.v1.metrics.stats import IterationStats from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks +from vllm_omni.edge.step_stats import StepStats from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.duplex.contracts import ( @@ -1383,7 +1384,12 @@ class Orchestrator: pending_get: asyncio.Task | None = None next_reconcile = _time.monotonic() + _ORCH_READER_RECONCILE_INTERVAL_S try: + _stats = StepStats.get() + _dispatch_t0: float | None = None while not self._shutdown_event.is_set(): + if _dispatch_t0 is not None: + _stats.add("orch.dispatch_ms", (_time.perf_counter() - _dispatch_t0) * 1000.0) + _dispatch_t0 = None if pending_get is None: pending_get = asyncio.create_task(ready_q.get(), name="orch-dispatch-get") done, _ = await asyncio.wait( @@ -1406,6 +1412,8 @@ class Orchestrator: continue kind, stage_id, replica_id, payload = pending_get.result() pending_get = None + if _stats.enabled: + _dispatch_t0 = _time.perf_counter() if kind == "error": # replica_id < 0 means the failure was raised by a poller diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index d8544986b..a32d7b383 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -191,6 +191,10 @@ class StageEngineCoreProc(EngineCoreProc): engine_index=dp_rank, **kwargs, ) + # Steady-state overhead attribution (no-op unless VLLM_OMNI_STEP_STATS_DIR). + from vllm_omni.edge.step_stats import StepStats, instrument_engine_core + + instrument_engine_core(engine_core) # Each subprocess corresponds to exactly one omni replica with # its own OmniMasterServer allocation, so the heartbeat client @@ -241,6 +245,8 @@ class StageEngineCoreProc(EngineCoreProc): engine_core._send_engine_dead() raise finally: + with contextlib.suppress(Exception): + StepStats.get().dump() signal.signal(signal.SIGTERM, signal.SIG_DFL) signal.signal(signal.SIGINT, signal.SIG_DFL) if signal_callback is not None: