move download_checkpoint_from_wandb from engine to runner

allenai · Jul 12, 2024 · 25cbe11 · 25cbe11
1 parent 9aedd25
commit 25cbe11
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/allenact/algorithms/onpolicy_sync/engine.py b/allenact/algorithms/onpolicy_sync/engine.py
@@ -1901,14 +1901,6 @@ def train(
         # noinspection PyBroadException
         try:
             if checkpoint_file_name is not None:
-                if "wandb://" == checkpoint_file_name[:8]:
-                    ckpt_dir = "/tmp/wandb_ckpts"
-                    os.makedirs(ckpt_dir, exist_ok=True)
-                    checkpoint_file_name = download_checkpoint_from_wandb(
-                        checkpoint_file_name,
-                        ckpt_dir,
-                        only_allow_one_ckpt=True
-                    )
                 self.checkpoint_load(checkpoint_file_name, restart_pipeline)
 
             self.run_pipeline(valid_on_initial_weights=valid_on_initial_weights)

diff --git a/allenact/algorithms/onpolicy_sync/runner.py b/allenact/algorithms/onpolicy_sync/runner.py
@@ -534,6 +534,15 @@ def start_train(
 
         worker_ids = self.local_worker_ids(TRAIN_MODE_STR)
 
+        if "wandb://" == checkpoint[:8]:
+            ckpt_dir = "/tmp/wandb_ckpts"
+            os.makedirs(ckpt_dir, exist_ok=True)
+            checkpoint = download_checkpoint_from_wandb(
+                checkpoint,
+                ckpt_dir,
+                only_allow_one_ckpt=True
+            )
+
         model_hash = None
         for trainer_id in worker_ids:
             training_kwargs = dict(