Add reproducible CollectiveSpec opportunity screen
This commit is contained in:
237
scripts/collectivespec/run_static_k_pilot.py
Normal file
237
scripts/collectivespec/run_static_k_pilot.py
Normal file
@@ -0,0 +1,237 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user