299 lines
11 KiB
Python
299 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the minimal CollectiveSpec P0 collective-phase premise experiment.
|
|
|
|
P0 is deliberately *not* a performance experiment. It holds the EAGLE
|
|
drafter at Kmax=3, injects a pre-fixed per-request verification horizon after
|
|
candidate generation, and records the scheduler and worker-side phase traces.
|
|
The two policy cells are a constant-K=3 control and a deterministic
|
|
request-static heterogeneous policy.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import datetime as dt
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--base-spec", type=Path, required=True)
|
|
parser.add_argument("--source-root", type=Path, required=True)
|
|
parser.add_argument("--source-trace", type=Path, required=True)
|
|
parser.add_argument("--output-root", type=Path, required=True)
|
|
parser.add_argument("--run-id", required=True)
|
|
parser.add_argument("--request-count", type=int, default=64)
|
|
parser.add_argument("--completion-tokens", type=int, default=64)
|
|
parser.add_argument("--port", type=int, required=True)
|
|
parser.add_argument("--seed", type=int, default=20260713)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def read_rows(source_trace: Path, count: int) -> list[dict[str, Any]]:
|
|
rows: list[dict[str, Any]] = []
|
|
with source_trace.open("r", encoding="utf-8") as handle:
|
|
for raw in handle:
|
|
if not raw.strip():
|
|
continue
|
|
row = json.loads(raw)
|
|
if not isinstance(row, dict):
|
|
continue
|
|
rows.append(row)
|
|
if len(rows) == count:
|
|
break
|
|
if len(rows) != count:
|
|
raise SystemExit(f"source trace contains {len(rows)} usable rows, need {count}")
|
|
return rows
|
|
|
|
|
|
def write_workload(args: argparse.Namespace) -> tuple[Path, Path]:
|
|
rows = read_rows(args.source_trace, args.request_count)
|
|
trace_path = args.output_root / "p0_trace.jsonl"
|
|
with trace_path.open("w", encoding="utf-8") as handle:
|
|
for index, row in enumerate(rows):
|
|
row = dict(row)
|
|
# A single burst makes both DP schedulers exercise heterogeneous
|
|
# verification plans concurrently. P0 does not measure QPS.
|
|
row["request_id"] = f"p0-{index:04d}"
|
|
row["sampling_u"] = 0.0
|
|
row["temperature"] = 0.0
|
|
row["timestamp"] = 0.0
|
|
handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
|
|
windows_path = args.output_root / "p0_windows.json"
|
|
windows = {
|
|
"kind": "collectivespec_p0_fixed_burst",
|
|
"source_trace": str(args.source_trace),
|
|
"source_trace_sha256": sha256(args.source_trace),
|
|
"windows": [
|
|
{
|
|
"window_id": "collectivespec_p0_burst",
|
|
"trace_file": str(trace_path),
|
|
"trace_type": "decode_only",
|
|
"window_start": 0.0,
|
|
"window_end": 1.0,
|
|
"num_requests": args.request_count,
|
|
"sampling_strategy": "all_requests_fixed_burst",
|
|
}
|
|
],
|
|
}
|
|
windows_path.write_text(
|
|
json.dumps(windows, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
return trace_path, windows_path
|
|
|
|
|
|
def derived_spec(
|
|
base: dict[str, Any],
|
|
*,
|
|
args: argparse.Namespace,
|
|
windows_path: Path,
|
|
policy_name: str,
|
|
) -> dict[str, Any]:
|
|
spec = json.loads(json.dumps(base))
|
|
spec["study_id"] = f"collectivespec-p0-{policy_name}-{args.run_id}"
|
|
spec["engine"]["cwd"] = str(args.source_root)
|
|
spec["engine"]["port"] = args.port
|
|
flags = spec["engine"]["base_flags"]
|
|
flags["port"] = args.port
|
|
flags["seed"] = args.seed
|
|
flags["scheduler-cls"] = "collectivespec_p0.scheduler.P0Scheduler"
|
|
flags["worker-cls"] = "collectivespec_p0.worker.P0Worker"
|
|
raw_speculative = flags.get("speculative-config")
|
|
if not isinstance(raw_speculative, str):
|
|
raise SystemExit("base spec must have speculative-config for P0")
|
|
speculative = json.loads(raw_speculative)
|
|
speculative["num_speculative_tokens"] = 3
|
|
flags["speculative-config"] = json.dumps(speculative, separators=(",", ":"))
|
|
|
|
trace = spec["trace"]
|
|
trace.update(
|
|
{
|
|
"completion_tokens_override": args.completion_tokens,
|
|
"early_stop_max_elapsed_s": 1800.0,
|
|
"early_stop_max_lag_s": None,
|
|
"max_concurrency": args.request_count,
|
|
"max_requests_per_probe": None,
|
|
"replay_time_scale": 1.0,
|
|
"request_mode": "decode_only",
|
|
"timestamp_field": "timestamp",
|
|
"u_field": "sampling_u",
|
|
"window_id": "collectivespec_p0_burst",
|
|
"windows_path": str(windows_path),
|
|
}
|
|
)
|
|
spec["slo"] = {"target_pass_rate": 1.0}
|
|
spec["search"] = {
|
|
"low": 0.0,
|
|
"high": 1.0,
|
|
"tolerance": 1e-6,
|
|
"max_probes": 1,
|
|
"sample_seed": args.seed,
|
|
}
|
|
return spec
|
|
|
|
|
|
def policy_payload(policy_name: str, seed: int) -> dict[str, Any]:
|
|
if policy_name == "control_k3":
|
|
return {
|
|
"version": "collectivespec-p0-v1-control-k3",
|
|
"seed": seed,
|
|
"max_k": 3,
|
|
"policy": "constant-k",
|
|
"constant_k": 3,
|
|
}
|
|
if policy_name == "heterogeneous":
|
|
return {
|
|
"version": "collectivespec-p0-v1-heterogeneous",
|
|
"seed": seed,
|
|
"max_k": 3,
|
|
"policy": "blake2b-request-static-k",
|
|
}
|
|
raise ValueError(f"unknown P0 policy {policy_name}")
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
if args.request_count < 8:
|
|
raise SystemExit("request-count must be at least 8 to exercise all K values")
|
|
if args.completion_tokens < 8:
|
|
raise SystemExit("completion-tokens must be at least 8")
|
|
if not args.base_spec.is_file():
|
|
raise SystemExit(f"base spec does not exist: {args.base_spec}")
|
|
if not args.source_root.is_dir():
|
|
raise SystemExit(f"source root does not exist: {args.source_root}")
|
|
if not args.source_trace.is_file():
|
|
raise SystemExit(f"source trace does not exist: {args.source_trace}")
|
|
overlay = args.source_root / "scripts" / "collectivespec" / "p0_overlay"
|
|
if not overlay.is_dir():
|
|
raise SystemExit(f"P0 overlay does not exist: {overlay}")
|
|
|
|
args.output_root.mkdir(parents=True, exist_ok=True)
|
|
trace_path, windows_path = write_workload(args)
|
|
base = json.loads(args.base_spec.read_text(encoding="utf-8"))
|
|
manifest = {
|
|
"kind": "collectivespec_p0_phase_trace",
|
|
"created_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
|
|
"base_spec": str(args.base_spec),
|
|
"base_spec_sha256": sha256(args.base_spec),
|
|
"source_root": str(args.source_root),
|
|
"source_trace": str(args.source_trace),
|
|
"source_trace_sha256": sha256(args.source_trace),
|
|
"materialized_trace": str(trace_path),
|
|
"materialized_trace_sha256": sha256(trace_path),
|
|
"windows_path": str(windows_path),
|
|
"overlay": str(overlay),
|
|
"run_id": args.run_id,
|
|
"topology_claim": "inherited_from_base_spec",
|
|
"request_count": args.request_count,
|
|
"completion_tokens": args.completion_tokens,
|
|
"policies": ["control_k3", "heterogeneous"],
|
|
"seed": args.seed,
|
|
"port": args.port,
|
|
"python": sys.executable,
|
|
}
|
|
(args.output_root / "manifest.json").write_text(
|
|
json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
|
|
)
|
|
|
|
results: list[dict[str, Any]] = []
|
|
for ordinal, policy_name in enumerate(manifest["policies"], start=1):
|
|
cell = args.output_root / policy_name
|
|
cell.mkdir(exist_ok=True)
|
|
policy_path = cell / "policy.json"
|
|
policy_path.write_text(
|
|
json.dumps(policy_payload(policy_name, args.seed), indent=2, sort_keys=True) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
spec_path = cell / "study_spec.json"
|
|
spec_path.write_text(
|
|
json.dumps(
|
|
derived_spec(
|
|
base,
|
|
args=args,
|
|
windows_path=windows_path,
|
|
policy_name=policy_name,
|
|
),
|
|
indent=2,
|
|
sort_keys=True,
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
store_root = cell / "store"
|
|
logs = cell / "p0_logs"
|
|
command = [
|
|
sys.executable,
|
|
"-m",
|
|
"aituner.cli",
|
|
"study",
|
|
"tune",
|
|
"--spec",
|
|
str(spec_path),
|
|
"--store-root",
|
|
str(store_root),
|
|
"--max-trials",
|
|
"1",
|
|
]
|
|
record = {
|
|
"ordinal": ordinal,
|
|
"policy": policy_name,
|
|
"policy_path": str(policy_path),
|
|
"spec_path": str(spec_path),
|
|
"p0_logs": str(logs),
|
|
"command": command,
|
|
}
|
|
print(json.dumps(record, sort_keys=True), flush=True)
|
|
if args.dry_run:
|
|
results.append({**record, "returncode": None})
|
|
continue
|
|
env = os.environ.copy()
|
|
inherited = env.get("PYTHONPATH", "")
|
|
env["PYTHONPATH"] = os.pathsep.join(
|
|
part for part in (str(overlay), str(args.source_root / "src"), inherited) if part
|
|
)
|
|
env["PYTHONDONTWRITEBYTECODE"] = "1"
|
|
env["COLLECTIVESPEC_P0_LOG_DIR"] = str(logs)
|
|
env["COLLECTIVESPEC_P0_POLICY_PATH"] = str(policy_path)
|
|
log_path = cell / "driver.log"
|
|
with log_path.open("w", encoding="utf-8") as handle:
|
|
completed = subprocess.run(
|
|
command,
|
|
cwd=args.source_root,
|
|
env=env,
|
|
text=True,
|
|
stdout=handle,
|
|
stderr=subprocess.STDOUT,
|
|
)
|
|
results.append({**record, "driver_log": str(log_path), "returncode": completed.returncode})
|
|
if completed.returncode:
|
|
break
|
|
result = {
|
|
"finished_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
|
|
"results": results,
|
|
"status": "ok" if results and all(item["returncode"] in (0, None) for item in results) else "failed",
|
|
}
|
|
(args.output_root / "driver_result.json").write_text(
|
|
json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8"
|
|
)
|
|
return 0 if result["status"] == "ok" else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|