exp: freeze raw KP traffic scaling control

3 files changed, 454 insertions, 0 deletions

diff --git a/experiments/analyze_kp_raw_traffic_scaling.py b/experiments/analyze_kp_raw_traffic_scaling.py
new file mode 100644
index 0000000..371406f
--- /dev/null
+++ b/experiments/analyze_kp_raw_traffic_scaling.py
@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+"""Audit the three-depth raw-KP mixed-traffic control."""
+import argparse
+import hashlib
+import json
+import math
+import os
+import sys
+
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ROOT)
+from experiments.crossover_hardware import assert_hardware_report  # noqa: E402
+from experiments.kp_raw_traffic_scaling import (  # noqa: E402
+    DEPTHS,
+    RESULT_ROOT,
+    jobs,
+    registry_sha256,
+)
+
+
+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 read_json(path):
+    with open(path, encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def finite(value):
+    return (
+        isinstance(value, (int, float))
+        and not isinstance(value, bool)
+        and math.isfinite(float(value))
+    )
+
+
+def audit_job(job, launch):
+    manifest_path = job["output"] + ".manifest.json"
+    if not os.path.isfile(manifest_path):
+        return {
+            "depth": job["depth"],
+            "architecture": job["architecture"],
+            "status": "missing_manifest",
+            "finite": False,
+            "final_validation_accuracy": None,
+        }
+    manifest = read_json(manifest_path)
+    assert manifest["source"] == launch["source"]
+    assert manifest["command"] == job["command"]
+    assert manifest["hardware_lock"]["physical_gpu_name"] == "NVIDIA RTX A6000"
+    assert_hardware_report(manifest["hardware_lock"], launch["hardware_policy"])
+    if manifest["status"] != "completed" or not manifest["output_exists"]:
+        return {
+            "depth": job["depth"],
+            "architecture": job["architecture"],
+            "status": manifest["status"],
+            "finite": False,
+            "final_validation_accuracy": None,
+            "driver_wall_seconds": manifest["driver_wall_seconds"],
+        }
+    assert sha256(job["output"]) == manifest["output_sha256"]
+    record = read_json(job["output"])
+    args = record["args"]
+    expected = {
+        "mode": "kp_traffic",
+        "traffic_rule": "raw",
+        "traffic_ratio": 4.0,
+        "traffic_seed": 5000,
+        "traffic_calibration_examples": 64,
+        "predictor_mode": "closed_form",
+        "predictor_warmup_steps": 1,
+        "predictor_every": 0,
+        "neutral_projection": 1,
+        "depth": job["depth"],
+        "width": 16,
+        "seed": 0,
+        "loader_seed": 0,
+        "split_seed": 2027,
+        "epochs": 200,
+        "eval_split": "validation",
+        "eval_every": 1,
+        "lr": 0.1,
+        "output_lr": 0.1,
+    }
+    for key, wanted in expected.items():
+        assert args[key] == wanted, (key, args[key], wanted)
+    assert record["provenance"] == {
+        "git_commit": launch["source"]["git_commit"],
+        "git_tracked_dirty": False,
+    }
+    assert record["split"]["train_examples"] == 45_000
+    assert record["split"]["validation_examples"] == 5_000
+    assert record["split"]["test_examples"] == 10_000
+    assert record["evaluation_protocol"]["test_evaluations"] == 0
+    assert record["counters"]["ordinary_examples"] == 9_000_000
+    assert record["counters"]["traffic_calibration_examples"] == 64
+    assert record["counters"]["predictor_update_examples"] == 64
+    assert record["counters"]["neutral_projection_examples"] == 9_000_000
+    assert record["counters"]["logical_batch_loss_queries"] == 0
+    final = record["final"]
+    complete = len(record["epochs"]) == 200
+    is_finite = bool(
+        complete
+        and final["finite"] is True
+        and finite(final["accuracy"])
+        and finite(final["loss"])
+    )
+    return {
+        "depth": job["depth"],
+        "architecture": job["architecture"],
+        "status": "completed" if complete else "incomplete_trajectory",
+        "finite": is_finite,
+        "final_validation_accuracy": (
+            float(final["accuracy"]) if finite(final["accuracy"]) else None
+        ),
+        "final_validation_loss": (
+            float(final["loss"]) if finite(final["loss"]) else None
+        ),
+        "first_nonfinite_step": record.get("first_nonfinite_step"),
+        "total_wall_seconds": record["timing"]["total_timed_wall_s"],
+        "peak_memory_allocated_bytes": record["hardware"][
+            "peak_memory_allocated_bytes"
+        ],
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--out", default=os.path.join(RESULT_ROOT, "r1_audit.json")
+    )
+    args = parser.parse_args()
+    launch_path = os.path.join(RESULT_ROOT, "r1_launch.json")
+    launch = read_json(launch_path)
+    registered = jobs()
+    assert launch["registry_sha256"] == registry_sha256(registered)
+    records = [audit_job(job, launch) for job in registered]
+    report = {
+        "audit_status": "passed",
+        "stage": "raw_kp_traffic_scaling_r1",
+        "complete_registry": True,
+        "num_expected_cells": len(DEPTHS),
+        "num_audited_cells": len(records),
+        "num_finite_cells": sum(record["finite"] for record in records),
+        "failure_retaining": True,
+        "test_policy": "none",
+        "records": records,
+    }
+    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
+    with open(args.out, "w", encoding="utf-8") as handle:
+        json.dump(report, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+    print(json.dumps(report, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()
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()
diff --git a/experiments/kp_raw_traffic_scaling_smoke.py b/experiments/kp_raw_traffic_scaling_smoke.py
new file mode 100644
index 0000000..77a5cf6
--- /dev/null
+++ b/experiments/kp_raw_traffic_scaling_smoke.py
@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+"""Registry smoke for the frozen raw-KP scaling control."""
+import os
+import sys
+
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from experiments.kp_raw_traffic_scaling import DEPTHS, jobs
+
+
+def value(command, flag):
+    return command[command.index(flag) + 1]
+
+
+def main():
+    records = jobs()
+    assert len(records) == 3
+    assert {record["depth"] for record in records} == set(DEPTHS)
+    assert len({record["output"] for record in records}) == 3
+    for record in records:
+        command = record["command"]
+        expected = {
+            "--mode": "kp_traffic",
+            "--traffic_rule": "raw",
+            "--traffic_ratio": "4",
+            "--traffic_seed": "5000",
+            "--traffic_calibration_examples": "64",
+            "--predictor_mode": "closed_form",
+            "--predictor_warmup_steps": "1",
+            "--predictor_every": "0",
+            "--neutral_projection": "1",
+            "--depth": str(record["depth"]),
+            "--width": "16",
+            "--seed": "0",
+            "--loader_seed": "0",
+            "--split_seed": "2027",
+            "--epochs": "200",
+            "--eval_split": "validation",
+            "--eval_every": "1",
+            "--lr": "0.1",
+            "--output_lr": "0.1",
+            "--lr_schedule": "step",
+            "--lr_milestones": "100,150",
+        }
+        for flag, wanted in expected.items():
+            assert value(command, flag) == wanted, (flag, command)
+        assert record["timeout_seconds"] == 48 * 60 * 60
+    print("raw-KP traffic scaling registry: 3/3 exact")
+
+
+if __name__ == "__main__":
+    main()