Support long code traces in Frontier replay

runs/frontier-prefill-kvgrowth-fix-v0/run_replay_v5.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Replay an S3-real cell with profile-v5-kvgrowth (CPU, no GPU needed).
+
+Adapted from runs/frontier-s3-real-v0/run_frontier_prefix_replay.py with one
+change: the attention profile input is overridden to profile-v5-kvgrowth so the
+predictor retrains with the KV-context grid. Everything else (curves, trace,
+config argv, prefix caching) is identical to the S3-real sim runs, so old-vs-new
+differs by exactly one variable.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import importlib.util
+import json
+import sys
+from pathlib import Path
+
+
+csv.field_size_limit(16 * 1024 * 1024)
+
+
+ROOT = Path(__file__).resolve().parent
+S3_REAL = ROOT.parent / "frontier-s3-real-v0"
+
+
+def load_s3_module():
+    spec = importlib.util.spec_from_file_location(
+        "s3_prefix_replay", S3_REAL / "run_frontier_prefix_replay.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    sys.path.insert(0, str(S3_REAL))  # trace_utils import
+    spec.loader.exec_module(module)
+    return module
+
+
+def main() -> None:
+    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", "tp1_mns16"), required=True)
+    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)
+    parser.add_argument(
+        "--reuse-model-cache",
+        action="store_true",
+        help="Load predictor models from the content-addressed cache when available.",
+    )
+    parser.add_argument(
+        "--attention-profile",
+        type=Path,
+        default=ROOT / "profiles/profile-v5-kvgrowth/attention.csv",
+    )
+    args = parser.parse_args()
+
+    profile = args.attention_profile.resolve()
+    if not profile.is_file():
+        raise SystemExit(f"attention profile missing: {profile}")
+
+    module = load_s3_module()
+    original_replace = module.replace_flag
+
+    def replace_and_override(argv: list[str], flag: str, value: str) -> None:
+        original_replace(argv, flag, value)
+        # Piggyback on the first replace_flag call (trace file) to inject the
+        # profile override exactly once per run.
+        atten_flag = "--random_forrest_execution_time_predictor_config_atten_input_file"
+        if flag.endswith("trace_file") and atten_flag in argv:
+            original_replace(argv, atten_flag, str(profile))
+            no_cache_flag = "--random_forrest_execution_time_predictor_config_no_cache"
+            if args.reuse_model_cache and no_cache_flag in argv:
+                argv.remove(no_cache_flag)
+
+    module.replace_flag = replace_and_override
+    # The shared S3 parser predates TP1; reuse this wrapper's validated namespace.
+    module.parse_args = lambda: args
+    sys.argv = [
+        "run_frontier_prefix_replay.py",
+        "--trace", str(args.trace),
+        "--output-root", str(args.output_root),
+        "--config", args.config,
+        "--label", args.label,
+        "--max-tokens", str(args.max_tokens),
+    ]
+    if args.duration_s:
+        sys.argv += ["--duration-s", str(args.duration_s)]
+    if args.cache_root:
+        sys.argv += ["--cache-root", str(args.cache_root)]
+    module.main()
+
+    manifest_path = args.output_root / "manifest.json"
+    manifest = json.loads(manifest_path.read_text())
+    manifest["attention_profile_override"] = str(profile)
+    manifest["attention_profile_sha256"] = module.sha256(profile)
+    manifest["reuse_model_cache"] = args.reuse_model_cache
+    manifest["schema"] = "frontier-prefill-kvgrowth-replay-v1"
+    manifest_path.write_text(json.dumps(manifest, indent=2))
+    print(f"replay done: {args.output_root}")
+
+
+if __name__ == "__main__":
+    main()