#!/usr/bin/env python3 """Replay the exact BC-8 simulator cells with the whole-layer decode curve.""" from __future__ import annotations import argparse import hashlib import json import os import subprocess import sys import time from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path ROOT = Path(__file__).resolve().parent REPO = ROOT.parents[1] EXPECTED_FRONTIER_COMMIT = "deadc4a321f0baaa534c6ebd17f974123733cdc2" CONFIGS = ("tp1_mns16", "tp2_mns16", "tp4_mns16", "tp4_mns32") SOURCE_MANIFEST = ( REPO / "runs/frontier-collective-joint-v0/counterfactual/joint-r2/manifest.json" ) JOINT_INPUTS = REPO / "runs/frontier-knee-sweep-v0/inputs" GRID = ROOT / "results/grid.json" WRAPPER = ROOT / "run_frontier_with_whole_decode_curve.py" CACHE_ROOT = REPO / "runs/frontier-knee-sweep-v0/cache" LOCAL_DEPENDENCY_ROOTS = ( REPO / "runs/frontier-collective-joint-v0/python-deps", Path("/home/gahow/.cache/uv/archive-v0/-_kzErLcPO5nASZFX8b9k"), Path("/home/gahow/.cache/uv/archive-v0/FbaBs_QJ9QKEbQ9V_4aIR"), Path("/home/gahow/.cache/uv/archive-v0/fuHsGXD0Lv_UjFC8yI4-7"), Path("/home/gahow/.cache/uv/archive-v0/jFGdqQLpB1eopfm9VxT3j"), Path("/home/gahow/.cache/uv/archive-v0/YWW6ExSJuPVvv4-qYQTin"), Path("/home/gahow/.cache/uv/archive-v0/3_qxZ5Ll-EpVAGZfbksfe"), ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--frontier-checkout", type=Path, required=True) parser.add_argument("--output-root", type=Path, default=ROOT / "replay/bc8") parser.add_argument("--jobs", type=int, default=2) return parser.parse_args() def sha256(path: Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() def write_json(path: Path, payload) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") def replace_flag(argv: list[str], flag: str, value: str) -> None: index = argv.index(flag) argv[index + 1] = value def validate(frontier: Path) -> None: commit = subprocess.check_output( ["git", "-C", str(frontier), "rev-parse", "HEAD"], text=True ).strip() status = subprocess.check_output( ["git", "-C", str(frontier), "status", "--porcelain"], text=True ).strip() if commit != EXPECTED_FRONTIER_COMMIT or status: raise ValueError( f"Frontier must be clean at {EXPECTED_FRONTIER_COMMIT}, " f"got commit={commit}, dirty={bool(status)}" ) required = ( SOURCE_MANIFEST, GRID, WRAPPER, CACHE_ROOT, *LOCAL_DEPENDENCY_ROOTS, ) missing = [str(path) for path in required if not path.exists()] if missing: raise ValueError(f"missing replay inputs: {missing}") def trace_for(config: str) -> Path: tp = int(config[2]) return ( JOINT_INPUTS / f"traces-per-gpu-low/tp{tp}/w0-short-fixed-uniform-none/" "rho0p02/public/frontier.csv" ) def run_one( config: str, *, frontier: Path, output_root: Path, templates: dict, ) -> dict: point = output_root / "raw" / config metrics_root = point / "metrics" expected = list(metrics_root.glob("**/system_metrics.json")) if len(expected) == 1 and (point / "usage.json").is_file(): return {"config": config, "status": "skipped_complete", "elapsed_s": 0.0} argv = list(templates[config]["argv"]) argv[0] = sys.executable argv[1] = str(WRAPPER.resolve()) replace_flag(argv, "--trace_request_generator_config_trace_file", str(trace_for(config))) replace_flag(argv, "--metrics_config_output_dir", str(metrics_root)) replace_flag(argv, "--metrics_config_run_id", f"decode_grid_bc8_{config}") replace_flag(argv, "--metrics_config_cache_dir", str(CACHE_ROOT / "model")) replace_flag(argv, "--vidur_cc_backend_config_cache_dir", str(CACHE_ROOT / "cc")) env = os.environ.copy() env.update( { "CUDA_VISIBLE_DEVICES": "", "PYTHONDONTWRITEBYTECODE": "1", "PYTHONPATH": os.pathsep.join( [str(frontier), *(str(path) for path in LOCAL_DEPENDENCY_ROOTS)] ), "FRONTIER_COLLECTIVE_CURVE": str( (JOINT_INPUTS / "collective-curve-b4-extrapolated.json").resolve() ), "FRONTIER_COLLECTIVE_CURVE_VARIANT": "drop_mean", "FRONTIER_FUSED_MOE_CURVE": str( (JOINT_INPUTS / "fused-moe-curve-b4-extrapolated.json").resolve() ), "FRONTIER_WHOLE_DECODE_GRID": str(GRID.resolve()), "FRONTIER_CURVE_USAGE": str((point / "usage.json").resolve()), } ) point.mkdir(parents=True, exist_ok=True) write_json(point / "command.json", argv) started = time.monotonic() with (point / "run.log").open("w") as output: completed = subprocess.run( argv, cwd=frontier, env=env, stdout=output, stderr=subprocess.STDOUT, check=False, ) elapsed = time.monotonic() - started metrics = list(metrics_root.glob("**/system_metrics.json")) status = ( "completed" if completed.returncode == 0 and len(metrics) == 1 and (point / "usage.json").is_file() else "failed" ) record = { "config": config, "status": status, "returncode": completed.returncode, "elapsed_s": elapsed, } write_json(point / "run-status.json", record) return record def main() -> None: args = parse_args() if args.jobs < 1: raise ValueError("--jobs must be positive") frontier = args.frontier_checkout.resolve() output_root = args.output_root.resolve() validate(frontier) source = json.loads(SOURCE_MANIFEST.read_text()) manifest = { "schema": "frontier-decode-grid-bc8-replay.v1", "frontier_checkout": str(frontier), "frontier_commit": EXPECTED_FRONTIER_COMMIT, "configs": list(CONFIGS), "rho_per_gpu": 0.02, "wrapper": str(WRAPPER.resolve()), "wrapper_sha256": sha256(WRAPPER), "whole_decode_grid": str(GRID.resolve()), "whole_decode_grid_sha256": sha256(GRID), "collective_curve_sha256": sha256( JOINT_INPUTS / "collective-curve-b4-extrapolated.json" ), "moe_curve_sha256": sha256( JOINT_INPUTS / "fused-moe-curve-b4-extrapolated.json" ), "traces": { config: { "path": str(trace_for(config).resolve()), "sha256": sha256(trace_for(config)), } for config in CONFIGS }, } write_json(output_root / "manifest.json", manifest) results = [] with ThreadPoolExecutor(max_workers=args.jobs) as pool: futures = { pool.submit( run_one, config, frontier=frontier, output_root=output_root, templates=source["cells"], ): config for config in CONFIGS } for future in as_completed(futures): result = future.result() results.append(result) print(json.dumps(result, sort_keys=True), flush=True) results.sort(key=lambda row: CONFIGS.index(row["config"])) write_json(output_root / "run-summary.json", results) failures = [row for row in results if row["status"] == "failed"] if failures: raise SystemExit(f"failed replay cells: {failures}") if __name__ == "__main__": main()