[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 <noreply@anthropic.com>
This commit is contained in:
tzhouamandClaude Fable 5.1 committed
1 parent
74acd66307
commit
37c0ea9a6b
16 files changed
+926
-16
No files matched your search
@@ -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())
|
||||
@@ -4,6 +4,10 @@ Committed JSON produced by the edge branch benchmarks (kept small; WAVs are neve
|
||||
|
||||
- `stream_latency/<model>/<deploy>/<timestamp>[_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/<timestamp>/summary.json` — `init_profile.py` comparison of initialization phases per configuration.
|
||||
- `hw_emulation/`, `scheduling/` — per-work-package comparison tables.
|
||||
- `hw_emulation/<timestamp>/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/<timestamp>/summary.json` — `cpu_thread_sweep.py` (CPU-only 8/16/32 threads + inductor variant, with per-step attribution).
|
||||
- `stream_latency/step_stats/<timestamp>/*.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`.
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+11
-6
@@ -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")
|
||||
|
||||
@@ -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 ``<model_dir>/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:
|
||||
|
||||
@@ -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)
|
||||
@@ -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=<dir>``; every process dumps
|
||||
``<dir>/<role>_<pid>.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)
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in new issue
Block a user