diff options
Diffstat (limited to 'experiments/kp_raw_traffic_scaling.py')
| -rw-r--r-- | experiments/kp_raw_traffic_scaling.py | 239 |
1 files changed, 239 insertions, 0 deletions
diff --git a/experiments/kp_raw_traffic_scaling.py b/experiments/kp_raw_traffic_scaling.py new file mode 100644 index 0000000..ad1beb7 --- /dev/null +++ b/experiments/kp_raw_traffic_scaling.py @@ -0,0 +1,239 @@ +#!/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() |