nshtrainer 0.44.1__py3-none-any.whl → 1.0.0b10__py3-none-any.whl

nshtrainer/callbacks/wandb_upload_code.py

@@ -26,7 +26,7 @@ class WandbUploadCodeCallbackConfig(CallbackConfigBase):
         return self.enabled
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         if not self:
             return
 

nshtrainer/callbacks/wandb_watch.py

@@ -33,7 +33,7 @@ class WandbWatchCallbackConfig(CallbackConfigBase):
         return self.enabled
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield WandbWatchCallback(self)
 
 

nshtrainer/config/__init__.py

@@ -7,8 +7,8 @@ from typing import TYPE_CHECKING
 # Config/alias imports
 
 if TYPE_CHECKING:
-    from nshtrainer import BaseConfig as BaseConfig
     from nshtrainer import MetricConfig as MetricConfig
+    from nshtrainer import TrainerConfig as TrainerConfig
     from nshtrainer._checkpoint.loader import (
         BestCheckpointStrategyConfig as BestCheckpointStrategyConfig,
     )
@@ -22,6 +22,7 @@ if TYPE_CHECKING:
     from nshtrainer._checkpoint.loader import (
         UserProvidedPathCheckpointStrategyConfig as UserProvidedPathCheckpointStrategyConfig,
     )
+    from nshtrainer._directory import DirectoryConfig as DirectoryConfig
     from nshtrainer._hf_hub import CallbackConfigBase as CallbackConfigBase
     from nshtrainer._hf_hub import (
         HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig,
@@ -51,6 +52,7 @@ if TYPE_CHECKING:
     from nshtrainer.callbacks import (
         LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
     )
+    from nshtrainer.callbacks import LogEpochCallbackConfig as LogEpochCallbackConfig
     from nshtrainer.callbacks import (
         NormLoggingCallbackConfig as NormLoggingCallbackConfig,
     )
@@ -66,7 +68,6 @@ if TYPE_CHECKING:
     from nshtrainer.callbacks import (
         SharedParametersCallbackConfig as SharedParametersCallbackConfig,
     )
-    from nshtrainer.callbacks import ThroughputMonitorConfig as ThroughputMonitorConfig
     from nshtrainer.callbacks import (
         WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
     )
@@ -77,6 +78,7 @@ if TYPE_CHECKING:
     from nshtrainer.callbacks.checkpoint._base import (
         BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
     )
+    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
@@ -90,9 +92,6 @@ if TYPE_CHECKING:
     from nshtrainer.lr_scheduler import (
         ReduceLROnPlateauConfig as ReduceLROnPlateauConfig,
     )
-    from nshtrainer.model import DirectoryConfig as DirectoryConfig
-    from nshtrainer.model import TrainerConfig as TrainerConfig
-    from nshtrainer.model.base import EnvironmentConfig as EnvironmentConfig
     from nshtrainer.nn import BaseNonlinearityConfig as BaseNonlinearityConfig
     from nshtrainer.nn import ELUNonlinearityConfig as ELUNonlinearityConfig
     from nshtrainer.nn import GELUNonlinearityConfig as GELUNonlinearityConfig
@@ -129,6 +128,7 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         CheckpointSavingConfig as CheckpointSavingConfig,
     )
+    from nshtrainer.trainer._config import EnvironmentConfig as EnvironmentConfig
     from nshtrainer.trainer._config import (
         GradientClippingConfig as GradientClippingConfig,
     )
@@ -179,272 +179,274 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "MetricConfig":
-            return importlib.import_module("nshtrainer").MetricConfig
-        if name == "BaseConfig":
-            return importlib.import_module("nshtrainer").BaseConfig
-        if name == "HuggingFaceHubAutoCreateConfig":
-            return importlib.import_module(
-                "nshtrainer._hf_hub"
-            ).HuggingFaceHubAutoCreateConfig
-        if name == "HuggingFaceHubConfig":
-            return importlib.import_module("nshtrainer._hf_hub").HuggingFaceHubConfig
-        if name == "CallbackConfigBase":
-            return importlib.import_module("nshtrainer._hf_hub").CallbackConfigBase
-        if name == "OptimizerConfigBase":
-            return importlib.import_module("nshtrainer.optimizer").OptimizerConfigBase
+        if name == "ActSaveConfig":
+            return importlib.import_module("nshtrainer.callbacks.actsave").ActSaveConfig
+        if name == "ActSaveLoggerConfig":
+            return importlib.import_module("nshtrainer.loggers").ActSaveLoggerConfig
         if name == "AdamWConfig":
             return importlib.import_module("nshtrainer.optimizer").AdamWConfig
-        if name == "DirectorySetupCallbackConfig":
+        if name == "AdvancedProfilerConfig":
+            return importlib.import_module("nshtrainer.profiler").AdvancedProfilerConfig
+        if name == "BaseCheckpointCallbackConfig":
             return importlib.import_module(
-                "nshtrainer.callbacks"
-            ).DirectorySetupCallbackConfig
-        if name == "DirectoryConfig":
-            return importlib.import_module("nshtrainer.model").DirectoryConfig
-        if name == "TrainerConfig":
-            return importlib.import_module("nshtrainer.model").TrainerConfig
-        if name == "EnvironmentConfig":
-            return importlib.import_module("nshtrainer.model.base").EnvironmentConfig
+                "nshtrainer.callbacks.checkpoint._base"
+            ).BaseCheckpointCallbackConfig
+        if name == "BaseLoggerConfig":
+            return importlib.import_module("nshtrainer.loggers").BaseLoggerConfig
         if name == "BaseNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").BaseNonlinearityConfig
-        if name == "MLPConfig":
-            return importlib.import_module("nshtrainer.nn").MLPConfig
-        if name == "PReLUConfig":
-            return importlib.import_module("nshtrainer.nn").PReLUConfig
-        if name == "LeakyReLUNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").LeakyReLUNonlinearityConfig
-        if name == "SwiGLUNonlinearityConfig":
+        if name == "BaseProfilerConfig":
+            return importlib.import_module("nshtrainer.profiler").BaseProfilerConfig
+        if name == "BestCheckpointCallbackConfig":
             return importlib.import_module(
-                "nshtrainer.nn.nonlinearity"
-            ).SwiGLUNonlinearityConfig
-        if name == "SoftsignNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SoftsignNonlinearityConfig
-        if name == "SiLUNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SiLUNonlinearityConfig
-        if name == "SigmoidNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SigmoidNonlinearityConfig
-        if name == "SoftplusNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SoftplusNonlinearityConfig
-        if name == "ELUNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").ELUNonlinearityConfig
-        if name == "SoftmaxNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SoftmaxNonlinearityConfig
-        if name == "GELUNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").GELUNonlinearityConfig
-        if name == "SwishNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SwishNonlinearityConfig
-        if name == "MishNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").MishNonlinearityConfig
-        if name == "TanhNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").TanhNonlinearityConfig
-        if name == "ReLUNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").ReLUNonlinearityConfig
-        if name == "LRSchedulerConfigBase":
+                "nshtrainer.callbacks"
+            ).BestCheckpointCallbackConfig
+        if name == "BestCheckpointStrategyConfig":
             return importlib.import_module(
-                "nshtrainer.lr_scheduler"
-            ).LRSchedulerConfigBase
-        if name == "LinearWarmupCosineDecayLRSchedulerConfig":
+                "nshtrainer._checkpoint.loader"
+            ).BestCheckpointStrategyConfig
+        if name == "CSVLoggerConfig":
+            return importlib.import_module("nshtrainer.loggers").CSVLoggerConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module("nshtrainer._hf_hub").CallbackConfigBase
+        if name == "CheckpointLoadingConfig":
             return importlib.import_module(
-                "nshtrainer.lr_scheduler"
-            ).LinearWarmupCosineDecayLRSchedulerConfig
-        if name == "ReduceLROnPlateauConfig":
+                "nshtrainer.trainer._config"
+            ).CheckpointLoadingConfig
+        if name == "CheckpointMetadata":
             return importlib.import_module(
-                "nshtrainer.lr_scheduler"
-            ).ReduceLROnPlateauConfig
-        if name == "BaseLoggerConfig":
-            return importlib.import_module("nshtrainer.loggers").BaseLoggerConfig
-        if name == "TensorboardLoggerConfig":
-            return importlib.import_module("nshtrainer.loggers").TensorboardLoggerConfig
-        if name == "WandbLoggerConfig":
-            return importlib.import_module("nshtrainer.loggers").WandbLoggerConfig
-        if name == "WandbUploadCodeCallbackConfig":
+                "nshtrainer._checkpoint.loader"
+            ).CheckpointMetadata
+        if name == "CheckpointSavingConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).CheckpointSavingConfig
+        if name == "DTypeConfig":
+            return importlib.import_module("nshtrainer.util.config").DTypeConfig
+        if name == "DebugFlagCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).WandbUploadCodeCallbackConfig
-        if name == "WandbWatchCallbackConfig":
+            ).DebugFlagCallbackConfig
+        if name == "DirectoryConfig":
+            return importlib.import_module("nshtrainer._directory").DirectoryConfig
+        if name == "DirectorySetupCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).WandbWatchCallbackConfig
-        if name == "CSVLoggerConfig":
-            return importlib.import_module("nshtrainer.loggers").CSVLoggerConfig
-        if name == "EnvironmentLinuxEnvironmentConfig":
+            ).DirectorySetupCallbackConfig
+        if name == "ELUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").ELUNonlinearityConfig
+        if name == "EMACallbackConfig":
+            return importlib.import_module("nshtrainer.callbacks").EMACallbackConfig
+        if name == "EarlyStoppingCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).EarlyStoppingCallbackConfig
+        if name == "EnvironmentCUDAConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentLinuxEnvironmentConfig
-        if name == "EnvironmentLSFInformationConfig":
+            ).EnvironmentCUDAConfig
+        if name == "EnvironmentClassInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentLSFInformationConfig
+            ).EnvironmentClassInformationConfig
+        if name == "EnvironmentConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).EnvironmentConfig
         if name == "EnvironmentGPUConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).EnvironmentGPUConfig
-        if name == "EnvironmentPackageConfig":
-            return importlib.import_module(
-                "nshtrainer.util._environment_info"
-            ).EnvironmentPackageConfig
         if name == "EnvironmentHardwareConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).EnvironmentHardwareConfig
-        if name == "EnvironmentSnapshotConfig":
-            return importlib.import_module(
-                "nshtrainer.util._environment_info"
-            ).EnvironmentSnapshotConfig
-        if name == "EnvironmentClassInformationConfig":
+        if name == "EnvironmentLSFInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentClassInformationConfig
-        if name == "GitRepositoryConfig":
+            ).EnvironmentLSFInformationConfig
+        if name == "EnvironmentLinuxEnvironmentConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).GitRepositoryConfig
-        if name == "EnvironmentCUDAConfig":
+            ).EnvironmentLinuxEnvironmentConfig
+        if name == "EnvironmentPackageConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentCUDAConfig
+            ).EnvironmentPackageConfig
         if name == "EnvironmentSLURMInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).EnvironmentSLURMInformationConfig
-        if name == "EpochsConfig":
-            return importlib.import_module("nshtrainer.util.config").EpochsConfig
-        if name == "StepsConfig":
-            return importlib.import_module("nshtrainer.util.config").StepsConfig
-        if name == "DTypeConfig":
-            return importlib.import_module("nshtrainer.util.config").DTypeConfig
-        if name == "CheckpointLoadingConfig":
+        if name == "EnvironmentSnapshotConfig":
             return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).CheckpointLoadingConfig
-        if name == "SanityCheckingConfig":
+                "nshtrainer.util._environment_info"
+            ).EnvironmentSnapshotConfig
+        if name == "EpochTimerCallbackConfig":
             return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).SanityCheckingConfig
-        if name == "OnExceptionCheckpointCallbackConfig":
+                "nshtrainer.callbacks"
+            ).EpochTimerCallbackConfig
+        if name == "EpochsConfig":
+            return importlib.import_module("nshtrainer.util.config").EpochsConfig
+        if name == "FiniteChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).OnExceptionCheckpointCallbackConfig
+            ).FiniteChecksCallbackConfig
+        if name == "GELUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").GELUNonlinearityConfig
+        if name == "GitRepositoryConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).GitRepositoryConfig
         if name == "GradientClippingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).GradientClippingConfig
-        if name == "LoggingConfig":
-            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
-        if name == "RLPSanityChecksCallbackConfig":
+        if name == "GradientSkippingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).RLPSanityChecksCallbackConfig
-        if name == "CheckpointSavingConfig":
+            ).GradientSkippingCallbackConfig
+        if name == "HuggingFaceHubAutoCreateConfig":
             return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).CheckpointSavingConfig
-        if name == "DebugFlagCallbackConfig":
+                "nshtrainer._hf_hub"
+            ).HuggingFaceHubAutoCreateConfig
+        if name == "HuggingFaceHubConfig":
+            return importlib.import_module("nshtrainer._hf_hub").HuggingFaceHubConfig
+        if name == "LRSchedulerConfigBase":
             return importlib.import_module(
-                "nshtrainer.callbacks"
-            ).DebugFlagCallbackConfig
+                "nshtrainer.lr_scheduler"
+            ).LRSchedulerConfigBase
         if name == "LastCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).LastCheckpointCallbackConfig
-        if name == "SharedParametersCallbackConfig":
+        if name == "LastCheckpointStrategyConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).LastCheckpointStrategyConfig
+        if name == "LeakyReLUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").LeakyReLUNonlinearityConfig
+        if name == "LinearWarmupCosineDecayLRSchedulerConfig":
+            return importlib.import_module(
+                "nshtrainer.lr_scheduler"
+            ).LinearWarmupCosineDecayLRSchedulerConfig
+        if name == "LogEpochCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).SharedParametersCallbackConfig
-        if name == "ReproducibilityConfig":
+            ).LogEpochCallbackConfig
+        if name == "LoggingConfig":
+            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
+        if name == "MLPConfig":
+            return importlib.import_module("nshtrainer.nn").MLPConfig
+        if name == "MetricConfig":
+            return importlib.import_module("nshtrainer").MetricConfig
+        if name == "MishNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").MishNonlinearityConfig
+        if name == "NormLoggingCallbackConfig":
             return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).ReproducibilityConfig
-        if name == "EarlyStoppingCallbackConfig":
+                "nshtrainer.callbacks"
+            ).NormLoggingCallbackConfig
+        if name == "OnExceptionCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).EarlyStoppingCallbackConfig
+            ).OnExceptionCheckpointCallbackConfig
         if name == "OptimizationConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).OptimizationConfig
-        if name == "BestCheckpointCallbackConfig":
+        if name == "OptimizerConfigBase":
+            return importlib.import_module("nshtrainer.optimizer").OptimizerConfigBase
+        if name == "PReLUConfig":
+            return importlib.import_module("nshtrainer.nn").PReLUConfig
+        if name == "PrintTableMetricsCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).BestCheckpointCallbackConfig
-        if name == "CheckpointMetadata":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).CheckpointMetadata
-        if name == "BestCheckpointStrategyConfig":
+            ).PrintTableMetricsCallbackConfig
+        if name == "PyTorchProfilerConfig":
+            return importlib.import_module("nshtrainer.profiler").PyTorchProfilerConfig
+        if name == "RLPSanityChecksCallbackConfig":
             return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).BestCheckpointStrategyConfig
-        if name == "LastCheckpointStrategyConfig":
+                "nshtrainer.callbacks"
+            ).RLPSanityChecksCallbackConfig
+        if name == "ReLUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").ReLUNonlinearityConfig
+        if name == "ReduceLROnPlateauConfig":
             return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).LastCheckpointStrategyConfig
-        if name == "UserProvidedPathCheckpointStrategyConfig":
+                "nshtrainer.lr_scheduler"
+            ).ReduceLROnPlateauConfig
+        if name == "ReproducibilityConfig":
             return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).UserProvidedPathCheckpointStrategyConfig
-        if name == "PrintTableMetricsCallbackConfig":
+                "nshtrainer.trainer._config"
+            ).ReproducibilityConfig
+        if name == "SanityCheckingConfig":
             return importlib.import_module(
-                "nshtrainer.callbacks"
-            ).PrintTableMetricsCallbackConfig
-        if name == "ThroughputMonitorConfig":
+                "nshtrainer.trainer._config"
+            ).SanityCheckingConfig
+        if name == "SharedParametersCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).ThroughputMonitorConfig
-        if name == "GradientSkippingCallbackConfig":
+            ).SharedParametersCallbackConfig
+        if name == "SiLUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SiLUNonlinearityConfig
+        if name == "SigmoidNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SigmoidNonlinearityConfig
+        if name == "SimpleProfilerConfig":
+            return importlib.import_module("nshtrainer.profiler").SimpleProfilerConfig
+        if name == "SoftmaxNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SoftmaxNonlinearityConfig
+        if name == "SoftplusNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SoftplusNonlinearityConfig
+        if name == "SoftsignNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SoftsignNonlinearityConfig
+        if name == "StepsConfig":
+            return importlib.import_module("nshtrainer.util.config").StepsConfig
+        if name == "SwiGLUNonlinearityConfig":
             return importlib.import_module(
-                "nshtrainer.callbacks"
-            ).GradientSkippingCallbackConfig
-        if name == "EMACallbackConfig":
-            return importlib.import_module("nshtrainer.callbacks").EMACallbackConfig
-        if name == "ActSaveConfig":
-            return importlib.import_module("nshtrainer.callbacks.actsave").ActSaveConfig
-        if name == "FiniteChecksCallbackConfig":
+                "nshtrainer.nn.nonlinearity"
+            ).SwiGLUNonlinearityConfig
+        if name == "SwishNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SwishNonlinearityConfig
+        if name == "TanhNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").TanhNonlinearityConfig
+        if name == "TensorboardLoggerConfig":
+            return importlib.import_module("nshtrainer.loggers").TensorboardLoggerConfig
+        if name == "TrainerConfig":
+            return importlib.import_module("nshtrainer").TrainerConfig
+        if name == "UserProvidedPathCheckpointStrategyConfig":
             return importlib.import_module(
-                "nshtrainer.callbacks"
-            ).FiniteChecksCallbackConfig
-        if name == "NormLoggingCallbackConfig":
+                "nshtrainer._checkpoint.loader"
+            ).UserProvidedPathCheckpointStrategyConfig
+        if name == "WandbLoggerConfig":
+            return importlib.import_module("nshtrainer.loggers").WandbLoggerConfig
+        if name == "WandbUploadCodeCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).NormLoggingCallbackConfig
-        if name == "EpochTimerCallbackConfig":
+            ).WandbUploadCodeCallbackConfig
+        if name == "WandbWatchCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).EpochTimerCallbackConfig
-        if name == "BaseCheckpointCallbackConfig":
-            return importlib.import_module(
-                "nshtrainer.callbacks.checkpoint._base"
-            ).BaseCheckpointCallbackConfig
-        if name == "BaseProfilerConfig":
-            return importlib.import_module("nshtrainer.profiler").BaseProfilerConfig
-        if name == "PyTorchProfilerConfig":
-            return importlib.import_module("nshtrainer.profiler").PyTorchProfilerConfig
-        if name == "AdvancedProfilerConfig":
-            return importlib.import_module("nshtrainer.profiler").AdvancedProfilerConfig
-        if name == "SimpleProfilerConfig":
-            return importlib.import_module("nshtrainer.profiler").SimpleProfilerConfig
-        if name == "OptimizerConfig":
-            return importlib.import_module("nshtrainer.optimizer").OptimizerConfig
-        if name == "LoggerConfig":
-            return importlib.import_module("nshtrainer.loggers").LoggerConfig
-        if name == "NonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").NonlinearityConfig
-        if name == "DurationConfig":
-            return importlib.import_module("nshtrainer.util.config").DurationConfig
-        if name == "LRSchedulerConfig":
-            return importlib.import_module("nshtrainer.lr_scheduler").LRSchedulerConfig
+            ).WandbWatchCallbackConfig
         if name == "CallbackConfig":
             return importlib.import_module("nshtrainer.callbacks").CallbackConfig
         if name == "CheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).CheckpointCallbackConfig
-        if name == "ProfilerConfig":
-            return importlib.import_module("nshtrainer.profiler").ProfilerConfig
         if name == "CheckpointLoadingStrategyConfig":
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"
             ).CheckpointLoadingStrategyConfig
+        if name == "DurationConfig":
+            return importlib.import_module("nshtrainer.util.config").DurationConfig
+        if name == "LRSchedulerConfig":
+            return importlib.import_module("nshtrainer.lr_scheduler").LRSchedulerConfig
+        if name == "LoggerConfig":
+            return importlib.import_module("nshtrainer.loggers").LoggerConfig
+        if name == "NonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").NonlinearityConfig
+        if name == "OptimizerConfig":
+            return importlib.import_module("nshtrainer.optimizer").OptimizerConfig
+        if name == "ProfilerConfig":
+            return importlib.import_module("nshtrainer.profiler").ProfilerConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 
@@ -456,10 +458,8 @@ from . import callbacks as callbacks
 from . import loggers as loggers
 from . import lr_scheduler as lr_scheduler
 from . import metrics as metrics
-from . import model as model
 from . import nn as nn
 from . import optimizer as optimizer
 from . import profiler as profiler
-from . import runner as runner
 from . import trainer as trainer
 from . import util as util

nshtrainer/config/_checkpoint/__init__.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer._checkpoint.loader import (
+        BestCheckpointStrategyConfig as BestCheckpointStrategyConfig,
+    )
+    from nshtrainer._checkpoint.loader import (
+        CheckpointLoadingConfig as CheckpointLoadingConfig,
+    )
+    from nshtrainer._checkpoint.loader import (
+        CheckpointLoadingStrategyConfig as CheckpointLoadingStrategyConfig,
+    )
+    from nshtrainer._checkpoint.loader import CheckpointMetadata as CheckpointMetadata
+    from nshtrainer._checkpoint.loader import (
+        LastCheckpointStrategyConfig as LastCheckpointStrategyConfig,
+    )
+    from nshtrainer._checkpoint.loader import MetricConfig as MetricConfig
+    from nshtrainer._checkpoint.loader import (
+        UserProvidedPathCheckpointStrategyConfig as UserProvidedPathCheckpointStrategyConfig,
+    )
+    from nshtrainer._checkpoint.metadata import EnvironmentConfig as EnvironmentConfig
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "BestCheckpointStrategyConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).BestCheckpointStrategyConfig
+        if name == "CheckpointLoadingConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).CheckpointLoadingConfig
+        if name == "CheckpointMetadata":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).CheckpointMetadata
+        if name == "EnvironmentConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.metadata"
+            ).EnvironmentConfig
+        if name == "LastCheckpointStrategyConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).LastCheckpointStrategyConfig
+        if name == "MetricConfig":
+            return importlib.import_module("nshtrainer._checkpoint.loader").MetricConfig
+        if name == "UserProvidedPathCheckpointStrategyConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).UserProvidedPathCheckpointStrategyConfig
+        if name == "CheckpointLoadingStrategyConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).CheckpointLoadingStrategyConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+# Submodule exports
+from . import loader as loader
+from . import metadata as metadata

nshtrainer/config/_checkpoint/loader/__init__.py

@@ -31,12 +31,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "MetricConfig":
-            return importlib.import_module("nshtrainer._checkpoint.loader").MetricConfig
         if name == "BestCheckpointStrategyConfig":
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"
             ).BestCheckpointStrategyConfig
+        if name == "CheckpointLoadingConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).CheckpointLoadingConfig
         if name == "CheckpointMetadata":
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"
@@ -45,10 +47,8 @@ else:
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"
             ).LastCheckpointStrategyConfig
-        if name == "CheckpointLoadingConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).CheckpointLoadingConfig
+        if name == "MetricConfig":
+            return importlib.import_module("nshtrainer._checkpoint.loader").MetricConfig
         if name == "UserProvidedPathCheckpointStrategyConfig":
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"