#!/usr/bin/env python3 """Frozen three-depth raw-KP control under four-RMS mixed traffic.""" import argparse import hashlib import json import os import subprocess import sys import time ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, ROOT) from experiments.crossover_hardware import ( # noqa: E402 DEFAULT_PROFILE, physical_gpu_report, policy_for_report, profile_choices, ) PROTOCOL = os.path.join(ROOT, "KP_RAW_TRAFFIC_SCALING.md") RESULT_ROOT = os.path.join(ROOT, "results", "kp_raw_traffic_scaling") DEPTHS = (20, 32, 56) RUN_ORDER = (56, 32, 20) def sha256(path): digest = hashlib.sha256() with open(path, "rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def git_output(*args): return subprocess.run( ["git", *args], cwd=ROOT, check=True, capture_output=True, text=True ).stdout.strip() def source_report(): if git_output("status", "--porcelain", "--untracked-files=no"): raise RuntimeError("raw-KP scaling requires clean tracked source") paths = ( PROTOCOL, os.path.abspath(__file__), os.path.join(ROOT, "experiments", "kp_raw_traffic_scaling_smoke.py"), os.path.join(ROOT, "experiments", "analyze_kp_raw_traffic_scaling.py"), os.path.join(ROOT, "experiments", "conv_run.py"), os.path.join(ROOT, "experiments", "crossover_hardware.py"), os.path.join(ROOT, "sdil", "conv.py"), os.path.join(ROOT, "sdil", "data.py"), ) for path in paths: relative = os.path.relpath(path, ROOT) subprocess.run( ["git", "ls-files", "--error-unmatch", relative], cwd=ROOT, check=True, capture_output=True, ) return { "git_commit": git_output("rev-parse", "HEAD"), "tracked_files": { os.path.relpath(path, ROOT): sha256(path) for path in paths }, } def raw_job(depth): name = f"raw-kp-traffic4-r1-d{depth}" output = os.path.join(RESULT_ROOT, "r1", name + ".json") command = [ sys.executable, "experiments/conv_run.py", "--mode", "kp_traffic", "--traffic_rule", "raw", "--out", output, "--device", "cuda", "--depth", str(depth), "--width", "16", "--seed", "0", "--loader_seed", "0", "--split_seed", "2027", "--batch_size", "128", "--epochs", "200", "--train_limit", "0", "--val_examples", "5000", "--eval_split", "validation", "--eval_every", "1", "--augment_train", "1", "--lr", "0.1", "--output_lr", "0.1", "--lr_schedule", "step", "--lr_milestones", "100,150", "--lr_gamma", "0.1", "--momentum", "0.9", "--weight_decay", "1e-4", "--normalization", "batchnorm", "--a_scale", "1", "--alignment_probe", "32", "--learn_P", "1", "--eta_P", "0.1", "--predictor_mode", "closed_form", "--predictor_warmup_steps", "1", "--predictor_every", "0", "--neutral_projection", "1", "--traffic_seed", "5000", "--traffic_ratio", "4", "--traffic_calibration_examples", "64", ] return { "stage": "raw_kp_traffic_scaling_r1", "method": "raw_kp_traffic4", "architecture": f"resnet{depth}", "depth": depth, "experiment_name": name, "output": output, "timeout_seconds": 48 * 60 * 60, "command": command, } def jobs(): values = [raw_job(depth) for depth in RUN_ORDER] if {job["depth"] for job in values} != set(DEPTHS): raise AssertionError("raw-KP scaling registry must cover three depths") return values def registry_sha256(values): encoded = json.dumps( values, sort_keys=True, separators=(",", ":") ).encode("utf-8") return hashlib.sha256(encoded).hexdigest() def ensure_launch(source, values, hardware_policy): path = os.path.join(RESULT_ROOT, "r1_launch.json") expected = { "stage": "raw_kp_traffic_scaling_r1", "source": source, "registry_sha256": registry_sha256(values), "num_jobs": len(values), "hardware_policy": hardware_policy, } if os.path.isfile(path): with open(path, encoding="utf-8") as handle: if json.load(handle) != expected: raise RuntimeError("raw-KP launch lock drift") return path os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, "w", encoding="utf-8") as handle: json.dump(expected, handle, indent=2, sort_keys=True) handle.write("\n") return path def assert_source_unchanged(source): if git_output("status", "--porcelain", "--untracked-files=no"): raise RuntimeError("tracked source changed after raw-KP launch") if git_output("rev-parse", "HEAD") != source["git_commit"]: raise RuntimeError("source commit changed after raw-KP launch") for relative, expected in source["tracked_files"].items(): if sha256(os.path.join(ROOT, relative)) != expected: raise RuntimeError(f"raw-KP source drift: {relative}") def run_job(job, source, gpu, dry_run): manifest_path = job["output"] + ".manifest.json" if os.path.exists(manifest_path): print(f"preserving {job['experiment_name']}", flush=True) return if os.path.exists(job["output"]): raise RuntimeError(f"orphaned raw-KP output: {job['output']}") print("RUN", " ".join(job["command"]), flush=True) if dry_run: return assert_source_unchanged(source) os.makedirs(os.path.dirname(job["output"]), exist_ok=True) started = time.time() try: result = subprocess.run( job["command"], cwd=ROOT, timeout=job["timeout_seconds"] ) return_code = result.returncode status = "completed" if return_code == 0 else "nonzero_exit" except subprocess.TimeoutExpired: return_code = None status = "timeout" output_exists = os.path.isfile(job["output"]) if status == "completed" and not output_exists: status = "missing_output" manifest = { **job, "source": source, "hardware_lock": gpu, "status": status, "return_code": return_code, "output_exists": output_exists, "output_sha256": sha256(job["output"]) if output_exists else None, "driver_wall_seconds": time.time() - started, "completed_unix_time": time.time(), "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), } with open(manifest_path, "w", encoding="utf-8") as handle: json.dump(manifest, handle, indent=2, sort_keys=True) handle.write("\n") print(f"DONE status={status} {job['experiment_name']}", flush=True) def main(): parser = argparse.ArgumentParser() parser.add_argument("--depth", type=int, choices=DEPTHS) parser.add_argument( "--hardware-profile", choices=profile_choices(), default=os.environ.get("SDIL_HARDWARE_PROFILE", DEFAULT_PROFILE), ) parser.add_argument("--dry-run", action="store_true") args = parser.parse_args() source = source_report() gpu = physical_gpu_report(args.hardware_profile, args.dry_run) complete = jobs() if not args.dry_run: launch = ensure_launch( source, complete, policy_for_report(args.hardware_profile, gpu) ) print(f"raw-KP launch lock: {launch}", flush=True) selected = complete if args.depth is not None: selected = [job for job in complete if job["depth"] == args.depth] for job in selected: run_job(job, source, gpu, args.dry_run) if __name__ == "__main__": main()