nshtrainer 0.26.2__tar.gz → 0.27.0__tar.gz

{nshtrainer-0.26.2 → nshtrainer-0.27.0}/PKG-INFO

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

{nshtrainer-0.26.2 → nshtrainer-0.27.0}/pyproject.toml

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

{nshtrainer-0.26.2 → nshtrainer-0.27.0}/src/nshtrainer/callbacks/checkpoint/_base.py

@@ -153,10 +153,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         # Save the new checkpoint
         filepath = self.resolve_checkpoint_path(self.current_metrics(trainer))
         trainer._nshtrainer_save_checkpoint(
-            filepath,
-            self.config.save_weights_only,
-            use_checkpoint_cache=None,
-            ckpt_cache_use_symlink=False,
+            filepath, self.config.save_weights_only, use_checkpoint_cache=None
         )
 
         if trainer.is_global_zero:

{nshtrainer-0.26.2 → nshtrainer-0.27.0}/src/nshtrainer/model/config.py

@@ -1012,7 +1012,7 @@ class TrainerConfig(C.Config):
     """If enabled, the model supports scaling the gradients of shared parameters that are registered using `LightningModuleBase.register_shared_parameters(...)`"""
     save_checkpoint_metadata: bool = True
     """If enabled, will save additional metadata whenever a checkpoint is saved."""
-    use_checkpoint_cache: bool = True
+    use_checkpoint_cache: bool = False
     """If enabled, will optimize the saving of duplicate checkpoints by creating symlinks instead of copying the file."""
 
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()

{nshtrainer-0.26.2 → nshtrainer-0.27.0}/src/nshtrainer/trainer/trainer.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import shutil
+from collections import defaultdict
 from collections.abc import Sequence
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
@@ -18,7 +19,6 @@ from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
 from typing_extensions import Unpack, assert_never, override
 
 from .._checkpoint.metadata import _write_checkpoint_metadata
-from .._checkpoint.saver import _link_checkpoint
 from ..callbacks.base import resolve_all_callbacks
 from ..model.config import (
     AcceleratorConfigProtocol,
@@ -280,6 +280,12 @@ class Trainer(LightningTrainer):
     if TYPE_CHECKING:
         callbacks: list[Callback]
 
+    def _nshtrainer_checkpoint_cache_get(self, key: tuple[int, int]):
+        return next(
+            (ckpt for ckpt in self._nshtrainer_checkpoint_cache[key] if ckpt.exists()),
+            None,
+        )
+
     @override
     def __init__(
         self,
@@ -287,7 +293,9 @@ class Trainer(LightningTrainer):
         /,
         **kwargs: Unpack[LightningTrainerKwargs],
     ):
-        self._nshtrainer_checkpoint_cache: dict[tuple[int, int], Path] = {}
+        self._nshtrainer_checkpoint_cache = defaultdict[tuple[int, int], list[Path]](
+            lambda: []
+        )
 
         self._pre_init(config)
 
@@ -417,7 +425,6 @@ class Trainer(LightningTrainer):
         weights_only: bool = False,
         storage_options: Any | None = None,
         use_checkpoint_cache: bool | None = None,
-        ckpt_cache_use_symlink: bool = False,
     ):
         lm = self._base_module
         root_config = cast(BaseConfig, lm.hparams)
@@ -433,7 +440,7 @@ class Trainer(LightningTrainer):
         if (
             use_checkpoint_cache
             and (
-                cached_path := self._nshtrainer_checkpoint_cache.get(
+                cached_path := self._nshtrainer_checkpoint_cache_get(
                     (self.current_epoch, self.global_step)
                 )
             )
@@ -442,10 +449,7 @@ class Trainer(LightningTrainer):
             # If we have a cached path, then we symlink it to the new path.
             log.info(f"Re-using cached checkpoint {cached_path} for {filepath}.")
             if self.is_global_zero:
-                if ckpt_cache_use_symlink:
-                    _link_checkpoint(cached_path, filepath, metadata=False)
-                else:
-                    shutil.copy(cached_path, filepath)
+                shutil.copy(cached_path, filepath)
             self.strategy.barrier("Trainer.save_checkpoint")
         else:
             super().save_checkpoint(filepath, weights_only, storage_options)
@@ -454,7 +458,7 @@ class Trainer(LightningTrainer):
         if use_checkpoint_cache and cached_path is None:
             self._nshtrainer_checkpoint_cache[
                 (self.current_epoch, self.global_step)
-            ] = filepath
+            ].append(filepath)
             log.debug(f"Checkpoint saved to cache: {filepath}")
 
         # Save the checkpoint metadata