nshtrainer 1.0.0b33__py3-none-any.whl → 1.0.0b37__py3-none-any.whl

nshtrainer/configs/callbacks/lr_monitor/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.lr_monitor import CallbackConfigBase as CallbackConfig
 from nshtrainer.callbacks.lr_monitor import (
     LearningRateMonitorConfig as LearningRateMonitorConfig,
 )
+from nshtrainer.callbacks.lr_monitor import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "LearningRateMonitorConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/norm_logging/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.norm_logging import CallbackConfigBase as CallbackConf
 from nshtrainer.callbacks.norm_logging import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
 )
+from nshtrainer.callbacks.norm_logging import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "NormLoggingCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/print_table/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.print_table import CallbackConfigBase as CallbackConfi
 from nshtrainer.callbacks.print_table import (
     PrintTableMetricsCallbackConfig as PrintTableMetricsCallbackConfig,
 )
+from nshtrainer.callbacks.print_table import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "PrintTableMetricsCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/rlp_sanity_checks/__init__.py

@@ -8,8 +8,12 @@ from nshtrainer.callbacks.rlp_sanity_checks import (
 from nshtrainer.callbacks.rlp_sanity_checks import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
+from nshtrainer.callbacks.rlp_sanity_checks import (
+    callback_registry as callback_registry,
+)
 
 __all__ = [
     "CallbackConfigBase",
     "RLPSanityChecksCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/shared_parameters/__init__.py

@@ -8,8 +8,12 @@ from nshtrainer.callbacks.shared_parameters import (
 from nshtrainer.callbacks.shared_parameters import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+from nshtrainer.callbacks.shared_parameters import (
+    callback_registry as callback_registry,
+)
 
 __all__ = [
     "CallbackConfigBase",
     "SharedParametersCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/timer/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.timer import CallbackConfigBase as CallbackConfigBase
 from nshtrainer.callbacks.timer import (
     EpochTimerCallbackConfig as EpochTimerCallbackConfig,
 )
+from nshtrainer.callbacks.timer import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "EpochTimerCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/wandb_upload_code/__init__.py

@@ -8,8 +8,12 @@ from nshtrainer.callbacks.wandb_upload_code import (
 from nshtrainer.callbacks.wandb_upload_code import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
 )
+from nshtrainer.callbacks.wandb_upload_code import (
+    callback_registry as callback_registry,
+)
 
 __all__ = [
     "CallbackConfigBase",
     "WandbUploadCodeCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/callbacks/wandb_watch/__init__.py

@@ -6,8 +6,10 @@ from nshtrainer.callbacks.wandb_watch import CallbackConfigBase as CallbackConfi
 from nshtrainer.callbacks.wandb_watch import (
     WandbWatchCallbackConfig as WandbWatchCallbackConfig,
 )
+from nshtrainer.callbacks.wandb_watch import callback_registry as callback_registry
 
 __all__ = [
     "CallbackConfigBase",
     "WandbWatchCallbackConfig",
+    "callback_registry",
 ]

nshtrainer/configs/loggers/__init__.py

@@ -3,11 +3,12 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.loggers import ActSaveLoggerConfig as ActSaveLoggerConfig
-from nshtrainer.loggers import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers import CSVLoggerConfig as CSVLoggerConfig
 from nshtrainer.loggers import LoggerConfig as LoggerConfig
+from nshtrainer.loggers import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.loggers import TensorboardLoggerConfig as TensorboardLoggerConfig
 from nshtrainer.loggers import WandbLoggerConfig as WandbLoggerConfig
+from nshtrainer.loggers import logger_registry as logger_registry
 from nshtrainer.loggers.wandb import CallbackConfigBase as CallbackConfigBase
 from nshtrainer.loggers.wandb import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
@@ -16,25 +17,26 @@ from nshtrainer.loggers.wandb import (
     WandbWatchCallbackConfig as WandbWatchCallbackConfig,
 )
 
-from . import _base as _base
 from . import actsave as actsave
+from . import base as base
 from . import csv as csv
 from . import tensorboard as tensorboard
 from . import wandb as wandb
 
 __all__ = [
     "ActSaveLoggerConfig",
-    "BaseLoggerConfig",
     "CSVLoggerConfig",
     "CallbackConfigBase",
     "LoggerConfig",
+    "LoggerConfigBase",
     "TensorboardLoggerConfig",
     "WandbLoggerConfig",
     "WandbUploadCodeCallbackConfig",
     "WandbWatchCallbackConfig",
-    "_base",
     "actsave",
+    "base",
     "csv",
+    "logger_registry",
     "tensorboard",
     "wandb",
 ]

nshtrainer/configs/loggers/actsave/__init__.py

@@ -3,9 +3,11 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.loggers.actsave import ActSaveLoggerConfig as ActSaveLoggerConfig
-from nshtrainer.loggers.actsave import BaseLoggerConfig as BaseLoggerConfig
+from nshtrainer.loggers.actsave import LoggerConfigBase as LoggerConfigBase
+from nshtrainer.loggers.actsave import logger_registry as logger_registry
 
 __all__ = [
     "ActSaveLoggerConfig",
-    "BaseLoggerConfig",
+    "LoggerConfigBase",
+    "logger_registry",
 ]

nshtrainer/configs/loggers/base/__init__.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.loggers.base import LoggerConfigBase as LoggerConfigBase
+from nshtrainer.loggers.base import logger_registry as logger_registry
+
+__all__ = [
+    "LoggerConfigBase",
+    "logger_registry",
+]

nshtrainer/configs/loggers/csv/__init__.py

@@ -2,10 +2,12 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.loggers.csv import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers.csv import CSVLoggerConfig as CSVLoggerConfig
+from nshtrainer.loggers.csv import LoggerConfigBase as LoggerConfigBase
+from nshtrainer.loggers.csv import logger_registry as logger_registry
 
 __all__ = [
-    "BaseLoggerConfig",
     "CSVLoggerConfig",
+    "LoggerConfigBase",
+    "logger_registry",
 ]

nshtrainer/configs/loggers/tensorboard/__init__.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.loggers.tensorboard import BaseLoggerConfig as BaseLoggerConfig
+from nshtrainer.loggers.tensorboard import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.loggers.tensorboard import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
+from nshtrainer.loggers.tensorboard import logger_registry as logger_registry
 
 __all__ = [
-    "BaseLoggerConfig",
+    "LoggerConfigBase",
     "TensorboardLoggerConfig",
+    "logger_registry",
 ]

nshtrainer/configs/loggers/wandb/__init__.py

@@ -2,8 +2,8 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.loggers.wandb import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers.wandb import CallbackConfigBase as CallbackConfigBase
+from nshtrainer.loggers.wandb import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.loggers.wandb import WandbLoggerConfig as WandbLoggerConfig
 from nshtrainer.loggers.wandb import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
@@ -11,11 +11,13 @@ from nshtrainer.loggers.wandb import (
 from nshtrainer.loggers.wandb import (
     WandbWatchCallbackConfig as WandbWatchCallbackConfig,
 )
+from nshtrainer.loggers.wandb import logger_registry as logger_registry
 
 __all__ = [
-    "BaseLoggerConfig",
     "CallbackConfigBase",
+    "LoggerConfigBase",
     "WandbLoggerConfig",
     "WandbUploadCodeCallbackConfig",
     "WandbWatchCallbackConfig",
+    "logger_registry",
 ]

nshtrainer/configs/lr_scheduler/__init__.py

@@ -8,12 +8,13 @@ from nshtrainer.lr_scheduler import (
 from nshtrainer.lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.lr_scheduler import LRSchedulerConfigBase as LRSchedulerConfigBase
 from nshtrainer.lr_scheduler import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
+from nshtrainer.lr_scheduler.base import lr_scheduler_registry as lr_scheduler_registry
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     DurationConfig as DurationConfig,
 )
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import MetricConfig as MetricConfig
 
-from . import _base as _base
+from . import base as base
 from . import linear_warmup_cosine as linear_warmup_cosine
 from . import reduce_lr_on_plateau as reduce_lr_on_plateau
 
@@ -24,7 +25,8 @@ __all__ = [
     "LinearWarmupCosineDecayLRSchedulerConfig",
     "MetricConfig",
     "ReduceLROnPlateauConfig",
-    "_base",
+    "base",
     "linear_warmup_cosine",
+    "lr_scheduler_registry",
     "reduce_lr_on_plateau",
 ]

nshtrainer/configs/lr_scheduler/base/__init__.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.lr_scheduler.base import LRSchedulerConfigBase as LRSchedulerConfigBase
+from nshtrainer.lr_scheduler.base import lr_scheduler_registry as lr_scheduler_registry
+
+__all__ = [
+    "LRSchedulerConfigBase",
+    "lr_scheduler_registry",
+]

nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py

@@ -11,9 +11,13 @@ from nshtrainer.lr_scheduler.linear_warmup_cosine import (
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     LRSchedulerConfigBase as LRSchedulerConfigBase,
 )
+from nshtrainer.lr_scheduler.linear_warmup_cosine import (
+    lr_scheduler_registry as lr_scheduler_registry,
+)
 
 __all__ = [
     "DurationConfig",
     "LRSchedulerConfigBase",
     "LinearWarmupCosineDecayLRSchedulerConfig",
+    "lr_scheduler_registry",
 ]

nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py

@@ -9,9 +9,13 @@ from nshtrainer.lr_scheduler.reduce_lr_on_plateau import MetricConfig as MetricC
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
     ReduceLROnPlateauConfig as ReduceLROnPlateauConfig,
 )
+from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
+    lr_scheduler_registry as lr_scheduler_registry,
+)
 
 __all__ = [
     "LRSchedulerConfigBase",
     "MetricConfig",
     "ReduceLROnPlateauConfig",
+    "lr_scheduler_registry",
 ]

nshtrainer/configs/nn/__init__.py

@@ -2,13 +2,13 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.nn import BaseNonlinearityConfig as BaseNonlinearityConfig
 from nshtrainer.nn import ELUNonlinearityConfig as ELUNonlinearityConfig
 from nshtrainer.nn import GELUNonlinearityConfig as GELUNonlinearityConfig
 from nshtrainer.nn import LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig
 from nshtrainer.nn import MishNonlinearityConfig as MishNonlinearityConfig
 from nshtrainer.nn import MLPConfig as MLPConfig
 from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
+from nshtrainer.nn import NonlinearityConfigBase as NonlinearityConfigBase
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
@@ -21,18 +21,19 @@ from nshtrainer.nn import TanhNonlinearityConfig as TanhNonlinearityConfig
 from nshtrainer.nn.nonlinearity import (
     SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig,
 )
+from nshtrainer.nn.nonlinearity import nonlinearity_registry as nonlinearity_registry
 
 from . import mlp as mlp
 from . import nonlinearity as nonlinearity
 
 __all__ = [
-    "BaseNonlinearityConfig",
     "ELUNonlinearityConfig",
     "GELUNonlinearityConfig",
     "LeakyReLUNonlinearityConfig",
     "MLPConfig",
     "MishNonlinearityConfig",
     "NonlinearityConfig",
+    "NonlinearityConfigBase",
     "PReLUConfig",
     "ReLUNonlinearityConfig",
     "SiLUNonlinearityConfig",
@@ -45,4 +46,5 @@ __all__ = [
     "TanhNonlinearityConfig",
     "mlp",
     "nonlinearity",
+    "nonlinearity_registry",
 ]

nshtrainer/configs/nn/mlp/__init__.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.nn.mlp import BaseNonlinearityConfig as BaseNonlinearityConfig
 from nshtrainer.nn.mlp import MLPConfig as MLPConfig
 from nshtrainer.nn.mlp import NonlinearityConfig as NonlinearityConfig
+from nshtrainer.nn.mlp import NonlinearityConfigBase as NonlinearityConfigBase
 
 __all__ = [
-    "BaseNonlinearityConfig",
     "MLPConfig",
     "NonlinearityConfig",
+    "NonlinearityConfigBase",
 ]

nshtrainer/configs/nn/nonlinearity/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.nn.nonlinearity import BaseNonlinearityConfig as BaseNonlinearityConfig
 from nshtrainer.nn.nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig
 from nshtrainer.nn.nonlinearity import GELUNonlinearityConfig as GELUNonlinearityConfig
 from nshtrainer.nn.nonlinearity import (
@@ -10,6 +9,7 @@ from nshtrainer.nn.nonlinearity import (
 )
 from nshtrainer.nn.nonlinearity import MishNonlinearityConfig as MishNonlinearityConfig
 from nshtrainer.nn.nonlinearity import NonlinearityConfig as NonlinearityConfig
+from nshtrainer.nn.nonlinearity import NonlinearityConfigBase as NonlinearityConfigBase
 from nshtrainer.nn.nonlinearity import PReLUConfig as PReLUConfig
 from nshtrainer.nn.nonlinearity import ReLUNonlinearityConfig as ReLUNonlinearityConfig
 from nshtrainer.nn.nonlinearity import (
@@ -32,14 +32,15 @@ from nshtrainer.nn.nonlinearity import (
     SwishNonlinearityConfig as SwishNonlinearityConfig,
 )
 from nshtrainer.nn.nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig
+from nshtrainer.nn.nonlinearity import nonlinearity_registry as nonlinearity_registry
 
 __all__ = [
-    "BaseNonlinearityConfig",
     "ELUNonlinearityConfig",
     "GELUNonlinearityConfig",
     "LeakyReLUNonlinearityConfig",
     "MishNonlinearityConfig",
     "NonlinearityConfig",
+    "NonlinearityConfigBase",
     "PReLUConfig",
     "ReLUNonlinearityConfig",
     "SiLUNonlinearityConfig",
@@ -50,4 +51,5 @@ __all__ = [
     "SwiGLUNonlinearityConfig",
     "SwishNonlinearityConfig",
     "TanhNonlinearityConfig",
+    "nonlinearity_registry",
 ]

nshtrainer/configs/optimizer/__init__.py

@@ -5,9 +5,11 @@ __codegen__ = True
 from nshtrainer.optimizer import AdamWConfig as AdamWConfig
 from nshtrainer.optimizer import OptimizerConfig as OptimizerConfig
 from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
+from nshtrainer.optimizer import optimizer_registry as optimizer_registry
 
 __all__ = [
     "AdamWConfig",
     "OptimizerConfig",
     "OptimizerConfigBase",
+    "optimizer_registry",
 ]

nshtrainer/configs/trainer/__init__.py

@@ -4,10 +4,10 @@ __codegen__ = True
 
 from nshtrainer.trainer import TrainerConfig as TrainerConfig
 from nshtrainer.trainer import accelerator_registry as accelerator_registry
+from nshtrainer.trainer import callback_registry as callback_registry
 from nshtrainer.trainer import plugin_registry as plugin_registry
 from nshtrainer.trainer._config import AcceleratorConfig as AcceleratorConfig
 from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
-from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.trainer._config import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
@@ -36,6 +36,7 @@ from nshtrainer.trainer._config import (
 )
 from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbackConfig
 from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
+from nshtrainer.trainer._config import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
@@ -55,9 +56,6 @@ from nshtrainer.trainer._config import StrategyConfig as StrategyConfig
 from nshtrainer.trainer._config import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
-from nshtrainer.trainer._config import (
-    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
-)
 from nshtrainer.trainer._config import WandbLoggerConfig as WandbLoggerConfig
 from nshtrainer.trainer.accelerator import CPUAcceleratorConfig as CPUAcceleratorConfig
 from nshtrainer.trainer.accelerator import (
@@ -135,7 +133,6 @@ __all__ = [
     "AcceleratorConfigBase",
     "ActSaveLoggerConfig",
     "AsyncCheckpointIOPlugin",
-    "BaseLoggerConfig",
     "BestCheckpointCallbackConfig",
     "BitsandbytesPluginConfig",
     "CPUAcceleratorConfig",
@@ -163,6 +160,7 @@ __all__ = [
     "LightningEnvironmentPlugin",
     "LogEpochCallbackConfig",
     "LoggerConfig",
+    "LoggerConfigBase",
     "MPIEnvironmentPlugin",
     "MPSAcceleratorConfig",
     "MetricConfig",
@@ -179,7 +177,6 @@ __all__ = [
     "StrategyConfig",
     "StrategyConfigBase",
     "TensorboardLoggerConfig",
-    "TimeCheckpointCallbackConfig",
     "TorchCheckpointIOPlugin",
     "TorchElasticEnvironmentPlugin",
     "TorchSyncBatchNormPlugin",
@@ -193,6 +190,7 @@ __all__ = [
     "_config",
     "accelerator",
     "accelerator_registry",
+    "callback_registry",
     "plugin",
     "plugin_registry",
     "strategy",

nshtrainer/configs/trainer/_config/__init__.py

@@ -4,7 +4,6 @@ __codegen__ = True
 
 from nshtrainer.trainer._config import AcceleratorConfig as AcceleratorConfig
 from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
-from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.trainer._config import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
@@ -33,6 +32,7 @@ from nshtrainer.trainer._config import (
 )
 from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbackConfig
 from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
+from nshtrainer.trainer._config import LoggerConfigBase as LoggerConfigBase
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
@@ -53,18 +53,12 @@ from nshtrainer.trainer._config import StrategyConfig as StrategyConfig
 from nshtrainer.trainer._config import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
-from nshtrainer.trainer._config import (
-    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
-)
 from nshtrainer.trainer._config import TrainerConfig as TrainerConfig
 from nshtrainer.trainer._config import WandbLoggerConfig as WandbLoggerConfig
-from nshtrainer.trainer._config import accelerator_registry as accelerator_registry
-from nshtrainer.trainer._config import plugin_registry as plugin_registry
 
 __all__ = [
     "AcceleratorConfig",
     "ActSaveLoggerConfig",
-    "BaseLoggerConfig",
     "BestCheckpointCallbackConfig",
     "CSVLoggerConfig",
     "CallbackConfig",
@@ -81,6 +75,7 @@ __all__ = [
     "LearningRateMonitorConfig",
     "LogEpochCallbackConfig",
     "LoggerConfig",
+    "LoggerConfigBase",
     "MetricConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
@@ -91,9 +86,6 @@ __all__ = [
     "SharedParametersCallbackConfig",
     "StrategyConfig",
     "TensorboardLoggerConfig",
-    "TimeCheckpointCallbackConfig",
     "TrainerConfig",
     "WandbLoggerConfig",
-    "accelerator_registry",
-    "plugin_registry",
 ]

nshtrainer/loggers/__init__.py

@@ -5,19 +5,14 @@ from typing import Annotated
 import nshconfig as C
 from typing_extensions import TypeAliasType
 
-from ._base import BaseLoggerConfig as BaseLoggerConfig
 from .actsave import ActSaveLoggerConfig as ActSaveLoggerConfig
+from .base import LoggerConfigBase as LoggerConfigBase
+from .base import logger_registry as logger_registry
 from .csv import CSVLoggerConfig as CSVLoggerConfig
 from .tensorboard import TensorboardLoggerConfig as TensorboardLoggerConfig
 from .wandb import WandbLoggerConfig as WandbLoggerConfig
 
 LoggerConfig = TypeAliasType(
     "LoggerConfig",
-    Annotated[
-        CSVLoggerConfig
-        | TensorboardLoggerConfig
-        | WandbLoggerConfig
-        | ActSaveLoggerConfig,
-        C.Field(discriminator="name"),
-    ],
+    Annotated[LoggerConfigBase, logger_registry.DynamicResolution()],
 )

nshtrainer/loggers/actsave.py

@@ -5,11 +5,14 @@ from typing import Any, Literal
 
 import numpy as np
 from lightning.pytorch.loggers import Logger
+from typing_extensions import final
 
-from ._base import BaseLoggerConfig
+from .base import LoggerConfigBase, logger_registry
 
 
-class ActSaveLoggerConfig(BaseLoggerConfig):
+@final
+@logger_registry.register
+class ActSaveLoggerConfig(LoggerConfigBase):
     name: Literal["actsave"] = "actsave"
 
     def create_logger(self, trainer_config):

nshtrainer/loggers/{_base.py → base.py}

@@ -10,7 +10,7 @@ if TYPE_CHECKING:
     from ..trainer._config import TrainerConfig
 
 
-class BaseLoggerConfig(C.Config, ABC):
+class LoggerConfigBase(C.Config, ABC):
     enabled: bool = True
     """Enable this logger."""
 
@@ -29,3 +29,6 @@ class BaseLoggerConfig(C.Config, ABC):
 
     def __bool__(self):
         return self.enabled
+
+
+logger_registry = C.Registry(LoggerConfigBase, discriminator="name")

nshtrainer/loggers/csv.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 from typing import Literal
 
-from typing_extensions import override
+from typing_extensions import final, override
 
-from ._base import BaseLoggerConfig
+from .base import LoggerConfigBase, logger_registry
 
 
-class CSVLoggerConfig(BaseLoggerConfig):
+@final
+@logger_registry.register
+class CSVLoggerConfig(LoggerConfigBase):
     name: Literal["csv"] = "csv"
 
     enabled: bool = True

nshtrainer/loggers/tensorboard.py

@@ -4,9 +4,9 @@ import logging
 from typing import Literal
 
 import nshconfig as C
-from typing_extensions import override
+from typing_extensions import final, override
 
-from ._base import BaseLoggerConfig
+from .base import LoggerConfigBase, logger_registry
 
 log = logging.getLogger(__name__)
 
@@ -30,7 +30,9 @@ def _tensorboard_available():
         return False
 
 
-class TensorboardLoggerConfig(BaseLoggerConfig):
+@final
+@logger_registry.register
+class TensorboardLoggerConfig(LoggerConfigBase):
     name: Literal["tensorboard"] = "tensorboard"
 
     enabled: bool = C.Field(default_factory=lambda: _tensorboard_available())

nshtrainer/loggers/wandb.py

@@ -7,12 +7,12 @@ from typing import TYPE_CHECKING, Literal
 import nshconfig as C
 from lightning.pytorch import Callback, LightningModule, Trainer
 from packaging import version
-from typing_extensions import assert_never, override
+from typing_extensions import assert_never, final, override
 
 from ..callbacks.base import CallbackConfigBase
 from ..callbacks.wandb_upload_code import WandbUploadCodeCallbackConfig
 from ..callbacks.wandb_watch import WandbWatchCallbackConfig
-from ._base import BaseLoggerConfig
+from .base import LoggerConfigBase, logger_registry
 
 if TYPE_CHECKING:
     from ..trainer._config import TrainerConfig
@@ -73,7 +73,9 @@ class FinishWandbOnTeardownCallback(Callback):
         wandb.finish()
 
 
-class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
+@final
+@logger_registry.register
+class WandbLoggerConfig(CallbackConfigBase, LoggerConfigBase):
     name: Literal["wandb"] = "wandb"
 
     enabled: bool = C.Field(default_factory=lambda: _wandb_available())

nshtrainer/lr_scheduler/__init__.py

@@ -5,8 +5,8 @@ from typing import Annotated
 import nshconfig as C
 from typing_extensions import TypeAliasType
 
-from ._base import LRSchedulerConfigBase as LRSchedulerConfigBase
-from ._base import LRSchedulerMetadata as LRSchedulerMetadata
+from .base import LRSchedulerConfigBase as LRSchedulerConfigBase
+from .base import LRSchedulerMetadata as LRSchedulerMetadata
 from .linear_warmup_cosine import (
     LinearWarmupCosineAnnealingLR as LinearWarmupCosineAnnealingLR,
 )