nshtrainer 0.13.1__tar.gz → 0.14.1__tar.gz

{nshtrainer-0.13.1 → nshtrainer-0.14.1}/PKG-INFO

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

{nshtrainer-0.13.1 → nshtrainer-0.14.1}/pyproject.toml

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

{nshtrainer-0.13.1 → nshtrainer-0.14.1}/src/nshtrainer/callbacks/early_stopping.py

@@ -73,9 +73,7 @@ class EarlyStopping(_EarlyStopping):
 
     @override
     @staticmethod
-    def _log_info(
-        trainer: Trainer | None, message: str, log_rank_zero_only: bool
-    ) -> None:
+    def _log_info(trainer: Trainer | None, message: str, log_rank_zero_only: bool):
         rank = _get_rank()
         if trainer is not None and trainer.world_size <= 1:
             rank = None

{nshtrainer-0.13.1 → nshtrainer-0.14.1}/src/nshtrainer/ll/__init__.py

@@ -23,7 +23,6 @@ from .log import pretty as pretty
 from .lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from .model import Base as Base
 from .model import BaseConfig as BaseConfig
-from .model import BaseLoggerConfig as BaseLoggerConfig
 from .model import BaseProfilerConfig as BaseProfilerConfig
 from .model import CheckpointLoadingConfig as CheckpointLoadingConfig
 from .model import CheckpointSavingConfig as CheckpointSavingConfig
@@ -48,7 +47,6 @@ from .model import PrimaryMetricConfig as PrimaryMetricConfig
 from .model import ReproducibilityConfig as ReproducibilityConfig
 from .model import SanityCheckingConfig as SanityCheckingConfig
 from .model import TrainerConfig as TrainerConfig
-from .model import WandbWatchConfig as WandbWatchConfig
 from .nn import TypedModuleDict as TypedModuleDict
 from .nn import TypedModuleList as TypedModuleList
 from .optimizer import OptimizerConfig as OptimizerConfig

nshtrainer-0.14.1/src/nshtrainer/loggers/__init__.py

@@ -0,0 +1,13 @@
+from typing import Annotated, TypeAlias
+
+import nshconfig as C
+
+from ._base import BaseLoggerConfig as BaseLoggerConfig
+from .csv import CSVLoggerConfig as CSVLoggerConfig
+from .tensorboard import TensorboardLoggerConfig as TensorboardLoggerConfig
+from .wandb import WandbLoggerConfig as WandbLoggerConfig
+
+LoggerConfig: TypeAlias = Annotated[
+    CSVLoggerConfig | TensorboardLoggerConfig | WandbLoggerConfig,
+    C.Field(discriminator="name"),
+]

nshtrainer-0.14.1/src/nshtrainer/loggers/_base.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+import nshconfig as C
+from lightning.pytorch.loggers import Logger
+
+if TYPE_CHECKING:
+    from ..model import BaseConfig
+
+
+class BaseLoggerConfig(C.Config, ABC):
+    enabled: bool = True
+    """Enable this logger."""
+
+    priority: int = 0
+    """Priority of the logger. Higher priority loggers are created first."""
+
+    log_dir: C.DirectoryPath | None = None
+    """Directory to save the logs to. If None, will use the default log directory for the trainer."""
+
+    @abstractmethod
+    def create_logger(self, root_config: "BaseConfig") -> Logger | None: ...
+
+    def disable_(self):
+        self.enabled = False
+        return self

nshtrainer-0.14.1/src/nshtrainer/loggers/csv.py

@@ -0,0 +1,40 @@
+from typing import Literal
+
+from typing_extensions import override
+
+from ._base import BaseLoggerConfig
+
+
+class CSVLoggerConfig(BaseLoggerConfig):
+    name: Literal["csv"] = "csv"
+
+    enabled: bool = True
+    """Enable CSV logging."""
+
+    priority: int = 0
+    """Priority of the logger. Higher priority loggers are created first."""
+
+    prefix: str = ""
+    """A string to put at the beginning of metric keys."""
+
+    flush_logs_every_n_steps: int = 100
+    """How often to flush logs to disk."""
+
+    @override
+    def create_logger(self, root_config):
+        if not self.enabled:
+            return None
+
+        from lightning.pytorch.loggers.csv_logs import CSVLogger
+
+        save_dir = root_config.directory._resolve_log_directory_for_logger(
+            root_config.id,
+            self,
+        )
+        return CSVLogger(
+            save_dir=save_dir,
+            name=root_config.run_name,
+            version=root_config.id,
+            prefix=self.prefix,
+            flush_logs_every_n_steps=self.flush_logs_every_n_steps,
+        )

nshtrainer-0.14.1/src/nshtrainer/loggers/tensorboard.py

@@ -0,0 +1,73 @@
+import logging
+from typing import Literal
+
+import nshconfig as C
+from typing_extensions import override
+
+from ._base import BaseLoggerConfig
+
+log = logging.getLogger(__name__)
+
+
+def _tensorboard_available():
+    try:
+        from lightning.fabric.loggers.tensorboard import (
+            _TENSORBOARD_AVAILABLE,
+            _TENSORBOARDX_AVAILABLE,
+        )
+
+        if not _TENSORBOARD_AVAILABLE and not _TENSORBOARDX_AVAILABLE:
+            log.warning(
+                "TensorBoard/TensorBoardX not found. Disabling TensorBoardLogger. "
+                "Please install TensorBoard with `pip install tensorboard` or "
+                "TensorBoardX with `pip install tensorboardx` to enable TensorBoard logging."
+            )
+            return False
+        return True
+    except ImportError:
+        return False
+
+
+class TensorboardLoggerConfig(BaseLoggerConfig):
+    name: Literal["tensorboard"] = "tensorboard"
+
+    enabled: bool = C.Field(default_factory=lambda: _tensorboard_available())
+    """Enable TensorBoard logging."""
+
+    priority: int = 2
+    """Priority of the logger. Higher priority loggers are created first."""
+
+    log_graph: bool = False
+    """
+    Adds the computational graph to tensorboard. This requires that
+        the user has defined the `self.example_input_array` attribute in their
+        model.
+    """
+
+    default_hp_metric: bool = True
+    """
+    Enables a placeholder metric with key `hp_metric` when `log_hyperparams` is
+        called without a metric (otherwise calls to log_hyperparams without a metric are ignored).
+    """
+
+    prefix: str = ""
+    """A string to put at the beginning of metric keys."""
+
+    @override
+    def create_logger(self, root_config):
+        if not self.enabled:
+            return None
+
+        from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
+
+        save_dir = root_config.directory._resolve_log_directory_for_logger(
+            root_config.id,
+            self,
+        )
+        return TensorBoardLogger(
+            save_dir=save_dir,
+            name=root_config.run_name,
+            version=root_config.id,
+            log_graph=self.log_graph,
+            default_hp_metric=self.default_hp_metric,
+        )

nshtrainer-0.14.1/src/nshtrainer/loggers/wandb.py

@@ -0,0 +1,163 @@
+import logging
+from typing import TYPE_CHECKING, Literal
+
+import nshconfig as C
+import pkg_resources
+from lightning.pytorch import Callback, LightningModule, Trainer
+from typing_extensions import override
+
+from ..callbacks import WandbWatchConfig
+from ..callbacks.base import CallbackConfigBase
+from ._base import BaseLoggerConfig
+
+if TYPE_CHECKING:
+    from ..model.config import BaseConfig
+
+log = logging.getLogger(__name__)
+
+
+def _project_name(
+    root_config: "BaseConfig",
+    default_project: str = "lightning_logs",
+):
+    # If the config has a project name, use that.
+    if project := root_config.project:
+        return project
+
+    # Otherwise, we should use the name of the module that the config is defined in,
+    #   if we can find it.
+    # If this isn't in a module, use the default project name.
+    if not (module := root_config.__module__):
+        return default_project
+
+    # If the module is a package, use the package name.
+    if not (module := module.split(".", maxsplit=1)[0].strip()):
+        return default_project
+
+    return module
+
+
+def _wandb_available():
+    try:
+        from lightning.pytorch.loggers.wandb import _WANDB_AVAILABLE
+
+        if not _WANDB_AVAILABLE:
+            log.warning("WandB not found. Disabling WandbLogger.")
+            return False
+        return True
+    except ImportError:
+        return False
+
+
+class FinishWandbOnTeardownCallback(Callback):
+    @override
+    def teardown(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        stage: str,
+    ):
+        try:
+            import wandb  # type: ignore
+        except ImportError:
+            return
+
+        if wandb.run is None:
+            return
+
+        wandb.finish()
+
+
+class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
+    name: Literal["wandb"] = "wandb"
+
+    enabled: bool = C.Field(default_factory=lambda: _wandb_available())
+    """Enable WandB logging."""
+
+    priority: int = 2
+    """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."""
+
+    log_model: bool | Literal["all"] = False
+    """
+    Whether to log the model checkpoints to wandb.
+    Valid values are:
+        - False: Do not log the model checkpoints.
+        - True: Log the latest model checkpoint.
+        - "all": Log all model checkpoints.
+    """
+
+    watch: WandbWatchConfig | None = WandbWatchConfig()
+    """WandB model watch configuration. Used to log model architecture, gradients, and parameters."""
+
+    offline: bool = False
+    """Whether to run WandB in offline mode."""
+
+    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.
+    """
+
+    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
+    def create_logger(self, root_config):
+        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 "
+                        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")
+                    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(
+            root_config.id,
+            self,
+        )
+        return WandbLogger(
+            save_dir=save_dir,
+            project=self.project or _project_name(root_config),
+            name=root_config.run_name,
+            version=root_config.id,
+            log_model=self.log_model,
+            notes=(
+                "\n".join(f"- {note}" for note in root_config.notes)
+                if root_config.notes
+                else None
+            ),
+            tags=root_config.tags,
+            offline=self.offline,
+        )
+
+    @override
+    def create_callbacks(self, root_config):
+        yield FinishWandbOnTeardownCallback()
+
+        if self.watch:
+            yield from self.watch.create_callbacks(root_config)

{nshtrainer-0.13.1 → nshtrainer-0.14.1}/src/nshtrainer/model/__init__.py

@@ -3,7 +3,6 @@ from typing_extensions import TypeAlias
 from .base import Base as Base
 from .base import LightningModuleBase as LightningModuleBase
 from .config import BaseConfig as BaseConfig
-from .config import BaseLoggerConfig as BaseLoggerConfig
 from .config import BaseProfilerConfig as BaseProfilerConfig
 from .config import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
 from .config import CheckpointLoadingConfig as CheckpointLoadingConfig
@@ -22,6 +21,5 @@ from .config import PrimaryMetricConfig as PrimaryMetricConfig
 from .config import ReproducibilityConfig as ReproducibilityConfig
 from .config import SanityCheckingConfig as SanityCheckingConfig
 from .config import TrainerConfig as TrainerConfig
-from .config import WandbWatchConfig as WandbWatchConfig
 
 ConfigList: TypeAlias = list[tuple[BaseConfig, type[LightningModuleBase]]]

{nshtrainer-0.13.1 → nshtrainer-0.14.1}/src/nshtrainer/model/config.py

@@ -19,7 +19,6 @@ 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
@@ -31,7 +30,6 @@ from lightning.pytorch.plugins.layer_sync import LayerSync
 from lightning.pytorch.plugins.precision.precision import Precision
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.strategies.strategy import Strategy
-from pydantic import DirectoryPath
 from typing_extensions import Self, TypedDict, TypeVar, override
 
 from .._checkpoint.loader import CheckpointLoadingConfig
@@ -41,9 +39,14 @@ from ..callbacks import (
     EarlyStoppingConfig,
     LastCheckpointCallbackConfig,
     OnExceptionCheckpointCallbackConfig,
-    WandbWatchConfig,
 )
 from ..callbacks.base import CallbackConfigBase
+from ..loggers import (
+    CSVLoggerConfig,
+    LoggerConfig,
+    TensorboardLoggerConfig,
+    WandbLoggerConfig,
+)
 from ..metrics import MetricConfig
 from ..util._environment_info import EnvironmentConfig
 
@@ -205,255 +208,6 @@ ProfilerConfig: TypeAlias = Annotated[
 ]
 
 
-class BaseLoggerConfig(C.Config, ABC):
-    enabled: bool = True
-    """Enable this logger."""
-
-    priority: int = 0
-    """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."""
-
-    @abstractmethod
-    def create_logger(self, root_config: "BaseConfig") -> Logger | None: ...
-
-    def disable_(self):
-        self.enabled = False
-        return self
-
-
-def _project_name(
-    root_config: "BaseConfig",
-    default_project: str = "lightning_logs",
-):
-    # If the config has a project name, use that.
-    if project := root_config.project:
-        return project
-
-    # Otherwise, we should use the name of the module that the config is defined in,
-    #   if we can find it.
-    # If this isn't in a module, use the default project name.
-    if not (module := root_config.__module__):
-        return default_project
-
-    # If the module is a package, use the package name.
-    if not (module := module.split(".", maxsplit=1)[0].strip()):
-        return default_project
-
-    return module
-
-
-def _wandb_available():
-    try:
-        from lightning.pytorch.loggers.wandb import _WANDB_AVAILABLE
-
-        if not _WANDB_AVAILABLE:
-            log.warning("WandB not found. Disabling WandbLogger.")
-            return False
-        return True
-    except ImportError:
-        return False
-
-
-class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
-    name: Literal["wandb"] = "wandb"
-
-    enabled: bool = C.Field(default_factory=lambda: _wandb_available())
-    """Enable WandB logging."""
-
-    priority: int = 2
-    """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."""
-
-    log_model: bool | Literal["all"] = False
-    """
-    Whether to log the model checkpoints to wandb.
-    Valid values are:
-        - False: Do not log the model checkpoints.
-        - True: Log the latest model checkpoint.
-        - "all": Log all model checkpoints.
-    """
-
-    watch: WandbWatchConfig = WandbWatchConfig()
-    """WandB model watch configuration. Used to log model architecture, gradients, and parameters."""
-
-    offline: bool = False
-    """Whether to run WandB in offline mode."""
-
-    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.
-    """
-
-    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
-    def create_logger(self, root_config):
-        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 "
-                        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")
-                    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(
-            root_config.id,
-            self,
-        )
-        return WandbLogger(
-            save_dir=save_dir,
-            project=self.project or _project_name(root_config),
-            name=root_config.run_name,
-            version=root_config.id,
-            log_model=self.log_model,
-            notes=(
-                "\n".join(f"- {note}" for note in root_config.notes)
-                if root_config.notes
-                else None
-            ),
-            tags=root_config.tags,
-            offline=self.offline,
-        )
-
-    @override
-    def create_callbacks(self, root_config):
-        if self.watch:
-            yield from self.watch.create_callbacks(root_config)
-
-
-class CSVLoggerConfig(BaseLoggerConfig):
-    name: Literal["csv"] = "csv"
-
-    enabled: bool = True
-    """Enable CSV logging."""
-
-    priority: int = 0
-    """Priority of the logger. Higher priority loggers are created first."""
-
-    prefix: str = ""
-    """A string to put at the beginning of metric keys."""
-
-    flush_logs_every_n_steps: int = 100
-    """How often to flush logs to disk."""
-
-    @override
-    def create_logger(self, root_config):
-        if not self.enabled:
-            return None
-
-        from lightning.pytorch.loggers.csv_logs import CSVLogger
-
-        save_dir = root_config.directory._resolve_log_directory_for_logger(
-            root_config.id,
-            self,
-        )
-        return CSVLogger(
-            save_dir=save_dir,
-            name=root_config.run_name,
-            version=root_config.id,
-            prefix=self.prefix,
-            flush_logs_every_n_steps=self.flush_logs_every_n_steps,
-        )
-
-
-def _tensorboard_available():
-    try:
-        from lightning.fabric.loggers.tensorboard import (
-            _TENSORBOARD_AVAILABLE,
-            _TENSORBOARDX_AVAILABLE,
-        )
-
-        if not _TENSORBOARD_AVAILABLE and not _TENSORBOARDX_AVAILABLE:
-            log.warning(
-                "TensorBoard/TensorBoardX not found. Disabling TensorBoardLogger. "
-                "Please install TensorBoard with `pip install tensorboard` or "
-                "TensorBoardX with `pip install tensorboardx` to enable TensorBoard logging."
-            )
-            return False
-        return True
-    except ImportError:
-        return False
-
-
-class TensorboardLoggerConfig(BaseLoggerConfig):
-    name: Literal["tensorboard"] = "tensorboard"
-
-    enabled: bool = C.Field(default_factory=lambda: _tensorboard_available())
-    """Enable TensorBoard logging."""
-
-    priority: int = 2
-    """Priority of the logger. Higher priority loggers are created first."""
-
-    log_graph: bool = False
-    """
-    Adds the computational graph to tensorboard. This requires that
-        the user has defined the `self.example_input_array` attribute in their
-        model.
-    """
-
-    default_hp_metric: bool = True
-    """
-    Enables a placeholder metric with key `hp_metric` when `log_hyperparams` is
-        called without a metric (otherwise calls to log_hyperparams without a metric are ignored).
-    """
-
-    prefix: str = ""
-    """A string to put at the beginning of metric keys."""
-
-    @override
-    def create_logger(self, root_config):
-        if not self.enabled:
-            return None
-
-        from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
-
-        save_dir = root_config.directory._resolve_log_directory_for_logger(
-            root_config.id,
-            self,
-        )
-        return TensorBoardLogger(
-            save_dir=save_dir,
-            name=root_config.run_name,
-            version=root_config.id,
-            log_graph=self.log_graph,
-            default_hp_metric=self.default_hp_metric,
-        )
-
-
-LoggerConfig: TypeAlias = Annotated[
-    WandbLoggerConfig | CSVLoggerConfig | TensorboardLoggerConfig,
-    C.Field(discriminator="name"),
-]
-
-
 class LoggingConfig(CallbackConfigBase):
     enabled: bool = True
     """Enable experiment tracking."""
@@ -474,29 +228,32 @@ class LoggingConfig(CallbackConfigBase):
     """If enabled, will automatically save logged metrics using ActSave (if nshutils is installed)."""
 
     @property
-    def wandb(self) -> WandbLoggerConfig | None:
+    def wandb(self):
         return next(
             (
                 logger
                 for logger in self.loggers
                 if isinstance(logger, WandbLoggerConfig)
             ),
+            None,
         )
 
     @property
-    def csv(self) -> CSVLoggerConfig | None:
+    def csv(self):
         return next(
             (logger for logger in self.loggers if isinstance(logger, CSVLoggerConfig)),
+            None,
         )
 
     @property
-    def tensorboard(self) -> TensorboardLoggerConfig | None:
+    def tensorboard(self):
         return next(
             (
                 logger
                 for logger in self.loggers
                 if isinstance(logger, TensorboardLoggerConfig)
             ),
+            None,
         )
 
     def create_loggers(self, root_config: "BaseConfig"):
@@ -509,9 +266,8 @@ class LoggingConfig(CallbackConfigBase):
         Returns:
             list[Logger]: A list of constructed loggers.
         """
-        loggers: list[Logger] = []
         if not self.enabled:
-            return loggers
+            return
 
         for logger_config in sorted(
             self.loggers,
@@ -522,8 +278,7 @@ class LoggingConfig(CallbackConfigBase):
                 continue
             if (logger := logger_config.create_logger(root_config)) is None:
                 continue
-            loggers.append(logger)
-        return loggers
+            yield logger
 
     @override
     def create_callbacks(self, root_config):

{nshtrainer-0.13.1 → nshtrainer-0.14.1}/src/nshtrainer/trainer/trainer.py

@@ -244,7 +244,7 @@ class Trainer(LightningTrainer):
             log.critical(f"Disabling logger because {config.trainer.logging.enabled=}.")
             kwargs["logger"] = False
         else:
-            _update_kwargs(logger=config.trainer.logging.create_loggers(config))
+            _update_kwargs(logger=list(config.trainer.logging.create_loggers(config)))
 
         if config.trainer.auto_determine_num_nodes:
             # When num_nodes is auto, we need to detect the number of nodes.