nshtrainer 0.11.5__py3-none-any.whl → 0.11.7__py3-none-any.whl

nshtrainer/_checkpoint/saver.py

@@ -8,11 +8,9 @@ from .metadata import _link_checkpoint_metadata, _remove_checkpoint_metadata
 
 
 def _link_checkpoint(
-    trainer: Trainer,
     filepath: str | Path | os.PathLike,
     linkpath: str | Path | os.PathLike,
     *,
-    barrier: bool,
     metadata: bool,
 ):
     if not isinstance(filepath, Path):
@@ -20,26 +18,23 @@ def _link_checkpoint(
     if not isinstance(linkpath, Path):
         linkpath = Path(linkpath)
 
-    if trainer.is_global_zero:
-        if linkpath.exists():
-            if linkpath.is_symlink() or linkpath.is_file():
-                linkpath.unlink()
-            elif linkpath.is_dir():
-                shutil.rmtree(linkpath)
-            _remove_checkpoint_metadata(linkpath)
+    if linkpath.exists():
+        if linkpath.is_symlink() or linkpath.is_file():
+            linkpath.unlink()
+        elif linkpath.is_dir():
+            shutil.rmtree(linkpath)
+        _remove_checkpoint_metadata(linkpath)
 
-        try:
-            target_path = filepath.relative_to(linkpath.parent)
-            linkpath.symlink_to(target_path)
-        except OSError:
-            # on Windows, special permissions are required to create symbolic links as a regular user
-            # fall back to copying the file
-            shutil.copy(filepath, linkpath)
+    try:
+        target_path = filepath.relative_to(linkpath.parent)
+        linkpath.symlink_to(target_path)
+    except OSError:
+        # on Windows, special permissions are required to create symbolic links as a regular user
+        # fall back to copying the file
+        shutil.copy(filepath, linkpath)
 
-        if metadata:
-            _link_checkpoint_metadata(filepath, linkpath)
-    if barrier:
-        trainer.strategy.barrier()
+    if metadata:
+        _link_checkpoint_metadata(filepath, linkpath)
 
 
 def _remove_checkpoint(
@@ -47,15 +42,10 @@ def _remove_checkpoint(
     filepath: str | Path | os.PathLike,
     *,
     metadata: bool,
-    barrier: bool,
 ):
     if not isinstance(filepath, Path):
         filepath = Path(filepath)
 
-    if trainer.is_global_zero:
-        trainer.strategy.remove_checkpoint(filepath)
-        if metadata:
-            _remove_checkpoint_metadata(filepath)
-
-    if barrier:
-        trainer.strategy.barrier()
+    trainer.strategy.remove_checkpoint(filepath)
+    if metadata:
+        _remove_checkpoint_metadata(filepath)

nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -72,6 +72,8 @@ class BestCheckpoint(Checkpoint):
         self.metric = metric
         self.dirpath = dirpath
 
+        self._last_global_step_saved = 0  # no need to save when no steps were taken
+
     @override
     def on_validation_end(self, trainer: Trainer, pl_module: LightningModule):
         self._save_best_checkpoint(trainer)
@@ -88,10 +90,6 @@ class BestCheckpoint(Checkpoint):
         filename = f"{self.PREFIX}{filename}{self.EXTENSION}"
         return self.dirpath / filename
 
-    def _remove_checkpoints(self, trainer: Trainer, ckpt_paths: list[Path]):
-        for ckpt_path in ckpt_paths:
-            _remove_checkpoint(trainer, ckpt_path, metadata=True, barrier=False)
-
     def _get_metric_value(self, metrics: dict[str, Any]):
         return metrics.get(
             self.metric.validation_monitor,
@@ -99,11 +97,16 @@ class BestCheckpoint(Checkpoint):
         )
 
     def _sorted_ckpts(self):
+        """
+        Get sorted checkpoints by the metric value.
+
+        Sort order: best -> worst
+        """
         ckpt_paths = list(self.dirpath.glob(f"{self.PREFIX}*{self.EXTENSION}"))
         return _sort_ckpts_by_metadata(
             ckpt_paths,
             key=lambda meta, _: self._get_metric_value(meta.metrics),
-            reverse=(self.metric.mode == "min"),
+            reverse=(self.metric.mode == "max"),
         )
 
     def _create_symlink(self, trainer: Trainer, best_ckpt_path: Path):
@@ -117,16 +120,15 @@ class BestCheckpoint(Checkpoint):
         if symlink_path.exists() and symlink_path.resolve() == best_ckpt_path:
             return
 
-        _link_checkpoint(
-            trainer,
-            best_ckpt_path,
-            symlink_path,
-            metadata=True,
-            barrier=False,
-        )
+        _link_checkpoint(best_ckpt_path, symlink_path, metadata=True)
         log.debug(f"Created best symlink: {symlink_path}")
 
     def _save_best_checkpoint(self, trainer: Trainer):
+        # Skip saving the checkpoint if we're not in the fitting state
+        if self._should_skip_saving_checkpoint(trainer):
+            return
+
+        # Get the current metric value
         if (current := self._get_metric_value(trainer.callback_metrics)) is None:
             log.warning(
                 f"Can't save best model, {self.metric.validation_monitor} not found in metrics"
@@ -152,21 +154,39 @@ class BestCheckpoint(Checkpoint):
         trainer.save_checkpoint(filepath, self.config.save_weights_only)
         log.debug(f"Saved best checkpoint: {filepath}")
 
-        # Remove worst checkpoint if we've reached save_top_k
-        # NOTE: We add 1 to save_top_k here because we have just saved a new checkpoint
-        if len(sorted_ckpts) + 1 > self.config._save_top_k_value:
+        if trainer.is_global_zero:
             # Get the sorted checkpoints again because now we have added a new checkpoint.
             # We could optimize this by adding the new checkpoint to the sorted list,
             # and then sorting it in place, but this is simpler.
             sorted_ckpts = self._sorted_ckpts()
-            self._remove_checkpoints(
-                trainer, [p for _, p in sorted_ckpts[self.config.save_top_k :]]
-            )
 
-        # Create symlink to best model
-        if sorted_ckpts:
-            _, best_ckpt_path = sorted_ckpts[0]
-            self._create_symlink(trainer, best_ckpt_path)
+            # Remove worst checkpoint if we've reached save_top_k
+            if (topk := self.config.save_top_k) != "all" and len(sorted_ckpts) > topk:
+                # NOTE: Sort order is best -> worst. Let's get the worst checkpoints.
+                for _, ckpt_path in sorted_ckpts[topk:]:
+                    _remove_checkpoint(trainer, ckpt_path, metadata=True)
+
+            # Create symlink to best model
+            if sorted_ckpts:
+                _, best_ckpt_path = sorted_ckpts[0]
+                self._create_symlink(trainer, best_ckpt_path)
+
+        # Update the last global step saved
+        self._last_global_step_saved = trainer.global_step
 
         # Barrier to ensure all processes have saved the checkpoint before continuing
         trainer.strategy.barrier()
+
+    def _should_skip_saving_checkpoint(self, trainer: Trainer) -> bool:
+        from lightning.pytorch.trainer.states import TrainerFn
+
+        return (
+            bool(
+                getattr(trainer, "fast_dev_run", False)
+            )  # disable checkpointing with fast_dev_run
+            or trainer.state.fn
+            != TrainerFn.FITTING  # don't save anything during non-fit
+            or trainer.sanity_checking  # don't save anything during sanity check
+            or self._last_global_step_saved
+            == trainer.global_step  # already saved at the last step
+        )

nshtrainer/callbacks/checkpoint/latest_epoch_checkpoint.py

@@ -51,6 +51,8 @@ class LatestEpochCheckpoint(Checkpoint):
         self.config = config
         self.dirpath = dirpath
 
+        self._last_global_step_saved = 0
+
     @override
     def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
         self._save_new_checkpoint(trainer)
@@ -67,53 +69,63 @@ class LatestEpochCheckpoint(Checkpoint):
         filename = f"{self.PREFIX}{filename}{self.EXTENSION}"
         return self.dirpath / filename
 
-    def _remove_checkpoints(self, trainer: Trainer, ckpt_paths: list[Path]):
-        for ckpt_path in ckpt_paths:
-            _remove_checkpoint(trainer, ckpt_path, metadata=True, barrier=False)
-
     def _remove_old_checkpoints(self, trainer: Trainer):
         if (latest_k := self.config.latest_k) == "all":
             return
 
-        # NOTE: We add 1 to the latest_k here because
-        # we're about to save a new checkpoint.
-        latest_k += 1
-
         # Get all configs, ignoring the latest symlink
-        ckpt_paths = list(self.dirpath.glob(f"{self.PREFIX}*{self.EXTENSION}"))
+        ckpts = list(self.dirpath.glob(f"{self.PREFIX}*{self.EXTENSION}"))
         # Ignore the latest symlink
         if (latest_symlink_filename := self._latest_symlink_filename()) is not None:
-            ckpt_paths = [p for p in ckpt_paths if p.name != latest_symlink_filename]
+            ckpts = [p for p in ckpts if p.name != latest_symlink_filename]
 
         # Sort by epoch, then step, then last modified
-        metadata_and_ckpt_paths = _sort_ckpts_by_metadata(
-            ckpt_paths,
+        ckpts = _sort_ckpts_by_metadata(
+            ckpts,
             key=lambda meta, p: (meta.epoch, meta.global_step, p.stat().st_mtime),
             reverse=True,
         )
 
         # Remove all but the latest k checkpoints
-        ckpts_to_remove = metadata_and_ckpt_paths[latest_k:]
-        self._remove_checkpoints(trainer, [p for _, p in ckpts_to_remove])
+        # NOTE: We add 1 to the latest_k here because
+        # we're about to save a new checkpoint.
+        for _, ckpt_path in ckpts[latest_k:]:
+            _remove_checkpoint(trainer, ckpt_path, metadata=True)
 
     def _save_new_checkpoint(self, trainer: Trainer):
-        # Remove old checkpoints
-        if trainer.is_global_zero:
-            self._remove_old_checkpoints(trainer)
-        trainer.strategy.barrier()
+        if self._should_skip_saving_checkpoint(trainer):
+            return
 
         # Save the new checkpoint
         filepath = self._ckpt_path(trainer)
         trainer.save_checkpoint(filepath, self.config.save_weights_only)
 
-        # Create the latest symlink
-        if (symlink_filename := self._latest_symlink_filename()) is not None:
-            symlink_path = self.dirpath / symlink_filename
-            _link_checkpoint(
-                trainer,
-                filepath,
-                symlink_path,
-                barrier=True,
-                metadata=True,
-            )
-            log.debug(f"Created latest symlink: {symlink_path}")
+        if trainer.is_global_zero:
+            # Remove old checkpoints
+            self._remove_old_checkpoints(trainer)
+
+            # Create the latest symlink
+            if (symlink_filename := self._latest_symlink_filename()) is not None:
+                symlink_path = self.dirpath / symlink_filename
+                _link_checkpoint(filepath, symlink_path, metadata=True)
+                log.debug(f"Created latest symlink: {symlink_path}")
+
+        # Set the last global step saved
+        self._last_global_step_saved = trainer.global_step
+
+        # Barrier to ensure all processes have saved the checkpoint before continuing
+        trainer.strategy.barrier()
+
+    def _should_skip_saving_checkpoint(self, trainer: Trainer) -> bool:
+        from lightning.pytorch.trainer.states import TrainerFn
+
+        return (
+            bool(
+                getattr(trainer, "fast_dev_run", False)
+            )  # disable checkpointing with fast_dev_run
+            or trainer.state.fn
+            != TrainerFn.FITTING  # don't save anything during non-fit
+            or trainer.sanity_checking  # don't save anything during sanity check
+            or self._last_global_step_saved
+            == trainer.global_step  # already saved at the last step
+        )

nshtrainer/callbacks/checkpoint/model_checkpoint.py

@@ -198,19 +198,10 @@ class ModelCheckpoint(_ModelCheckpoint):
 
     @override
     def _link_checkpoint(self, trainer: Trainer, filepath: str, linkpath: str):  # pyright: ignore[reportIncompatibleMethodOverride]
-        return _link_checkpoint(
-            trainer,
-            filepath,
-            linkpath,
-            barrier=True,
-            metadata=True,
-        )
+        if trainer.is_global_zero:
+            _link_checkpoint(filepath, linkpath, metadata=True)
+        trainer.strategy.barrier()
 
     @override
     def _remove_checkpoint(self, trainer: Trainer, filepath: str):
-        return _ckpt_saver_remove_checkpoint(
-            trainer,
-            filepath,
-            metadata=True,
-            barrier=False,
-        )
+        _ckpt_saver_remove_checkpoint(trainer, filepath, metadata=True)

{nshtrainer-0.11.5.dist-info → nshtrainer-0.11.7.dist-info}/METADATA

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

{nshtrainer-0.11.5.dist-info → nshtrainer-0.11.7.dist-info}/RECORD

@@ -1,7 +1,7 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
 nshtrainer/_checkpoint/loader.py,sha256=_3jBf-k-fJCFfmU8wjDwbnE9rb4WoKYEyQiKGsBOCi4,13777
 nshtrainer/_checkpoint/metadata.py,sha256=M9eAZ2xMs36Z1G1xULu9MHZhsHxN8_9mNt3Iv7wuq-I,5069
-nshtrainer/_checkpoint/saver.py,sha256=z_c7a91O4Bh4lZZjqJgxT3w25qFlJsOopV3cpJtkHk8,1655
+nshtrainer/_checkpoint/saver.py,sha256=TuSAP39DOOVvSnSukQ9RitMV60JnDg6L27fMRc2uVJc,1358
 nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ4b-PoWM,120
 nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
 nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
@@ -11,9 +11,9 @@ nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHm
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,3555
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=zrEVCGFikfkt0iOMceOFzXsZG2-6QrqY79RKBCS7bu4,738
-nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=w99O5GWRcV89XBe4j__v2TvNEHys0x_r3tSTr-6Lhec,6154
-nshtrainer/callbacks/checkpoint/latest_epoch_checkpoint.py,sha256=NES-acaslPBiZQIMAdk_YwtnBrkm_y_BJQ8Ian0UKP0,4294
-nshtrainer/callbacks/checkpoint/model_checkpoint.py,sha256=mLFMbNzeMiBer3BCb7o3ucswKpOCQlYyN3wdB92N-LY,6884
+nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=BUO0sWqlwfyxD1UeII5DZ-01SGLiawJAEsL8HjGX4XA,7018
+nshtrainer/callbacks/checkpoint/latest_epoch_checkpoint.py,sha256=CQ0IqhuPI7zAxpQLy48kK8qqfVfwXEJoHGRqI4h8xNk,4819
+nshtrainer/callbacks/checkpoint/model_checkpoint.py,sha256=JS1z2YuEiQxk61HgZU1jySzF_pzdfXYO54_qHo-q3CQ,6776
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=s8tOHrnb_uVqLVeV2K38ZszXrXPTEGdDVfXuXgo_KDQ,3277
 nshtrainer/callbacks/early_stopping.py,sha256=LGn3rdbvkFfUo9kwMzK4eMGlPAqD9uFdowDx6VdfozQ,3761
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
@@ -82,6 +82,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.11.5.dist-info/METADATA,sha256=KHgvYOhQXbc37awWeLbpbdVQbSEU4J7KoC7Lr5286KE,860
-nshtrainer-0.11.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.11.5.dist-info/RECORD,,
+nshtrainer-0.11.7.dist-info/METADATA,sha256=htPbfKNDbqr1taf0bEvSbl-hQOaPfjKkJhvrWoMz2r0,860
+nshtrainer-0.11.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.11.7.dist-info/RECORD,,