Files
aituner/runs/frontier-fidelity-envelope-v1/run_frontier_qwen30_exact_trace_surface.py

425 lines
16 KiB
Python

#!/usr/bin/env python3
"""Freeze Frontier's Qwen30 exact production-trace response surface."""
from __future__ import annotations
import argparse
import csv
import importlib.util
import json
import math
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Any
TARGET_PASS_RATE = 0.95
TPOT_SLOS_MS = (50.0, 100.0, 150.0, 180.0)
WINDOW_SECONDS = 600.0
BASE_RUNNER = (
Path(__file__).resolve().parents[1]
/ "frontier-phase-factorial-v0/run_frontier_qwen30_prefill_surface.py"
)
def load_base():
spec = importlib.util.spec_from_file_location("qwen30_prefill_surface_base", BASE_RUNNER)
if spec is None or spec.loader is None:
raise ImportError(BASE_RUNNER)
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
BASE = load_base()
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--frontier-source", type=Path, required=True)
parser.add_argument("--replayserve-root", type=Path, required=True)
parser.add_argument("--profile-root", type=Path, required=True)
parser.add_argument("--python-deps", type=Path, required=True)
parser.add_argument("--output-root", type=Path, required=True)
parser.add_argument(
"--trace",
action="append",
required=True,
help="Frozen trace anchor as LABEL=PATH; repeat in increasing load order.",
)
parser.add_argument("--config", action="append")
parser.add_argument(
"--rate-contract",
choices=("trace-window", "uniform-spacing"),
default="trace-window",
)
parser.add_argument(
"--prefix-caching",
action=argparse.BooleanOptionalAction,
default=True,
)
parser.add_argument(
"--cc-backend", choices=("analytical", "vidur"), default="vidur"
)
parser.add_argument("--allreduce-csv", type=Path)
parser.add_argument("--timeout-seconds", type=float, default=1800.0)
parser.add_argument("--resume", action="store_true")
return parser.parse_args()
def ttft_limit_ms(input_tokens: int) -> float:
return 1000.0 + 1000.0 * input_tokens / 8000.0
def percentile(values: list[float], fraction: float) -> float | None:
if not values:
return None
ordered = sorted(values)
return ordered[math.ceil(fraction * len(ordered)) - 1]
def parse_trace(
specification: str, *, rate_contract: str = "trace-window"
) -> dict[str, Any]:
if "=" not in specification:
raise ValueError(f"trace must be LABEL=PATH: {specification}")
label, raw_path = specification.split("=", 1)
if not label or "/" in label:
raise ValueError(f"invalid trace label: {label!r}")
path = Path(raw_path).resolve()
if not path.is_file():
raise FileNotFoundError(path)
with path.open(newline="") as source:
rows = list(csv.DictReader(source))
if not rows:
raise ValueError(f"empty trace: {path}")
required = {
"arrived_at",
"num_prefill_tokens",
"num_decode_tokens",
"session_id",
"block_hash_ids",
}
if not required.issubset(rows[0]):
raise ValueError(f"trace columns missing: {required - set(rows[0])}")
arrivals = [float(row["arrived_at"]) for row in rows]
if any(not math.isfinite(value) or value < 0 for value in arrivals):
raise ValueError(f"invalid arrival in {path}")
if any(right < left for left, right in zip(arrivals, arrivals[1:])):
raise ValueError(f"arrival order drift in {path}")
if rate_contract == "trace-window":
offered_request_rate = len(rows) / WINDOW_SECONDS
elif rate_contract == "uniform-spacing":
if len(rows) < 2 or arrivals[-1] <= arrivals[0]:
raise ValueError("uniform-spacing traces require at least two arrivals")
intervals = [right - left for left, right in zip(arrivals, arrivals[1:])]
expected_interval = (arrivals[-1] - arrivals[0]) / (len(arrivals) - 1)
if any(abs(value - expected_interval) > 1e-9 for value in intervals):
raise ValueError(f"non-uniform fixed trace: {path}")
offered_request_rate = 1.0 / expected_interval
else:
raise ValueError(f"unknown rate contract: {rate_contract}")
shapes = [
(int(row["num_prefill_tokens"]), int(row["num_decode_tokens"]))
for row in rows
]
if any(prefill <= 0 or decode <= 0 or prefill + decode > 40960 for prefill, decode in shapes):
raise ValueError(f"out-of-contract shape in {path}")
return {
"label": label,
"path": path,
"sha256": BASE.sha256(path),
"requests": len(rows),
"offered_request_rate": offered_request_rate,
"first_arrival_s": arrivals[0],
"last_arrival_s": arrivals[-1],
"shapes": shapes,
}
def find_request_metrics(run_dir: Path) -> Path:
matches = list((run_dir / "metrics").rglob("request_metrics.csv"))
if len(matches) != 1:
raise RuntimeError(f"expected one request_metrics.csv, got {matches}")
return matches[0]
def score(path: Path, expected_shapes: list[tuple[int, int]]) -> dict[str, Any]:
with path.open(newline="") as source:
rows = list(csv.DictReader(source))
if len(rows) != len(expected_shapes):
raise ValueError(f"request count mismatch: {len(rows)} != {len(expected_shapes)}")
rows.sort(key=lambda row: int(row["Request Id"]))
request_metrics = []
for index, (row, expected) in enumerate(zip(rows, expected_shapes, strict=True)):
if int(row["Request Id"]) != index:
raise ValueError("Frontier request ID/order drift")
shape = (
int(float(row["request_num_prefill_tokens"])),
int(float(row["request_num_decode_tokens"])),
)
if shape != expected:
raise ValueError(f"request shape drift at {index}: {shape} != {expected}")
ttft = float(row["ttft"])
e2e = float(row["request_e2e_time"])
tpot = (e2e - ttft) / (shape[1] - 1) if shape[1] > 1 else None
values = [ttft, e2e] + ([] if tpot is None else [tpot])
if not all(math.isfinite(value) and value >= 0 for value in values):
raise ValueError(f"invalid latency at request {index}")
request_metrics.append(
{
"request_id": index,
"input_tokens": shape[0],
"output_tokens": shape[1],
"ttft_ms": ttft,
"ttft_limit_ms": ttft_limit_ms(shape[0]),
"tpot_ms": tpot,
"e2e_ms": e2e,
}
)
slos = {}
for limit in TPOT_SLOS_MS:
passed = sum(
row["ttft_ms"] <= row["ttft_limit_ms"]
and (row["tpot_ms"] is None or row["tpot_ms"] <= limit)
for row in request_metrics
)
pass_rate = passed / len(request_metrics)
slos[f"tpot_{int(limit)}ms"] = {
"passed": passed,
"pass_rate": pass_rate,
"feasible": pass_rate >= TARGET_PASS_RATE,
}
ttfts = [float(row["ttft_ms"]) for row in request_metrics]
tpots = [float(row["tpot_ms"]) for row in request_metrics if row["tpot_ms"] is not None]
return {
"ttft_p50_ms": percentile(ttfts, 0.50),
"ttft_p95_ms": percentile(ttfts, 0.95),
"tpot_p50_ms": percentile(tpots, 0.50),
"tpot_p95_ms": percentile(tpots, 0.95),
"slos": slos,
}
def main() -> None:
args = parse_args()
for name in (
"frontier_source",
"replayserve_root",
"profile_root",
"python_deps",
"output_root",
):
setattr(args, name, getattr(args, name).resolve())
if args.allreduce_csv is not None:
args.allreduce_csv = args.allreduce_csv.resolve()
traces = [
parse_trace(specification, rate_contract=args.rate_contract)
for specification in args.trace
]
if len({trace["label"] for trace in traces}) != len(traces):
raise ValueError("trace labels must be unique")
if any(
right["offered_request_rate"] <= left["offered_request_rate"]
for left, right in zip(traces, traces[1:])
):
raise ValueError("trace anchors must be supplied in increasing load order")
selected = list(BASE.GRID)
if args.config:
wanted = set(args.config)
selected = [config for config in BASE.GRID if config.name in wanted]
if {config.name for config in selected} != wanted:
raise ValueError(f"unknown configs: {wanted - {config.name for config in selected}}")
paths = BASE.profile_paths(args.profile_root)
coverage = BASE.validate_profile(paths)
builder = BASE.load_module(
"qwen30_exact_trace_frontier_builder",
args.replayserve_root / "tools/run_frontier_sweep.py",
)
frontier_head = subprocess.run(
["git", "-C", str(args.frontier_source), "rev-parse", "HEAD"],
check=True,
text=True,
stdout=subprocess.PIPE,
).stdout.strip()
environment = os.environ.copy()
pythonpath = [str(args.python_deps), str(args.frontier_source)]
if environment.get("PYTHONPATH"):
pythonpath.append(environment["PYTHONPATH"])
environment.update(
{
"PYTHONPATH": ":".join(pythonpath),
"CUDA_VISIBLE_DEVICES": "",
"NVIDIA_VISIBLE_DEVICES": "void",
"WANDB_DISABLED": "true",
"VIDUR_DISABLE_WANDB": "1",
"FRONTIER_LOG_LEVEL": "WARNING",
"PYTHONDONTWRITEBYTECODE": "1",
}
)
config_results = []
for config in selected:
loads = []
config_knobs = BASE.knobs(config, paths, args.output_root / "cache")
config_knobs["enable_prefix_caching"] = args.prefix_caching
for trace in traces:
run_dir = args.output_root / "runs" / config.name / trace["label"]
result_path = run_dir / "result.json"
if args.resume and result_path.is_file():
loads.append(json.loads(result_path.read_text()))
continue
run_dir.mkdir(parents=True, exist_ok=True)
command = builder.build_frontier_command(
python_bin="/usr/bin/python3",
trace_file=trace["path"],
metrics_root=run_dir / "metrics",
run_id=f"qwen30_trace_{config.name}_{trace['label']}",
knobs=config_knobs,
)
command = BASE.configure_cc_command(
command,
backend=args.cc_backend,
allreduce_csv=args.allreduce_csv,
cache=args.output_root / "cc-cache",
)
BASE.write_json(run_dir / "command.json", command)
started = time.time()
with (run_dir / "stdout.log").open("w") as stdout, (
run_dir / "stderr.log"
).open("w") as stderr:
completed = subprocess.run(
command,
cwd=args.frontier_source,
env=environment,
stdout=stdout,
stderr=stderr,
timeout=args.timeout_seconds,
check=False,
)
if completed.returncode != 0:
raise RuntimeError(
f"Frontier failed for {config.name}/{trace['label']}: {completed.returncode}"
)
metrics = find_request_metrics(run_dir)
result = {
"status": "completed",
"config": {"tp": config.tp, "mns": config.mns, "name": config.name},
"trace_label": trace["label"],
"offered_request_rate": trace["offered_request_rate"],
"offered_request_rate_per_gpu": trace["offered_request_rate"] / config.tp,
"request_count": trace["requests"],
"elapsed_seconds": time.time() - started,
"trace_sha256": trace["sha256"],
"request_metrics_sha256": BASE.sha256(metrics),
"score": score(metrics, trace["shapes"]),
}
BASE.write_json(result_path, result)
loads.append(result)
print(
json.dumps(
{
"config": config.name,
"trace": trace["label"],
"rate": trace["offered_request_rate"],
"primary": result["score"]["slos"]["tpot_150ms"],
},
sort_keys=True,
),
flush=True,
)
config_results.append(
{
"config": {"tp": config.tp, "mns": config.mns, "name": config.name},
"loads": loads,
}
)
rankings = {}
for slo in (f"tpot_{int(value)}ms" for value in TPOT_SLOS_MS):
records = []
for item in config_results:
feasible = [
load["offered_request_rate"]
for load in item["loads"]
if load["score"]["slos"][slo]["feasible"]
]
capacity = max(feasible, default=None)
records.append(
{
"config": item["config"],
"maximum_tested_feasible_request_rate": capacity,
"maximum_tested_feasible_request_rate_per_gpu": (
capacity / item["config"]["tp"] if capacity is not None else None
),
"lower_censored": capacity is None,
"upper_censored": capacity == traces[-1]["offered_request_rate"],
}
)
records.sort(
key=lambda row: (
-(
row["maximum_tested_feasible_request_rate_per_gpu"]
if row["maximum_tested_feasible_request_rate_per_gpu"] is not None
else -1
),
row["config"]["name"],
)
)
rankings[slo] = records
manifest = {
"schema": "frontier-qwen30-exact-trace-surface-v1",
"status": "frozen_before_real" if selected == list(BASE.GRID) else "partial_not_decision_bearing",
"contract": {
"window_seconds": (
WINDOW_SECONDS if args.rate_contract == "trace-window" else None
),
"rate_contract": args.rate_contract,
"prefix_caching": args.prefix_caching,
"arrival": "original_trace_timestamp_and_order",
"input_output": "exact_source_values",
"ttft_slo": "1000ms + 1000ms * input_tokens / 8000",
"tpot_slos_ms": TPOT_SLOS_MS,
"primary_tpot_slo_ms": 150.0,
"target_pass_rate": TARGET_PASS_RATE,
},
"frontier": {
"source": str(args.frontier_source),
"git_head": frontier_head,
"git_status_short": subprocess.run(
["git", "-C", str(args.frontier_source), "status", "--short"],
check=True,
text=True,
stdout=subprocess.PIPE,
).stdout,
},
"profiles": {
"root": str(args.profile_root),
"coverage": coverage,
"sha256": {name: BASE.sha256(path) for name, path in paths.items()},
},
"collective": {
"backend": args.cc_backend,
"allreduce_csv": str(args.allreduce_csv) if args.allreduce_csv else None,
"allreduce_csv_sha256": BASE.sha256(args.allreduce_csv) if args.allreduce_csv else None,
},
"traces": [
{key: value for key, value in trace.items() if key != "shapes"}
for trace in traces
],
"config_results": config_results,
"rankings": rankings,
}
BASE.write_json(args.output_root / "frontier_trace_surface_frozen.json", manifest)
print(args.output_root / "frontier_trace_surface_frozen.json")
if __name__ == "__main__":
main()