nshtrainer 0.44.0__py3-none-any.whl → 1.0.0b9__py3-none-any.whl

nshtrainer/_hf_hub.py

@@ -7,19 +7,19 @@ import re
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
 
 import nshconfig as C
 from nshrunner._env import SNAPSHOT_DIR
 from typing_extensions import assert_never, override
 
 from ._callback import NTCallbackBase
-from .callbacks.base import CallbackConfigBase
+from .callbacks.base import CallbackConfigBase, CallbackMetadataConfig
 
 if TYPE_CHECKING:
     from huggingface_hub import HfApi  # noqa: F401
 
-    from .model.base import BaseConfig
+    from .trainer._config import TrainerConfig
 
 
 log = logging.getLogger(__name__)
@@ -42,6 +42,8 @@ class HuggingFaceHubAutoCreateConfig(C.Config):
 class HuggingFaceHubConfig(CallbackConfigBase):
     """Configuration options for Hugging Face Hub integration."""
 
+    metadata: ClassVar[CallbackMetadataConfig] = {"ignore_if_exists": True}
+
     enabled: bool = False
     """Enable Hugging Face Hub integration."""
 
@@ -82,7 +84,7 @@ class HuggingFaceHubConfig(CallbackConfigBase):
         return self.enabled
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         # Attempt to login. If it fails, we'll log a warning or error based on the configuration.
         try:
             api = _api(self.token)
@@ -107,7 +109,7 @@ class HuggingFaceHubConfig(CallbackConfigBase):
                 case _:
                     assert_never(self.on_login_error)
 
-        yield self.with_metadata(HFHubCallback(self), ignore_if_exists=True)
+        yield HFHubCallback(self)
 
 
 def _api(token: str | None = None):
@@ -138,19 +140,20 @@ def _api(token: str | None = None):
     return api
 
 
-def _repo_name(api: "HfApi", root_config: "BaseConfig"):
+def _repo_name(api: HfApi, trainer_config: TrainerConfig):
     username = None
-    if (ac := root_config.trainer.hf_hub.auto_create) and ac.namespace:
+    if (ac := trainer_config.hf_hub.auto_create) and ac.namespace:
         username = ac.namespace
     elif (username := api.whoami().get("name", None)) is None:
         raise ValueError("Could not get username from Hugging Face Hub.")
 
     # Sanitize the project (if it exists), run_name, and id
     parts = []
-    if root_config.project:
-        parts.append(re.sub(r"[^a-zA-Z0-9-]", "-", root_config.project))
-    parts.append(re.sub(r"[^a-zA-Z0-9-]", "-", root_config.run_name))
-    parts.append(re.sub(r"[^a-zA-Z0-9-]", "-", root_config.id))
+    if trainer_config.project:
+        parts.append(re.sub(r"[^a-zA-Z0-9-]", "-", trainer_config.project))
+    if trainer_config.full_name:
+        parts.append(re.sub(r"[^a-zA-Z0-9-]", "-", trainer_config.full_name))
+    parts.append(re.sub(r"[^a-zA-Z0-9-]", "-", trainer_config.id))
 
     # Combine parts and ensure it starts and ends with alphanumeric characters
     repo_name = "-".join(parts)
@@ -179,14 +182,10 @@ class _Upload:
     path_in_repo: Path
 
     @classmethod
-    def from_local_path(
-        cls,
-        local_path: Path,
-        root_config: "BaseConfig",
-    ):
+    def from_local_path(cls, local_path: Path, trainer_config: TrainerConfig):
         # Resolve the checkpoint directory
-        checkpoint_dir = root_config.directory.resolve_subdirectory(
-            root_config.id, "checkpoint"
+        checkpoint_dir = trainer_config.directory.resolve_subdirectory(
+            trainer_config.id, "checkpoint"
         )
 
         try:
@@ -224,8 +223,7 @@ class HFHubCallback(NTCallbackBase):
 
     @override
     def setup(self, trainer, pl_module, stage):
-        root_config = cast("BaseConfig", pl_module.hparams)
-        self._repo_id = _repo_name(self.api, root_config)
+        self._repo_id = _repo_name(self.api, trainer.hparams)
 
         if not self.config or not trainer.is_global_zero:
             return
@@ -234,7 +232,7 @@ class HFHubCallback(NTCallbackBase):
         self._create_repo_if_not_exists()
 
         # Upload the config and code
-        self._save_config(root_config)
+        self._save_config(trainer.hparams)
         self._save_code()
 
     @override
@@ -248,10 +246,9 @@ class HFHubCallback(NTCallbackBase):
             return
 
         with self._with_error_handling("save checkpoints"):
-            root_config = cast("BaseConfig", pl_module.hparams)
             self._save_checkpoint(
-                _Upload.from_local_path(ckpt_path, root_config),
-                _Upload.from_local_path(metadata_path, root_config)
+                _Upload.from_local_path(ckpt_path, trainer.hparams),
+                _Upload.from_local_path(metadata_path, trainer.hparams)
                 if metadata_path is not None
                 else None,
             )
@@ -300,10 +297,12 @@ class HFHubCallback(NTCallbackBase):
                     f"Error checking repository '{self.repo_id}'", exc_info=True
                 )
 
-    def _save_config(self, root_config: "BaseConfig"):
+    def _save_config(self, trainer_config: TrainerConfig):
         with self._with_error_handling("upload config"):
             self.api.upload_file(
-                path_or_fileobj=root_config.model_dump_json(indent=4).encode("utf-8"),
+                path_or_fileobj=trainer_config.model_dump_json(indent=4).encode(
+                    "utf-8"
+                ),
                 path_in_repo="config.json",
                 repo_id=self.repo_id,
                 repo_type="model",

nshtrainer/callbacks/__init__.py

@@ -6,7 +6,7 @@ import nshconfig as C
 
 from . import checkpoint as checkpoint
 from .base import CallbackConfigBase as CallbackConfigBase
-from .checkpoint import BestCheckpoint as BestCheckpoint
+from .checkpoint import BestCheckpointCallback as BestCheckpointCallback
 from .checkpoint import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
 from .checkpoint import LastCheckpointCallback as LastCheckpointCallback
 from .checkpoint import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
@@ -49,7 +49,6 @@ from .shared_parameters import SharedParametersCallback as SharedParametersCallb
 from .shared_parameters import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
-from .throughput_monitor import ThroughputMonitorConfig as ThroughputMonitorConfig
 from .timer import EpochTimerCallback as EpochTimerCallback
 from .timer import EpochTimerCallbackConfig as EpochTimerCallbackConfig
 from .wandb_upload_code import WandbUploadCodeCallback as WandbUploadCodeCallback
@@ -62,7 +61,6 @@ from .wandb_watch import WandbWatchCallbackConfig as WandbWatchCallbackConfig
 CallbackConfig = Annotated[
     DebugFlagCallbackConfig
     | EarlyStoppingCallbackConfig
-    | ThroughputMonitorConfig
     | EpochTimerCallbackConfig
     | PrintTableMetricsCallbackConfig
     | FiniteChecksCallbackConfig

nshtrainer/callbacks/actsave.py

@@ -4,11 +4,9 @@ import contextlib
 from pathlib import Path
 from typing import Literal
 
-from lightning.pytorch import LightningModule, Trainer
-from lightning.pytorch.callbacks.callback import Callback
-from nshutils import ActSave
 from typing_extensions import TypeAlias, override
 
+from .._callback import NTCallbackBase
 from .base import CallbackConfigBase
 
 Stage: TypeAlias = Literal["train", "validation", "test", "predict"]
@@ -25,15 +23,17 @@ class ActSaveConfig(CallbackConfigBase):
         return self.enabled
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield ActSaveCallback(
             self,
             self.save_dir
-            or root_config.directory.resolve_subdirectory(root_config.id, "activation"),
+            or trainer_config.directory.resolve_subdirectory(
+                trainer_config.id, "activation"
+            ),
         )
 
 
-class ActSaveCallback(Callback):
+class ActSaveCallback(NTCallbackBase):
     def __init__(self, config: ActSaveConfig, save_dir: Path):
         super().__init__()
 
@@ -43,20 +43,20 @@ class ActSaveCallback(Callback):
         self._active_contexts: dict[Stage, contextlib._GeneratorContextManager] = {}
 
     @override
-    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
+    def setup(self, trainer, pl_module, stage) -> None:
         super().setup(trainer, pl_module, stage)
 
         if not self.config:
             return
 
+        from nshutils import ActSave
+
         context = ActSave.enabled(self.save_dir)
         context.__enter__()
         self._enabled_context = context
 
     @override
-    def teardown(
-        self, trainer: Trainer, pl_module: LightningModule, stage: str
-    ) -> None:
+    def teardown(self, trainer, pl_module, stage) -> None:
         super().teardown(trainer, pl_module, stage)
 
         if not self.config:
@@ -66,10 +66,12 @@ class ActSaveCallback(Callback):
             self._enabled_context.__exit__(None, None, None)
             self._enabled_context = None
 
-    def _on_start(self, stage: Stage, trainer: Trainer, pl_module: LightningModule):
+    def _on_start(self, stage: Stage, trainer, pl_module):
         if not self.config:
             return
 
+        from nshutils import ActSave
+
         # If we have an active context manager for this stage, exit it
         if active_contexts := self._active_contexts.get(stage):
             active_contexts.__exit__(None, None, None)
@@ -79,7 +81,7 @@ class ActSaveCallback(Callback):
         context.__enter__()
         self._active_contexts[stage] = context
 
-    def _on_end(self, stage: Stage, trainer: Trainer, pl_module: LightningModule):
+    def _on_end(self, stage: Stage, trainer, pl_module):
         if not self.config:
             return
 
@@ -88,33 +90,33 @@ class ActSaveCallback(Callback):
             active_contexts.__exit__(None, None, None)
 
     @override
-    def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule):
+    def on_train_epoch_start(self, trainer, pl_module):
         return self._on_start("train", trainer, pl_module)
 
     @override
-    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+    def on_train_epoch_end(self, trainer, pl_module):
         return self._on_end("train", trainer, pl_module)
 
     @override
-    def on_validation_epoch_start(self, trainer: Trainer, pl_module: LightningModule):
+    def on_validation_epoch_start(self, trainer, pl_module):
         return self._on_start("validation", trainer, pl_module)
 
     @override
-    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+    def on_validation_epoch_end(self, trainer, pl_module):
         return self._on_end("validation", trainer, pl_module)
 
     @override
-    def on_test_epoch_start(self, trainer: Trainer, pl_module: LightningModule):
+    def on_test_epoch_start(self, trainer, pl_module):
         return self._on_start("test", trainer, pl_module)
 
     @override
-    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+    def on_test_epoch_end(self, trainer, pl_module):
         return self._on_end("test", trainer, pl_module)
 
     @override
-    def on_predict_epoch_start(self, trainer: Trainer, pl_module: LightningModule):
+    def on_predict_epoch_start(self, trainer, pl_module):
         return self._on_start("predict", trainer, pl_module)
 
     @override
-    def on_predict_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+    def on_predict_epoch_end(self, trainer, pl_module):
         return self._on_end("predict", trainer, pl_module)

nshtrainer/callbacks/base.py

@@ -11,7 +11,7 @@ from lightning.pytorch import Callback
 from typing_extensions import TypedDict, Unpack
 
 if TYPE_CHECKING:
-    from ..model.config import BaseConfig
+    from ..trainer._config import TrainerConfig
 
 
 class CallbackMetadataConfig(TypedDict, total=False):
@@ -49,15 +49,15 @@ class CallbackConfigBase(C.Config, ABC):
 
     @abstractmethod
     def create_callbacks(
-        self, root_config: "BaseConfig"
+        self, trainer_config: TrainerConfig
     ) -> Iterable[Callback | CallbackWithMetadata]: ...
 
 
 # region Config resolution helpers
 def _create_callbacks_with_metadata(
-    config: CallbackConfigBase, root_config: "BaseConfig"
+    config: CallbackConfigBase, trainer_config: TrainerConfig
 ) -> Iterable[CallbackWithMetadata]:
-    for callback in config.create_callbacks(root_config):
+    for callback in config.create_callbacks(trainer_config):
         if isinstance(callback, CallbackWithMetadata):
             yield callback
             continue
@@ -102,16 +102,16 @@ def _process_and_filter_callbacks(
     return [callback.callback for callback in callbacks]
 
 
-def resolve_all_callbacks(root_config: "BaseConfig"):
+def resolve_all_callbacks(trainer_config: TrainerConfig):
     callback_configs = [
         config
-        for config in root_config._nshtrainer_all_callback_configs()
+        for config in trainer_config._nshtrainer_all_callback_configs()
         if config is not None
     ]
     callbacks = _process_and_filter_callbacks(
         callback
         for callback_config in callback_configs
-        for callback in _create_callbacks_with_metadata(callback_config, root_config)
+        for callback in _create_callbacks_with_metadata(callback_config, trainer_config)
     )
     return callbacks
 

nshtrainer/callbacks/checkpoint/__init__.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from .best_checkpoint import BestCheckpoint as BestCheckpoint
+from .best_checkpoint import BestCheckpointCallback as BestCheckpointCallback
 from .best_checkpoint import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )

nshtrainer/callbacks/checkpoint/_base.py

@@ -16,7 +16,8 @@ from ..._checkpoint.saver import _link_checkpoint, _remove_checkpoint
 from ..base import CallbackConfigBase
 
 if TYPE_CHECKING:
-    from ...model.config import BaseConfig
+    from ...trainer._config import TrainerConfig
+
 
 log = logging.getLogger(__name__)
 
@@ -41,18 +42,20 @@ class BaseCheckpointCallbackConfig(CallbackConfigBase, ABC):
     @abstractmethod
     def create_checkpoint(
         self,
-        root_config: "BaseConfig",
+        trainer_config: TrainerConfig,
         dirpath: Path,
     ) -> "CheckpointBase | None": ...
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         dirpath = Path(
             self.dirpath
-            or root_config.directory.resolve_subdirectory(root_config.id, "checkpoint")
+            or trainer_config.directory.resolve_subdirectory(
+                trainer_config.id, "checkpoint"
+            )
         )
 
-        if (callback := self.create_checkpoint(root_config, dirpath)) is not None:
+        if (callback := self.create_checkpoint(trainer_config, dirpath)) is not None:
             yield callback
 
 

nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -28,10 +28,10 @@ class BestCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
     """
 
     @override
-    def create_checkpoint(self, root_config, dirpath):
+    def create_checkpoint(self, trainer_config, dirpath):
         # Resolve metric
         if (metric := self.metric) is None and (
-            metric := root_config.primary_metric
+            metric := trainer_config.primary_metric
         ) is None:
             error_msg = (
                 "No metric provided and no primary metric found in the root config. "
@@ -43,11 +43,11 @@ class BestCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
                 log.warning(error_msg)
                 return None
 
-        return BestCheckpoint(self, dirpath, metric)
+        return BestCheckpointCallback(self, dirpath, metric)
 
 
 @final
-class BestCheckpoint(CheckpointBase[BestCheckpointCallbackConfig]):
+class BestCheckpointCallback(CheckpointBase[BestCheckpointCallbackConfig]):
     @property
     def _metric_name_normalized(self):
         return self.metric.name.replace("/", "_").replace(" ", "_").replace(".", "_")

nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -18,7 +18,7 @@ class LastCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
     name: Literal["last_checkpoint"] = "last_checkpoint"
 
     @override
-    def create_checkpoint(self, root_config, dirpath):
+    def create_checkpoint(self, trainer_config, dirpath):
         return LastCheckpointCallback(self, dirpath)
 
 

nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py

@@ -54,13 +54,13 @@ class OnExceptionCheckpointCallbackConfig(CallbackConfigBase):
     """Checkpoint filename. This must not include the extension. If `None`, `on_exception_{id}_{timestamp}` is used."""
 
     @override
-    def create_callbacks(self, root_config):
-        dirpath = self.dirpath or root_config.directory.resolve_subdirectory(
-            root_config.id, "checkpoint"
+    def create_callbacks(self, trainer_config):
+        dirpath = self.dirpath or trainer_config.directory.resolve_subdirectory(
+            trainer_config.id, "checkpoint"
         )
 
         if not (filename := self.filename):
-            filename = f"on_exception_{root_config.id}"
+            filename = f"on_exception_{trainer_config.id}"
         yield OnExceptionCheckpointCallback(
             self, dirpath=Path(dirpath), filename=filename
         )

nshtrainer/callbacks/debug_flag.py

@@ -1,17 +1,13 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Literal, cast
+from typing import Literal
 
-from lightning.pytorch import LightningModule, Trainer
-from lightning.pytorch.callbacks import Callback
 from typing_extensions import override
 
+from .._callback import NTCallbackBase
 from .base import CallbackConfigBase
 
-if TYPE_CHECKING:
-    from ..model.config import BaseConfig
-
 log = logging.getLogger(__name__)
 
 
@@ -25,14 +21,14 @@ class DebugFlagCallbackConfig(CallbackConfigBase):
         return self.enabled
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         if not self:
             return
 
         yield DebugFlagCallback(self)
 
 
-class DebugFlagCallback(Callback):
+class DebugFlagCallback(NTCallbackBase):
     """
     Sets the debug flag to true in the following circumstances:
     - fast_dev_run is enabled
@@ -46,27 +42,26 @@ class DebugFlagCallback(Callback):
         self.config = config
         del config
 
+        self._debug = False
+
     @override
-    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str):
+    def setup(self, trainer, pl_module, stage):
         if not getattr(trainer, "fast_dev_run", False):
             return
 
-        hparams = cast("BaseConfig", pl_module.hparams)
-        if not hparams.debug:
+        if not trainer.debug:
             log.critical("Fast dev run detected, setting debug flag to True.")
-        hparams.debug = True
+        trainer.debug = True
 
     @override
-    def on_sanity_check_start(self, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast("BaseConfig", pl_module.hparams)
-        self._debug = hparams.debug
+    def on_sanity_check_start(self, trainer, pl_module):
+        self._debug = trainer.debug
         if not self._debug:
             log.critical("Enabling debug flag during sanity check routine.")
-        hparams.debug = True
+        trainer.debug = True
 
     @override
-    def on_sanity_check_end(self, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast("BaseConfig", pl_module.hparams)
+    def on_sanity_check_end(self, trainer, pl_module):
         if not self._debug:
             log.critical("Sanity check routine complete, disabling debug flag.")
-        hparams.debug = self._debug
+        trainer.debug = self._debug

nshtrainer/callbacks/directory_setup.py

@@ -5,9 +5,9 @@ import os
 from pathlib import Path
 from typing import Literal
 
-from lightning.pytorch import Callback
 from typing_extensions import override
 
+from .._callback import NTCallbackBase
 from .base import CallbackConfigBase
 
 log = logging.getLogger(__name__)
@@ -55,14 +55,14 @@ class DirectorySetupCallbackConfig(CallbackConfigBase):
     def __bool__(self):
         return self.enabled
 
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         if not self:
             return
 
         yield DirectorySetupCallback(self)
 
 
-class DirectorySetupCallback(Callback):
+class DirectorySetupCallback(NTCallbackBase):
     @override
     def __init__(self, config: DirectorySetupCallbackConfig):
         super().__init__()
@@ -76,12 +76,7 @@ class DirectorySetupCallback(Callback):
 
         # Create a symlink to the root folder for the Runner
         if self.config.create_symlink_to_nshrunner_root:
-            # Resolve the base dir
-            from ..model.config import BaseConfig
-
-            assert isinstance(
-                config := pl_module.hparams, BaseConfig
-            ), f"Expected a BaseConfig, got {type(config)}"
-
-            base_dir = config.directory.resolve_run_root_directory(config.id)
+            base_dir = trainer.hparams.directory.resolve_run_root_directory(
+                trainer.hparams.id
+            )
             _create_symlink_to_nshrunner(base_dir)

nshtrainer/callbacks/early_stopping.py

@@ -48,12 +48,12 @@ class EarlyStoppingCallbackConfig(CallbackConfigBase):
     """
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         if (metric := self.metric) is None and (
-            metric := root_config.primary_metric
+            metric := trainer_config.primary_metric
         ) is None:
             raise ValueError(
-                "Either `metric` or `root_config.primary_metric` must be set to use EarlyStopping."
+                "Either `metric` or `trainer_config.primary_metric` must be set to use EarlyStopping."
             )
 
         yield EarlyStoppingCallback(self, metric)

nshtrainer/callbacks/ema.py

@@ -376,7 +376,7 @@ class EMACallbackConfig(CallbackConfigBase):
     """Offload weights to CPU."""
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield EMACallback(
             decay=self.decay,
             validate_original_weights=self.validate_original_weights,

nshtrainer/callbacks/finite_checks.py

@@ -70,7 +70,7 @@ class FiniteChecksCallbackConfig(CallbackConfigBase):
     """Whether to check for None gradients"""
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield FiniteChecksCallback(
             nonfinite_grads=self.nonfinite_grads,
             none_grads=self.none_grads,

nshtrainer/callbacks/gradient_skipping.py

@@ -95,5 +95,5 @@ class GradientSkippingCallbackConfig(CallbackConfigBase):
     """
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield GradientSkippingCallback(self)

nshtrainer/callbacks/log_epoch.py

@@ -17,7 +17,7 @@ class LogEpochCallbackConfig(CallbackConfigBase):
     name: Literal["log_epoch"] = "log_epoch"
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield LogEpochCallback()
 
 

nshtrainer/callbacks/norm_logging.py

@@ -182,7 +182,7 @@ class NormLoggingCallbackConfig(CallbackConfigBase):
         )
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         if not self:
             return
 

nshtrainer/callbacks/print_table.py

@@ -88,5 +88,5 @@ class PrintTableMetricsCallbackConfig(CallbackConfigBase):
     """List of patterns to filter the metrics to be displayed. If None, all metrics are displayed."""
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield PrintTableMetricsCallback(metric_patterns=self.metric_patterns)

nshtrainer/callbacks/rlp_sanity_checks.py

@@ -36,7 +36,7 @@ class RLPSanityChecksCallbackConfig(CallbackConfigBase):
     def __bool__(self):
         return self.enabled
 
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         if not self:
             return
 

nshtrainer/callbacks/shared_parameters.py

@@ -30,7 +30,7 @@ class SharedParametersCallbackConfig(CallbackConfigBase):
     name: Literal["shared_parameters"] = "shared_parameters"
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield SharedParametersCallback(self)
 
 

nshtrainer/callbacks/timer.py

@@ -155,5 +155,5 @@ class EpochTimerCallbackConfig(CallbackConfigBase):
     name: Literal["epoch_timer"] = "epoch_timer"
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         yield EpochTimerCallback()