nshtrainer 1.0.0b24__py3-none-any.whl → 1.0.0b26__py3-none-any.whl

nshtrainer/.nshconfig.generated.json

@@ -0,0 +1,6 @@
+{
+    "module": "nshtrainer",
+    "output": "configs",
+    "typed_dicts": null,
+    "json_schemas": null
+}

nshtrainer/_checkpoint/metadata.py

@@ -55,7 +55,7 @@ class CheckpointMetadata(C.Config):
     metrics: dict[str, Any]
     environment: EnvironmentConfig
 
-    hparams: Any
+    hparams: Any | None
 
     @classmethod
     def from_file(cls, path: Path):

nshtrainer/callbacks/__init__.py

@@ -14,6 +14,8 @@ from .checkpoint import OnExceptionCheckpointCallback as OnExceptionCheckpointCa
 from .checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
+from .checkpoint import TimeCheckpointCallback as TimeCheckpointCallback
+from .checkpoint import TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig
 from .debug_flag import DebugFlagCallback as DebugFlagCallback
 from .debug_flag import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from .directory_setup import DirectorySetupCallback as DirectorySetupCallback
@@ -71,6 +73,7 @@ CallbackConfig = Annotated[
     | BestCheckpointCallbackConfig
     | LastCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig
+    | TimeCheckpointCallbackConfig
     | SharedParametersCallbackConfig
     | RLPSanityChecksCallbackConfig
     | WandbWatchCallbackConfig

nshtrainer/callbacks/checkpoint/__init__.py

@@ -14,3 +14,7 @@ from .on_exception_checkpoint import (
 from .on_exception_checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
+from .time_checkpoint import TimeCheckpointCallback as TimeCheckpointCallback
+from .time_checkpoint import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)

nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -7,8 +7,7 @@ from typing import Literal
 from lightning.pytorch import LightningModule, Trainer
 from typing_extensions import final, override
 
-from nshtrainer._checkpoint.metadata import CheckpointMetadata
-
+from ..._checkpoint.metadata import CheckpointMetadata
 from ...metrics._config import MetricConfig
 from ._base import BaseCheckpointCallbackConfig, CheckpointBase
 

nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -6,8 +6,7 @@ from typing import Literal
 from lightning.pytorch import LightningModule, Trainer
 from typing_extensions import final, override
 
-from nshtrainer._checkpoint.metadata import CheckpointMetadata
-
+from ..._checkpoint.metadata import CheckpointMetadata
 from ._base import BaseCheckpointCallbackConfig, CheckpointBase
 
 log = logging.getLogger(__name__)

nshtrainer/callbacks/checkpoint/time_checkpoint.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import logging
+import time
+from datetime import timedelta
+from pathlib import Path
+from typing import Any, Literal
+
+from lightning.pytorch import LightningModule, Trainer
+from typing_extensions import final, override
+
+from nshtrainer._checkpoint.metadata import CheckpointMetadata
+
+from ._base import BaseCheckpointCallbackConfig, CheckpointBase
+
+log = logging.getLogger(__name__)
+
+
+@final
+class TimeCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
+    name: Literal["time_checkpoint"] = "time_checkpoint"
+
+    interval: timedelta = timedelta(hours=12)
+    """Time interval between checkpoints."""
+
+    @override
+    def create_checkpoint(self, trainer_config, dirpath):
+        return TimeCheckpointCallback(self, dirpath)
+
+
+@final
+class TimeCheckpointCallback(CheckpointBase[TimeCheckpointCallbackConfig]):
+    def __init__(self, config: TimeCheckpointCallbackConfig, dirpath: Path):
+        super().__init__(config, dirpath)
+        self.start_time = time.time()
+        self.last_checkpoint_time = self.start_time
+        self.interval_seconds = config.interval.total_seconds()
+
+    @override
+    def name(self):
+        return "time"
+
+    @override
+    def default_filename(self):
+        return "epoch{epoch}-step{step}-duration{train_duration}"
+
+    @override
+    def topk_sort_key(self, metadata: CheckpointMetadata):
+        return metadata.checkpoint_timestamp
+
+    @override
+    def topk_sort_reverse(self):
+        return True
+
+    def _should_checkpoint(self) -> bool:
+        current_time = time.time()
+        elapsed_time = current_time - self.last_checkpoint_time
+        return elapsed_time >= self.interval_seconds
+
+    def _format_duration(self, seconds: float) -> str:
+        """Format duration in seconds to a human-readable string."""
+        td = timedelta(seconds=int(seconds))
+        days = td.days
+        hours, remainder = divmod(td.seconds, 3600)
+        minutes, seconds = divmod(remainder, 60)
+
+        parts = []
+        if days > 0:
+            parts.append(f"{days}d")
+        if hours > 0:
+            parts.append(f"{hours}h")
+        if minutes > 0:
+            parts.append(f"{minutes}m")
+        if seconds > 0 or not parts:
+            parts.append(f"{seconds}s")
+
+        return "_".join(parts)
+
+    @override
+    def current_metrics(self, trainer: Trainer) -> dict[str, Any]:
+        metrics = super().current_metrics(trainer)
+        train_duration = time.time() - self.start_time
+        metrics["train_duration"] = self._format_duration(train_duration)
+        return metrics
+
+    @override
+    def on_train_batch_end(
+        self, trainer: Trainer, pl_module: LightningModule, *args, **kwargs
+    ):
+        if self._should_checkpoint():
+            self.save_checkpoints(trainer)
+            self.last_checkpoint_time = time.time()
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        """Save the timer state for checkpoint resumption.
+
+        Returns:
+            Dictionary containing the start time and last checkpoint time.
+        """
+        return {
+            "start_time": self.start_time,
+            "last_checkpoint_time": self.last_checkpoint_time,
+        }
+
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """Restore the timer state when resuming from a checkpoint.
+
+        Args:
+            state_dict: Dictionary containing the previously saved timer state.
+        """
+        self.start_time = state_dict["start_time"]
+        self.last_checkpoint_time = state_dict["last_checkpoint_time"]

nshtrainer/callbacks/print_table.py

@@ -49,14 +49,14 @@ class PrintTableMetricsCallback(Callback):
             }
         self.metrics.append(metrics_dict)
 
-        from rich.console import Console
+        from rich.console import Console  # type: ignore[reportMissingImports] # noqa
 
         console = Console()
         table = self.create_metrics_table()
         console.print(table)
 
     def create_metrics_table(self):
-        from rich.table import Table
+        from rich.table import Table  # type: ignore[reportMissingImports] # noqa
 
         table = Table(show_header=True, header_style="bold magenta")
 

nshtrainer/configs/__init__.py

@@ -14,7 +14,6 @@ from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
 from nshtrainer.callbacks import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
-from nshtrainer.callbacks import CallbackConfig as CallbackConfig
 from nshtrainer.callbacks import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from nshtrainer.callbacks import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
@@ -58,13 +57,11 @@ from nshtrainer.callbacks.checkpoint._base import (
 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 TensorboardLoggerConfig as TensorboardLoggerConfig
 from nshtrainer.loggers import WandbLoggerConfig as WandbLoggerConfig
 from nshtrainer.lr_scheduler import (
     LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig,
 )
-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.nn import BaseNonlinearityConfig as BaseNonlinearityConfig
@@ -73,7 +70,6 @@ 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 PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
@@ -87,23 +83,21 @@ from nshtrainer.nn.nonlinearity import (
     SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig,
 )
 from nshtrainer.optimizer import AdamWConfig as AdamWConfig
-from nshtrainer.optimizer import OptimizerConfig as OptimizerConfig
 from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
 from nshtrainer.profiler import AdvancedProfilerConfig as AdvancedProfilerConfig
 from nshtrainer.profiler import BaseProfilerConfig as BaseProfilerConfig
-from nshtrainer.profiler import ProfilerConfig as ProfilerConfig
 from nshtrainer.profiler import PyTorchProfilerConfig as PyTorchProfilerConfig
 from nshtrainer.profiler import SimpleProfilerConfig as SimpleProfilerConfig
-from nshtrainer.trainer._config import (
-    CheckpointCallbackConfig as CheckpointCallbackConfig,
-)
+from nshtrainer.trainer._config import AcceleratorConfigBase as AcceleratorConfigBase
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
 from nshtrainer.trainer._config import EnvironmentConfig as EnvironmentConfig
 from nshtrainer.trainer._config import GradientClippingConfig as GradientClippingConfig
 from nshtrainer.trainer._config import (
     LearningRateMonitorConfig as LearningRateMonitorConfig,
 )
+from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
+from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.util._environment_info import (
     EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
 )
@@ -133,7 +127,6 @@ from nshtrainer.util._environment_info import (
 )
 from nshtrainer.util._environment_info import GitRepositoryConfig as GitRepositoryConfig
 from nshtrainer.util.config import DTypeConfig as DTypeConfig
-from nshtrainer.util.config import DurationConfig as DurationConfig
 from nshtrainer.util.config import EpochsConfig as EpochsConfig
 from nshtrainer.util.config import StepsConfig as StepsConfig
 
@@ -149,3 +142,95 @@ from . import optimizer as optimizer
 from . import profiler as profiler
 from . import trainer as trainer
 from . import util as util
+
+__all__ = [
+    "AcceleratorConfigBase",
+    "ActSaveConfig",
+    "ActSaveLoggerConfig",
+    "AdamWConfig",
+    "AdvancedProfilerConfig",
+    "BaseCheckpointCallbackConfig",
+    "BaseLoggerConfig",
+    "BaseNonlinearityConfig",
+    "BaseProfilerConfig",
+    "BestCheckpointCallbackConfig",
+    "CSVLoggerConfig",
+    "CallbackConfigBase",
+    "CheckpointMetadata",
+    "CheckpointSavingConfig",
+    "DTypeConfig",
+    "DebugFlagCallbackConfig",
+    "DirectoryConfig",
+    "DirectorySetupCallbackConfig",
+    "ELUNonlinearityConfig",
+    "EMACallbackConfig",
+    "EarlyStoppingCallbackConfig",
+    "EnvironmentCUDAConfig",
+    "EnvironmentClassInformationConfig",
+    "EnvironmentConfig",
+    "EnvironmentGPUConfig",
+    "EnvironmentHardwareConfig",
+    "EnvironmentLSFInformationConfig",
+    "EnvironmentLinuxEnvironmentConfig",
+    "EnvironmentPackageConfig",
+    "EnvironmentSLURMInformationConfig",
+    "EnvironmentSnapshotConfig",
+    "EpochTimerCallbackConfig",
+    "EpochsConfig",
+    "FiniteChecksCallbackConfig",
+    "GELUNonlinearityConfig",
+    "GitRepositoryConfig",
+    "GradientClippingConfig",
+    "GradientSkippingCallbackConfig",
+    "HuggingFaceHubAutoCreateConfig",
+    "HuggingFaceHubConfig",
+    "LRSchedulerConfigBase",
+    "LastCheckpointCallbackConfig",
+    "LeakyReLUNonlinearityConfig",
+    "LearningRateMonitorConfig",
+    "LinearWarmupCosineDecayLRSchedulerConfig",
+    "LogEpochCallbackConfig",
+    "MLPConfig",
+    "MetricConfig",
+    "MishNonlinearityConfig",
+    "NormLoggingCallbackConfig",
+    "OnExceptionCheckpointCallbackConfig",
+    "OptimizerConfigBase",
+    "PReLUConfig",
+    "PluginConfigBase",
+    "PrintTableMetricsCallbackConfig",
+    "PyTorchProfilerConfig",
+    "RLPSanityChecksCallbackConfig",
+    "ReLUNonlinearityConfig",
+    "ReduceLROnPlateauConfig",
+    "SanityCheckingConfig",
+    "SharedParametersCallbackConfig",
+    "SiLUNonlinearityConfig",
+    "SigmoidNonlinearityConfig",
+    "SimpleProfilerConfig",
+    "SoftmaxNonlinearityConfig",
+    "SoftplusNonlinearityConfig",
+    "SoftsignNonlinearityConfig",
+    "StepsConfig",
+    "StrategyConfigBase",
+    "SwiGLUNonlinearityConfig",
+    "SwishNonlinearityConfig",
+    "TanhNonlinearityConfig",
+    "TensorboardLoggerConfig",
+    "TrainerConfig",
+    "WandbLoggerConfig",
+    "WandbUploadCodeCallbackConfig",
+    "WandbWatchCallbackConfig",
+    "_checkpoint",
+    "_directory",
+    "_hf_hub",
+    "callbacks",
+    "loggers",
+    "lr_scheduler",
+    "metrics",
+    "nn",
+    "optimizer",
+    "profiler",
+    "trainer",
+    "util",
+]

nshtrainer/configs/_checkpoint/__init__.py

@@ -6,3 +6,9 @@ from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMeta
 from nshtrainer._checkpoint.metadata import EnvironmentConfig as EnvironmentConfig
 
 from . import metadata as metadata
+
+__all__ = [
+    "CheckpointMetadata",
+    "EnvironmentConfig",
+    "metadata",
+]

nshtrainer/configs/_checkpoint/metadata/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
 from nshtrainer._checkpoint.metadata import EnvironmentConfig as EnvironmentConfig
+
+__all__ = [
+    "CheckpointMetadata",
+    "EnvironmentConfig",
+]

nshtrainer/configs/_directory/__init__.py

@@ -6,4 +6,8 @@ from nshtrainer._directory import DirectoryConfig as DirectoryConfig
 from nshtrainer._directory import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
 )
-from nshtrainer._directory import LoggerConfig as LoggerConfig
+
+__all__ = [
+    "DirectoryConfig",
+    "DirectorySetupCallbackConfig",
+]

nshtrainer/configs/_hf_hub/__init__.py

@@ -7,3 +7,9 @@ from nshtrainer._hf_hub import (
     HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig,
 )
 from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
+
+__all__ = [
+    "CallbackConfigBase",
+    "HuggingFaceHubAutoCreateConfig",
+    "HuggingFaceHubConfig",
+]

nshtrainer/configs/callbacks/__init__.py

@@ -5,7 +5,6 @@ __codegen__ = True
 from nshtrainer.callbacks import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
-from nshtrainer.callbacks import CallbackConfig as CallbackConfig
 from nshtrainer.callbacks import CallbackConfigBase as CallbackConfigBase
 from nshtrainer.callbacks import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from nshtrainer.callbacks import (
@@ -73,3 +72,47 @@ from . import shared_parameters as shared_parameters
 from . import timer as timer
 from . import wandb_upload_code as wandb_upload_code
 from . import wandb_watch as wandb_watch
+
+__all__ = [
+    "ActSaveConfig",
+    "BaseCheckpointCallbackConfig",
+    "BestCheckpointCallbackConfig",
+    "CallbackConfigBase",
+    "CheckpointMetadata",
+    "DebugFlagCallbackConfig",
+    "DirectorySetupCallbackConfig",
+    "EMACallbackConfig",
+    "EarlyStoppingCallbackConfig",
+    "EpochTimerCallbackConfig",
+    "FiniteChecksCallbackConfig",
+    "GradientSkippingCallbackConfig",
+    "LastCheckpointCallbackConfig",
+    "LearningRateMonitorConfig",
+    "LogEpochCallbackConfig",
+    "MetricConfig",
+    "NormLoggingCallbackConfig",
+    "OnExceptionCheckpointCallbackConfig",
+    "PrintTableMetricsCallbackConfig",
+    "RLPSanityChecksCallbackConfig",
+    "SharedParametersCallbackConfig",
+    "WandbUploadCodeCallbackConfig",
+    "WandbWatchCallbackConfig",
+    "actsave",
+    "base",
+    "checkpoint",
+    "debug_flag",
+    "directory_setup",
+    "early_stopping",
+    "ema",
+    "finite_checks",
+    "gradient_skipping",
+    "log_epoch",
+    "lr_monitor",
+    "norm_logging",
+    "print_table",
+    "rlp_sanity_checks",
+    "shared_parameters",
+    "timer",
+    "wandb_upload_code",
+    "wandb_watch",
+]

nshtrainer/configs/callbacks/actsave/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig
 from nshtrainer.callbacks.actsave import CallbackConfigBase as CallbackConfigBase
+
+__all__ = [
+    "ActSaveConfig",
+    "CallbackConfigBase",
+]

nshtrainer/configs/callbacks/base/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.callbacks.base import CallbackConfigBase as CallbackConfigBase
+
+__all__ = [
+    "CallbackConfigBase",
+]

nshtrainer/configs/callbacks/checkpoint/__init__.py

@@ -26,3 +26,17 @@ from . import _base as _base
 from . import best_checkpoint as best_checkpoint
 from . import last_checkpoint as last_checkpoint
 from . import on_exception_checkpoint as on_exception_checkpoint
+
+__all__ = [
+    "BaseCheckpointCallbackConfig",
+    "BestCheckpointCallbackConfig",
+    "CallbackConfigBase",
+    "CheckpointMetadata",
+    "LastCheckpointCallbackConfig",
+    "MetricConfig",
+    "OnExceptionCheckpointCallbackConfig",
+    "_base",
+    "best_checkpoint",
+    "last_checkpoint",
+    "on_exception_checkpoint",
+]

nshtrainer/configs/callbacks/checkpoint/_base/__init__.py

@@ -11,3 +11,9 @@ from nshtrainer.callbacks.checkpoint._base import (
 from nshtrainer.callbacks.checkpoint._base import (
     CheckpointMetadata as CheckpointMetadata,
 )
+
+__all__ = [
+    "BaseCheckpointCallbackConfig",
+    "CallbackConfigBase",
+    "CheckpointMetadata",
+]

nshtrainer/configs/callbacks/checkpoint/best_checkpoint/__init__.py

@@ -12,3 +12,10 @@ from nshtrainer.callbacks.checkpoint.best_checkpoint import (
     CheckpointMetadata as CheckpointMetadata,
 )
 from nshtrainer.callbacks.checkpoint.best_checkpoint import MetricConfig as MetricConfig
+
+__all__ = [
+    "BaseCheckpointCallbackConfig",
+    "BestCheckpointCallbackConfig",
+    "CheckpointMetadata",
+    "MetricConfig",
+]

nshtrainer/configs/callbacks/checkpoint/last_checkpoint/__init__.py

@@ -11,3 +11,9 @@ from nshtrainer.callbacks.checkpoint.last_checkpoint import (
 from nshtrainer.callbacks.checkpoint.last_checkpoint import (
     LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
 )
+
+__all__ = [
+    "BaseCheckpointCallbackConfig",
+    "CheckpointMetadata",
+    "LastCheckpointCallbackConfig",
+]

nshtrainer/configs/callbacks/checkpoint/on_exception_checkpoint/__init__.py

@@ -8,3 +8,8 @@ from nshtrainer.callbacks.checkpoint.on_exception_checkpoint import (
 from nshtrainer.callbacks.checkpoint.on_exception_checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "OnExceptionCheckpointCallbackConfig",
+]

nshtrainer/configs/callbacks/debug_flag/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.debug_flag import CallbackConfigBase as CallbackConfig
 from nshtrainer.callbacks.debug_flag import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "DebugFlagCallbackConfig",
+]

nshtrainer/configs/callbacks/directory_setup/__init__.py

@@ -8,3 +8,8 @@ from nshtrainer.callbacks.directory_setup import (
 from nshtrainer.callbacks.directory_setup import (
     DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "DirectorySetupCallbackConfig",
+]

nshtrainer/configs/callbacks/early_stopping/__init__.py

@@ -7,3 +7,9 @@ from nshtrainer.callbacks.early_stopping import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
 from nshtrainer.callbacks.early_stopping import MetricConfig as MetricConfig
+
+__all__ = [
+    "CallbackConfigBase",
+    "EarlyStoppingCallbackConfig",
+    "MetricConfig",
+]

nshtrainer/configs/callbacks/ema/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.callbacks.ema import CallbackConfigBase as CallbackConfigBase
 from nshtrainer.callbacks.ema import EMACallbackConfig as EMACallbackConfig
+
+__all__ = [
+    "CallbackConfigBase",
+    "EMACallbackConfig",
+]

nshtrainer/configs/callbacks/finite_checks/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.finite_checks import CallbackConfigBase as CallbackCon
 from nshtrainer.callbacks.finite_checks import (
     FiniteChecksCallbackConfig as FiniteChecksCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "FiniteChecksCallbackConfig",
+]

nshtrainer/configs/callbacks/gradient_skipping/__init__.py

@@ -8,3 +8,8 @@ from nshtrainer.callbacks.gradient_skipping import (
 from nshtrainer.callbacks.gradient_skipping import (
     GradientSkippingCallbackConfig as GradientSkippingCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "GradientSkippingCallbackConfig",
+]

nshtrainer/configs/callbacks/log_epoch/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.callbacks.log_epoch import CallbackConfigBase as CallbackConfigB
 from nshtrainer.callbacks.log_epoch import (
     LogEpochCallbackConfig as LogEpochCallbackConfig,
 )
+
+__all__ = [
+    "CallbackConfigBase",
+    "LogEpochCallbackConfig",
+]

nshtrainer/configs/callbacks/lr_monitor/__init__.py

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

nshtrainer/configs/callbacks/norm_logging/__init__.py

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

nshtrainer/configs/callbacks/print_table/__init__.py

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

nshtrainer/configs/callbacks/rlp_sanity_checks/__init__.py

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

nshtrainer/configs/callbacks/shared_parameters/__init__.py

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