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

nshtrainer/configs/callbacks/timer/__init__.py

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

nshtrainer/configs/callbacks/wandb_upload_code/__init__.py

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

nshtrainer/configs/callbacks/wandb_watch/__init__.py

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

nshtrainer/configs/loggers/__init__.py

@@ -5,7 +5,6 @@ __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 TensorboardLoggerConfig as TensorboardLoggerConfig
 from nshtrainer.loggers import WandbLoggerConfig as WandbLoggerConfig
 from nshtrainer.loggers.wandb import CallbackConfigBase as CallbackConfigBase
@@ -21,3 +20,19 @@ from . import actsave as actsave
 from . import csv as csv
 from . import tensorboard as tensorboard
 from . import wandb as wandb
+
+__all__ = [
+    "ActSaveLoggerConfig",
+    "BaseLoggerConfig",
+    "CSVLoggerConfig",
+    "CallbackConfigBase",
+    "TensorboardLoggerConfig",
+    "WandbLoggerConfig",
+    "WandbUploadCodeCallbackConfig",
+    "WandbWatchCallbackConfig",
+    "_base",
+    "actsave",
+    "csv",
+    "tensorboard",
+    "wandb",
+]

nshtrainer/configs/loggers/_base/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.loggers._base import BaseLoggerConfig as BaseLoggerConfig
+
+__all__ = [
+    "BaseLoggerConfig",
+]

nshtrainer/configs/loggers/actsave/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.loggers.actsave import ActSaveLoggerConfig as ActSaveLoggerConfig
 from nshtrainer.loggers.actsave import BaseLoggerConfig as BaseLoggerConfig
+
+__all__ = [
+    "ActSaveLoggerConfig",
+    "BaseLoggerConfig",
+]

nshtrainer/configs/loggers/csv/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.loggers.csv import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers.csv import CSVLoggerConfig as CSVLoggerConfig
+
+__all__ = [
+    "BaseLoggerConfig",
+    "CSVLoggerConfig",
+]

nshtrainer/configs/loggers/tensorboard/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.loggers.tensorboard import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers.tensorboard import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
+
+__all__ = [
+    "BaseLoggerConfig",
+    "TensorboardLoggerConfig",
+]

nshtrainer/configs/loggers/wandb/__init__.py

@@ -11,3 +11,11 @@ from nshtrainer.loggers.wandb import (
 from nshtrainer.loggers.wandb import (
     WandbWatchCallbackConfig as WandbWatchCallbackConfig,
 )
+
+__all__ = [
+    "BaseLoggerConfig",
+    "CallbackConfigBase",
+    "WandbLoggerConfig",
+    "WandbUploadCodeCallbackConfig",
+    "WandbWatchCallbackConfig",
+]

nshtrainer/configs/lr_scheduler/__init__.py

@@ -5,14 +5,20 @@ __codegen__ = True
 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.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 linear_warmup_cosine as linear_warmup_cosine
 from . import reduce_lr_on_plateau as reduce_lr_on_plateau
+
+__all__ = [
+    "LRSchedulerConfigBase",
+    "LinearWarmupCosineDecayLRSchedulerConfig",
+    "MetricConfig",
+    "ReduceLROnPlateauConfig",
+    "_base",
+    "linear_warmup_cosine",
+    "reduce_lr_on_plateau",
+]

nshtrainer/configs/lr_scheduler/_base/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.lr_scheduler._base import LRSchedulerConfigBase as LRSchedulerConfigBase
+
+__all__ = [
+    "LRSchedulerConfigBase",
+]

nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.lr_scheduler.linear_warmup_cosine import (
-    DurationConfig as DurationConfig,
-)
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig,
 )
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     LRSchedulerConfigBase as LRSchedulerConfigBase,
 )
+
+__all__ = [
+    "LRSchedulerConfigBase",
+    "LinearWarmupCosineDecayLRSchedulerConfig",
+]

nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py

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

nshtrainer/configs/metrics/__init__.py

@@ -5,3 +5,8 @@ __codegen__ = True
 from nshtrainer.metrics import MetricConfig as MetricConfig
 
 from . import _config as _config
+
+__all__ = [
+    "MetricConfig",
+    "_config",
+]

nshtrainer/configs/metrics/_config/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.metrics._config import MetricConfig as MetricConfig
+
+__all__ = [
+    "MetricConfig",
+]

nshtrainer/configs/nn/__init__.py

@@ -8,7 +8,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
@@ -24,3 +23,24 @@ from nshtrainer.nn.nonlinearity import (
 
 from . import mlp as mlp
 from . import nonlinearity as nonlinearity
+
+__all__ = [
+    "BaseNonlinearityConfig",
+    "ELUNonlinearityConfig",
+    "GELUNonlinearityConfig",
+    "LeakyReLUNonlinearityConfig",
+    "MLPConfig",
+    "MishNonlinearityConfig",
+    "PReLUConfig",
+    "ReLUNonlinearityConfig",
+    "SiLUNonlinearityConfig",
+    "SigmoidNonlinearityConfig",
+    "SoftmaxNonlinearityConfig",
+    "SoftplusNonlinearityConfig",
+    "SoftsignNonlinearityConfig",
+    "SwiGLUNonlinearityConfig",
+    "SwishNonlinearityConfig",
+    "TanhNonlinearityConfig",
+    "mlp",
+    "nonlinearity",
+]

nshtrainer/configs/nn/mlp/__init__.py

@@ -4,4 +4,8 @@ __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
+
+__all__ = [
+    "BaseNonlinearityConfig",
+    "MLPConfig",
+]

nshtrainer/configs/nn/nonlinearity/__init__.py

@@ -9,7 +9,6 @@ from nshtrainer.nn.nonlinearity import (
     LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig,
 )
 from nshtrainer.nn.nonlinearity import MishNonlinearityConfig as MishNonlinearityConfig
-from nshtrainer.nn.nonlinearity import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn.nonlinearity import PReLUConfig as PReLUConfig
 from nshtrainer.nn.nonlinearity import ReLUNonlinearityConfig as ReLUNonlinearityConfig
 from nshtrainer.nn.nonlinearity import (
@@ -32,3 +31,21 @@ from nshtrainer.nn.nonlinearity import (
     SwishNonlinearityConfig as SwishNonlinearityConfig,
 )
 from nshtrainer.nn.nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig
+
+__all__ = [
+    "BaseNonlinearityConfig",
+    "ELUNonlinearityConfig",
+    "GELUNonlinearityConfig",
+    "LeakyReLUNonlinearityConfig",
+    "MishNonlinearityConfig",
+    "PReLUConfig",
+    "ReLUNonlinearityConfig",
+    "SiLUNonlinearityConfig",
+    "SigmoidNonlinearityConfig",
+    "SoftmaxNonlinearityConfig",
+    "SoftplusNonlinearityConfig",
+    "SoftsignNonlinearityConfig",
+    "SwiGLUNonlinearityConfig",
+    "SwishNonlinearityConfig",
+    "TanhNonlinearityConfig",
+]

nshtrainer/configs/optimizer/__init__.py

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

nshtrainer/configs/profiler/__init__.py

@@ -4,7 +4,6 @@ __codegen__ = True
 
 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
 
@@ -12,3 +11,14 @@ from . import _base as _base
 from . import advanced as advanced
 from . import pytorch as pytorch
 from . import simple as simple
+
+__all__ = [
+    "AdvancedProfilerConfig",
+    "BaseProfilerConfig",
+    "PyTorchProfilerConfig",
+    "SimpleProfilerConfig",
+    "_base",
+    "advanced",
+    "pytorch",
+    "simple",
+]

nshtrainer/configs/profiler/_base/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.profiler._base import BaseProfilerConfig as BaseProfilerConfig
+
+__all__ = [
+    "BaseProfilerConfig",
+]

nshtrainer/configs/profiler/advanced/__init__.py

@@ -6,3 +6,8 @@ from nshtrainer.profiler.advanced import (
     AdvancedProfilerConfig as AdvancedProfilerConfig,
 )
 from nshtrainer.profiler.advanced import BaseProfilerConfig as BaseProfilerConfig
+
+__all__ = [
+    "AdvancedProfilerConfig",
+    "BaseProfilerConfig",
+]

nshtrainer/configs/profiler/pytorch/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.profiler.pytorch import BaseProfilerConfig as BaseProfilerConfig
 from nshtrainer.profiler.pytorch import PyTorchProfilerConfig as PyTorchProfilerConfig
+
+__all__ = [
+    "BaseProfilerConfig",
+    "PyTorchProfilerConfig",
+]

nshtrainer/configs/profiler/simple/__init__.py

@@ -4,3 +4,8 @@ __codegen__ = True
 
 from nshtrainer.profiler.simple import BaseProfilerConfig as BaseProfilerConfig
 from nshtrainer.profiler.simple import SimpleProfilerConfig as SimpleProfilerConfig
+
+__all__ = [
+    "BaseProfilerConfig",
+    "SimpleProfilerConfig",
+]

nshtrainer/configs/trainer/__init__.py

@@ -3,16 +3,13 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.trainer import TrainerConfig as TrainerConfig
+from nshtrainer.trainer._config import AcceleratorConfigBase as AcceleratorConfigBase
 from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
 from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.trainer._config import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
-from nshtrainer.trainer._config import CallbackConfig as CallbackConfig
 from nshtrainer.trainer._config import CallbackConfigBase as CallbackConfigBase
-from nshtrainer.trainer._config import (
-    CheckpointCallbackConfig as CheckpointCallbackConfig,
-)
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
 from nshtrainer.trainer._config import CSVLoggerConfig as CSVLoggerConfig
 from nshtrainer.trainer._config import (
@@ -32,7 +29,6 @@ from nshtrainer.trainer._config import (
     LearningRateMonitorConfig as LearningRateMonitorConfig,
 )
 from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbackConfig
-from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
@@ -40,7 +36,7 @@ from nshtrainer.trainer._config import (
 from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
-from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
+from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
@@ -48,6 +44,7 @@ from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingCon
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.trainer._config import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
@@ -55,3 +52,35 @@ from nshtrainer.trainer._config import WandbLoggerConfig as WandbLoggerConfig
 
 from . import _config as _config
 from . import trainer as trainer
+
+__all__ = [
+    "AcceleratorConfigBase",
+    "ActSaveLoggerConfig",
+    "BaseLoggerConfig",
+    "BestCheckpointCallbackConfig",
+    "CSVLoggerConfig",
+    "CallbackConfigBase",
+    "CheckpointSavingConfig",
+    "DebugFlagCallbackConfig",
+    "DirectoryConfig",
+    "EarlyStoppingCallbackConfig",
+    "EnvironmentConfig",
+    "GradientClippingConfig",
+    "HuggingFaceHubConfig",
+    "LastCheckpointCallbackConfig",
+    "LearningRateMonitorConfig",
+    "LogEpochCallbackConfig",
+    "MetricConfig",
+    "NormLoggingCallbackConfig",
+    "OnExceptionCheckpointCallbackConfig",
+    "PluginConfigBase",
+    "RLPSanityChecksCallbackConfig",
+    "SanityCheckingConfig",
+    "SharedParametersCallbackConfig",
+    "StrategyConfigBase",
+    "TensorboardLoggerConfig",
+    "TrainerConfig",
+    "WandbLoggerConfig",
+    "_config",
+    "trainer",
+]

nshtrainer/configs/trainer/_config/__init__.py

@@ -2,16 +2,13 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.trainer._config import AcceleratorConfigBase as AcceleratorConfigBase
 from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
 from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.trainer._config import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
-from nshtrainer.trainer._config import CallbackConfig as CallbackConfig
 from nshtrainer.trainer._config import CallbackConfigBase as CallbackConfigBase
-from nshtrainer.trainer._config import (
-    CheckpointCallbackConfig as CheckpointCallbackConfig,
-)
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
 from nshtrainer.trainer._config import CSVLoggerConfig as CSVLoggerConfig
 from nshtrainer.trainer._config import (
@@ -31,7 +28,6 @@ from nshtrainer.trainer._config import (
     LearningRateMonitorConfig as LearningRateMonitorConfig,
 )
 from nshtrainer.trainer._config import LogEpochCallbackConfig as LogEpochCallbackConfig
-from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
 from nshtrainer.trainer._config import MetricConfig as MetricConfig
 from nshtrainer.trainer._config import (
     NormLoggingCallbackConfig as NormLoggingCallbackConfig,
@@ -39,7 +35,7 @@ from nshtrainer.trainer._config import (
 from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
-from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
+from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
@@ -47,8 +43,39 @@ from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingCon
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.trainer._config import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
 from nshtrainer.trainer._config import TrainerConfig as TrainerConfig
 from nshtrainer.trainer._config import WandbLoggerConfig as WandbLoggerConfig
+
+__all__ = [
+    "AcceleratorConfigBase",
+    "ActSaveLoggerConfig",
+    "BaseLoggerConfig",
+    "BestCheckpointCallbackConfig",
+    "CSVLoggerConfig",
+    "CallbackConfigBase",
+    "CheckpointSavingConfig",
+    "DebugFlagCallbackConfig",
+    "DirectoryConfig",
+    "EarlyStoppingCallbackConfig",
+    "EnvironmentConfig",
+    "GradientClippingConfig",
+    "HuggingFaceHubConfig",
+    "LastCheckpointCallbackConfig",
+    "LearningRateMonitorConfig",
+    "LogEpochCallbackConfig",
+    "MetricConfig",
+    "NormLoggingCallbackConfig",
+    "OnExceptionCheckpointCallbackConfig",
+    "PluginConfigBase",
+    "RLPSanityChecksCallbackConfig",
+    "SanityCheckingConfig",
+    "SharedParametersCallbackConfig",
+    "StrategyConfigBase",
+    "TensorboardLoggerConfig",
+    "TrainerConfig",
+    "WandbLoggerConfig",
+]

nshtrainer/configs/trainer/trainer/__init__.py

@@ -2,5 +2,14 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.trainer.trainer import AcceleratorConfigBase as AcceleratorConfigBase
 from nshtrainer.trainer.trainer import EnvironmentConfig as EnvironmentConfig
+from nshtrainer.trainer.trainer import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.trainer.trainer import TrainerConfig as TrainerConfig
+
+__all__ = [
+    "AcceleratorConfigBase",
+    "EnvironmentConfig",
+    "StrategyConfigBase",
+    "TrainerConfig",
+]

nshtrainer/configs/util/__init__.py

@@ -32,9 +32,27 @@ 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
 
 from . import _environment_info as _environment_info
 from . import config as config
+
+__all__ = [
+    "DTypeConfig",
+    "EnvironmentCUDAConfig",
+    "EnvironmentClassInformationConfig",
+    "EnvironmentConfig",
+    "EnvironmentGPUConfig",
+    "EnvironmentHardwareConfig",
+    "EnvironmentLSFInformationConfig",
+    "EnvironmentLinuxEnvironmentConfig",
+    "EnvironmentPackageConfig",
+    "EnvironmentSLURMInformationConfig",
+    "EnvironmentSnapshotConfig",
+    "EpochsConfig",
+    "GitRepositoryConfig",
+    "StepsConfig",
+    "_environment_info",
+    "config",
+]

nshtrainer/configs/util/_environment_info/__init__.py

@@ -31,3 +31,17 @@ from nshtrainer.util._environment_info import (
     EnvironmentSnapshotConfig as EnvironmentSnapshotConfig,
 )
 from nshtrainer.util._environment_info import GitRepositoryConfig as GitRepositoryConfig
+
+__all__ = [
+    "EnvironmentCUDAConfig",
+    "EnvironmentClassInformationConfig",
+    "EnvironmentConfig",
+    "EnvironmentGPUConfig",
+    "EnvironmentHardwareConfig",
+    "EnvironmentLSFInformationConfig",
+    "EnvironmentLinuxEnvironmentConfig",
+    "EnvironmentPackageConfig",
+    "EnvironmentSLURMInformationConfig",
+    "EnvironmentSnapshotConfig",
+    "GitRepositoryConfig",
+]

nshtrainer/configs/util/config/__init__.py

@@ -3,9 +3,16 @@ from __future__ import annotations
 __codegen__ = True
 
 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
 
 from . import dtype as dtype
 from . import duration as duration
+
+__all__ = [
+    "DTypeConfig",
+    "EpochsConfig",
+    "StepsConfig",
+    "dtype",
+    "duration",
+]

nshtrainer/configs/util/config/dtype/__init__.py

@@ -3,3 +3,7 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.util.config.dtype import DTypeConfig as DTypeConfig
+
+__all__ = [
+    "DTypeConfig",
+]

nshtrainer/configs/util/config/duration/__init__.py

@@ -2,6 +2,10 @@ from __future__ import annotations
 
 __codegen__ = True
 
-from nshtrainer.util.config.duration import DurationConfig as DurationConfig
 from nshtrainer.util.config.duration import EpochsConfig as EpochsConfig
 from nshtrainer.util.config.duration import StepsConfig as StepsConfig
+
+__all__ = [
+    "EpochsConfig",
+    "StepsConfig",
+]

nshtrainer/model/mixins/logger.py

@@ -64,6 +64,32 @@ class LoggerLightningModuleMixin(mixin_base_type(LightningModule)):
 
         self._logger_prefix_stack = deque[_LogContextKwargs]()
 
+    @property
+    def logging_enabled(self) -> bool:
+        # Logging is disabled in barebones mode.
+        if (trainer := self._trainer) is not None and trainer.barebones:
+            # Warn the user once that logging is disabled in barebones mode.
+            if not hasattr(self, "_barebones_logging_warned"):
+                rank_zero_warn(
+                    "Logging is disabled in barebones mode. "
+                    "This is to reduce the overhead of logging in barebones mode. "
+                    "If you want to enable logging, set `barebones=False` in the Trainer.",
+                )
+                self._barebones_logging_warned = True
+            return False
+
+        # If no loggers are registered, then logging is disabled.
+        if not self.logger:
+            return False
+
+        # Check if the topmost non-null context is disabled
+        for context in reversed(self._logger_prefix_stack):
+            if context.disabled is not None:
+                return not context.disabled
+
+        # Otherwise, logging is enabled.
+        return True
+
     @contextmanager
     def log_context(
         self,
@@ -121,22 +147,14 @@ class LoggerLightningModuleMixin(mixin_base_type(LightningModule)):
         metric_attribute: str | None = None,
         rank_zero_only: bool = False,
     ) -> None:
+        # If logging is disabled, then do nothing.
+        if not self.logging_enabled:
+            return
+
         # join all prefixes
         prefix = "".join(c.prefix for c in self._logger_prefix_stack if c.prefix)
         name = f"{prefix}{name}"
 
-        # check for disabled context:
-        # if the topmost non-null context is disabled, then we don't log
-        for c in reversed(self._logger_prefix_stack):
-            if c.disabled is not None:
-                if c.disabled:
-                    rank_zero_warn(
-                        f"Skipping logging of {name} due to disabled context"
-                    )
-                    return
-                else:
-                    break
-
         fn_kwargs = _LogContextKwargs()
         for c in self._logger_prefix_stack:
             fn_kwargs = fn_kwargs.copy_from(c)