238 lines
8.7 KiB
Python
238 lines
8.7 KiB
Python
#!/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())
|