Add reproducible community vLLM prefill grid
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user