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

nshtrainer/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from . import _experimental as _experimental
 from . import callbacks as callbacks
 from . import config as config

nshtrainer/_callback.py

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

nshtrainer/_checkpoint/loader.py

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

nshtrainer/_checkpoint/metadata.py

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

nshtrainer/_checkpoint/saver.py

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

nshtrainer/_directory.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import logging
 from pathlib import Path
 
 import nshconfig as C
 
-from .callbacks.directory_setup import DirectorySetupConfig
+from .callbacks.directory_setup import DirectorySetupCallbackConfig
 from .loggers import LoggerConfig
 
 log = logging.getLogger(__name__)
@@ -32,7 +34,7 @@ class DirectoryConfig(C.Config):
     profile: Path | None = None
     """Directory to save profiling information to. If None, will use nshtrainer/{id}/profile/."""
 
-    setup_callback: DirectorySetupConfig = DirectorySetupConfig()
+    setup_callback: DirectorySetupCallbackConfig = DirectorySetupCallbackConfig()
     """Configuration for the directory setup PyTorch Lightning callback."""
 
     def resolve_run_root_directory(self, run_id: str) -> Path:

nshtrainer/_experimental/__init__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from lightning.fabric.utilities.throughput import measure_flops as measure_flops

nshtrainer/_hf_hub.py

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

nshtrainer/callbacks/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Annotated
 
 import nshconfig as C
@@ -6,60 +8,74 @@ from . import checkpoint as checkpoint
 from .base import CallbackConfigBase as CallbackConfigBase
 from .checkpoint import BestCheckpoint as BestCheckpoint
 from .checkpoint import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
-from .checkpoint import LastCheckpoint as LastCheckpoint
+from .checkpoint import LastCheckpointCallback as LastCheckpointCallback
 from .checkpoint import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
-from .checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
+from .checkpoint import OnExceptionCheckpointCallback as OnExceptionCheckpointCallback
 from .checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
 from .debug_flag import DebugFlagCallback as DebugFlagCallback
 from .debug_flag import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from .directory_setup import DirectorySetupCallback as DirectorySetupCallback
-from .directory_setup import DirectorySetupConfig as DirectorySetupConfig
-from .early_stopping import EarlyStopping as EarlyStopping
-from .early_stopping import EarlyStoppingConfig as EarlyStoppingConfig
-from .ema import EMA as EMA
-from .ema import EMAConfig as EMAConfig
+from .directory_setup import (
+    DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
+)
+from .early_stopping import EarlyStoppingCallback as EarlyStoppingCallback
+from .early_stopping import EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig
+from .ema import EMACallback as EMACallback
+from .ema import EMACallbackConfig as EMACallbackConfig
 from .finite_checks import FiniteChecksCallback as FiniteChecksCallback
-from .finite_checks import FiniteChecksConfig as FiniteChecksConfig
-from .gradient_skipping import GradientSkipping as GradientSkipping
-from .gradient_skipping import GradientSkippingConfig as GradientSkippingConfig
+from .finite_checks import FiniteChecksCallbackConfig as FiniteChecksCallbackConfig
+from .gradient_skipping import GradientSkippingCallback as GradientSkippingCallback
+from .gradient_skipping import (
+    GradientSkippingCallbackConfig as GradientSkippingCallbackConfig,
+)
 from .interval import EpochIntervalCallback as EpochIntervalCallback
 from .interval import IntervalCallback as IntervalCallback
 from .interval import StepIntervalCallback as StepIntervalCallback
 from .log_epoch import LogEpochCallback as LogEpochCallback
+from .log_epoch import LogEpochCallbackConfig as LogEpochCallbackConfig
 from .norm_logging import NormLoggingCallback as NormLoggingCallback
-from .norm_logging import NormLoggingConfig as NormLoggingConfig
+from .norm_logging import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from .print_table import PrintTableMetricsCallback as PrintTableMetricsCallback
-from .print_table import PrintTableMetricsConfig as PrintTableMetricsConfig
+from .print_table import (
+    PrintTableMetricsCallbackConfig as PrintTableMetricsCallbackConfig,
+)
 from .rlp_sanity_checks import RLPSanityChecksCallback as RLPSanityChecksCallback
-from .rlp_sanity_checks import RLPSanityChecksConfig as RLPSanityChecksConfig
+from .rlp_sanity_checks import (
+    RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
+)
 from .shared_parameters import SharedParametersCallback as SharedParametersCallback
-from .shared_parameters import SharedParametersConfig as SharedParametersConfig
+from .shared_parameters import (
+    SharedParametersCallbackConfig as SharedParametersCallbackConfig,
+)
 from .throughput_monitor import ThroughputMonitorConfig as ThroughputMonitorConfig
-from .timer import EpochTimer as EpochTimer
-from .timer import EpochTimerConfig as EpochTimerConfig
+from .timer import EpochTimerCallback as EpochTimerCallback
+from .timer import EpochTimerCallbackConfig as EpochTimerCallbackConfig
 from .wandb_upload_code import WandbUploadCodeCallback as WandbUploadCodeCallback
-from .wandb_upload_code import WandbUploadCodeConfig as WandbUploadCodeConfig
+from .wandb_upload_code import (
+    WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
+)
 from .wandb_watch import WandbWatchCallback as WandbWatchCallback
-from .wandb_watch import WandbWatchConfig as WandbWatchConfig
+from .wandb_watch import WandbWatchCallbackConfig as WandbWatchCallbackConfig
 
 CallbackConfig = Annotated[
     DebugFlagCallbackConfig
-    | EarlyStoppingConfig
+    | EarlyStoppingCallbackConfig
     | ThroughputMonitorConfig
-    | EpochTimerConfig
-    | PrintTableMetricsConfig
-    | FiniteChecksConfig
-    | NormLoggingConfig
-    | GradientSkippingConfig
-    | EMAConfig
+    | EpochTimerCallbackConfig
+    | PrintTableMetricsCallbackConfig
+    | FiniteChecksCallbackConfig
+    | NormLoggingCallbackConfig
+    | GradientSkippingCallbackConfig
+    | LogEpochCallbackConfig
+    | EMACallbackConfig
     | BestCheckpointCallbackConfig
     | LastCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig
-    | SharedParametersConfig
-    | RLPSanityChecksConfig
-    | WandbWatchConfig
-    | WandbUploadCodeConfig,
+    | SharedParametersCallbackConfig
+    | RLPSanityChecksCallbackConfig
+    | WandbWatchCallbackConfig
+    | WandbUploadCodeCallbackConfig,
     C.Field(discriminator="name"),
 ]

nshtrainer/callbacks/_throughput_monitor_callback.py

@@ -12,6 +12,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
 import time
 from collections import deque
 from typing import (

nshtrainer/callbacks/actsave.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 from pathlib import Path
 from typing import Literal

nshtrainer/callbacks/base.py

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

nshtrainer/callbacks/checkpoint/__init__.py

@@ -1,12 +1,16 @@
+from __future__ import annotations
+
 from .best_checkpoint import BestCheckpoint as BestCheckpoint
 from .best_checkpoint import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
-from .last_checkpoint import LastCheckpoint as LastCheckpoint
+from .last_checkpoint import LastCheckpointCallback as LastCheckpointCallback
 from .last_checkpoint import (
     LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
 )
-from .on_exception_checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
+from .on_exception_checkpoint import (
+    OnExceptionCheckpointCallback as OnExceptionCheckpointCallback,
+)
 from .on_exception_checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )

nshtrainer/callbacks/checkpoint/_base.py

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

nshtrainer/callbacks/checkpoint/best_checkpoint.py

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

nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 
@@ -17,11 +19,11 @@ class LastCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
 
     @override
     def create_checkpoint(self, root_config, dirpath):
-        return LastCheckpoint(self, dirpath)
+        return LastCheckpointCallback(self, dirpath)
 
 
 @final
-class LastCheckpoint(CheckpointBase[LastCheckpointCallbackConfig]):
+class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
     @override
     def name(self):
         return "last"

nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import datetime
 import logging
@@ -59,10 +61,12 @@ class OnExceptionCheckpointCallbackConfig(CallbackConfigBase):
 
         if not (filename := self.filename):
             filename = f"on_exception_{root_config.id}"
-        yield OnExceptionCheckpoint(self, dirpath=Path(dirpath), filename=filename)
+        yield OnExceptionCheckpointCallback(
+            self, dirpath=Path(dirpath), filename=filename
+        )
 
 
-class OnExceptionCheckpoint(_OnExceptionCheckpoint):
+class OnExceptionCheckpointCallback(_OnExceptionCheckpoint):
     @override
     def __init__(
         self,

nshtrainer/callbacks/debug_flag.py

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

nshtrainer/callbacks/directory_setup.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import os
 from pathlib import Path
@@ -41,7 +43,7 @@ def _create_symlink_to_nshrunner(base_dir: Path):
     symlink_path.symlink_to(session_dir)
 
 
-class DirectorySetupConfig(CallbackConfigBase):
+class DirectorySetupCallbackConfig(CallbackConfigBase):
     name: Literal["directory_setup"] = "directory_setup"
 
     enabled: bool = True
@@ -62,7 +64,7 @@ class DirectorySetupConfig(CallbackConfigBase):
 
 class DirectorySetupCallback(Callback):
     @override
-    def __init__(self, config: DirectorySetupConfig):
+    def __init__(self, config: DirectorySetupCallbackConfig):
         super().__init__()
 
         self.config = config

nshtrainer/callbacks/early_stopping.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import math
 from typing import Literal
@@ -14,7 +16,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class EarlyStoppingConfig(CallbackConfigBase):
+class EarlyStoppingCallbackConfig(CallbackConfigBase):
     name: Literal["early_stopping"] = "early_stopping"
 
     metric: MetricConfig | None = None
@@ -54,11 +56,11 @@ class EarlyStoppingConfig(CallbackConfigBase):
                 "Either `metric` or `root_config.primary_metric` must be set to use EarlyStopping."
             )
 
-        yield EarlyStopping(self, metric)
+        yield EarlyStoppingCallback(self, metric)
 
 
-class EarlyStopping(_EarlyStopping):
-    def __init__(self, config: EarlyStoppingConfig, metric: MetricConfig):
+class EarlyStoppingCallback(_EarlyStopping):
+    def __init__(self, config: EarlyStoppingCallbackConfig, metric: MetricConfig):
         self.config = config
         self.metric = metric
         del config, metric

nshtrainer/callbacks/ema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import copy
 import threading
@@ -13,7 +15,7 @@ from typing_extensions import override
 from .base import CallbackConfigBase
 
 
-class EMA(Callback):
+class EMACallback(Callback):
     """
     Implements Exponential Moving Averaging (EMA).
 
@@ -358,7 +360,7 @@ class EMAOptimizer(torch.optim.Optimizer):
         self.rebuild_ema_params = True
 
 
-class EMAConfig(CallbackConfigBase):
+class EMACallbackConfig(CallbackConfigBase):
     name: Literal["ema"] = "ema"
 
     decay: float
@@ -375,7 +377,7 @@ class EMAConfig(CallbackConfigBase):
 
     @override
     def create_callbacks(self, root_config):
-        yield EMA(
+        yield EMACallback(
             decay=self.decay,
             validate_original_weights=self.validate_original_weights,
             every_n_steps=self.every_n_steps,

nshtrainer/callbacks/finite_checks.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 
@@ -58,7 +60,7 @@ class FiniteChecksCallback(Callback):
         )
 
 
-class FiniteChecksConfig(CallbackConfigBase):
+class FiniteChecksCallbackConfig(CallbackConfigBase):
     name: Literal["finite_checks"] = "finite_checks"
 
     nonfinite_grads: bool = True

nshtrainer/callbacks/gradient_skipping.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Any, Literal, Protocol, runtime_checkable
 
@@ -18,8 +20,8 @@ class HasGradSkippedSteps(Protocol):
     grad_skipped_steps: Any
 
 
-class GradientSkipping(Callback):
-    def __init__(self, config: "GradientSkippingConfig"):
+class GradientSkippingCallback(Callback):
+    def __init__(self, config: "GradientSkippingCallbackConfig"):
         super().__init__()
         self.config = config
 
@@ -73,7 +75,7 @@ class GradientSkipping(Callback):
         )
 
 
-class GradientSkippingConfig(CallbackConfigBase):
+class GradientSkippingCallbackConfig(CallbackConfigBase):
     name: Literal["gradient_skipping"] = "gradient_skipping"
 
     threshold: float
@@ -94,4 +96,4 @@ class GradientSkippingConfig(CallbackConfigBase):
 
     @override
     def create_callbacks(self, root_config):
-        yield GradientSkipping(self)
+        yield GradientSkippingCallback(self)

nshtrainer/callbacks/interval.py

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

nshtrainer/callbacks/log_epoch.py

@@ -1,14 +1,26 @@
+from __future__ import annotations
+
 import logging
 import math
-from typing import Any
+from typing import Any, Literal
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
 from typing_extensions import override
 
+from .base import CallbackConfigBase
+
 log = logging.getLogger(__name__)
 
 
+class LogEpochCallbackConfig(CallbackConfigBase):
+    name: Literal["log_epoch"] = "log_epoch"
+
+    @override
+    def create_callbacks(self, root_config):
+        yield LogEpochCallback()
+
+
 class LogEpochCallback(Callback):
     def __init__(self, metric_name: str = "computed_epoch"):
         super().__init__()

nshtrainer/callbacks/norm_logging.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal, cast
 
@@ -96,7 +98,7 @@ def compute_norm(
 
 
 class NormLoggingCallback(Callback):
-    def __init__(self, config: "NormLoggingConfig"):
+    def __init__(self, config: "NormLoggingCallbackConfig"):
         super().__init__()
 
         self.config = config
@@ -155,7 +157,7 @@ class NormLoggingCallback(Callback):
                 )
 
 
-class NormLoggingConfig(CallbackConfigBase):
+class NormLoggingCallbackConfig(CallbackConfigBase):
     name: Literal["norm_logging"] = "norm_logging"
 
     log_grad_norm: bool | str | float = False

nshtrainer/callbacks/print_table.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import fnmatch
 import importlib.util
@@ -74,7 +76,7 @@ class PrintTableMetricsCallback(Callback):
         return table
 
 
-class PrintTableMetricsConfig(CallbackConfigBase):
+class PrintTableMetricsCallbackConfig(CallbackConfigBase):
     """Configuration class for PrintTableMetricsCallback."""
 
     name: Literal["print_table_metrics"] = "print_table_metrics"

nshtrainer/callbacks/rlp_sanity_checks.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import Mapping
 from typing import Literal, cast
@@ -16,7 +18,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class RLPSanityChecksConfig(CallbackConfigBase):
+class RLPSanityChecksCallbackConfig(CallbackConfigBase):
     """
     If enabled, will do some sanity checks if the `ReduceLROnPlateau` scheduler is used:
         - If the ``interval`` is step, it makes sure that validation is called every ``frequency`` steps.
@@ -43,7 +45,7 @@ class RLPSanityChecksConfig(CallbackConfigBase):
 
 class RLPSanityChecksCallback(Callback):
     @override
-    def __init__(self, config: RLPSanityChecksConfig):
+    def __init__(self, config: RLPSanityChecksCallbackConfig):
         super().__init__()
 
         self.config = config

nshtrainer/callbacks/shared_parameters.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import Iterable
 from typing import Literal, Protocol, TypeAlias, runtime_checkable
@@ -17,7 +19,7 @@ def _parameters_to_names(parameters: Iterable[nn.Parameter], model: nn.Module):
     return [mapping[id(p)] for p in parameters]
 
 
-class SharedParametersConfig(CallbackConfigBase):
+class SharedParametersCallbackConfig(CallbackConfigBase):
     """A callback that allows scaling the gradients of shared parameters that
     are registered in the ``self.shared_parameters`` list of the root module.
 
@@ -43,7 +45,7 @@ class ModuleWithSharedParameters(Protocol):
 
 class SharedParametersCallback(Callback):
     @override
-    def __init__(self, config: SharedParametersConfig):
+    def __init__(self, config: SharedParametersCallbackConfig):
         super().__init__()
 
         self.config = config

nshtrainer/callbacks/throughput_monitor.py

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

nshtrainer/callbacks/timer.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import time
 from typing import Any, Literal
@@ -12,7 +14,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class EpochTimer(Callback):
+class EpochTimerCallback(Callback):
     def __init__(self):
         super().__init__()
 
@@ -149,9 +151,9 @@ class EpochTimer(Callback):
         self._total_batches = state_dict["total_batches"]
 
 
-class EpochTimerConfig(CallbackConfigBase):
+class EpochTimerCallbackConfig(CallbackConfigBase):
     name: Literal["epoch_timer"] = "epoch_timer"
 
     @override
     def create_callbacks(self, root_config):
-        yield EpochTimer()
+        yield EpochTimerCallback()

nshtrainer/callbacks/wandb_upload_code.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import os
 from pathlib import Path
@@ -14,7 +16,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class WandbUploadCodeConfig(CallbackConfigBase):
+class WandbUploadCodeCallbackConfig(CallbackConfigBase):
     name: Literal["wandb_upload_code"] = "wandb_upload_code"
 
     enabled: bool = True
@@ -32,7 +34,7 @@ class WandbUploadCodeConfig(CallbackConfigBase):
 
 
 class WandbUploadCodeCallback(Callback):
-    def __init__(self, config: WandbUploadCodeConfig):
+    def __init__(self, config: WandbUploadCodeCallbackConfig):
         super().__init__()
 
         self.config = config

nshtrainer/callbacks/wandb_watch.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal, Protocol, cast, runtime_checkable
 
@@ -12,7 +14,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class WandbWatchConfig(CallbackConfigBase):
+class WandbWatchCallbackConfig(CallbackConfigBase):
     name: Literal["wandb_watch"] = "wandb_watch"
 
     enabled: bool = True
@@ -41,7 +43,7 @@ class _HasWandbLogModuleProtocol(Protocol):
 
 
 class WandbWatchCallback(Callback):
-    def __init__(self, config: WandbWatchConfig):
+    def __init__(self, config: WandbWatchCallbackConfig):
         super().__init__()
 
         self.config = config