nshtrainer 0.30.1__py3-none-any.whl → 0.32.0__py3-none-any.whl

nshtrainer/trainer/trainer.py

@@ -18,10 +18,8 @@ from typing_extensions import Unpack, assert_never, override
 
 from .._checkpoint.metadata import _write_checkpoint_metadata
 from ..callbacks.base import resolve_all_callbacks
-from ..model.config import (
+from ._config import (
     AcceleratorConfigProtocol,
-    BaseConfig,
-    BaseProfilerConfig,
     LightningTrainerKwargs,
     StrategyConfigProtocol,
 )
@@ -29,6 +27,9 @@ from ._runtime_callback import RuntimeTrackerCallback, Stage
 from .checkpoint_connector import _CheckpointConnector
 from .signal_connector import _SignalConnector
 
+if TYPE_CHECKING:
+    from ..model.config import BaseConfig
+
 log = logging.getLogger(__name__)
 
 
@@ -58,14 +59,14 @@ def _is_bf16_supported_no_emulation():
 
 class Trainer(LightningTrainer):
     @classmethod
-    def _pre_init(cls, config: BaseConfig):
+    def _pre_init(cls, config: "BaseConfig"):
         if (precision := config.trainer.set_float32_matmul_precision) is not None:
             torch.set_float32_matmul_precision(precision)
 
     @classmethod
     def _update_kwargs(
         cls,
-        config: BaseConfig,
+        config: "BaseConfig",
         kwargs_ctor: LightningTrainerKwargs,
     ):
         kwargs: LightningTrainerKwargs = {
@@ -217,18 +218,16 @@ class Trainer(LightningTrainer):
                 gradient_clip_val=grad_clip_config.value,
             )
 
-        if profiler := config.trainer.profiler:
-            # If the profiler is an ProfilerConfig instance, then we instantiate it.
-            if isinstance(profiler, BaseProfilerConfig):
-                profiler = profiler.create_profiler(config)
-                # Make sure that the profiler is an instance of `Profiler`.
-                if not isinstance(profiler, Profiler):
-                    raise ValueError(f"{profiler=} is not an instance of `{Profiler}`.")
-
+        if profiler_config := config.trainer.profiler:
+            if (profiler := profiler_config.create_profiler(config)) is None:
+                log.warning(f"Profiler config {profiler_config=} returned None.")
+            # Make sure that the profiler is an instance of `Profiler`.
+            elif not isinstance(profiler, Profiler):
+                raise ValueError(f"{profiler=} is not an instance of `{Profiler}`.")
             # Otherwise, if the profiler is a string (e.g., "simpe", "advanced", "pytorch"),
             #   then we just pass it through.
-            # kwargs["profiler"] = profiler
-            _update_kwargs(profiler=profiler)
+            else:
+                _update_kwargs(profiler=profiler)
 
         if callbacks := resolve_all_callbacks(config):
             _update_kwargs(callbacks=callbacks)
@@ -281,7 +280,7 @@ class Trainer(LightningTrainer):
     @override
     def __init__(
         self,
-        config: BaseConfig,
+        config: "BaseConfig",
         /,
         **kwargs: Unpack[LightningTrainerKwargs],
     ):
@@ -424,7 +423,7 @@ class Trainer(LightningTrainer):
         # Save the checkpoint metadata
         metadata_path = None
         lm = self._base_module
-        root_config = cast(BaseConfig, lm.hparams)
+        root_config = cast("BaseConfig", lm.hparams)
         if root_config.trainer.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
             if (

nshtrainer/{config → util/config}/__init__.py

@@ -1,3 +1,4 @@
+from . import duration as duration
 from .duration import Duration as Duration
 from .duration import Epochs as Epochs
 from .duration import Steps as Steps

{nshtrainer-0.30.1.dist-info → nshtrainer-0.32.0.dist-info}/METADATA

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

{nshtrainer-0.30.1.dist-info → nshtrainer-0.32.0.dist-info}/RECORD

@@ -1,11 +1,12 @@
-nshtrainer/__init__.py,sha256=sUb2yNdkHHhrKWCeWA5QKIA1Xx3jkO1QGD5Pa-HvgbA,614
+nshtrainer/__init__.py,sha256=flMI50Hj1Ie8c1YMSUQ759AqtNBQLT_zHaV2J9EUmOs,573
 nshtrainer/_callback.py,sha256=A1zLsTy4b_wOYnInLLXGSRdHzT2yNa6mPEql-ozm0u0,1013
 nshtrainer/_checkpoint/loader.py,sha256=5vjg-OFChXJjgiOVv8vnV8nwTscfdDtEdxQRz6uPfDE,14158
 nshtrainer/_checkpoint/metadata.py,sha256=5D4PgKodzhLsmQvuF3xxkH49epKaegxi4wh_ImDTtns,4737
 nshtrainer/_checkpoint/saver.py,sha256=MbX_WjkDtHHAf9Ms-KXDlknkjiPXVoGIe2ciO28AdZ0,1264
+nshtrainer/_directory.py,sha256=RjnW6vKTeKlz2vQWT3cG0Jje5BkFXA7HpUubDhcSiq4,2993
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
 nshtrainer/_hf_hub.py,sha256=0bkXkqhve5D1onMW-fCfuvVKlTn0i6jv_6uMNgZ7OHQ,12974
-nshtrainer/callbacks/__init__.py,sha256=4qocBDzQbLLhhbIEfvbA3SQB_Dy9ZJH7keMwPay-ZS8,2359
+nshtrainer/callbacks/__init__.py,sha256=1SBLpMsx7BzgimO35MwQViYBcbgxlkyvTMz1JKUKK-0,3060
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=NpjeKmonJ1Kaz5_39XSn3LlDwvbGjk6WV8BpHSNCvI4,3508
@@ -14,6 +15,8 @@ nshtrainer/callbacks/checkpoint/_base.py,sha256=vvlwuD-20NozYVIolGGShmUdkkNYeuwN
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=8BHgLAd3Tuzf5sup0guEAKF1jJiAwYsjdKBFYZw98ac,2171
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CWWv0cSwQ1VAX26N7hAyMxbNCk26Keh39oQguBEK5To,1102
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
+nshtrainer/callbacks/debug_flag.py,sha256=Mo69CtJqPWMlFBvgBEuYls8Vfp5v1QFiyMRTiMStdec,2059
+nshtrainer/callbacks/directory_setup.py,sha256=c0uY0oTqLcQ3egInHO7G6BeQQgk_xvOLoHH8FR-9U0U,2629
 nshtrainer/callbacks/early_stopping.py,sha256=VWuJz0oN87b6SwBeVc32YNpeJr1wts8K45k8JJJmG9I,4617
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
 nshtrainer/callbacks/finite_checks.py,sha256=gJC_RUr3ais3FJI0uB6wUZnDdE3WRwCix3ppA3PwQXA,2077
@@ -22,15 +25,16 @@ nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvL
 nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
 nshtrainer/callbacks/norm_logging.py,sha256=T2psu8mYsw9iahPKT6aUPjkGrZ4TIzm6_UUUmE09GJs,6274
 nshtrainer/callbacks/print_table.py,sha256=_FdAHhqylWGk4Z0c2FrLFeiMA4jhfA_beZRK_BHpzmE,2837
+nshtrainer/callbacks/rlp_sanity_checks.py,sha256=c30G9jAu42QLLIS5LnusdSnI3wqyIHgOUFDRcKESuNI,9935
+nshtrainer/callbacks/shared_parameters.py,sha256=fqlDweFDXPV_bfcAWpRgaJIad9i5AehYDtuJjDtUum4,2922
 nshtrainer/callbacks/throughput_monitor.py,sha256=H_ocXErZxUO3dxFk8Tx_VQdpI9E_Ztvqof5WtFevLyQ,1838
 nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
 nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
-nshtrainer/config/__init__.py,sha256=v9RtlM1Pqj_4fCDfskgxEtiGtbWH3Tj7lqNsKCDQ4gk,119
-nshtrainer/config/duration.py,sha256=f_obz0eorkktI3HzAuIawABDkvuL4lDqCxcPb3UW7Q4,692
+nshtrainer/config.py,sha256=W6nAmn5Y1GVZto9vkx4v8i5XdikMSdVYDiq7kbDEWAg,5900
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=dGBTDDtlBU6c-ZlVQOCnTW7SjTB5hczWsOWEdUWjvkA,4385
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
-nshtrainer/ll/__init__.py,sha256=6UTt2apSD8tOZw3M7hyd-33v4RKSpNNATlWFbW4cNnU,2523
+nshtrainer/ll/__init__.py,sha256=L-aTi1V1bbvnZjOro8NvI393zbHQSFR9movWSRK9Mds,2477
 nshtrainer/ll/_experimental.py,sha256=oBQCKOEVYoxuUU9eLb-Fg2B2mzZD7SA0zfAO6lmWZ88,53
 nshtrainer/ll/actsave.py,sha256=2lbiseSrjcwFT6AiyLNWarTWl1bnzliVWlu1iOfnP30,209
 nshtrainer/ll/callbacks.py,sha256=AxyUmc8aGRSjx6WwwgXYCmdJ73rwLuEAEH0AGRosojQ,49
@@ -38,7 +42,7 @@ nshtrainer/ll/config.py,sha256=fKumJf42HY2FITX1QUM1OTXkYD6U2np2ciyd4PFRPZ8,145
 nshtrainer/ll/data.py,sha256=zRG0FRje-jtSHximVzkHIHzpwsyQxpHCoACFihNKLPM,44
 nshtrainer/ll/log.py,sha256=d4BB3TyM8imK65EXOiOeUTF0zFM1ropbe7Vq3DeB0xU,140
 nshtrainer/ll/lr_scheduler.py,sha256=7xjhN6L69BCUzFhcy33NtMtPuCzHiB611zVWFg92lQ0,52
-nshtrainer/ll/model.py,sha256=cxFQfFc-2mAYBGwDpP8m5tjQBs7M47cZ6JoPXksPaoI,473
+nshtrainer/ll/model.py,sha256=Cw8Vq8IUL6YU1fTUcOIZsXcNJ3XyKgQY4YENIsL9H7c,996
 nshtrainer/ll/nn.py,sha256=8qiRDFwojIxkB7-LtNWk4mLL2tJbaskHYofDsOIHiNg,42
 nshtrainer/ll/optimizer.py,sha256=3T-VZtT73jVvwCNJGDjgGEbzs-1LFTzMQH-SB_58mSo,49
 nshtrainer/ll/runner.py,sha256=B0m5VEhNKIjF1aFmqPkonkQxDoRL2jeHZGsV3zwhSVE,117
@@ -51,44 +55,46 @@ nshtrainer/loggers/__init__.py,sha256=C_xk0A3_qKbNdTmzK85AgjRHFD3w-jPRS2ig-iPhfE
 nshtrainer/loggers/_base.py,sha256=xiZKEK0ALJkcqf4OpVNRY0QbZsamR_WR7x7m_68YHXQ,705
 nshtrainer/loggers/csv.py,sha256=D_lYyd94bZ8jAgnRo-ARtFgVcInaD9zktxtsUD9RWCI,1052
 nshtrainer/loggers/tensorboard.py,sha256=wL2amRSdP68zbslZvBeM0ZQBnjF3hIKsz-_lBbdomaM,2216
-nshtrainer/loggers/wandb.py,sha256=FPwbf618AYmuPzHdhd1ZFhJ8qDjwTUiSe7cm7g3KCyM,5112
+nshtrainer/loggers/wandb.py,sha256=8B2BMMzILRSUEiCkmp_fBpcXs69euRKViTiaV__DJZk,5128
 nshtrainer/lr_scheduler/__init__.py,sha256=uEvgaFAs-4s_bAEMaildy0GT6OvgpgOEKTuzqutESHE,736
 nshtrainer/lr_scheduler/_base.py,sha256=7xOIuxQ86YHbFWG5a3gX46emQj1WN_LaY4-i0Q1TDBg,3659
-nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=Fyontbfu4k2932xZenE63QL4CrVGWANXdTeq63dUko0,5347
+nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=YQm84Sb4SWrofpBwa39DCslJvu2uorjbpWaGWyys1l4,5352
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-WKArwxxf7ZQizhmbCw,2782
 nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
 nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ,1288
-nshtrainer/model/__init__.py,sha256=VyRziPT3YilP6xjLi_StsSqtlvn7N4LOMzgukRsOnF8,1380
-nshtrainer/model/base.py,sha256=oQVolDk81acy4OlckwQEBHuX2gCaVSYiIA0JaDIfhQ4,17517
-nshtrainer/model/config.py,sha256=zcCLcqvg4u7Zg6SLtCnqdIfiW8I0eART47lf1LCYl-A,43326
-nshtrainer/model/modules/callback.py,sha256=1z6gUDBd35KG3phGzRekgZM6SIk-wj5Uo6APN4YhRR0,8549
-nshtrainer/model/modules/debug.py,sha256=Yy7XEdPou9BkCsD5hJchwJGmCVGrfUru5g9VjPM4uAw,1120
-nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
-nshtrainer/model/modules/logger.py,sha256=CJWSmNT8SV5GLtfml-qGYenqRPXcNOMsJRGEavAd8Hw,5464
-nshtrainer/model/modules/profiler.py,sha256=rQ_jRMcM1Z2AIROZlRnBRHM5rkTpq67afZPD6CIRfXs,825
-nshtrainer/model/modules/rlp_sanity_checks.py,sha256=I_ralr2ThQ-D_FkVQTwbdXLLlgHJEr7-s01I5wSDjps,8893
-nshtrainer/model/modules/shared_parameters.py,sha256=ZiRKkZXr6RwdwLCdZCJPl3dXe7bnT8Z9yTeRK5bXBGk,2687
+nshtrainer/model/__init__.py,sha256=2i_VEy6u_Y1LUGKljHXWeekvhnUcanZM2QyaaBM1Bmw,261
+nshtrainer/model/base.py,sha256=1zVY8ybZTzVKhpp7sUC0t360Ut3YmdGxAW5PZAIBSyw,18535
+nshtrainer/model/config.py,sha256=Q4Wong6w3cp_Sq7s8iZdABKF-LZBbSCFn_TQPYkhkrI,6572
+nshtrainer/model/mixins/logger.py,sha256=xOymSTofukEYZGkGojXsMEO__ZlBI5lIPZVmlotMEX8,5291
 nshtrainer/nn/__init__.py,sha256=0QPFl02a71WZQjLMGOlFNMmsYP5aa1q3eABHmnWH58Q,1427
 nshtrainer/nn/mlp.py,sha256=V0FrScpIUdg_IgIO8GMtIsGEtmHjwF14i2IWxmZrsqg,5952
 nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,2329
 nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
 nshtrainer/nn/nonlinearity.py,sha256=4sYE4MN5zojc-go1k0PYtqssVRuXrM7D4tbpIXp5K-E,6078
 nshtrainer/optimizer.py,sha256=kuJEA1pvB3y1FcsfhAoOJujVqEZqFHlmYO8GW6JeA1g,1527
+nshtrainer/profiler/__init__.py,sha256=RQYkqQBVWuVvfdtAJIk2x5bNsXownklT87Mr_j-uXjw,474
+nshtrainer/profiler/_base.py,sha256=YF5lsJBIl9qts9GLW5Z62JuYeo4SnIArhlFwTGkfTb4,897
+nshtrainer/profiler/advanced.py,sha256=44asloha0aGUW8YwjQt3lm3ve8H-N6mM4QgseUSLT30,1112
+nshtrainer/profiler/pytorch.py,sha256=tGeRvoPP5ulWX2RkfXrQvMBoki1T95dpz5p8mwyon1I,2709
+nshtrainer/profiler/simple.py,sha256=MbMfsJvligd0mtGiltxJ0T8MQVDP9T9BzQZFwswl66Y,957
 nshtrainer/runner.py,sha256=USAjrExHkN5oVNVunsoPnLxfQrEHSaa54S3RipOe544,3605
 nshtrainer/scripts/find_packages.py,sha256=ixYivZobumyyGsf2B9oYMLyLTRcBzY_vUv-u3bNW-hs,1424
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
+nshtrainer/trainer/_config.py,sha256=ZIodM5Ek1lpkWFhQ_VfmKR7q1mZFFwtjfx8FH72H8WM,29174
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=r0ir4xYSdf_jebM0x09qaO6nJsvsiRQDyM0fs80ppOQ,2347
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
-nshtrainer/trainer/trainer.py,sha256=L4nYXq6Gts2sS9CQGenwEcvMET4L5vO5c60KM5Hm8Do,17544
+nshtrainer/trainer/trainer.py,sha256=iYueHW-m8fHyC8SQuXmpgxq_-GUa7pAJik7rDFPXmy0,17499
 nshtrainer/util/_environment_info.py,sha256=CFUUZYjXhBLWGc0jtPNOaZgYMueUDEHpEaWFA1f3GoY,24213
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
+nshtrainer/util/config/__init__.py,sha256=6iCFLhujhbOi7Q694e--Sq-ethiGoGHShm699GPV8Zg,154
+nshtrainer/util/config/duration.py,sha256=f_obz0eorkktI3HzAuIawABDkvuL4lDqCxcPb3UW7Q4,692
 nshtrainer/util/environment.py,sha256=AeW_kLl-N70wmb6L_JLz1wRj0kA70xs6RCmc9iUqczE,4159
 nshtrainer/util/path.py,sha256=VkpuhR4GaZtSFBVqbGAvfjcrU-PR8xwiGzzwFNOWP9c,2995
 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.30.1.dist-info/METADATA,sha256=LV0wQlmotpfC3qO76dFVCbS26bEl-9YMiTetEeqVQsU,916
-nshtrainer-0.30.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.30.1.dist-info/RECORD,,
+nshtrainer-0.32.0.dist-info/METADATA,sha256=pe-TVRS0ZmZ9kx5NBQ8-0C6m4ZzaH_MalJZmh31mUNQ,916
+nshtrainer-0.32.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.32.0.dist-info/RECORD,,

nshtrainer/model/modules/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 _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():
-            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

nshtrainer/model/modules/debug.py

@@ -1,42 +0,0 @@
-import logging
-
-import torch
-import torch.distributed
-
-log = logging.getLogger(__name__)
-
-
-class DebugModuleMixin:
-    @torch.jit.unused
-    def breakpoint(self, rank_zero_only: bool = True):
-        if (
-            not rank_zero_only
-            or not torch.distributed.is_initialized()
-            or torch.distributed.get_rank() == 0
-        ):
-            breakpoint()
-
-        if rank_zero_only and torch.distributed.is_initialized():
-            _ = torch.distributed.barrier()
-
-    @torch.jit.unused
-    def ensure_finite(
-        self,
-        tensor: torch.Tensor,
-        name: str | None = None,
-        throw: bool = False,
-    ):
-        name_parts: list[str] = ["Tensor"]
-        if name is not None:
-            name_parts.append(name)
-        name = " ".join(name_parts)
-
-        not_finite = ~torch.isfinite(tensor)
-        if not_finite.any():
-            msg = f"{name} has {not_finite.sum().item()}/{not_finite.numel()} non-finite values."
-            if throw:
-                raise RuntimeError(msg)
-            else:
-                log.warning(msg)
-            return False
-        return True

nshtrainer/model/modules/distributed.py

@@ -1,70 +0,0 @@
-from typing import Any, Literal, cast
-
-import torch.distributed
-from lightning.pytorch import LightningModule
-from torch.distributed import ReduceOp
-from typing_extensions import TypeVar
-
-from ...util.typing_utils import mixin_base_type
-
-T = TypeVar("T", infer_variance=True)
-
-ReduceOpStr = Literal[
-    "avg",
-    "mean",
-    "band",
-    "bor",
-    "bxor",
-    "max",
-    "min",
-    "premul_sum",
-    "product",
-    "sum",
-]
-VALID_REDUCE_OPS = (
-    "avg",
-    "mean",
-    "band",
-    "bor",
-    "bxor",
-    "max",
-    "min",
-    "premul_sum",
-    "product",
-    "sum",
-)
-
-
-class DistributedMixin(mixin_base_type(LightningModule)):
-    def all_gather_object(
-        self,
-        object: T,
-        group: torch.distributed.ProcessGroup | None = None,
-    ) -> list[T]:
-        if (
-            not torch.distributed.is_available()
-            or not torch.distributed.is_initialized()
-        ):
-            return [object]
-
-        object_list = [cast(T, None) for _ in range(self.trainer.world_size)]
-        torch.distributed.all_gather_object(object_list, object, group=group)
-        return object_list
-
-    def barrier(self, name: str | None = None):
-        self.trainer.strategy.barrier(name=name)
-
-    def reduce(
-        self,
-        tensor: torch.Tensor,
-        reduce_op: ReduceOp.RedOpType | ReduceOpStr,
-        group: Any | None = None,
-    ) -> torch.Tensor:
-        if isinstance(reduce_op, str):
-            # validate reduce_op
-            if reduce_op not in VALID_REDUCE_OPS:
-                raise ValueError(
-                    f"reduce_op must be one of {VALID_REDUCE_OPS}, got {reduce_op}"
-                )
-
-        return self.trainer.strategy.reduce(tensor, group=group, reduce_op=reduce_op)

nshtrainer/model/modules/profiler.py

@@ -1,24 +0,0 @@
-from lightning.pytorch import LightningDataModule, LightningModule
-from lightning.pytorch.profilers import PassThroughProfiler
-
-from ...util.typing_utils import mixin_base_type
-
-
-class ProfilerMixin(mixin_base_type(LightningModule)):
-    @property
-    def profiler(self):
-        if not isinstance(self, (LightningModule, LightningDataModule)):
-            raise TypeError(
-                "`profiler` can only be used on LightningModule or LightningDataModule"
-            )
-
-        if (trainer := self.trainer) is None:
-            raise RuntimeError("trainer is not defined")
-
-        if not hasattr(trainer, "profiler"):
-            raise RuntimeError("trainer does not have profiler")
-
-        if (profiler := getattr(trainer, "profiler")) is None:
-            profiler = PassThroughProfiler()
-
-        return profiler