Files
aituner/scripts/collectivespec/summarize_engine_metrics.py

136 lines
4.8 KiB
Python

#!/usr/bin/env python3
"""Extract per-DP-window speculative metrics from a vLLM engine log.
vLLM emits an ``Engine NNN`` throughput line followed by the corresponding
``SpecDecoding metrics`` line. This script keeps that adjacency explicit and
does not infer request-level outcomes from the aggregate metrics.
"""
from __future__ import annotations
import argparse
import json
import math
import re
import statistics
from collections import defaultdict
from pathlib import Path
from typing import Any
ENGINE = re.compile(
r"INFO (?P<clock>\d\d-\d\d \d\d:\d\d:\d\d\.\d+) .*?"
r"Engine (?P<engine>\d+): .*?"
r"Avg decode step duration: (?P<decode_ms>[0-9.]+) ms .*?"
r"Running: (?P<running>\d+) reqs, Waiting: (?P<waiting>\d+) reqs"
)
SPEC = re.compile(
r"INFO (?P<clock>\d\d-\d\d \d\d:\d\d:\d\d\.\d+) .*?"
r"SpecDecoding metrics: Mean acceptance length: (?P<mean_accept>[0-9.]+), .*?"
r"Avg Draft acceptance rate: (?P<accept_pct>[0-9.]+)%"
)
def percentile(values: list[float], q: float) -> float | None:
if not values:
return None
values = sorted(values)
index = (len(values) - 1) * q
lower, upper = math.floor(index), math.ceil(index)
if lower == upper:
return values[lower]
return values[lower] + (values[upper] - values[lower]) * (index - lower)
def describe(values: list[float]) -> dict[str, Any]:
return {
"n": len(values),
"mean": statistics.fmean(values) if values else None,
"min": min(values) if values else None,
"p50": percentile(values, 0.5),
"p95": percentile(values, 0.95),
"max": max(values) if values else None,
"distinct_value_count": len(set(values)),
}
def extract(path: Path) -> dict[str, Any]:
pending: tuple[str, dict[str, Any]] | None = None
rows: list[dict[str, Any]] = []
for line in path.read_text(errors="replace").splitlines():
engine = ENGINE.search(line)
if engine:
pending = (
engine.group("clock")[:17],
{
"clock": engine.group("clock"),
"engine": int(engine.group("engine")),
"decode_ms": float(engine.group("decode_ms")),
"running": int(engine.group("running")),
"waiting": int(engine.group("waiting")),
},
)
continue
spec = SPEC.search(line)
if not spec or pending is None:
continue
clock, row = pending
if spec.group("clock")[:17] != clock:
pending = None
continue
row["mean_accept_length"] = float(spec.group("mean_accept"))
row["accept_pct"] = float(spec.group("accept_pct"))
rows.append(row)
pending = None
by_engine: dict[int, list[dict[str, Any]]] = defaultdict(list)
for row in rows:
by_engine[row["engine"]].append(row)
summaries: list[dict[str, Any]] = []
for engine, items in sorted(by_engine.items()):
summaries.append(
{
"engine": engine,
"window_count": len(items),
"accept_pct": describe([item["accept_pct"] for item in items]),
"mean_accept_length": describe([item["mean_accept_length"] for item in items]),
"decode_ms": describe([item["decode_ms"] for item in items]),
"running_requests": describe([float(item["running"]) for item in items]),
}
)
acceptance = [row["accept_pct"] for row in rows]
return {
"source": str(path),
"records": rows,
"per_engine": summaries,
"data_sanity": {
"accept_pct": {
"n": len(acceptance),
"min": min(acceptance) if acceptance else None,
"max": max(acceptance) if acceptance else None,
"distinct_value_count": len(set(acceptance)),
"within_0_100": all(0 <= item <= 100 for item in acceptance),
},
"invariants": {
"non_negative_decode_ms": all(row["decode_ms"] >= 0 for row in rows),
"non_negative_running": all(row["running"] >= 0 for row in rows),
"all_records_have_engine": all("engine" in row for row in rows),
},
},
}
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--engine-log", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
result = extract(args.engine_log)
args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
print(json.dumps({"per_engine": result["per_engine"], "data_sanity": result["data_sanity"]}, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())