704 lines
24 KiB
Python
704 lines
24 KiB
Python
#!/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()
|