nshtrainer 0.11.11__tar.gz → 0.11.13__tar.gz

{nshtrainer-0.11.11 → nshtrainer-0.11.13}/PKG-INFO

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

{nshtrainer-0.11.11 → nshtrainer-0.11.13}/pyproject.toml

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

{nshtrainer-0.11.11 → nshtrainer-0.11.13}/src/nshtrainer/_checkpoint/metadata.py

@@ -137,7 +137,11 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
     linked_path = linked_checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
     try:
         try:
-            linked_path.symlink_to(path)
+            # linked_path.symlink_to(path)
+            # We should store the path as a relative path
+            # to the metadata file to avoid issues with
+            # moving the checkpoint directory
+            linked_path.symlink_to(path.relative_to(linked_path.parent))
         except OSError:
             # on Windows, special permissions are required to create symbolic links as a regular user
             # fall back to copying the file

{nshtrainer-0.11.11 → nshtrainer-0.11.13}/src/nshtrainer/_checkpoint/saver.py

@@ -12,22 +12,25 @@ def _link_checkpoint(
     linkpath: str | Path | os.PathLike,
     *,
     metadata: bool,
+    remove_existing: bool = True,
 ):
     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)
+    if remove_existing:
+        if linkpath.exists():
+            if linkpath.is_symlink() or linkpath.is_file():
+                linkpath.unlink()
+            elif linkpath.is_dir():
+                shutil.rmtree(linkpath)
+
+        if metadata:
+            _remove_checkpoint_metadata(linkpath)
 
     try:
-        target_path = filepath.relative_to(linkpath.parent)
-        linkpath.symlink_to(target_path)
+        linkpath.symlink_to(filepath.relative_to(linkpath.parent))
     except OSError:
         # on Windows, special permissions are required to create symbolic links as a regular user
         # fall back to copying the file

{nshtrainer-0.11.11 → nshtrainer-0.11.13}/src/nshtrainer/callbacks/checkpoint/_base.py

@@ -79,6 +79,9 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
     @abstractmethod
     def topk_sort_key(self, metadata: CheckpointMetadata) -> Any: ...
 
+    @abstractmethod
+    def topk_sort_reverse(self) -> bool: ...
+
     def symlink_path(self):
         if not self.config.save_symlink:
             return None
@@ -102,7 +105,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         ]
 
         # Sort by the topk sort key
-        metas = sorted(metas, key=self.topk_sort_key)
+        metas = sorted(metas, key=self.topk_sort_key, reverse=self.topk_sort_reverse())
 
         # Now, the metas are sorted from the best to the worst,
         # so we can remove the worst checkpoints
@@ -145,15 +148,15 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         trainer.save_checkpoint(filepath, self.config.save_weights_only)
 
         if trainer.is_global_zero:
-            # Remove old checkpoints
-            self.remove_old_checkpoints(trainer)
-
             # Create the latest symlink
             if (symlink_filename := self.symlink_path()) is not None:
                 symlink_path = self.dirpath / symlink_filename
                 _link_checkpoint(filepath, symlink_path, metadata=True)
                 log.debug(f"Created latest symlink: {symlink_path}")
 
+            # Remove old checkpoints
+            self.remove_old_checkpoints(trainer)
+
         # Barrier to ensure all processes have saved the checkpoint,
         # deleted the old checkpoints, and created the symlink before continuing
         trainer.strategy.barrier()

{nshtrainer-0.11.11 → nshtrainer-0.11.13}/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -64,6 +64,10 @@ class BestCheckpoint(CheckpointBase[BestCheckpointCallbackConfig]):
             float("-inf" if self.metric.mode == "max" else "inf"),
         )
 
+    @override
+    def topk_sort_reverse(self):
+        return self.metric.mode == "max"
+
     # Events
     @override
     def on_validation_end(self, trainer: Trainer, pl_module: LightningModule):

{nshtrainer-0.11.11 → nshtrainer-0.11.13}/src/nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -34,6 +34,10 @@ class LastCheckpoint(CheckpointBase[LastCheckpointCallbackConfig]):
     def topk_sort_key(self, metadata: CheckpointMetadata):
         return metadata.checkpoint_timestamp
 
+    @override
+    def topk_sort_reverse(self):
+        return True
+
     @override
     def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
         self.save_checkpoints(trainer)

{nshtrainer-0.11.11 → nshtrainer-0.11.13}/src/nshtrainer/model/config.py

@@ -3,7 +3,6 @@ import logging
 import os
 import string
 import time
-import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
 from datetime import timedelta
@@ -50,10 +49,6 @@ from ..util._environment_info import EnvironmentConfig
 log = logging.getLogger(__name__)
 
 
-class IdSeedWarning(Warning):
-    pass
-
-
 class BaseProfilerConfig(C.Config, ABC):
     dirpath: str | Path | None = None
     """
@@ -1478,35 +1473,12 @@ class BaseConfig(C.Config):
     _rng: ClassVar[np.random.Generator | None] = None
 
     @staticmethod
-    def generate_id(
-        *,
-        length: int = 8,
-        ignore_rng: bool = False,
-    ) -> str:
+    def generate_id(*, length: int = 8) -> str:
         """
         Generate a random ID of specified length.
 
-        Args:
-            length (int): The length of the generated ID. Default is 8.
-            ignore_rng (bool): If True, ignore the global random number generator and use a new one. Default is False.
-
-        Returns:
-            str: The generated random ID.
-
-        Raises:
-            IdSeedWarning: If the global random number generator is None and ignore_rng is False.
-
-        Notes:
-            - The generated IDs will not be reproducible if the global random number generator is None and ignore_rng is False.
-            - To ensure reproducibility, call BaseConfig.set_seed(...) before generating any IDs.
         """
-        rng = BaseConfig._rng if not ignore_rng else np.random.default_rng()
-        if rng is None:
-            warnings.warn(
-                "BaseConfig._rng is None. The generated IDs will not be reproducible. "
-                + "To fix this, call BaseConfig.set_seed(...) before generating any IDs.",
-                category=IdSeedWarning,
-            )
+        if (rng := BaseConfig._rng) is None:
             rng = np.random.default_rng()
 
         alphabet = list(string.ascii_lowercase + string.digits)