nshtrainer 1.0.0b33__tar.gz → 1.0.0b37__tar.gz

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/PKG-INFO

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

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/pyproject.toml

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

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/__init__.py

@@ -15,6 +15,7 @@ from .model import LightningModuleBase as LightningModuleBase
 from .trainer import Trainer as Trainer
 from .trainer import TrainerConfig as TrainerConfig
 from .trainer import accelerator_registry as accelerator_registry
+from .trainer import callback_registry as callback_registry
 from .trainer import plugin_registry as plugin_registry
 
 try:

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/_directory.py

@@ -81,7 +81,9 @@ class DirectoryConfig(C.Config):
 
         # Save to nshtrainer/{id}/log/{logger name}
         log_dir = self.resolve_subdirectory(run_id, "log")
-        log_dir = log_dir / logger.name
+        log_dir = log_dir / getattr(logger, "name")
+        # ^ NOTE: Logger must have a `name` attribute, as this is
+        # the discriminator for the logger registry
         log_dir.mkdir(exist_ok=True)
 
         return log_dir

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/_hf_hub.py

@@ -14,7 +14,11 @@ from nshrunner._env import SNAPSHOT_DIR
 from typing_extensions import assert_never, override
 
 from ._callback import NTCallbackBase
-from .callbacks.base import CallbackConfigBase, CallbackMetadataConfig
+from .callbacks.base import (
+    CallbackConfigBase,
+    CallbackMetadataConfig,
+    callback_registry,
+)
 
 if TYPE_CHECKING:
     from huggingface_hub import HfApi  # noqa: F401
@@ -39,9 +43,12 @@ class HuggingFaceHubAutoCreateConfig(C.Config):
         return self.enabled
 
 
+@callback_registry.register
 class HuggingFaceHubConfig(CallbackConfigBase):
     """Configuration options for Hugging Face Hub integration."""
 
+    name: Literal["hf_hub"] = "hf_hub"
+
     metadata: ClassVar[CallbackMetadataConfig] = {"ignore_if_exists": True}
 
     enabled: bool = False

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/__init__.py

@@ -2,10 +2,13 @@ from __future__ import annotations
 
 from typing import Annotated
 
-import nshconfig as C
+from typing_extensions import TypeAliasType
 
 from . import checkpoint as checkpoint
+from .actsave import ActSaveCallback as ActSaveCallback
+from .actsave import ActSaveConfig as ActSaveConfig
 from .base import CallbackConfigBase as CallbackConfigBase
+from .base import callback_registry as callback_registry
 from .checkpoint import BestCheckpointCallback as BestCheckpointCallback
 from .checkpoint import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
 from .checkpoint import LastCheckpointCallback as LastCheckpointCallback
@@ -14,8 +17,6 @@ from .checkpoint import OnExceptionCheckpointCallback as OnExceptionCheckpointCa
 from .checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
-from .checkpoint import TimeCheckpointCallback as TimeCheckpointCallback
-from .checkpoint import TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig
 from .debug_flag import DebugFlagCallback as DebugFlagCallback
 from .debug_flag import DebugFlagCallbackConfig as DebugFlagCallbackConfig
 from .directory_setup import DirectorySetupCallback as DirectorySetupCallback
@@ -37,6 +38,8 @@ 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 .lr_monitor import LearningRateMonitor as LearningRateMonitor
+from .lr_monitor import LearningRateMonitorConfig as LearningRateMonitorConfig
 from .norm_logging import NormLoggingCallback as NormLoggingCallback
 from .norm_logging import NormLoggingCallbackConfig as NormLoggingCallbackConfig
 from .print_table import PrintTableMetricsCallback as PrintTableMetricsCallback
@@ -60,23 +63,7 @@ from .wandb_upload_code import (
 from .wandb_watch import WandbWatchCallback as WandbWatchCallback
 from .wandb_watch import WandbWatchCallbackConfig as WandbWatchCallbackConfig
 
-CallbackConfig = Annotated[
-    DebugFlagCallbackConfig
-    | EarlyStoppingCallbackConfig
-    | EpochTimerCallbackConfig
-    | PrintTableMetricsCallbackConfig
-    | FiniteChecksCallbackConfig
-    | NormLoggingCallbackConfig
-    | GradientSkippingCallbackConfig
-    | LogEpochCallbackConfig
-    | EMACallbackConfig
-    | BestCheckpointCallbackConfig
-    | LastCheckpointCallbackConfig
-    | OnExceptionCheckpointCallbackConfig
-    | TimeCheckpointCallbackConfig
-    | SharedParametersCallbackConfig
-    | RLPSanityChecksCallbackConfig
-    | WandbWatchCallbackConfig
-    | WandbUploadCodeCallbackConfig,
-    C.Field(discriminator="name"),
-]
+CallbackConfig = TypeAliasType(
+    "CallbackConfig",
+    Annotated[CallbackConfigBase, callback_registry.DynamicResolution()],
+)

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/actsave.py

@@ -4,15 +4,19 @@ import contextlib
 from pathlib import Path
 from typing import Literal
 
-from typing_extensions import TypeAliasType, override
+from typing_extensions import TypeAliasType, final, override
 
 from .._callback import NTCallbackBase
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 Stage = TypeAliasType("Stage", Literal["train", "validation", "test", "predict"])
 
 
+@final
+@callback_registry.register
 class ActSaveConfig(CallbackConfigBase):
+    name: Literal["act_save"] = "act_save"
+
     enabled: bool = True
     """Enable activation saving."""
 

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/base.py

@@ -55,6 +55,9 @@ class CallbackConfigBase(C.Config, ABC):
     ) -> Iterable[Callback | CallbackWithMetadata]: ...
 
 
+callback_registry = C.Registry(CallbackConfigBase, discriminator="name")
+
+
 # region Config resolution helpers
 def _create_callbacks_with_metadata(
     config: CallbackConfigBase, trainer_config: TrainerConfig

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/checkpoint/__init__.py

@@ -14,7 +14,3 @@ from .on_exception_checkpoint import (
 from .on_exception_checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
-from .time_checkpoint import TimeCheckpointCallback as TimeCheckpointCallback
-from .time_checkpoint import (
-    TimeCheckpointCallbackConfig as TimeCheckpointCallbackConfig,
-)

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -9,12 +9,14 @@ from typing_extensions import final, override
 
 from ..._checkpoint.metadata import CheckpointMetadata
 from ...metrics._config import MetricConfig
+from ..base import callback_registry
 from ._base import BaseCheckpointCallbackConfig, CheckpointBase
 
 log = logging.getLogger(__name__)
 
 
 @final
+@callback_registry.register
 class BestCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
     name: Literal["best_checkpoint"] = "best_checkpoint"
 

nshtrainer-1.0.0b33/src/nshtrainer/callbacks/checkpoint/time_checkpoint.py → nshtrainer-1.0.0b37/src/nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -9,36 +9,41 @@ from typing import Any, Literal
 from lightning.pytorch import LightningModule, Trainer
 from typing_extensions import final, override
 
-from nshtrainer._checkpoint.metadata import CheckpointMetadata
-
+from ..._checkpoint.metadata import CheckpointMetadata
+from ..base import callback_registry
 from ._base import BaseCheckpointCallbackConfig, CheckpointBase
 
 log = logging.getLogger(__name__)
 
 
 @final
-class TimeCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
-    name: Literal["time_checkpoint"] = "time_checkpoint"
+@callback_registry.register
+class LastCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
+    name: Literal["last_checkpoint"] = "last_checkpoint"
+
+    save_on_time_interval: bool = True
+    """Whether to save checkpoints based on time interval."""
 
     interval: timedelta = timedelta(hours=12)
-    """Time interval between checkpoints."""
+    """Time interval between checkpoints when save_on_time_interval is True."""
 
     @override
     def create_checkpoint(self, trainer_config, dirpath):
-        return TimeCheckpointCallback(self, dirpath)
+        return LastCheckpointCallback(self, dirpath)
 
 
 @final
-class TimeCheckpointCallback(CheckpointBase[TimeCheckpointCallbackConfig]):
-    def __init__(self, config: TimeCheckpointCallbackConfig, dirpath: Path):
+class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
+    def __init__(self, config: LastCheckpointCallbackConfig, dirpath: Path):
         super().__init__(config, dirpath)
         self.start_time = time.time()
         self.last_checkpoint_time = self.start_time
         self.interval_seconds = config.interval.total_seconds()
+        self.save_on_time_interval = config.save_on_time_interval
 
     @override
     def name(self):
-        return "time"
+        return "last"
 
     @override
     def default_filename(self):
@@ -53,6 +58,8 @@ class TimeCheckpointCallback(CheckpointBase[TimeCheckpointCallbackConfig]):
         return True
 
     def _should_checkpoint(self) -> bool:
+        if not self.save_on_time_interval:
+            return False
         current_time = time.time()
         elapsed_time = current_time - self.last_checkpoint_time
         return elapsed_time >= self.interval_seconds
@@ -85,30 +92,23 @@ class TimeCheckpointCallback(CheckpointBase[TimeCheckpointCallbackConfig]):
 
     @override
     def on_train_batch_end(
-        self, trainer: Trainer, pl_module: LightningModule, *args, **kwargs
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        *args,
+        **kwargs,
     ):
-        if self._should_checkpoint():
-            self.save_checkpoints(trainer)
-            self.last_checkpoint_time = time.time()
+        if not self._should_checkpoint():
+            return
+        self.save_checkpoints(trainer)
 
     @override
-    def state_dict(self) -> dict[str, Any]:
-        """Save the timer state for checkpoint resumption.
-
-        Returns:
-            Dictionary containing the start time and last checkpoint time.
-        """
-        return {
-            "start_time": self.start_time,
-            "last_checkpoint_time": self.last_checkpoint_time,
-        }
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        self.save_checkpoints(trainer)
 
     @override
-    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
-        """Restore the timer state when resuming from a checkpoint.
-
-        Args:
-            state_dict: Dictionary containing the previously saved timer state.
-        """
-        self.start_time = state_dict["start_time"]
-        self.last_checkpoint_time = state_dict["last_checkpoint_time"]
+    def save_checkpoints(self, trainer):
+        super().save_checkpoints(trainer)
+
+        if self.save_on_time_interval:
+            self.last_checkpoint_time = time.time()

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py

@@ -9,9 +9,9 @@ from typing import Any, Literal
 
 from lightning.pytorch import Trainer as LightningTrainer
 from lightning.pytorch.callbacks import OnExceptionCheckpoint as _OnExceptionCheckpoint
-from typing_extensions import override
+from typing_extensions import final, override
 
-from ..base import CallbackConfigBase
+from ..base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
@@ -44,6 +44,8 @@ def _monkey_patch_disable_barrier(trainer: LightningTrainer):
         log.warning("Reverted monkey-patched barrier.")
 
 
+@final
+@callback_registry.register
 class OnExceptionCheckpointCallbackConfig(CallbackConfigBase):
     name: Literal["on_exception_checkpoint"] = "on_exception_checkpoint"
 

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/debug_flag.py

@@ -3,14 +3,16 @@ from __future__ import annotations
 import logging
 from typing import Literal
 
-from typing_extensions import override
+from typing_extensions import final, override
 
 from .._callback import NTCallbackBase
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
 class DebugFlagCallbackConfig(CallbackConfigBase):
     name: Literal["debug_flag"] = "debug_flag"
 

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/directory_setup.py

@@ -5,14 +5,35 @@ import os
 from pathlib import Path
 from typing import Literal
 
-from typing_extensions import override
+from typing_extensions import final, override
 
 from .._callback import NTCallbackBase
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
+class DirectorySetupCallbackConfig(CallbackConfigBase):
+    name: Literal["directory_setup"] = "directory_setup"
+
+    enabled: bool = True
+    """Whether to enable the directory setup callback."""
+
+    create_symlink_to_nshrunner_root: bool = True
+    """Should we create a symlink to the root folder for the Runner (if we're in one)?"""
+
+    def __bool__(self):
+        return self.enabled
+
+    def create_callbacks(self, trainer_config):
+        if not self:
+            return
+
+        yield DirectorySetupCallback(self)
+
+
 def _create_symlink_to_nshrunner(base_dir: Path):
     # Resolve the current nshrunner session directory
     if not (session_dir := os.environ.get("NSHRUNNER_SESSION_DIR")):
@@ -43,25 +64,6 @@ def _create_symlink_to_nshrunner(base_dir: Path):
     symlink_path.symlink_to(session_dir)
 
 
-class DirectorySetupCallbackConfig(CallbackConfigBase):
-    name: Literal["directory_setup"] = "directory_setup"
-
-    enabled: bool = True
-    """Whether to enable the directory setup callback."""
-
-    create_symlink_to_nshrunner_root: bool = True
-    """Should we create a symlink to the root folder for the Runner (if we're in one)?"""
-
-    def __bool__(self):
-        return self.enabled
-
-    def create_callbacks(self, trainer_config):
-        if not self:
-            return
-
-        yield DirectorySetupCallback(self)
-
-
 class DirectorySetupCallback(NTCallbackBase):
     @override
     def __init__(self, config: DirectorySetupCallbackConfig):

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/early_stopping.py

@@ -8,14 +8,16 @@ from lightning.fabric.utilities.rank_zero import _get_rank
 from lightning.pytorch import Trainer
 from lightning.pytorch.callbacks import EarlyStopping as _EarlyStopping
 from lightning.pytorch.utilities.rank_zero import rank_prefixed_message
-from typing_extensions import override
+from typing_extensions import final, override
 
 from ..metrics._config import MetricConfig
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
 class EarlyStoppingCallbackConfig(CallbackConfigBase):
     name: Literal["early_stopping"] = "early_stopping"
 

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/ema.py

@@ -10,9 +10,36 @@ import lightning.pytorch as pl
 import torch
 from lightning.pytorch import Callback
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
-from typing_extensions import override
+from typing_extensions import final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
+
+
+@final
+@callback_registry.register
+class EMACallbackConfig(CallbackConfigBase):
+    name: Literal["ema"] = "ema"
+
+    decay: float
+    """The exponential decay used when calculating the moving average. Has to be between 0-1."""
+
+    validate_original_weights: bool = False
+    """Validate the original weights, as apposed to the EMA weights."""
+
+    every_n_steps: int = 1
+    """Apply EMA every N steps."""
+
+    cpu_offload: bool = False
+    """Offload weights to CPU."""
+
+    @override
+    def create_callbacks(self, trainer_config):
+        yield EMACallback(
+            decay=self.decay,
+            validate_original_weights=self.validate_original_weights,
+            every_n_steps=self.every_n_steps,
+            cpu_offload=self.cpu_offload,
+        )
 
 
 class EMACallback(Callback):
@@ -358,28 +385,3 @@ class EMAOptimizer(torch.optim.Optimizer):
     def add_param_group(self, param_group):
         self.optimizer.add_param_group(param_group)
         self.rebuild_ema_params = True
-
-
-class EMACallbackConfig(CallbackConfigBase):
-    name: Literal["ema"] = "ema"
-
-    decay: float
-    """The exponential decay used when calculating the moving average. Has to be between 0-1."""
-
-    validate_original_weights: bool = False
-    """Validate the original weights, as apposed to the EMA weights."""
-
-    every_n_steps: int = 1
-    """Apply EMA every N steps."""
-
-    cpu_offload: bool = False
-    """Offload weights to CPU."""
-
-    @override
-    def create_callbacks(self, trainer_config):
-        yield EMACallback(
-            decay=self.decay,
-            validate_original_weights=self.validate_original_weights,
-            every_n_steps=self.every_n_steps,
-            cpu_offload=self.cpu_offload,
-        )

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/finite_checks.py

@@ -5,13 +5,32 @@ from typing import Literal
 
 import torch
 from lightning.pytorch import Callback, LightningModule, Trainer
-from typing_extensions import override
+from typing_extensions import final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
+class FiniteChecksCallbackConfig(CallbackConfigBase):
+    name: Literal["finite_checks"] = "finite_checks"
+
+    nonfinite_grads: bool = True
+    """Whether to check for non-finite (i.e. NaN or Inf) gradients"""
+
+    none_grads: bool = True
+    """Whether to check for None gradients"""
+
+    @override
+    def create_callbacks(self, trainer_config):
+        yield FiniteChecksCallback(
+            nonfinite_grads=self.nonfinite_grads,
+            none_grads=self.none_grads,
+        )
+
+
 def finite_checks(
     module: LightningModule,
     nonfinite_grads: bool = True,
@@ -58,20 +77,3 @@ class FiniteChecksCallback(Callback):
             nonfinite_grads=self._nonfinite_grads,
             none_grads=self._none_grads,
         )
-
-
-class FiniteChecksCallbackConfig(CallbackConfigBase):
-    name: Literal["finite_checks"] = "finite_checks"
-
-    nonfinite_grads: bool = True
-    """Whether to check for non-finite (i.e. NaN or Inf) gradients"""
-
-    none_grads: bool = True
-    """Whether to check for None gradients"""
-
-    @override
-    def create_callbacks(self, trainer_config):
-        yield FiniteChecksCallback(
-            nonfinite_grads=self.nonfinite_grads,
-            none_grads=self.none_grads,
-        )

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/gradient_skipping.py

@@ -7,21 +7,47 @@ import torch
 import torchmetrics
 from lightning.pytorch import Callback, LightningModule, Trainer
 from torch.optim import Optimizer
-from typing_extensions import override
+from typing_extensions import final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 from .norm_logging import compute_norm
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
+class GradientSkippingCallbackConfig(CallbackConfigBase):
+    name: Literal["gradient_skipping"] = "gradient_skipping"
+
+    threshold: float
+    """Threshold to use for gradient skipping."""
+
+    norm_type: str | float = 2.0
+    """Norm type to use for gradient skipping."""
+
+    start_after_n_steps: int | None = 100
+    """Number of steps to wait before starting gradient skipping."""
+
+    skip_non_finite: bool = False
+    """
+    If False, it doesn't skip steps with non-finite norms. This is useful when using AMP, as AMP checks for NaN/Inf grads to adjust the loss scale. Otherwise, skips steps with non-finite norms.
+
+    Should almost always be False, especially when using AMP (unless you know what you're doing!).
+    """
+
+    @override
+    def create_callbacks(self, trainer_config):
+        yield GradientSkippingCallback(self)
+
+
 @runtime_checkable
 class HasGradSkippedSteps(Protocol):
     grad_skipped_steps: Any
 
 
 class GradientSkippingCallback(Callback):
-    def __init__(self, config: "GradientSkippingCallbackConfig"):
+    def __init__(self, config: GradientSkippingCallbackConfig):
         super().__init__()
         self.config = config
 
@@ -73,27 +99,3 @@ class GradientSkippingCallback(Callback):
             on_step=True,
             on_epoch=False,
         )
-
-
-class GradientSkippingCallbackConfig(CallbackConfigBase):
-    name: Literal["gradient_skipping"] = "gradient_skipping"
-
-    threshold: float
-    """Threshold to use for gradient skipping."""
-
-    norm_type: str | float = 2.0
-    """Norm type to use for gradient skipping."""
-
-    start_after_n_steps: int | None = 100
-    """Number of steps to wait before starting gradient skipping."""
-
-    skip_non_finite: bool = False
-    """
-    If False, it doesn't skip steps with non-finite norms. This is useful when using AMP, as AMP checks for NaN/Inf grads to adjust the loss scale. Otherwise, skips steps with non-finite norms.
-
-    Should almost always be False, especially when using AMP (unless you know what you're doing!).
-    """
-
-    @override
-    def create_callbacks(self, trainer_config):
-        yield GradientSkippingCallback(self)

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/log_epoch.py

@@ -6,13 +6,15 @@ from typing import Any, Literal
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
-from typing_extensions import override
+from typing_extensions import final, override
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
 
 
+@final
+@callback_registry.register
 class LogEpochCallbackConfig(CallbackConfigBase):
     name: Literal["log_epoch"] = "log_epoch"
 

{nshtrainer-1.0.0b33 → nshtrainer-1.0.0b37}/src/nshtrainer/callbacks/lr_monitor.py

@@ -3,11 +3,16 @@ from __future__ import annotations
 from typing import Literal
 
 from lightning.pytorch.callbacks import LearningRateMonitor
+from typing_extensions import final
 
-from .base import CallbackConfigBase
+from .base import CallbackConfigBase, callback_registry
 
 
+@final
+@callback_registry.register
 class LearningRateMonitorConfig(CallbackConfigBase):
+    name: Literal["learning_rate_monitor"] = "learning_rate_monitor"
+
     logging_interval: Literal["step", "epoch"] | None = None
     """
     Set to 'epoch' or 'step' to log 'lr' of all optimizers at the same interval, set to None to log at individual interval according to the 'interval' key of each scheduler. Defaults to None.