run_all: chosen_lr.json shared via Hub between job stages

run_all.py

@@ -52,15 +52,27 @@ def stage_lrsweep(args):
         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():
+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()
+    lr = chosen_lr(args.hub_repo)
     seeds = (0, 1, 2)
     archs = [("fixed", 4, "1"), ("rec", 1, "4"), ("rec", 2, "2"),
              ("rec", 1, "dynamic")]
@@ -77,7 +89,7 @@ def stage1(args):
 
 
 def stage2(args):
-    lr = chosen_lr()
+    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"):