analysis: summarize workload simulator regimes

runs/frontier-workload-regime-taxonomy-v0/analyze_simulator_regimes.py

@@ -0,0 +1,274 @@
+#!/usr/bin/env python3
+"""Summarize workload-regime Frontier metrics and closed-loop batch state."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+
+METRICS = (
+    "ttft_mean_ms",
+    "ttft_p90_ms",
+    "tpot_mean_ms",
+    "tpot_p90_ms",
+    "e2e_mean_ms",
+    "e2e_p90_ms",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--sim-root", type=Path, required=True)
+    parser.add_argument("--workload-manifest", type=Path, required=True)
+    parser.add_argument("--json-output", type=Path, required=True)
+    parser.add_argument("--markdown-output", type=Path, required=True)
+    return parser.parse_args()
+
+
+def summarize_ledger(path: Path) -> dict[str, float | int | None]:
+    total_time = 0.0
+    weighted_batch = 0.0
+    decode_time = 0.0
+    decode_weighted_batch = 0.0
+    prefill_time = 0.0
+    prefill_weighted_batch = 0.0
+    max_batch = 0
+    batches = 0
+    moe_time = 0.0
+    collective_time = 0.0
+    for line in path.open():
+        if not line.strip():
+            continue
+        row = json.loads(line)
+        request_tokens = [int(value) for value in row["request_num_tokens"]]
+        batch = len(request_tokens)
+        execution = row["execution_time"]
+        duration = float(execution["total_time_ms"])
+        components = execution["component_ledger_ms"]
+        total_time += duration
+        weighted_batch += batch * duration
+        max_batch = max(max_batch, batch)
+        batches += 1
+        if request_tokens and all(value == 1 for value in request_tokens):
+            decode_time += duration
+            decode_weighted_batch += batch * duration
+        else:
+            prefill_time += duration
+            prefill_weighted_batch += batch * duration
+        moe_time += sum(
+            float(components.get(name, 0.0))
+            for name in (
+                "moe_gating_linear_time",
+                "moe_gating_routing_topk_time",
+                "moe_grouped_gemm_time",
+                "moe_shuffling_time",
+            )
+        )
+        collective_time += sum(
+            float(components.get(name, 0.0))
+            for name in (
+                "attention_all_reduce_time",
+                "mlp_all_reduce_time",
+                "moe_tensor_parallel_allgather_time",
+                "expert_parallel_communication_time",
+            )
+        )
+    return {
+        "batches": batches,
+        "max_batch_size": max_batch,
+        "time_weighted_batch_size": weighted_batch / total_time if total_time else None,
+        "decode_time_weighted_batch_size": (
+            decode_weighted_batch / decode_time if decode_time else None
+        ),
+        "prefill_mixed_time_weighted_batch_size": (
+            prefill_weighted_batch / prefill_time if prefill_time else None
+        ),
+        "decode_time_fraction": decode_time / total_time if total_time else None,
+        "moe_time_fraction": moe_time / total_time if total_time else None,
+        "collective_time_fraction": (
+            collective_time / total_time if total_time else None
+        ),
+    }
+
+
+def find_one(root: Path, pattern: str) -> Path:
+    matches = list(root.glob(pattern))
+    if len(matches) != 1:
+        raise ValueError(f"expected one {pattern} under {root}, got {matches}")
+    return matches[0]
+
+
+def load_rows(sim_root: Path, workload_manifest: Path) -> list[dict[str, Any]]:
+    suite = json.loads(workload_manifest.read_text())
+    cases = {case["public_csv_sha256"]: case for case in suite["cases"]}
+    rows = []
+    for result_path in sorted(sim_root.glob("tp*-prefix-*/runs/*/*/result.json")):
+        result = json.loads(result_path.read_text())
+        case = cases.get(result.get("trace_sha256"))
+        if case is None:
+            raise ValueError(f"unknown trace hash in {result_path}")
+        row: dict[str, Any] = {
+            "family": case["family"],
+            "rho": case["rho"],
+            "prefix_caching": case["prefix_caching"],
+            "tp": result["config"]["tp"],
+            "mns": result["config"]["mns"],
+            "config": result["config"]["name"],
+            "status": result["status"],
+            "failure_kind": result.get("failure_kind"),
+            "result_path": str(result_path),
+        }
+        if result["status"] == "completed":
+            row.update({name: result["score"].get(name) for name in METRICS})
+            ledger = find_one(
+                result_path.parent,
+                "metrics/**/frontier_stage_batch_ledger.jsonl",
+            )
+            row["ledger_path"] = str(ledger)
+            row.update(summarize_ledger(ledger))
+        rows.append(row)
+    return rows
+
+
+def best_configs(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    groups: dict[tuple[str, float], list[dict[str, Any]]] = defaultdict(list)
+    for row in rows:
+        if row["status"] == "completed":
+            groups[(row["family"], row["rho"])].append(row)
+    best = []
+    for (family, rho), records in sorted(groups.items()):
+        output: dict[str, Any] = {"family": family, "rho": rho}
+        for metric in METRICS:
+            valid = [row for row in records if row.get(metric) is not None]
+            if valid:
+                winner = min(valid, key=lambda row: (row[metric], row["config"]))
+                output[metric] = {
+                    "config": winner["config"],
+                    "value_ms": winner[metric],
+                }
+        best.append(output)
+    return best
+
+
+def state_knees(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
+    for row in rows:
+        value = row.get("decode_time_weighted_batch_size")
+        if row["status"] == "completed" and value is not None:
+            groups[(row["family"], row["config"])].append(row)
+    knees = []
+    for (family, config), records in sorted(groups.items()):
+        records.sort(key=lambda row: row["rho"])
+        candidates = []
+        for left, right in zip(records, records[1:]):
+            delta = (
+                right["decode_time_weighted_batch_size"]
+                - left["decode_time_weighted_batch_size"]
+            )
+            candidates.append((abs(delta), delta, left, right))
+        if candidates:
+            _, delta, left, right = max(candidates, key=lambda item: item[0])
+            knees.append(
+                {
+                    "family": family,
+                    "config": config,
+                    "rho_interval": [left["rho"], right["rho"]],
+                    "decode_batch_before": left[
+                        "decode_time_weighted_batch_size"
+                    ],
+                    "decode_batch_after": right[
+                        "decode_time_weighted_batch_size"
+                    ],
+                    "decode_batch_delta": delta,
+                }
+            )
+    return knees
+
+
+def render_markdown(summary: dict[str, Any]) -> str:
+    lines = [
+        "# Frontier simulator workload-regime discovery",
+        "",
+        f"- Expected cells: {summary['coverage']['expected']}",
+        f"- Observed cells: {summary['coverage']['observed']}",
+        f"- Completed: {summary['coverage']['completed']}",
+        f"- Failed: {summary['coverage']['failed']}",
+        "",
+        "This table reports simulator predictions only. It does not classify a workload as reliable without matched real-system evidence.",
+        "",
+        "## Simulator winners",
+        "",
+        "| Family | rho | TTFT mean | TTFT p90 | TPOT mean | TPOT p90 |",
+        "|---|---:|---|---|---|---|",
+    ]
+    for row in summary["best_configs"]:
+        values = []
+        for metric in (
+            "ttft_mean_ms",
+            "ttft_p90_ms",
+            "tpot_mean_ms",
+            "tpot_p90_ms",
+        ):
+            winner = row.get(metric)
+            values.append(
+                "—"
+                if winner is None
+                else f"{winner['config']} ({winner['value_ms']:.2f} ms)"
+            )
+        lines.append(
+            f"| {row['family']} | {row['rho']:.2f} | " + " | ".join(values) + " |"
+        )
+    lines.extend(
+        [
+            "",
+            "## Largest decode-batch transition per config",
+            "",
+            "| Family | Config | rho interval | Batch before | Batch after | Delta |",
+            "|---|---|---|---:|---:|---:|",
+        ]
+    )
+    for row in summary["state_knees"]:
+        lines.append(
+            "| {family} | {config} | {left:.2f}–{right:.2f} | {before:.2f} | {after:.2f} | {delta:+.2f} |".format(
+                family=row["family"],
+                config=row["config"],
+                left=row["rho_interval"][0],
+                right=row["rho_interval"][1],
+                before=row["decode_batch_before"],
+                after=row["decode_batch_after"],
+                delta=row["decode_batch_delta"],
+            )
+        )
+    return "\n".join(lines) + "\n"
+
+
+def main() -> None:
+    args = parse_args()
+    rows = load_rows(args.sim_root, args.workload_manifest)
+    completed = sum(row["status"] == "completed" for row in rows)
+    summary = {
+        "schema": "frontier-workload-regime-simulator-analysis-v1",
+        "coverage": {
+            "expected": 7 * 5 * 12,
+            "observed": len(rows),
+            "completed": completed,
+            "failed": len(rows) - completed,
+        },
+        "rows": rows,
+        "best_configs": best_configs(rows),
+        "state_knees": state_knees(rows),
+    }
+    for output in (args.json_output, args.markdown_output):
+        output.parent.mkdir(parents=True, exist_ok=True)
+    args.json_output.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n")
+    args.markdown_output.write_text(render_markdown(summary))
+    print(args.markdown_output)
+
+
+if __name__ == "__main__":
+    main()