diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/fixed_cohort_rank.py b/runs/frontier-multicase-sufficiency-v0/best_effort/fixed_cohort_rank.py new file mode 100644 index 0000000..78d4f58 --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/fixed_cohort_rank.py @@ -0,0 +1,1404 @@ +#!/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", + "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() diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/test_fixed_cohort_rank.py b/runs/frontier-multicase-sufficiency-v0/best_effort/test_fixed_cohort_rank.py new file mode 100644 index 0000000..966a912 --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/test_fixed_cohort_rank.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import argparse +import csv +import importlib.util +import json +import sys +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("fixed_cohort_rank.py") +SPEC = importlib.util.spec_from_file_location("fixed_cohort_rank", SCRIPT) +MODULE = importlib.util.module_from_spec(SPEC) +assert SPEC.loader is not None +sys.modules[SPEC.name] = MODULE +SPEC.loader.exec_module(MODULE) + + +def _write_trace(path: Path) -> None: + lengths = [100, 1500, 3000, 6000, 12000, 24000] + rows = [] + for index in range(120): + rows.append( + { + "timestamp": index * 0.5, + "sampling_u": (index + 1) / 121, + "input_length": lengths[index % len(lengths)], + "output_length": 10, + "prompt": f"raw prompt {index}", + } + ) + path.write_text("".join(json.dumps(row) + "\n" for row in rows)) + + +def test_prepare_protocol_keeps_identical_cohort_and_scales_only_arrivals( + tmp_path: Path, +) -> None: + trace = tmp_path / "trace.jsonl" + _write_trace(trace) + manifest_path = MODULE.prepare_protocol( + argparse.Namespace( + trace=trace, + expected_trace_sha256="", + output_root=tmp_path / "protocol", + cohort_size=24, + seed=7, + rate=[0.2, 0.4], + ) + ) + manifest = json.loads(manifest_path.read_text()) + assert manifest["selection"]["cohort_bin_quotas"] == [4, 4, 4, 4, 4, 4] + assert manifest["load_contract"]["binary_search"] is False + assert manifest["load_contract"]["monotonicity_assumed"] is False + + materialized = [] + for record in manifest["rates"].values(): + with Path(record["path"]).open(newline="") as source: + rows = list(csv.DictReader(source)) + materialized.append(rows) + duration = float(rows[-1]["arrived_at"]) - float(rows[0]["arrived_at"]) + assert len(rows) / duration == record["offered_request_rate"] + assert [row["source_row_index"] for row in materialized[0]] == [ + row["source_row_index"] for row in materialized[1] + ] + assert [row["num_prefill_tokens"] for row in materialized[0]] == [ + row["num_prefill_tokens"] for row in materialized[1] + ] + assert MODULE.audit_protocol(manifest_path)["status"] == "passed" + + +def test_score_requests_applies_all_slos_to_same_outcomes() -> None: + requests = [ + {"input_tokens": 100, "success": True, "ttft_ms": 900.0}, + {"input_tokens": 16000, "success": True, "ttft_ms": 2500.0}, + {"input_tokens": 24000, "success": False, "ttft_ms": None}, + ] + score = MODULE.score_requests(requests) + assert score["scores"]["linear_8k_primary"]["passed_request_count"] == 2 + assert score["scores"]["legacy_step_1s_2s"]["passed_request_count"] == 1 + assert score["ttft_ms"]["count"] == 2 + + +def test_rank_surface_records_nonmonotone_feasibility_without_binary_search() -> None: + config_results = [ + { + "config": {"name": "a", "tp": 4}, + "loads": [ + { + "offered_request_rate": 0.2, + "scores": {"linear_8k_primary": {"feasible": False, "slo_pass_rate": 0.9}}, + }, + { + "offered_request_rate": 0.4, + "scores": {"linear_8k_primary": {"feasible": True, "slo_pass_rate": 0.95}}, + }, + ], + }, + { + "config": {"name": "b", "tp": 8}, + "loads": [ + { + "offered_request_rate": 0.2, + "scores": {"linear_8k_primary": {"feasible": True, "slo_pass_rate": 1.0}}, + }, + { + "offered_request_rate": 0.4, + "scores": {"linear_8k_primary": {"feasible": False, "slo_pass_rate": 0.9}}, + }, + ], + }, + ] + ranking = MODULE.rank_surface(config_results, slo_name="linear_8k_primary") + a = next(item for item in ranking if item["config"]["name"] == "a") + assert a["maximum_tested_feasible_request_rate"] == 0.4 + assert a["monotonicity_violations"] == [[0.2, 0.4]] + + +def test_average_ranks_use_midrank_for_capacity_ties() -> None: + assert MODULE._average_ranks({"a": 0.2, "b": 0.2, "c": 0.1}) == { + "a": 1.5, + "b": 1.5, + "c": 3.0, + } + + +def test_community_study_uses_affine_primary_slo(tmp_path: Path) -> None: + payload = MODULE._community_study_payload( + tp=8, + repo=tmp_path, + python=tmp_path / "python", + vllm=tmp_path / "vllm", + model=tmp_path / "model", + trace=tmp_path / "trace.jsonl", + windows=tmp_path / "windows.json", + port=18918, + ) + assert payload["slo"]["ttft_rule"] == { + "kind": "linear_ms", + "intercept_ms": 1000.0, + "per_token_ms": 0.125, + } + assert payload["trace"]["request_mode"] == "raw_completion" + assert payload["trace"]["completion_tokens_override"] == 1