aituner/scripts/collectivespec/summarize_static_k_pilot.py

#!/usr/bin/env python3
"""Summarize immutable artifacts from ``run_static_k_pilot.py``.

The output intentionally distinguishes an absent/failed result from an
SLO-infeasible probe, and includes a small data-sanity block for review.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from typing import Any


def number(value: Any) -> float | None:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    value = float(value)
    return value if math.isfinite(value) else None


def first_result(store: Path) -> Path | None:
    candidates = sorted(store.glob("*/trials/trial-*/result.json"))
    return candidates[0] if candidates else None


def summarize_result(k: int, path: Path | None) -> dict[str, Any]:
    record: dict[str, Any] = {"k": k, "label": "NoSpec" if k == 0 else f"K={k}"}
    if path is None:
        record.update({"status": "missing_result"})
        return record

    payload = json.loads(path.read_text())
    probes = payload.get("probes") if isinstance(payload.get("probes"), list) else []
    feasible = [probe for probe in probes if isinstance(probe, dict) and probe.get("feasible")]
    best_payload = feasible[-1].get("payload", {}) if feasible else {}
    latency = best_payload.get("latency_summary", {}) if isinstance(best_payload, dict) else {}
    tpot = latency.get("tpot_ms", {}) if isinstance(latency, dict) else {}
    record.update(
        {
            "status": str(payload.get("status", "unknown")),
            "result_path": str(path),
            "best_sampling_u": number(payload.get("best_sampling_u")),
            "best_request_rate": number(payload.get("best_request_rate")),
            "best_pass_rate": number(payload.get("best_pass_rate")),
            "best_request_count": number(payload.get("best_request_count")),
            "probe_count": len(probes),
            "feasible_probe_count": len(feasible),
            "probe_thresholds": [number(probe.get("threshold")) for probe in probes if isinstance(probe, dict)],
            "best_tpot_ms": {
                metric: number(tpot.get(metric)) for metric in ("mean", "p50", "p90", "p95", "p99")
            },
        }
    )
    return record


def sanity(records: list[dict[str, Any]]) -> dict[str, Any]:
    completed = [r for r in records if r.get("status") == "completed"]
    values = [r["best_sampling_u"] for r in completed if number(r.get("best_sampling_u")) is not None]
    pass_rates = [r["best_pass_rate"] for r in completed if number(r.get("best_pass_rate")) is not None]
    distinct = len(set(values))
    return {
        "n_configurations": len(records),
        "n_completed": len(completed),
        "sampling_u": {
            "n": len(values),
            "min": min(values) if values else None,
            "max": max(values) if values else None,
            "distinct_value_count": distinct,
            "all_identical": len(values) > 1 and distinct == 1,
            "non_negative": all(value >= 0 for value in values),
        },
        "pass_rate": {
            "n": len(pass_rates),
            "min": min(pass_rates) if pass_rates else None,
            "max": max(pass_rates) if pass_rates else None,
            "distinct_value_count": len(set(pass_rates)),
            "within_0_1": all(0 <= value <= 1 for value in pass_rates),
        },
        "invariants": {
            "one_result_per_k": len({r["k"] for r in records}) == len(records),
            "all_best_request_rates_non_negative": all(
                (value := number(r.get("best_request_rate"))) is None or value >= 0 for r in completed
            ),
        },
    }


def to_markdown(records: list[dict[str, Any]], checks: dict[str, Any]) -> str:
    lines = [
        "# CollectiveSpec static-K screening summary",
        "",
        "| configuration | status | max feasible sampling_u | request rate | pass rate | p95 TPOT (ms) | probes |",
        "|---|---:|---:|---:|---:|---:|---:|",
    ]
    for r in records:
        tpot = r.get("best_tpot_ms") or {}
        fmt = lambda value: "—" if value is None else f"{value:.6g}"
        lines.append(
            "| {label} | {status} | {u} | {rate} | {pass_rate} | {p95} | {probes} |".format(
                label=r["label"],
                status=r["status"],
                u=fmt(r.get("best_sampling_u")),
                rate=fmt(r.get("best_request_rate")),
                pass_rate=fmt(r.get("best_pass_rate")),
                p95=fmt(tpot.get("p95")),
                probes=r.get("probe_count", "—"),
            )
        )
    lines.extend(["", "## Data sanity", "", "```json", json.dumps(checks, indent=2, sort_keys=True), "``", ""])
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--output-json", type=Path, required=True)
    parser.add_argument("--output-md", type=Path, required=True)
    args = parser.parse_args()

    manifest = json.loads((args.root / "manifest.json").read_text())
    k_values = [int(k) for k in manifest["k_values"]]
    records = [summarize_result(k, first_result(args.root / "stores" / f"k{k}")) for k in sorted(k_values)]
    checks = sanity(records)
    output = {"manifest": manifest, "records": records, "data_sanity": checks}
    args.output_json.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n")
    args.output_md.write_text(to_markdown(records, checks))
    print(json.dumps(output, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())