106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
#!/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()
|