From 9db02a0485e1289524111354e5c7fa35ef714a79 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 9 Jul 2026 12:50:36 -0700 Subject: [PATCH 1/2] fail if no checkpoint found --- ScaFFold/cli.py | 19 +++++++++++++++++++ ScaFFold/utils/checkpointing.py | 7 ++++++- ScaFFold/utils/config_utils.py | 1 + ScaFFold/utils/trainer.py | 4 +++- 4 files changed, 29 insertions(+), 2 deletions(-) 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..c0f58b0 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -145,7 +145,7 @@ def restore_training_state(self, snapshot: Dict[str, Any]) -> None: if self.optimizer: self.optimizer.zero_grad(set_to_none=True) - def load_from_checkpoint(self) -> int: + def load_from_checkpoint(self, require_checkpoint: bool = False) -> int: """Load the latest checkpoint. Returns start_epoch (default 1).""" self.wait_for_save() # Safety: don't load while writing @@ -159,6 +159,11 @@ def load_from_checkpoint(self) -> int: candidate = self._broadcast_obj(candidate) if not candidate: + if require_checkpoint: + raise FileNotFoundError( + "Restart requested but no checkpoint was found. " + f"Expected {self.last_ckpt_path} or {self.best_ckpt_path}." + ) return 1 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"]) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index fc5af6e..e3a41f8 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -318,7 +318,9 @@ def cleanup_or_resume(self): self.start_epoch = 1 else: # Load checkpoint via manager - self.start_epoch = self.checkpoint_manager.load_from_checkpoint() + self.start_epoch = self.checkpoint_manager.load_from_checkpoint( + require_checkpoint=getattr(self.config, "restart", False) + ) # Restore extra metadata if needed (e.g. mask values) if "train_mask_values" in self.checkpoint_manager.restored_extras: From 534cc38f0a54fbcc76023ac6033e5b978dba6aed Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 16 Jul 2026 12:49:32 -0700 Subject: [PATCH 2/2] remove extra arg --- ScaFFold/utils/checkpointing.py | 12 +++++------- ScaFFold/utils/trainer.py | 4 +--- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index c0f58b0..65c62f3 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -145,7 +145,7 @@ def restore_training_state(self, snapshot: Dict[str, Any]) -> None: if self.optimizer: self.optimizer.zero_grad(set_to_none=True) - def load_from_checkpoint(self, require_checkpoint: bool = False) -> int: + def load_from_checkpoint(self) -> int: """Load the latest checkpoint. Returns start_epoch (default 1).""" self.wait_for_save() # Safety: don't load while writing @@ -159,12 +159,10 @@ def load_from_checkpoint(self, require_checkpoint: bool = False) -> int: candidate = self._broadcast_obj(candidate) if not candidate: - if require_checkpoint: - raise FileNotFoundError( - "Restart requested but no checkpoint was found. " - f"Expected {self.last_ckpt_path} or {self.best_ckpt_path}." - ) - 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/trainer.py b/ScaFFold/utils/trainer.py index e3a41f8..fc5af6e 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -318,9 +318,7 @@ def cleanup_or_resume(self): self.start_epoch = 1 else: # Load checkpoint via manager - self.start_epoch = self.checkpoint_manager.load_from_checkpoint( - require_checkpoint=getattr(self.config, "restart", False) - ) + self.start_epoch = self.checkpoint_manager.load_from_checkpoint() # Restore extra metadata if needed (e.g. mask values) if "train_mask_values" in self.checkpoint_manager.restored_extras: