nshtrainer 0.10.5__py3-none-any.whl → 0.10.7__py3-none-any.whl

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/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/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(
             (

{nshtrainer-0.10.5.dist-info → nshtrainer-0.10.7.dist-info}/METADATA

@@ -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.dist-info → nshtrainer-0.10.7.dist-info}/RECORD

@@ -14,7 +14,7 @@ nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,1
 nshtrainer/callbacks/finite_checks.py,sha256=AO5fa51uANAjAkeJfTquOjK6W_4RSU5Kky3f5jmAPlQ,2084
 nshtrainer/callbacks/gradient_skipping.py,sha256=fSJpjgHbztFKz7w3qFuCHZpmbEt9BCLAy-sU0B4xJQI,3474
 nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvLfaw,8080
-nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=UnwgGIc2reD7cTnUeIlDHo1LeAkgLEZFNvy2NGvUfRQ,2838
+nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=zCRAUsqW-2PaoIwVKlXOqdh2uF_B_YUUTmQO1wSomR8,2489
 nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
 nshtrainer/callbacks/model_checkpoint.py,sha256=N0raLsHlCVSbO3QU5eNFUXUDqxxW3C73oQwceMnFE_k,5955
 nshtrainer/callbacks/norm_logging.py,sha256=EWyrfkp8iHjQi9iAAXHxb0xStw2RwkdpKG2_gLarQRA,6281
@@ -52,7 +52,7 @@ nshtrainer/metrics/_config.py,sha256=hWWS4IXENRyH3RmJ7z1Wx1n3Lt1sNMlGOrcU6PW15o0
 nshtrainer/model/__init__.py,sha256=TbexTxiE20WHYg5q3L88Hysk4LlHeKk_isv33aSBREA,1918
 nshtrainer/model/_environment.py,sha256=JCFxxwMhkviiMDkqIXJmiuepqiSYIlcoSQM7Y2H2KX4,23036
 nshtrainer/model/base.py,sha256=Bmw-t70TydDbE9P0ee-lTibGoUhrCx5Qke-upa7FGVM,17512
-nshtrainer/model/config.py,sha256=OsVba02cmEYVf6V-A6ljV7VMAW5XZO6GWNRk8ktUw2o,53177
+nshtrainer/model/config.py,sha256=6pAqDUk1eBloR3vZmtsWVdMrKeT2V3UvOn5UZ7YhZ_Q,53283
 nshtrainer/model/modules/callback.py,sha256=JF59U9-CjJsAIspEhTJbVaGN0wGctZG7UquE3IS7R8A,6408
 nshtrainer/model/modules/debug.py,sha256=DTVty8cKnzj1GCULRyGx_sWTTsq9NLi30dzqjRTnuCU,1127
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
@@ -73,12 +73,12 @@ nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=xoqI2dcPnlNFPPLVIU6dBOvRPC9PtfX5qu__xV1lx0Y,2124
 nshtrainer/trainer/signal_connector.py,sha256=llwc8pdKAWxREFpjdi14Bpy8rGVMEJsmJx_s2p4gI8E,10689
-nshtrainer/trainer/trainer.py,sha256=n3T9Iz3eaDostxEdjapWImAsVMxyU9WBdhlPl0THX-g,16785
+nshtrainer/trainer/trainer.py,sha256=tFyzIsF8c-FABTH6wwDOR9y8kydVJqeVO7PDNFMvhSU,16950
 nshtrainer/util/environment.py,sha256=_SEtiQ_s5bL5pllUlf96AOUv15kNvCPvocVC13S7mIk,4166
 nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.10.5.dist-info/METADATA,sha256=nBa8n5rSpkY6MWYI-2JjmJzWKWa7gY0NOs737jsXXsU,695
-nshtrainer-0.10.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.10.5.dist-info/RECORD,,
+nshtrainer-0.10.7.dist-info/METADATA,sha256=IQ6IEecsAvygnoV5P6_mkG9RjRGnb_cFuOf2Ic2HLIY,695
+nshtrainer-0.10.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.10.7.dist-info/RECORD,,