nshtrainer 0.11.6__tar.gz → 0.11.7__tar.gz

{nshtrainer-0.11.6 → nshtrainer-0.11.7}/PKG-INFO

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

{nshtrainer-0.11.6 → nshtrainer-0.11.7}/pyproject.toml

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

nshtrainer-0.11.7/src/nshtrainer/_checkpoint/saver.py

@@ -0,0 +1,51 @@
+import os
+import shutil
+from pathlib import Path
+
+from lightning.pytorch import Trainer
+
+from .metadata import _link_checkpoint_metadata, _remove_checkpoint_metadata
+
+
+def _link_checkpoint(
+    filepath: str | Path | os.PathLike,
+    linkpath: str | Path | os.PathLike,
+    *,
+    metadata: bool,
+):
+    if not isinstance(filepath, Path):
+        filepath = Path(filepath)
+    if not isinstance(linkpath, Path):
+        linkpath = Path(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)
+
+    if metadata:
+        _link_checkpoint_metadata(filepath, linkpath)
+
+
+def _remove_checkpoint(
+    trainer: Trainer,
+    filepath: str | Path | os.PathLike,
+    *,
+    metadata: bool,
+):
+    if not isinstance(filepath, Path):
+        filepath = Path(filepath)
+
+    trainer.strategy.remove_checkpoint(filepath)
+    if metadata:
+        _remove_checkpoint_metadata(filepath)

{nshtrainer-0.11.6 → nshtrainer-0.11.7}/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -90,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,
@@ -101,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):
@@ -119,13 +120,7 @@ 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):
@@ -159,21 +154,22 @@ 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

{nshtrainer-0.11.6 → nshtrainer-0.11.7}/src/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-0.11.6 → nshtrainer-0.11.7}/src/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.6/src/nshtrainer/_checkpoint/saver.py

@@ -1,61 +0,0 @@
-import os
-import shutil
-from pathlib import Path
-
-from lightning.pytorch import Trainer
-
-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):
-        filepath = Path(filepath)
-    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)
-
-        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()
-
-
-def _remove_checkpoint(
-    trainer: Trainer,
-    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()