nshtrainer 0.10.15__tar.gz → 0.10.17__tar.gz

{nshtrainer-0.10.15 → nshtrainer-0.10.17}/PKG-INFO

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

{nshtrainer-0.10.15 → nshtrainer-0.10.17}/pyproject.toml

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

{nshtrainer-0.10.15 → nshtrainer-0.10.17}/src/nshtrainer/_checkpoint/metadata.py

@@ -105,10 +105,8 @@ def _write_checkpoint_metadata(
 
 
 def _remove_checkpoint_metadata(checkpoint_path: Path):
-    for path in (
-        checkpoint_path.with_suffix(METADATA_PATH_SUFFIX),
-        checkpoint_path.with_suffix(HPARAMS_PATH_SUFFIX),
-    ):
+    for suffix in (METADATA_PATH_SUFFIX, HPARAMS_PATH_SUFFIX):
+        path = checkpoint_path.with_suffix(suffix)
         try:
             path.unlink(missing_ok=True)
         except Exception as e:
@@ -122,11 +120,9 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
     _remove_checkpoint_metadata(linked_checkpoint_path)
 
     # Link the metadata files to the new checkpoint
-    for path in (
-        checkpoint_path.with_suffix(METADATA_PATH_SUFFIX),
-        checkpoint_path.with_suffix(HPARAMS_PATH_SUFFIX),
-    ):
-        linked_path = linked_checkpoint_path.with_suffix(path.suffix)
+    for suffix in (METADATA_PATH_SUFFIX, HPARAMS_PATH_SUFFIX):
+        path = checkpoint_path.with_suffix(suffix)
+        linked_path = linked_checkpoint_path.with_suffix(suffix)
         try:
             try:
                 linked_path.symlink_to(path)

{nshtrainer-0.10.15 → nshtrainer-0.10.17}/src/nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -64,7 +64,7 @@ class LatestEpochCheckpoint(Checkpoint):
         filename = self.config.filename.format(
             epoch=trainer.current_epoch, step=trainer.global_step
         )
-        filename = f"{self.PREFIX}{filename}.{self.EXTENSION}"
+        filename = f"{self.PREFIX}{filename}{self.EXTENSION}"
         return self.dirpath / filename
 
     def _remove_checkpoints(self, trainer: Trainer, ckpt_paths: list[Path]):
@@ -95,7 +95,9 @@ class LatestEpochCheckpoint(Checkpoint):
 
     def _save_new_checkpoint(self, trainer: Trainer):
         # Remove old checkpoints
-        self._remove_old_checkpoints(trainer)
+        if trainer.is_global_zero:
+            self._remove_old_checkpoints(trainer)
+        trainer.strategy.barrier()
 
         # Save the new checkpoint
         filepath = self._ckpt_path(trainer)

{nshtrainer-0.10.15 → nshtrainer-0.10.17}/src/nshtrainer/model/config.py

@@ -288,7 +288,7 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
     offline: bool = False
     """Whether to run WandB in offline mode."""
 
-    use_wandb_core: bool = False
+    use_wandb_core: bool = True
     """Whether to use the new `wandb-core` backend for WandB.
     `wandb-core` is a new backend for WandB that is faster and more efficient than the old backend.
     """