#!/usr/bin/env python3 """Run an isolated, reproducible static-speculation screening study. This driver deliberately does not implement CollectiveSpec. It establishes whether the current static-K deployment has enough headroom to justify such a system. Each K gets an immutable derived StudySpec and a separate store. """ 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("--output-root", type=Path, required=True) parser.add_argument("--run-id", required=True) parser.add_argument("--k-values", nargs="+", type=int, default=[0, 1, 2, 3]) parser.add_argument( "--order", nargs="+", type=int, help="Sequential launch order. Defaults to --k-values order.", ) parser.add_argument("--sampling-u-low", type=float, default=0.005) parser.add_argument("--sampling-u-high", type=float, default=0.020) parser.add_argument("--tolerance", type=float, default=0.003) parser.add_argument("--max-probes", type=int, default=3) parser.add_argument( "--completion-tokens-override", type=int, help="Force every request to this exact completion length for a controlled screen.", ) parser.add_argument( "--trace-windows-path", type=Path, help="Override trace.windows_path in every derived StudySpec.", ) parser.add_argument( "--trace-window-id", help="Override trace.window_id in every derived StudySpec.", ) parser.add_argument( "--seed", type=int, help="Explicit engine seed recorded in every derived StudySpec.", ) parser.add_argument( "--port", type=int, help="Serving port. Defaults to the base StudySpec's engine/base flag port.", ) parser.add_argument( "--continue-on-error", action="store_true", help="Record a failed K and proceed to the next independent K.", ) parser.add_argument("--dry-run", action="store_true") return parser.parse_args() def derived_spec( base: dict[str, Any], args: argparse.Namespace, k: int ) -> dict[str, Any]: spec = json.loads(json.dumps(base)) label = f"collectivespec-static-k{k}-screen-{args.run_id}" spec["study_id"] = label spec["engine"]["cwd"] = str(args.source_root) flags = spec["engine"]["base_flags"] port = args.port if port is None: port = int(flags.get("port", spec["engine"]["port"])) spec["engine"]["port"] = port flags["port"] = port if k == 0: # vLLM's EAGLE config rejects num_speculative_tokens=0. Removing the # entire flag is the real no-speculation baseline. flags.pop("speculative-config", None) else: raw_config = flags.get("speculative-config") if raw_config is None: raise ValueError("base spec has no speculative-config for K > 0") config = json.loads(raw_config) config["num_speculative_tokens"] = k flags["speculative-config"] = json.dumps(config, separators=(",", ":")) search = spec["search"] search["low"] = args.sampling_u_low search["high"] = args.sampling_u_high search["tolerance"] = args.tolerance search["max_probes"] = args.max_probes if args.completion_tokens_override is not None: spec["trace"]["completion_tokens_override"] = args.completion_tokens_override if args.trace_windows_path is not None: spec["trace"]["windows_path"] = str(args.trace_windows_path) if args.trace_window_id is not None: spec["trace"]["window_id"] = args.trace_window_id if args.seed is not None: flags["seed"] = args.seed return spec def main() -> int: args = parse_args() 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 any(k < 0 for k in args.k_values): raise SystemExit("K must be non-negative") if not 0 <= args.sampling_u_low < args.sampling_u_high <= 1: raise SystemExit("sampling-u bounds must satisfy 0 <= low < high <= 1") if args.max_probes < 1: raise SystemExit("max-probes must be positive") if args.completion_tokens_override is not None and args.completion_tokens_override < 0: raise SystemExit("completion token override must be non-negative") if (args.trace_windows_path is None) != (args.trace_window_id is None): raise SystemExit( "--trace-windows-path and --trace-window-id must be provided together" ) if args.trace_windows_path is not None and not args.trace_windows_path.is_file(): raise SystemExit(f"trace windows file does not exist: {args.trace_windows_path}") order = args.order or args.k_values if sorted(order) != sorted(args.k_values) or len(order) != len(args.k_values): raise SystemExit("--order must be a permutation of --k-values") base = json.loads(args.base_spec.read_text()) args.output_root.mkdir(parents=True, exist_ok=True) specs_dir = args.output_root / "specs" specs_dir.mkdir(exist_ok=True) manifest = { "kind": "collectivespec_static_k_screen", "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), "run_id": args.run_id, "k_values": args.k_values, "launch_order": order, "sampling_u": [args.sampling_u_low, args.sampling_u_high], "tolerance": args.tolerance, "max_probes": args.max_probes, "completion_tokens_override": args.completion_tokens_override, "trace_windows_path": ( str(args.trace_windows_path) if args.trace_windows_path is not None else None ), "trace_window_id": args.trace_window_id, "seed": args.seed, "port": args.port if args.port is not None else int(base["engine"]["base_flags"].get("port", base["engine"]["port"])), "python": sys.executable, } (args.output_root / "manifest.json").write_text( json.dumps(manifest, indent=2, sort_keys=True) + "\n" ) failures: list[dict[str, Any]] = [] for ordinal, k in enumerate(order, start=1): spec = derived_spec(base, args, k) spec_path = specs_dir / f"{ordinal:02d}_k{k}.json" spec_path.write_text(json.dumps(spec, indent=2, sort_keys=True) + "\n") store_root = args.output_root / "stores" / f"k{k}" command = [ sys.executable, "-m", "aituner.cli", "study", "tune", "--spec", str(spec_path), "--store-root", str(store_root), "--max-trials", "1", ] print(f"[{ordinal}/{len(order)}] K={k} command={json.dumps(command)}", flush=True) if args.dry_run: continue log_path = args.output_root / "logs" / f"{ordinal:02d}_k{k}.log" log_path.parent.mkdir(exist_ok=True) env = os.environ.copy() env["PYTHONPATH"] = str(args.source_root / "src") env["PYTHONDONTWRITEBYTECODE"] = "1" with log_path.open("w") as log: completed = subprocess.run( command, cwd=args.source_root, env=env, text=True, stdout=log, stderr=subprocess.STDOUT, ) print( f"[{ordinal}/{len(order)}] K={k} returncode={completed.returncode} log={log_path}", flush=True, ) if completed.returncode: failures.append({"k": k, "returncode": completed.returncode, "log": str(log_path)}) if not args.continue_on_error: break result = { "finished_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(), "failures": failures, "status": "ok" if not failures else "completed_with_failures", } (args.output_root / "driver_result.json").write_text( json.dumps(result, indent=2, sort_keys=True) + "\n" ) return 0 if not failures else 1 if __name__ == "__main__": raise SystemExit(main())