sskj/experiments/dsv4_p800_256k_4k_probe/parse_results.py

#!/usr/bin/env python3
"""Parse P800 SGLang benchmark outputs and produce results.json + report.md.

Usage:
    python3 parse_results.py <result_root>
"""

import json
import os
import re
import sys
from pathlib import Path
from collections import OrderedDict

METRIC_PATTERNS = OrderedDict(
    [
        ("successful_requests", [r"Successful requests:\s+(\d+)"]),
        ("benchmark_duration_s", [r"Benchmark duration \(s\):\s+([\d.]+)"]),
        ("request_throughput", [r"Request throughput \(req/s\):\s+([\d.]+)"]),
        ("input_token_throughput", [r"Input token throughput \(tok/s\):\s+([\d.]+)"]),
        ("output_token_throughput", [r"Output token throughput \(tok/s\):\s+([\d.]+)"]),
        ("total_token_throughput", [r"Total token throughput \(tok/s\):\s+([\d.]+)"]),
        ("ttft_mean", [r"Mean TTFT \(ms\):\s+([\d.]+)"]),
        ("ttft_p50", [r"Median TTFT \(ms\):\s+([\d.]+)", r"P50 TTFT \(ms\):\s+([\d.]+)"]),
        ("ttft_p90", [r"P90 TTFT \(ms\):\s+([\d.]+)"]),
        ("ttft_p99", [r"P99 TTFT \(ms\):\s+([\d.]+)"]),
        ("tpot_mean", [r"Mean TPOT \(ms\):\s+([\d.]+)"]),
        ("tpot_p50", [r"Median TPOT \(ms\):\s+([\d.]+)", r"P50 TPOT \(ms\):\s+([\d.]+)"]),
        ("tpot_p90", [r"P90 TPOT \(ms\):\s+([\d.]+)"]),
        ("tpot_p99", [r"P99 TPOT \(ms\):\s+([\d.]+)"]),
        ("e2e_mean", [r"Mean E2E \(ms\):\s+([\d.]+)"]),
        ("e2e_p50", [r"Median E2E \(ms\):\s+([\d.]+)", r"P50 E2E \(ms\):\s+([\d.]+)"]),
        ("e2e_p90", [r"P90 E2E \(ms\):\s+([\d.]+)"]),
        ("e2e_p99", [r"P99 E2E \(ms\):\s+([\d.]+)"]),
    ]
)


def find_float(text: str, patterns: list[str]) -> float | None:
    for pat in patterns:
        m = re.search(pat, text)
        if m:
            try:
                return float(m.group(1))
            except ValueError:
                return None
    return None


def parse_summary_log(log_text: str) -> dict:
    result = {}
    for key, patterns in METRIC_PATTERNS.items():
        result[key] = find_float(log_text, patterns)
    return result


def percentile(values: list[float], p: float) -> float:
    if not values:
        return 0.0
    sorted_values = sorted(values)
    k = (len(sorted_values) - 1) * p / 100.0
    f = int(k)
    c = min(f + 1, len(sorted_values) - 1)
    if f == c:
        return sorted_values[f]
    return sorted_values[f] * (c - k) + sorted_values[c] * (k - f)


def parse_jsonl(jsonl_path: Path) -> tuple[list[dict], dict]:
    """Parse sglang.bench_serving --output-file output.

    Newer sglang writes one aggregate JSON object. Older versions write one
    JSON object per request. Returns (raw_requests, aggregates).
    """
    text = jsonl_path.read_text(encoding="utf-8", errors="replace").strip()
    if not text:
        return [], {}

    # Try single aggregate JSON object first.
    try:
        data = json.loads(text)
        if isinstance(data, dict) and "mean_e2e_latency_ms" in data:
            aggregates = {
                "success": data.get("total_output_tokens", 0) > 0 and 1 or 0,
                "failed": 0,
                "input_tokens": data.get("total_input_tokens", 0),
                "output_tokens": data.get("total_output_tokens", 0),
                "latencies": {
                    "e2e_ms": {
                        "mean": data.get("mean_e2e_latency_ms"),
                        "p50": data.get("median_e2e_latency_ms"),
                        "p90": data.get("p90_e2e_latency_ms"),
                        "p95": None,
                        "p99": data.get("p99_e2e_latency_ms"),
                    },
                    "ttft_ms": {
                        "mean": data.get("mean_ttft_ms"),
                        "p50": data.get("median_ttft_ms"),
                        "p90": None,
                        "p95": None,
                        "p99": data.get("p99_ttft_ms"),
                    },
                    "tpot_ms": {
                        "mean": data.get("mean_tpot_ms"),
                        "p50": data.get("median_tpot_ms"),
                        "p90": None,
                        "p95": None,
                        "p99": data.get("p99_tpot_ms"),
                    },
                    "itl_ms": {
                        "mean": data.get("mean_itl_ms"),
                        "p50": data.get("median_itl_ms"),
                        "p90": None,
                        "p95": data.get("p95_itl_ms"),
                        "p99": data.get("p99_itl_ms"),
                    },
                },
            }
            return [data], aggregates
    except json.JSONDecodeError:
        pass

    # Fall back to JSONL per-request parsing.
    raw_requests = []
    ttfts = []
    tpots = []
    itls = []
    e2es = []
    input_tokens = []
    output_tokens = []
    success = 0
    failed = 0

    for line in text.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
        except json.JSONDecodeError:
            continue

        raw_requests.append(req)

        ttft = req.get("ttft") or req.get("ttft_ms") or 0
        tpot = req.get("tpot") or req.get("tpot_ms") or 0
        itl = req.get("inter_token_latency") or req.get("itl") or req.get("itl_ms") or 0
        e2e = req.get("e2e_latency") or req.get("e2e") or req.get("e2e_ms") or 0
        in_tok = req.get("input_tokens") or req.get("prompt_tokens") or 0
        out_tok = req.get("output_tokens") or req.get("completion_tokens") or 0

        if ttft:
            ttfts.append(float(ttft))
        if tpot:
            tpots.append(float(tpot))
        if itl:
            itls.append(float(itl))
        if e2e:
            e2es.append(float(e2e))
        if in_tok:
            input_tokens.append(int(in_tok))
        if out_tok:
            output_tokens.append(int(out_tok))

        if req.get("success", True):
            success += 1
        else:
            failed += 1

    def latency_bucket(values: list[float]) -> dict:
        if not values:
            return {"mean": None, "p50": None, "p90": None, "p95": None, "p99": None}
        return {
            "mean": round(sum(values) / len(values), 2),
            "p50": round(percentile(values, 50), 2),
            "p90": round(percentile(values, 90), 2),
            "p95": round(percentile(values, 95), 2),
            "p99": round(percentile(values, 99), 2),
        }

    aggregates = {
        "success": success,
        "failed": failed,
        "input_tokens": sum(input_tokens),
        "output_tokens": sum(output_tokens),
        "latencies": {
            "e2e_ms": latency_bucket(e2es),
            "ttft_ms": latency_bucket(ttfts),
            "tpot_ms": latency_bucket(tpots),
            "itl_ms": latency_bucket(itls),
        },
    }

    return raw_requests, aggregates


def parse_scenario(result_root: Path, raw_file: Path) -> dict | None:
    """Parse one raw output file into a scenario dict."""
    # Filename: {chip}_{engine}_{MMDD}_{concurrency}_{input_len}_{output_len}.jsonl
    parts = raw_file.stem.split("_")
    if len(parts) < 6:
        return None
    try:
        concurrency = int(parts[-3])
        input_len = int(parts[-2])
        output_len = int(parts[-1])
    except ValueError:
        return None

    log_file = result_root / "logs" / f"sglang_c{concurrency}_i{input_len}_o{output_len}.log"
    summary = {}
    if log_file.exists():
        summary = parse_summary_log(log_file.read_text(encoding="utf-8", errors="replace"))

    raw_requests, aggregates = parse_jsonl(raw_file)

    # For newer sglang aggregate JSON, success/failed are not present in the
    # raw output file. Override them from the human-readable summary log when
    # it is available.
    if summary.get("successful_requests") is not None:
        aggregates["success"] = int(summary["successful_requests"])
        # The summary log only reports successes; assume failures are zero
        # unless the aggregate JSON already provided a non-zero failed count.
        if not aggregates.get("failed"):
            aggregates["failed"] = 0

    scenario = {
        "name": f"c{concurrency}_i{input_len}_o{output_len}",
        "concurrency": concurrency,
        "input_len": input_len,
        "output_len": output_len,
        "success": aggregates["success"],
        "failed": aggregates["failed"],
        "duration_s": summary.get("benchmark_duration_s"),
        "request_throughput": summary.get("request_throughput"),
        "input_token_throughput": summary.get("input_token_throughput"),
        "output_token_throughput": summary.get("output_token_throughput"),
        "total_token_throughput": summary.get("total_token_throughput"),
        "accept_length": None,
        "latencies": aggregates["latencies"],
        "raw_requests": raw_requests[:100] if len(raw_requests) <= 100 else None,
    }

    return scenario


def main() -> None:
    result_root = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("results")
    raw_dir = result_root / "raw_outputs"
    json_path = result_root / "results.json"
    report_path = result_root / "report.md"

    if not json_path.exists():
        raise SystemExit(f"metadata results.json not found: {json_path}")

    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    scenarios = []
    if raw_dir.exists():
        for raw_file in sorted(raw_dir.glob("*.jsonl")):
            scenario = parse_scenario(result_root, raw_file)
            if scenario:
                scenarios.append(scenario)

    data["scenarios"] = scenarios

    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

    # Generate report.md
    with open(report_path, "w", encoding="utf-8") as f:
        meta = data["metadata"]
        f.write(f"# Benchmark Report: {meta['experiment']}\n\n")
        f.write("## Metadata\n\n")
        f.write(f"- **Run ID**: {meta['run_id']}\n")
        f.write(f"- **Timestamp**: {meta['timestamp']}\n")
        f.write(f"- **Chip/Accelerator**: {meta.get('accelerator', '')} / {meta.get('chip', '')}\n")
        f.write(f"- **Engine/Backend**: {meta.get('engine', '')} / {meta.get('backend', '')}\n")
        f.write(f"- **Hardware**: {meta.get('hardware', '')}\n")
        f.write(f"- **Model**: {meta['model']}\n")
        f.write(f"- **Git Commit**: {meta.get('git_commit', 'unknown')}\n")
        f.write("\n## Results\n\n")
        f.write(
            "| Scenario | Conc | In/Out | Success | Failed | Req/s | OutTok/s | "
            "TTFT p50 | TTFT p99 | TPOT p50 | TPOT p99 | E2E p99 |\n"
        )
        f.write(
            "|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|\n"
        )
        for s in scenarios:
            lat = s["latencies"]
            f.write(
                f"| {s['name']} "
                f"| {s['concurrency']} "
                f"| {s['input_len']}/{s['output_len']} "
                f"| {s['success']} "
                f"| {s['failed']} "
                f"| {s.get('request_throughput') or ''} "
                f"| {s.get('output_token_throughput') or ''} "
                f"| {lat['ttft_ms']['p50'] or ''} "
                f"| {lat['ttft_ms']['p99'] or ''} "
                f"| {lat['tpot_ms']['p50'] or ''} "
                f"| {lat['tpot_ms']['p99'] or ''} "
                f"| {lat['e2e_ms']['p99'] or ''} |\n"
            )

    print(f"Updated {json_path}")
    print(f"Wrote {report_path}")


if __name__ == "__main__":
    main()