#!/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())