nshtrainer 0.31.0__tar.gz → 0.32.0__tar.gz

{nshtrainer-0.31.0 → nshtrainer-0.32.0}/PKG-INFO

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

{nshtrainer-0.31.0 → nshtrainer-0.32.0}/pyproject.toml

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

{nshtrainer-0.31.0 → nshtrainer-0.32.0}/src/nshtrainer/callbacks/__init__.py

@@ -12,6 +12,8 @@ from .checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
 from .checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
+from .debug_flag import DebugFlagCallback as DebugFlagCallback
+from .debug_flag import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from .directory_setup import DirectorySetupCallback as DirectorySetupCallback
 from .directory_setup import DirectorySetupConfig as DirectorySetupConfig
 from .early_stopping import EarlyStopping as EarlyStopping
@@ -41,7 +43,8 @@ from .wandb_watch import WandbWatchCallback as WandbWatchCallback
 from .wandb_watch import WandbWatchConfig as WandbWatchConfig
 
 CallbackConfig = Annotated[
-    EarlyStoppingConfig
+    DebugFlagCallbackConfig
+    | EarlyStoppingConfig
     | ThroughputMonitorConfig
     | EpochTimerConfig
     | PrintTableMetricsConfig

nshtrainer-0.32.0/src/nshtrainer/callbacks/debug_flag.py

@@ -0,0 +1,72 @@
+import logging
+from typing import TYPE_CHECKING, Literal, cast
+
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+from typing_extensions import override
+
+from nshtrainer.model.config import BaseConfig
+
+from .base import CallbackConfigBase
+
+if TYPE_CHECKING:
+    from ..model.config import BaseConfig
+
+log = logging.getLogger(__name__)
+
+
+class DebugFlagCallbackConfig(CallbackConfigBase):
+    name: Literal["debug_flag"] = "debug_flag"
+
+    enabled: bool = True
+    """Whether to enable the callback."""
+
+    def __bool__(self):
+        return self.enabled
+
+    @override
+    def create_callbacks(self, root_config):
+        if not self:
+            return
+
+        yield DebugFlagCallback(self)
+
+
+class DebugFlagCallback(Callback):
+    """
+    Sets the debug flag to true in the following circumstances:
+    - fast_dev_run is enabled
+    - sanity check is running
+    """
+
+    @override
+    def __init__(self, config: DebugFlagCallbackConfig):
+        super().__init__()
+
+        self.config = config
+        del config
+
+    @override
+    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str):
+        if not getattr(trainer, "fast_dev_run", False):
+            return
+
+        hparams = cast("BaseConfig", pl_module.hparams)
+        if not hparams.debug:
+            log.critical("Fast dev run detected, setting debug flag to True.")
+        hparams.debug = True
+
+    @override
+    def on_sanity_check_start(self, trainer: Trainer, pl_module: LightningModule):
+        hparams = cast("BaseConfig", pl_module.hparams)
+        self._debug = hparams.debug
+        if not self._debug:
+            log.critical("Enabling debug flag during sanity check routine.")
+        hparams.debug = True
+
+    @override
+    def on_sanity_check_end(self, trainer: Trainer, pl_module: LightningModule):
+        hparams = cast("BaseConfig", pl_module.hparams)
+        if not self._debug:
+            log.critical("Sanity check routine complete, disabling debug flag.")
+        hparams.debug = self._debug

{nshtrainer-0.31.0 → nshtrainer-0.32.0}/src/nshtrainer/model/base.py

@@ -7,8 +7,7 @@ from typing import IO, TYPE_CHECKING, Any, Generic, Literal, cast
 import torch
 import torch.distributed
 from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
-from lightning.pytorch import LightningModule, Trainer
-from lightning.pytorch.callbacks import Callback
+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
@@ -16,7 +15,6 @@ from typing_extensions import Self, TypeVar, 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.logger import LoggerLightningModuleMixin
 
 log = logging.getLogger(__name__)
@@ -24,39 +22,6 @@ log = logging.getLogger(__name__)
 THparams = TypeVar("THparams", bound=BaseConfig, infer_variance=True)
 
 
-class DebugFlagCallback(Callback):
-    """
-    Sets the debug flag to true in the following circumstances:
-    - fast_dev_run is enabled
-    - sanity check is running
-    """
-
-    @override
-    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str):
-        if not getattr(trainer, "fast_dev_run", False):
-            return
-
-        hparams = cast(BaseConfig, pl_module.hparams)
-        if not hparams.debug:
-            log.critical("Fast dev run detected, setting debug flag to True.")
-        hparams.debug = True
-
-    @override
-    def on_sanity_check_start(self, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast(BaseConfig, pl_module.hparams)
-        self._debug = hparams.debug
-        if not self._debug:
-            log.critical("Enabling debug flag during sanity check routine.")
-        hparams.debug = True
-
-    @override
-    def on_sanity_check_end(self, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast(BaseConfig, pl_module.hparams)
-        if not self._debug:
-            log.critical("Sanity check routine complete, disabling debug flag.")
-        hparams.debug = self._debug
-
-
 T = TypeVar("T", infer_variance=True)
 
 ReduceOpStr = Literal[
@@ -88,7 +53,6 @@ VALID_REDUCE_OPS = (
 class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
     _RLPSanityCheckModuleMixin,
     LoggerLightningModuleMixin,
-    CallbackModuleMixin,
     LightningModule,
     ABC,
     Generic[THparams],
@@ -288,10 +252,8 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
         hparams.environment = EnvironmentConfig.from_current_environment(hparams, self)
         hparams = self.pre_init_update_hparams(hparams)
 
-        super().__init__(hparams)
-
+        super().__init__()
         self.save_hyperparameters(hparams)
-        self.register_callback(lambda: DebugFlagCallback())
 
     def zero_loss(self):
         """

{nshtrainer-0.31.0 → nshtrainer-0.32.0}/src/nshtrainer/trainer/_config.py

@@ -35,6 +35,7 @@ from ..callbacks import (
     OnExceptionCheckpointCallbackConfig,
 )
 from ..callbacks.base import CallbackConfigBase
+from ..callbacks.debug_flag import DebugFlagCallbackConfig
 from ..callbacks.rlp_sanity_checks import RLPSanityChecksConfig
 from ..callbacks.shared_parameters import SharedParametersConfig
 from ..loggers import (
@@ -751,6 +752,11 @@ class TrainerConfig(C.Config):
     """If enabled, will automatically set the default root dir to [cwd/lightning_logs/<id>/]. There is basically no reason to disable this."""
     save_checkpoint_metadata: bool = True
     """If enabled, will save additional metadata whenever a checkpoint is saved."""
+    auto_set_debug_flag: DebugFlagCallbackConfig | None = DebugFlagCallbackConfig()
+    """If enabled, will automatically set the debug flag to True if:
+    - The trainer is running in fast_dev_run mode.
+    - The trainer is running a sanity check (which happens before starting the training routine).
+    """
 
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()
     """
@@ -775,4 +781,7 @@ class TrainerConfig(C.Config):
         yield self.logging
         yield self.optimizer
         yield self.hf_hub
+        yield self.shared_parameters
+        yield self.reduce_lr_on_plateau_sanity_checking
+        yield self.auto_set_debug_flag
         yield from self.callbacks

nshtrainer-0.31.0/src/nshtrainer/model/mixins/callback.py

@@ -1,206 +0,0 @@
-import logging
-from collections.abc import Callable, Iterable, Sequence
-from typing import Any, TypeAlias, cast, final, overload
-
-from lightning.pytorch import Callback, LightningModule
-from lightning.pytorch.callbacks import LambdaCallback
-from typing_extensions import override
-
-from ...util.typing_utils import mixin_base_type
-
-log = logging.getLogger(__name__)
-
-CallbackFn: TypeAlias = Callable[[], Callback | Iterable[Callback] | None]
-
-
-class CallbackRegistrarModuleMixin:
-    @override
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self._nshtrainer_callbacks: list[CallbackFn] = []
-
-    @overload
-    def register_callback(
-        self, callback: Callback | Iterable[Callback] | CallbackFn | None = None, /
-    ): ...
-
-    @overload
-    def register_callback(
-        self,
-        /,
-        *,
-        setup: Callable | None = None,
-        teardown: Callable | None = None,
-        on_fit_start: Callable | None = None,
-        on_fit_end: Callable | None = None,
-        on_sanity_check_start: Callable | None = None,
-        on_sanity_check_end: Callable | None = None,
-        on_train_batch_start: Callable | None = None,
-        on_train_batch_end: Callable | None = None,
-        on_train_epoch_start: Callable | None = None,
-        on_train_epoch_end: Callable | None = None,
-        on_validation_epoch_start: Callable | None = None,
-        on_validation_epoch_end: Callable | None = None,
-        on_test_epoch_start: Callable | None = None,
-        on_test_epoch_end: Callable | None = None,
-        on_validation_batch_start: Callable | None = None,
-        on_validation_batch_end: Callable | None = None,
-        on_test_batch_start: Callable | None = None,
-        on_test_batch_end: Callable | None = None,
-        on_train_start: Callable | None = None,
-        on_train_end: Callable | None = None,
-        on_validation_start: Callable | None = None,
-        on_validation_end: Callable | None = None,
-        on_test_start: Callable | None = None,
-        on_test_end: Callable | None = None,
-        on_exception: Callable | None = None,
-        on_save_checkpoint: Callable | None = None,
-        on_load_checkpoint: Callable | None = None,
-        on_before_backward: Callable | None = None,
-        on_after_backward: Callable | None = None,
-        on_before_optimizer_step: Callable | None = None,
-        on_before_zero_grad: Callable | None = None,
-        on_predict_start: Callable | None = None,
-        on_predict_end: Callable | None = None,
-        on_predict_batch_start: Callable | None = None,
-        on_predict_batch_end: Callable | None = None,
-        on_predict_epoch_start: Callable | None = None,
-        on_predict_epoch_end: Callable | None = None,
-    ): ...
-
-    def register_callback(
-        self,
-        callback: Callback | Iterable[Callback] | CallbackFn | None = None,
-        /,
-        *,
-        setup: Callable | None = None,
-        teardown: Callable | None = None,
-        on_fit_start: Callable | None = None,
-        on_fit_end: Callable | None = None,
-        on_sanity_check_start: Callable | None = None,
-        on_sanity_check_end: Callable | None = None,
-        on_train_batch_start: Callable | None = None,
-        on_train_batch_end: Callable | None = None,
-        on_train_epoch_start: Callable | None = None,
-        on_train_epoch_end: Callable | None = None,
-        on_validation_epoch_start: Callable | None = None,
-        on_validation_epoch_end: Callable | None = None,
-        on_test_epoch_start: Callable | None = None,
-        on_test_epoch_end: Callable | None = None,
-        on_validation_batch_start: Callable | None = None,
-        on_validation_batch_end: Callable | None = None,
-        on_test_batch_start: Callable | None = None,
-        on_test_batch_end: Callable | None = None,
-        on_train_start: Callable | None = None,
-        on_train_end: Callable | None = None,
-        on_validation_start: Callable | None = None,
-        on_validation_end: Callable | None = None,
-        on_test_start: Callable | None = None,
-        on_test_end: Callable | None = None,
-        on_exception: Callable | None = None,
-        on_save_checkpoint: Callable | None = None,
-        on_load_checkpoint: Callable | None = None,
-        on_before_backward: Callable | None = None,
-        on_after_backward: Callable | None = None,
-        on_before_optimizer_step: Callable | None = None,
-        on_before_zero_grad: Callable | None = None,
-        on_predict_start: Callable | None = None,
-        on_predict_end: Callable | None = None,
-        on_predict_batch_start: Callable | None = None,
-        on_predict_batch_end: Callable | None = None,
-        on_predict_epoch_start: Callable | None = None,
-        on_predict_epoch_end: Callable | None = None,
-    ):
-        if callback is None:
-            callback = LambdaCallback(
-                setup=setup,
-                teardown=teardown,
-                on_fit_start=on_fit_start,
-                on_fit_end=on_fit_end,
-                on_sanity_check_start=on_sanity_check_start,
-                on_sanity_check_end=on_sanity_check_end,
-                on_train_batch_start=on_train_batch_start,
-                on_train_batch_end=on_train_batch_end,
-                on_train_epoch_start=on_train_epoch_start,
-                on_train_epoch_end=on_train_epoch_end,
-                on_validation_epoch_start=on_validation_epoch_start,
-                on_validation_epoch_end=on_validation_epoch_end,
-                on_test_epoch_start=on_test_epoch_start,
-                on_test_epoch_end=on_test_epoch_end,
-                on_validation_batch_start=on_validation_batch_start,
-                on_validation_batch_end=on_validation_batch_end,
-                on_test_batch_start=on_test_batch_start,
-                on_test_batch_end=on_test_batch_end,
-                on_train_start=on_train_start,
-                on_train_end=on_train_end,
-                on_validation_start=on_validation_start,
-                on_validation_end=on_validation_end,
-                on_test_start=on_test_start,
-                on_test_end=on_test_end,
-                on_exception=on_exception,
-                on_save_checkpoint=on_save_checkpoint,
-                on_load_checkpoint=on_load_checkpoint,
-                on_before_backward=on_before_backward,
-                on_after_backward=on_after_backward,
-                on_before_optimizer_step=on_before_optimizer_step,
-                on_before_zero_grad=on_before_zero_grad,
-                on_predict_start=on_predict_start,
-                on_predict_end=on_predict_end,
-                on_predict_batch_start=on_predict_batch_start,
-                on_predict_batch_end=on_predict_batch_end,
-                on_predict_epoch_start=on_predict_epoch_start,
-                on_predict_epoch_end=on_predict_epoch_end,
-            )
-
-        if not callable(callback):
-            callback_ = cast(CallbackFn, lambda: callback)
-        else:
-            callback_ = callback
-
-        self._nshtrainer_callbacks.append(callback_)
-
-
-class CallbackModuleMixin(
-    CallbackRegistrarModuleMixin,
-    mixin_base_type(LightningModule),
-):
-    def _nshtrainer_gather_all_callbacks(self):
-        modules: list[Any] = []
-        if isinstance(self, CallbackRegistrarModuleMixin):
-            modules.append(self)
-        if (
-            datamodule := getattr(self.trainer, "datamodule", None)
-        ) is not None and isinstance(datamodule, CallbackRegistrarModuleMixin):
-            modules.append(datamodule)
-        modules.extend(
-            module
-            for module in self.children()
-            if isinstance(module, CallbackRegistrarModuleMixin)
-        )
-        for module in modules:
-            yield from module._nshtrainer_callbacks
-
-    @final
-    @override
-    def configure_callbacks(self):
-        callbacks = super().configure_callbacks()
-        if not isinstance(callbacks, Sequence):
-            callbacks = [callbacks]
-
-        callbacks = list(callbacks)
-        for callback_fn in self._nshtrainer_gather_all_callbacks():
-            callback_result = callback_fn()
-            if callback_result is None:
-                continue
-
-            if not isinstance(callback_result, Iterable):
-                callback_result = [callback_result]
-
-            for callback in callback_result:
-                log.info(
-                    f"Registering {callback.__class__.__qualname__} callback {callback}"
-                )
-                callbacks.append(callback)
-
-        return callbacks