From a6f101b09e5ba1089ee982fadef303d6b0b8b344 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Wed, 15 Jul 2026 19:39:30 +0800 Subject: [PATCH] Add reproducible Frontier prefill grid freeze --- .../best_effort/frontier_prefill_grid.py | 703 ++++++++++++++++++ .../best_effort/test_frontier_prefill_grid.py | 66 ++ 2 files changed, 769 insertions(+) create mode 100644 runs/frontier-multicase-sufficiency-v0/best_effort/frontier_prefill_grid.py create mode 100644 runs/frontier-multicase-sufficiency-v0/best_effort/test_frontier_prefill_grid.py diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_prefill_grid.py b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_prefill_grid.py new file mode 100644 index 0000000..ce6bec5 --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_prefill_grid.py @@ -0,0 +1,703 @@ +#!/usr/bin/env python3 +"""Freeze a profile-only Frontier ranking for the Qwen235B prefill grid.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import os +import subprocess +import sys +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Iterable + + +SCHEMA = "frontier-qwen235b-prefill-grid-v1" +MODEL = "Qwen3-235B-A22B-FP8" +TRACE_SHA256 = "f878e9af18f94dcfaced94a8e1e6b20a2f7d97d64aa862448025660dbbd965b2" +PROFILE_RELATIVE = Path("compute/h20") / MODEL +NETWORK_RELATIVE = Path("network/h20_nccl/all_reduce.csv") +SEARCH_LOW = 0.0 +SEARCH_HIGH = 0.125 +SEARCH_PROBES = 6 +WINDOW_DURATION_S = 600.0 +MAX_INPUT_TOKENS = 32768 +MAX_MODEL_TOKENS = 40960 +BLOCK_SIZE_TOKENS = 16 +MOE_ROUTING_SEED = 42 + + +@dataclass(frozen=True) +class GridConfig: + tp: int + mns: int + mbt: int + moe_tp: int + moe_ep: int + num_gpu_blocks: int + + @property + def name(self) -> str: + return f"tp{self.tp}_mns{self.mns}_mbt{self.mbt}" + + @property + def gpu_count(self) -> int: + return self.tp + + +GRID = tuple( + GridConfig( + tp=tp, + mns=mns, + mbt=mbt, + moe_tp=4 if tp == 4 else 1, + moe_ep=1 if tp == 4 else 8, + # Measured by community vLLM on the same checkpoint/runtime. The real + # grid will use --num-gpu-blocks-override with these same values. + num_gpu_blocks=26101 if tp == 4 else 62351, + ) + for tp in (4, 8) + for mns in (64, 128) + for mbt in (8192, 16384) +) + + +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 order_hash(values: Iterable[object]) -> str: + payload = "\n".join(str(value) for value in values).encode() + return hashlib.sha256(payload).hexdigest() + + +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 sha256_bytes(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +def frontier_source_fingerprint(source: Path) -> dict[str, Any]: + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=source, + check=True, + stdout=subprocess.PIPE, + text=True, + ).stdout.strip() + status = subprocess.run( + ["git", "status", "--porcelain=v1"], + cwd=source, + check=True, + stdout=subprocess.PIPE, + text=True, + ).stdout + diff = subprocess.run( + ["git", "diff", "--binary", "HEAD"], + cwd=source, + check=True, + stdout=subprocess.PIPE, + ).stdout + model_config = source / "data" / "config" / "models" / f"{MODEL}.json" + if not model_config.is_file(): + raise FileNotFoundError(f"Frontier model config is missing: {model_config}") + return { + "commit": commit, + "status_porcelain": status.splitlines(), + "tracked_diff_sha256": sha256_bytes(diff), + "model_config": { + "path": str(model_config.resolve()), + "sha256": sha256(model_config), + }, + } + + +def anchor_key(anchor: float) -> str: + return f"{anchor:.12f}" + + +def anchor_filename(anchor: float) -> str: + return f"u_{anchor_key(anchor).replace('.', 'p')}.csv" + + +def binary_search_lattice() -> list[float]: + intervals = [(SEARCH_LOW, SEARCH_HIGH)] + anchors: list[float] = [] + for _ in range(SEARCH_PROBES): + next_intervals = [] + for low, high in intervals: + midpoint = (low + high) / 2.0 + anchors.append(midpoint) + next_intervals.extend(((low, midpoint), (midpoint, high))) + intervals = next_intervals + return sorted(set(anchors)) + + +def ttft_slo_ms(input_tokens: int) -> float: + return 1000.0 if input_tokens <= 8191 else 2000.0 + + +def load_trace_rows(path: Path) -> list[dict[str, Any]]: + rows = [] + with path.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 <= MAX_INPUT_TOKENS: + continue + timestamp = float(raw["timestamp"]) + sampling_u = float(raw["sampling_u"]) + rows.append( + { + "arrived_at": timestamp, + "num_prefill_tokens": input_tokens, + "num_decode_tokens": 1, + "source_row_index": source_index, + "source_request_id": str( + raw.get("request_id") or raw.get("id") or source_index + ), + "sampling_u": sampling_u, + "slo_ttft_ms": ttft_slo_ms(input_tokens), + } + ) + # This reproduces AITuner's stable arrival-only sort. Tied timestamps retain + # source order; sampling_u must not become a secondary key. + rows.sort(key=lambda row: row["arrived_at"]) + return rows + + +def prepare_traces(trace: Path, output_root: Path, expected_sha256: str) -> Path: + actual_sha256 = sha256(trace) + if expected_sha256 and actual_sha256 != expected_sha256: + raise ValueError( + f"trace SHA256 mismatch: expected={expected_sha256}, actual={actual_sha256}" + ) + rows = load_trace_rows(trace) + if not rows: + raise ValueError("no trace rows remain after input-length filtering") + if any( + rows[index]["arrived_at"] > rows[index + 1]["arrived_at"] + for index in range(len(rows) - 1) + ): + raise ValueError("materialized trace is not monotonic by arrival") + + 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", + "sampling_u", + "slo_ttft_ms", + ] + anchors: dict[str, Any] = {} + for anchor in binary_search_lattice(): + selected = [row for row in rows if row["sampling_u"] <= anchor] + target = trace_dir / anchor_filename(anchor) + with target.open("w", encoding="utf-8", newline="") as output: + writer = csv.DictWriter(output, fieldnames=fields) + writer.writeheader() + writer.writerows(selected) + key = anchor_key(anchor) + anchors[key] = { + "anchor": anchor, + "path": str(target.resolve()), + "sha256": sha256(target), + "request_count": len(selected), + "request_rate": len(selected) / WINDOW_DURATION_S, + "source_row_order_sha256": order_hash( + row["source_row_index"] for row in selected + ), + "arrival_order_sha256": order_hash( + f"{row['arrived_at']:.12f}" for row in selected + ), + "input_length_order_sha256": order_hash( + row["num_prefill_tokens"] for row in selected + ), + } + + manifest = { + "schema": SCHEMA, + "source": { + "path": str(trace.resolve()), + "sha256": actual_sha256, + "filtered_request_count": len(rows), + }, + "selection_contract": { + "input_tokens": [0, MAX_INPUT_TOKENS], + "completion_tokens_override": 1, + "stable_sort_key": "arrived_at_only", + "selection": "sampling_u <= anchor", + "window_duration_s": WINDOW_DURATION_S, + "search_low": SEARCH_LOW, + "search_high": SEARCH_HIGH, + "search_probes": SEARCH_PROBES, + "lattice_anchor_count": len(anchors), + "ttft_slo_ms": {"input_le_8191": 1000, "otherwise": 2000}, + "target_pass_rate": 0.95, + }, + "anchors": anchors, + } + manifest_path = output_root / "trace_manifest.json" + write_json(manifest_path, manifest) + return manifest_path + + +def resolve_profile_paths(profile_root: Path) -> dict[str, Path]: + paths = { + "linear": profile_root / PROFILE_RELATIVE / "linear_op.csv", + "attention": profile_root / PROFILE_RELATIVE / "attention.csv", + "moe": profile_root / PROFILE_RELATIVE / "moe.csv", + "all_reduce": profile_root / NETWORK_RELATIVE, + "manifest": profile_root / "profile_manifest.json", + } + missing = [f"{name}={path}" for name, path in paths.items() if not path.is_file()] + if missing: + raise FileNotFoundError("missing profile inputs: " + ", ".join(missing)) + return paths + + +def build_command( + *, + python: Path, + frontier_source: Path, + profile_root: Path, + profile_paths: dict[str, Path], + trace: Path, + config: GridConfig, + probe_dir: Path, + run_id: str, + cache_root: Path, +) -> list[str]: + return [ + str(python), + "-m", + "frontier.main", + "--simulation_mode", + "offline", + "--sys_arch", + "co-location", + "--cluster_config_num_replicas", + "1", + "--replica_config_model_name", + MODEL, + "--replica_config_attn_tensor_parallel_size", + str(config.tp), + "--replica_config_attn_data_parallel_size", + "1", + "--replica_config_moe_tensor_parallel_size", + str(config.moe_tp), + "--replica_config_moe_expert_parallel_size", + str(config.moe_ep), + "--replica_config_total_expert_num", + "128", + "--replica_config_router_topk", + "8", + "--replica_config_moe_routing_mode", + "simulation", + "--replica_config_moe_routing_seed", + str(MOE_ROUTING_SEED), + "--replica_config_num_pipeline_stages", + "1", + "--replica_config_device", + "h20", + "--replica_config_network_device", + "h20_dgx", + "--cc_backend_config_type", + "vidur", + "--vidur_cc_backend_config_profiling_data_dir", + str(profile_root), + "--vidur_cc_backend_config_cache_dir", + str(cache_root / "collectives"), + "--vidur_cc_backend_config_all_reduce_input_file", + str(profile_paths["all_reduce"]), + "--replica_scheduler_config_type", + "vllm_v1", + "--decode_cuda_graph_mode", + "none", + "--vllm_v1_scheduler_config_batch_size_cap", + str(config.mns), + "--vllm_v1_scheduler_config_block_size", + str(BLOCK_SIZE_TOKENS), + "--vllm_v1_scheduler_config_num_blocks", + str(config.num_gpu_blocks), + "--vllm_v1_scheduler_config_num_blocks_mode", + "explicit", + "--vllm_v1_scheduler_config_max_tokens_in_batch", + str(config.mbt), + "--vllm_v1_scheduler_config_enable_chunked_prefill", + "--no-vllm_v1_scheduler_config_enable_prefix_caching", + "--request_generator_config_type", + "trace_replay", + "--trace_request_generator_config_trace_file", + str(trace), + "--trace_request_generator_config_time_scale_factor", + "1", + "--trace_request_generator_config_prefill_scale_factor", + "1", + "--trace_request_generator_config_decode_scale_factor", + "1", + "--trace_request_generator_config_max_tokens", + str(MAX_MODEL_TOKENS), + "--no-random_forrest_execution_time_predictor_config_enable_dummy_mode", + "--random_forrest_execution_time_predictor_config_linear_op_input_file", + str(profile_paths["linear"]), + "--random_forrest_execution_time_predictor_config_atten_input_file", + str(profile_paths["attention"]), + "--random_forrest_execution_time_predictor_config_moe_input_file", + str(profile_paths["moe"]), + "--random_forrest_execution_time_predictor_config_all_reduce_input_file", + str(profile_paths["all_reduce"]), + "--random_forrest_execution_time_predictor_config_prediction_max_prefill_chunk_size", + "16384", + "--random_forrest_execution_time_predictor_config_prediction_max_tokens_per_request", + str(MAX_INPUT_TOKENS + 1), + "--random_forrest_execution_time_predictor_config_prediction_max_batch_size", + "128", + "--random_forrest_execution_time_predictor_config_skip_cpu_overhead_modeling", + "--metrics_config_cache_dir", + str(cache_root / "execution"), + "--metrics_config_output_dir", + str(probe_dir / "metrics"), + "--metrics_config_run_id", + run_id, + "--metrics_config_write_metrics", + "--metrics_config_store_request_metrics", + "--no-metrics_config_store_plots", + "--no-metrics_config_enable_chrome_trace", + "--no-metrics_config_write_json_trace", + ] + + +def score_request_metrics(trace: Path, request_metrics: Path) -> dict[str, Any]: + with trace.open(encoding="utf-8", newline="") as source: + trace_rows = list(csv.DictReader(source)) + with request_metrics.open(encoding="utf-8", newline="") as source: + metric_rows = list(csv.DictReader(source)) + if len(metric_rows) != len(trace_rows): + raise ValueError( + f"request count mismatch: trace={len(trace_rows)}, metrics={len(metric_rows)}" + ) + metrics_by_id = {int(row["Request Id"]): row for row in metric_rows} + expected_ids = set(range(len(trace_rows))) + if set(metrics_by_id) != expected_ids: + raise ValueError("Frontier Request Ids do not match trace row positions") + + passed = 0 + ttfts = [] + for request_id, trace_row in enumerate(trace_rows): + ttft = float(metrics_by_id[request_id]["ttft"]) + threshold = float(trace_row["slo_ttft_ms"]) + ttfts.append(ttft) + passed += int(ttft <= threshold) + count = len(trace_rows) + pass_rate = passed / count if count else 0.0 + ordered = sorted(ttfts) + + def percentile(fraction: float) -> float | None: + if not ordered: + return None + index = round((len(ordered) - 1) * fraction) + return ordered[index] + + return { + "request_count": count, + "passed_request_count": passed, + "slo_pass_rate": pass_rate, + "feasible": pass_rate >= 0.95, + "ttft_ms": { + "min": min(ordered) if ordered else None, + "p50": percentile(0.50), + "p95": percentile(0.95), + "p99": percentile(0.99), + "max": max(ordered) if ordered else None, + }, + } + + +def find_request_metrics(probe_dir: Path) -> Path: + candidates = list((probe_dir / "metrics").glob("**/request_metrics.csv")) + if len(candidates) != 1: + raise ValueError( + f"expected one request_metrics.csv under {probe_dir}, got {candidates}" + ) + return candidates[0] + + +def run_probe( + *, + python: Path, + frontier_source: Path, + profile_root: Path, + profile_paths: dict[str, Path], + trace_record: dict[str, Any], + config: GridConfig, + probe_index: int, + output_root: Path, + cache_root: Path, +) -> dict[str, Any]: + anchor = float(trace_record["anchor"]) + probe_dir = output_root / "runs" / config.name / f"probe_{probe_index}_{anchor_filename(anchor)[:-4]}" + result_path = probe_dir / "result.json" + if result_path.is_file(): + result = json.loads(result_path.read_text()) + if result.get("status") == "completed": + return result + probe_dir.mkdir(parents=True, exist_ok=True) + trace = Path(trace_record["path"]) + run_id = f"{config.name}_probe{probe_index}_{anchor_filename(anchor)[:-4]}" + command = build_command( + python=python, + frontier_source=frontier_source, + profile_root=profile_root, + profile_paths=profile_paths, + trace=trace, + config=config, + probe_dir=probe_dir, + run_id=run_id, + cache_root=cache_root, + ) + write_json(probe_dir / "command.json", command) + environment = os.environ.copy() + environment.update( + { + "PYTHONPATH": str(frontier_source), + "WANDB_DISABLED": "true", + "VIDUR_DISABLE_WANDB": "1", + } + ) + started = time.time() + with (probe_dir / "stdout.log").open("w", encoding="utf-8") as output: + completed = subprocess.run( + command, + cwd=frontier_source, + env=environment, + stdout=output, + stderr=subprocess.STDOUT, + check=False, + ) + elapsed = time.time() - started + if completed.returncode != 0: + result = { + "status": "failed", + "returncode": completed.returncode, + "elapsed_seconds": elapsed, + "config": asdict(config), + "anchor": anchor, + "trace_sha256": trace_record["sha256"], + } + write_json(result_path, result) + raise RuntimeError(f"Frontier probe failed: {config.name}, anchor={anchor}") + + request_metrics = find_request_metrics(probe_dir) + score = score_request_metrics(trace, request_metrics) + result = { + "status": "completed", + "elapsed_seconds": elapsed, + "config": asdict(config), + "anchor": anchor, + "trace_path": str(trace), + "trace_sha256": trace_record["sha256"], + "request_metrics_path": str(request_metrics), + "request_metrics_sha256": sha256(request_metrics), + "request_rate": score["request_count"] / WINDOW_DURATION_S, + "request_rate_per_gpu": score["request_count"] + / WINDOW_DURATION_S + / config.gpu_count, + **score, + } + write_json(result_path, result) + return result + + +def selected_configs(names: list[str] | None) -> list[GridConfig]: + if not names: + return list(GRID) + by_name = {config.name: config for config in GRID} + unknown = sorted(set(names) - set(by_name)) + if unknown: + raise ValueError(f"unknown configs: {unknown}; available={sorted(by_name)}") + return [by_name[name] for name in names] + + +def run_grid(args: argparse.Namespace) -> None: + trace_manifest = json.loads(args.trace_manifest.read_text()) + if trace_manifest.get("schema") != SCHEMA: + raise ValueError(f"unexpected trace manifest schema: {trace_manifest.get('schema')}") + profile_paths = resolve_profile_paths(args.profile_root) + args.output_root.mkdir(parents=True, exist_ok=True) + cache_root = args.output_root / "cache" + configs = selected_configs(args.config) + run_manifest = { + "schema": SCHEMA, + "frontier": { + "source": str(args.frontier_source.resolve()), + "python": str(args.python.resolve()), + "fingerprint": frontier_source_fingerprint(args.frontier_source), + }, + "trace_manifest": { + "path": str(args.trace_manifest.resolve()), + "sha256": sha256(args.trace_manifest), + }, + "profiles": { + 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, + }, + }, + "contract": { + "metric": "maximum SLO-feasible request_rate_per_gpu", + "target_pass_rate": 0.95, + "cpu_overhead_modeling": "skipped_no_community_vllm_native_records", + "end_to_end_calibration": False, + "moe_routing_mode": "simulation", + "moe_routing_seed": MOE_ROUTING_SEED, + "kv_blocks_source": "community_vllm_measured_and_fixed_per_topology", + "network_device": "h20_dgx", + }, + "configs": [asdict(config) | {"name": config.name} for config in configs], + } + write_json(args.output_root / "run_manifest.json", run_manifest) + + all_results = {} + for config in configs: + low = SEARCH_LOW + high = SEARCH_HIGH + probes = [] + best = None + for probe_index in range(SEARCH_PROBES): + anchor = (low + high) / 2.0 + record = trace_manifest["anchors"].get(anchor_key(anchor)) + if record is None: + raise ValueError(f"trace manifest lacks anchor {anchor_key(anchor)}") + result = run_probe( + python=args.python, + frontier_source=args.frontier_source, + profile_root=args.profile_root, + profile_paths=profile_paths, + trace_record=record, + config=config, + probe_index=probe_index, + output_root=args.output_root, + cache_root=cache_root, + ) + probes.append(result) + print( + json.dumps( + { + "config": config.name, + "probe": probe_index, + "anchor": anchor, + "pass_rate": result["slo_pass_rate"], + "feasible": result["feasible"], + "elapsed_seconds": result["elapsed_seconds"], + }, + sort_keys=True, + ), + flush=True, + ) + if result["feasible"]: + low = anchor + best = result + else: + high = anchor + summary = { + "config": asdict(config) | {"name": config.name}, + "probes": probes, + "capacity_interval": [low, high], + "best_feasible": best, + } + write_json(args.output_root / "results" / f"{config.name}.json", summary) + all_results[config.name] = summary + + ranked = sorted( + all_results.values(), + key=lambda item: ( + -( + item["best_feasible"]["request_rate_per_gpu"] + if item["best_feasible"] is not None + else -1.0 + ), + item["config"]["name"], + ), + ) + freeze = { + "schema": SCHEMA, + "run_manifest_sha256": sha256(args.output_root / "run_manifest.json"), + "ranking": [ + { + "rank": index + 1, + "config": item["config"], + "capacity_interval_sampling_u": item["capacity_interval"], + "best_feasible": item["best_feasible"], + } + for index, item in enumerate(ranked) + ], + } + write_json(args.output_root / "frontier_ranking_frozen.json", freeze) + print(json.dumps(freeze, indent=2), flush=True) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(dest="command", required=True) + prepare = subparsers.add_parser("prepare") + prepare.add_argument("--trace", type=Path, required=True) + prepare.add_argument("--output-root", type=Path, required=True) + prepare.add_argument("--expected-trace-sha256", default=TRACE_SHA256) + + run = subparsers.add_parser("run") + run.add_argument("--frontier-source", type=Path, required=True) + run.add_argument("--python", type=Path, required=True) + run.add_argument("--profile-root", type=Path, required=True) + run.add_argument("--trace-manifest", type=Path, required=True) + run.add_argument("--output-root", type=Path, required=True) + run.add_argument("--tp4-capacity-artifact", type=Path, required=True) + run.add_argument("--tp8-capacity-artifact", type=Path, required=True) + run.add_argument("--config", action="append") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.command == "prepare": + manifest = prepare_traces( + args.trace, args.output_root, args.expected_trace_sha256 + ) + print(manifest) + return + if args.command == "run": + run_grid(args) + return + raise AssertionError(args.command) + + +if __name__ == "__main__": + main() diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/test_frontier_prefill_grid.py b/runs/frontier-multicase-sufficiency-v0/best_effort/test_frontier_prefill_grid.py new file mode 100644 index 0000000..3ace7eb --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/test_frontier_prefill_grid.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import csv +import importlib.util +import json +import sys +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("frontier_prefill_grid.py") +SPEC = importlib.util.spec_from_file_location("frontier_prefill_grid", 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 test_binary_search_lattice_covers_six_probe_tree() -> None: + anchors = MODULE.binary_search_lattice() + assert len(anchors) == 63 + assert anchors[0] == 0.001953125 + assert anchors[-1] == 0.123046875 + assert 0.0625 in anchors + + +def test_prepare_preserves_stable_arrival_order(tmp_path: Path) -> None: + trace = tmp_path / "trace.jsonl" + rows = [ + {"timestamp": 1.0, "sampling_u": 0.02, "input_length": 9000, "id": "a"}, + {"timestamp": 0.5, "sampling_u": 0.03, "input_length": 100, "id": "b"}, + {"timestamp": 1.0, "sampling_u": 0.01, "input_length": 200, "id": "c"}, + {"timestamp": 0.2, "sampling_u": 0.01, "input_length": 40000, "id": "x"}, + ] + trace.write_text("".join(json.dumps(row) + "\n" for row in rows)) + manifest_path = MODULE.prepare_traces(trace, tmp_path / "out", "") + manifest = json.loads(manifest_path.read_text()) + record = manifest["anchors"][MODULE.anchor_key(0.0625)] + with Path(record["path"]).open(newline="") as source: + selected = list(csv.DictReader(source)) + assert [row["source_request_id"] for row in selected] == ["b", "a", "c"] + assert [float(row["slo_ttft_ms"]) for row in selected] == [1000, 2000, 1000] + + +def test_score_uses_trace_slo_and_request_id_mapping(tmp_path: Path) -> None: + trace = tmp_path / "trace.csv" + with trace.open("w", newline="") as output: + writer = csv.DictWriter( + output, + fieldnames=["arrived_at", "num_prefill_tokens", "num_decode_tokens", "slo_ttft_ms"], + ) + writer.writeheader() + writer.writerows( + [ + {"arrived_at": 0, "num_prefill_tokens": 10, "num_decode_tokens": 1, "slo_ttft_ms": 1000}, + {"arrived_at": 1, "num_prefill_tokens": 9000, "num_decode_tokens": 1, "slo_ttft_ms": 2000}, + ] + ) + metrics = tmp_path / "request_metrics.csv" + with metrics.open("w", newline="") as output: + writer = csv.DictWriter(output, fieldnames=["Request Id", "ttft"]) + writer.writeheader() + writer.writerows([{"Request Id": 1, "ttft": 1900}, {"Request Id": 0, "ttft": 1100}]) + score = MODULE.score_request_metrics(trace, metrics) + assert score["passed_request_count"] == 1 + assert score["slo_pass_rate"] == 0.5 + assert score["feasible"] is False