#!/usr/bin/env python3 """Replay a remapped trace with S1 joint injection and prefix caching enabled.""" from __future__ import annotations import argparse import csv import json import os import subprocess from collections import Counter from pathlib import Path from typing import Any from trace_utils import distribution, sha256, write_json ROOT = Path(__file__).resolve().parent REPO = ROOT.parents[1] REFERENCE = REPO / "runs/frontier-collective-joint-v0/counterfactual/joint-r2/manifest.json" EXPECTED_FRONTIER_COMMIT = "deadc4a321f0baaa534c6ebd17f974123733cdc2" MAX_CURVE_BATCH = 32 csv.field_size_limit(16 * 1024 * 1024) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--trace", type=Path, required=True) parser.add_argument("--output-root", type=Path, required=True) parser.add_argument("--config", choices=("tp4_mns16", "tp2_mns16"), default="tp4_mns16") parser.add_argument("--label", required=True) parser.add_argument("--max-tokens", type=int, required=True) parser.add_argument("--duration-s", type=float) parser.add_argument("--cache-root", type=Path) return parser.parse_args() def replace_flag(argv: list[str], flag: str, value: str) -> None: try: index = argv.index(flag) except ValueError as error: raise ValueError(f"template command is missing {flag}") from error argv[index + 1] = value def prepare_curves(reference: dict[str, Any], output_root: Path) -> dict[str, Any]: inputs = output_root / "inputs" inputs.mkdir(parents=True) outputs: dict[str, Any] = {} for kind, reference_key, filename in ( ("collective", "collective_curve", "collective-curve-b4-extrapolated.json"), ("moe", "moe_curve", "fused-moe-curve-b4-extrapolated.json"), ): payload = json.loads(Path(reference[reference_key]).read_text()) curves = ( payload["curve_variants_ms_per_step"].values() if kind == "collective" and "curve_variants_ms_per_step" in payload else (payload,) ) for curve in curves: for points in curve.values(): b4 = points["4"] for batch in range(5, MAX_CURVE_BATCH + 1): points[str(batch)] = b4 destination = inputs / filename write_json(destination, payload) outputs[kind] = str(destination) outputs[f"{kind}_sha256"] = sha256(destination) return outputs def find_one(root: Path, name: str) -> Path: matches = list(root.glob(f"**/{name}")) if len(matches) != 1: raise ValueError(f"expected one {name} below {root}, found {matches}") return matches[0] def read_csv(path: Path) -> list[dict[str, str]]: with path.open(newline="") as stream: return list(csv.DictReader(stream)) def summarize(trace: Path, metrics_root: Path, duration_s: float) -> dict[str, Any]: trace_rows = read_csv(trace) metrics_path = find_one(metrics_root, "request_metrics.csv") metric_rows = read_csv(metrics_path) if len(trace_rows) != len(metric_rows): raise ValueError(f"request count mismatch: trace={len(trace_rows)}, metrics={len(metric_rows)}") system_path = find_one(metrics_root, "system_metrics.json") system = json.loads(system_path.read_text()) ledger_path = find_one(metrics_root, "frontier_stage_batch_ledger.jsonl") waiting_ms = [float(row["request_waiting_time_total"]) for row in metric_rows] ttft_ms = [float(row["ttft"]) for row in metric_rows] # TPOT is undefined for one-token outputs because there is no inter-token # interval. Frontier records those cells as an empty CSV value. tpot_ms = [ float(row["tpot"]) for row in metric_rows if row.get("tpot") is not None and row["tpot"].strip() ] e2e_ms = [float(row["request_e2e_time"]) for row in metric_rows] cached_tokens = [int(float(row.get("request_cached_prefill_tokens", 0))) for row in metric_rows] query_blocks = [int(float(row.get("request_prefix_cache_query_blocks", 0))) for row in metric_rows] hit_blocks = [int(float(row.get("request_prefix_cache_hit_blocks", 0))) for row in metric_rows] arrivals = [float(row["arrived_at"]) for row in trace_rows] prefill = [int(row["num_prefill_tokens"]) for row in trace_rows] decode = [int(row["num_decode_tokens"]) for row in trace_rows] completion_times = [arrival + e2e / 1000 for arrival, e2e in zip(arrivals, e2e_ms)] batch_hist: Counter[int] = Counter() with ledger_path.open() as stream: for line in stream: if not line.strip(): continue row = json.loads(line) tokens = row.get("request_num_tokens") or [] if tokens and all(int(token) == 1 for token in tokens): batch_hist[len(tokens)] += 1 stages = sum(batch_hist.values()) prefix = system.get("prefix_cache_statistics") or { "block_size_tokens": 16, "requests": len(metric_rows), "total_cached_prefill_tokens": sum(cached_tokens), "total_query_blocks": sum(query_blocks), "total_hit_blocks": sum(hit_blocks), "hit_ratio": sum(hit_blocks) / sum(query_blocks) if sum(query_blocks) else 0.0, } return { "schema": "frontier-s3-real-prefix-replay-summary-v1", "requests": len(trace_rows), "duration_s": duration_s, "offered_load": { "requests_per_s": len(trace_rows) / duration_s, "prefill_tokens_per_s_raw": sum(prefill) / duration_s, "prefill_tokens_per_s_after_prefix": (sum(prefill) - sum(cached_tokens)) / duration_s, "decode_tokens_per_s": sum(decode) / duration_s, }, "latency_ms": { "waiting": distribution(waiting_ms), "ttft": distribution(ttft_ms), "tpot": distribution(tpot_ms), "e2e": distribution(e2e_ms), }, "drain": { "last_arrival_s": max(arrivals), "last_completion_s": max(completion_times), "tail_after_last_arrival_s": max(completion_times) - max(arrivals), }, "decode_batch": { "stages": stages, "histogram": dict(sorted(batch_hist.items())), "share_gt_1": sum(count for batch, count in batch_hist.items() if batch > 1) / stages if stages else 0.0, "share_gt_4": sum(count for batch, count in batch_hist.items() if batch > 4) / stages if stages else 0.0, "max": max(batch_hist, default=0), }, "prefix_cache": prefix, "artifacts": { "request_metrics": str(metrics_path), "system_metrics": str(system_path), "stage_ledger": str(ledger_path), }, } def main() -> None: args = parse_args() args.trace = args.trace.resolve() args.output_root = args.output_root.resolve() args.cache_root = (args.cache_root or args.output_root.parent / "cache").resolve() if args.output_root.exists(): raise ValueError(f"refusing to overwrite {args.output_root}") reference = json.loads(REFERENCE.read_text()) frontier = Path(reference["frontier_checkout"]) commit = subprocess.check_output(["git", "-C", str(frontier), "rev-parse", "HEAD"], text=True).strip() if commit != EXPECTED_FRONTIER_COMMIT or commit != reference["frontier_commit"]: raise ValueError(f"Frontier commit drift: {commit}") if subprocess.check_output(["git", "-C", str(frontier), "status", "--porcelain"], text=True).strip(): raise ValueError("Frontier checkout must be clean") args.output_root.mkdir(parents=True) curves = prepare_curves(reference, args.output_root) cell = reference["cells"][args.config] argv = list(cell["argv"]) replace_flag(argv, "--trace_request_generator_config_trace_file", str(args.trace)) replace_flag(argv, "--trace_request_generator_config_max_tokens", str(args.max_tokens)) replace_flag(argv, "--random_forrest_execution_time_predictor_config_prediction_max_tokens_per_request", str(args.max_tokens)) replace_flag(argv, "--metrics_config_output_dir", str(args.output_root / "metrics")) replace_flag(argv, "--metrics_config_run_id", args.label) replace_flag(argv, "--metrics_config_cache_dir", str(args.cache_root / "model")) replace_flag(argv, "--vidur_cc_backend_config_cache_dir", str(args.cache_root / "cc")) argv.extend(("--vllm_v1_scheduler_config_enable_prefix_caching", "--log_level", "warning", "--cluster_event_log_level", "WARNING")) usage = args.output_root / "usage.json" env = os.environ.copy() env.update( { "CUDA_VISIBLE_DEVICES": "", "PYTHONDONTWRITEBYTECODE": "1", "PYTHONPATH": os.pathsep.join([str(frontier), *reference["python_dependency_roots"]]), "FRONTIER_COLLECTIVE_CURVE": curves["collective"], "FRONTIER_COLLECTIVE_CURVE_VARIANT": reference["collective_curve_variant"], "FRONTIER_FUSED_MOE_CURVE": curves["moe"], "FRONTIER_CURVE_USAGE": str(usage), } ) manifest = { "schema": "frontier-s3-real-prefix-replay-v1", "frontier_commit": commit, "reference_manifest": str(REFERENCE.resolve()), "reference_manifest_sha256": sha256(REFERENCE), "trace": str(args.trace), "trace_sha256": sha256(args.trace), "config": args.config, "prefix_caching": True, "block_size": 16, "max_tokens": args.max_tokens, "argv": argv, "curves": curves, } write_json(args.output_root / "manifest.json", manifest) log_path = args.output_root / "sim.log" with log_path.open("w") as log: completed = subprocess.run(argv, cwd=frontier, env=env, stdout=log, stderr=subprocess.STDOUT, check=False) if completed.returncode: raise SystemExit(f"sim failed with exit code {completed.returncode}; see {log_path}") with args.trace.open(newline="") as stream: trace_rows = list(csv.DictReader(stream)) duration_s = args.duration_s or max(float(row["arrived_at"]) for row in trace_rows) or 1.0 summary = summarize(args.trace, args.output_root / "metrics", duration_s) write_json(args.output_root / "summary.json", summary) print(json.dumps(summary, sort_keys=True)) if __name__ == "__main__": main()