103 lines
3.8 KiB
Python
103 lines
3.8 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])-(?P<version>v[0-9]+)"
|
|
)
|
|
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(
|
|
"--workload-mode",
|
|
choices=("prefill_decode", "prefill_only"),
|
|
default="prefill_decode",
|
|
)
|
|
parser.add_argument("--cell-version", default="v3")
|
|
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 match.group("version") != args.cell_version
|
|
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",
|
|
"workload_mode": args.workload_mode,
|
|
"cell_version": args.cell_version,
|
|
"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()
|