Files
aituner/runs/frontier-multicase-sufficiency-v0/best_effort/fixed_cohort_rank.py

1406 lines
54 KiB
Python

#!/usr/bin/env python3
"""Blind Frontier/community-vLLM rank comparison on one fixed prompt cohort.
The earlier grid selected a different ``sampling_u`` subset at every load. That
confounds offered load with prompt mix and makes binary-search monotonicity an
untested assumption. This runner selects one length-stratified cohort, changes
load only by scaling its complete arrival timeline, evaluates every discrete
load, freezes Frontier first, and only then permits the community-vLLM run.
"""
from __future__ import annotations
import argparse
import csv
import hashlib
import json
import math
import os
import random
import statistics
import subprocess
import sys
import time
from dataclasses import asdict, replace
from pathlib import Path
from typing import Any, Iterable
SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SCRIPT_DIR.parents[2]
for search_path in (SCRIPT_DIR, REPO_ROOT / "src"):
if str(search_path) not in sys.path:
sys.path.insert(0, str(search_path))
import community_prefill_grid as community_grid # noqa: E402
import frontier_prefill_grid as frontier_grid # noqa: E402
from aituner.engine import build_launch_recipe # noqa: E402
from aituner.spec import ConfigPatch, load_study_spec # noqa: E402
from aituner.trace import TraceRequest, load_trace_requests # noqa: E402
from aituner.worker import ( # noqa: E402
_ignore_sigterm_if_main,
_install_sigterm_as_keyboardinterrupt,
_replay_requests,
_restore_sigterm,
_terminate_process_tree,
_wait_for_server_or_exit,
)
PROTOCOL_SCHEMA = "qwen235b-prefill-fixed-cohort-protocol-v1"
FRONTIER_SCHEMA = "frontier-qwen235b-prefill-fixed-cohort-v1"
COMMUNITY_SCHEMA = "community-qwen235b-prefill-fixed-cohort-v1"
COMPARISON_SCHEMA = "frontier-community-qwen235b-rank-comparison-v1"
COHORT_SIZE = 64
COHORT_SEED = 2026071501
CONFIG_ORDER_SEED = 2026071502
RATE_ORDER_SEED = 2026071503
INPUT_BINS = (0, 1024, 2048, 4096, 8192, 16384, 32769)
OFFERED_RATES = (0.15, 0.25, 0.35, 0.50, 0.75, 1.00, 1.50)
TARGET_PASS_RATE = 0.95
EXPECTED_QUANT_SIGNATURE = "method=fp8|act=dynamic|serialized=True|block=128x128"
SLO_VARIANTS: dict[str, dict[str, Any]] = {
"linear_8k_primary": {
"description": "TTFT <= 1000 ms + input_tokens / 8000 tokens/s",
"kind": "linear_ms",
"intercept_ms": 1000.0,
"per_token_ms": 0.125,
},
"linear_6k": {
"description": "TTFT <= 1000 ms + input_tokens / 6000 tokens/s",
"kind": "linear_ms",
"intercept_ms": 1000.0,
"per_token_ms": 1.0 / 6.0,
},
"linear_10k": {
"description": "TTFT <= 1000 ms + input_tokens / 10000 tokens/s",
"kind": "linear_ms",
"intercept_ms": 1000.0,
"per_token_ms": 0.1,
},
"legacy_step_1s_2s": {
"description": "Original strict 1 s/2 s step SLO (null-capacity sensitivity)",
"kind": "step_ms",
"buckets": [
{"max_input_tokens": 8191, "threshold_ms": 1000.0},
{"threshold_ms": 2000.0},
],
},
}
PRIMARY_SLO = "linear_8k_primary"
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for chunk in iter(lambda: source.read(1 << 20), b""):
digest.update(chunk)
return digest.hexdigest()
def sha256_text(value: str) -> str:
return hashlib.sha256(value.encode()).hexdigest()
def order_hash(values: Iterable[object]) -> str:
return sha256_text("\n".join(str(value) for value in values))
def write_json(path: Path, payload: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
os.replace(temporary, path)
def rate_key(rate: float) -> str:
return f"r_{rate:.3f}".replace(".", "p")
def _percentile(values: list[float], fraction: float) -> float | None:
if not values:
return None
ordered = sorted(values)
index = min(len(ordered) - 1, max(0, math.ceil(fraction * len(ordered)) - 1))
return float(ordered[index])
def _latency_summary(values: list[float]) -> dict[str, Any]:
return {
"count": len(values),
"mean": statistics.fmean(values) if values else None,
"p50": _percentile(values, 0.50),
"p95": _percentile(values, 0.95),
"p99": _percentile(values, 0.99),
"max": max(values) if values else None,
}
def slo_threshold_ms(variant: dict[str, Any], input_tokens: int) -> float:
if variant["kind"] == "linear_ms":
return float(variant["intercept_ms"]) + float(variant["per_token_ms"]) * input_tokens
if variant["kind"] != "step_ms":
raise ValueError(f"unsupported SLO kind: {variant['kind']}")
for bucket in variant["buckets"]:
ceiling = bucket.get("max_input_tokens")
if ceiling is None or input_tokens <= int(ceiling):
return float(bucket["threshold_ms"])
raise AssertionError("step SLO must have a terminal bucket")
def _input_bin(input_tokens: int) -> int:
for index, (low, high) in enumerate(zip(INPUT_BINS, INPUT_BINS[1:])):
if low <= input_tokens < high:
return index
raise ValueError(f"input length outside protocol bins: {input_tokens}")
def _allocate_quotas(counts: list[int], cohort_size: int) -> list[int]:
total = sum(counts)
if cohort_size <= 0 or cohort_size > total:
raise ValueError(f"invalid cohort size {cohort_size} for {total} rows")
exact = [cohort_size * count / total for count in counts]
quotas = [min(count, math.floor(value)) for count, value in zip(counts, exact)]
remaining = cohort_size - sum(quotas)
order = sorted(
range(len(counts)),
key=lambda index: (-(exact[index] - math.floor(exact[index])), index),
)
while remaining:
advanced = False
for index in order:
if quotas[index] >= counts[index]:
continue
quotas[index] += 1
remaining -= 1
advanced = True
if not remaining:
break
if not advanced:
raise AssertionError("could not allocate complete cohort")
return quotas
def load_source_rows(trace: Path) -> list[dict[str, Any]]:
rows = []
with trace.open(encoding="utf-8") as source:
for source_index, line in enumerate(source):
if not line.strip():
continue
raw = json.loads(line)
input_tokens = int(raw["input_length"])
if not 0 <= input_tokens <= 32768:
continue
prompt = raw.get("prompt")
if not isinstance(prompt, str) or not prompt:
raise ValueError(f"row {source_index} lacks raw prompt")
rows.append(
{
"source_row_index": source_index,
"source_request_id": str(
raw.get("request_id") or raw.get("id") or source_index
),
"source_arrival_s": float(raw["timestamp"]),
"sampling_u": float(raw["sampling_u"]),
"input_tokens": input_tokens,
"prompt_sha256": sha256_text(prompt),
"input_bin": _input_bin(input_tokens),
}
)
return rows
def select_cohort(
rows: list[dict[str, Any]], *, cohort_size: int, seed: int
) -> tuple[list[dict[str, Any]], list[int]]:
by_bin = [[] for _ in range(len(INPUT_BINS) - 1)]
for row in rows:
by_bin[int(row["input_bin"])].append(row)
quotas = _allocate_quotas([len(items) for items in by_bin], cohort_size)
selected = []
for bin_index, (items, quota) in enumerate(zip(by_bin, quotas)):
ranked = sorted(
items,
key=lambda row: sha256_text(
f"{seed}|{bin_index}|{row['source_row_index']}|{row['prompt_sha256']}"
),
)
selected.extend(ranked[:quota])
selected.sort(key=lambda row: (row["source_arrival_s"], row["source_row_index"]))
if len(selected) != cohort_size:
raise AssertionError("cohort selection returned wrong size")
return selected, quotas
def prepare_protocol(args: argparse.Namespace) -> Path:
trace = args.trace.resolve()
actual_sha = sha256(trace)
if args.expected_trace_sha256 and actual_sha != args.expected_trace_sha256:
raise ValueError(
f"trace SHA256 mismatch: expected={args.expected_trace_sha256}, actual={actual_sha}"
)
rows = load_source_rows(trace)
cohort, quotas = select_cohort(rows, cohort_size=args.cohort_size, seed=args.seed)
first_arrival = float(cohort[0]["source_arrival_s"])
source_span = float(cohort[-1]["source_arrival_s"]) - first_arrival
if source_span <= 0:
raise ValueError("selected cohort has no arrival span")
output_root = args.output_root.resolve()
trace_dir = output_root / "traces"
trace_dir.mkdir(parents=True, exist_ok=True)
fields = [
"arrived_at",
"num_prefill_tokens",
"num_decode_tokens",
"source_row_index",
"source_request_id",
"source_arrival_s",
"sampling_u",
"slo_ttft_ms",
]
rates = {}
for rate in args.rate:
duration_s = len(cohort) / rate
target = trace_dir / f"{rate_key(rate)}.csv"
with target.open("w", encoding="utf-8", newline="") as output:
writer = csv.DictWriter(output, fieldnames=fields)
writer.writeheader()
for row in cohort:
arrival_s = (
(float(row["source_arrival_s"]) - first_arrival)
/ source_span
* duration_s
)
writer.writerow(
{
"arrived_at": f"{arrival_s:.12f}",
"num_prefill_tokens": row["input_tokens"],
"num_decode_tokens": 1,
"source_row_index": row["source_row_index"],
"source_request_id": row["source_request_id"],
"source_arrival_s": f"{row['source_arrival_s']:.12f}",
"sampling_u": f"{row['sampling_u']:.12f}",
"slo_ttft_ms": f"{slo_threshold_ms(SLO_VARIANTS[PRIMARY_SLO], int(row['input_tokens'])):.6f}",
}
)
rates[rate_key(rate)] = {
"offered_request_rate": rate,
"duration_s": duration_s,
"path": str(target.resolve()),
"sha256": sha256(target),
"request_count": len(cohort),
"request_rate_recomputed": len(cohort) / duration_s,
"source_row_order_sha256": order_hash(
row["source_row_index"] for row in cohort
),
"input_length_order_sha256": order_hash(
row["input_tokens"] for row in cohort
),
}
bin_counts = [0] * (len(INPUT_BINS) - 1)
for row in rows:
bin_counts[int(row["input_bin"])] += 1
manifest = {
"schema": PROTOCOL_SCHEMA,
"created_unix_s": time.time(),
"source": {
"path": str(trace),
"sha256": actual_sha,
"eligible_request_count": len(rows),
},
"selection": {
"method": "length_stratified_smallest_sha256",
"seed": args.seed,
"cohort_size": len(cohort),
"input_bin_edges": list(INPUT_BINS),
"source_bin_counts": bin_counts,
"cohort_bin_quotas": quotas,
"stable_order": ["source_arrival_s", "source_row_index"],
"cohort_source_row_order_sha256": order_hash(
row["source_row_index"] for row in cohort
),
"source_arrival_span_s": source_span,
},
"load_contract": {
"only_mutated_variable": "uniform_arrival_timeline_scale",
"request_count_per_point": len(cohort),
"completion_tokens_override": 1,
"target_pass_rate": TARGET_PASS_RATE,
"binary_search": False,
"monotonicity_assumed": False,
"offered_rates": list(args.rate),
},
"slo_variants": SLO_VARIANTS,
"primary_slo": PRIMARY_SLO,
"cohort": cohort,
"rates": rates,
}
manifest_path = output_root / "protocol_manifest.json"
write_json(manifest_path, manifest)
audit_protocol(manifest_path)
print(manifest_path)
return manifest_path
def audit_protocol(manifest_path: Path) -> dict[str, Any]:
manifest = json.loads(manifest_path.read_text())
if manifest.get("schema") != PROTOCOL_SCHEMA:
raise ValueError(f"unexpected protocol schema: {manifest.get('schema')}")
expected_ids = [str(row["source_row_index"]) for row in manifest["cohort"]]
expected_lengths = [int(row["input_tokens"]) for row in manifest["cohort"]]
rate_checks = []
for record in manifest["rates"].values():
path = Path(record["path"])
if sha256(path) != record["sha256"]:
raise ValueError(f"trace hash mismatch: {path}")
with path.open(newline="") as source:
rows = list(csv.DictReader(source))
ids = [row["source_row_index"] for row in rows]
lengths = [int(row["num_prefill_tokens"]) for row in rows]
arrivals = [float(row["arrived_at"]) for row in rows]
if ids != expected_ids or lengths != expected_lengths:
raise ValueError(f"cohort/order drift in {path}")
if arrivals != sorted(arrivals) or abs(arrivals[0]) > 1e-9:
raise ValueError(f"invalid arrival sequence in {path}")
recomputed = len(rows) / (arrivals[-1] - arrivals[0])
if not math.isclose(
recomputed, float(record["offered_request_rate"]), rel_tol=1e-9
):
raise ValueError(f"offered-rate mismatch in {path}: {recomputed}")
rate_checks.append(
{"rate": record["offered_request_rate"], "recomputed_rate": recomputed}
)
result = {
"status": "passed",
"protocol_manifest_sha256": sha256(manifest_path),
"cohort_size": len(expected_ids),
"rate_checks": rate_checks,
}
write_json(manifest_path.parent / "protocol_audit.json", result)
return result
def _read_trace(path: Path) -> list[dict[str, Any]]:
with path.open(encoding="utf-8", newline="") as source:
return list(csv.DictReader(source))
def score_requests(requests: list[dict[str, Any]]) -> dict[str, Any]:
scores = {}
for name, variant in SLO_VARIANTS.items():
passed = 0
for request in requests:
threshold = slo_threshold_ms(variant, int(request["input_tokens"]))
passed += int(
bool(request["success"])
and request.get("ttft_ms") is not None
and float(request["ttft_ms"]) <= threshold
)
count = len(requests)
pass_rate = passed / count if count else 0.0
scores[name] = {
"request_count": count,
"passed_request_count": passed,
"slo_pass_rate": pass_rate,
"target_pass_rate": TARGET_PASS_RATE,
"feasible": pass_rate >= TARGET_PASS_RATE,
}
ttfts = [
float(request["ttft_ms"])
for request in requests
if request.get("ttft_ms") is not None
]
return {"scores": scores, "ttft_ms": _latency_summary(ttfts)}
def _capacity_record(
results: list[dict[str, Any]], *, config: dict[str, Any], slo_name: str
) -> dict[str, Any]:
ordered = sorted(results, key=lambda item: float(item["offered_request_rate"]))
feasible = [
float(item["offered_request_rate"])
for item in ordered
if item["scores"][slo_name]["feasible"]
]
vector = [bool(item["scores"][slo_name]["feasible"]) for item in ordered]
violations = []
for lower_index in range(len(vector)):
for higher_index in range(lower_index + 1, len(vector)):
if not vector[lower_index] and vector[higher_index]:
violations.append(
[
float(ordered[lower_index]["offered_request_rate"]),
float(ordered[higher_index]["offered_request_rate"]),
]
)
capacity = max(feasible) if feasible else None
tp = int(config["tp"])
return {
"config": config,
"slo": slo_name,
"tested_rates": [float(item["offered_request_rate"]) for item in ordered],
"pass_rates": [float(item["scores"][slo_name]["slo_pass_rate"]) for item in ordered],
"feasibility": vector,
"maximum_tested_feasible_request_rate": capacity,
"maximum_tested_feasible_request_rate_per_gpu": (
capacity / tp if capacity is not None else None
),
"lower_censored": capacity is None,
"upper_censored": bool(feasible) and capacity == max(item["offered_request_rate"] for item in ordered),
"monotonicity_violations": violations,
}
def rank_surface(
config_results: list[dict[str, Any]], *, slo_name: str
) -> list[dict[str, Any]]:
records = [
_capacity_record(item["loads"], config=item["config"], slo_name=slo_name)
for item in config_results
]
records.sort(
key=lambda item: (
-(
item["maximum_tested_feasible_request_rate_per_gpu"]
if item["maximum_tested_feasible_request_rate_per_gpu"] is not None
else -1.0
),
item["config"]["name"],
)
)
previous_value = object()
previous_rank = 0
for index, record in enumerate(records, start=1):
value = record["maximum_tested_feasible_request_rate_per_gpu"]
if value != previous_value:
previous_rank = index
previous_value = value
record["rank"] = previous_rank
return records
def _validate_profile_contract(profile_paths: dict[str, Path]) -> dict[str, Any]:
payload = json.loads(profile_paths["manifest"].read_text())
contract = payload.get("contract") or {}
if contract.get("quant_signature") != EXPECTED_QUANT_SIGNATURE:
raise ValueError(f"unexpected FP8 profile contract: {contract}")
if contract.get("profiling_precision") != "BF16":
raise ValueError(f"unexpected output/accumulation precision: {contract}")
for name, path in (
("linear_op.csv", profile_paths["linear"]),
("attention.csv", profile_paths["attention"]),
("moe.csv", profile_paths["moe"]),
):
expected = payload["outputs"][name]["sha256"]
if sha256(path) != expected:
raise ValueError(f"profile output hash mismatch: {path}")
return contract
def run_frontier_load(
*,
args: argparse.Namespace,
profile_paths: dict[str, Path],
config: Any,
rate_record: dict[str, Any],
) -> dict[str, Any]:
rate = float(rate_record["offered_request_rate"])
load_dir = args.output_root / "runs" / config.name / rate_key(rate)
result_path = load_dir / "result.json"
if result_path.is_file():
result = json.loads(result_path.read_text())
if result.get("status") == "completed":
return result
load_dir.mkdir(parents=True, exist_ok=True)
trace = Path(rate_record["path"])
command = frontier_grid.build_command(
python=args.python,
frontier_source=args.frontier_source,
profile_root=args.profile_root,
profile_paths=profile_paths,
trace=trace,
config=config,
probe_dir=load_dir,
run_id=f"{config.name}_{rate_key(rate)}",
cache_root=args.output_root / "cache",
)
write_json(load_dir / "command.json", command)
environment = os.environ.copy()
environment.update(
{
"PYTHONPATH": str(args.frontier_source),
"WANDB_DISABLED": "true",
"VIDUR_DISABLE_WANDB": "1",
}
)
started = time.time()
with (load_dir / "stdout.log").open("w", encoding="utf-8") as output:
completed = subprocess.run(
command,
cwd=args.frontier_source,
env=environment,
stdout=output,
stderr=subprocess.STDOUT,
check=False,
)
if completed.returncode != 0:
raise RuntimeError(
f"Frontier failed: config={config.name}, rate={rate}, rc={completed.returncode}"
)
request_metrics = frontier_grid.find_request_metrics(load_dir)
trace_rows = _read_trace(trace)
with request_metrics.open(encoding="utf-8", newline="") as source:
metric_rows = list(csv.DictReader(source))
metrics_by_id = {int(row["Request Id"]): row for row in metric_rows}
if set(metrics_by_id) != set(range(len(trace_rows))):
raise ValueError(f"Frontier Request Id mismatch: {request_metrics}")
requests = [
{
"request_id": row["source_row_index"],
"input_tokens": int(row["num_prefill_tokens"]),
"success": True,
"ttft_ms": float(metrics_by_id[index]["ttft"]),
}
for index, row in enumerate(trace_rows)
]
score = score_requests(requests)
result = {
"status": "completed",
"config": asdict(config) | {"name": config.name},
"offered_request_rate": rate,
"request_rate_per_gpu": rate / config.gpu_count,
"elapsed_seconds": time.time() - started,
"trace_path": str(trace),
"trace_sha256": sha256(trace),
"request_metrics_path": str(request_metrics),
"request_metrics_sha256": sha256(request_metrics),
"requests": requests,
**score,
}
write_json(result_path, result)
return result
def run_frontier(args: argparse.Namespace) -> Path:
protocol = json.loads(args.protocol_manifest.read_text())
if protocol.get("schema") != PROTOCOL_SCHEMA:
raise ValueError("invalid fixed-cohort protocol manifest")
audit_protocol(args.protocol_manifest)
profile_paths = frontier_grid.resolve_profile_paths(args.profile_root)
profile_contract = _validate_profile_contract(profile_paths)
configs = frontier_grid.selected_configs(args.config)
args.output_root.mkdir(parents=True, exist_ok=True)
run_manifest = {
"schema": FRONTIER_SCHEMA,
"created_unix_s": time.time(),
"protocol_manifest": {
"path": str(args.protocol_manifest.resolve()),
"sha256": sha256(args.protocol_manifest),
},
"frontier": {
"source": str(args.frontier_source.resolve()),
"python": str(args.python.resolve()),
"fingerprint": frontier_grid.frontier_source_fingerprint(
args.frontier_source, args.frontier_commit
),
},
"profiles": {
"contract": profile_contract,
"files": {
name: {"path": str(path.resolve()), "sha256": sha256(path)}
for name, path in profile_paths.items()
},
},
"kv_capacity_evidence": {
"tp4": {
"path": str(args.tp4_capacity_artifact.resolve()),
"sha256": sha256(args.tp4_capacity_artifact),
"num_gpu_blocks": 26101,
},
"tp8": {
"path": str(args.tp8_capacity_artifact.resolve()),
"sha256": sha256(args.tp8_capacity_artifact),
"num_gpu_blocks": 62351,
},
},
"best_effort_contract": {
"arrival_semantics": "preserve_materialized_trace_arrivals",
"quantization": "block_fp8_w8a8_dynamic_with_bf16_output_accumulation",
"execution_profiles": "community_vllm_0.10.2_serving_entrypoints",
"moe_ep_prefill": "global_routing_then_slowest_local_ep_lane",
"moe_ep_critical_lane_cache": "exact_(total_routed_tokens,layer_id)_memoization",
"kv_blocks": "measured_community_vllm_and_fixed_per_topology",
"cpu_overhead_modeling": "disabled_no_native_community_vllm_records",
"end_to_end_action_specific_calibration": False,
"real_serving_results_observed_before_freeze": False,
},
"configs": [asdict(config) | {"name": config.name} for config in configs],
}
write_json(args.output_root / "run_manifest.json", run_manifest)
config_results = []
rate_records = sorted(
protocol["rates"].values(), key=lambda item: item["offered_request_rate"]
)
for config in configs:
loads = []
for record in rate_records:
result = run_frontier_load(
args=args,
profile_paths=profile_paths,
config=config,
rate_record=record,
)
loads.append(result)
print(
json.dumps(
{
"system": "Frontier",
"config": config.name,
"rate": result["offered_request_rate"],
"pass_rate": result["scores"][PRIMARY_SLO]["slo_pass_rate"],
"feasible": result["scores"][PRIMARY_SLO]["feasible"],
"elapsed_seconds": result["elapsed_seconds"],
},
sort_keys=True,
),
flush=True,
)
summary = {"config": asdict(config) | {"name": config.name}, "loads": loads}
write_json(args.output_root / "results" / f"{config.name}.json", summary)
config_results.append(summary)
freeze = {
"schema": FRONTIER_SCHEMA,
"status": "frozen_before_community_run",
"created_unix_s": time.time(),
"run_manifest_sha256": sha256(args.output_root / "run_manifest.json"),
"protocol_manifest_sha256": sha256(args.protocol_manifest),
"primary_slo": PRIMARY_SLO,
"rankings": {
name: rank_surface(config_results, slo_name=name) for name in SLO_VARIANTS
},
"config_results": config_results,
}
path = args.output_root / "frontier_ranking_frozen.json"
write_json(path, freeze)
print(path)
return path
def _community_study_payload(
*,
tp: int,
repo: Path,
python: Path,
vllm: Path,
model: Path,
trace: Path,
windows: Path,
port: int,
) -> dict[str, Any]:
payload = community_grid.study_payload(
tp=tp,
repo=repo,
python=python,
vllm=vllm,
model=model,
trace=trace,
windows=windows,
port=port,
)
payload["study_id"] = f"community-qwen235b-prefill-fixed-cohort-tp{tp}-v1"
payload["slo"] = {
"target_pass_rate": TARGET_PASS_RATE,
"ttft_rule": {
key: value
for key, value in SLO_VARIANTS[PRIMARY_SLO].items()
if key != "description"
},
}
# The dedicated runner supplies exact fixed-cohort arrivals and disables
# SLO-based early stopping; these trace values only define raw prompt loading.
payload["trace"]["early_stop_max_lag_s"] = None
payload["trace"]["early_stop_max_elapsed_s"] = None
return payload
def _validate_frontier_freeze(path: Path, protocol_manifest: Path) -> dict[str, Any]:
freeze = json.loads(path.read_text())
if freeze.get("schema") != FRONTIER_SCHEMA:
raise ValueError(f"unexpected Frontier freeze schema: {freeze.get('schema')}")
if freeze.get("status") != "frozen_before_community_run":
raise ValueError("Frontier result is not marked frozen")
if freeze.get("protocol_manifest_sha256") != sha256(protocol_manifest):
raise ValueError("Frontier and community protocol manifests differ")
if len(freeze.get("config_results") or []) != len(frontier_grid.GRID):
raise ValueError("Frontier freeze is incomplete")
expected_rates = set(OFFERED_RATES)
for item in freeze["config_results"]:
rates = {float(load["offered_request_rate"]) for load in item["loads"]}
if rates != expected_rates:
raise ValueError(f"Frontier freeze lacks rates for {item['config']['name']}")
return freeze
def prepare_community(args: argparse.Namespace) -> Path:
protocol = json.loads(args.protocol_manifest.read_text())
if protocol.get("schema") != PROTOCOL_SCHEMA:
raise ValueError("invalid fixed-cohort protocol manifest")
_validate_frontier_freeze(args.frontier_freeze, args.protocol_manifest)
repo = args.repo.resolve()
python = args.python.resolve()
vllm = args.vllm.resolve()
model = args.model.resolve()
windows = args.windows.resolve()
trace = Path(protocol["source"]["path"])
for path in (repo, python, vllm, model, windows, trace):
if not path.exists():
raise FileNotFoundError(path)
if sha256(trace) != protocol["source"]["sha256"]:
raise ValueError("raw prompt trace changed after protocol freeze")
output_root = args.output_root.resolve()
output_root.mkdir(parents=True, exist_ok=True)
studies = {}
for tp in (4, 8):
path = output_root / "studies" / f"tp{tp}" / "study.json"
write_json(
path,
_community_study_payload(
tp=tp,
repo=repo,
python=python,
vllm=vllm,
model=model,
trace=trace,
windows=windows,
port=args.port,
),
)
studies[str(tp)] = {"path": str(path), "sha256": sha256(path)}
configs = list(community_grid.GRID)
random.Random(args.config_order_seed).shuffle(configs)
records = []
for execution_index, config in enumerate(configs, start=1):
rates = list(OFFERED_RATES)
random.Random(f"{args.rate_order_seed}|{config.name}").shuffle(rates)
records.append(
{
"execution_index": execution_index,
"config": asdict(config) | {"name": config.name},
"rate_order": rates,
"study_path": studies[str(config.tp)]["path"],
"result_path": str(output_root / "results" / f"{config.name}.json"),
"engine_log_path": str(output_root / "runs" / config.name / "engine.log"),
}
)
total_replay_seconds = sum(
float(record["duration_s"]) for record in protocol["rates"].values()
) * len(configs)
manifest = {
"schema": COMMUNITY_SCHEMA,
"created_unix_s": time.time(),
"repository": community_grid.git_fingerprint(repo),
"protocol_manifest": {
"path": str(args.protocol_manifest.resolve()),
"sha256": sha256(args.protocol_manifest),
},
"frontier_freeze": {
"path": str(args.frontier_freeze.resolve()),
"sha256": sha256(args.frontier_freeze),
},
"runtime": {"python": str(python), "vllm": str(vllm)},
"model": {
"path": str(model),
"config_sha256": sha256(model / "config.json"),
},
"studies": studies,
"execution_contract": {
"request_mode": "raw_completion",
"completion_tokens_override": 1,
"prefix_caching": False,
"client_max_concurrency": community_grid.CLIENT_MAX_CONCURRENCY,
"slo_early_stop": False,
"one_server_launch_per_config": True,
"rate_order_randomized_within_config": True,
"config_order_seed": args.config_order_seed,
"rate_order_seed": args.rate_order_seed,
"estimated_replay_wall_seconds_excluding_8_model_loads": total_replay_seconds,
},
"configs": records,
}
path = output_root / "run_manifest.json"
write_json(path, manifest)
print(path)
return path
def _materialize_real_requests(
*,
all_requests: list[TraceRequest],
rate_record: dict[str, Any],
cohort: list[dict[str, Any]],
) -> list[TraceRequest]:
by_id = {request.row_id: request for request in all_requests}
expected_prompt_hash = {
str(row["source_row_index"]): row["prompt_sha256"] for row in cohort
}
trace_rows = _read_trace(Path(rate_record["path"]))
requests = []
for row in trace_rows:
request_id = row["source_row_index"]
if request_id not in by_id:
raise ValueError(f"cohort request is absent from raw trace loader: {request_id}")
request = by_id[request_id]
prompt = request.body.get("prompt")
if not isinstance(prompt, str) or sha256_text(prompt) != expected_prompt_hash[request_id]:
raise ValueError(f"raw prompt fingerprint changed: request={request_id}")
requests.append(replace(request, arrival_s=float(row["arrived_at"])))
return requests
def _real_request_records(
requests: list[TraceRequest], outcomes: list[Any]
) -> list[dict[str, Any]]:
outcomes_by_id = {outcome.request_id: outcome for outcome in outcomes}
records = []
for request in requests:
outcome = outcomes_by_id.get(request.row_id)
records.append(
{
"request_id": request.row_id,
"arrival_s": request.arrival_s,
"input_tokens": int(request.prompt_tokens_hint or 0),
"success": bool(outcome and outcome.success),
"ttft_ms": outcome.ttft_ms if outcome is not None else None,
"tpot_ms": outcome.tpot_ms if outcome is not None else None,
"completion_tokens": outcome.completion_tokens if outcome is not None else None,
"completion_tokens_source": (
outcome.completion_tokens_source if outcome is not None else ""
),
"error": outcome.error if outcome is not None else "missing_outcome",
}
)
return records
def run_community_config(
*,
manifest: dict[str, Any],
record: dict[str, Any],
protocol: dict[str, Any],
) -> dict[str, Any]:
result_path = Path(record["result_path"])
if result_path.is_file():
result = json.loads(result_path.read_text())
if result.get("status") == "completed":
return result
config = record["config"]
study_path = Path(record["study_path"])
study = load_study_spec(study_path)
_, all_requests = load_trace_requests(study, study_spec_path=study_path)
recipe = build_launch_recipe(
study.engine,
ConfigPatch(
flag_patch={
"max-num-seqs": int(config["mns"]),
"max-num-batched-tokens": int(config["mbt"]),
}
),
)
run_dir = Path(record["engine_log_path"]).parent
run_dir.mkdir(parents=True, exist_ok=True)
write_json(run_dir / "engine_command.json", recipe.argv)
marker_env = {
"AITUNER_STUDY_ID": study.study_id,
"AITUNER_TRIAL_ID": f"fixed-cohort-{config['name']}",
}
loads = []
load_by_rate = {
float(item["offered_request_rate"]): item
for item in protocol["rates"].values()
}
started = time.time()
with Path(record["engine_log_path"]).open("w", encoding="utf-8") as engine_log:
process = subprocess.Popen( # noqa: S603
recipe.argv,
cwd=recipe.cwd,
env={**recipe.env, **marker_env},
stdout=engine_log,
stderr=subprocess.STDOUT,
text=True,
start_new_session=True,
)
previous_sigterm = _install_sigterm_as_keyboardinterrupt()
try:
_wait_for_server_or_exit(
process,
base_url=recipe.base_url,
healthcheck_path=recipe.healthcheck_path,
ready_timeout_s=recipe.ready_timeout_s,
)
for rate in record["rate_order"]:
rate = float(rate)
load_dir = run_dir / rate_key(rate)
load_result_path = load_dir / "result.json"
if load_result_path.is_file():
load_result = json.loads(load_result_path.read_text())
if load_result.get("status") == "completed":
loads.append(load_result)
continue
load_dir.mkdir(parents=True, exist_ok=True)
rate_record = load_by_rate[rate]
requests = _materialize_real_requests(
all_requests=all_requests,
rate_record=rate_record,
cohort=protocol["cohort"],
)
load_started = time.time()
outcomes, early_stopped, early_stop_reason = _replay_requests(
requests,
base_url=recipe.base_url,
timeout_s=recipe.request_timeout_s,
max_concurrency=study.trace.max_concurrency,
# target=0 makes the mathematical SLO early-stop bound
# unreachable while retaining the existing replay engine.
target_pass_rate=0.0,
max_lag_s=None,
max_elapsed_s=max(float(rate_record["duration_s"]) + 900.0, 1200.0),
evaluate_outcome=lambda outcome: type(
"NoEarlyStopEvaluation", (), {"passed": bool(outcome.success)}
)(),
)
request_records = _real_request_records(requests, outcomes)
score = score_requests(request_records)
load_result = {
"status": "completed",
"config": config,
"offered_request_rate": rate,
"request_rate_per_gpu": rate / int(config["tp"]),
"elapsed_seconds": time.time() - load_started,
"trace_path": rate_record["path"],
"trace_sha256": rate_record["sha256"],
"early_stopped": early_stopped,
"early_stop_reason": early_stop_reason,
"requests": request_records,
**score,
}
if early_stopped:
raise RuntimeError(
f"complete replay stopped early: {config['name']} rate={rate}: "
f"{early_stop_reason}"
)
write_json(load_result_path, load_result)
loads.append(load_result)
print(
json.dumps(
{
"system": "community_vllm",
"config": config["name"],
"rate": rate,
"pass_rate": score["scores"][PRIMARY_SLO]["slo_pass_rate"],
"feasible": score["scores"][PRIMARY_SLO]["feasible"],
"elapsed_seconds": load_result["elapsed_seconds"],
},
sort_keys=True,
),
flush=True,
)
time.sleep(2.0)
finally:
_ignore_sigterm_if_main()
_terminate_process_tree(process, timeout_s=30.0, marker_env=marker_env)
_restore_sigterm(previous_sigterm)
loads.sort(key=lambda item: float(item["offered_request_rate"]))
result = {
"schema": COMMUNITY_SCHEMA,
"status": "completed",
"config": config,
"elapsed_seconds": time.time() - started,
"loads": loads,
}
write_json(result_path, result)
return result
def run_community(args: argparse.Namespace) -> Path:
manifest = json.loads(args.manifest.read_text())
if manifest.get("schema") != COMMUNITY_SCHEMA:
raise ValueError(f"unexpected community manifest schema: {manifest.get('schema')}")
protocol_path = Path(manifest["protocol_manifest"]["path"])
frontier_freeze = Path(manifest["frontier_freeze"]["path"])
if sha256(protocol_path) != manifest["protocol_manifest"]["sha256"]:
raise ValueError("protocol changed after community manifest preparation")
if sha256(frontier_freeze) != manifest["frontier_freeze"]["sha256"]:
raise ValueError("Frontier freeze changed after community manifest preparation")
_validate_frontier_freeze(frontier_freeze, protocol_path)
protocol = json.loads(protocol_path.read_text())
config_results = []
for record in sorted(manifest["configs"], key=lambda item: item["execution_index"]):
result = run_community_config(
manifest=manifest,
record=record,
protocol=protocol,
)
config_results.append(result)
print(
json.dumps(
{
"config": result["config"]["name"],
"config_completed": True,
"elapsed_seconds": result["elapsed_seconds"],
},
sort_keys=True,
),
flush=True,
)
freeze = {
"schema": COMMUNITY_SCHEMA,
"status": "completed",
"created_unix_s": time.time(),
"run_manifest_sha256": sha256(args.manifest),
"protocol_manifest_sha256": sha256(protocol_path),
"frontier_freeze_sha256": sha256(frontier_freeze),
"primary_slo": PRIMARY_SLO,
"rankings": {
name: rank_surface(config_results, slo_name=name) for name in SLO_VARIANTS
},
"config_results": config_results,
}
path = args.manifest.parent / "community_ranking_frozen.json"
write_json(path, freeze)
print(path)
return path
def _rank_map(ranking: list[dict[str, Any]]) -> dict[str, int]:
return {item["config"]["name"]: int(item["rank"]) for item in ranking}
def _capacity_map(ranking: list[dict[str, Any]]) -> dict[str, float | None]:
return {
item["config"]["name"]: item["maximum_tested_feasible_request_rate_per_gpu"]
for item in ranking
}
def _pearson(left: list[float], right: list[float]) -> float | None:
if len(left) < 2:
return None
left_mean = statistics.fmean(left)
right_mean = statistics.fmean(right)
numerator = sum((a - left_mean) * (b - right_mean) for a, b in zip(left, right))
denominator = math.sqrt(
sum((a - left_mean) ** 2 for a in left)
* sum((b - right_mean) ** 2 for b in right)
)
return numerator / denominator if denominator else None
def _average_ranks(capacities: dict[str, float | None]) -> dict[str, float]:
ordered = sorted(
capacities,
key=lambda name: (
-(
capacities[name]
if capacities[name] is not None
else float("-inf")
),
name,
),
)
ranks = {}
start = 0
while start < len(ordered):
value = capacities[ordered[start]]
end = start + 1
while end < len(ordered) and capacities[ordered[end]] == value:
end += 1
average_rank = ((start + 1) + end) / 2.0
for name in ordered[start:end]:
ranks[name] = average_rank
start = end
return ranks
def _pairwise_accuracy(
frontier_capacity: dict[str, float | None], real_capacity: dict[str, float | None]
) -> dict[str, Any]:
names = sorted(set(frontier_capacity) & set(real_capacity))
comparable = 0
correct = 0
frontier_ties = 0
real_ties = 0
for left_index, left in enumerate(names):
for right in names[left_index + 1 :]:
f_left = frontier_capacity[left]
f_right = frontier_capacity[right]
r_left = real_capacity[left]
r_right = real_capacity[right]
if None in (f_left, f_right, r_left, r_right):
continue
f_sign = (f_left > f_right) - (f_left < f_right)
r_sign = (r_left > r_right) - (r_left < r_right)
frontier_ties += int(f_sign == 0)
real_ties += int(r_sign == 0)
if f_sign == 0 or r_sign == 0:
continue
comparable += 1
correct += int(f_sign == r_sign)
return {
"comparable_non_tied_pairs": comparable,
"correct_pairs": correct,
"accuracy": correct / comparable if comparable else None,
"frontier_tied_pairs": frontier_ties,
"real_tied_pairs": real_ties,
}
def _request_residuals(
frontier_results: list[dict[str, Any]], real_results: list[dict[str, Any]]
) -> dict[str, Any]:
def index(results: list[dict[str, Any]]) -> dict[tuple[str, float, str], dict[str, Any]]:
indexed = {}
for config in results:
name = config["config"]["name"]
for load in config["loads"]:
rate = float(load["offered_request_rate"])
for request in load["requests"]:
indexed[(name, rate, str(request["request_id"]))] = request
return indexed
frontier_index = index(frontier_results)
real_index = index(real_results)
keys = sorted(set(frontier_index) & set(real_index))
errors = []
absolute_percentage_errors = []
by_bin: dict[str, list[float]] = {}
for key in keys:
frontier_request = frontier_index[key]
real_request = real_index[key]
if not real_request["success"] or real_request.get("ttft_ms") is None:
continue
predicted = float(frontier_request["ttft_ms"])
observed = float(real_request["ttft_ms"])
error = predicted - observed
errors.append(error)
if observed > 0:
absolute_percentage_errors.append(abs(error) / observed)
input_tokens = int(real_request["input_tokens"])
bin_index = _input_bin(input_tokens)
label = f"[{INPUT_BINS[bin_index]},{INPUT_BINS[bin_index + 1]})"
by_bin.setdefault(label, []).append(error)
return {
"matched_successful_requests": len(errors),
"signed_error_ms_mean": statistics.fmean(errors) if errors else None,
"mae_ms": statistics.fmean(abs(value) for value in errors) if errors else None,
"mape": (
statistics.fmean(absolute_percentage_errors)
if absolute_percentage_errors
else None
),
"by_input_bin": {
label: {
"count": len(values),
"signed_error_ms_mean": statistics.fmean(values),
"mae_ms": statistics.fmean(abs(value) for value in values),
}
for label, values in sorted(by_bin.items())
},
}
def _pass_rate_residuals(
frontier_results: list[dict[str, Any]], real_results: list[dict[str, Any]]
) -> dict[str, Any]:
def index(
results: list[dict[str, Any]], slo_name: str
) -> dict[tuple[str, float], float]:
return {
(config["config"]["name"], float(load["offered_request_rate"])): float(
load["scores"][slo_name]["slo_pass_rate"]
)
for config in results
for load in config["loads"]
}
variants = {}
for slo_name in SLO_VARIANTS:
frontier_index = index(frontier_results, slo_name)
real_index = index(real_results, slo_name)
keys = sorted(set(frontier_index) & set(real_index))
errors = [frontier_index[key] - real_index[key] for key in keys]
variants[slo_name] = {
"matched_config_load_points": len(errors),
"signed_error_mean": statistics.fmean(errors) if errors else None,
"mae": statistics.fmean(abs(value) for value in errors) if errors else None,
"max_absolute_error": max((abs(value) for value in errors), default=None),
}
return variants
def compare(args: argparse.Namespace) -> Path:
frontier = json.loads(args.frontier_freeze.read_text())
community = json.loads(args.community_freeze.read_text())
if frontier.get("schema") != FRONTIER_SCHEMA:
raise ValueError("unexpected Frontier freeze")
if community.get("schema") != COMMUNITY_SCHEMA or community.get("status") != "completed":
raise ValueError("unexpected/incomplete community freeze")
if frontier["protocol_manifest_sha256"] != community["protocol_manifest_sha256"]:
raise ValueError("comparison inputs used different protocols")
variants = {}
for slo_name in SLO_VARIANTS:
frontier_ranking = frontier["rankings"][slo_name]
real_ranking = community["rankings"][slo_name]
frontier_ranks = _rank_map(frontier_ranking)
real_ranks = _rank_map(real_ranking)
names = sorted(set(frontier_ranks) & set(real_ranks))
frontier_capacity = _capacity_map(frontier_ranking)
real_capacity = _capacity_map(real_ranking)
frontier_average_ranks = _average_ranks(frontier_capacity)
real_average_ranks = _average_ranks(real_capacity)
best_frontier = max(
(value for value in frontier_capacity.values() if value is not None),
default=None,
)
best_real = max(
(value for value in real_capacity.values() if value is not None),
default=None,
)
frontier_top_set = [
name for name, value in frontier_capacity.items() if value == best_frontier
]
real_top_set = [name for name, value in real_capacity.items() if value == best_real]
selected_real_capacities = [
real_capacity[name]
for name in frontier_top_set
if real_capacity[name] is not None
]
regret_values = [
(best_real - value) / best_real
for value in selected_real_capacities
if best_real not in (None, 0)
]
variants[slo_name] = {
"frontier_top1_set": frontier_top_set,
"real_top1_set": real_top_set,
"top1_set_intersection": sorted(set(frontier_top_set) & set(real_top_set)),
"top1_set_match": set(frontier_top_set) == set(real_top_set),
"top1_regret_fraction_best_tie_break": min(regret_values)
if regret_values
else None,
"top1_regret_fraction_worst_tie_break": max(regret_values)
if regret_values
else None,
"spearman_rank_correlation": _pearson(
[frontier_average_ranks[name] for name in names],
[real_average_ranks[name] for name in names],
),
"pairwise_ordering": _pairwise_accuracy(frontier_capacity, real_capacity),
"frontier_ranking": frontier_ranking,
"real_ranking": real_ranking,
}
result = {
"schema": COMPARISON_SCHEMA,
"created_unix_s": time.time(),
"frontier_freeze": {
"path": str(args.frontier_freeze.resolve()),
"sha256": sha256(args.frontier_freeze),
},
"community_freeze": {
"path": str(args.community_freeze.resolve()),
"sha256": sha256(args.community_freeze),
},
"primary_slo": PRIMARY_SLO,
"variants": variants,
"request_level_ttft_residuals": _request_residuals(
frontier["config_results"], community["config_results"]
),
"pass_rate_surface_residuals": _pass_rate_residuals(
frontier["config_results"], community["config_results"]
),
}
write_json(args.output, result)
print(args.output)
return args.output
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(dest="command", required=True)
prepare = subparsers.add_parser("prepare-protocol")
prepare.add_argument("--trace", type=Path, required=True)
prepare.add_argument("--expected-trace-sha256", default=frontier_grid.TRACE_SHA256)
prepare.add_argument("--output-root", type=Path, required=True)
prepare.add_argument("--cohort-size", type=int, default=COHORT_SIZE)
prepare.add_argument("--seed", type=int, default=COHORT_SEED)
prepare.add_argument("--rate", action="append", type=float, default=[])
audit = subparsers.add_parser("audit-protocol")
audit.add_argument("--manifest", type=Path, required=True)
frontier = subparsers.add_parser("run-frontier")
frontier.add_argument("--frontier-source", type=Path, required=True)
frontier.add_argument(
"--frontier-commit", default="d9cfeb6d8791fbf2f295dd9744c56a666171776e"
)
frontier.add_argument("--python", type=Path, required=True)
frontier.add_argument("--profile-root", type=Path, required=True)
frontier.add_argument("--protocol-manifest", type=Path, required=True)
frontier.add_argument("--output-root", type=Path, required=True)
frontier.add_argument("--tp4-capacity-artifact", type=Path, required=True)
frontier.add_argument("--tp8-capacity-artifact", type=Path, required=True)
frontier.add_argument("--config", action="append")
community = subparsers.add_parser("prepare-community")
community.add_argument("--repo", type=Path, required=True)
community.add_argument("--python", type=Path, required=True)
community.add_argument("--vllm", type=Path, required=True)
community.add_argument("--model", type=Path, required=True)
community.add_argument("--windows", type=Path, required=True)
community.add_argument("--protocol-manifest", type=Path, required=True)
community.add_argument("--frontier-freeze", type=Path, required=True)
community.add_argument("--output-root", type=Path, required=True)
community.add_argument("--port", type=int, default=18918)
community.add_argument("--config-order-seed", type=int, default=CONFIG_ORDER_SEED)
community.add_argument("--rate-order-seed", type=int, default=RATE_ORDER_SEED)
run_real = subparsers.add_parser("run-community")
run_real.add_argument("--manifest", type=Path, required=True)
compare_parser = subparsers.add_parser("compare")
compare_parser.add_argument("--frontier-freeze", type=Path, required=True)
compare_parser.add_argument("--community-freeze", type=Path, required=True)
compare_parser.add_argument("--output", type=Path, required=True)
return parser.parse_args()
def main() -> None:
args = parse_args()
if args.command == "prepare-protocol":
if not args.rate:
args.rate = list(OFFERED_RATES)
if sorted(set(args.rate)) != list(args.rate):
raise ValueError("--rate values must be unique and ascending")
prepare_protocol(args)
elif args.command == "audit-protocol":
print(json.dumps(audit_protocol(args.manifest), indent=2))
elif args.command == "run-frontier":
run_frontier(args)
elif args.command == "prepare-community":
prepare_community(args)
elif args.command == "run-community":
run_community(args)
elif args.command == "compare":
compare(args)
else:
raise AssertionError(args.command)
if __name__ == "__main__":
main()