Summarize vLLM scheduler step profiles

runs/frontier-qwen30-vllm020-profile-v1/summarize_opprof.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+"""Summarize prompt-free vLLM scheduler/operator-profiler step records."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import statistics
+from collections import Counter, defaultdict
+from pathlib import Path
+from typing import Any
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cell-root", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    return parser.parse_args()
+
+
+def percentile(values: list[float], fraction: float) -> float:
+    ordered = sorted(values)
+    index = min(len(ordered) - 1, max(0, math.ceil(fraction * len(ordered)) - 1))
+    return ordered[index]
+
+
+def stats(values: list[float]) -> dict[str, float]:
+    return {
+        "min": min(values),
+        "mean": statistics.fmean(values),
+        "p50": percentile(values, 0.50),
+        "p95": percentile(values, 0.95),
+        "p99": percentile(values, 0.99),
+        "max": max(values),
+    }
+
+
+def summarize_cell(path: Path) -> dict[str, Any]:
+    streams = sorted(path.glob("opprof/*.jsonl"))
+    if len(streams) != 1:
+        raise ValueError(f"expected one opprof stream in {path}, got {len(streams)}")
+    groups: dict[tuple[str, str], dict[str, list[float]]] = defaultdict(
+        lambda: defaultdict(list)
+    )
+    phase_counts: Counter[str] = Counter()
+    graph_counts: Counter[str] = Counter()
+    prefix_queries = 0
+    prefix_hits = 0
+    total_records = 0
+    dropped_records_max = 0
+    with streams[0].open() as handle:
+        for line in handle:
+            record = json.loads(line)
+            if not record.get("model_executed", False):
+                continue
+            prefill_tokens = int(record.get("prefill_tokens", 0))
+            decode_tokens = int(record.get("decode_tokens", 0))
+            if prefill_tokens and decode_tokens:
+                phase = "true_mixed"
+            elif prefill_tokens:
+                phase = "pure_prefill"
+            elif decode_tokens:
+                phase = "pure_decode"
+            else:
+                phase = "empty"
+            graph = str((record.get("cudagraph") or {}).get("runtime_mode", "UNKNOWN"))
+            duration_ms = (
+                int(record["complete_mono_ns"]) - int(record["submit_mono_ns"])
+            ) / 1e6
+            if duration_ms < 0:
+                raise ValueError(f"negative duration in {streams[0]}")
+            group = groups[(phase, graph)]
+            group["submit_to_complete_ms"].append(duration_ms)
+            group["scheduled_requests"].append(int(record["scheduled_requests"]))
+            group["prefill_tokens"].append(prefill_tokens)
+            group["decode_tokens"].append(decode_tokens)
+            group["total_tokens"].append(prefill_tokens + decode_tokens)
+            group["queue_waiting"].append(int(record["queues"]["waiting"]))
+            group["kv_usage"].append(float(record["kv"]["usage"]))
+            local_prefix = (record.get("prefix") or {}).get("local") or {}
+            prefix_queries += int(local_prefix.get("queries", 0))
+            prefix_hits += int(local_prefix.get("hits", 0))
+            dropped_records_max = max(
+                dropped_records_max, int(record.get("dropped_records_before", 0))
+            )
+            phase_counts[phase] += 1
+            graph_counts[graph] += 1
+            total_records += 1
+    if total_records == 0:
+        raise ValueError(f"empty opprof stream: {streams[0]}")
+    summarized_groups = []
+    for (phase, graph), metrics in sorted(groups.items()):
+        summarized_groups.append(
+            {
+                "phase": phase,
+                "cudagraph_runtime_mode": graph,
+                "steps": len(metrics["submit_to_complete_ms"]),
+                **{name: stats(values) for name, values in metrics.items()},
+            }
+        )
+    cell = path.name
+    tp_text, mns_text = cell.split("_")
+    return {
+        "cell": cell,
+        "tensor_parallel_size": int(tp_text.removeprefix("tp")),
+        "max_num_seqs": int(mns_text.removeprefix("mns")),
+        "stream": str(streams[0]),
+        "records": total_records,
+        "phase_counts": dict(phase_counts),
+        "cudagraph_runtime_mode_counts": dict(graph_counts),
+        "prefix": {
+            "queries": prefix_queries,
+            "hits": prefix_hits,
+            "hit_rate": prefix_hits / prefix_queries if prefix_queries else 0.0,
+        },
+        "dropped_records_before_max": dropped_records_max,
+        "groups": summarized_groups,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    cells = [summarize_cell(path) for path in sorted(args.cell_root.glob("tp*_mns*"))]
+    if not cells:
+        raise SystemExit(f"no cells found in {args.cell_root}")
+    payload = {
+        "schema_version": "qwen30_vllm020_opprof_summary.v1",
+        "measurement_semantics": (
+            "complete_mono_ns - submit_mono_ns from the existing vLLM opprof record; "
+            "this includes the instrumented submit-to-completion interval and is not "
+            "claimed to be a single-kernel CUDA-event duration"
+        ),
+        "contains_prompt_text": False,
+        "cell_root": str(args.cell_root),
+        "cells": cells,
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+    print(json.dumps({"cells": len(cells), "records": sum(c["records"] for c in cells)}))
+
+
+if __name__ == "__main__":
+    main()