nshtrainer 0.8.7__py3-none-any.whl → 0.10.0__py3-none-any.whl

nshtrainer/__init__.py

@@ -2,13 +2,14 @@ from . import _experimental as _experimental
 from . import callbacks as callbacks
 from . import data as data
 from . import lr_scheduler as lr_scheduler
+from . import metrics as metrics
 from . import model as model
 from . import nn as nn
 from . import optimizer as optimizer
+from .metrics import MetricConfig as MetricConfig
 from .model import Base as Base
 from .model import BaseConfig as BaseConfig
 from .model import ConfigList as ConfigList
 from .model import LightningModuleBase as LightningModuleBase
-from .model import MetricConfig as MetricConfig
 from .runner import Runner as Runner
 from .trainer import Trainer as Trainer

nshtrainer/callbacks/__init__.py

@@ -14,15 +14,27 @@ from .interval import EpochIntervalCallback as EpochIntervalCallback
 from .interval import IntervalCallback as IntervalCallback
 from .interval import StepIntervalCallback as StepIntervalCallback
 from .latest_epoch_checkpoint import LatestEpochCheckpoint as LatestEpochCheckpoint
+from .latest_epoch_checkpoint import (
+    LatestEpochCheckpointCallbackConfig as LatestEpochCheckpointCallbackConfig,
+)
 from .log_epoch import LogEpochCallback as LogEpochCallback
+from .model_checkpoint import ModelCheckpoint as ModelCheckpoint
+from .model_checkpoint import (
+    ModelCheckpointCallbackConfig as ModelCheckpointCallbackConfig,
+)
 from .norm_logging import NormLoggingCallback as NormLoggingCallback
 from .norm_logging import NormLoggingConfig as NormLoggingConfig
 from .on_exception_checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
+from .on_exception_checkpoint import (
+    OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+)
 from .print_table import PrintTableMetricsCallback as PrintTableMetricsCallback
 from .print_table import PrintTableMetricsConfig as PrintTableMetricsConfig
 from .throughput_monitor import ThroughputMonitorConfig as ThroughputMonitorConfig
 from .timer import EpochTimer as EpochTimer
 from .timer import EpochTimerConfig as EpochTimerConfig
+from .wandb_watch import WandbWatchCallback as WandbWatchCallback
+from .wandb_watch import WandbWatchConfig as WandbWatchConfig
 
 CallbackConfig = Annotated[
     ThroughputMonitorConfig
@@ -31,6 +43,10 @@ CallbackConfig = Annotated[
     | FiniteChecksConfig
     | NormLoggingConfig
     | GradientSkippingConfig
-    | EMAConfig,
+    | EMAConfig
+    | ModelCheckpointCallbackConfig
+    | LatestEpochCheckpointCallbackConfig
+    | OnExceptionCheckpointCallbackConfig
+    | WandbWatchConfig,
     C.Field(discriminator="name"),
 ]

nshtrainer/{actsave/_callback.py → callbacks/actsave.py}

@@ -1,28 +1,87 @@
 import contextlib
-from typing import TYPE_CHECKING, Literal, cast
+from pathlib import Path
+from typing import Literal
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks.callback import Callback
-from nshutils.actsave import ActSave
 from typing_extensions import TypeAlias, override
 
-if TYPE_CHECKING:
-    from ..model.config import BaseConfig
+from .base import CallbackConfigBase
+
+try:
+    from nshutils import ActSave  # type: ignore
+except ImportError:
+    ActSave = None
 
 Stage: TypeAlias = Literal["train", "validation", "test", "predict"]
 
 
+class ActSaveConfig(CallbackConfigBase):
+    enabled: bool = True
+    """Enable activation saving."""
+
+    save_dir: Path | None = None
+    """Directory to save activations to. If None, will use the activation directory set in `config.directory`."""
+
+    def __bool__(self):
+        return self.enabled
+
+    @override
+    def create_callbacks(self, root_config):
+        yield ActSaveCallback(
+            self,
+            self.save_dir
+            or root_config.directory.resolve_subdirectory(root_config.id, "activation"),
+        )
+
+
 class ActSaveCallback(Callback):
-    def __init__(self):
+    def __init__(self, config: ActSaveConfig, save_dir: Path):
         super().__init__()
 
+        self.config = config
+        self.save_dir = save_dir
+        self._enabled_context: contextlib._GeneratorContextManager | None = None
         self._active_contexts: dict[Stage, contextlib._GeneratorContextManager] = {}
 
+    @override
+    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
+        super().setup(trainer, pl_module, stage)
+
+        if not self.config:
+            return
+
+        if ActSave is None:
+            raise ImportError(
+                "ActSave is not installed. Please install nshutils to use the ActSaveCallback."
+            )
+
+        context = ActSave.enabled(self.save_dir)
+        context.__enter__()
+        self._enabled_context = context
+
+    @override
+    def teardown(
+        self, trainer: Trainer, pl_module: LightningModule, stage: str
+    ) -> None:
+        super().teardown(trainer, pl_module, stage)
+
+        if not self.config:
+            return
+
+        if self._enabled_context is not None:
+            self._enabled_context.__exit__(None, None, None)
+            self._enabled_context = None
+
     def _on_start(self, stage: Stage, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast("BaseConfig", pl_module.hparams)
-        if not hparams.trainer.actsave:
+        if not self.config:
             return
 
+        if ActSave is None:
+            raise ImportError(
+                "ActSave is not installed. Please install nshutils to use the ActSaveCallback."
+            )
+
         # 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)
@@ -33,12 +92,11 @@ class ActSaveCallback(Callback):
         self._active_contexts[stage] = context
 
     def _on_end(self, stage: Stage, trainer: Trainer, pl_module: LightningModule):
-        hparams = cast("BaseConfig", pl_module.hparams)
-        if not hparams.trainer.actsave:
+        if not self.config:
             return
 
         # If we have an active context manager for this stage, exit it
-        if active_contexts := self._active_contexts.get(stage):
+        if (active_contexts := self._active_contexts.pop(stage, None)) is not None:
             active_contexts.__exit__(None, None, None)
 
     @override

nshtrainer/callbacks/base.py

@@ -46,16 +46,16 @@ class CallbackConfigBase(C.Config, ABC):
         )
 
     @abstractmethod
-    def construct_callbacks(
+    def create_callbacks(
         self, root_config: "BaseConfig"
     ) -> Iterable[Callback | CallbackWithMetadata]: ...
 
 
 # region Config resolution helpers
-def _construct_callbacks_with_metadata(
+def _create_callbacks_with_metadata(
     config: CallbackConfigBase, root_config: "BaseConfig"
 ) -> Iterable[CallbackWithMetadata]:
-    for callback in config.construct_callbacks(root_config):
+    for callback in config.create_callbacks(root_config):
         if isinstance(callback, CallbackWithMetadata):
             yield callback
             continue
@@ -99,12 +99,14 @@ def _process_and_filter_callbacks(
 
 def resolve_all_callbacks(root_config: "BaseConfig"):
     callback_configs = [
-        config for config in root_config.ll_all_callback_configs() if config is not None
+        config
+        for config in root_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 _construct_callbacks_with_metadata(callback_config, root_config)
+        for callback in _create_callbacks_with_metadata(callback_config, root_config)
     )
     return callbacks
 

nshtrainer/callbacks/ema.py

@@ -374,7 +374,7 @@ class EMAConfig(CallbackConfigBase):
     """Offload weights to CPU."""
 
     @override
-    def construct_callbacks(self, root_config):
+    def create_callbacks(self, root_config):
         yield EMA(
             decay=self.decay,
             validate_original_weights=self.validate_original_weights,

nshtrainer/callbacks/finite_checks.py

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

nshtrainer/callbacks/gradient_skipping.py

@@ -99,5 +99,5 @@ class GradientSkippingConfig(CallbackConfigBase):
     """
 
     @override
-    def construct_callbacks(self, root_config):
+    def create_callbacks(self, root_config):
         yield GradientSkipping(self)

nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -1,35 +1,54 @@
 import logging
 from pathlib import Path
+from typing import Literal
 
-from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Checkpoint
 from typing_extensions import override
 
+from .base import CallbackConfigBase
+
 log = logging.getLogger(__name__)
 
 
+class LatestEpochCheckpointCallbackConfig(CallbackConfigBase):
+    kind: Literal["latest_epoch_checkpoint"] = "latest_epoch_checkpoint"
+
+    dirpath: str | Path | None = None
+    """Directory path to save the checkpoint file."""
+
+    filename: str = "latest_epoch{epoch:02d}_step{step:04d}.ckpt"
+    """Checkpoint filename. This must not include the extension."""
+
+    save_weights_only: bool = False
+    """Whether to save only the model's weights or the entire model object."""
+
+    latest_symlink_filename: str | None = "latest.ckpt"
+    """Filename for the latest symlink. If None, no symlink will be created."""
+
+    @override
+    def create_callbacks(self, root_config):
+        dirpath = self.dirpath or root_config.directory.resolve_subdirectory(
+            root_config.id, "checkpoint"
+        )
+        dirpath = Path(dirpath)
+
+        yield LatestEpochCheckpoint(self, dirpath)
+
+
 class LatestEpochCheckpoint(Checkpoint):
-    DEFAULT_FILENAME = "latest_epoch{epoch:02d}_step{step:04d}.ckpt"
-
-    def __init__(
-        self,
-        dirpath: _PATH,
-        filename: str | None = None,
-        save_weights_only: bool = False,
-    ):
+    def __init__(self, config: LatestEpochCheckpointCallbackConfig, dirpath: Path):
         super().__init__()
 
-        self._dirpath = Path(dirpath)
-        self._filename = filename or self.DEFAULT_FILENAME
-        self._save_weights_only = save_weights_only
+        self.config = config
+        self.dirpath = dirpath
 
         # Also, we hold a reference to the last checkpoint path
         # to be able to remove it when a new checkpoint is saved.
         self._last_ckpt_path: Path | None = None
 
     def _ckpt_path(self, trainer: Trainer):
-        return self._dirpath / self._filename.format(
+        return self.dirpath / self.config.filename.format(
             epoch=trainer.current_epoch, step=trainer.global_step
         )
 
@@ -41,5 +60,22 @@ class LatestEpochCheckpoint(Checkpoint):
 
         # Save the new checkpoint
         filepath = self._ckpt_path(trainer)
-        trainer.save_checkpoint(filepath, self._save_weights_only)
+        trainer.save_checkpoint(filepath, self.config.save_weights_only)
         self._last_ckpt_path = filepath
+
+        # Create the latest symlink
+        if (symlink_filename := self.config.latest_symlink_filename) is not None:
+            symlink_path = self.dirpath / symlink_filename
+            if symlink_path.exists():
+                symlink_path.unlink()
+            symlink_path.symlink_to(filepath.name)
+            log.info(f"Created latest symlink: {symlink_path}")
+
+    def latest_checkpoint(self):
+        if (symlink_filename := self.config.latest_symlink_filename) is None:
+            return None
+
+        if not (symlink_path := self.dirpath / symlink_filename).exists():
+            return None
+
+        return symlink_path

nshtrainer/callbacks/model_checkpoint.py

@@ -0,0 +1,187 @@
+import re
+from datetime import timedelta
+from logging import getLogger
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+from lightning.pytorch.callbacks.model_checkpoint import (
+    ModelCheckpoint as _ModelCheckpoint,
+)
+from typing_extensions import override
+
+from ..metrics import MetricConfig
+from .base import CallbackConfigBase
+
+if TYPE_CHECKING:
+    from ..model.config import BaseConfig
+
+log = getLogger(__name__)
+
+
+def _convert_string(input_string: str):
+    # Find all variables enclosed in curly braces
+    variables = re.findall(r"\{(.*?)\}", input_string)
+
+    # Replace each variable with its corresponding key-value pair
+    output_string = input_string
+    for variable in variables:
+        # If the name is something like {variable:format}, we shouldn't process the format.
+        key_name = variable
+        if ":" in variable:
+            key_name, _ = variable.split(":", 1)
+            continue
+
+        # Replace '/' with '_' in the key name
+        key_name = key_name.replace("/", "_")
+        output_string = output_string.replace(
+            f"{{{variable}}}", f"{key_name}={{{variable}}}"
+        )
+
+    return output_string
+
+
+class ModelCheckpointCallbackConfig(CallbackConfigBase):
+    """Arguments for the ModelCheckpoint callback."""
+
+    kind: Literal["model_checkpoint"] = "model_checkpoint"
+
+    dirpath: str | Path | None = None
+    """
+    Directory path to save the model file. If `None`, we save to the checkpoint directory set in `config.directory`.
+    """
+
+    filename: str | None = None
+    """
+    Checkpoint filename.
+        If None, a default template is used (see :attr:`ModelCheckpoint.CHECKPOINT_JOIN_CHAR`).
+    """
+
+    metric: MetricConfig | None = None
+    """
+    Metric to monitor for saving checkpoints.
+        If None, the primary metric of the runner will be used, if available.
+    """
+
+    verbose: bool = False
+    """Verbosity mode. If True, print additional information about checkpoints."""
+
+    save_last: Literal[True, False, "link"] | None = "link"
+    """
+    Whether to save the last checkpoint.
+        If True, saves a copy of the last checkpoint separately.
+        If "link", creates a symbolic link to the last checkpoint.
+    """
+
+    save_top_k: int = 1
+    """
+    Number of best models to save.
+        If -1, all models are saved.
+        If 0, no models are saved.
+    """
+
+    save_weights_only: bool = False
+    """Whether to save only the model's weights or the entire model object."""
+
+    auto_insert_metric_name: bool = True
+    """Whether to automatically insert the metric name in the checkpoint filename."""
+
+    every_n_train_steps: int | None = None
+    """
+    Number of training steps between checkpoints.
+        If None or 0, no checkpoints are saved during training.
+    """
+
+    train_time_interval: timedelta | None = None
+    """
+    Time interval between checkpoints during training.
+        If None, no checkpoints are saved during training based on time.
+    """
+
+    every_n_epochs: int | None = None
+    """
+    Number of epochs between checkpoints.
+        If None or 0, no checkpoints are saved at the end of epochs.
+    """
+
+    save_on_train_epoch_end: bool | None = None
+    """
+    Whether to run checkpointing at the end of the training epoch.
+        If False, checkpointing runs at the end of the validation.
+    """
+
+    enable_version_counter: bool = True
+    """Whether to append a version to the existing file name."""
+
+    auto_append_metric: bool = True
+    """If enabled, this will automatically add "-{monitor}" to the filename."""
+
+    def metric_or_default(self, root_config: "BaseConfig"):
+        if self.metric is not None:
+            return self.metric
+        if root_config.primary_metric is not None:
+            return root_config.primary_metric
+        raise ValueError("Primary metric must be provided if metric is not specified.")
+
+    def resolve_filename(self, root_config: "BaseConfig"):
+        metric = self.metric_or_default(root_config)
+
+        filename = self.filename
+        if not filename:
+            filename = "{epoch}-{step}"
+        if self.auto_append_metric:
+            filename = f"{filename}-{{{metric.validation_monitor}}}"
+
+        if self.auto_insert_metric_name and filename:
+            new_filename = _convert_string(filename)
+            log.critical(
+                f"Updated ModelCheckpoint filename: {filename} -> {new_filename}"
+            )
+            filename = new_filename
+
+        return filename
+
+    @override
+    def create_callbacks(self, root_config):
+        dirpath = self.dirpath or root_config.directory.resolve_subdirectory(
+            root_config.id, "checkpoint"
+        )
+
+        metric = self.metric_or_default(root_config)
+        filename = self.resolve_filename(root_config)
+
+        yield ModelCheckpoint(
+            self,
+            dirpath=Path(dirpath),
+            filename=filename,
+            metric=metric,
+        )
+
+
+class ModelCheckpoint(_ModelCheckpoint):
+    @override
+    def __init__(
+        self,
+        config: ModelCheckpointCallbackConfig,
+        dirpath: Path,
+        filename: str,
+        metric: MetricConfig,
+    ):
+        self.config = config
+        del config
+
+        super().__init__(
+            dirpath=dirpath,
+            filename=filename,
+            monitor=metric.validation_monitor,
+            mode=metric.mode,
+            verbose=self.config.verbose,
+            save_last=self.config.save_last,
+            save_top_k=self.config.save_top_k,
+            save_weights_only=self.config.save_weights_only,
+            auto_insert_metric_name=False,
+            every_n_train_steps=self.config.every_n_train_steps,
+            train_time_interval=self.config.train_time_interval,
+            every_n_epochs=self.config.every_n_epochs,
+            save_on_train_epoch_end=self.config.save_on_train_epoch_end,
+            enable_version_counter=self.config.enable_version_counter,
+        )

nshtrainer/callbacks/norm_logging.py

@@ -180,7 +180,7 @@ class NormLoggingConfig(CallbackConfigBase):
         )
 
     @override
-    def construct_callbacks(self, root_config):
+    def create_callbacks(self, root_config):
         if not self:
             return
 

nshtrainer/callbacks/on_exception_checkpoint.py

@@ -1,16 +1,82 @@
+import contextlib
 import datetime
 import logging
 import os
-from typing import Any
+from pathlib import Path
+from typing import Any, Literal
 
-from lightning.pytorch import Trainer
+from lightning.pytorch import Trainer as LightningTrainer
 from lightning.pytorch.callbacks import OnExceptionCheckpoint as _OnExceptionCheckpoint
 from typing_extensions import override
 
+from .base import CallbackConfigBase
+
 log = logging.getLogger(__name__)
 
 
+@contextlib.contextmanager
+def _monkey_patch_disable_barrier(trainer: LightningTrainer):
+    """
+    Monkey-patch the strategy instance to make the barrier operation a no-op.
+
+    We do this because `save_checkpoint` calls `barrier`. This is okay in most
+    cases, but when we want to save a checkpoint in the case of an exception,
+    `barrier` causes a deadlock. So we monkey-patch the strategy instance to
+    make the barrier operation a no-op.
+    """
+
+    # We monkey-patch the barrier method to do nothing.
+    original_barrier = trainer.strategy.barrier
+
+    def new_barrier(*args, **kwargs):
+        log.warning("Monkey-patched no-op barrier.")
+        pass
+
+    trainer.strategy.barrier = new_barrier
+    log.warning("Monkey-patched barrier to no-op.")
+
+    try:
+        yield
+    finally:
+        trainer.strategy.barrier = original_barrier
+        log.warning("Reverted monkey-patched barrier.")
+
+
+class OnExceptionCheckpointCallbackConfig(CallbackConfigBase):
+    kind: Literal["on_exception_checkpoint"] = "on_exception_checkpoint"
+
+    dirpath: str | Path | None = None
+    """Directory path to save the checkpoint file."""
+
+    filename: str | None = None
+    """Checkpoint filename. This must not include the extension. If `None`, `on_exception_{id}_{timestamp}` is used."""
+
+    @override
+    def create_callbacks(self, root_config):
+        from ..callbacks.on_exception_checkpoint import OnExceptionCheckpoint
+
+        dirpath = self.dirpath or root_config.directory.resolve_subdirectory(
+            root_config.id, "checkpoint"
+        )
+
+        if not (filename := self.filename):
+            filename = f"on_exception_{root_config.id}"
+        yield OnExceptionCheckpoint(self, dirpath=Path(dirpath), filename=filename)
+
+
 class OnExceptionCheckpoint(_OnExceptionCheckpoint):
+    @override
+    def __init__(
+        self,
+        config: OnExceptionCheckpointCallbackConfig,
+        dirpath: Path,
+        filename: str,
+    ):
+        self.config = config
+        del config
+
+        super().__init__(dirpath, filename)
+
     @property
     @override
     def ckpt_path(self) -> str:
@@ -22,23 +88,11 @@ class OnExceptionCheckpoint(_OnExceptionCheckpoint):
         return f"{ckpt_path}_{timestamp}{ext}"
 
     @override
-    def on_exception(self, trainer: Trainer, *_: Any, **__: Any) -> None:
-        # We override this to checkpoint the model manually,
-        # without calling the dist barrier.
-
-        # trainer.save_checkpoint(self.ckpt_path)
-
-        if trainer.model is None:
-            raise AttributeError(
-                "Saving a checkpoint is only possible if a model is attached to the Trainer. Did you call"
-                " `Trainer.save_checkpoint()` before calling `Trainer.{fit,validate,test,predict}`?"
-            )
-        checkpoint = trainer._checkpoint_connector.dump_checkpoint(weights_only=False)
-        trainer.strategy.save_checkpoint(
-            checkpoint, self.ckpt_path, storage_options=None
-        )
-        # self.strategy.barrier("Trainer.save_checkpoint") # <-- This is disabled
-
-    @override
-    def teardown(self, trainer: Trainer, *_: Any, **__: Any) -> None:
-        trainer.strategy.remove_checkpoint(self.ckpt_path)
+    def on_exception(self, trainer: LightningTrainer, *args: Any, **kwargs: Any):
+        # Monkey-patch the strategy instance to make the barrier operation a no-op.
+        # We do this because `save_checkpoint` calls `barrier`. This is okay in most
+        # cases, but when we want to save a checkpoint in the case of an exception,
+        # `barrier` causes a deadlock. So we monkey-patch the strategy instance to
+        # make the barrier operation a no-op.
+        with _monkey_patch_disable_barrier(trainer):
+            return super().on_exception(trainer, *args, **kwargs)

nshtrainer/callbacks/print_table.py

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

nshtrainer/callbacks/throughput_monitor.py

@@ -52,5 +52,5 @@ class ThroughputMonitorConfig(CallbackConfigBase):
     """Number of batches to use for a rolling average."""
 
     @override
-    def construct_callbacks(self, root_config):
+    def create_callbacks(self, root_config):
         yield ThroughputMonitor(window_size=self.window_size)

nshtrainer/callbacks/timer.py

@@ -153,5 +153,5 @@ class EpochTimerConfig(CallbackConfigBase):
     name: Literal["epoch_timer"] = "epoch_timer"
 
     @override
-    def construct_callbacks(self, root_config):
+    def create_callbacks(self, root_config):
         yield EpochTimer()

nshtrainer/callbacks/wandb_watch.py

@@ -99,5 +99,5 @@ class WandbWatchConfig(CallbackConfigBase):
         return self.enabled
 
     @override
-    def construct_callbacks(self, root_config):
+    def create_callbacks(self, root_config):
         yield WandbWatchCallback(self)

nshtrainer/ll/__init__.py

@@ -21,7 +21,6 @@ from .log import init_python_logging as init_python_logging
 from .log import lovely as lovely
 from .log import pretty as pretty
 from .lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
-from .model import ActSaveConfig as ActSaveConfig
 from .model import Base as Base
 from .model import BaseConfig as BaseConfig
 from .model import BaseLoggerConfig as BaseLoggerConfig