456 lines
15 KiB
Python
456 lines
15 KiB
Python
#!/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-v2"
|
|
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
|
|
CLIENT_MAX_CONCURRENCY = 256
|
|
|
|
|
|
@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}-v2",
|
|
"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",
|
|
# Keep the request generator above the largest server-side MNS so
|
|
# MNS=128 is exercised by vLLM rather than clipped by the client.
|
|
# The extra headroom also leaves an explicit engine waiting queue
|
|
# under overload instead of moving that queue into the replay loop.
|
|
"max_concurrency": CLIENT_MAX_CONCURRENCY,
|
|
"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}-v2",
|
|
"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,
|
|
"client_max_concurrency": CLIENT_MAX_CONCURRENCY,
|
|
"max_server_mns": max(config.mns for config in GRID),
|
|
"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()
|