nshtrainer 0.34.1__py3-none-any.whl → 0.35.0__py3-none-any.whl

nshtrainer/data/balanced_batch_sampler.py

@@ -81,6 +81,22 @@ class BalancedBatchSampler(BatchSampler):
     ):
         super().__init__(sampler, batch_size, drop_last=drop_last)
 
+        # Validate the dataset
+        dataset = self._unwrap_dataset(self.distributed_sampler.dataset)
+        # Dataset much either implement `data_sizes`, or we need to provide a custom
+        # implementation of the dataset sizes function.
+        if isinstance(dataset, DatasetWithSizes):
+            log.critical(f"BalancedBatchSampler: Resolved dataset to {type(dataset)}")
+
+        elif self._data_sizes_fn is not None:
+            log.critical("BalancedBatchSampler: Using custom data_sizes_fn")
+        else:
+            raise ValueError(
+                "Dataset must implement the `data_sizes` method, "
+                "or a custom data_sizes_fn must be provided "
+                "to the BalancedBatchSampler."
+            )
+
         self._device = device
         self._data_sizes_fn = data_sizes_fn
 
@@ -97,7 +113,6 @@ class BalancedBatchSampler(BatchSampler):
         # Dataset much either implement `data_sizes`, or we need to provide a custom
         # implementation of the dataset sizes function.
         if isinstance(dataset, DatasetWithSizes):
-            log.critical(f"BalancedBatchSampler: Resolved dataset to {type(dataset)}")
             return dataset.data_sizes(indices)
 
         if (data_sizes_fn := self._data_sizes_fn) is not None:

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:
+    @override
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._nshtrainer_callbacks: list[CallbackFn] = []
+
+    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.1.dist-info → nshtrainer-0.35.0.dist-info}/METADATA

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

{nshtrainer-0.34.1.dist-info → nshtrainer-0.35.0.dist-info}/RECORD

@@ -32,7 +32,7 @@ nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50
 nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
 nshtrainer/config.py,sha256=6U7B-kCIMrfEnF_y92RuBm1WfASW7k05Zsm2uHBzRrk,8205
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
-nshtrainer/data/balanced_batch_sampler.py,sha256=WAjhbO9EsZ_UadhdW3obBsjvEDMc2V-irpjegqIb7AI,4791
+nshtrainer/data/balanced_batch_sampler.py,sha256=ybMJF-CguaZ17fLEweZ5suaGOiHOMEm3Bn8rQfGTzGQ,5445
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
 nshtrainer/ll/__init__.py,sha256=L-aTi1V1bbvnZjOro8NvI393zbHQSFR9movWSRK9Mds,2477
 nshtrainer/ll/_experimental.py,sha256=oBQCKOEVYoxuUU9eLb-Fg2B2mzZD7SA0zfAO6lmWZ88,53
@@ -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=lvX9Q2ErETXmGFd79CscSAOJAlTWq-mwMKVC0d0uH1c,2324
 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.1.dist-info/METADATA,sha256=c_iXv-CQLl6kig2u3lmrP4EDSOdMvq8L2WewYiFL-8Q,916
-nshtrainer-0.34.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.34.1.dist-info/RECORD,,
+nshtrainer-0.35.0.dist-info/METADATA,sha256=NBZegh-RUfnkVt_ERUPdH7fdCFZriQZXoMskq_8HB60,916
+nshtrainer-0.35.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.35.0.dist-info/RECORD,,