nshtrainer 0.9.1__py3-none-any.whl → 0.10.1__py3-none-any.whl

nshtrainer/model/modules/logger.py

@@ -11,10 +11,14 @@ from lightning.pytorch.utilities.types import _METRIC
 from lightning_utilities.core.rank_zero import rank_zero_warn
 from typing_extensions import override
 
-from ...actsave import ActSave
 from ...util.typing_utils import mixin_base_type
 from ..config import BaseConfig
 
+try:
+    from nshutils import ActSave  # type: ignore
+except ImportError:
+    ActSave = None
+
 
 @dataclass(frozen=True, kw_only=True)
 class _LogContext:
@@ -155,14 +159,15 @@ class LoggerLightningModuleMixin(LoggerModuleMixin, mixin_base_type(LightningMod
 
     def _logger_actsave(self, name: str, value: _METRIC) -> None:
         hparams = cast(BaseConfig, self.hparams)
-        if (
-            not hparams.trainer.actsave
-            or not hparams.trainer.actsave.auto_save_logged_metrics
-        ):
+        if not hparams.trainer.logging.actsave_logged_metrics:
+            return
+
+        if ActSave is None:
+            rank_zero_warn("ActSave is not available, skipping logging of metrics")
             return
 
         ActSave.save(
-            {
+            lambda: {
                 f"logger.{name}": lambda: value.compute()
                 if isinstance(value, torchmetrics.Metric)
                 else value

nshtrainer/runner.py

@@ -5,6 +5,7 @@ from typing import Generic
 
 from nshrunner import RunInfo
 from nshrunner import Runner as _Runner
+from nshrunner._submit import screen
 from nshrunner.snapshot import SnapshotArgType
 from typing_extensions import TypeVar, TypeVarTuple, Unpack, override
 
@@ -89,6 +90,7 @@ class Runner(
     def fast_dev_run_session(
         self,
         runs: Iterable[tuple[TConfig, Unpack[TArguments]]],
+        options: screen.ScreenJobKwargs = {},
         n_batches: int = 1,
         *,
         snapshot: SnapshotArgType,
@@ -99,10 +101,7 @@ class Runner(
         ]
         | None = None,
         activate_venv: bool = True,
-        session_name: str = "nshrunner",
-        attach: bool = True,
         print_command: bool = True,
-        pause_before_exit: bool = False,
     ):
         transforms = transforms or []
         transforms.append(
@@ -110,13 +109,11 @@ class Runner(
         )
         return self.session(
             runs,
+            options,
             snapshot=snapshot,
             setup_commands=setup_commands,
             env=env,
             transforms=transforms,
             activate_venv=activate_venv,
-            session_name=session_name,
-            attach=attach,
             print_command=print_command,
-            pause_before_exit=pause_before_exit,
         )

nshtrainer/trainer/_runtime_callback.py

@@ -0,0 +1,120 @@
+import datetime
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any, Literal, TypeAlias
+
+from lightning.pytorch.callbacks.callback import Callback
+from typing_extensions import override
+
+log = logging.getLogger(__name__)
+
+Stage: TypeAlias = Literal["train", "validate", "test", "predict"]
+ALL_STAGES = ("train", "validate", "test", "predict")
+
+
+@dataclass
+class TimeInfo:
+    datetime: datetime.datetime
+    monotonic: float
+
+
+class RuntimeTrackerCallback(Callback):
+    def __init__(self):
+        super().__init__()
+        self._start_time: dict[Stage, TimeInfo] = {}
+        self._end_time: dict[Stage, TimeInfo] = {}
+        self._offsets = {stage: datetime.timedelta() for stage in ALL_STAGES}
+
+    def start_time(self, stage: Stage) -> TimeInfo | None:
+        """Return the start time of a particular stage"""
+        return self._start_time.get(stage)
+
+    def end_time(self, stage: Stage) -> TimeInfo | None:
+        """Return the end time of a particular stage"""
+        return self._end_time.get(stage)
+
+    def time_elapsed(self, stage: Stage) -> datetime.timedelta:
+        """Return the time elapsed for a particular stage"""
+        start = self.start_time(stage)
+        end = self.end_time(stage)
+        offset = self._offsets[stage]
+        if start is None:
+            return offset
+        if end is None:
+            current = TimeInfo(datetime.datetime.now(), time.monotonic())
+            return (
+                datetime.timedelta(seconds=current.monotonic - start.monotonic) + offset
+            )
+        return datetime.timedelta(seconds=end.monotonic - start.monotonic) + offset
+
+    def _record_time(self, stage: Stage, time_dict: dict[Stage, TimeInfo]):
+        time_dict[stage] = TimeInfo(datetime.datetime.now(), time.monotonic())
+
+    @override
+    def on_train_start(self, trainer, pl_module):
+        self._record_time("train", self._start_time)
+
+    @override
+    def on_train_end(self, trainer, pl_module):
+        self._record_time("train", self._end_time)
+
+    @override
+    def on_validation_start(self, trainer, pl_module):
+        self._record_time("validate", self._start_time)
+
+    @override
+    def on_validation_end(self, trainer, pl_module):
+        self._record_time("validate", self._end_time)
+
+    @override
+    def on_test_start(self, trainer, pl_module):
+        self._record_time("test", self._start_time)
+
+    @override
+    def on_test_end(self, trainer, pl_module):
+        self._record_time("test", self._end_time)
+
+    @override
+    def on_predict_start(self, trainer, pl_module):
+        self._record_time("predict", self._start_time)
+
+    @override
+    def on_predict_end(self, trainer, pl_module):
+        self._record_time("predict", self._end_time)
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "time_elapsed": {
+                stage: self.time_elapsed(stage).total_seconds() for stage in ALL_STAGES
+            },
+            "start_times": {
+                stage: (info.datetime.isoformat(), info.monotonic)
+                for stage, info in self._start_time.items()
+            },
+            "end_times": {
+                stage: (info.datetime.isoformat(), info.monotonic)
+                for stage, info in self._end_time.items()
+            },
+        }
+
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]):
+        time_elapsed: dict[Stage, float] = state_dict.get("time_elapsed", {})
+        for stage in ALL_STAGES:
+            self._offsets[stage] = datetime.timedelta(
+                seconds=time_elapsed.get(stage, 0)
+            )
+
+        start_times: dict[Stage, tuple[str, float]] = state_dict.get("start_times", {})
+        for stage, (dt_str, monotonic) in start_times.items():
+            self._start_time[stage] = TimeInfo(
+                datetime.datetime.fromisoformat(dt_str), monotonic
+            )
+
+        end_times: dict[Stage, tuple[str, float]] = state_dict.get("end_times", {})
+        for stage, (dt_str, monotonic) in end_times.items():
+            self._end_time[stage] = TimeInfo(
+                datetime.datetime.fromisoformat(dt_str), monotonic
+            )

nshtrainer/trainer/checkpoint_connector.py

@@ -0,0 +1,63 @@
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+from lightning.pytorch.trainer.connectors.checkpoint_connector import (
+    _CheckpointConnector,
+)
+from lightning.pytorch.trainer.states import TrainerFn
+from typing_extensions import override
+
+from .._checkpoint.loader import CheckpointLoadingConfig, _resolve_checkpoint
+
+if TYPE_CHECKING:
+    from ..model.config import BaseConfig
+log = logging.getLogger(__name__)
+
+
+class CheckpointConnector(_CheckpointConnector):
+    def __resolve_auto_ckpt_path(
+        self,
+        ckpt_path: str | Path | None,
+        state_fn: TrainerFn,
+    ):
+        from .trainer import Trainer
+
+        # If this isn't an `nshtrainer` trainer (which I don't know why it wouldn't be),
+        # then we just default to the parent class's implementation of `_parse_ckpt_path`.
+        trainer = self.trainer
+        if not isinstance(trainer, Trainer):
+            return None
+
+        # Now, resolve the checkpoint loader config.
+        root_config = cast("BaseConfig", trainer._base_module.config)
+        if (ckpt_loader_config := root_config.trainer.checkpoint_loading) == "auto":
+            ckpt_loader_config = CheckpointLoadingConfig.auto(ckpt_path, state_fn)
+        log.debug(f"Checkpoint loader config: {ckpt_loader_config}")
+
+        # Use the config to resolve the checkpoint.
+        if (
+            ckpt_path := _resolve_checkpoint(ckpt_loader_config, root_config, trainer)
+        ) is None:
+            log.info(
+                "No checkpoint found for the current trainer state. "
+                "Training will start from scratch."
+            )
+
+        log.info(f"Loading checkpoint from: {ckpt_path}")
+        return ckpt_path
+
+    @override
+    def _parse_ckpt_path(
+        self,
+        state_fn: TrainerFn,
+        ckpt_path: str | Path | None,
+        model_provided: bool,
+        model_connected: bool,
+    ):
+        if (p := self.__resolve_auto_ckpt_path(ckpt_path, state_fn)) is not None:
+            return p
+
+        return super()._parse_ckpt_path(
+            state_fn, ckpt_path, model_provided, model_connected
+        )

nshtrainer/trainer/signal_connector.py

@@ -11,6 +11,7 @@ from pathlib import Path
 from types import FrameType
 from typing import Any, TypeAlias
 
+import nshrunner as nr
 import torch.utils.data
 from lightning.fabric.plugins.environments.lsf import LSFEnvironment
 from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
@@ -27,20 +28,22 @@ _HANDLER: TypeAlias = Callable[[_SIGNUM, FrameType], Any] | int | signal.Handler
 
 
 def _resolve_requeue_signals():
-    signals: list[signal.Signals] = []
-
-    if timeout_signal_name := os.environ.get("NSHRUNNER_TIMEOUT_SIGNAL"):
-        signals.append(signal.Signals[timeout_signal_name])
-
-    if preempt_signal_name := os.environ.get("NSHRUNNER_PREEMPT_SIGNAL"):
-        signals.append(signal.Signals[preempt_signal_name])
+    if (session := nr.Session.from_current_session()) is None:
+        return None
 
+    signals: list[signal.Signals] = []
+    if session.submit_timeout_signal:
+        signals.append(session.submit_timeout_signal)
+    if session.submit_preempt_signal:
+        signals.append(session.submit_preempt_signal)
     return signals
 
 
 class _SignalConnector(_LightningSignalConnector):
-    def _auto_requeue_signals(self) -> list[signal.Signals]:
-        signals = _resolve_requeue_signals()
+    def _auto_requeue_signals(self) -> list[signal.Signals] | None:
+        if not (signals := _resolve_requeue_signals()):
+            return None
+
         signals_set = set(signals)
         valid_signals: set[signal.Signals] = signal.valid_signals()
         assert signals_set.issubset(

nshtrainer/trainer/trainer.py

@@ -3,7 +3,7 @@ import logging
 import os
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 from lightning.fabric.plugins.environments.lsf import LSFEnvironment
@@ -11,11 +11,13 @@ from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
 from lightning.pytorch import LightningModule
 from lightning.pytorch import Trainer as LightningTrainer
+from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.profilers import Profiler
+from lightning.pytorch.trainer.states import TrainerFn
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
 from typing_extensions import Unpack, assert_never, override
 
-from ..actsave import ActSave
+from .._checkpoint.metadata import _write_checkpoint_metadata
 from ..callbacks.base import resolve_all_callbacks
 from ..model.config import (
     AcceleratorConfigProtocol,
@@ -24,6 +26,7 @@ from ..model.config import (
     LightningTrainerKwargs,
     StrategyConfigProtocol,
 )
+from ._runtime_callback import RuntimeTrackerCallback, Stage
 from .signal_connector import _SignalConnector
 
 log = logging.getLogger(__name__)
@@ -168,12 +171,12 @@ class Trainer(LightningTrainer):
 
         if (accelerator := config.trainer.accelerator) is not None:
             if isinstance(accelerator, AcceleratorConfigProtocol):
-                accelerator = accelerator.construct_accelerator()
+                accelerator = accelerator.create_accelerator()
             _update_kwargs(accelerator=accelerator)
 
         if (strategy := config.trainer.strategy) is not None:
             if isinstance(strategy, StrategyConfigProtocol):
-                strategy = strategy.construct_strategy()
+                strategy = strategy.create_strategy()
             _update_kwargs(strategy=strategy)
 
         if (precision := config.trainer.precision) is not None:
@@ -220,7 +223,7 @@ class Trainer(LightningTrainer):
         if profiler := config.trainer.profiler:
             # If the profiler is an ProfilerConfig instance, then we instantiate it.
             if isinstance(profiler, BaseProfilerConfig):
-                profiler = profiler.construct_profiler(config)
+                profiler = profiler.create_profiler(config)
                 # Make sure that the profiler is an instance of `Profiler`.
                 if not isinstance(profiler, Profiler):
                     raise ValueError(f"{profiler=} is not an instance of `{Profiler}`.")
@@ -236,7 +239,7 @@ class Trainer(LightningTrainer):
         if plugin_configs := config.trainer.plugins:
             _update_kwargs(
                 plugins=[
-                    plugin_config.construct_plugin() for plugin_config in plugin_configs
+                    plugin_config.create_plugin() for plugin_config in plugin_configs
                 ]
             )
 
@@ -244,7 +247,7 @@ class Trainer(LightningTrainer):
             log.critical(f"Disabling logger because {config.trainer.logging.enabled=}.")
             kwargs["logger"] = False
         else:
-            _update_kwargs(logger=config.trainer.logging.construct_loggers(config))
+            _update_kwargs(logger=config.trainer.logging.create_loggers(config))
 
         if config.trainer.auto_determine_num_nodes:
             # When num_nodes is auto, we need to detect the number of nodes.
@@ -275,6 +278,9 @@ class Trainer(LightningTrainer):
 
         return kwargs
 
+    if TYPE_CHECKING:
+        callbacks: list[Callback]
+
     @override
     def __init__(
         self,
@@ -282,12 +288,14 @@ class Trainer(LightningTrainer):
         /,
         **kwargs: Unpack[LightningTrainerKwargs],
     ):
-        self._ll_config = config
         kwargs = self._update_kwargs(config, kwargs)
         log.critical(f"LightningTrainer.__init__ with {kwargs=}.")
 
         super().__init__(**kwargs)
 
+        # Add our own start time callback to measure the start time.
+        self.callbacks.append(RuntimeTrackerCallback())
+
         # Replace the signal connector with our own.
         self._signal_connector = _SignalConnector(self)
 
@@ -296,34 +304,89 @@ class Trainer(LightningTrainer):
             log_dir = str(Path(log_dir).resolve())
         log.critical(f"LightningTrainer log directory: {self.log_dir}.")
 
-        # Checkpoint loading
-        if (
-            ckpt_loading := self._ll_config.trainer.checkpoint_loading
-        ) and ckpt_loading.path:
-            self.ckpt_path = ckpt_loading.path
+    def __runtime_tracker(self):
+        return next(
+            (
+                callback
+                for callback in self.callbacks
+                if isinstance(callback, RuntimeTrackerCallback)
+            ),
+            None,
+        )
+
+    def __current_stage(self) -> Stage:
+        match self.state.fn:
+            case None:
+                raise ValueError(
+                    "Trainer state function is not set. "
+                    "You must call `fit`, `validate`, `test`, or `predict`, "
+                    "or explicitly provide a stage."
+                )
+            case TrainerFn.FITTING:
+                return "train"
+            case TrainerFn.VALIDATING:
+                return "validate"
+            case TrainerFn.TESTING:
+                return "test"
+            case TrainerFn.PREDICTING:
+                return "predict"
+            case _:
+                assert_never(self.state.fn)
+
+    def start_time(self, stage: Stage | None = None):
+        """Return the start time of the run"""
+        if (tracker := self.__runtime_tracker()) is None:
+            raise ValueError(
+                "RuntimeTrackerCallback is not set. Cannot get start time."
+            )
+        if stage is None:
+            stage = self.__current_stage()
 
-    @contextlib.contextmanager
-    def _actsave_context(self, model: LightningModule):
-        hparams = cast(BaseConfig, model.hparams)
-        if not (actsave_config := hparams.trainer.actsave):
-            yield
-            return
+        return tracker.start_time(stage)
+
+    def end_time(self, stage: Stage | None = None):
+        """Return the end time of the run"""
+        if (tracker := self.__runtime_tracker()) is None:
+            raise ValueError(
+                "RuntimeTrackerCallback is not set. Cannot get start time."
+            )
+        if stage is None:
+            stage = self.__current_stage()
+
+        return tracker.end_time(stage)
+
+    def time_elapsed(self, stage: Stage | None = None):
+        """Return the time elapsed for the run"""
+        if (tracker := self.__runtime_tracker()) is None:
+            raise ValueError(
+                "RuntimeTrackerCallback is not set. Cannot get start time."
+            )
+        if stage is None:
+            stage = self.__current_stage()
 
-        # Enter actsave context
-        with ActSave.enabled(actsave_config.resolve_save_dir(hparams)):
-            yield
+        return tracker.time_elapsed(stage)
+
+    @property
+    def _base_module(self):
+        if self.lightning_module is None:
+            raise ValueError("LightningModule is not set.")
+
+        from ..model.base import LightningModuleBase
+
+        if not isinstance(self.lightning_module, LightningModuleBase):
+            raise ValueError(
+                f"LightningModule is not an instance of {LightningModuleBase}."
+            )
+
+        return self.lightning_module
 
     @override
     def _run(
         self, model: LightningModule, ckpt_path: str | Path | None = None
     ) -> _EVALUATE_OUTPUT | _PREDICT_OUTPUT | None:
-        """
-        Two things done here:
-            1. Lightning doesn't support gradient clipping with manual optimization.
-            We patch the `Trainer._run` method to throw if gradient clipping is enabled
-            and `model.automatic_optimization` is False.
-
-            2. We actually set up actsave here.
+        """Lightning doesn't support gradient clipping with manual optimization.
+        We patch the `Trainer._run` method to throw if gradient clipping is enabled
+        and `model.automatic_optimization` is False.
         """
 
         if not model.automatic_optimization and (
@@ -336,5 +399,22 @@ class Trainer(LightningTrainer):
                 "or disable automatic gradient clipping. "
             )
 
-        with self._actsave_context(model):
-            return super()._run(model, ckpt_path)
+        return super()._run(model, ckpt_path)
+
+    @override
+    def save_checkpoint(
+        self,
+        filepath: str | Path,
+        weights_only: bool = False,
+        storage_options: Any | None = None,
+    ):
+        filepath = Path(filepath)
+        ret_val = super().save_checkpoint(filepath, weights_only, storage_options)
+
+        # Save the checkpoint metadata
+        lm = self._base_module
+        if lm.config.trainer.save_checkpoint_metadata and self.is_global_zero:
+            # Generate the metadata and write to disk
+            _write_checkpoint_metadata(self, lm, filepath)
+
+        return ret_val

{nshtrainer-0.9.1.dist-info → nshtrainer-0.10.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.9.1
+Version: 0.10.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -9,11 +9,13 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: GitPython
 Requires-Dist: lightning
 Requires-Dist: nshconfig
 Requires-Dist: nshrunner
 Requires-Dist: nshutils
 Requires-Dist: numpy
+Requires-Dist: psutil
 Requires-Dist: pytorch-lightning
 Requires-Dist: torch
 Requires-Dist: torchmetrics

{nshtrainer-0.9.1.dist-info → nshtrainer-0.10.1.dist-info}/RECORD

@@ -1,32 +1,34 @@
-nshtrainer/__init__.py,sha256=nbZHdfTk0oWqsJgrSzdgk2DSf4CGhdZn79esoJGauO8,548
+nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
+nshtrainer/_checkpoint/loader.py,sha256=48flPr1XgQHOgIPaCrRqOEvRuG0SZuV3cQ1vgHLqFqI,11025
+nshtrainer/_checkpoint/metadata.py,sha256=C7je_soYyEbZjiq7p2_pSVFkgcXnz2J2H5sMy8oskx0,3051
 nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ4b-PoWM,120
 nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
 nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
 nshtrainer/_experimental/flops/module_tracker.py,sha256=bUL-IRTd0aF_DwmXkZjHZAA31p4ZEhyqhc26XWKQUUY,4922
-nshtrainer/actsave/__init__.py,sha256=_ZuwgRtF1-ekouXNvtZCAS1g_IDYGB4NX8BFSGNGBT8,119
-nshtrainer/actsave/_callback.py,sha256=mnHOtuG9vtHEzz9q4vCvDNC6VvjZsgb4MSSuOoUDh3M,2778
-nshtrainer/callbacks/__init__.py,sha256=I6W33ityL9Ko8jjqHh3WH_8miV59SAe9LxInhoqX5XE,1665
+nshtrainer/callbacks/__init__.py,sha256=ifXQRwtccznl4lMKwKLSuuAQC4bKFBgfzQ4rx9gOqjE,2345
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
-nshtrainer/callbacks/base.py,sha256=LrcRUV02bZEKXRIRvhHT9qsvw_kwoWiAdQkVMyKc5NU,3542
+nshtrainer/callbacks/actsave.py,sha256=aY6T_NAzaFAVU8WMHOXnWL5wd2bi8eVxeU2S0iAs70c,4446
+nshtrainer/callbacks/base.py,sha256=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,3555
 nshtrainer/callbacks/early_stopping.py,sha256=jriSU761wf_qTJ9Bos0D3h5aDvZHYpRqK62Ne8aWp5I,3768
-nshtrainer/callbacks/ema.py,sha256=zKCtvzZFo0ORlwNZHjaMk-sJoxrlTtFWOzR-yGy95W0,12134
-nshtrainer/callbacks/finite_checks.py,sha256=kX3TIJsxyqx0GuLJfYsqVgKU27zwjG9Z8324lyCFtwM,2087
-nshtrainer/callbacks/gradient_skipping.py,sha256=ModaIXpb69LbA8TpEXKRLdr4Sq7-l0CWnN6fvpaV188,3477
+nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
+nshtrainer/callbacks/finite_checks.py,sha256=AO5fa51uANAjAkeJfTquOjK6W_4RSU5Kky3f5jmAPlQ,2084
+nshtrainer/callbacks/gradient_skipping.py,sha256=fSJpjgHbztFKz7w3qFuCHZpmbEt9BCLAy-sU0B4xJQI,3474
 nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvLfaw,8080
-nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=ZT0bn7X0BZbQXbk6fos47NsbbhD4Z9c9YmFqdcUEqus,1503
+nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=p0zeDK3PLWWl485e9o08ywEEARCfuZ5it47tNCtR4ec,2838
 nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
-nshtrainer/callbacks/norm_logging.py,sha256=IMrK0WiVSDFyspwyPpwELMK4mmd5Jpx4enAW_GsWbi4,6284
-nshtrainer/callbacks/on_exception_checkpoint.py,sha256=eDyB7qkpPdAaKjAY2uFMMY8Nht6TGeuDnsgHuKtp8eA,1615
-nshtrainer/callbacks/print_table.py,sha256=FcA-CBWwMf9c1NNRinvYpZC400RNQxuP28bJfgniT3Q,2840
-nshtrainer/callbacks/throughput_monitor.py,sha256=YQLdpX3LGybIiD814yT9yCCVSEXRWf8WwsvVaN5aDBE,1848
-nshtrainer/callbacks/timer.py,sha256=sDXPPcdDKu5xnuK_bjr8plIq9MBuluNJ42Mt9LvPZzc,4610
-nshtrainer/callbacks/wandb_watch.py,sha256=pUpMsNxd03ex1rzOmFw2HzGOXjnQGaH84m8cc2dXo4g,2937
+nshtrainer/callbacks/model_checkpoint.py,sha256=4zYycpXHGRyL4svWLP6GmG3WJs5m3B5PRCOzXC3m_qg,5955
+nshtrainer/callbacks/norm_logging.py,sha256=EWyrfkp8iHjQi9iAAXHxb0xStw2RwkdpKG2_gLarQRA,6281
+nshtrainer/callbacks/on_exception_checkpoint.py,sha256=zna_QF_x4HwD7Es5XxrHLDED43NU1GpcDNoL139HEOs,3355
+nshtrainer/callbacks/print_table.py,sha256=_FdAHhqylWGk4Z0c2FrLFeiMA4jhfA_beZRK_BHpzmE,2837
+nshtrainer/callbacks/throughput_monitor.py,sha256=4EF3b79HdHiRgBGIFDyD4O1oywb5h1tV8nml7NuuDjU,1845
+nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
+nshtrainer/callbacks/wandb_watch.py,sha256=bicXS3nZfPGoN7Owu1XIBS-1bw7yeIJdYJTnRN0dp2E,2934
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=bcJBcQjh1hB1yKF_xSlT9AtEWv0BJjYc1CuH2BF-ea8,4392
 nshtrainer/data/transform.py,sha256=JeGxvytQly8hougrsdMmKG8gJ6qvFPDglJCO4Tp6STk,1795
-nshtrainer/ll/__init__.py,sha256=nxYPtoFOFAvzkD6O3EIuwCiRi_LedYa_EH-RIfDG91s,2685
+nshtrainer/ll/__init__.py,sha256=dD0ISxHJ2lg1HLSM0b3db7TBlsPpQCtChnuYO-c2oqI,2635
 nshtrainer/ll/_experimental.py,sha256=oBQCKOEVYoxuUU9eLb-Fg2B2mzZD7SA0zfAO6lmWZ88,53
-nshtrainer/ll/actsave.py,sha256=QJ7yJIqvabpZzumX7PLPzkh6dfqY-zxiEdzv48VtZEY,123
+nshtrainer/ll/actsave.py,sha256=2lbiseSrjcwFT6AiyLNWarTWl1bnzliVWlu1iOfnP30,209
 nshtrainer/ll/callbacks.py,sha256=AxyUmc8aGRSjx6WwwgXYCmdJ73rwLuEAEH0AGRosojQ,49
 nshtrainer/ll/config.py,sha256=fKumJf42HY2FITX1QUM1OTXkYD6U2np2ciyd4PFRPZ8,145
 nshtrainer/ll/data.py,sha256=zRG0FRje-jtSHximVzkHIHzpwsyQxpHCoACFihNKLPM,44
@@ -45,13 +47,16 @@ nshtrainer/lr_scheduler/__init__.py,sha256=uEvgaFAs-4s_bAEMaildy0GT6OvgpgOEKTuzq
 nshtrainer/lr_scheduler/_base.py,sha256=7xOIuxQ86YHbFWG5a3gX46emQj1WN_LaY4-i0Q1TDBg,3659
 nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=mn6cyizyI_stkXtg6zxIEGF9btIxMRWigUHUTlUYCSw,5221
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-WKArwxxf7ZQizhmbCw,2782
-nshtrainer/model/__init__.py,sha256=y32Hla-5whpzLL2BtCJpBakSp8o-1nQbpO0j_-xq_Po,1864
-nshtrainer/model/base.py,sha256=YtqnjiMf0cLVjFEQuOLm5WwCkVnZftiHlIdCrxdax3s,21297
-nshtrainer/model/config.py,sha256=-I_HLTTwqWimnnoKJ64oBEq3x31CZj9rwrg9MnFzs38,68215
+nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
+nshtrainer/metrics/_config.py,sha256=hWWS4IXENRyH3RmJ7z1Wx1n3Lt1sNMlGOrcU6PW15o0,1104
+nshtrainer/model/__init__.py,sha256=TbexTxiE20WHYg5q3L88Hysk4LlHeKk_isv33aSBREA,1918
+nshtrainer/model/_environment.py,sha256=s3JFnigbssFRJTwH33K7DcAYVhLOFCC1OZgFNXJgjuw,22317
+nshtrainer/model/base.py,sha256=Bmw-t70TydDbE9P0ee-lTibGoUhrCx5Qke-upa7FGVM,17512
+nshtrainer/model/config.py,sha256=B1XkKYbhpAm6RmF4n4eR66hMh-kCXwIQB2pQuhR9TZE,53177
 nshtrainer/model/modules/callback.py,sha256=JF59U9-CjJsAIspEhTJbVaGN0wGctZG7UquE3IS7R8A,6408
 nshtrainer/model/modules/debug.py,sha256=DTVty8cKnzj1GCULRyGx_sWTTsq9NLi30dzqjRTnuCU,1127
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
-nshtrainer/model/modules/logger.py,sha256=XEeo3QrplTNKZqfl6iWZf3fze3R4YOeOvs-RKVHFoQs,5527
+nshtrainer/model/modules/logger.py,sha256=YYhehQysqTjuVFcd_EREYDh57CIlezidFBS2Ohp_xKo,5661
 nshtrainer/model/modules/profiler.py,sha256=rQ_jRMcM1Z2AIROZlRnBRHM5rkTpq67afZPD6CIRfXs,825
 nshtrainer/model/modules/rlp_sanity_checks.py,sha256=o6gUceFwsuDHmL8eLOYuT3JGXFzq_qc4awl2RWaBygU,8900
 nshtrainer/model/modules/shared_parameters.py,sha256=mD5wrlBE3c025vzVdTpnSyC8yxzuI-aUWMmPhqPT0a0,2694
@@ -61,17 +66,19 @@ nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,
 nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
 nshtrainer/nn/nonlinearity.py,sha256=owtU4kh4G98psD0axOJWVfBhm-OtJVgFM-TXSHmbNPU,3625
 nshtrainer/optimizer.py,sha256=kuJEA1pvB3y1FcsfhAoOJujVqEZqFHlmYO8GW6JeA1g,1527
-nshtrainer/runner.py,sha256=7EumpnBkdNWjSNT9Gm-pkxAJ3W6-iMC-yae-WNeZcLw,3771
+nshtrainer/runner.py,sha256=6qfE5FBONzD79kVHuWYKEvK0J_Qi5dMBbHQhRMmnIhE,3649
 nshtrainer/scripts/check_env.py,sha256=IMl6dSqsLYppI0XuCsVq8lK4bYqXwY9KHJkzsShz4Kg,806
 nshtrainer/scripts/find_packages.py,sha256=FbdlfmAefttFSMfaT0A46a-oHLP_ioaQKihwBfBeWeA,1467
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
-nshtrainer/trainer/signal_connector.py,sha256=JSP8W2PSdzwO3iWX1WOL1l8dufh2dKgUWeJ2gEWCppg,10626
-nshtrainer/trainer/trainer.py,sha256=eYEYfY9v70MuorHcSf8nqM7f2CkmUHhpPcjCk4FJD7k,14034
+nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
+nshtrainer/trainer/checkpoint_connector.py,sha256=xoqI2dcPnlNFPPLVIU6dBOvRPC9PtfX5qu__xV1lx0Y,2124
+nshtrainer/trainer/signal_connector.py,sha256=llwc8pdKAWxREFpjdi14Bpy8rGVMEJsmJx_s2p4gI8E,10689
+nshtrainer/trainer/trainer.py,sha256=n3T9Iz3eaDostxEdjapWImAsVMxyU9WBdhlPl0THX-g,16785
 nshtrainer/util/environment.py,sha256=_SEtiQ_s5bL5pllUlf96AOUv15kNvCPvocVC13S7mIk,4166
 nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.9.1.dist-info/METADATA,sha256=3s9luSztUNVhu3t_sSmOw3HhwuVVUoiLhQwlxBiaaSg,647
-nshtrainer-0.9.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.9.1.dist-info/RECORD,,
+nshtrainer-0.10.1.dist-info/METADATA,sha256=O8wMPb0ksoZajyes8dsq4IIjsfP_jQaxGYpW3rYE9Ro,695
+nshtrainer-0.10.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.10.1.dist-info/RECORD,,

nshtrainer/actsave/__init__.py

@@ -1,3 +0,0 @@
-from nshutils.actsave import *  # type: ignore # noqa: F403
-
-from ._callback import ActSaveCallback as ActSaveCallback