Check for checkpoints

ScaFFold/cli.py

@@ -302,6 +302,25 @@ def main():
 
     comm.Barrier()
     combined_config = comm.bcast(combined_config, root=0)
+
+    if combined_config.get("restart", False):
+        run_dir = combined_config.get("run_dir")
+        if not run_dir:
+            raise ValueError("--restart requires --run-dir")
+
+        checkpoint_dir = Path(run_dir) / combined_config.get(
+            "checkpoint_dir", "checkpoints"
+        )
+        expected_checkpoints = (
+            checkpoint_dir / "checkpoint_last.pth",
+            checkpoint_dir / "checkpoint_best.pth",
+        )
+        if not any(path.exists() for path in expected_checkpoints):
+            expected = " or ".join(str(path) for path in expected_checkpoints)
+            raise FileNotFoundError(
+                f"Restart requested but no checkpoint was found. Expected {expected}."
+            )
+
     if rank == 0:
         print(f"combined_config = {combined_config}")
 

ScaFFold/utils/checkpointing.py

@@ -159,7 +159,10 @@ def load_from_checkpoint(self) -> int:
 
         candidate = self._broadcast_obj(candidate)
         if not candidate:
-            return 1
+            raise FileNotFoundError(
+                "Restart requested but no checkpoint was found. "
+                f"Expected {self.last_ckpt_path} or {self.best_ckpt_path}."
+            )
 
         self._log(f"Loading checkpoint from {candidate}")
 

ScaFFold/utils/config_utils.py

@@ -60,6 +60,7 @@ def __init__(self, config_dict):
         self.more_determinism = bool(config_dict["more_determinism"])
         self.datagen_from_scratch = bool(config_dict["datagen_from_scratch"])
         self.train_from_scratch = bool(config_dict["train_from_scratch"])
+        self.restart = bool(config_dict.get("restart", False))
         self.val_split = config_dict["val_split"]
         self.seed = config_dict["seed"]
         self.dist = bool(config_dict["dist"])