run_all.py

"""Job driver: run experiment stages sequentially inside one job.

Stages:
  lrsweep  - 6 short runs (3k steps): {fixed-4, rec-1x4} x lr {1e-4,3e-4,1e-3}
  stage1   - matched-FLOP family at the chosen lr: fixed-4, rec-1x4, rec-2x2,
             rec-1xdyn (each 3 seeds, 25k steps) + their evaluations
  stage2   - iso-param pair: rec-4x4 (3 seeds), fixed-4-long (100k steps,
             2 seeds) + evaluations
  eval     - evaluate every local checkpoint not yet evaluated
"""

import argparse
import glob
import json
import os
import subprocess
import sys

STEPS = 25000
LR_FILE = "out/chosen_lr.json"


def run(cmd):
    print("+", " ".join(cmd), flush=True)
    subprocess.run(cmd, check=True)


def run_eval(args, ckpt, out_dir="results"):
    """Run evaluate.py on one checkpoint, then upload the JSON to the Hub.
    Without the upload the eval JSONs die with the job's ephemeral disk."""
    name = os.path.splitext(os.path.basename(ckpt))[0]
    out = f"{out_dir}/{name}_eval.json"
    if os.path.exists(out):
        return
    run([sys.executable, "evaluate.py", "--ckpt", ckpt, "--out-dir", out_dir])
    if args.hub_repo:
        run([sys.executable, "-c",
             f"from huggingface_hub import HfApi; HfApi().upload_file("
             f"path_or_fileobj={out!r}, "
             f"path_in_repo='results/{name}_eval.json', "
             f"repo_id='{args.hub_repo}', repo_type='model')"])


def train_cmd(a, run_name, steps, lr, hub_repo):
    cmd = [sys.executable, "train.py", "--arch", a["arch"],
           "--depth", str(a["depth"]), "--iters", a["iters"],
           "--seed", str(a["seed"]), "--steps", str(steps), "--lr", str(lr),
           "--run-name", run_name, "--out-dir", "out"]
    if hub_repo:
        cmd += ["--hub-repo", hub_repo, "--space-id",
                os.environ.get("TRACKIO_SPACE_ID", "")]
    return cmd


def stage_lrsweep(args):
    for lr in (1e-4, 3e-4, 1e-3):
        for a in ({"arch": "fixed", "depth": 4, "iters": "1", "seed": 0},
                  {"arch": "rec", "depth": 1, "iters": "4", "seed": 0}):
            name = f"lr{lr:g}-{a['arch']}{a['depth']}x{a['iters']}"
            run(train_cmd(a, name, 3000, lr, args.hub_repo))
    # pick lr by mean id accuracy over the last eval of each pair
    scores = {}
    for lr in (1e-4, 3e-4, 1e-3):
        tot = 0.0
        for a in ("fixed4x1", "rec1x4"):
            h = json.load(open(f"out/lr{lr:g}-{a}_history.json"))
            tot += h["id_acc"][-1]["mean"]
        scores[str(lr)] = tot / 2
    best = max(scores, key=scores.get)
    json.dump({"scores": scores, "chosen": best}, open(LR_FILE, "w"))
    if args.hub_repo:
        run([sys.executable, "-c", "import sys; from huggingface_hub import "
             "HfApi; HfApi().upload_file(path_or_fileobj='out/chosen_lr.json', "
             "path_in_repo='results/chosen_lr.json', "
             f"repo_id='{args.hub_repo}', repo_type='model')"])
    print("[lrsweep] scores", scores, "-> chosen", best)


def chosen_lr(hub_repo=None):
    if not os.path.exists(LR_FILE):
        # stage1/stage2 run in fresh containers: fetch the lr choice made by
        # the lrsweep job from the Hub repo
        from huggingface_hub import hf_hub_download
        p = hf_hub_download(repo_id=hub_repo, repo_type="model",
                            filename="results/chosen_lr.json", local_dir="out")
        os.replace(os.path.join("out", "results", "chosen_lr.json"), LR_FILE)
    return json.load(open(LR_FILE))["chosen"]


def stage1(args):
    lr = chosen_lr(args.hub_repo)
    seeds = (0, 1, 2)
    archs = [("fixed", 4, "1"), ("rec", 1, "4"), ("rec", 2, "2"),
             ("rec", 1, "dynamic")]
    for arch, depth, iters in archs:
        for s in seeds:
            name = f"{arch}{depth}x{iters}-seed{s}"
            ck = f"out/{name}.pt"
            if os.path.exists(ck):
                continue
            run(train_cmd({"arch": arch, "depth": depth, "iters": iters,
                           "seed": s}, name, STEPS, lr, args.hub_repo))
            run_eval(args, ck)


def stage2(args):
    lr = chosen_lr(args.hub_repo)
    for s in (0, 1, 2):
        name = f"rec4x4-seed{s}"
        if not os.path.exists(f"out/{name}.pt"):
            run(train_cmd({"arch": "rec", "depth": 4, "iters": "4",
                           "seed": s}, name, STEPS, lr, args.hub_repo))
        run_eval(args, f"out/{name}.pt")
    for s in (0, 1):
        name = f"fixed4-long-seed{s}"
        if not os.path.exists(f"out/{name}.pt"):
            run(train_cmd({"arch": "fixed", "depth": 4, "iters": "1",
                           "seed": s}, name, 4 * STEPS, lr, args.hub_repo))
        run_eval(args, f"out/{name}.pt")


def stage3(args):
    """Aligned-unroll experiment: train with R = k (one unroll per hop),
    testing whether the models' implicit R~k tuning can be made explicit
    and extended beyond the training range."""
    lr = chosen_lr(args.hub_repo)
    for s in (0, 1, 2):
        name = f"rec1aligned-seed{s}"
        if not os.path.exists(f"out/{name}.pt"):
            run(train_cmd({"arch": "rec", "depth": 1, "iters": "aligned",
                           "seed": s}, name, STEPS, lr, args.hub_repo))
        run_eval(args, f"out/{name}.pt")


def stage4(args):
    """Aligned objective at 4x the step budget (1 seed): distinguishes
    'harder optimization, needs more steps' from 'fundamentally harder'
    after the 25k-step aligned runs plateaued near 0.10 id_acc."""
    lr = chosen_lr(args.hub_repo)
    name = "rec1aligned-long-seed0"
    if not os.path.exists(f"out/{name}.pt"):
        run(train_cmd({"arch": "rec", "depth": 1, "iters": "aligned",
                       "seed": 0}, name, 4 * STEPS, lr, args.hub_repo))
    run_eval(args, f"out/{name}.pt")


def stage_eval(args):
    if not glob.glob("out/*.pt"):
        # recovery mode (fresh container): pull every checkpoint from the Hub,
        # plus existing eval JSONs so the skip-if-present logic works
        from huggingface_hub import snapshot_download
        snapshot_download(repo_id=args.hub_repo, repo_type="model",
                          allow_patterns=["ckpts/*.pt", "results/*_eval.json"],
                          local_dir=".")
    cks = sorted(glob.glob("out/*.pt") + glob.glob("ckpts/*.pt"))
    for ck in cks:
        name = os.path.splitext(os.path.basename(ck))[0]
        if name.startswith(("lr", "smoke")):
            continue  # not part of the main comparison grid
        if os.path.exists(f"results/{name}_eval.json"):
            continue
        run_eval(args, ck)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("stage", choices=["lrsweep", "stage1", "stage2", "stage3",
                                     "stage4", "eval"])
    p.add_argument("--hub-repo", default=None)
    a = p.parse_args()
    if a.hub_repo:
        os.environ.setdefault("TRACKIO_HUB_REPO", a.hub_repo)
    {"lrsweep": stage_lrsweep, "stage1": stage1, "stage2": stage2,
     "stage3": stage3, "stage4": stage4,
     "eval": stage_eval}[a.stage](a)


if __name__ == "__main__":
    main()