gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support long code traces in Frontier replay

Browse Source
This commit is contained in:
Gahow Wang
2026-07-24 00:57:08 +08:00
parent b80d3f03de
commit 7aed90f9e6
4 changed files with 646 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

235
runs/frontier-s3-real-v0/run_frontier_prefix_replay.py Normal file
View File

@@ -0,0 +1,235 @@
#!/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()