nshtrainer 0.34.1__tar.gz → 0.35.0__tar.gz

{nshtrainer-0.34.1 → nshtrainer-0.35.0}/PKG-INFO

@@ -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 → nshtrainer-0.35.0}/pyproject.toml

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

{nshtrainer-0.34.1 → nshtrainer-0.35.0}/src/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-0.34.1 → nshtrainer-0.35.0}/src/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-0.35.0/src/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