nshtrainer 0.10.15__py3-none-any.whl → 0.10.16__py3-none-any.whl

nshtrainer/_checkpoint/metadata.py

@@ -105,10 +105,8 @@ def _write_checkpoint_metadata(
 
 
 def _remove_checkpoint_metadata(checkpoint_path: Path):
-    for path in (
-        checkpoint_path.with_suffix(METADATA_PATH_SUFFIX),
-        checkpoint_path.with_suffix(HPARAMS_PATH_SUFFIX),
-    ):
+    for suffix in (METADATA_PATH_SUFFIX, HPARAMS_PATH_SUFFIX):
+        path = checkpoint_path.with_suffix(suffix)
         try:
             path.unlink(missing_ok=True)
         except Exception as e:
@@ -122,11 +120,9 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
     _remove_checkpoint_metadata(linked_checkpoint_path)
 
     # Link the metadata files to the new checkpoint
-    for path in (
-        checkpoint_path.with_suffix(METADATA_PATH_SUFFIX),
-        checkpoint_path.with_suffix(HPARAMS_PATH_SUFFIX),
-    ):
-        linked_path = linked_checkpoint_path.with_suffix(path.suffix)
+    for suffix in (METADATA_PATH_SUFFIX, HPARAMS_PATH_SUFFIX):
+        path = checkpoint_path.with_suffix(suffix)
+        linked_path = linked_checkpoint_path.with_suffix(suffix)
         try:
             try:
                 linked_path.symlink_to(path)

nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -64,7 +64,7 @@ class LatestEpochCheckpoint(Checkpoint):
         filename = self.config.filename.format(
             epoch=trainer.current_epoch, step=trainer.global_step
         )
-        filename = f"{self.PREFIX}{filename}.{self.EXTENSION}"
+        filename = f"{self.PREFIX}{filename}{self.EXTENSION}"
         return self.dirpath / filename
 
     def _remove_checkpoints(self, trainer: Trainer, ckpt_paths: list[Path]):
@@ -95,7 +95,9 @@ class LatestEpochCheckpoint(Checkpoint):
 
     def _save_new_checkpoint(self, trainer: Trainer):
         # Remove old checkpoints
-        self._remove_old_checkpoints(trainer)
+        if trainer.is_global_zero:
+            self._remove_old_checkpoints(trainer)
+        trainer.strategy.barrier()
 
         # Save the new checkpoint
         filepath = self._ckpt_path(trainer)

{nshtrainer-0.10.15.dist-info → nshtrainer-0.10.16.dist-info}/METADATA

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

{nshtrainer-0.10.15.dist-info → nshtrainer-0.10.16.dist-info}/RECORD

@@ -1,6 +1,6 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
 nshtrainer/_checkpoint/loader.py,sha256=48flPr1XgQHOgIPaCrRqOEvRuG0SZuV3cQ1vgHLqFqI,11025
-nshtrainer/_checkpoint/metadata.py,sha256=soK9tXVs6EOpzhlnIxTEF51KmdkaCDUj0Rdyid3uREk,5640
+nshtrainer/_checkpoint/metadata.py,sha256=GlhlAyJh5gcp3R8l2Y3eAUQtQzBnitFlB0xdx-khEUQ,5579
 nshtrainer/_checkpoint/saver.py,sha256=z_c7a91O4Bh4lZZjqJgxT3w25qFlJsOopV3cpJtkHk8,1655
 nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ4b-PoWM,120
 nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
@@ -15,7 +15,7 @@ nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,1
 nshtrainer/callbacks/finite_checks.py,sha256=gJC_RUr3ais3FJI0uB6wUZnDdE3WRwCix3ppA3PwQXA,2077
 nshtrainer/callbacks/gradient_skipping.py,sha256=pqu5AELx4ctJxR2Y7YSSiGd5oGauVCTZFCEIIS6s88w,3665
 nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvLfaw,8080
-nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=zUeYAGfeQby0R6IwQBJH3lng-MD0vkckdX4aIOm-VIc,4146
+nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=5JC-JCdgWNnunl0jv4Q9LhkEspLAn0x8VpCMJZi7-ow,4219
 nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
 nshtrainer/callbacks/model_checkpoint.py,sha256=8D0wWLhr_KiksAA1fjfIuby42Mq6XokCvAnVUhjADd8,6564
 nshtrainer/callbacks/norm_logging.py,sha256=T2psu8mYsw9iahPKT6aUPjkGrZ4TIzm6_UUUmE09GJs,6274
@@ -79,6 +79,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 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.15.dist-info/METADATA,sha256=lBdMigvT3LEgOyWtMBwaRvru8XRTU8K5GQ-ll3kqwE8,696
-nshtrainer-0.10.15.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.10.15.dist-info/RECORD,,
+nshtrainer-0.10.16.dist-info/METADATA,sha256=8jgjZDL82cNf_ys1xKUuqfKXAol8m2dWYB909W239fk,696
+nshtrainer-0.10.16.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.10.16.dist-info/RECORD,,