nshtrainer 0.11.3__tar.gz → 0.11.5__tar.gz

{nshtrainer-0.11.3 → nshtrainer-0.11.5}/PKG-INFO

@@ -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 → nshtrainer-0.11.5}/pyproject.toml

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

{nshtrainer-0.11.3 → nshtrainer-0.11.5}/src/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()