sskj/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/compare.py

#!/usr/bin/env python3
"""Cross TP×DP configuration comparison for dsv4_h200_sglang_tp_dp_matrix.

Usage:
    python3 compare.py --run-root results/<run_id> [--output comparison.md]
"""
import argparse
import json
import re
from collections import defaultdict
from pathlib import Path


def load_result(result_root: Path) -> dict:
    path = result_root / "results.json"
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def slo_status(ttft_p95_ms: float, tpot_mean_ms: float,
               ttft_limit_ms: float = 3000.0, tpot_limit_ms: float = 50.0) -> str:
    ttft_ok = ttft_p95_ms < ttft_limit_ms
    tpot_ok = tpot_mean_ms < tpot_limit_ms
    if ttft_ok and tpot_ok:
        return "PASS"
    if ttft_ok or tpot_ok:
        return "PARTIAL"
    return "FAIL"


def gpu_memory_str(gpu: dict | None) -> str:
    if not gpu:
        return "-"
    peak = gpu.get("peak_used_mb", 0)
    total = gpu.get("memory_total_mb", 0)
    if total:
        return f"{peak:.0f}/{total:.0f} ({100*peak/total:.1f}%)"
    return f"{peak:.0f}"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-root", type=Path, required=True)
    parser.add_argument("-o", "--output", type=Path, default=Path("comparison.md"))
    parser.add_argument("--ttft-limit", type=float, default=3000.0)
    parser.add_argument("--tpot-limit", type=float, default=50.0)
    args = parser.parse_args()

    # Discover configurations: tp*_dp* directories.
    configs = []
    for subdir in sorted(args.run_root.iterdir()):
        if not subdir.is_dir():
            continue
        name = subdir.name
        if not (name.startswith("tp") and "_dp" in name):
            continue
        results_json = subdir / "results.json"
        if not results_json.exists():
            continue
        configs.append((name, load_result(subdir)))

    if not configs:
        print(f"No tp*_dp* results found under {args.run_root}")
        return

    model = configs[0][1].get("metadata", {}).get("model", "unknown")
    hardware = configs[0][1].get("metadata", {}).get("hardware", "unknown")

    # Group by scenario name.
    by_scenario: dict[str, dict[str, dict]] = defaultdict(dict)
    skipped: dict[str, dict[str, str]] = defaultdict(dict)
    for label, data in configs:
        for s in data.get("scenarios", []):
            key = s["name"]
            if s.get("status") == "skipped_oom":
                skipped[key][label] = s.get("note", "skipped")
            else:
                by_scenario[key][label] = s

    with open(args.output, "w", encoding="utf-8") as f:
        f.write(f"# SGLang TP×DP matrix comparison ({hardware})\n\n")
        f.write("## Summary\n\n")
        f.write(f"- Model: `{model}`\n")
        f.write(f"- Hardware: {hardware}\n")
        f.write("- Backend: SGLang (Docker)\n")
        f.write("- Benchmark client: `sglang.benchmark.serving`\n")
        f.write(f"- SLO reference: TTFT P95 < {args.ttft_limit}ms, TPOT mean < {args.tpot_limit}ms\n\n")

        # Configuration overview.
        f.write("### Configurations\n\n")
        f.write("| Config | TP | DP | GPUs/replica | Notes |\n")
        f.write("|---|---:|---:|---:|---|\n")
        for label, data in configs:
            cfg = data.get("config", {})
            tp = cfg.get("tp", "?")
            dp = cfg.get("dp", "?")
            f.write(f"| {label} | {tp} | {dp} | {tp} | server args recorded per ISL in results.json |\n")
        f.write("\n")

        # Side-by-side table.
        f.write("## Side-by-side results\n\n")
        headers = [
            "Scenario", "ISL", "DSL", "Config", "Conc", "Req/s", "OutTok/s",
            "TTFT P95(ms)", "TTFT P99(ms)", "TPOT Mean(ms)", "TPOT P95(ms)",
            "TPOT P99(ms)", "E2E P99(ms)", "Peak GPU mem", "SLO"
        ]
        f.write("| " + " | ".join(headers) + " |\n")
        f.write("|" + "|".join(["---"] * len(headers)) + "|\n")

        for scenario_name in sorted(by_scenario.keys(), key=lambda x: tuple(map(int, re.findall(r"\d+", x)))):
            _, isl, dsl = re.findall(r"\d+", scenario_name)
            for label, data in configs:
                s = by_scenario[scenario_name].get(label)
                if s is None:
                    if scenario_name in skipped and label in skipped[scenario_name]:
                        note = skipped[scenario_name][label]
                        f.write(f"| {scenario_name} | {isl} | {dsl} | {label} | - | - | - | - | - | - | - | - | - | - | {note} |\n")
                    continue
                cfg = s["config"]
                m = s["metrics"]
                status = slo_status(m["ttft_ms"]["p95"], m["tpot_ms"]["mean"], args.ttft_limit, args.tpot_limit)
                gpu = m.get("gpu_memory")
                f.write(
                    f"| {scenario_name} | {isl} | {dsl} | {label} | {cfg['concurrency']} | "
                    f"{m['request_throughput']:.2f} | {m['output_token_throughput']:.2f} | "
                    f"{m['ttft_ms']['p95']:.2f} | {m['ttft_ms']['p99']:.2f} | "
                    f"{m['tpot_ms']['mean']:.2f} | {m['tpot_ms']['p95']:.2f} | {m['tpot_ms']['p99']:.2f} | "
                    f"{m['e2e_ms']['p99']:.2f} | {gpu_memory_str(gpu)} | {status} |\n"
                )

        # Best throughput per ISL/DSL.
        f.write("\n## Best throughput per (ISL, DSL)\n\n")
        f.write("| ISL | DSL | Best Config | Concurrency | OutTok/s | TTFT P95(ms) | TPOT Mean(ms) | SLO |\n")
        f.write("|---:|---:|---|---:|---:|---:|---:|---:|\n")
        best_by_shape: dict[tuple[int, int], tuple[float, str, dict]] = {}
        for scenario_name, backends in by_scenario.items():
            _, isl, dsl = re.findall(r"\d+", scenario_name)
            isl_i, dsl_i = int(isl), int(dsl)
            for label, s in backends.items():
                m = s["metrics"]
                out_tok = m["output_token_throughput"]
                if (isl_i, dsl_i) not in best_by_shape or out_tok > best_by_shape[(isl_i, dsl_i)][0]:
                    best_by_shape[(isl_i, dsl_i)] = (out_tok, label, s)
        for (isl_i, dsl_i), (out_tok, label, s) in sorted(best_by_shape.items()):
            m = s["metrics"]
            status = slo_status(m["ttft_ms"]["p95"], m["tpot_ms"]["mean"], args.ttft_limit, args.tpot_limit)
            f.write(
                f"| {isl_i} | {dsl_i} | {label} | {s['config']['concurrency']} | "
                f"{out_tok:.2f} | {m['ttft_ms']['p95']:.2f} | {m['tpot_ms']['mean']:.2f} | {status} |\n"
            )

        f.write("\n## Notes\n\n")
        f.write("- SLO check uses TTFT P95 and TPOT mean.\n")
        f.write("- A PARTIAL indicates one of the two metrics is out of target; FAIL indicates both are out.\n")
        f.write("- `Peak GPU mem` shows peak used / total MB and utilization percentage.\n")
        f.write("- Optional (P) combinations that failed are marked as skipped/OOM and do not break the run.\n")

    print(f"Wrote comparison to {args.output}")


if __name__ == "__main__":
    main()