repro_bundle/code/jobsrc/job_driver.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["torch", "numpy", "scipy", "huggingface_hub"]
# ///
"""HF Job driver for the L-SR1 analytic reproduction.

Runs (on one GPU): meta-training of all analytic optimizers, the projection
ablation, evaluation (Newton alignment + performance profiles + convergence),
and the efficiency microbenchmark. All artifacts are written to $OUT (a mounted
HF bucket) so they survive the job.

Launch (from the repo root):
  hf jobs uv run --flavor l40sx1 --timeout 3h --secrets HF_TOKEN \
    -v ./jobsrc:/code:ro \
    -v hf://buckets/txus/lsr1-artifacts/run1:/out:rw \
    -e OUT=/out -e CODE=/code \
    ./jobsrc/job_driver.py
"""
import os
import subprocess
import sys
import time

CODE = os.environ.get("CODE", "/code")
OUT = os.environ.get("OUT", "/out")
os.makedirs(OUT, exist_ok=True)

# Scale knobs (overridable by env for a cheap dry-run).
QUAD_ITERS = int(os.environ.get("QUAD_ITERS", 10000))
HARD_ITERS = int(os.environ.get("HARD_ITERS", 6000))
BATCH = int(os.environ.get("BATCH", 128))
DEVICE = os.environ.get("DEVICE", "cuda")
EVAL_QUICK = int(os.environ.get("EVAL_QUICK", 0))


def run(cmd):
    print(f"\n>>> {' '.join(str(c) for c in cmd)}", flush=True)
    t0 = time.time()
    r = subprocess.run(cmd, cwd=CODE)
    print(f"<<< exit={r.returncode} ({time.time()-t0:.1f}s)", flush=True)
    if r.returncode != 0:
        print(f"!! command failed (continuing): {cmd}", flush=True)
    return r.returncode


def train(preset, out, iters, projection=1, extra=None):
    cmd = [sys.executable, f"{CODE}/train_analytic.py", "--preset", preset,
           "--iters", str(iters), "--batch", str(BATCH), "--device", DEVICE,
           "--projection", str(projection), "--log-every", "200",
           "--out", f"{OUT}/{out}"]
    if extra:
        cmd += extra
    return run(cmd)


def main():
    t0 = time.time()
    print(f"=== L-SR1 job start | device={DEVICE} quad_iters={QUAD_ITERS} "
          f"hard_iters={HARD_ITERS} ===", flush=True)

    # --- Claim 5 (analytic) + Claim 4 model: quadratics at N=2, buffer L=8 -----
    train("quad2", "quad2_proj1.pt", QUAD_ITERS, projection=1)

    # --- Claim 2 ablation. Two views:
    #   (a) matched trust region -> quantitative on vs off comparison.
    #   (b) no trust region -> show projection is needed for stable training.
    train("quad2", "quad2tr_proj1.pt", QUAD_ITERS, projection=1,
          extra=["--trust-region", "3.0"])
    train("quad2", "quad2tr_proj0.pt", QUAD_ITERS, projection=0,
          extra=["--trust-region", "3.0"])
    # (the projection-off / no-trust-region divergence demo is already captured
    #  locally + in the prior run's logs; skipped here to save GPU time.)

    # --- Claim 4: performance-profile models (per family) ----------------------
    train("quad100", "quad100_proj1.pt", QUAD_ITERS, projection=1)
    train("rosen100", "rosen100_proj1.pt", HARD_ITERS, projection=1)
    train("rastr100", "rastr100_proj1.pt", HARD_ITERS, projection=1)

    # --- Evaluation ------------------------------------------------------------
    run([sys.executable, f"{CODE}/eval_analytic.py", "--device", DEVICE,
         "--quick", str(EVAL_QUICK),
         "--quad2", f"{OUT}/quad2tr_proj1.pt",
         "--quad2-noproj", f"{OUT}/quad2tr_proj0.pt",
         "--quad100", f"{OUT}/quad100_proj1.pt",
         "--rosen100", f"{OUT}/rosen100_proj1.pt",
         "--rastr100", f"{OUT}/rastr100_proj1.pt",
         "--out", f"{OUT}/eval.json"])

    # --- Claim 6 efficiency microbenchmark ------------------------------------
    run([sys.executable, f"{CODE}/efficiency.py", "--device", DEVICE,
         "--batch", "256", "--dim", "256", "--buffer", "4",
         "--out", f"{OUT}/efficiency.json"])

    print(f"\n=== L-SR1 job done in {(time.time()-t0)/60:.1f} min ===", flush=True)
    print("outputs:", sorted(os.listdir(OUT)), flush=True)


if __name__ == "__main__":
    main()