nshtrainer 0.10.5__tar.gz → 0.10.7__tar.gz

{nshtrainer-0.10.5 → nshtrainer-0.10.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.10.5
+Version: 0.10.7
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.10.5 → nshtrainer-0.10.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.10.5"
+version = "0.10.7"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-0.10.5 → nshtrainer-0.10.7}/src/nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -43,10 +43,6 @@ class LatestEpochCheckpoint(Checkpoint):
         self.config = config
         self.dirpath = dirpath
 
-        # Also, we hold a reference to the last checkpoint path
-        # to be able to remove it when a new checkpoint is saved.
-        self._last_ckpt_path: Path | None = None
-
     def _ckpt_path(self, trainer: Trainer):
         return self.dirpath / self.config.filename.format(
             epoch=trainer.current_epoch, step=trainer.global_step
@@ -54,20 +50,17 @@ class LatestEpochCheckpoint(Checkpoint):
 
     @override
     def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
-        # Remove the last checkpoint if it exists
-        if self._last_ckpt_path is not None:
-            trainer.strategy.remove_checkpoint(self._last_ckpt_path)
-
         # Save the new checkpoint
         filepath = self._ckpt_path(trainer)
         trainer.save_checkpoint(filepath, self.config.save_weights_only)
-        self._last_ckpt_path = filepath
 
         # Create the latest symlink
-        if (symlink_filename := self.config.latest_symlink_filename) is not None:
+        if (
+            trainer.is_global_zero
+            and (symlink_filename := self.config.latest_symlink_filename) is not None
+        ):
             symlink_path = self.dirpath / symlink_filename
-            if symlink_path.exists():
-                symlink_path.unlink()
+            symlink_path.unlink(missing_ok=True)
             symlink_path.symlink_to(filepath.name)
             log.info(f"Created latest symlink: {symlink_path}")
 

{nshtrainer-0.10.5 → nshtrainer-0.10.7}/src/nshtrainer/model/config.py

@@ -1121,6 +1121,9 @@ class SanityCheckingConfig(C.Config):
 
 
 class TrainerConfig(C.Config):
+    ckpt_path: str | Path | None = None
+    """Path to a checkpoint to load and resume training from."""
+
     checkpoint_loading: CheckpointLoadingConfig | Literal["auto"] = "auto"
     """Checkpoint loading configuration options."""
 

{nshtrainer-0.10.5 → nshtrainer-0.10.7}/src/nshtrainer/trainer/trainer.py

@@ -304,6 +304,10 @@ class Trainer(LightningTrainer):
             log_dir = str(Path(log_dir).resolve())
         log.critical(f"LightningTrainer log directory: {self.log_dir}.")
 
+        # Set the checkpoint
+        if (ckpt_path := config.trainer.ckpt_path) is not None:
+            self.ckpt_path = str(Path(ckpt_path).resolve().absolute())
+
     def __runtime_tracker(self):
         return next(
             (