path: root/experiments/kp_raw_traffic_scaling.py

#!/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()