diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 9249aac..20b933f 100644 --- a/ScaFFold/cli.py +++ b/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}") diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 2a06a3c..65c62f3 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/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}") diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 9cb632c..46ff425 100644 --- a/ScaFFold/utils/config_utils.py +++ b/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"])