tzhouamandClaude Fable 5.1
073d31b9cd
* vllm_omni/engine/init_timeline.py: InitTimeline phase collector; engine phases (stages_init, g3_device_lock, g2_spawn/ready) and worker phases (proc_entry, engine_core_init, init_device, load_model/load_weights, talker_mtp_init, attn_prewarm, cudagraph_capture, talker_mtp_capture, code2wav_cudagraph, predictor_setup_compile/warmup/cudagraph) recorded via worker_phase/timed_phase; workers append JSONL under VLLM_OMNI_INIT_TIMELINE_DIR (inherited at spawn) and the engine merges them. Enabled by VLLM_OMNI_INIT_TIMELINE=<json> or log_stats; table logged at INFO, partial table on orchestrator startup timeout, suggested timeouts logged. * omni_base.py: named deploy profiles default to stage_init_timeout=900 / init_timeout=1500 unless given explicitly. * gpu_model_runner._init_talker_mtp tolerates max_cudagraph_capture_size=None; gpu_ar_model_runner.shutdown skips CUDA graph teardown on CPU. * benchmarks/tts/init_profile.py: runs stream_latency_bench --init-only per configuration (default, cold_cache, eager_stage1, eager_both, parallel_stage_init, edge, ...) and tabulates phase durations. * stream_latency_bench: --parallel-stage-init / --deploy-profile flags; the resource poller no longer shadows Thread._stop and skips GPU sampling on CPU. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
111 lines
4.1 KiB
Python
111 lines
4.1 KiB
Python
"""Tests for vllm_omni/engine/init_timeline.py (CPU, no engine)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.engine import init_timeline as it
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_timeline(monkeypatch):
|
|
it.InitTimeline.reset()
|
|
monkeypatch.delenv(it.ENV_DIR, raising=False)
|
|
monkeypatch.delenv(it.ENV_STAGE_ID, raising=False)
|
|
monkeypatch.delenv(it.ENV_REPLICA_ID, raising=False)
|
|
yield
|
|
it.InitTimeline.reset()
|
|
|
|
|
|
def test_disabled_timeline_records_nothing():
|
|
tl = it.InitTimeline.start(enabled=False)
|
|
with tl.phase("x"):
|
|
pass
|
|
tl.mark("m")
|
|
assert tl.phases == []
|
|
assert it.ENV_DIR not in os.environ
|
|
with it.worker_phase("w"):
|
|
pass # no-op without env / enabled timeline
|
|
|
|
|
|
def test_phase_nesting_and_durations(tmp_path):
|
|
tl = it.InitTimeline.start(enabled=True, run_dir=tmp_path / "run")
|
|
with tl.phase("outer", stage_id=0) as info:
|
|
time.sleep(0.01)
|
|
with tl.phase("inner", stage_id=0, n=3):
|
|
time.sleep(0.01)
|
|
info["k"] = "v"
|
|
names = [p.name for p in tl.phases]
|
|
assert names == ["inner", "outer"] # inner closes first
|
|
outer = tl.phases[1]
|
|
assert outer.dur_s >= 0.02 and outer.extra == {"k": "v"}
|
|
assert tl.phases[0].extra == {"n": 3}
|
|
assert os.environ[it.ENV_DIR] == str(tmp_path / "run")
|
|
|
|
|
|
def test_worker_helpers_write_jsonl_and_ingest(tmp_path, monkeypatch):
|
|
run_dir = tmp_path / "run"
|
|
tl = it.InitTimeline.start(enabled=True, run_dir=run_dir)
|
|
# Simulate a worker: no in-process timeline, env set by the parent.
|
|
it.InitTimeline.reset()
|
|
monkeypatch.setenv(it.ENV_STAGE_ID, "1")
|
|
monkeypatch.setenv(it.ENV_REPLICA_ID, "2")
|
|
with it.worker_phase("load_weights", gib=1.5):
|
|
time.sleep(0.005)
|
|
it.worker_mark("proc_entry")
|
|
files = list(run_dir.glob("worker_*.jsonl"))
|
|
assert len(files) == 1
|
|
rows = [json.loads(line) for line in files[0].read_text().splitlines()]
|
|
assert {r["name"] for r in rows} == {"load_weights", "proc_entry"}
|
|
assert all(r["stage_id"] == 1 and r["replica_id"] == 2 and r["role"] == "worker" for r in rows)
|
|
# Back in the engine process: ingest (twice -> idempotent).
|
|
it.InitTimeline._current = tl
|
|
assert tl.ingest_workers() == 2
|
|
assert tl.ingest_workers() == 0
|
|
assert [p.name for p in tl.phases] == sorted(
|
|
[p.name for p in tl.phases], key=lambda n: [p.start for p in tl.phases if p.name == n][0]
|
|
)
|
|
|
|
|
|
def test_table_dump_and_suggested_timeouts(tmp_path):
|
|
tl = it.InitTimeline(enabled=True, run_dir=None, t0=1000.0)
|
|
tl.record(it.InitPhase("g2_spawn", 1000.0, 1010.0, stage_id=0, role="engine", pid=1))
|
|
tl.record(it.InitPhase("load_weights", 1002.0, 1200.0, stage_id=0, role="worker", pid=2))
|
|
tl.record(it.InitPhase("load_weights", 1210.0, 1260.0, stage_id=1, role="worker", pid=3))
|
|
tl.t_end = 1500.0
|
|
assert tl.total_s == pytest.approx(500.0)
|
|
table = tl.format_table()
|
|
assert "load_weights" in table and "g2_spawn" in table and "total=500.00s" in table
|
|
d = tl.to_dict()
|
|
assert d["schema"] == 1 and len(d["phases"]) == 3
|
|
out = tmp_path / "tl.json"
|
|
tl.dump(out)
|
|
assert json.loads(out.read_text())["total_s"] == pytest.approx(500.0)
|
|
# stage 0 span = 1000..1200 = 200 s -> 300 (== default floor); total 500 s -> 750
|
|
sug = tl.suggest_timeouts()
|
|
assert sug == {"stage_init_timeout": 300, "init_timeout": 750}
|
|
# short runs floor at the defaults
|
|
short = it.InitTimeline(enabled=True, t0=0.0)
|
|
short.record(it.InitPhase("x", 0.0, 1.0, stage_id=0))
|
|
short.t_end = 2.0
|
|
assert short.suggest_timeouts() == {"stage_init_timeout": 300, "init_timeout": 600}
|
|
|
|
|
|
def test_timed_phase_decorator_records_in_engine_process(tmp_path):
|
|
tl = it.InitTimeline.start(enabled=True, run_dir=tmp_path / "run")
|
|
|
|
class Runner:
|
|
@it.timed_phase("capture")
|
|
def capture(self, x):
|
|
return x * 2
|
|
|
|
assert Runner().capture(21) == 42
|
|
assert [p.name for p in tl.phases] == ["capture"]
|
|
assert tl.phases[0].role == "engine"
|