256 lines
11 KiB
Python
256 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Exact C1 anchor replay using the pinned AITuner trace/worker/SLO paths."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import dataclasses
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
AITUNER_ROOT = Path(os.environ.get("AITUNER_ROOT", Path(__file__).resolve().parents[2]))
|
|
sys.path.insert(0, str(AITUNER_ROOT / "src"))
|
|
os.environ.setdefault("AITUNER_CODEX_BASE_URL", "http://127.0.0.1:1")
|
|
|
|
from aituner.slo import evaluate_request, summarize_evaluations # noqa: E402
|
|
from aituner.spec import load_study_spec # noqa: E402
|
|
from aituner.trace import load_trace_requests, select_requests_for_threshold # noqa: E402
|
|
from aituner.worker import _probe_drain_deadline, _replay_requests # noqa: E402
|
|
|
|
|
|
def atomic_json(path: Path, payload: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = path.with_suffix(path.suffix + ".tmp")
|
|
tmp.write_text(json.dumps(payload, sort_keys=True, indent=2) + "\n")
|
|
os.replace(tmp, path)
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
h = hashlib.sha256()
|
|
with path.open("rb") as f:
|
|
for chunk in iter(lambda: f.read(1 << 20), b""):
|
|
h.update(chunk)
|
|
return h.hexdigest()
|
|
|
|
|
|
def numeric(values: list[float | int | None]) -> dict[str, Any]:
|
|
finite = [float(x) for x in values if x is not None and math.isfinite(float(x))]
|
|
return {
|
|
"n": len(values), "finite_n": len(finite), "missing_n": len(values) - len(finite),
|
|
"min": min(finite) if finite else None, "max": max(finite) if finite else None,
|
|
"distinct_n": len(set(finite)),
|
|
}
|
|
|
|
|
|
def load_selected(study_path: Path, anchor: float):
|
|
study = load_study_spec(study_path)
|
|
window, requests = load_trace_requests(study, study_spec_path=study_path)
|
|
selected = select_requests_for_threshold(requests, threshold=anchor)
|
|
return study, window, requests, selected
|
|
|
|
|
|
def selected_summary(selected, duration_s: float, tp: int) -> dict[str, Any]:
|
|
ids = "\n".join(item.row_id for item in selected)
|
|
arrival = "\n".join(f"{item.arrival_s:.12f}" for item in selected)
|
|
lengths = "\n".join(str(item.prompt_tokens_hint) for item in selected)
|
|
return {
|
|
"count": len(selected),
|
|
"offered_req_s": len(selected) / duration_s,
|
|
"offered_req_s_per_gpu": len(selected) / duration_s / tp,
|
|
"request_id_order_sha256": hashlib.sha256(ids.encode()).hexdigest(),
|
|
"arrival_order_sha256": hashlib.sha256(arrival.encode()).hexdigest(),
|
|
"raw_length_order_sha256": hashlib.sha256(lengths.encode()).hexdigest(),
|
|
"arrival_s": numeric([item.arrival_s for item in selected]),
|
|
"raw_input_tokens": numeric([item.prompt_tokens_hint for item in selected]),
|
|
"long_gt4096": sum(int(item.prompt_tokens_hint or 0) > 4096 for item in selected),
|
|
}
|
|
|
|
|
|
def run_replay(args: argparse.Namespace, *, warmup: bool) -> dict[str, Any]:
|
|
study_path = Path(args.study)
|
|
study, window, _requests, selected = load_selected(study_path, args.anchor)
|
|
if warmup:
|
|
first = selected[:16]
|
|
if not any(int(item.prompt_tokens_hint or 0) > 4096 for item in first):
|
|
long_item = next(item for item in selected if int(item.prompt_tokens_hint or 0) > 4096)
|
|
first = [*selected[:15], long_item]
|
|
first = sorted({item.row_id: item for item in first}.values(), key=lambda item: item.arrival_s)
|
|
if len(first) < 16:
|
|
raise RuntimeError("warmup set has fewer than 16 unique requests")
|
|
start = first[0].arrival_s
|
|
selected = [dataclasses.replace(item, arrival_s=item.arrival_s - start) for item in first]
|
|
|
|
duration_s = float(window.window_end - window.window_start)
|
|
interval_start_mono_ns = time.monotonic_ns()
|
|
interval_start_wall_ns = time.time_ns()
|
|
outcomes, early_stopped, early_stop_reason = _replay_requests(
|
|
selected,
|
|
base_url=args.base_url,
|
|
timeout_s=study.engine.request_timeout_s,
|
|
max_concurrency=study.trace.max_concurrency,
|
|
target_pass_rate=(0.0 if warmup else study.slo.target_pass_rate),
|
|
max_lag_s=study.trace.early_stop_max_lag_s,
|
|
max_elapsed_s=(
|
|
120.0 if warmup else _probe_drain_deadline(
|
|
selected, study.slo, ceiling=study.trace.early_stop_max_elapsed_s
|
|
)
|
|
),
|
|
evaluate_outcome=lambda outcome: evaluate_request(outcome, study.slo),
|
|
drain_inflight_on_early_stop=True,
|
|
)
|
|
interval_end_mono_ns = time.monotonic_ns()
|
|
interval_end_wall_ns = time.time_ns()
|
|
evaluations, slo_summary = summarize_evaluations(outcomes, study.slo)
|
|
by_id = {item.row_id: item for item in selected}
|
|
details = []
|
|
for outcome, evaluation in zip(outcomes, evaluations):
|
|
request = by_id[outcome.request_id]
|
|
details.append({
|
|
"request_id": outcome.request_id,
|
|
"sampling_u": request.sampling_u,
|
|
"arrival_s": request.arrival_s,
|
|
"raw_input_tokens": request.prompt_tokens_hint,
|
|
"success": outcome.success,
|
|
"ttft_ms": outcome.ttft_ms,
|
|
"tpot_ms": outcome.tpot_ms,
|
|
"completion_tokens": outcome.completion_tokens,
|
|
"completion_tokens_source": outcome.completion_tokens_source,
|
|
"completed_mono_ns": outcome.completed_mono_ns,
|
|
"completed_elapsed_s": (
|
|
(outcome.completed_mono_ns - interval_start_mono_ns) / 1e9
|
|
if outcome.completed_mono_ns is not None else None
|
|
),
|
|
"slo_pass": evaluation.passed,
|
|
"reasons": evaluation.reasons,
|
|
"error": outcome.error,
|
|
})
|
|
out = Path(args.result_dir)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
with (out / "requests.jsonl").open("w") as f:
|
|
for item in details:
|
|
f.write(json.dumps(item, sort_keys=True) + "\n")
|
|
summary = selected_summary(selected, duration_s, args.tp)
|
|
exact = sum(
|
|
item.success and item.completion_tokens_source == "usage" and item.completion_tokens == 128
|
|
for item in outcomes
|
|
)
|
|
result = {
|
|
"schema": 1,
|
|
"kind": "warmup" if warmup else "anchor",
|
|
"cell": args.cell,
|
|
"anchor": args.anchor,
|
|
"tp": args.tp,
|
|
"mns": args.mns,
|
|
"study_sha256": sha256_file(study_path),
|
|
"interval": {
|
|
"start_mono_ns": interval_start_mono_ns, "end_mono_ns": interval_end_mono_ns,
|
|
"start_wall_ns": interval_start_wall_ns, "end_wall_ns": interval_end_wall_ns,
|
|
"elapsed_s": (interval_end_mono_ns - interval_start_mono_ns) / 1e9,
|
|
},
|
|
"selection": summary,
|
|
"observed_count": len(outcomes),
|
|
"exact_output_count": exact,
|
|
"slo_pass_count": slo_summary["slo_pass_count"],
|
|
"pass_rate": slo_summary["slo_pass_rate"],
|
|
"feasible": bool(slo_summary["feasible"]),
|
|
"early_stopped": early_stopped,
|
|
"early_stop_reason": early_stop_reason,
|
|
"ttft_ms": numeric([item.ttft_ms for item in outcomes]),
|
|
"tpot_ms": numeric([item.tpot_ms for item in outcomes]),
|
|
"invariants": {
|
|
"selected_nonempty": bool(selected),
|
|
"outcomes_cover_selected": len(outcomes) == len(selected),
|
|
"exact_output_or_failed": all(
|
|
(not item.success) or (
|
|
item.completion_tokens_source == "usage" and item.completion_tokens == 128
|
|
) for item in outcomes
|
|
),
|
|
"raw_lengths_present": all(item.prompt_tokens_hint is not None for item in selected),
|
|
"arrival_nondecreasing": all(
|
|
b.arrival_s >= a.arrival_s for a, b in zip(selected, selected[1:])
|
|
),
|
|
"warmup_16": (len(outcomes) >= 16 if warmup else True),
|
|
"warmup_exact_16": (exact >= 16 if warmup else True),
|
|
"warmup_long": (
|
|
any(int(item.prompt_tokens_hint or 0) > 4096 for item in selected)
|
|
if warmup else True
|
|
),
|
|
},
|
|
}
|
|
atomic_json(out / "result.json", result)
|
|
print(json.dumps({k: result[k] for k in ("cell", "anchor", "kind", "pass_rate", "feasible")}))
|
|
if not all(result["invariants"].values()):
|
|
raise RuntimeError(f"client invariants failed: {result['invariants']}")
|
|
return result
|
|
|
|
|
|
def preflight(args: argparse.Namespace) -> None:
|
|
ground = json.loads(Path(args.ground_truth).read_text())
|
|
studies = {1: Path(args.primary_study), 2: Path(args.primary_study), 4: Path(args.tp4_study)}
|
|
loaded = {}
|
|
mismatches = []
|
|
values = []
|
|
for cell in ground["cells"]:
|
|
tp = int(cell["tensor_parallel_size"])
|
|
if tp not in loaded:
|
|
_study, _window, requests, _selected = load_selected(studies[tp], 0.0)
|
|
loaded[tp] = requests
|
|
for historical in cell["probe_history"]:
|
|
selected = select_requests_for_threshold(
|
|
loaded[tp], threshold=float(historical["sampling_u"])
|
|
)
|
|
values.append(len(selected))
|
|
if len(selected) != int(historical["request_count"]):
|
|
mismatches.append({
|
|
"cell": cell["cell_id"], "anchor": historical["sampling_u"],
|
|
"expected": historical["request_count"], "actual": len(selected),
|
|
})
|
|
result = {
|
|
"schema": 1, "observations": len(values), "mismatches": mismatches,
|
|
"request_counts": numeric(values),
|
|
"invariants": {"observations_92": len(values) == 92, "counts_match": not mismatches},
|
|
}
|
|
atomic_json(Path(args.out), result)
|
|
print(json.dumps(result, sort_keys=True))
|
|
if not all(result["invariants"].values()):
|
|
raise RuntimeError("preflight count reconstruction failed")
|
|
|
|
|
|
def parser() -> argparse.ArgumentParser:
|
|
p = argparse.ArgumentParser()
|
|
sub = p.add_subparsers(dest="command", required=True)
|
|
pf = sub.add_parser("preflight")
|
|
pf.add_argument("--ground-truth", required=True)
|
|
pf.add_argument("--primary-study", required=True)
|
|
pf.add_argument("--tp4-study", required=True)
|
|
pf.add_argument("--out", required=True)
|
|
for name in ("warmup", "run-anchor"):
|
|
q = sub.add_parser(name)
|
|
q.add_argument("--study", required=True)
|
|
q.add_argument("--cell", required=True)
|
|
q.add_argument("--anchor", type=float, required=True)
|
|
q.add_argument("--tp", type=int, required=True)
|
|
q.add_argument("--mns", type=int, required=True)
|
|
q.add_argument("--base-url", required=True)
|
|
q.add_argument("--result-dir", required=True)
|
|
return p
|
|
|
|
|
|
def main() -> None:
|
|
args = parser().parse_args()
|
|
if args.command == "preflight":
|
|
preflight(args)
|
|
else:
|
|
run_replay(args, warmup=args.command == "warmup")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|