nshtrainer 0.22.1__py3-none-any.whl → 0.23.0__py3-none-any.whl

nshtrainer/_checkpoint/saver.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import shutil
 from pathlib import Path
@@ -7,6 +8,8 @@ from lightning.pytorch import Trainer
 from ..util.path import get_relative_path
 from .metadata import _link_checkpoint_metadata, _remove_checkpoint_metadata
 
+log = logging.getLogger(__name__)
+
 
 def _link_checkpoint(
     filepath: str | Path | os.PathLike,
@@ -19,11 +22,14 @@ def _link_checkpoint(
     linkpath = Path(linkpath)
 
     if remove_existing:
-        if linkpath.exists():
-            if linkpath.is_symlink() or linkpath.is_file():
-                linkpath.unlink()
-            elif linkpath.is_dir():
-                shutil.rmtree(linkpath)
+        try:
+            if linkpath.exists():
+                if linkpath.is_dir():
+                    shutil.rmtree(linkpath, ignore_errors=True)
+                else:
+                    linkpath.unlink(missing_ok=True)
+        except Exception:
+            log.exception(f"Failed to remove {linkpath}")
 
         if metadata:
             _remove_checkpoint_metadata(linkpath)

nshtrainer/callbacks/checkpoint/_base.py

@@ -11,6 +11,7 @@ from typing_extensions import TypeVar, override
 
 from ..._checkpoint.metadata import CheckpointMetadata, _metadata_path
 from ..._checkpoint.saver import _link_checkpoint, _remove_checkpoint
+from ...util.path import find_symlinks
 from ..base import CallbackConfigBase
 
 if TYPE_CHECKING:
@@ -116,9 +117,47 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
                 )
                 continue
 
-            _remove_checkpoint(trainer, old_ckpt_path, metadata=True)
+            self._remove_checkpoint_with_link_support(
+                trainer, old_ckpt_path, metadata=True
+            )
             log.debug(f"Removed old checkpoint: {old_ckpt_path}")
 
+    def _remove_checkpoint_with_link_support(
+        self,
+        trainer: Trainer,
+        path: Path,
+        metadata: bool,
+    ):
+        # Find all the symlinks to the checkpoint
+        ckpt_callbacks: list[CheckpointBase] = [
+            callback
+            for callback in trainer.checkpoint_callbacks
+            if isinstance(callback, CheckpointBase) and callback is not self
+        ]
+        symlink_paths = find_symlinks(
+            path,
+            *[callback.dirpath for callback in ckpt_callbacks],
+            glob_pattern=f"*.{self.extension()}",
+        )
+
+        # If there are no symlinks, just remove the checkpoint
+        if not symlink_paths:
+            _remove_checkpoint(trainer, path, metadata=metadata)
+            return
+
+        log.debug(
+            f"Removing checkpoint with symlinks: {path}, symlinks: {symlink_paths}"
+        )
+
+        # For the first symlink, we can just move the checkpoint file
+        # to the symlink path. For the rest, we need to make new symlinks.
+        new_target = symlink_paths.pop(0)
+        path.rename(new_target)
+        log.debug(f"New symlink target: {new_target}")
+
+        for symlink_path in symlink_paths:
+            _link_checkpoint(new_target, symlink_path, metadata=False)
+
     def current_metrics(self, trainer: Trainer) -> dict[str, Any]:
         current_metrics: dict[str, Any] = {
             "epoch": trainer.current_epoch,

nshtrainer/trainer/trainer.py

@@ -441,6 +441,7 @@ class Trainer(LightningTrainer):
             log.info(f"Re-using cached checkpoint {cached_path} for {filepath}.")
             if self.is_global_zero:
                 _link_checkpoint(cached_path, filepath, metadata=False)
+            self.strategy.barrier("Trainer.save_checkpoint")
         else:
             super().save_checkpoint(filepath, weights_only, storage_options)
 

nshtrainer/util/path.py

@@ -27,3 +27,26 @@ def get_relative_path(source: _Path, destination: _Path):
     down = os.sep.join(destination_parts[i:])
 
     return Path(os.path.normpath(os.path.join(up, down)))
+
+
+def find_symlinks(
+    target_file: _Path,
+    *search_directories: _Path,
+    glob_pattern: str = "*",
+):
+    target_file = Path(target_file).resolve()
+    symlinks: list[Path] = []
+
+    for search_directory in search_directories:
+        search_directory = Path(search_directory)
+        for path in search_directory.rglob(glob_pattern):
+            if path.is_symlink():
+                try:
+                    link_target = path.resolve()
+                    if link_target.samefile(target_file):
+                        symlinks.append(path)
+                except FileNotFoundError:
+                    # Handle broken symlinks
+                    pass
+
+    return symlinks

{nshtrainer-0.22.1.dist-info → nshtrainer-0.23.0.dist-info}/METADATA

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

{nshtrainer-0.22.1.dist-info → nshtrainer-0.23.0.dist-info}/RECORD

@@ -2,7 +2,7 @@ nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
 nshtrainer/_callback.py,sha256=A1zLsTy4b_wOYnInLLXGSRdHzT2yNa6mPEql-ozm0u0,1013
 nshtrainer/_checkpoint/loader.py,sha256=5vjg-OFChXJjgiOVv8vnV8nwTscfdDtEdxQRz6uPfDE,14158
 nshtrainer/_checkpoint/metadata.py,sha256=E4tfiGzhnn65X95P0Y6K2d_YfPWqvHZoF0FF1-smEJc,5221
-nshtrainer/_checkpoint/saver.py,sha256=6W-Rbc3QGuhcF_mcwN8v31uEjLQCsZvt8CPuqPs4m5g,1342
+nshtrainer/_checkpoint/saver.py,sha256=fvRKGI5aeXtsHBOIO4cwGe__wmO-6DiD0-744VASYA4,1500
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
 nshtrainer/_hf_hub.py,sha256=iqhXH54RhSqmot_K3UCVcHTC_TC81_YY7cwvHGHXXlw,16782
 nshtrainer/callbacks/__init__.py,sha256=4qocBDzQbLLhhbIEfvbA3SQB_Dy9ZJH7keMwPay-ZS8,2359
@@ -10,7 +10,7 @@ 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=NpjeKmonJ1Kaz5_39XSn3LlDwvbGjk6WV8BpHSNCvI4,3508
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=g-3zIthupERKqWZQw-A_busQPaPRkto6iHBV-M7nK1Y,527
-nshtrainer/callbacks/checkpoint/_base.py,sha256=r6IPpl3sGUmxBNv80y9r326lTrPAIVSU3Fu-3LrYH2s,6691
+nshtrainer/callbacks/checkpoint/_base.py,sha256=zKN-6n61aGze-Hf8MBY1Surh6B-xDwNSApqQJtPcTUs,8048
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=DJiLo7NDzd-lp-O3v7Cv8WejyjXPV_6_RmfltKO9fvE,2165
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CqB_8Xv32rtpLCaEEPi6DbRZm4ph5TWS-LfqIHXUIUA,1097
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
@@ -78,15 +78,15 @@ nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=r0ir4xYSdf_jebM0x09qaO6nJsvsiRQDyM0fs80ppOQ,2347
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
-nshtrainer/trainer/trainer.py,sha256=KXsvAhgVgYjmYfoqzH_qoQXqd6nVx7-vs9ObQJpwbIk,19140
+nshtrainer/trainer/trainer.py,sha256=Leh3ADxoYsRWlJFIW20netohLcKx0XxUrRhD9LM4jws,19201
 nshtrainer/util/_environment_info.py,sha256=gIdq9TJgzGCdcVzZxjHcwYasJ_HmEGVHbvE-KJVVtWs,24187
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
 nshtrainer/util/environment.py,sha256=AeW_kLl-N70wmb6L_JLz1wRj0kA70xs6RCmc9iUqczE,4159
-nshtrainer/util/path.py,sha256=A_Ocag3_hbwns_zAxFDlH-5eVHWFlcy2DKxHQ7jddvk,837
+nshtrainer/util/path.py,sha256=WbPWXpu5LIDocQihQC3-72qxN1sa6-d1kPOmKDR-NC8,1520
 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.22.1.dist-info/METADATA,sha256=XF3QXKbeAN7I5vYHNbjExlV_6CF8QgPqPYFsCxs52rA,935
-nshtrainer-0.22.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.22.1.dist-info/RECORD,,
+nshtrainer-0.23.0.dist-info/METADATA,sha256=wkbqsz6A4d0h1u-8CCZwfYYmqLm7YjirdnS-fTA-mkI,935
+nshtrainer-0.23.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.23.0.dist-info/RECORD,,