agentic-pd-hybrid/scripts/analysis/paired_compare.py

#!/usr/bin/env python3
"""Paired latency comparison with bootstrap CI.

Implements docs/EVALUATION_PROTOCOL_ZH.md §2.2 (M2 fix): when comparing
mechanism A vs B on the same trace, the only honest comparison is paired
on same-trial-mask. This script joins two metrics.jsonl by request_id,
keeps the rows where BOTH sides succeeded, and reports paired deltas
with 95% bootstrap CIs.

Out vs the existing `compare_no_error.py`:
  - works on raw metrics.jsonl, not pre-aggregated summary.json
  - bootstrap CIs (not just point estimates)
  - reports paired-mask size + per-side failure counts so the reader
    sees how many rows were dropped from the comparison

Usage:
  scripts/analysis/paired_compare.py \
      --baseline outputs/run-dp/request-metrics.jsonl \
      --candidate outputs/run-kvc/request-metrics.jsonl
  scripts/analysis/paired_compare.py ... --bootstrap 5000 --seed 42
  scripts/analysis/paired_compare.py ... --json > paired.json

stdlib only — no scipy/numpy. Runs without GPU and without SGLang.
"""

from __future__ import annotations

import argparse
import json
import math
import random
import sys
from pathlib import Path


def _load(path: Path) -> dict[str, dict]:
    out: dict[str, dict] = {}
    with path.open() as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            rid = row.get("request_id")
            if rid is None:
                continue
            out[rid] = row
    return out


def _ok(row: dict) -> bool:
    return row.get("error") is None and row.get("latency_s") is not None


def _quantile(values: list[float], q: float) -> float:
    if not values:
        return float("nan")
    s = sorted(values)
    if len(s) == 1:
        return s[0]
    pos = (len(s) - 1) * q
    lo = math.floor(pos)
    hi = math.ceil(pos)
    if lo == hi:
        return s[lo]
    return s[lo] + (s[hi] - s[lo]) * (pos - lo)


def _stats(deltas: list[float]) -> dict[str, float]:
    if not deltas:
        return {"mean": float("nan"), "p50": float("nan"), "p90": float("nan"), "p99": float("nan")}
    return {
        "mean": sum(deltas) / len(deltas),
        "p50": _quantile(deltas, 0.50),
        "p90": _quantile(deltas, 0.90),
        "p99": _quantile(deltas, 0.99),
    }


def _bootstrap_ci(
    deltas: list[float], statistic, n_boot: int, rng: random.Random
) -> tuple[float, float]:
    """Return (lo, hi) 95% CI for `statistic(deltas)`."""
    if len(deltas) < 2:
        return (float("nan"), float("nan"))
    n = len(deltas)
    samples = []
    for _ in range(n_boot):
        # resample with replacement
        resample = [deltas[rng.randrange(n)] for _ in range(n)]
        samples.append(statistic(resample))
    samples.sort()
    lo = samples[int(0.025 * (n_boot - 1))]
    hi = samples[int(0.975 * (n_boot - 1))]
    return (lo, hi)


def compare(
    baseline: dict[str, dict],
    candidate: dict[str, dict],
    *,
    metric: str,
    n_boot: int,
    seed: int,
) -> dict:
    common_ids = set(baseline.keys()) & set(candidate.keys())
    paired_ids = [
        rid for rid in common_ids if _ok(baseline[rid]) and _ok(candidate[rid])
    ]
    paired_ids.sort()

    base_only_fail = sum(1 for rid in common_ids if not _ok(baseline[rid]))
    cand_only_fail = sum(1 for rid in common_ids if not _ok(candidate[rid]))

    deltas = []
    wins = losses = ties = 0
    for rid in paired_ids:
        b = baseline[rid].get(metric)
        c = candidate[rid].get(metric)
        if b is None or c is None:
            continue
        d = float(c) - float(b)
        deltas.append(d)
        if d < 0:
            wins += 1
        elif d > 0:
            losses += 1
        else:
            ties += 1

    rng = random.Random(seed)
    stats = _stats(deltas)
    ci_mean = _bootstrap_ci(deltas, lambda x: sum(x) / len(x), n_boot, rng)
    ci_p50 = _bootstrap_ci(deltas, lambda x: _quantile(x, 0.50), n_boot, rng)
    ci_p90 = _bootstrap_ci(deltas, lambda x: _quantile(x, 0.90), n_boot, rng)

    return {
        "metric": metric,
        "baseline_size": len(baseline),
        "candidate_size": len(candidate),
        "intersection_size": len(common_ids),
        "paired_size": len(paired_ids),
        "baseline_fail_in_common": base_only_fail,
        "candidate_fail_in_common": cand_only_fail,
        "delta_stats": stats,
        "delta_mean_ci95": ci_mean,
        "delta_p50_ci95": ci_p50,
        "delta_p90_ci95": ci_p90,
        "wins_candidate": wins,
        "losses_candidate": losses,
        "ties": ties,
    }


def _fmt(x: float, w: int = 6) -> str:
    if x is None or (isinstance(x, float) and math.isnan(x)):
        return "  nan "
    return f"{x:+{w}.3f}"


def render(result: dict) -> str:
    s = result["delta_stats"]
    mlo, mhi = result["delta_mean_ci95"]
    p5lo, p5hi = result["delta_p50_ci95"]
    p9lo, p9hi = result["delta_p90_ci95"]
    n = result["paired_size"]
    lines = [
        f"# paired comparison ({result['metric']})",
        "",
        f"baseline rows:           {result['baseline_size']}",
        f"candidate rows:          {result['candidate_size']}",
        f"intersection (rid):      {result['intersection_size']}",
        f"paired (both ok):        {result['paired_size']}",
        f"  baseline fails in common:  {result['baseline_fail_in_common']}",
        f"  candidate fails in common: {result['candidate_fail_in_common']}",
        "",
        "## delta (candidate - baseline)  — negative = candidate is faster",
        "",
        "| stat | value | 95% CI |",
        "|---|---:|---:|",
        f"| mean | {_fmt(s['mean'])} | [{_fmt(mlo)}, {_fmt(mhi)}] |",
        f"| p50  | {_fmt(s['p50'])}  | [{_fmt(p5lo)}, {_fmt(p5hi)}] |",
        f"| p90  | {_fmt(s['p90'])}  | [{_fmt(p9lo)}, {_fmt(p9hi)}] |",
        f"| p99  | {_fmt(s['p99'])}  | — |",
        "",
        f"win/loss/tie: {result['wins_candidate']} / {result['losses_candidate']} / {result['ties']} (of {n})",
    ]
    return "\n".join(lines)


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    p.add_argument("--baseline", required=True, type=Path)
    p.add_argument("--candidate", required=True, type=Path)
    p.add_argument(
        "--metric",
        default="latency_s",
        choices=["latency_s", "ttft_s", "tpot_s"],
        help="which per-request field to compare (default: latency_s)",
    )
    p.add_argument("--bootstrap", type=int, default=2000)
    p.add_argument("--seed", type=int, default=20260512)
    p.add_argument("--json", action="store_true")
    args = p.parse_args()

    baseline = _load(args.baseline)
    candidate = _load(args.candidate)
    if not baseline or not candidate:
        print("empty input on one side", file=sys.stderr)
        sys.exit(1)

    result = compare(
        baseline, candidate,
        metric=args.metric, n_boot=args.bootstrap, seed=args.seed,
    )

    if args.json:
        json.dump(result, sys.stdout, indent=2, default=lambda x: None if isinstance(x, float) and math.isnan(x) else x)
        sys.stdout.write("\n")
    else:
        print(render(result))


if __name__ == "__main__":
    main()