nshtrainer 0.10.13__py3-none-any.whl → 0.10.14__py3-none-any.whl

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/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
     """

{nshtrainer-0.10.13.dist-info → nshtrainer-0.10.14.dist-info}/METADATA

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

{nshtrainer-0.10.13.dist-info → nshtrainer-0.10.14.dist-info}/RECORD

@@ -23,7 +23,7 @@ nshtrainer/callbacks/on_exception_checkpoint.py,sha256=x42BYZ2ejf2rhqPLCmT5nyWKh
 nshtrainer/callbacks/print_table.py,sha256=_FdAHhqylWGk4Z0c2FrLFeiMA4jhfA_beZRK_BHpzmE,2837
 nshtrainer/callbacks/throughput_monitor.py,sha256=H_ocXErZxUO3dxFk8Tx_VQdpI9E_Ztvqof5WtFevLyQ,1838
 nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
-nshtrainer/callbacks/wandb_watch.py,sha256=EJ93mtJlph4BZsXh8HJPNiw2VNSm2N6TOwpCwqRAeKI,2923
+nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=dGBTDDtlBU6c-ZlVQOCnTW7SjTB5hczWsOWEdUWjvkA,4385
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
@@ -52,7 +52,7 @@ nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy0
 nshtrainer/metrics/_config.py,sha256=hWWS4IXENRyH3RmJ7z1Wx1n3Lt1sNMlGOrcU6PW15o0,1104
 nshtrainer/model/__init__.py,sha256=NpvyQHmGaHB8xdraHmm8l7kDHLmvJSgBNQKkfYqtgyI,1454
 nshtrainer/model/base.py,sha256=AXRfEsFAT0Ln7zjYVPU5NgtHS_c8FZM-M4pyLamO7OA,17516
-nshtrainer/model/config.py,sha256=65UDzt3ZZFUQaHMlK7f9wzwyGH3cDyHGtjZ2eOjHvVo,53360
+nshtrainer/model/config.py,sha256=z6kSkTirvRsyW3YIDTG1uAmK4fCC-gNAQrMi7Osxiow,54643
 nshtrainer/model/modules/callback.py,sha256=K0-cyEtBcQhI7Q2e-AGTE8T-GghUPY9DYmneU6ULV6g,6401
 nshtrainer/model/modules/debug.py,sha256=Yy7XEdPou9BkCsD5hJchwJGmCVGrfUru5g9VjPM4uAw,1120
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
@@ -79,6 +79,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.10.13.dist-info/METADATA,sha256=HpGl8_E6q2l2nQrIzU5ibNEyUXj8adF8cxMzouUSpAg,696
-nshtrainer-0.10.13.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.10.13.dist-info/RECORD,,
+nshtrainer-0.10.14.dist-info/METADATA,sha256=7aheATImk1o69ugMxCBAbWhgWfRgRKYbP_QeFaqUGbM,696
+nshtrainer-0.10.14.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.10.14.dist-info/RECORD,,