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

nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -106,6 +106,26 @@ class BestCheckpoint(Checkpoint):
             reverse=(self.metric.mode == "min"),
         )
 
+    def _create_symlink(self, trainer: Trainer, best_ckpt_path: Path):
+        # Resolve the symlink filename
+        if (symlink_filename := self._best_symlink_filename()) is None:
+            return
+
+        # If the symlink already exists and points to the best checkpoint,
+        # then we don't need to create a new symlink.
+        symlink_path = self.dirpath / symlink_filename
+        if symlink_path.exists() and symlink_path.resolve() == best_ckpt_path:
+            return
+
+        _link_checkpoint(
+            trainer,
+            best_ckpt_path,
+            symlink_path,
+            metadata=True,
+            barrier=False,
+        )
+        log.debug(f"Created best symlink: {symlink_path}")
+
     def _save_best_checkpoint(self, trainer: Trainer):
         if (current := self._get_metric_value(trainer.callback_metrics)) is None:
             log.warning(
@@ -130,6 +150,7 @@ class BestCheckpoint(Checkpoint):
         # Save the current model
         filepath = self._ckpt_path(trainer)
         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
@@ -143,13 +164,9 @@ class BestCheckpoint(Checkpoint):
             )
 
         # Create symlink to best model
-        if (symlink_filename := self._best_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 best symlink: {symlink_path}")
+        if sorted_ckpts:
+            _, best_ckpt_path = sorted_ckpts[0]
+            self._create_symlink(trainer, best_ckpt_path)
+
+        # Barrier to ensure all processes have saved the checkpoint before continuing
+        trainer.strategy.barrier()

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

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

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

@@ -11,7 +11,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=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,3555
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=zrEVCGFikfkt0iOMceOFzXsZG2-6QrqY79RKBCS7bu4,738
-nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=9qp6z4AmriakkiMbay2_QR97RxNLejHAdvJPU__8Y_c,5519
+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/on_exception_checkpoint.py,sha256=s8tOHrnb_uVqLVeV2K38ZszXrXPTEGdDVfXuXgo_KDQ,3277
@@ -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.3.dist-info/METADATA,sha256=lJqbTAjie2SQ9HT_yPyHIfrwTJQ_ognnuqD9YJ_QNPQ,860
-nshtrainer-0.11.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.11.3.dist-info/RECORD,,
+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,,