nshtrainer 0.42.0__py3-none-any.whl → 0.43.0__py3-none-any.whl

nshtrainer/config/util/_environment_info/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -43,30 +45,30 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "EnvironmentPackageConfig":
-            return importlib.import_module(
-                "nshtrainer.util._environment_info"
-            ).EnvironmentPackageConfig
-        if name == "EnvironmentSnapshotConfig":
+        if name == "EnvironmentLinuxEnvironmentConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentSnapshotConfig
+            ).EnvironmentLinuxEnvironmentConfig
         if name == "EnvironmentLSFInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).EnvironmentLSFInformationConfig
-        if name == "EnvironmentLinuxEnvironmentConfig":
+        if name == "EnvironmentGPUConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentLinuxEnvironmentConfig
-        if name == "EnvironmentSLURMInformationConfig":
+            ).EnvironmentGPUConfig
+        if name == "EnvironmentPackageConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentSLURMInformationConfig
-        if name == "EnvironmentConfig":
+            ).EnvironmentPackageConfig
+        if name == "EnvironmentHardwareConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentConfig
+            ).EnvironmentHardwareConfig
+        if name == "EnvironmentSnapshotConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentSnapshotConfig
         if name == "EnvironmentClassInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
@@ -75,18 +77,18 @@ else:
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).GitRepositoryConfig
-        if name == "EnvironmentCUDAConfig":
+        if name == "EnvironmentConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentCUDAConfig
-        if name == "EnvironmentGPUConfig":
+            ).EnvironmentConfig
+        if name == "EnvironmentCUDAConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentGPUConfig
-        if name == "EnvironmentHardwareConfig":
+            ).EnvironmentCUDAConfig
+        if name == "EnvironmentSLURMInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentHardwareConfig
+            ).EnvironmentSLURMInformationConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/util/config/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING

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

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING

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

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING

nshtrainer/data/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from . import transform as dataset_transform
 from .balanced_batch_sampler import BalancedBatchSampler as BalancedBatchSampler
 from .datamodule import LightningDataModuleBase as LightningDataModuleBase

nshtrainer/data/balanced_batch_sampler.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import heapq
 import logging
 from typing import Any, Protocol, runtime_checkable

nshtrainer/data/datamodule.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from lightning.pytorch import LightningDataModule
 
 

nshtrainer/data/transform.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 from collections.abc import Callable
 from typing import Any, cast

nshtrainer/ll/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import TypeAlias
 
 from . import _experimental as _experimental

nshtrainer/ll/_experimental.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer._experimental import *  # noqa: F403

nshtrainer/ll/actsave.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshutils.actsave import *  # type: ignore # noqa: F403
 
 from nshtrainer.callbacks.actsave import ActSaveCallback as ActSaveCallback

nshtrainer/ll/callbacks.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.callbacks import *  # noqa: F403

nshtrainer/ll/config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshconfig import *  # type: ignore  # noqa: F403
 from nshconfig import Config as TypedConfig  # type: ignore  # noqa: F401
 

nshtrainer/ll/data.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.data import *  # noqa: F403

nshtrainer/ll/log.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshutils import init_python_logging as init_python_logging
 from nshutils import lovely as lovely
 from nshutils import pretty as pretty

nshtrainer/ll/lr_scheduler.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.lr_scheduler import *  # noqa: F403

nshtrainer/ll/model.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshtrainer.model import *  # noqa: F403
 
 from ..trainer._config import CheckpointLoadingConfig as CheckpointLoadingConfig

nshtrainer/ll/nn.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.nn import *  # noqa: F403

nshtrainer/ll/optimizer.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.optimizer import *  # noqa: F403

nshtrainer/ll/runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshrunner import SnapshotConfig as SnapshotConfig
 
 from nshtrainer.runner import *  # type: ignore # noqa: F403

nshtrainer/ll/snapshot.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshsnap import *  # pyright: ignore[reportWildcardImportFromLibrary] # noqa: F403

nshtrainer/ll/snoop.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshutils import snoop as snoop

nshtrainer/ll/trainer.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.trainer import *  # noqa: F403

nshtrainer/ll/typecheck.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshutils.typecheck import *  # type: ignore # noqa: F403

nshtrainer/ll/util.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.util import *  # noqa: F403

nshtrainer/loggers/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Annotated, TypeAlias
 
 import nshconfig as C

nshtrainer/loggers/_base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 

nshtrainer/loggers/csv.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Literal
 
 from typing_extensions import override

nshtrainer/loggers/tensorboard.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 

nshtrainer/loggers/wandb.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import importlib.metadata
 import logging
 from typing import TYPE_CHECKING, Literal
@@ -8,8 +10,8 @@ from packaging import version
 from typing_extensions import assert_never, override
 
 from ..callbacks.base import CallbackConfigBase
-from ..callbacks.wandb_upload_code import WandbUploadCodeConfig
-from ..callbacks.wandb_watch import WandbWatchConfig
+from ..callbacks.wandb_upload_code import WandbUploadCodeCallbackConfig
+from ..callbacks.wandb_watch import WandbWatchCallbackConfig
 from ._base import BaseLoggerConfig
 
 if TYPE_CHECKING:
@@ -92,10 +94,10 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
     - "none" or False: Do not log any checkpoints
     """
 
-    log_code: WandbUploadCodeConfig | None = WandbUploadCodeConfig()
+    log_code: WandbUploadCodeCallbackConfig | None = WandbUploadCodeCallbackConfig()
     """WandB code upload configuration. Used to upload code to WandB."""
 
-    watch: WandbWatchConfig | None = WandbWatchConfig()
+    watch: WandbWatchCallbackConfig | None = WandbWatchCallbackConfig()
     """WandB model watch configuration. Used to log model architecture, gradients, and parameters."""
 
     offline: bool = False

nshtrainer/lr_scheduler/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Annotated, TypeAlias
 
 import nshconfig as C

nshtrainer/lr_scheduler/_base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 from abc import ABC, abstractmethod
 from collections.abc import Mapping

nshtrainer/lr_scheduler/linear_warmup_cosine.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 import warnings
 from typing import Literal

nshtrainer/lr_scheduler/reduce_lr_on_plateau.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import TYPE_CHECKING, Literal, cast
 
 from lightning.pytorch.utilities.types import LRSchedulerConfigType

nshtrainer/metrics/__init__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from ._config import MetricConfig as MetricConfig

nshtrainer/metrics/_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import builtins
 from typing import Literal
 

nshtrainer/model/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .base import LightningModuleBase as LightningModuleBase
 from .config import BaseConfig as BaseConfig
 from .config import DirectoryConfig as DirectoryConfig

nshtrainer/model/base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import inspect
 import logging
 from abc import ABC, abstractmethod

nshtrainer/model/config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import logging
 import os

nshtrainer/model/mixins/callback.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import Callable, Iterable, Sequence
 from typing import Any, TypeAlias, cast, final

nshtrainer/model/mixins/logger.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections import deque
 from collections.abc import Callable, Generator
 from contextlib import contextmanager

nshtrainer/nn/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .mlp import MLP as MLP
 from .mlp import MLPConfig as MLPConfig
 from .mlp import MLPConfigDict as MLPConfigDict

nshtrainer/nn/mlp.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 from collections.abc import Callable, Sequence
 from typing import Literal, Protocol, runtime_checkable

nshtrainer/nn/module_dict.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import Iterable, Mapping
 from typing import Generic, cast
 

nshtrainer/nn/module_list.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import Iterable, Iterator
 from typing import Generic, TypeVar, overload
 

nshtrainer/nn/nonlinearity.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import Annotated, Literal
 

nshtrainer/optimizer.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from typing import Annotated, Any, Literal, TypeAlias

nshtrainer/profiler/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Annotated, TypeAlias
 
 import nshconfig as C

nshtrainer/profiler/_base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path

nshtrainer/profiler/advanced.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 

nshtrainer/profiler/pytorch.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Any, Literal
 

nshtrainer/profiler/simple.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 

nshtrainer/runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import logging
 from collections.abc import Callable, Iterable, Mapping, Sequence

nshtrainer/scripts/find_packages.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import argparse
 import ast
 import glob

nshtrainer/trainer/__init__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from .trainer import Trainer as Trainer

nshtrainer/trainer/_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import Iterable, Sequence
 from datetime import timedelta
@@ -30,14 +32,15 @@ from .._hf_hub import HuggingFaceHubConfig
 from ..callbacks import (
     BestCheckpointCallbackConfig,
     CallbackConfig,
-    EarlyStoppingConfig,
+    EarlyStoppingCallbackConfig,
     LastCheckpointCallbackConfig,
     OnExceptionCheckpointCallbackConfig,
 )
 from ..callbacks.base import CallbackConfigBase
 from ..callbacks.debug_flag import DebugFlagCallbackConfig
-from ..callbacks.rlp_sanity_checks import RLPSanityChecksConfig
-from ..callbacks.shared_parameters import SharedParametersConfig
+from ..callbacks.log_epoch import LogEpochCallbackConfig
+from ..callbacks.rlp_sanity_checks import RLPSanityChecksCallbackConfig
+from ..callbacks.shared_parameters import SharedParametersCallbackConfig
 from ..loggers import (
     CSVLoggerConfig,
     LoggerConfig,
@@ -65,7 +68,7 @@ class LoggingConfig(CallbackConfigBase):
 
     log_lr: bool | Literal["step", "epoch"] = True
     """If enabled, will register a `LearningRateMonitor` callback to log the learning rate to the logger."""
-    log_epoch: bool = True
+    log_epoch: LogEpochCallbackConfig | None = LogEpochCallbackConfig()
     """If enabled, will log the fractional epoch number to the logger."""
 
     actsave_logged_metrics: bool = False
@@ -136,9 +139,7 @@ class LoggingConfig(CallbackConfigBase):
             yield LearningRateMonitor(logging_interval=logging_interval)
 
         if self.log_epoch:
-            from ..callbacks.log_epoch import LogEpochCallback
-
-            yield LogEpochCallback()
+            yield from self.log_epoch.create_callbacks(root_config)
 
         for logger in self.loggers:
             if not logger or not isinstance(logger, CallbackConfigBase):
@@ -172,9 +173,9 @@ class OptimizationConfig(CallbackConfigBase):
 
     @override
     def create_callbacks(self, root_config):
-        from ..callbacks.norm_logging import NormLoggingConfig
+        from ..callbacks.norm_logging import NormLoggingCallbackConfig
 
-        yield from NormLoggingConfig(
+        yield from NormLoggingCallbackConfig(
             log_grad_norm=self.log_grad_norm,
             log_grad_norm_per_param=self.log_grad_norm_per_param,
             log_param_norm=self.log_param_norm,
@@ -564,8 +565,8 @@ class TrainerConfig(C.Config):
     reproducibility: ReproducibilityConfig = ReproducibilityConfig()
     """Reproducibility configuration options."""
 
-    reduce_lr_on_plateau_sanity_checking: RLPSanityChecksConfig | None = (
-        RLPSanityChecksConfig()
+    reduce_lr_on_plateau_sanity_checking: RLPSanityChecksCallbackConfig | None = (
+        RLPSanityChecksCallbackConfig()
     )
     """
     If enabled, will do some sanity checks if the `ReduceLROnPlateau` scheduler is used:
@@ -573,7 +574,7 @@ class TrainerConfig(C.Config):
         - If the `interval` is epoch, it makes sure that validation is called every `frequency` epochs.
     """
 
-    early_stopping: EarlyStoppingConfig | None = None
+    early_stopping: EarlyStoppingCallbackConfig | None = None
     """Early stopping configuration options."""
 
     profiler: ProfilerConfig | None = None
@@ -741,7 +742,9 @@ class TrainerConfig(C.Config):
     automatic selection based on the chosen accelerator. Default: ``"auto"``.
     """
 
-    shared_parameters: SharedParametersConfig | None = SharedParametersConfig()
+    shared_parameters: SharedParametersCallbackConfig | None = (
+        SharedParametersCallbackConfig()
+    )
     """If enabled, the model supports scaling the gradients of shared parameters that
     are registered in the self.shared_parameters list. This is useful for models that
     share parameters across multiple modules (e.g., in a GPT model) and want to

nshtrainer/trainer/_runtime_callback.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import datetime
 import logging
 import time

nshtrainer/trainer/checkpoint_connector.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from pathlib import Path
 from typing import TYPE_CHECKING, cast

nshtrainer/trainer/signal_connector.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import os
 import platform

nshtrainer/trainer/trainer.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import os
 from collections.abc import Sequence

nshtrainer/util/_environment_info.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import getpass
 import importlib.metadata
 import inspect

nshtrainer/util/bf16.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import torch
 
 

nshtrainer/util/config/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from . import duration as duration
 from .dtype import DTypeConfig as DTypeConfig
 from .duration import DurationConfig as DurationConfig

nshtrainer/util/config/dtype.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import TYPE_CHECKING, Literal, TypeAlias
 
 import nshconfig as C

nshtrainer/util/config/duration.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 from typing import Annotated, Literal
 

nshtrainer/util/environment.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import os
 from contextlib import contextmanager

nshtrainer/util/path.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import hashlib
 import logging
 import os

nshtrainer/util/seed.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 
 import lightning.fabric.utilities.seed as LS

nshtrainer/util/slurm.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 class SlurmParseException(Exception):
     pass
 

nshtrainer/util/typed.py

@@ -1,2 +1,4 @@
+from __future__ import annotations
+
 from ..nn.module_dict import TypedModuleDict as TypedModuleDict
 from ..nn.module_list import TypedModuleList as TypedModuleList