nshtrainer 1.0.0b25__py3-none-any.whl → 1.0.0b27__py3-none-any.whl

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/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/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"),
+    ],
+)

nshtrainer/model/mixins/callback.py

@@ -2,18 +2,20 @@ from __future__ import annotations
 
 import logging
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, TypeAlias, cast
+from typing import Any, cast
 
 from lightning.pytorch import Callback, LightningModule
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 from ..._callback import NTCallbackBase
 from ...util.typing_utils import mixin_base_type
 
 log = logging.getLogger(__name__)
 
-_Callback = Callback | NTCallbackBase
-CallbackFn: TypeAlias = Callable[[], _Callback | Iterable[_Callback] | None]
+_Callback = TypeAliasType("_Callback", Callback | NTCallbackBase)
+CallbackFn = TypeAliasType(
+    "CallbackFn", Callable[[], _Callback | Iterable[_Callback] | None]
+)
 
 
 class CallbackRegistrarModuleMixin:

nshtrainer/optimizer.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from typing import Annotated, Any, Literal, TypeAlias
+from typing import Annotated, Any, Literal
 
 import nshconfig as C
 import torch.nn as nn
 from torch.optim import Optimizer
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 
 class OptimizerConfigBase(C.Config, ABC):
@@ -57,4 +57,6 @@ class AdamWConfig(OptimizerConfigBase):
         )
 
 
-OptimizerConfig: TypeAlias = Annotated[AdamWConfig, C.Field(discriminator="name")]
+OptimizerConfig = TypeAliasType(
+    "OptimizerConfig", Annotated[AdamWConfig, C.Field(discriminator="name")]
+)

nshtrainer/profiler/__init__.py

@@ -1,15 +1,19 @@
 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 BaseProfilerConfig as BaseProfilerConfig
 from .advanced import AdvancedProfilerConfig as AdvancedProfilerConfig
 from .pytorch import PyTorchProfilerConfig as PyTorchProfilerConfig
 from .simple import SimpleProfilerConfig as SimpleProfilerConfig
 
-ProfilerConfig: TypeAlias = Annotated[
-    SimpleProfilerConfig | AdvancedProfilerConfig | PyTorchProfilerConfig,
-    C.Field(discriminator="name"),
-]
+ProfilerConfig = TypeAliasType(
+    "ProfilerConfig",
+    Annotated[
+        SimpleProfilerConfig | AdvancedProfilerConfig | PyTorchProfilerConfig,
+        C.Field(discriminator="name"),
+    ],
+)

nshtrainer/trainer/_config.py

@@ -5,6 +5,7 @@ import logging
 import os
 import string
 import time
+from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
 from datetime import timedelta
 from pathlib import Path
@@ -13,9 +14,6 @@ from typing import (
     Any,
     ClassVar,
     Literal,
-    Protocol,
-    TypeAlias,
-    runtime_checkable,
 )
 
 import nshconfig as C
@@ -30,7 +28,7 @@ from lightning.pytorch.plugins.layer_sync import LayerSync
 from lightning.pytorch.plugins.precision.precision import Precision
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.strategies.strategy import Strategy
-from typing_extensions import TypedDict, TypeVar, override
+from typing_extensions import TypeAliasType, TypedDict, override
 
 from .._directory import DirectoryConfig
 from .._hf_hub import HuggingFaceHubConfig
@@ -43,6 +41,7 @@ from ..callbacks import (
     OnExceptionCheckpointCallbackConfig,
 )
 from ..callbacks.base import CallbackConfigBase
+from ..callbacks.checkpoint.time_checkpoint import TimeCheckpointCallbackConfig
 from ..callbacks.debug_flag import DebugFlagCallbackConfig
 from ..callbacks.log_epoch import LogEpochCallbackConfig
 from ..callbacks.lr_monitor import LearningRateMonitorConfig
@@ -72,71 +71,82 @@ class GradientClippingConfig(C.Config):
     """Norm type to use for gradient clipping."""
 
 
-TPlugin = TypeVar(
-    "TPlugin",
-    Precision,
-    ClusterEnvironment,
-    CheckpointIO,
-    LayerSync,
-    infer_variance=True,
+Plugin = TypeAliasType(
+    "Plugin", Precision | ClusterEnvironment | CheckpointIO | LayerSync
 )
 
 
-@runtime_checkable
-class PluginConfigProtocol(Protocol[TPlugin]):
-    def create_plugin(self) -> TPlugin: ...
+class PluginConfigBase(C.Config, ABC):
+    @abstractmethod
+    def create_plugin(self) -> Plugin: ...
 
 
-@runtime_checkable
-class AcceleratorConfigProtocol(Protocol):
+plugin_registry = C.Registry(PluginConfigBase, discriminator="name")
+
+
+class AcceleratorConfigBase(C.Config, ABC):
+    @abstractmethod
     def create_accelerator(self) -> Accelerator: ...
 
 
-@runtime_checkable
-class StrategyConfigProtocol(Protocol):
+accelerator_registry = C.Registry(AcceleratorConfigBase, discriminator="name")
+
+
+class StrategyConfigBase(C.Config, ABC):
+    @abstractmethod
     def create_strategy(self) -> Strategy: ...
 
 
-AcceleratorLiteral: TypeAlias = Literal[
-    "cpu", "gpu", "tpu", "ipu", "hpu", "mps", "auto"
-]
-
-StrategyLiteral: TypeAlias = Literal[
-    "auto",
-    "ddp",
-    "ddp_find_unused_parameters_false",
-    "ddp_find_unused_parameters_true",
-    "ddp_spawn",
-    "ddp_spawn_find_unused_parameters_false",
-    "ddp_spawn_find_unused_parameters_true",
-    "ddp_fork",
-    "ddp_fork_find_unused_parameters_false",
-    "ddp_fork_find_unused_parameters_true",
-    "ddp_notebook",
-    "dp",
-    "deepspeed",
-    "deepspeed_stage_1",
-    "deepspeed_stage_1_offload",
-    "deepspeed_stage_2",
-    "deepspeed_stage_2_offload",
-    "deepspeed_stage_3",
-    "deepspeed_stage_3_offload",
-    "deepspeed_stage_3_offload_nvme",
-    "fsdp",
-    "fsdp_cpu_offload",
-    "single_xla",
-    "xla_fsdp",
-    "xla",
-    "single_tpu",
-]
-
-
-CheckpointCallbackConfig: TypeAlias = Annotated[
-    BestCheckpointCallbackConfig
-    | LastCheckpointCallbackConfig
-    | OnExceptionCheckpointCallbackConfig,
-    C.Field(discriminator="name"),
-]
+strategy_registry = C.Registry(StrategyConfigBase, discriminator="name")
+
+
+AcceleratorLiteral = TypeAliasType(
+    "AcceleratorLiteral", Literal["cpu", "gpu", "tpu", "ipu", "hpu", "mps", "auto"]
+)
+
+StrategyLiteral = TypeAliasType(
+    "StrategyLiteral",
+    Literal[
+        "auto",
+        "ddp",
+        "ddp_find_unused_parameters_false",
+        "ddp_find_unused_parameters_true",
+        "ddp_spawn",
+        "ddp_spawn_find_unused_parameters_false",
+        "ddp_spawn_find_unused_parameters_true",
+        "ddp_fork",
+        "ddp_fork_find_unused_parameters_false",
+        "ddp_fork_find_unused_parameters_true",
+        "ddp_notebook",
+        "dp",
+        "deepspeed",
+        "deepspeed_stage_1",
+        "deepspeed_stage_1_offload",
+        "deepspeed_stage_2",
+        "deepspeed_stage_2_offload",
+        "deepspeed_stage_3",
+        "deepspeed_stage_3_offload",
+        "deepspeed_stage_3_offload_nvme",
+        "fsdp",
+        "fsdp_cpu_offload",
+        "single_xla",
+        "xla_fsdp",
+        "xla",
+        "single_tpu",
+    ],
+)
+
+
+CheckpointCallbackConfig = TypeAliasType(
+    "CheckpointCallbackConfig",
+    Annotated[
+        BestCheckpointCallbackConfig
+        | LastCheckpointCallbackConfig
+        | OnExceptionCheckpointCallbackConfig
+        | TimeCheckpointCallbackConfig,
+        C.Field(discriminator="name"),
+    ],
+)
 
 
 class CheckpointSavingConfig(CallbackConfigBase):
@@ -147,6 +157,7 @@ class CheckpointSavingConfig(CallbackConfigBase):
         BestCheckpointCallbackConfig(throw_on_no_metric=False),
         LastCheckpointCallbackConfig(),
         OnExceptionCheckpointCallbackConfig(),
+        TimeCheckpointCallbackConfig(interval=timedelta(hours=12)),
     ]
     """Checkpoint callback configurations."""
 
@@ -420,6 +431,9 @@ class SanityCheckingConfig(C.Config):
     """
 
 
+@plugin_registry.rebuild_on_registers
+@strategy_registry.rebuild_on_registers
+@accelerator_registry.rebuild_on_registers
 class TrainerConfig(C.Config):
     # region Active Run Configuration
     id: str = C.Field(default_factory=lambda: TrainerConfig.generate_id())
@@ -564,7 +578,9 @@ class TrainerConfig(C.Config):
     Default: ``False``.
     """
 
-    plugins: list[PluginConfigProtocol] | None = None
+    plugins: (
+        list[Annotated[PluginConfigBase, plugin_registry.DynamicResolution()]] | None
+    ) = None
     """
     Plugins allow modification of core behavior like ddp and amp, and enable custom lightning plugins.
         Default: ``None``.
@@ -724,13 +740,21 @@ class TrainerConfig(C.Config):
     Default: ``True``.
     """
 
-    accelerator: AcceleratorConfigProtocol | AcceleratorLiteral | None = None
+    accelerator: (
+        Annotated[AcceleratorConfigBase, accelerator_registry.DynamicResolution()]
+        | AcceleratorLiteral
+        | None
+    ) = None
     """Supports passing different accelerator types ("cpu", "gpu", "tpu", "ipu", "hpu", "mps", "auto")
     as well as custom accelerator instances.
     Default: ``"auto"``.
     """
 
-    strategy: StrategyConfigProtocol | StrategyLiteral | None = None
+    strategy: (
+        Annotated[StrategyConfigBase, strategy_registry.DynamicResolution()]
+        | StrategyLiteral
+        | None
+    ) = None
     """Supports different training strategies with aliases as well custom strategies.
     Default: ``"auto"``.
     """

nshtrainer/trainer/_runtime_callback.py

@@ -4,14 +4,14 @@ import datetime
 import logging
 import time
 from dataclasses import dataclass
-from typing import Any, Literal, TypeAlias
+from typing import Any, Literal
 
 from lightning.pytorch.callbacks.callback import Callback
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 log = logging.getLogger(__name__)
 
-Stage: TypeAlias = Literal["train", "validate", "test", "predict"]
+Stage = TypeAliasType("Stage", Literal["train", "validate", "test", "predict"])
 ALL_STAGES = ("train", "validate", "test", "predict")
 
 

nshtrainer/trainer/signal_connector.py

@@ -12,7 +12,7 @@ from collections import defaultdict
 from collections.abc import Callable
 from pathlib import Path
 from types import FrameType
-from typing import Any, TypeAlias
+from typing import Any
 
 import nshrunner as nr
 import torch.utils.data
@@ -22,12 +22,14 @@ from lightning.pytorch.trainer.connectors.signal_connector import _HandlersCompo
 from lightning.pytorch.trainer.connectors.signal_connector import (
     _SignalConnector as _LightningSignalConnector,
 )
-from typing_extensions import override
+from typing_extensions import TypeAliasType, override
 
 log = logging.getLogger(__name__)
 
-_SIGNUM = int | signal.Signals
-_HANDLER: TypeAlias = Callable[[_SIGNUM, FrameType], Any] | int | signal.Handlers | None
+_SIGNUM = TypeAliasType("_SIGNUM", int | signal.Signals)
+_HANDLER = TypeAliasType(
+    "_HANDLER", Callable[[_SIGNUM, FrameType], Any] | int | signal.Handlers | None
+)
 _IS_WINDOWS = platform.system() == "Windows"
 
 

nshtrainer/trainer/trainer.py

@@ -23,9 +23,9 @@ from ..callbacks.base import resolve_all_callbacks
 from ..util._environment_info import EnvironmentConfig
 from ..util.bf16 import is_bf16_supported_no_emulation
 from ._config import (
-    AcceleratorConfigProtocol,
+    AcceleratorConfigBase,
     LightningTrainerKwargs,
-    StrategyConfigProtocol,
+    StrategyConfigBase,
     TrainerConfig,
 )
 from ._runtime_callback import RuntimeTrackerCallback, Stage
@@ -171,12 +171,12 @@ class Trainer(LightningTrainer):
             _update_kwargs(use_distributed_sampler=use_distributed_sampler)
 
         if (accelerator := hparams.accelerator) is not None:
-            if isinstance(accelerator, AcceleratorConfigProtocol):
+            if isinstance(accelerator, AcceleratorConfigBase):
                 accelerator = accelerator.create_accelerator()
             _update_kwargs(accelerator=accelerator)
 
         if (strategy := hparams.strategy) is not None:
-            if isinstance(strategy, StrategyConfigProtocol):
+            if isinstance(strategy, StrategyConfigBase):
                 strategy = strategy.create_strategy()
             _update_kwargs(strategy=strategy)
 

nshtrainer/util/_useful_types.py

@@ -7,7 +7,14 @@ from collections.abc import Set as AbstractSet
 from os import PathLike
 from typing import Any, TypeVar, overload
 
-from typing_extensions import Buffer, Literal, Protocol, SupportsIndex, TypeAlias
+from typing_extensions import (
+    Buffer,
+    Literal,
+    Protocol,
+    SupportsIndex,
+    TypeAlias,
+    TypeAliasType,
+)
 
 _KT = TypeVar("_KT")
 _KT_co = TypeVar("_KT_co", covariant=True)
@@ -60,7 +67,9 @@ class SupportsAllComparisons(
 ): ...
 
 
-SupportsRichComparison: TypeAlias = SupportsDunderLT[Any] | SupportsDunderGT[Any]
+SupportsRichComparison = TypeAliasType(
+    "SupportsRichComparison", SupportsDunderLT[Any] | SupportsDunderGT[Any]
+)
 SupportsRichComparisonT = TypeVar(
     "SupportsRichComparisonT", bound=SupportsRichComparison
 )

nshtrainer/util/config/dtype.py

@@ -1,57 +1,60 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal, TypeAlias
+from typing import TYPE_CHECKING, Literal
 
 import nshconfig as C
 import torch
-from typing_extensions import assert_never
+from typing_extensions import TypeAliasType, assert_never
 
 from ..bf16 import is_bf16_supported_no_emulation
 
 if TYPE_CHECKING:
     from ...trainer._config import TrainerConfig
 
-DTypeName: TypeAlias = Literal[
-    "float32",
-    "float",
-    "float64",
-    "double",
-    "float16",
-    "bfloat16",
-    "float8_e4m3fn",
-    "float8_e4m3fnuz",
-    "float8_e5m2",
-    "float8_e5m2fnuz",
-    "half",
-    "uint8",
-    "uint16",
-    "uint32",
-    "uint64",
-    "int8",
-    "int16",
-    "short",
-    "int32",
-    "int",
-    "int64",
-    "long",
-    "complex32",
-    "complex64",
-    "chalf",
-    "cfloat",
-    "complex128",
-    "cdouble",
-    "quint8",
-    "qint8",
-    "qint32",
-    "bool",
-    "quint4x2",
-    "quint2x4",
-    "bits1x8",
-    "bits2x4",
-    "bits4x2",
-    "bits8",
-    "bits16",
-]
+DTypeName = TypeAliasType(
+    "DTypeName",
+    Literal[
+        "float32",
+        "float",
+        "float64",
+        "double",
+        "float16",
+        "bfloat16",
+        "float8_e4m3fn",
+        "float8_e4m3fnuz",
+        "float8_e5m2",
+        "float8_e5m2fnuz",
+        "half",
+        "uint8",
+        "uint16",
+        "uint32",
+        "uint64",
+        "int8",
+        "int16",
+        "short",
+        "int32",
+        "int",
+        "int64",
+        "long",
+        "complex32",
+        "complex64",
+        "chalf",
+        "cfloat",
+        "complex128",
+        "cdouble",
+        "quint8",
+        "qint8",
+        "qint32",
+        "bool",
+        "quint4x2",
+        "quint2x4",
+        "bits1x8",
+        "bits2x4",
+        "bits4x2",
+        "bits8",
+        "bits16",
+    ],
+)
 
 
 class DTypeConfig(C.Config):

nshtrainer/util/path.py

@@ -6,11 +6,12 @@ import os
 import platform
 import shutil
 from pathlib import Path
-from typing import TypeAlias
+
+from typing_extensions import TypeAliasType
 
 log = logging.getLogger(__name__)
 
-_Path: TypeAlias = str | Path | os.PathLike
+_Path = TypeAliasType("_Path", str | Path)
 
 
 def get_relative_path(source: _Path, destination: _Path):

{nshtrainer-1.0.0b25.dist-info → nshtrainer-1.0.0b27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 1.0.0b25
+Version: 1.0.0b27
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -9,6 +9,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Provides-Extra: extra
 Requires-Dist: GitPython ; extra == "extra"
 Requires-Dist: huggingface-hub ; extra == "extra"