Add Qwen235 Frontier latency surface runner

runs/frontier-fidelity-envelope-v1/run_frontier_qwen235_v020_surface.py

@@ -0,0 +1,233 @@
+#!/usr/bin/env python3
+"""Run Frontier on the four-cell Qwen235 vLLM 0.20 latency surface."""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import os
+import subprocess
+import sys
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+
+HERE = Path(__file__).resolve().parent
+Q30_RUNNER = HERE / "run_frontier_qwen30_exact_trace_surface.py"
+
+
+def load_q30():
+    spec = importlib.util.spec_from_file_location("q30_exact_surface", Q30_RUNNER)
+    if spec is None or spec.loader is None:
+        raise ImportError(Q30_RUNNER)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+Q30 = load_q30()
+BASE = Q30.BASE
+MODEL = "Qwen3-235B-A22B"
+
+
+@dataclass(frozen=True)
+class Config:
+    tp: int
+    mns: int
+    moe_tp: int
+    moe_ep: int
+
+    @property
+    def name(self) -> str:
+        return f"tp{self.tp}_ep{self.moe_ep}_mns{self.mns}"
+
+
+GRID = tuple(
+    Config(tp, mns, 4 if tp == 4 else 1, 1 if tp == 4 else 8)
+    for tp in (4, 8)
+    for mns in (64, 128)
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--frontier-source", type=Path, required=True)
+    parser.add_argument("--replayserve-root", type=Path, required=True)
+    parser.add_argument("--profile-root", type=Path, required=True)
+    parser.add_argument("--python-deps", type=Path, required=True)
+    parser.add_argument("--output-root", type=Path, required=True)
+    parser.add_argument("--runtime-contract", type=Path, required=True)
+    parser.add_argument("--trace", action="append", required=True)
+    parser.add_argument("--config", action="append")
+    parser.add_argument("--prefix-caching", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--allreduce-csv", type=Path, required=True)
+    parser.add_argument("--timeout-seconds", type=float, default=7200)
+    parser.add_argument("--resume", action="store_true")
+    return parser.parse_args()
+
+
+def profile_paths(root: Path) -> dict[str, Path]:
+    paths = {
+        "linear": root / "linear_op.csv",
+        "attention": root / "attention.csv",
+        "moe": root / "moe.csv",
+        "linear_kernel": root / "linear_op_kernel_only.csv",
+        "attention_kernel": root / "attention_kernel_only.csv",
+        "moe_kernel": root / "moe_kernel_only.csv",
+        "manifest": root / "manifest.json",
+    }
+    missing = [str(path) for path in paths.values() if not path.is_file()]
+    if missing:
+        raise FileNotFoundError(missing)
+    manifest = json.loads(paths["manifest"].read_text())
+    for name, path in paths.items():
+        if name == "manifest":
+            continue
+        expected = manifest.get("outputs", {}).get(path.name)
+        if expected != BASE.sha256(path):
+            raise ValueError(f"profile hash mismatch: {path}")
+    return paths
+
+
+def knobs(config: Config, paths: dict[str, Path], contract: dict, cache: Path, prefix: bool):
+    resolved = contract[config.name]
+    return {
+        "simulation_mode": "online",
+        "sys_arch": "co-location",
+        "num_replicas": 1,
+        "cluster_scheduler": "sticky_round_robin",
+        "model_name": MODEL,
+        "device": "h20",
+        "network_device": "h20_dgx",
+        "attn_tensor_parallel_size": config.tp,
+        "attn_data_parallel_size": 1,
+        "moe_tensor_parallel_size": config.moe_tp,
+        "moe_expert_parallel_size": config.moe_ep,
+        "num_pipeline_stages": 1,
+        "replica_scheduler": "vllm_v1",
+        "decode_cuda_graph_mode": "piecewise",
+        "batch_size_cap": config.mns,
+        "max_tokens_in_batch": 8192,
+        "long_prefill_token_threshold": 0,
+        "block_size": 16,
+        "num_blocks_mode": "explicit",
+        "num_blocks": int(resolved["num_gpu_blocks"]),
+        "gpu_memory_utilization": 0.80,
+        "non_kv_cache_overhead_bytes": 0,
+        "trace_max_tokens": 40960,
+        "cache_dir": str(cache / config.name),
+        "enable_prefix_caching": prefix,
+        "enable_dummy_mode": False,
+        "linear_op_input_file": str(paths["linear"]),
+        "atten_input_file": str(paths["attention"]),
+        "moe_input_file": str(paths["moe"]),
+        "linear_op_kernel_only_input_file": str(paths["linear_kernel"]),
+        "atten_kernel_only_input_file": str(paths["attention_kernel"]),
+        "moe_kernel_only_input_file": str(paths["moe_kernel"]),
+        "prediction_max_prefill_chunk_size": 8192,
+        "prediction_max_tokens_per_request": 40960,
+        "prediction_max_batch_size": max(int(v) for v in resolved["capture_sizes"]),
+        "no_cache": True,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    for name in ("frontier_source", "replayserve_root", "profile_root", "python_deps", "output_root", "runtime_contract", "allreduce_csv"):
+        setattr(args, name, getattr(args, name).resolve())
+    paths = profile_paths(args.profile_root)
+    contract = json.loads(args.runtime_contract.read_text())["configs"]
+    traces = [
+        Q30.parse_trace(spec, rate_contract="trace-window", prefix_caching=args.prefix_caching)
+        for spec in args.trace
+    ]
+    selected = list(GRID)
+    if args.config:
+        wanted = set(args.config)
+        selected = [config for config in GRID if config.name in wanted]
+        if {config.name for config in selected} != wanted:
+            raise ValueError(f"unknown configs: {wanted - {config.name for config in selected}}")
+
+    builder = BASE.load_module(
+        "qwen235_frontier_builder", args.replayserve_root / "tools/run_frontier_sweep.py"
+    )
+    environment = os.environ.copy()
+    environment.update(
+        {
+            "PYTHONPATH": ":".join([str(args.python_deps), str(args.frontier_source)]),
+            "CUDA_VISIBLE_DEVICES": "",
+            "NVIDIA_VISIBLE_DEVICES": "void",
+            "WANDB_DISABLED": "true",
+            "VIDUR_DISABLE_WANDB": "1",
+            "FRONTIER_LOG_LEVEL": "WARNING",
+            "PYTHONDONTWRITEBYTECODE": "1",
+        }
+    )
+    results = []
+    for config in selected:
+        config_knobs = knobs(config, paths, contract, args.output_root / "cache", args.prefix_caching)
+        for trace in traces:
+            run_dir = args.output_root / "runs" / config.name / trace["label"]
+            result_path = run_dir / "result.json"
+            if args.resume and result_path.is_file():
+                results.append(json.loads(result_path.read_text()))
+                continue
+            run_dir.mkdir(parents=True, exist_ok=True)
+            command = builder.build_frontier_command(
+                python_bin="/usr/bin/python3",
+                trace_file=trace["path"],
+                metrics_root=run_dir / "metrics",
+                run_id=f"qwen235_v020_{config.name}_{trace['label']}",
+                knobs=config_knobs,
+            )
+            command.extend(["--cudagraph_capture_sizes", *(str(v) for v in contract[config.name]["capture_sizes"])])
+            command = BASE.configure_cc_command(
+                command,
+                backend="vidur",
+                allreduce_csv=args.allreduce_csv,
+                cache=args.output_root / "cc-cache",
+            )
+            BASE.write_json(run_dir / "command.json", command)
+            started = time.time()
+            with (run_dir / "stdout.log").open("w") as stdout, (run_dir / "stderr.log").open("w") as stderr:
+                completed = subprocess.run(
+                    command,
+                    cwd=args.frontier_source,
+                    env=environment,
+                    stdout=stdout,
+                    stderr=stderr,
+                    timeout=args.timeout_seconds,
+                    check=False,
+                )
+            result = {
+                "config": asdict(config) | {"name": config.name},
+                "trace": Q30.trace_manifest_entry(trace),
+                "elapsed_seconds": time.time() - started,
+                "returncode": completed.returncode,
+            }
+            if completed.returncode == 0:
+                metrics = Q30.find_request_metrics(run_dir)
+                result.update(status="completed", metrics=Q30.score(metrics, trace["shapes"]), request_metrics_sha256=BASE.sha256(metrics))
+            else:
+                stderr_text = (run_dir / "stderr.log").read_text(errors="replace")
+                result.update(status="failed", failure_class=Q30.classify_frontier_failure(stderr_text))
+            BASE.write_json(result_path, result)
+            results.append(result)
+            print(json.dumps({"config": config.name, "trace": trace["label"], "status": result["status"]}, sort_keys=True), flush=True)
+
+    manifest = {
+        "schema": "qwen235-v020-frontier-latency-surface-v1",
+        "frontier_commit": subprocess.check_output(["git", "-C", str(args.frontier_source), "rev-parse", "HEAD"], text=True).strip(),
+        "profiles": {name: BASE.sha256(path) for name, path in paths.items()},
+        "runtime_contract_sha256": BASE.sha256(args.runtime_contract),
+        "prefix_caching": args.prefix_caching,
+        "results": results,
+    }
+    BASE.write_json(args.output_root / "frontier_surface.json", manifest)
+
+
+if __name__ == "__main__":
+    main()