Add reproducible CollectiveSpec opportunity screen

scripts/collectivespec/summarize_engine_metrics.py

@@ -0,0 +1,135 @@
+#!/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())