Files
aituner/runs/frontier-code-trace-v0/summarize_calibration.py

89 lines
3.4 KiB
Python

#!/usr/bin/env python3
"""Compact code-trace Frontier rho calibration into a reviewable artifact."""
from __future__ import annotations
import argparse
import hashlib
import json
import re
from pathlib import Path
CELL_RE = re.compile(r"r(?P<rho>[0-9p]+)-tp(?P<tp>[24])-v3")
ATTENTION_FLAG = "--random_forrest_execution_time_predictor_config_atten_input_file"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--sim-root", type=Path, required=True)
parser.add_argument("--profile-manifest", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
return parser.parse_args()
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def main() -> None:
args = parse_args()
profile = json.loads(args.profile_manifest.read_text())
expected_profile_sha = profile["output_sha256"]
cells = []
for root in sorted(args.sim_root.iterdir()):
match = CELL_RE.fullmatch(root.name)
if match is None or not (root / "summary.json").is_file():
continue
manifest = json.loads((root / "manifest.json").read_text())
observed_profile_sha = manifest.get("attention_profile_sha256")
if observed_profile_sha is None:
argv = manifest["argv"]
profile_path = Path(argv[argv.index(ATTENTION_FLAG) + 1])
observed_profile_sha = sha256(profile_path)
if observed_profile_sha != expected_profile_sha:
raise ValueError(
f"{root}: attention profile {observed_profile_sha} "
f"!= profile-v6 {expected_profile_sha}"
)
summary = json.loads((root / "summary.json").read_text())
drain_fraction = (
summary["drain"]["tail_after_last_arrival_s"] / summary["duration_s"]
)
cells.append(
{
"topology": f"tp{match.group('tp')}_mns16",
"rho": float(match.group("rho").replace("p", ".")),
"requests": summary["requests"],
"offered_load": summary["offered_load"],
"waiting_ms": summary["latency_ms"]["waiting"],
"ttft_ms": summary["latency_ms"]["ttft"],
"tpot_ms": summary["latency_ms"]["tpot"],
"decode_batch": summary["decode_batch"],
"drain_tail_s": summary["drain"]["tail_after_last_arrival_s"],
"drain_fraction": drain_fraction,
"subcritical_gate": drain_fraction <= 0.1,
"prefix_cache_hit_ratio": summary["prefix_cache"]["hit_ratio"],
"trace_sha256": manifest["trace_sha256"],
}
)
if not cells:
raise ValueError(f"no completed v3 calibration cells below {args.sim_root}")
payload = {
"schema": "frontier-code-trace-calibration-summary-v1",
"attention_profile_sha256": expected_profile_sha,
"subcritical_rule": "drain tail <= 10% of the 3660s arrival window",
"cells": cells,
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
print(json.dumps({"cells": len(cells), "output": str(args.output)}, sort_keys=True))
if __name__ == "__main__":
main()