nshtrainer 0.9.1__py3-none-any.whl → 0.10.1__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/_checkpoint/loader.py

@@ -0,0 +1,319 @@
+import logging
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias, overload
+
+import nshconfig as C
+from lightning.pytorch import Trainer as LightningTrainer
+from lightning.pytorch.trainer.states import TrainerFn
+from typing_extensions import assert_never
+
+from ..metrics._config import MetricConfig
+from .metadata import METADATA_PATH_SUFFIX, CheckpointMetadata
+
+if TYPE_CHECKING:
+    from ..model.config import BaseConfig
+
+log = logging.getLogger(__name__)
+
+
+class BestCheckpointStrategyConfig(C.Config):
+    name: Literal["best"] = "best"
+
+    metric: MetricConfig | None = None
+    """The metric to use for selecting the best checkpoint. If `None`, the primary metric will be used."""
+
+    additional_candidates: Iterable[Path] = []
+    """Additional checkpoint candidates to consider when selecting the last checkpoint."""
+
+
+class UserProvidedPathCheckpointStrategyConfig(C.Config):
+    name: Literal["user_provided_path"] = "user_provided_path"
+
+    path: Path
+    """The path to the checkpoint to load."""
+
+    on_error: Literal["warn", "raise"] = "warn"
+    """The behavior when the checkpoint does not belong to the current run.
+
+    - `warn`: Log a warning and skip the checkpoint.
+    - `raise`: Raise an error.
+    """
+
+
+class LastCheckpointStrategyConfig(C.Config):
+    name: Literal["last"] = "last"
+
+    criterion: Literal["global_step", "runtime"] = "global_step"
+    """The criterion to use for selecting the last checkpoint.
+
+    - `global_step`: The checkpoint with the highest global step will be selected.
+    - `runtime`: The checkpoint with the highest runtime will be selected.
+    """
+
+    additional_candidates: Iterable[Path] = []
+    """Additional checkpoint candidates to consider when selecting the last checkpoint."""
+
+
+CheckpointLoadingStrategyConfig: TypeAlias = Annotated[
+    BestCheckpointStrategyConfig
+    | LastCheckpointStrategyConfig
+    | UserProvidedPathCheckpointStrategyConfig,
+    C.Field(discriminator="name"),
+]
+
+
+class CheckpointLoadingConfig(C.Config):
+    strategies: Sequence[CheckpointLoadingStrategyConfig]
+    """The strategies to use for loading checkpoints.
+
+    The order of the strategies determines the priority of the strategies.
+    The first strategy that resolves a checkpoint will be used.
+    """
+
+    include_hpc: bool
+    """Whether to include checkpoints from HPC pre-emption."""
+
+    @classmethod
+    def _auto_train(cls, ckpt: Literal["best", "last"] | str | Path | None):
+        if ckpt is None:
+            ckpt = "last"
+        match ckpt:
+            case "best":
+                return cls(
+                    strategies=[BestCheckpointStrategyConfig()],
+                    include_hpc=True,
+                )
+            case "last":
+                return cls(
+                    strategies=[LastCheckpointStrategyConfig()],
+                    include_hpc=True,
+                )
+            case Path() | str():
+                ckpt = Path(ckpt)
+                return cls(
+                    strategies=[
+                        LastCheckpointStrategyConfig(additional_candidates=[ckpt]),
+                        UserProvidedPathCheckpointStrategyConfig(path=ckpt),
+                    ],
+                    include_hpc=True,
+                )
+            case _:
+                assert_never(ckpt)
+
+    @classmethod
+    def _auto_eval(cls, ckpt: Literal["best", "last"] | str | Path | None):
+        if ckpt is None:
+            raise ValueError("Checkpoint path must be provided for evaluation.")
+
+        match ckpt:
+            case "best":
+                return cls(
+                    strategies=[BestCheckpointStrategyConfig()],
+                    include_hpc=False,
+                )
+            case "last":
+                return cls(
+                    strategies=[LastCheckpointStrategyConfig()],
+                    include_hpc=False,
+                )
+            case Path() | str():
+                ckpt = Path(ckpt)
+                return cls(
+                    strategies=[UserProvidedPathCheckpointStrategyConfig(path=ckpt)],
+                    include_hpc=False,
+                )
+            case _:
+                assert_never(ckpt)
+
+    @classmethod
+    def auto(
+        cls,
+        ckpt: Literal["best", "last"] | str | Path | None,
+        trainer_mode: TrainerFn,
+    ):
+        match trainer_mode:
+            case TrainerFn.FITTING:
+                return cls._auto_train(ckpt)
+            case TrainerFn.VALIDATING | TrainerFn.TESTING | TrainerFn.PREDICTING:
+                return cls._auto_eval(ckpt)
+            case _:
+                assert_never(trainer_mode)
+
+
+@dataclass
+class _CkptCandidate:
+    meta: CheckpointMetadata
+    meta_path: Path
+
+    @property
+    def ckpt_path(self):
+        return self.meta_path.with_name(self.meta.checkpoint_filename)
+
+
+@overload
+def _load_ckpt_meta(
+    path: Path,
+    root_config: "BaseConfig",
+    on_error: Literal["warn"] = "warn",
+) -> _CkptCandidate | None: ...
+@overload
+def _load_ckpt_meta(
+    path: Path,
+    root_config: "BaseConfig",
+    on_error: Literal["raise"],
+) -> _CkptCandidate: ...
+def _load_ckpt_meta(
+    path: Path,
+    root_config: "BaseConfig",
+    on_error: Literal["warn", "raise"] = "warn",
+):
+    meta = CheckpointMetadata.from_file(path)
+    if root_config.id != meta.run_id:
+        error_msg = f"Skipping checkpoint {path} because it belongs to a different run"
+        match on_error:
+            case "warn":
+                log.warn(error_msg)
+            case "raise":
+                raise ValueError(error_msg)
+            case _:
+                assert_never(on_error)
+        return None
+    return _CkptCandidate(meta, path)
+
+
+def _checkpoint_candidates(
+    root_config: "BaseConfig",
+    trainer: LightningTrainer,
+    *,
+    include_hpc: bool = True,
+):
+    # Load the checkpoint directory, and throw if it doesn't exist.
+    # This indicates a non-standard setup, and we don't want to guess
+    # where the checkpoints are.
+    ckpt_dir = root_config.directory.resolve_subdirectory(root_config.id, "checkpoint")
+    if not ckpt_dir.is_dir():
+        raise FileNotFoundError(
+            f"Checkpoint directory {ckpt_dir} not found. "
+            "Please ensure that the checkpoint directory exists."
+        )
+
+    # Load all checkpoints in the directory.
+    # We can do this by looking for metadata files.
+    for path in ckpt_dir.glob(f"*{METADATA_PATH_SUFFIX}"):
+        if (meta := _load_ckpt_meta(path, root_config)) is not None:
+            yield meta
+
+    # If we have a pre-empted checkpoint, load it
+    if include_hpc and (hpc_path := trainer._checkpoint_connector._hpc_resume_path):
+        hpc_meta_path = Path(hpc_path).with_suffix(METADATA_PATH_SUFFIX)
+        if (meta := _load_ckpt_meta(hpc_meta_path, root_config)) is not None:
+            yield meta
+
+
+def _additional_candidates(
+    additional_candidates: Iterable[Path], root_config: "BaseConfig"
+):
+    for path in additional_candidates:
+        if (
+            meta := _load_ckpt_meta(path.with_suffix(METADATA_PATH_SUFFIX), root_config)
+        ) is None:
+            continue
+        yield meta
+
+
+def _resolve_checkpoint(
+    config: CheckpointLoadingConfig,
+    root_config: "BaseConfig",
+    trainer: LightningTrainer,
+):
+    # We lazily load the checkpoint candidates to avoid loading them
+    # if they are not needed.
+    _ckpt_candidates: list[_CkptCandidate] | None = None
+
+    def ckpt_candidates():
+        nonlocal _ckpt_candidates, root_config, trainer
+
+        if _ckpt_candidates is None:
+            _ckpt_candidates = list(
+                _checkpoint_candidates(
+                    root_config, trainer, include_hpc=config.include_hpc
+                )
+            )
+        return _ckpt_candidates
+
+    # Iterate over the strategies and try to resolve the checkpoint.
+    for strategy in config.strategies:
+        match strategy:
+            case UserProvidedPathCheckpointStrategyConfig():
+                meta = _load_ckpt_meta(
+                    strategy.path.with_suffix(METADATA_PATH_SUFFIX),
+                    root_config,
+                    on_error=strategy.on_error,
+                )
+                if meta is None:
+                    continue
+                return meta.ckpt_path
+            case BestCheckpointStrategyConfig():
+                candidates = [
+                    *ckpt_candidates(),
+                    *_additional_candidates(
+                        strategy.additional_candidates, root_config
+                    ),
+                ]
+                if not candidates:
+                    log.warn(
+                        "No checkpoint candidates found for `best` checkpoint strategy."
+                    )
+                    continue
+
+                if (metric := strategy.metric or root_config.primary_metric) is None:
+                    log.warn(
+                        "No metric specified for `best` checkpoint strategy, "
+                        "and no primary metric is set in the configuration. "
+                        "Skipping strategy."
+                    )
+                    continue
+
+                # Find the best checkpoint based on the metric.
+                def metric_value(ckpt: _CkptCandidate):
+                    assert metric is not None
+                    if (
+                        value := ckpt.meta.metrics.get(metric.validation_monitor)
+                    ) is None:
+                        raise ValueError(
+                            f"Metric {metric.validation_monitor} not found in checkpoint metadata. "
+                            f"Available metrics: {ckpt.meta.metrics.keys()}"
+                        )
+                    return value
+
+                best_candidate = metric.best(candidates, key=metric_value)
+                return best_candidate.ckpt_path
+            case LastCheckpointStrategyConfig():
+                candidates = [
+                    *ckpt_candidates(),
+                    *_additional_candidates(
+                        strategy.additional_candidates, root_config
+                    ),
+                ]
+                if not candidates:
+                    log.warn(
+                        "No checkpoint candidates found for `last` checkpoint strategy."
+                    )
+                    continue
+
+                # Find the last checkpoint based on the criterion.
+                def criterion_value(ckpt: _CkptCandidate):
+                    match strategy.criterion:
+                        case "global_step":
+                            return ckpt.meta.global_step
+                        case "runtime":
+                            return ckpt.meta.training_time.total_seconds()
+                        case _:
+                            assert_never(strategy.criterion)
+
+                last_candidate = max(candidates, key=criterion_value)
+                return last_candidate.ckpt_path
+            case _:
+                assert_never(strategy)

nshtrainer/_checkpoint/metadata.py

@@ -0,0 +1,102 @@
+import copy
+import datetime
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import nshconfig as C
+import numpy as np
+import torch
+
+from ..model._environment import EnvironmentConfig
+
+if TYPE_CHECKING:
+    from ..model import BaseConfig, LightningModuleBase
+    from ..trainer.trainer import Trainer
+
+log = logging.getLogger(__name__)
+
+
+METADATA_PATH_SUFFIX = ".metadata.json"
+HPARAMS_PATH_SUFFIX = ".hparams.json"
+
+
+class CheckpointMetadata(C.Config):
+    checkpoint_path: Path
+    checkpoint_filename: str
+
+    run_id: str
+    name: str
+    project: str | None
+    checkpoint_timestamp: datetime.datetime
+    start_timestamp: datetime.datetime | None
+
+    epoch: int
+    global_step: int
+    training_time: datetime.timedelta
+    metrics: dict[str, Any]
+    environment: EnvironmentConfig
+
+    @classmethod
+    def from_file(cls, path: Path):
+        return cls.model_validate_json(path.read_text())
+
+
+def _generate_checkpoint_metadata(
+    config: "BaseConfig", trainer: "Trainer", checkpoint_path: Path
+):
+    checkpoint_timestamp = datetime.datetime.now()
+    start_timestamp = trainer.start_time()
+    training_time = trainer.time_elapsed()
+
+    metrics: dict[str, Any] = {}
+    for name, metric in copy.deepcopy(trainer.callback_metrics).items():
+        match metric:
+            case torch.Tensor() | np.ndarray():
+                metrics[name] = metric.detach().cpu().item()
+            case _:
+                metrics[name] = metric
+
+    return CheckpointMetadata(
+        checkpoint_path=checkpoint_path,
+        checkpoint_filename=checkpoint_path.name,
+        run_id=config.id,
+        name=config.run_name,
+        project=config.project,
+        checkpoint_timestamp=checkpoint_timestamp,
+        start_timestamp=start_timestamp.datetime
+        if start_timestamp is not None
+        else None,
+        epoch=trainer.current_epoch,
+        global_step=trainer.global_step,
+        training_time=training_time,
+        metrics=metrics,
+        environment=config.environment,
+    )
+
+
+def _write_checkpoint_metadata(
+    trainer: "Trainer",
+    model: "LightningModuleBase",
+    checkpoint_path: Path,
+):
+    config = cast("BaseConfig", model.config)
+    metadata = _generate_checkpoint_metadata(config, trainer, checkpoint_path)
+
+    # Write the metadata to the checkpoint directory
+    try:
+        metadata_path = checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
+        metadata_path.write_text(metadata.model_dump_json(indent=4))
+    except Exception as e:
+        log.warning(f"Failed to write metadata to {checkpoint_path}: {e}")
+    else:
+        log.info(f"Checkpoint metadata written to {checkpoint_path}")
+
+    # Write the hparams to the checkpoint directory
+    try:
+        hparams_path = checkpoint_path.with_suffix(HPARAMS_PATH_SUFFIX)
+        hparams_path.write_text(config.model_dump_json(indent=4))
+    except Exception as e:
+        log.warning(f"Failed to write hparams to {checkpoint_path}: {e}")
+    else:
+        log.info(f"Checkpoint metadata written to {checkpoint_path}")

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)