nshtrainer 1.0.0b26__tar.gz → 1.0.0b28__tar.gz

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/PKG-INFO

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

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "1.0.0-beta26"
+version = "1.0.0-beta28"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/callbacks/actsave.py

@@ -4,12 +4,12 @@ import contextlib
 from pathlib import Path
 from typing import Literal
 
-from typing_extensions import TypeAlias, override
+from typing_extensions import TypeAliasType, override
 
 from .._callback import NTCallbackBase
 from .base import CallbackConfigBase
 
-Stage: TypeAlias = Literal["train", "validation", "test", "predict"]
+Stage = TypeAliasType("Stage", Literal["train", "validation", "test", "predict"])
 
 
 class ActSaveConfig(CallbackConfigBase):

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/callbacks/base.py

@@ -4,11 +4,11 @@ from abc import ABC, abstractmethod
 from collections import Counter
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar, TypeAlias
+from typing import TYPE_CHECKING, ClassVar
 
 import nshconfig as C
 from lightning.pytorch import Callback
-from typing_extensions import TypedDict, Unpack
+from typing_extensions import TypeAliasType, TypedDict, Unpack
 
 if TYPE_CHECKING:
     from ..trainer._config import TrainerConfig
@@ -30,7 +30,9 @@ class CallbackWithMetadata:
     metadata: CallbackMetadataConfig
 
 
-ConstructedCallback: TypeAlias = Callback | CallbackWithMetadata
+ConstructedCallback = TypeAliasType(
+    "ConstructedCallback", Callback | CallbackWithMetadata
+)
 
 
 class CallbackConfigBase(C.Config, ABC):

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/callbacks/shared_parameters.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 
 import logging
 from collections.abc import Iterable
-from typing import Literal, Protocol, TypeAlias, runtime_checkable
+from typing import Literal, Protocol, runtime_checkable
 
 import torch.nn as nn
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 from .base import CallbackConfigBase
 
@@ -34,7 +34,9 @@ class SharedParametersCallbackConfig(CallbackConfigBase):
         yield SharedParametersCallback(self)
 
 
-SharedParametersList: TypeAlias = list[tuple[nn.Parameter, int | float]]
+SharedParametersList = TypeAliasType(
+    "SharedParametersList", list[tuple[nn.Parameter, int | float]]
+)
 
 
 @runtime_checkable

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/__init__.py

@@ -14,6 +14,7 @@ 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,
@@ -46,6 +47,9 @@ from nshtrainer.callbacks import (
 from nshtrainer.callbacks import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+from nshtrainer.callbacks import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
 from nshtrainer.callbacks import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
 )
@@ -57,11 +61,13 @@ 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
@@ -70,6 +76,7 @@ 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
@@ -83,12 +90,17 @@ 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 AcceleratorConfigBase as AcceleratorConfigBase
+from nshtrainer.trainer._config import (
+    CheckpointCallbackConfig as CheckpointCallbackConfig,
+)
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
 from nshtrainer.trainer._config import EnvironmentConfig as EnvironmentConfig
 from nshtrainer.trainer._config import GradientClippingConfig as GradientClippingConfig
@@ -127,6 +139,7 @@ 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
 
@@ -155,13 +168,16 @@ __all__ = [
     "BaseProfilerConfig",
     "BestCheckpointCallbackConfig",
     "CSVLoggerConfig",
+    "CallbackConfig",
     "CallbackConfigBase",
+    "CheckpointCallbackConfig",
     "CheckpointMetadata",
     "CheckpointSavingConfig",
     "DTypeConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
     "DirectorySetupCallbackConfig",
+    "DurationConfig",
     "ELUNonlinearityConfig",
     "EMACallbackConfig",
     "EarlyStoppingCallbackConfig",
@@ -184,21 +200,26 @@ __all__ = [
     "GradientSkippingCallbackConfig",
     "HuggingFaceHubAutoCreateConfig",
     "HuggingFaceHubConfig",
+    "LRSchedulerConfig",
     "LRSchedulerConfigBase",
     "LastCheckpointCallbackConfig",
     "LeakyReLUNonlinearityConfig",
     "LearningRateMonitorConfig",
     "LinearWarmupCosineDecayLRSchedulerConfig",
     "LogEpochCallbackConfig",
+    "LoggerConfig",
     "MLPConfig",
     "MetricConfig",
     "MishNonlinearityConfig",
+    "NonlinearityConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
+    "OptimizerConfig",
     "OptimizerConfigBase",
     "PReLUConfig",
     "PluginConfigBase",
     "PrintTableMetricsCallbackConfig",
+    "ProfilerConfig",
     "PyTorchProfilerConfig",
     "RLPSanityChecksCallbackConfig",
     "ReLUNonlinearityConfig",
@@ -217,6 +238,7 @@ __all__ = [
     "SwishNonlinearityConfig",
     "TanhNonlinearityConfig",
     "TensorboardLoggerConfig",
+    "TimeCheckpointCallbackConfig",
     "TrainerConfig",
     "WandbLoggerConfig",
     "WandbUploadCodeCallbackConfig",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/_directory/__init__.py

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

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/callbacks/__init__.py

@@ -5,6 +5,7 @@ __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 (
@@ -38,6 +39,9 @@ from nshtrainer.callbacks import (
 from nshtrainer.callbacks import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
+from nshtrainer.callbacks import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
 from nshtrainer.callbacks import (
     WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
 )
@@ -77,6 +81,7 @@ __all__ = [
     "ActSaveConfig",
     "BaseCheckpointCallbackConfig",
     "BestCheckpointCallbackConfig",
+    "CallbackConfig",
     "CallbackConfigBase",
     "CheckpointMetadata",
     "DebugFlagCallbackConfig",
@@ -95,6 +100,7 @@ __all__ = [
     "PrintTableMetricsCallbackConfig",
     "RLPSanityChecksCallbackConfig",
     "SharedParametersCallbackConfig",
+    "TimeCheckpointCallbackConfig",
     "WandbUploadCodeCallbackConfig",
     "WandbWatchCallbackConfig",
     "actsave",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/callbacks/checkpoint/__init__.py

@@ -11,6 +11,9 @@ from nshtrainer.callbacks.checkpoint import (
 from nshtrainer.callbacks.checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
+from nshtrainer.callbacks.checkpoint import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
 from nshtrainer.callbacks.checkpoint._base import (
     BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
 )
@@ -26,6 +29,7 @@ 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
+from . import time_checkpoint as time_checkpoint
 
 __all__ = [
     "BaseCheckpointCallbackConfig",
@@ -35,8 +39,10 @@ __all__ = [
     "LastCheckpointCallbackConfig",
     "MetricConfig",
     "OnExceptionCheckpointCallbackConfig",
+    "TimeCheckpointCallbackConfig",
     "_base",
     "best_checkpoint",
     "last_checkpoint",
     "on_exception_checkpoint",
+    "time_checkpoint",
 ]

nshtrainer-1.0.0b28/src/nshtrainer/configs/callbacks/checkpoint/time_checkpoint/__init__.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.callbacks.checkpoint.time_checkpoint import (
+    BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
+)
+from nshtrainer.callbacks.checkpoint.time_checkpoint import (
+    CheckpointMetadata as CheckpointMetadata,
+)
+from nshtrainer.callbacks.checkpoint.time_checkpoint import (
+    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
+)
+
+__all__ = [
+    "BaseCheckpointCallbackConfig",
+    "CheckpointMetadata",
+    "TimeCheckpointCallbackConfig",
+]

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/loggers/__init__.py

@@ -5,6 +5,7 @@ __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
@@ -26,6 +27,7 @@ __all__ = [
     "BaseLoggerConfig",
     "CSVLoggerConfig",
     "CallbackConfigBase",
+    "LoggerConfig",
     "TensorboardLoggerConfig",
     "WandbLoggerConfig",
     "WandbUploadCodeCallbackConfig",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/lr_scheduler/__init__.py

@@ -5,8 +5,12 @@ __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
@@ -14,6 +18,8 @@ from . import linear_warmup_cosine as linear_warmup_cosine
 from . import reduce_lr_on_plateau as reduce_lr_on_plateau
 
 __all__ = [
+    "DurationConfig",
+    "LRSchedulerConfig",
     "LRSchedulerConfigBase",
     "LinearWarmupCosineDecayLRSchedulerConfig",
     "MetricConfig",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py

@@ -2,6 +2,9 @@ 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,
 )
@@ -10,6 +13,7 @@ from nshtrainer.lr_scheduler.linear_warmup_cosine import (
 )
 
 __all__ = [
+    "DurationConfig",
     "LRSchedulerConfigBase",
     "LinearWarmupCosineDecayLRSchedulerConfig",
 ]

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/nn/__init__.py

@@ -8,6 +8,7 @@ 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
@@ -31,6 +32,7 @@ __all__ = [
     "LeakyReLUNonlinearityConfig",
     "MLPConfig",
     "MishNonlinearityConfig",
+    "NonlinearityConfig",
     "PReLUConfig",
     "ReLUNonlinearityConfig",
     "SiLUNonlinearityConfig",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/nn/mlp/__init__.py

@@ -4,8 +4,10 @@ __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",
+    "NonlinearityConfig",
 ]

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/nn/nonlinearity/__init__.py

@@ -9,6 +9,7 @@ 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 (
@@ -38,6 +39,7 @@ __all__ = [
     "GELUNonlinearityConfig",
     "LeakyReLUNonlinearityConfig",
     "MishNonlinearityConfig",
+    "NonlinearityConfig",
     "PReLUConfig",
     "ReLUNonlinearityConfig",
     "SiLUNonlinearityConfig",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/optimizer/__init__.py

@@ -3,9 +3,11 @@ 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",
+    "OptimizerConfig",
     "OptimizerConfigBase",
 ]

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/profiler/__init__.py

@@ -4,6 +4,7 @@ __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
 
@@ -15,6 +16,7 @@ from . import simple as simple
 __all__ = [
     "AdvancedProfilerConfig",
     "BaseProfilerConfig",
+    "ProfilerConfig",
     "PyTorchProfilerConfig",
     "SimpleProfilerConfig",
     "_base",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/trainer/__init__.py

@@ -9,7 +9,11 @@ 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 (
@@ -29,6 +33,7 @@ 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,
@@ -37,6 +42,7 @@ from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
 from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
@@ -48,6 +54,9 @@ from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
 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 . import _config as _config
@@ -59,7 +68,9 @@ __all__ = [
     "BaseLoggerConfig",
     "BestCheckpointCallbackConfig",
     "CSVLoggerConfig",
+    "CallbackConfig",
     "CallbackConfigBase",
+    "CheckpointCallbackConfig",
     "CheckpointSavingConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
@@ -70,15 +81,18 @@ __all__ = [
     "LastCheckpointCallbackConfig",
     "LearningRateMonitorConfig",
     "LogEpochCallbackConfig",
+    "LoggerConfig",
     "MetricConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "PluginConfigBase",
+    "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
     "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "StrategyConfigBase",
     "TensorboardLoggerConfig",
+    "TimeCheckpointCallbackConfig",
     "TrainerConfig",
     "WandbLoggerConfig",
     "_config",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/trainer/_config/__init__.py

@@ -8,7 +8,11 @@ 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 (
@@ -28,6 +32,7 @@ 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,
@@ -36,6 +41,7 @@ from nshtrainer.trainer._config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
 from nshtrainer.trainer._config import PluginConfigBase as PluginConfigBase
+from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
@@ -47,6 +53,9 @@ from nshtrainer.trainer._config import StrategyConfigBase as StrategyConfigBase
 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
 
@@ -56,7 +65,9 @@ __all__ = [
     "BaseLoggerConfig",
     "BestCheckpointCallbackConfig",
     "CSVLoggerConfig",
+    "CallbackConfig",
     "CallbackConfigBase",
+    "CheckpointCallbackConfig",
     "CheckpointSavingConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
@@ -67,15 +78,18 @@ __all__ = [
     "LastCheckpointCallbackConfig",
     "LearningRateMonitorConfig",
     "LogEpochCallbackConfig",
+    "LoggerConfig",
     "MetricConfig",
     "NormLoggingCallbackConfig",
     "OnExceptionCheckpointCallbackConfig",
     "PluginConfigBase",
+    "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
     "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "StrategyConfigBase",
     "TensorboardLoggerConfig",
+    "TimeCheckpointCallbackConfig",
     "TrainerConfig",
     "WandbLoggerConfig",
 ]

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/util/__init__.py

@@ -32,6 +32,7 @@ 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
 
@@ -40,6 +41,7 @@ from . import config as config
 
 __all__ = [
     "DTypeConfig",
+    "DurationConfig",
     "EnvironmentCUDAConfig",
     "EnvironmentClassInformationConfig",
     "EnvironmentConfig",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/util/config/__init__.py

@@ -3,6 +3,7 @@ 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
 
@@ -11,6 +12,7 @@ from . import duration as duration
 
 __all__ = [
     "DTypeConfig",
+    "DurationConfig",
     "EpochsConfig",
     "StepsConfig",
     "dtype",

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/configs/util/config/duration/__init__.py

@@ -2,10 +2,12 @@ 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__ = [
+    "DurationConfig",
     "EpochsConfig",
     "StepsConfig",
 ]

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/loggers/__init__.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from typing import Annotated, TypeAlias
+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
@@ -10,7 +11,13 @@ from .csv import CSVLoggerConfig as CSVLoggerConfig
 from .tensorboard import TensorboardLoggerConfig as TensorboardLoggerConfig
 from .wandb import WandbLoggerConfig as WandbLoggerConfig
 
-LoggerConfig: TypeAlias = Annotated[
-    CSVLoggerConfig | TensorboardLoggerConfig | WandbLoggerConfig | ActSaveLoggerConfig,
-    C.Field(discriminator="name"),
-]
+LoggerConfig = TypeAliasType(
+    "LoggerConfig",
+    Annotated[
+        CSVLoggerConfig
+        | TensorboardLoggerConfig
+        | WandbLoggerConfig
+        | ActSaveLoggerConfig,
+        C.Field(discriminator="name"),
+    ],
+)

{nshtrainer-1.0.0b26 → nshtrainer-1.0.0b28}/src/nshtrainer/lr_scheduler/__init__.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from typing import Annotated, TypeAlias
+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
@@ -15,7 +16,10 @@ from .linear_warmup_cosine import (
 from .reduce_lr_on_plateau import ReduceLROnPlateau as ReduceLROnPlateau
 from .reduce_lr_on_plateau import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
 
-LRSchedulerConfig: TypeAlias = Annotated[
-    LinearWarmupCosineDecayLRSchedulerConfig | ReduceLROnPlateauConfig,
-    C.Field(discriminator="name"),
-]
+LRSchedulerConfig = TypeAliasType(
+    "LRSchedulerConfig",
+    Annotated[
+        LinearWarmupCosineDecayLRSchedulerConfig | ReduceLROnPlateauConfig,
+        C.Field(discriminator="name"),
+    ],
+)