nshtrainer 0.34.2__py3-none-any.whl → 0.35.1__py3-none-any.whl

nshtrainer/model/base.py

@@ -15,6 +15,7 @@ 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__)
@@ -53,6 +54,7 @@ VALID_REDUCE_OPS = (
 class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
     _RLPSanityCheckModuleMixin,
     LoggerLightningModuleMixin,
+    CallbackModuleMixin,
     LightningModule,
     ABC,
     Generic[THparams],

nshtrainer/model/mixins/callback.py

@@ -0,0 +1,74 @@
+import logging
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any, TypeAlias, cast, final
+
+from lightning.pytorch import Callback, LightningModule
+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:
+    @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_)
+
+
+class CallbackModuleMixin(
+    CallbackRegistrarModuleMixin, mixin_base_type(LightningModule)
+):
+    def _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._gather_all_callbacks():
+            if (callback_result := callback_fn()) 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

{nshtrainer-0.34.2.dist-info → nshtrainer-0.35.1.dist-info}/METADATA

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

{nshtrainer-0.34.2.dist-info → nshtrainer-0.35.1.dist-info}/RECORD

@@ -63,8 +63,9 @@ nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-
 nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
 nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ,1288
 nshtrainer/model/__init__.py,sha256=2i_VEy6u_Y1LUGKljHXWeekvhnUcanZM2QyaaBM1Bmw,261
-nshtrainer/model/base.py,sha256=1zVY8ybZTzVKhpp7sUC0t360Ut3YmdGxAW5PZAIBSyw,18535
+nshtrainer/model/base.py,sha256=NasbYZJBuEly6Hm9t9HVZk-CUHmy4T7p1v-Ye981XA4,18609
 nshtrainer/model/config.py,sha256=Q4Wong6w3cp_Sq7s8iZdABKF-LZBbSCFn_TQPYkhkrI,6572
+nshtrainer/model/mixins/callback.py,sha256=rbe8P22iEjPkH1df6rfEo3Txw7EwSz6Dkm0TWO_AysM,2419
 nshtrainer/model/mixins/logger.py,sha256=xOymSTofukEYZGkGojXsMEO__ZlBI5lIPZVmlotMEX8,5291
 nshtrainer/nn/__init__.py,sha256=0QPFl02a71WZQjLMGOlFNMmsYP5aa1q3eABHmnWH58Q,1427
 nshtrainer/nn/mlp.py,sha256=V0FrScpIUdg_IgIO8GMtIsGEtmHjwF14i2IWxmZrsqg,5952
@@ -97,6 +98,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.34.2.dist-info/METADATA,sha256=DQyYTUO0wpboH1gy3nSRJV6EsWCpY7Kb_ldD8v4BQFY,916
-nshtrainer-0.34.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.34.2.dist-info/RECORD,,
+nshtrainer-0.35.1.dist-info/METADATA,sha256=LJUDrvicSUUgu2IfOlKoySGMoIrFMce0ryo4Fvrzwbs,916
+nshtrainer-0.35.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.35.1.dist-info/RECORD,,