run_eval now uploads eval JSONs to the Hub; stage_eval recovery covers ckpts/

run_all.py

@@ -25,10 +25,20 @@ def run(cmd):
     subprocess.run(cmd, check=True)
 
 
-def run_eval(ckpt, out_dir="results"):
-    """Run evaluate.py on one checkpoint and upload the JSON to the Hub
-    (evaluate.py reads HUB_REPO from the environment, set in main())."""
+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):
@@ -90,8 +100,7 @@ def stage1(args):
                 continue
             run(train_cmd({"arch": arch, "depth": depth, "iters": iters,
                            "seed": s}, name, STEPS, lr, args.hub_repo))
-            run([sys.executable, "evaluate.py", "--ckpt", ck,
-                 "--out-dir", "results"])
+            run_eval(args, ck)
 
 
 def stage2(args):
@@ -101,15 +110,13 @@ def stage2(args):
         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([sys.executable, "evaluate.py", "--ckpt", f"out/{name}.pt",
-             "--out-dir", "results"])
+        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([sys.executable, "evaluate.py", "--ckpt", f"out/{name}.pt",
-             "--out-dir", "results"])
+        run_eval(args, f"out/{name}.pt")
 
 
 def stage_eval(args):
@@ -118,12 +125,12 @@ def stage_eval(args):
         from huggingface_hub import snapshot_download
         snapshot_download(repo_id=args.hub_repo, repo_type="model",
                           allow_patterns=["ckpts/*.pt"], local_dir=".")
-    for ck in sorted(glob.glob("out/*.pt")):
+    cks = sorted(glob.glob("out/*.pt") + glob.glob("ckpts/*.pt"))
+    for ck in cks:
         name = os.path.splitext(os.path.basename(ck))[0]
         if os.path.exists(f"results/{name}_eval.json"):
             continue
-        run([sys.executable, "evaluate.py", "--ckpt", ck,
-             "--out-dir", "results"])
+        run_eval(args, ck)
 
 
 def main():
@@ -132,7 +139,6 @@ def main():
     p.add_argument("--hub-repo", default=None)
     a = p.parse_args()
     if a.hub_repo:
-        os.environ["HUB_REPO"] = a.hub_repo  # evaluate.py uploads when set
         os.environ.setdefault("TRACKIO_HUB_REPO", a.hub_repo)
     {"lrsweep": stage_lrsweep, "stage1": stage1, "stage2": stage2,
      "eval": stage_eval}[a.stage](a)