diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/community_prefill_grid.py b/runs/frontier-multicase-sufficiency-v0/best_effort/community_prefill_grid.py new file mode 100644 index 0000000..ce347ac --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/community_prefill_grid.py @@ -0,0 +1,448 @@ +#!/usr/bin/env python3 +"""Run the community-vLLM side of the frozen Qwen235B prefill grid.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import random +import subprocess +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + + +SCHEMA = "community-vllm-qwen235b-prefill-grid-v1" +MODEL = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-235B-A22B-FP8" +SERVED_MODEL = "qwen3-235b-community-prefill" +TRACE = ( + "/home/admin/cpfs/wjh/aituner/aituner/trace_windows/traces/" + "thinking_w20260327_1000.jsonl" +) +WINDOWS = "/home/admin/cpfs/wjh/aituner/aituner/trace_windows/windows.json" +TRACE_SHA256 = "f878e9af18f94dcfaced94a8e1e6b20a2f7d97d64aa862448025660dbbd965b2" +ORDER_SEED = 20260715 + + +@dataclass(frozen=True) +class GridConfig: + tp: int + mns: int + mbt: int + expert_parallel: bool + num_gpu_blocks: int + + @property + def name(self) -> str: + return f"tp{self.tp}_mns{self.mns}_mbt{self.mbt}" + + +GRID = tuple( + GridConfig( + tp=tp, + mns=mns, + mbt=mbt, + expert_parallel=tp == 8, + 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 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 search_spec() -> dict[str, Any]: + return { + "low": 0.0, + "high": 0.125, + "tolerance": 0.001, + "max_probes": 6, + "sample_seed": 20260325, + "inherit_incumbent_floor": False, + "auto_high": { + "enabled": False, + "max_sampling_u": 1.0, + "require_human_confirmation_beyond_trace": True, + }, + } + + +def study_payload( + *, + tp: int, + repo: Path, + python: Path, + vllm: Path, + model: Path, + trace: Path, + windows: Path, + port: int, +) -> dict[str, Any]: + base_flags: dict[str, Any] = { + "host": "127.0.0.1", + "port": port, + "served-model-name": SERVED_MODEL, + "tensor-parallel-size": tp, + "disable-custom-all-reduce": True, + "quantization": "fp8", + "gpu-memory-utilization": 0.80, + "num-gpu-blocks-override": 26101 if tp == 4 else 62351, + "kv-cache-dtype": "auto", + "max-model-len": 40960, + "max-num-batched-tokens": 8192, + "max-num-seqs": 64, + "enable-prefix-caching": False, + "enable-chunked-prefill": True, + "enforce-eager": True, + "disable-log-requests": True, + } + if tp == 8: + base_flags["enable-expert-parallel"] = True + return { + "study_id": f"community-qwen235b-prefill-tp{tp}-v1", + "hardware": { + "gpu_count": tp, + "gpu_model": "NVIDIA H20", + "host_candidates": ["dash0"], + }, + "model": {"model_id": str(model), "served_model_name": SERVED_MODEL}, + "engine": { + "engine_name": "vllm", + "engine_version": "0.10.2-community", + "exec_path": str(vllm), + "cwd": str(repo), + "host": "127.0.0.1", + "port": port, + "ready_timeout_s": 1800, + "request_timeout_s": 1800, + "healthcheck_path": "/v1/models", + "launch_args": ["serve", str(model)], + "base_envs": { + "CUDA_VISIBLE_DEVICES": ",".join(str(index) for index in range(tp)), + }, + "base_flags": base_flags, + "tunable_envs": [], + "tunable_flags": ["max-num-seqs", "max-num-batched-tokens"], + "python_executable": str(python), + }, + "trace": { + "windows_path": str(windows), + "window_id": "thinking_w20260327_1000", + "trace_file_override": str(trace), + "request_mode": "raw_completion", + "completion_tokens_override": 1, + "u_field": "sampling_u", + "timestamp_field": "timestamp", + "max_concurrency": 64, + "input_length_filter": { + "min_input_tokens": 0, + "max_input_tokens": 32768, + }, + "replay_time_scale": 1.0, + "early_stop_max_lag_s": 180.0, + "early_stop_max_elapsed_s": 1200.0, + "restart_engine_after_early_stop": False, + "adaptive_stop": {"enabled": False}, + }, + "slo": { + "target_pass_rate": 0.95, + "ttft_rule": { + "kind": "step_ms", + "buckets": [ + {"max_input_tokens": 8191, "threshold_ms": 1000}, + {"max_input_tokens": 32767, "threshold_ms": 2000}, + {"threshold_ms": 2000}, + ], + }, + }, + "search": search_spec(), + "llm": {"use_harness": False}, + } + + +def git_fingerprint(repo: Path) -> dict[str, Any]: + def run(*args: str) -> str: + return subprocess.run( + ["git", *args], + cwd=repo, + check=True, + stdout=subprocess.PIPE, + text=True, + ).stdout.strip() + + return { + "commit": run("rev-parse", "HEAD"), + "status_porcelain": run("status", "--porcelain=v1").splitlines(), + } + + +def prepare(args: argparse.Namespace) -> None: + repo = args.repo.resolve() + python = args.python.resolve() + vllm = args.vllm.resolve() + model = args.model.resolve() + trace = args.trace.resolve() + windows = args.windows.resolve() + for path in (repo, python, vllm, model, trace, windows): + if not path.exists(): + raise FileNotFoundError(path) + actual_trace_sha = sha256(trace) + if actual_trace_sha != args.expected_trace_sha256: + raise ValueError( + f"trace SHA256 mismatch: expected={args.expected_trace_sha256}, " + f"actual={actual_trace_sha}" + ) + output = args.output_root.resolve() + output.mkdir(parents=True, exist_ok=True) + studies: dict[int, dict[str, str]] = {} + for tp, port in ((4, args.port), (8, args.port)): + study_dir = output / "studies" / f"tp{tp}" + study_path = study_dir / "study.json" + pointer_path = study_dir / "study_spec.source" + write_json( + study_path, + study_payload( + tp=tp, + repo=repo, + python=python, + vllm=vllm, + model=model, + trace=trace, + windows=windows, + port=port, + ), + ) + pointer_path.write_text(str(study_path) + "\n") + studies[tp] = { + "study_path": str(study_path), + "study_sha256": sha256(study_path), + "pointer_path": str(pointer_path), + } + + order = list(GRID) + random.Random(args.order_seed).shuffle(order) + trials = [] + for index, config in enumerate(order, start=1): + trial_dir = output / "trials" / config.name + trial_path = trial_dir / "trial_spec.json" + trial = { + "study_id": f"community-qwen235b-prefill-tp{config.tp}-v1", + "trial_id": config.name, + "config_patch": { + "env_patch": {}, + "flag_patch": { + "max-num-seqs": config.mns, + "max-num-batched-tokens": config.mbt, + }, + }, + "search": search_spec(), + "study_spec_path": studies[config.tp]["pointer_path"], + "artifact_dir": str(trial_dir), + "probe_log_path": str(trial_dir / "probe_history.json"), + "engine_log_path": str(trial_dir / "engine.log"), + "result_path": str(trial_dir / "result.json"), + "search_evidence": { + "enabled": False, + "original_high": 0.125, + "effective_high": 0.125, + "trace_max_sampling_u": None, + "max_sampling_u": 1.0, + "require_human_confirmation_beyond_trace": True, + "reason": "auto_high_disabled", + }, + } + write_json(trial_path, trial) + trials.append( + { + "execution_index": index, + "config": asdict(config) | {"name": config.name}, + "trial_spec_path": str(trial_path), + "trial_spec_sha256": sha256(trial_path), + } + ) + manifest = { + "schema": SCHEMA, + "created_unix_s": time.time(), + "repository": git_fingerprint(repo), + "runtime": {"python": str(python), "vllm": str(vllm)}, + "model": { + "path": str(model), + "config_sha256": sha256(model / "config.json"), + }, + "trace": {"path": str(trace), "sha256": actual_trace_sha}, + "contract": { + "metric": "maximum SLO-feasible request_rate_per_gpu", + "request_mode": "raw_completion", + "completion_tokens_override": 1, + "target_pass_rate": 0.95, + "search": search_spec(), + "prefix_caching": False, + "chunked_prefill": True, + "cuda_graphs": False, + "custom_all_reduce": False, + "kv_blocks_source": "community_vllm_measured_and_fixed_per_topology", + }, + "order": {"method": "python_random_shuffle", "seed": args.order_seed}, + "studies": studies, + "trials": trials, + } + write_json(output / "run_manifest.json", manifest) + print(output / "run_manifest.json") + + +def run(args: argparse.Namespace) -> None: + manifest = json.loads(args.manifest.read_text()) + if manifest.get("schema") != SCHEMA: + raise ValueError(f"unexpected manifest schema: {manifest.get('schema')}") + repo = args.repo.resolve() + env = os.environ.copy() + env["PYTHONPATH"] = str(repo / "src") + for trial in manifest["trials"]: + trial_path = Path(trial["trial_spec_path"]) + trial_dir = trial_path.parent + result_path = trial_dir / "result.json" + if result_path.is_file(): + result = json.loads(result_path.read_text()) + if result.get("status") == "completed": + print(json.dumps({"config": trial["config"]["name"], "skipped": True})) + continue + command = [ + str(args.python.resolve()), + "-m", + "aituner.cli", + "worker", + "run-trial", + "--trial-spec", + str(trial_path), + ] + write_json(trial_dir / "worker_command.json", command) + started = time.time() + with (trial_dir / "worker.log").open("w", encoding="utf-8") as output: + completed = subprocess.run( + command, + cwd=repo, + env=env, + stdout=output, + stderr=subprocess.STDOUT, + check=False, + ) + record = { + "config": trial["config"]["name"], + "returncode": completed.returncode, + "elapsed_seconds": time.time() - started, + } + print(json.dumps(record), flush=True) + if completed.returncode != 0: + raise RuntimeError(f"community trial failed: {record}") + assemble(args.manifest) + + +def capacity_interval(probes: list[dict[str, Any]]) -> list[float]: + low, high = 0.0, 0.125 + for probe in probes: + midpoint = float(probe["threshold"]) + if probe["feasible"]: + low = midpoint + else: + high = midpoint + return [low, high] + + +def assemble(manifest_path: Path) -> Path: + manifest = json.loads(manifest_path.read_text()) + records = [] + for trial in manifest["trials"]: + result_path = Path(trial["trial_spec_path"]).parent / "result.json" + if not result_path.is_file(): + raise FileNotFoundError(result_path) + result = json.loads(result_path.read_text()) + if result.get("status") != "completed": + raise ValueError(f"trial did not complete: {result_path}") + tp = int(trial["config"]["tp"]) + request_rate = result.get("best_request_rate") + records.append( + { + "config": trial["config"], + "result_path": str(result_path), + "result_sha256": sha256(result_path), + "capacity_interval_sampling_u": capacity_interval(result["probes"]), + "best_sampling_u": result.get("best_sampling_u"), + "best_request_rate": request_rate, + "best_request_rate_per_gpu": ( + float(request_rate) / tp if request_rate is not None else None + ), + "best_pass_rate": result.get("best_pass_rate"), + } + ) + records.sort( + key=lambda item: ( + -(item["best_request_rate_per_gpu"] or -1.0), + item["config"]["name"], + ) + ) + freeze = { + "schema": SCHEMA, + "run_manifest_sha256": sha256(manifest_path), + "ranking": [dict(record, rank=index + 1) for index, record in enumerate(records)], + } + path = manifest_path.parent / "community_ranking_frozen.json" + write_json(path, freeze) + print(path) + return path + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(dest="command", required=True) + prepare_parser = subparsers.add_parser("prepare") + prepare_parser.add_argument("--repo", type=Path, required=True) + prepare_parser.add_argument("--python", type=Path, required=True) + prepare_parser.add_argument("--vllm", type=Path, required=True) + prepare_parser.add_argument("--model", type=Path, default=Path(MODEL)) + prepare_parser.add_argument("--trace", type=Path, default=Path(TRACE)) + prepare_parser.add_argument("--windows", type=Path, default=Path(WINDOWS)) + prepare_parser.add_argument("--expected-trace-sha256", default=TRACE_SHA256) + prepare_parser.add_argument("--output-root", type=Path, required=True) + prepare_parser.add_argument("--order-seed", type=int, default=ORDER_SEED) + prepare_parser.add_argument("--port", type=int, default=18918) + run_parser = subparsers.add_parser("run") + run_parser.add_argument("--repo", type=Path, required=True) + run_parser.add_argument("--python", type=Path, required=True) + run_parser.add_argument("--manifest", type=Path, required=True) + assemble_parser = subparsers.add_parser("assemble") + assemble_parser.add_argument("--manifest", type=Path, required=True) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.command == "prepare": + prepare(args) + elif args.command == "run": + run(args) + elif args.command == "assemble": + assemble(args.manifest) + else: + raise AssertionError(args.command) + + +if __name__ == "__main__": + main() diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/test_community_prefill_grid.py b/runs/frontier-multicase-sufficiency-v0/best_effort/test_community_prefill_grid.py new file mode 100644 index 0000000..0587fad --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/test_community_prefill_grid.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("community_prefill_grid.py") +SPEC = importlib.util.spec_from_file_location("community_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_grid_matches_frozen_frontier_surface() -> None: + assert len(MODULE.GRID) == 8 + assert {item.name for item in MODULE.GRID} == { + f"tp{tp}_mns{mns}_mbt{mbt}" + for tp in (4, 8) + for mns in (64, 128) + for mbt in (8192, 16384) + } + assert {item.num_gpu_blocks for item in MODULE.GRID if item.tp == 4} == {26101} + assert {item.num_gpu_blocks for item in MODULE.GRID if item.tp == 8} == {62351} + + +def test_study_uses_raw_prompt_and_fixed_serving_contract(tmp_path: Path) -> None: + payload = MODULE.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["trace"]["request_mode"] == "raw_completion" + assert payload["trace"]["completion_tokens_override"] == 1 + assert payload["trace"]["replay_time_scale"] == 1.0 + assert payload["engine"]["base_flags"]["enable-expert-parallel"] is True + assert payload["engine"]["base_flags"]["disable-custom-all-reduce"] is True + assert payload["engine"]["base_flags"]["num-gpu-blocks-override"] == 62351 + assert payload["engine"]["base_flags"]["enable-prefix-caching"] is False + + +def test_capacity_interval_replays_binary_decisions() -> None: + probes = [ + {"threshold": 0.0625, "feasible": False}, + {"threshold": 0.03125, "feasible": False}, + {"threshold": 0.015625, "feasible": True}, + {"threshold": 0.0234375, "feasible": False}, + {"threshold": 0.01953125, "feasible": True}, + {"threshold": 0.021484375, "feasible": True}, + ] + assert MODULE.capacity_interval(probes) == [0.021484375, 0.0234375] + + +def test_manifest_json_round_trip(tmp_path: Path) -> None: + path = tmp_path / "x.json" + MODULE.write_json(path, {"schema": MODULE.SCHEMA}) + assert json.loads(path.read_text()) == {"schema": MODULE.SCHEMA}