nshtrainer 0.10.12__tar.gz → 0.10.14__tar.gz

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

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

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

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

{nshtrainer-0.10.12 → nshtrainer-0.10.14}/src/nshtrainer/_checkpoint/saver.py

@@ -36,7 +36,8 @@ def _link_checkpoint(
             # fall back to copying the file
             shutil.copy(filepath, linkpath)
 
-        _link_checkpoint_metadata(filepath, linkpath)
+        if metadata:
+            _link_checkpoint_metadata(filepath, linkpath)
     if barrier:
         trainer.strategy.barrier()
 
@@ -44,9 +45,17 @@ def _link_checkpoint(
 def _remove_checkpoint(
     trainer: Trainer,
     filepath: str | Path | os.PathLike,
-    remove_metadata: bool = True,
+    *,
+    metadata: bool,
+    barrier: bool,
 ):
     if not isinstance(filepath, Path):
         filepath = Path(filepath)
-    trainer.strategy.remove_checkpoint(filepath)
-    _remove_checkpoint_metadata(filepath)
+
+    if trainer.is_global_zero:
+        trainer.strategy.remove_checkpoint(filepath)
+        if metadata:
+            _remove_checkpoint_metadata(filepath)
+
+    if barrier:
+        trainer.strategy.barrier()

{nshtrainer-0.10.12 → nshtrainer-0.10.14}/src/nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -69,7 +69,7 @@ class LatestEpochCheckpoint(Checkpoint):
 
     def _remove_checkpoints(self, trainer: Trainer, ckpt_paths: list[Path]):
         for ckpt_path in ckpt_paths:
-            _remove_checkpoint(trainer, ckpt_path, remove_metadata=True)
+            _remove_checkpoint(trainer, ckpt_path, metadata=True, barrier=False)
 
     def _remove_old_checkpoints(self, trainer: Trainer):
         if (latest_k := self.config.latest_k) == "all":

{nshtrainer-0.10.12 → nshtrainer-0.10.14}/src/nshtrainer/callbacks/model_checkpoint.py

@@ -202,4 +202,4 @@ class ModelCheckpoint(_ModelCheckpoint):
 
     @override
     def _remove_checkpoint(self, trainer: Trainer, filepath: str):
-        return _remove_checkpoint(trainer, filepath, remove_metadata=True)
+        return _remove_checkpoint(trainer, filepath, metadata=True, barrier=False)

{nshtrainer-0.10.12 → nshtrainer-0.10.14}/src/nshtrainer/callbacks/wandb_watch.py

@@ -12,13 +12,36 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
+class WandbWatchConfig(CallbackConfigBase):
+    name: Literal["wandb_watch"] = "wandb_watch"
+
+    enabled: bool = True
+    """Enable watching the model for wandb."""
+
+    log: str | None = None
+    """Log type for wandb."""
+
+    log_graph: bool = True
+    """Whether to log the graph for wandb."""
+
+    log_freq: int = 100
+    """Log frequency for wandb."""
+
+    def __bool__(self):
+        return self.enabled
+
+    @override
+    def create_callbacks(self, root_config):
+        yield WandbWatchCallback(self)
+
+
 @runtime_checkable
 class _HasWandbLogModuleProtocol(Protocol):
     def wandb_log_module(self) -> nn.Module | None: ...
 
 
 class WandbWatchCallback(Callback):
-    def __init__(self, config: "WandbWatchConfig"):
+    def __init__(self, config: WandbWatchConfig):
         super().__init__()
 
         self.config = config
@@ -78,26 +101,3 @@ class WandbWatchCallback(Callback):
             log_graph=self.config.log_graph,
         )
         setattr(pl_module, "_model_watched", True)
-
-
-class WandbWatchConfig(CallbackConfigBase):
-    name: Literal["wandb_watch"] = "wandb_watch"
-
-    enabled: bool = True
-    """Enable watching the model for wandb."""
-
-    log: str | None = None
-    """Log type for wandb."""
-
-    log_graph: bool = True
-    """Whether to log the graph for wandb."""
-
-    log_freq: int = 100
-    """Log frequency for wandb."""
-
-    def __bool__(self):
-        return self.enabled
-
-    @override
-    def create_callbacks(self, root_config):
-        yield WandbWatchCallback(self)

{nshtrainer-0.10.12 → nshtrainer-0.10.14}/src/nshtrainer/model/config.py

@@ -20,6 +20,7 @@ from typing import (
 
 import nshconfig as C
 import numpy as np
+import pkg_resources
 import torch
 from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
@@ -213,7 +214,7 @@ class BaseLoggerConfig(C.Config, ABC):
     """Enable this logger."""
 
     priority: int = 0
-    """Priority of the logger. Higher values are logged first."""
+    """Priority of the logger. Higher priority loggers are created first."""
 
     log_dir: DirectoryPath | None = None
     """Directory to save the logs to. If None, will use the default log directory for the trainer."""
@@ -266,7 +267,8 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
     """Enable WandB logging."""
 
     priority: int = 2
-    """Priority of the logger. Higher values are logged first."""
+    """Priority of the logger. Higher priority loggers are created first,
+    and the highest priority logger is the "main" logger for PyTorch Lightning."""
 
     project: str | None = None
     """WandB project name to use for the logger. If None, will use the root config's project name."""
@@ -286,8 +288,17 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
     offline: bool = False
     """Whether to run WandB in offline mode."""
 
-    def offline_(self):
-        self.offline = True
+    use_wandb_core: bool = False
+    """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.
+    """
+
+    def offline_(self, value: bool = True):
+        self.offline = value
+        return self
+
+    def core_(self, value: bool = True):
+        self.use_wandb_core = value
         return self
 
     @override
@@ -295,6 +306,25 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
         if not self.enabled:
             return None
 
+        # If `wandb-core` is enabled, we should use the new backend.
+        if self.use_wandb_core:
+            try:
+                import wandb  # type: ignore
+
+                # The minimum version that supports the new backend is 0.17.5
+                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."
+                    )
+                else:
+                    wandb.require("core")
+                    log.critical("Using the `wandb-core` backend for WandB.")
+            except ImportError:
+                pass
+
         from lightning.pytorch.loggers.wandb import WandbLogger
 
         save_dir = root_config.directory._resolve_log_directory_for_logger(
@@ -329,7 +359,7 @@ class CSVLoggerConfig(BaseLoggerConfig):
     """Enable CSV logging."""
 
     priority: int = 0
-    """Priority of the logger. Higher values are logged first."""
+    """Priority of the logger. Higher priority loggers are created first."""
 
     prefix: str = ""
     """A string to put at the beginning of metric keys."""
@@ -383,7 +413,7 @@ class TensorboardLoggerConfig(BaseLoggerConfig):
     """Enable TensorBoard logging."""
 
     priority: int = 2
-    """Priority of the logger. Higher values are logged first."""
+    """Priority of the logger. Higher priority loggers are created first."""
 
     log_graph: bool = False
     """