nshtrainer 0.44.1__py3-none-any.whl → 1.0.0b10__py3-none-any.whl

nshtrainer/loggers/wandb.py

@@ -15,23 +15,24 @@ from ..callbacks.wandb_watch import WandbWatchCallbackConfig
 from ._base import BaseLoggerConfig
 
 if TYPE_CHECKING:
-    from ..model.config import BaseConfig
+    from ..trainer._config import TrainerConfig
+
 
 log = logging.getLogger(__name__)
 
 
 def _project_name(
-    root_config: "BaseConfig",
+    trainer_config: TrainerConfig,
     default_project: str = "lightning_logs",
 ):
     # If the config has a project name, use that.
-    if project := root_config.project:
+    if project := trainer_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__):
+    if not (module := trainer_config.__module__):
         return default_project
 
     # If the module is a package, use the package name.
@@ -129,7 +130,7 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
                 assert_never(self.log_model)
 
     @override
-    def create_logger(self, root_config):
+    def create_logger(self, trainer_config):
         if not self.enabled:
             return None
 
@@ -171,31 +172,31 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
 
         from lightning.pytorch.loggers.wandb import WandbLogger
 
-        save_dir = root_config.directory._resolve_log_directory_for_logger(
-            root_config.id,
+        save_dir = trainer_config.directory._resolve_log_directory_for_logger(
+            trainer_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,
+            project=self.project or _project_name(trainer_config),
+            name=trainer_config.full_name,
+            version=trainer_config.id,
             log_model=self._lightning_log_model,
             notes=(
-                "\n".join(f"- {note}" for note in root_config.notes)
-                if root_config.notes
+                "\n".join(f"- {note}" for note in trainer_config.notes)
+                if trainer_config.notes
                 else None
             ),
-            tags=root_config.tags,
+            tags=trainer_config.tags,
             offline=self.offline,
         )
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield FinishWandbOnTeardownCallback()
 
         if self.watch:
-            yield from self.watch.create_callbacks(root_config)
+            yield from self.watch.create_callbacks(trainer_config)
 
         if self.log_code:
-            yield from self.log_code.create_callbacks(root_config)
+            yield from self.log_code.create_callbacks(trainer_config)

nshtrainer/lr_scheduler/reduce_lr_on_plateau.py

@@ -1,17 +1,14 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal, cast
+from typing import Literal
 
 from lightning.pytorch.utilities.types import LRSchedulerConfigType
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 from typing_extensions import override
 
-from ..model.config import MetricConfig
+from ..metrics._config import MetricConfig
 from ._base import LRSchedulerConfigBase, LRSchedulerMetadata
 
-if TYPE_CHECKING:
-    from ..model.base import BaseConfig
-
 
 class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
     """Reduce learning rate when a metric has stopped improving."""
@@ -48,9 +45,14 @@ class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
         self, optimizer, lightning_module
     ) -> LRSchedulerConfigType:
         if (metric := self.metric) is None:
-            lm_config = cast("BaseConfig", lightning_module.config)
+            from ..trainer import Trainer
+
+            assert isinstance(
+                trainer := lightning_module.trainer, Trainer
+            ), "The trainer must be a `nshtrainer.Trainer` instance."
+
             assert (
-                metric := lm_config.primary_metric
+                metric := trainer.hparams.primary_metric
             ) is not None, "Primary metric must be provided if metric is not specified."
 
         lr_scheduler = ReduceLROnPlateau(

nshtrainer/model/__init__.py

@@ -1,7 +1,3 @@
 from __future__ import annotations
 
 from .base import LightningModuleBase as LightningModuleBase
-from .config import BaseConfig as BaseConfig
-from .config import DirectoryConfig as DirectoryConfig
-from .config import MetricConfig as MetricConfig
-from .config import TrainerConfig as TrainerConfig

nshtrainer/model/base.py

@@ -1,28 +1,25 @@
 from __future__ import annotations
 
-import inspect
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import MutableMapping
-from typing import IO, TYPE_CHECKING, Any, Generic, Literal, cast
+from collections.abc import Mapping
+from typing import Any, Generic, Literal, cast
 
+import nshconfig as C
 import torch
 import torch.distributed
-from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
 from lightning.pytorch import LightningModule
 from lightning.pytorch.profilers import PassThroughProfiler, Profiler
-from lightning.pytorch.utilities.types import STEP_OUTPUT
-from typing_extensions import Self, TypeVar, override
+from typing_extensions import Never, TypeVar, deprecated, override
 
 from ..callbacks.rlp_sanity_checks import _RLPSanityCheckModuleMixin
-from ..util._environment_info import EnvironmentConfig
-from .config import BaseConfig
 from .mixins.callback import CallbackModuleMixin
+from .mixins.debug import _DebugModuleMixin, _trainer
 from .mixins.logger import LoggerLightningModuleMixin
 
 log = logging.getLogger(__name__)
 
-THparams = TypeVar("THparams", bound=BaseConfig, infer_variance=True)
+THparams = TypeVar("THparams", bound=C.Config, infer_variance=True)
 
 
 T = TypeVar("T", infer_variance=True)
@@ -53,7 +50,8 @@ VALID_REDUCE_OPS = (
 )
 
 
-class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
+class LightningModuleBase(
+    _DebugModuleMixin,
     _RLPSanityCheckModuleMixin,
     LoggerLightningModuleMixin,
     CallbackModuleMixin,
@@ -61,21 +59,36 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
     ABC,
     Generic[THparams],
 ):
-    # region Config
-    @torch.jit.unused
-    @property
-    def config(self) -> THparams:
-        return self.hparams
-
+    # region Debug
     @property
     def debug(self) -> bool:
         if torch.jit.is_scripting():
             return False
-        return self.config.debug
 
-    # endregion
+        if (trainer := self._trainer) is None:
+            return False
 
-    # region Debug
+        from ..trainer import Trainer
+
+        if not isinstance(trainer, Trainer):
+            return False
+
+        return trainer.debug
+
+    @debug.setter
+    def debug(self, value: bool):
+        if torch.jit.is_scripting():
+            return
+
+        if (trainer := self._trainer) is None:
+            return
+
+        from ..trainer import Trainer
+
+        if not isinstance(trainer, Trainer):
+            return
+
+        trainer.debug = value
 
     @torch.jit.unused
     def breakpoint(self, rank_zero_only: bool = True):
@@ -146,7 +159,7 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
         return object_list
 
     def barrier(self, name: str | None = None):
-        self.trainer.strategy.barrier(name=name)
+        return self.trainer.strategy.barrier(name=name)
 
     def reduce(
         self,
@@ -170,7 +183,7 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
     @override
     def __repr__(self):
         parts: list[str] = []
-        parts.append(f"config={self.hparams.concise_repr()}")
+        parts.append(f"hparams={repr(self.hparams)}")
         parts.append(f"device={self.device}")
         if self.debug:
             parts.append("debug=True")
@@ -178,85 +191,46 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
         parts_str = ", ".join(parts)
         return f"{self.__class__.__name__}({parts_str})"
 
-    @classmethod
-    def _validate_class_for_ckpt_loading(cls):
-        # Make sure that the `__init__` method takes a single argument, `hparams`.
-        if (init_fn := getattr(cls, "__init__", None)) is None:
-            return
+    @property
+    @override
+    def hparams(self) -> THparams:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return cast(THparams, super().hparams)
 
-        if not inspect.isfunction(init_fn):
-            raise TypeError(f"__init__ must be a function: {init_fn}")
+    @property
+    @override
+    def hparams_initial(self):  # pyright: ignore[reportIncompatibleMethodOverride]
+        hparams = cast(THparams, super().hparams_initial)
+        hparams_dict = {"model": hparams.model_dump(mode="json")}
+        if (trainer := self._trainer) is not None:
+            from ..trainer import Trainer
 
-        parameters = dict(inspect.signature(init_fn).parameters)
-        # Remove the "self" parameter.
-        _ = parameters.pop("self", None)
-        if len(parameters) != 1:
-            raise TypeError(
-                f"__init__ must take a single argument, got {len(parameters)}: {init_fn}"
-            )
+            if isinstance(trainer, Trainer):
+                hparams_dict["trainer"] = trainer.hparams.model_dump(mode="json")
 
-        if "hparams" not in parameters:
-            raise TypeError(
-                f"__init__'s argument must be named 'hparams', got {parameters}"
-            )
+        return cast(Never, hparams_dict)
 
-    hparams: THparams  # pyright: ignore[reportIncompatibleMethodOverride]
-    hparams_initial: THparams  # pyright: ignore[reportIncompatibleMethodOverride]
+    @property
+    @deprecated("Use `hparams` instead")
+    def config(self):
+        return cast(Never, self.hparams)
 
     @classmethod
     @abstractmethod
-    def config_cls(cls) -> type[THparams]: ...
-
-    @classmethod
-    def load_checkpoint(
-        cls,
-        checkpoint_path: _PATH | IO,
-        hparams: THparams | MutableMapping[str, Any] | None = None,
-        map_location: _MAP_LOCATION_TYPE = None,
-        strict: bool = True,
-    ) -> Self:
-        if strict:
-            cls._validate_class_for_ckpt_loading()
-
-        kwargs: dict[str, Any] = {}
-        if hparams is not None:
-            kwargs["hparams"] = hparams
-
-        return super().load_from_checkpoint(
-            checkpoint_path,
-            map_location=map_location,
-            hparams_file=None,
-            strict=strict,
-            **kwargs,
-        )
-
-    def pre_init_update_hparams_dict(self, hparams: MutableMapping[str, Any]):
-        """
-        Override this method to update the hparams dictionary before it is used to create the hparams object.
-        Mapping-based parameters are passed to the constructor of the hparams object when we're loading the model from a checkpoint.
-        """
-        return hparams
-
-    def pre_init_update_hparams(self, hparams: THparams):
-        """
-        Override this method to update the hparams object before it is used to create the hparams_initial object.
-        """
-        return hparams
+    def hparams_cls(cls) -> type[THparams]: ...
 
     @override
-    def __init__(self, hparams: THparams | MutableMapping[str, Any]):
-        if not isinstance(hparams, BaseConfig):
-            if not isinstance(hparams, MutableMapping):
-                raise TypeError(
-                    f"hparams must be a BaseConfig or a MutableMapping: {type(hparams)}"
-                )
-
-            hparams = self.pre_init_update_hparams_dict(hparams)
-            hparams = self.config_cls().model_validate(hparams)
-        hparams.environment = EnvironmentConfig.from_current_environment(hparams, self)
-        hparams = self.pre_init_update_hparams(hparams)
-
+    def __init__(self, hparams: THparams | Mapping[str, Any]):
         super().__init__()
+
+        # Validate and save hyperparameters
+        hparams_cls = self.hparams_cls()
+        if isinstance(hparams, Mapping):
+            hparams = hparams_cls.model_validate(hparams)
+        elif not isinstance(hparams, hparams_cls):
+            raise TypeError(
+                f"Expected hparams to be either a Mapping or an instance of {hparams_cls}, got {type(hparams)}"
+            )
+        hparams = hparams.model_deep_validate()
         self.save_hyperparameters(hparams)
 
     def zero_loss(self):
@@ -267,260 +241,3 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
         loss = sum((0.0 * v).sum() for v in self.parameters() if v.requires_grad)
         loss = cast(torch.Tensor, loss)
         return loss
-
-    if TYPE_CHECKING:
-
-        @override
-        def training_step(  # pyright: ignore[reportIncompatibleMethodOverride]
-            self,
-            batch: Any,
-            batch_idx: int,
-        ) -> Any:
-            r"""Here you compute and return the training loss and some additional metrics for e.g. the progress bar or
-            logger.
-
-            Args:
-                batch: The output of your data iterable, normally a :class:`~torch.utils.data.DataLoader`.
-                batch_idx: The index of this batch.
-                dataloader_idx: The index of the dataloader that produced this batch.
-                    (only if multiple dataloaders used)
-
-            Return:
-                - :class:`~torch.Tensor` - The loss tensor
-                - ``dict`` - A dictionary which can include any keys, but must include the key ``'loss'`` in the case of
-                automatic optimization.
-                - ``None`` - In automatic optimization, this will skip to the next batch (but is not supported for
-                multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning
-                the loss is not required.
-
-            In this step you'd normally do the forward pass and calculate the loss for a batch.
-            You can also do fancier things like multiple forward passes or something model specific.
-
-            Example::
-
-                def training_step(self, batch, batch_idx):
-                    x, y, z = batch
-                    out = self.encoder(x)
-                    loss = self.loss(out, x)
-                    return loss
-
-            To use multiple optimizers, you can switch to 'manual optimization' and control their stepping:
-
-            .. code-block:: python
-
-                def __init__(self):
-                    super().__init__()
-                    self.automatic_optimization = False
-
-
-                # Multiple optimizers (e.g.: GANs)
-                def training_step(self, batch, batch_idx):
-                    opt1, opt2 = self.optimizers()
-
-                    # do training_step with encoder
-                    ...
-                    opt1.step()
-                    # do training_step with decoder
-                    ...
-                    opt2.step()
-
-            Note:
-                When ``accumulate_grad_batches`` > 1, the loss returned here will be automatically
-                normalized by ``accumulate_grad_batches`` internally.
-
-            """
-            raise NotImplementedError
-
-        @override
-        def validation_step(  # pyright: ignore[reportIncompatibleMethodOverride]
-            self,
-            batch: Any,
-            batch_idx: int,
-        ) -> STEP_OUTPUT:
-            r"""Operates on a single batch of data from the validation set. In this step you'd might generate examples or
-            calculate anything of interest like accuracy.
-
-            Args:
-                batch: The output of your data iterable, normally a :class:`~torch.utils.data.DataLoader`.
-                batch_idx: The index of this batch.
-                dataloader_idx: The index of the dataloader that produced this batch.
-                    (only if multiple dataloaders used)
-
-            Return:
-                - :class:`~torch.Tensor` - The loss tensor
-                - ``dict`` - A dictionary. Can include any keys, but must include the key ``'loss'``.
-                - ``None`` - Skip to the next batch.
-
-            .. code-block:: python
-
-                # if you have one val dataloader:
-                def validation_step(self, batch, batch_idx): ...
-
-
-                # if you have multiple val dataloaders:
-                def validation_step(self, batch, batch_idx, dataloader_idx=0): ...
-
-            Examples::
-
-                # CASE 1: A single validation dataset
-                def validation_step(self, batch, batch_idx):
-                    x, y = batch
-
-                    # implement your own
-                    out = self(x)
-                    loss = self.loss(out, y)
-
-                    # log 6 example images
-                    # or generated text... or whatever
-                    sample_imgs = x[:6]
-                    grid = torchvision.utils.make_grid(sample_imgs)
-                    self.logger.experiment.add_image('example_images', grid, 0)
-
-                    # calculate acc
-                    labels_hat = torch.argmax(out, dim=1)
-                    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-
-                    # log the outputs!
-                    self.log_dict({'val_loss': loss, 'val_acc': val_acc})
-
-            If you pass in multiple val dataloaders, :meth:`validation_step` will have an additional argument. We recommend
-            setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.
-
-            .. code-block:: python
-
-                # CASE 2: multiple validation dataloaders
-                def validation_step(self, batch, batch_idx, dataloader_idx=0):
-                    # dataloader_idx tells you which dataset this is.
-                    ...
-
-            Note:
-                If you don't need to validate you don't need to implement this method.
-
-            Note:
-                When the :meth:`validation_step` is called, the model has been put in eval mode
-                and PyTorch gradients have been disabled. At the end of validation,
-                the model goes back to training mode and gradients are enabled.
-
-            """
-            raise NotImplementedError
-
-        @override
-        def test_step(  # pyright: ignore[reportIncompatibleMethodOverride]
-            self,
-            batch: Any,
-            batch_idx: int,
-        ) -> STEP_OUTPUT:
-            r"""Operates on a single batch of data from the test set. In this step you'd normally generate examples or
-            calculate anything of interest such as accuracy.
-
-            Args:
-                batch: The output of your data iterable, normally a :class:`~torch.utils.data.DataLoader`.
-                batch_idx: The index of this batch.
-                dataloader_idx: The index of the dataloader that produced this batch.
-                    (only if multiple dataloaders used)
-
-            Return:
-                - :class:`~torch.Tensor` - The loss tensor
-                - ``dict`` - A dictionary. Can include any keys, but must include the key ``'loss'``.
-                - ``None`` - Skip to the next batch.
-
-            .. code-block:: python
-
-                # if you have one test dataloader:
-                def test_step(self, batch, batch_idx): ...
-
-
-                # if you have multiple test dataloaders:
-                def test_step(self, batch, batch_idx, dataloader_idx=0): ...
-
-            Examples::
-
-                # CASE 1: A single test dataset
-                def test_step(self, batch, batch_idx):
-                    x, y = batch
-
-                    # implement your own
-                    out = self(x)
-                    loss = self.loss(out, y)
-
-                    # log 6 example images
-                    # or generated text... or whatever
-                    sample_imgs = x[:6]
-                    grid = torchvision.utils.make_grid(sample_imgs)
-                    self.logger.experiment.add_image('example_images', grid, 0)
-
-                    # calculate acc
-                    labels_hat = torch.argmax(out, dim=1)
-                    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-
-                    # log the outputs!
-                    self.log_dict({'test_loss': loss, 'test_acc': test_acc})
-
-            If you pass in multiple test dataloaders, :meth:`test_step` will have an additional argument. We recommend
-            setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.
-
-            .. code-block:: python
-
-                # CASE 2: multiple test dataloaders
-                def test_step(self, batch, batch_idx, dataloader_idx=0):
-                    # dataloader_idx tells you which dataset this is.
-                    ...
-
-            Note:
-                If you don't need to test you don't need to implement this method.
-
-            Note:
-                When the :meth:`test_step` is called, the model has been put in eval mode and
-                PyTorch gradients have been disabled. At the end of the test epoch, the model goes back
-                to training mode and gradients are enabled.
-
-            """
-            raise NotImplementedError
-
-    @override
-    def predict_step(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self,
-        batch: Any,
-        batch_idx: int,
-    ) -> STEP_OUTPUT:
-        """Step function called during :meth:`~lightning.pytorch.trainer.trainer.Trainer.predict`. By default, it calls
-        :meth:`~lightning.pytorch.core.LightningModule.forward`. Override to add any processing logic.
-
-        The :meth:`~lightning.pytorch.core.LightningModule.predict_step` is used
-        to scale inference on multi-devices.
-
-        To prevent an OOM error, it is possible to use :class:`~lightning.pytorch.callbacks.BasePredictionWriter`
-        callback to write the predictions to disk or database after each batch or on epoch end.
-
-        The :class:`~lightning.pytorch.callbacks.BasePredictionWriter` should be used while using a spawn
-        based accelerator. This happens for ``Trainer(strategy="ddp_spawn")``
-        or training on 8 TPU cores with ``Trainer(accelerator="tpu", devices=8)`` as predictions won't be returned.
-
-        Args:
-            batch: The output of your data iterable, normally a :class:`~torch.utils.data.DataLoader`.
-            batch_idx: The index of this batch.
-            dataloader_idx: The index of the dataloader that produced this batch.
-                (only if multiple dataloaders used)
-
-        Return:
-            Predicted output (optional).
-
-        Example ::
-
-            class MyModel(LightningModule):
-
-                def predict_step(self, batch, batch_idx, dataloader_idx=0):
-                    return self(batch)
-
-            dm = ...
-            model = MyModel()
-            trainer = Trainer(accelerator="gpu", devices=2)
-            predictions = trainer.predict(model, dm)
-
-        """
-        prediction = self(batch)
-        return {
-            "prediction": prediction,
-            "batch": batch,
-            "batch_idx": batch_idx,
-        }

nshtrainer/model/mixins/callback.py

@@ -2,16 +2,18 @@ from __future__ import annotations
 
 import logging
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, TypeAlias, cast, final
+from typing import Any, TypeAlias, cast
 
 from lightning.pytorch import Callback, LightningModule
 from typing_extensions import override
 
+from ..._callback import NTCallbackBase
 from ...util.typing_utils import mixin_base_type
 
 log = logging.getLogger(__name__)
 
-CallbackFn: TypeAlias = Callable[[], Callback | Iterable[Callback] | None]
+_Callback = Callback | NTCallbackBase
+CallbackFn: TypeAlias = Callable[[], _Callback | Iterable[_Callback] | None]
 
 
 class CallbackRegistrarModuleMixin:
@@ -23,7 +25,7 @@ class CallbackRegistrarModuleMixin:
 
     def register_callback(
         self,
-        callback: Callback | Iterable[Callback] | CallbackFn | None = None,
+        callback: _Callback | Iterable[_Callback] | CallbackFn | None = None,
     ):
         if not callable(callback):
             callback_ = cast(CallbackFn, lambda: callback)
@@ -34,8 +36,26 @@ class CallbackRegistrarModuleMixin:
 
 
 class CallbackModuleMixin(
-    CallbackRegistrarModuleMixin, mixin_base_type(LightningModule)
+    CallbackRegistrarModuleMixin,
+    mixin_base_type(LightningModule),
 ):
+    @property
+    def _nshtrainer_callbacks(self) -> list[CallbackFn]:
+        if not hasattr(self, "_private_nshtrainer_callbacks_list"):
+            self._private_nshtrainer_callbacks_list = []
+        return self._private_nshtrainer_callbacks_list
+
+    def register_callback(
+        self,
+        callback: _Callback | Iterable[_Callback] | CallbackFn | None = None,
+    ):
+        if not callable(callback):
+            callback_ = cast(CallbackFn, lambda: callback)
+        else:
+            callback_ = callback
+
+        self._nshtrainer_callbacks.append(callback_)
+
     def _gather_all_callbacks(self):
         modules: list[Any] = []
         if isinstance(self, CallbackRegistrarModuleMixin):
@@ -52,7 +72,6 @@ class CallbackModuleMixin(
         for module in modules:
             yield from module._nshtrainer_callbacks
 
-    @final
     @override
     def configure_callbacks(self):
         callbacks = super().configure_callbacks()