nshtrainer 0.10.14__tar.gz → 0.10.16__tar.gz

{nshtrainer-0.10.14 → nshtrainer-0.10.16}/PKG-INFO

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

{nshtrainer-0.10.14 → nshtrainer-0.10.16}/pyproject.toml

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

{nshtrainer-0.10.14 → nshtrainer-0.10.16}/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.14 → nshtrainer-0.10.16}/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.14 → nshtrainer-0.10.16}/src/nshtrainer/model/config.py

@@ -315,9 +315,10 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
                 if pkg_resources.parse_version(
                     wandb.__version__
                 ) < pkg_resources.parse_version("0.17.5"):
-                    log.warning(
-                        "The version of WandB installed does not support the `wandb-core` backend. "
-                        "Unable to use the `wandb-core` backend for WandB."
+                    raise ValueError(
+                        "The version of WandB installed does not support the `wandb-core` backend "
+                        f"(expected version >= 0.17.5, found version {wandb.__version__}). "
+                        "Please either upgrade to a newer version of WandB or disable the `use_wandb_core` option."
                     )
                 else:
                     wandb.require("core")