nshtrainer 0.8.7__py3-none-any.whl → 0.10.0__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/_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 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/trainer/_checkpoint_resolver.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 ._checkpoint_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/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_resolver 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
+        )