nshtrainer 0.11.1__tar.gz → 0.11.3__tar.gz

{nshtrainer-0.11.1 → nshtrainer-0.11.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.11.1
+Version: 0.11.3
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.11.1 → nshtrainer-0.11.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.11.1"
+version = "0.11.3"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-0.11.1 → nshtrainer-0.11.3}/src/nshtrainer/_checkpoint/loader.py

@@ -133,6 +133,68 @@ class CheckpointLoadingConfig(C.Config):
         ckpt: Literal["best", "last"] | str | Path | None,
         trainer_mode: TrainerFn,
     ):
+        """
+        Automatically create a CheckpointLoadingConfig based on the provided checkpoint option and trainer mode.
+
+        This method provides a convenient way to generate a checkpoint loading configuration
+        tailored to different training and evaluation scenarios.
+
+        Parameters:
+        -----------
+        ckpt : Literal["best", "last"] | str | Path | None
+            Specifies the checkpoint loading preference:
+            - "best": Use the best checkpoint based on the primary metric.
+            - "last": Use the most recent checkpoint.
+            - str or Path: Path to a specific checkpoint file.
+            - None: Defaults to "last" for training, raises an error for evaluation.
+
+        trainer_mode : TrainerFn
+            The mode in which the trainer is operating. This affects how the configuration is created.
+            - TrainerFn.FITTING: Used for training scenarios.
+            - TrainerFn.VALIDATING, TrainerFn.TESTING, TrainerFn.PREDICTING: Used for evaluation scenarios.
+
+        Returns:
+        --------
+        CheckpointLoadingConfig
+            A configuration object for checkpoint loading based on the given parameters.
+
+        Behavior:
+        ---------
+        1. For training (TrainerFn.FITTING):
+        - Includes HPC pre-emption checkpoints.
+        - If ckpt is None, defaults to "last".
+        - For "best" or "last", creates a single-strategy configuration that loads the best or last checkpoint.
+        - For a specific path, creates a two-strategy configuration:
+            a) Tries to load the checkpoint as the last checkpoint.
+            b) Falls back to loading it as a user-provided path.
+
+        2. For evaluation (VALIDATING, TESTING, PREDICTING):
+        - Does not include HPC pre-emption checkpoints.
+        - Requires ckpt to be specified (raises ValueError if None).
+        - Creates a single-strategy configuration based on the ckpt value.
+
+        Raises:
+        -------
+        ValueError
+            If ckpt is None during evaluation modes.
+
+        Examples:
+        ---------
+        # Training mode, use last checkpoint
+        config = CheckpointLoadingConfig.auto("last", TrainerFn.FITTING)
+
+        # Evaluation mode, use best checkpoint
+        config = CheckpointLoadingConfig.auto("best", TrainerFn.TESTING)
+
+        # Training mode, use specific checkpoint
+        config = CheckpointLoadingConfig.auto("/path/to/checkpoint.ckpt", TrainerFn.FITTING)
+
+        Notes:
+        ------
+        - The method internally calls _auto_train or _auto_eval based on the trainer_mode.
+        - The resulting configuration always includes strategies as a sequence, even if there's only one strategy.
+        """
+        # Implementation remains the same...
         match trainer_mode:
             case TrainerFn.FITTING:
                 return cls._auto_train(ckpt)

{nshtrainer-0.11.1 → nshtrainer-0.11.3}/src/nshtrainer/_checkpoint/metadata.py

@@ -43,6 +43,16 @@ class CheckpointMetadata(C.Config):
     def from_file(cls, path: Path):
         return cls.model_validate_json(path.read_text())
 
+    @classmethod
+    def from_ckpt_path(cls, checkpoint_path: Path):
+        if not (
+            metadata_path := checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
+        ).exists():
+            raise FileNotFoundError(
+                f"Metadata file not found for checkpoint: {checkpoint_path}"
+            )
+        return cls.from_file(metadata_path)
+
 
 def _generate_checkpoint_metadata(
     config: "BaseConfig", trainer: "Trainer", checkpoint_path: Path
@@ -136,36 +146,13 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
             log.debug(f"Linked {path} to {linked_path}")
 
 
-def _checkpoint_sort_key_fn(key: Callable[[CheckpointMetadata, Path], Any]):
-    def sort_key_fn(checkpoint_path: Path):
-        if not (p := checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)).exists():
-            raise FileNotFoundError(f"Metadata file not found: {p}")
-
-        nonlocal key
-        return key(CheckpointMetadata.from_file(p), p)
-
-    return sort_key_fn
-
-
 def _sort_ckpts_by_metadata(
     checkpoint_paths: list[Path],
     key: Callable[[CheckpointMetadata, Path], Any],
-    fallback_key: Callable[[Path], Any],
+    reverse: bool = False,
 ):
-    # First, let's make sure all the metadata files exist.
-    # If not, use the fallback function to sort the checkpoints.
-    no_metadata_paths: list[Path] = []
-    for path in checkpoint_paths:
-        if (path.with_suffix(METADATA_PATH_SUFFIX)).exists():
-            continue
-
-        no_metadata_paths.append(path)
-
-    if no_metadata_paths:
-        log.warning(
-            f"Metadata file not found on {len(no_metadata_paths)} checkpoints: {no_metadata_paths}\n"
-            "Falling back to sorting by last modified time."
-        )
-        return sorted(checkpoint_paths, key=fallback_key)
-
-    return sorted(checkpoint_paths, key=_checkpoint_sort_key_fn(key))
+    return sorted(
+        [(CheckpointMetadata.from_ckpt_path(path), path) for path in checkpoint_paths],
+        key=lambda args_tuple: key(*args_tuple),
+        reverse=reverse,
+    )

{nshtrainer-0.11.1 → nshtrainer-0.11.3}/src/nshtrainer/callbacks/__init__.py

@@ -2,7 +2,20 @@ from typing import Annotated
 
 import nshconfig as C
 
+from . import checkpoint as checkpoint
 from .base import CallbackConfigBase as CallbackConfigBase
+from .checkpoint import BestCheckpoint as BestCheckpoint
+from .checkpoint import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
+from .checkpoint import LatestEpochCheckpoint as LatestEpochCheckpoint
+from .checkpoint import (
+    LatestEpochCheckpointCallbackConfig as LatestEpochCheckpointCallbackConfig,
+)
+from .checkpoint import ModelCheckpoint as ModelCheckpoint
+from .checkpoint import ModelCheckpointCallbackConfig as ModelCheckpointCallbackConfig
+from .checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
+from .checkpoint import (
+    OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+)
 from .early_stopping import EarlyStopping as EarlyStopping
 from .ema import EMA as EMA
 from .ema import EMAConfig as EMAConfig
@@ -13,21 +26,9 @@ from .gradient_skipping import GradientSkippingConfig as GradientSkippingConfig
 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
@@ -44,6 +45,7 @@ CallbackConfig = Annotated[
     | NormLoggingConfig
     | GradientSkippingConfig
     | EMAConfig
+    | BestCheckpointCallbackConfig
     | ModelCheckpointCallbackConfig
     | LatestEpochCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig

nshtrainer-0.11.3/src/nshtrainer/callbacks/checkpoint/__init__.py

@@ -0,0 +1,16 @@
+from .best_checkpoint import BestCheckpoint as BestCheckpoint
+from .best_checkpoint import (
+    BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
+)
+from .latest_epoch_checkpoint import LatestEpochCheckpoint as LatestEpochCheckpoint
+from .latest_epoch_checkpoint import (
+    LatestEpochCheckpointCallbackConfig as LatestEpochCheckpointCallbackConfig,
+)
+from .model_checkpoint import ModelCheckpoint as ModelCheckpoint
+from .model_checkpoint import (
+    ModelCheckpointCallbackConfig as ModelCheckpointCallbackConfig,
+)
+from .on_exception_checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
+from .on_exception_checkpoint import (
+    OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+)

nshtrainer-0.11.3/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -0,0 +1,155 @@
+import logging
+from pathlib import Path
+from typing import Any, Literal
+
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Checkpoint
+from typing_extensions import override
+
+from ..._checkpoint.metadata import _sort_ckpts_by_metadata
+from ..._checkpoint.saver import _link_checkpoint, _remove_checkpoint
+from ...metrics._config import MetricConfig
+from ..base import CallbackConfigBase
+
+log = logging.getLogger(__name__)
+
+
+class BestCheckpointCallbackConfig(CallbackConfigBase):
+    name: Literal["best_checkpoint"] = "best_checkpoint"
+
+    dirpath: str | Path | None = None
+    """Directory path to save the checkpoint file."""
+
+    filename: str = "epoch{epoch:02d}_step{step:04d}"
+    """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."""
+
+    metric: MetricConfig | None = None
+    """Metric to monitor, or `None` to use the default metric."""
+
+    best_symlink_filename: str | None = "best"
+    """Filename for the best symlink. If None, no symlink will be created."""
+
+    save_top_k: int | Literal["all"] = 1
+    """The number of best checkpoints to keep."""
+
+    @override
+    def create_callbacks(self, root_config):
+        dirpath = Path(
+            self.dirpath
+            or root_config.directory.resolve_subdirectory(root_config.id, "checkpoint")
+        )
+
+        # Resolve metric
+        if (metric := self.metric) is None and (
+            metric := root_config.primary_metric
+        ) is None:
+            raise ValueError(
+                "No metric provided and no primary metric found in the root config"
+            )
+
+        yield BestCheckpoint(self, metric, dirpath)
+
+    @property
+    def _save_top_k_value(self):
+        return float("inf" if self.save_top_k == "all" else self.save_top_k)
+
+
+class BestCheckpoint(Checkpoint):
+    PREFIX = "best_"
+    EXTENSION = ".ckpt"
+
+    def __init__(
+        self,
+        config: BestCheckpointCallbackConfig,
+        metric: MetricConfig,
+        dirpath: Path,
+    ):
+        super().__init__()
+        self.config = config
+        self.metric = metric
+        self.dirpath = dirpath
+
+    @override
+    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule):
+        self._save_best_checkpoint(trainer)
+
+    def _best_symlink_filename(self):
+        if (filename := self.config.best_symlink_filename) is None:
+            return None
+        return f"{filename}{self.EXTENSION}"
+
+    def _ckpt_path(self, trainer: Trainer):
+        filename = self.config.filename.format(
+            epoch=trainer.current_epoch, step=trainer.global_step
+        )
+        filename = f"{self.PREFIX}{filename}{self.EXTENSION}"
+        return self.dirpath / filename
+
+    def _remove_checkpoints(self, trainer: Trainer, ckpt_paths: list[Path]):
+        for ckpt_path in ckpt_paths:
+            _remove_checkpoint(trainer, ckpt_path, metadata=True, barrier=False)
+
+    def _get_metric_value(self, metrics: dict[str, Any]):
+        return metrics.get(
+            self.metric.validation_monitor,
+            float("-inf" if self.metric.mode == "max" else "inf"),
+        )
+
+    def _sorted_ckpts(self):
+        ckpt_paths = list(self.dirpath.glob(f"{self.PREFIX}*{self.EXTENSION}"))
+        return _sort_ckpts_by_metadata(
+            ckpt_paths,
+            key=lambda meta, _: self._get_metric_value(meta.metrics),
+            reverse=(self.metric.mode == "min"),
+        )
+
+    def _save_best_checkpoint(self, trainer: Trainer):
+        if (current := self._get_metric_value(trainer.callback_metrics)) is None:
+            log.warning(
+                f"Can't save best model, {self.metric.validation_monitor} not found in metrics"
+            )
+            return
+
+        # Get sorted checkpoints
+        sorted_ckpts = self._sorted_ckpts()
+
+        # If the current model is worse than the worst checkpoint,
+        # and we have already saved the maximum number of checkpoints,
+        # then don't save the current model.
+        if len(
+            sorted_ckpts
+        ) >= self.config._save_top_k_value and not self.metric.is_better(
+            current,
+            self._get_metric_value(sorted_ckpts[-1][0].metrics),
+        ):
+            return
+
+        # Save the current model
+        filepath = self._ckpt_path(trainer)
+        trainer.save_checkpoint(filepath, self.config.save_weights_only)
+
+        # Remove worst checkpoint if we've reached save_top_k
+        # NOTE: We add 1 to save_top_k here because we have just saved a new checkpoint
+        if len(sorted_ckpts) + 1 > self.config._save_top_k_value:
+            # Get the sorted checkpoints again because now we have added a new checkpoint.
+            # We could optimize this by adding the new checkpoint to the sorted list,
+            # and then sorting it in place, but this is simpler.
+            sorted_ckpts = self._sorted_ckpts()
+            self._remove_checkpoints(
+                trainer, [p for _, p in sorted_ckpts[self.config.save_top_k :]]
+            )
+
+        # Create symlink to best model
+        if (symlink_filename := self._best_symlink_filename()) is not None:
+            symlink_path = self.dirpath / symlink_filename
+            _link_checkpoint(
+                trainer,
+                filepath,
+                symlink_path,
+                barrier=True,
+                metadata=True,
+            )
+            log.debug(f"Created best symlink: {symlink_path}")

{nshtrainer-0.11.1/src/nshtrainer/callbacks → nshtrainer-0.11.3/src/nshtrainer/callbacks/checkpoint}/latest_epoch_checkpoint.py

@@ -6,9 +6,9 @@ from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Checkpoint
 from typing_extensions import override
 
-from .._checkpoint.metadata import _sort_ckpts_by_metadata
-from .._checkpoint.saver import _link_checkpoint, _remove_checkpoint
-from .base import CallbackConfigBase
+from ..._checkpoint.metadata import _sort_ckpts_by_metadata
+from ..._checkpoint.saver import _link_checkpoint, _remove_checkpoint
+from ..base import CallbackConfigBase
 
 log = logging.getLogger(__name__)
 
@@ -75,6 +75,10 @@ class LatestEpochCheckpoint(Checkpoint):
         if (latest_k := self.config.latest_k) == "all":
             return
 
+        # NOTE: We add 1 to the latest_k here because
+        # we're about to save a new checkpoint.
+        latest_k += 1
+
         # Get all configs, ignoring the latest symlink
         ckpt_paths = list(self.dirpath.glob(f"{self.PREFIX}*{self.EXTENSION}"))
         # Ignore the latest symlink
@@ -82,16 +86,15 @@ class LatestEpochCheckpoint(Checkpoint):
             ckpt_paths = [p for p in ckpt_paths if p.name != latest_symlink_filename]
 
         # Sort by epoch, then step, then last modified
-        ckpt_paths = _sort_ckpts_by_metadata(
+        metadata_and_ckpt_paths = _sort_ckpts_by_metadata(
             ckpt_paths,
             key=lambda meta, p: (meta.epoch, meta.global_step, p.stat().st_mtime),
-            fallback_key=lambda p: p.stat().st_mtime,
-            # ^ Called if metadata is not found on all checkpoints
+            reverse=True,
         )
 
         # Remove all but the latest k checkpoints
-        ckpts_to_remove = ckpt_paths[:-latest_k]
-        self._remove_checkpoints(trainer, ckpts_to_remove)
+        ckpts_to_remove = metadata_and_ckpt_paths[latest_k:]
+        self._remove_checkpoints(trainer, [p for _, p in ckpts_to_remove])
 
     def _save_new_checkpoint(self, trainer: Trainer):
         # Remove old checkpoints
@@ -113,4 +116,4 @@ class LatestEpochCheckpoint(Checkpoint):
                 barrier=True,
                 metadata=True,
             )
-            log.info(f"Created latest symlink: {symlink_path}")
+            log.debug(f"Created latest symlink: {symlink_path}")

{nshtrainer-0.11.1/src/nshtrainer/callbacks → nshtrainer-0.11.3/src/nshtrainer/callbacks/checkpoint}/model_checkpoint.py

@@ -10,12 +10,13 @@ from lightning.pytorch.callbacks.model_checkpoint import (
 )
 from typing_extensions import override
 
-from .._checkpoint.saver import _link_checkpoint, _remove_checkpoint
-from ..metrics import MetricConfig
-from .base import CallbackConfigBase
+from ..._checkpoint.saver import _link_checkpoint
+from ..._checkpoint.saver import _remove_checkpoint as _ckpt_saver_remove_checkpoint
+from ...metrics import MetricConfig
+from ..base import CallbackConfigBase
 
 if TYPE_CHECKING:
-    from ..model.config import BaseConfig
+    from ...model.config import BaseConfig
 
 log = logging.getLogger(__name__)
 
@@ -74,10 +75,10 @@ class ModelCheckpointCallbackConfig(CallbackConfigBase):
         If "link", creates a symbolic link to the last checkpoint.
     """
 
-    save_top_k: int = 1
+    save_top_k: int | Literal["all"] = 1
     """
     Number of best models to save.
-        If -1, all models are saved.
+        If "all" or -1, all models are saved.
         If 0, no models are saved.
     """
 
@@ -158,6 +159,11 @@ class ModelCheckpointCallbackConfig(CallbackConfigBase):
             metric=metric,
         )
 
+    def _save_top_k_model_ckpt_input(self):
+        if self.save_top_k == "all":
+            return -1
+        return self.save_top_k
+
 
 class ModelCheckpoint(_ModelCheckpoint):
     CHECKPOINT_NAME_LAST = "best"
@@ -180,7 +186,7 @@ class ModelCheckpoint(_ModelCheckpoint):
             mode=metric.mode,
             verbose=self.config.verbose,
             save_last=self.config.save_last,
-            save_top_k=self.config.save_top_k,
+            save_top_k=self.config._save_top_k_model_ckpt_input(),
             save_weights_only=self.config.save_weights_only,
             auto_insert_metric_name=False,
             every_n_train_steps=self.config.every_n_train_steps,
@@ -202,4 +208,9 @@ class ModelCheckpoint(_ModelCheckpoint):
 
     @override
     def _remove_checkpoint(self, trainer: Trainer, filepath: str):
-        return _remove_checkpoint(trainer, filepath, metadata=True, barrier=False)
+        return _ckpt_saver_remove_checkpoint(
+            trainer,
+            filepath,
+            metadata=True,
+            barrier=False,
+        )

{nshtrainer-0.11.1/src/nshtrainer/callbacks → nshtrainer-0.11.3/src/nshtrainer/callbacks/checkpoint}/on_exception_checkpoint.py

@@ -9,7 +9,7 @@ from lightning.pytorch import Trainer as LightningTrainer
 from lightning.pytorch.callbacks import OnExceptionCheckpoint as _OnExceptionCheckpoint
 from typing_extensions import override
 
-from .base import CallbackConfigBase
+from ..base import CallbackConfigBase
 
 log = logging.getLogger(__name__)
 
@@ -53,8 +53,6 @@ class OnExceptionCheckpointCallbackConfig(CallbackConfigBase):
 
     @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"
         )

{nshtrainer-0.11.1 → nshtrainer-0.11.3}/src/nshtrainer/metrics/_config.py

@@ -3,6 +3,8 @@ from typing import Literal
 
 import nshconfig as C
 
+from ..util._useful_types import SupportsRichComparisonT
+
 
 class MetricConfig(C.Config):
     name: str
@@ -35,3 +37,6 @@ class MetricConfig(C.Config):
     @property
     def best(self):
         return builtins.min if self.mode == "min" else builtins.max
+
+    def is_better(self, a: SupportsRichComparisonT, b: SupportsRichComparisonT) -> bool:
+        return self.best(a, b) == a

{nshtrainer-0.11.1 → nshtrainer-0.11.3}/src/nshtrainer/model/config.py

@@ -37,6 +37,7 @@ from typing_extensions import Self, TypedDict, TypeVar, override
 
 from .._checkpoint.loader import CheckpointLoadingConfig
 from ..callbacks import (
+    BestCheckpointCallbackConfig,
     CallbackConfig,
     LatestEpochCheckpointCallbackConfig,
     ModelCheckpointCallbackConfig,
@@ -771,6 +772,7 @@ class ReproducibilityConfig(C.Config):
 
 CheckpointCallbackConfig: TypeAlias = Annotated[
     ModelCheckpointCallbackConfig
+    | BestCheckpointCallbackConfig
     | LatestEpochCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig,
     C.Field(discriminator="name"),
@@ -782,7 +784,8 @@ class CheckpointSavingConfig(CallbackConfigBase):
     """Enable checkpoint saving."""
 
     checkpoint_callbacks: Sequence[CheckpointCallbackConfig] = [
-        ModelCheckpointCallbackConfig(),
+        # ModelCheckpointCallbackConfig(),
+        BestCheckpointCallbackConfig(),
         LatestEpochCheckpointCallbackConfig(),
         OnExceptionCheckpointCallbackConfig(),
     ]

{nshtrainer-0.11.1 → nshtrainer-0.11.3}/src/nshtrainer/trainer/checkpoint_connector.py

@@ -3,7 +3,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, cast
 
 from lightning.pytorch.trainer.connectors.checkpoint_connector import (
-    _CheckpointConnector,
+    _CheckpointConnector as _LightningCheckpointConnector,
 )
 from lightning.pytorch.trainer.states import TrainerFn
 from typing_extensions import override
@@ -15,7 +15,7 @@ if TYPE_CHECKING:
 log = logging.getLogger(__name__)
 
 
-class CheckpointConnector(_CheckpointConnector):
+class _CheckpointConnector(_LightningCheckpointConnector):
     def __resolve_auto_ckpt_path(
         self,
         ckpt_path: str | Path | None,

{nshtrainer-0.11.1 → nshtrainer-0.11.3}/src/nshtrainer/trainer/trainer.py

@@ -26,6 +26,7 @@ from ..model.config import (
     StrategyConfigProtocol,
 )
 from ._runtime_callback import RuntimeTrackerCallback, Stage
+from .checkpoint_connector import _CheckpointConnector
 from .signal_connector import _SignalConnector
 
 log = logging.getLogger(__name__)
@@ -297,6 +298,9 @@ class Trainer(LightningTrainer):
         # Replace the signal connector with our own.
         self._signal_connector = _SignalConnector(self)
 
+        # Replace the checkpoint connector with our own.
+        self._checkpoint_connector = _CheckpointConnector(self)
+
         # Print out the log dir, so that we can easily find it in the logs.
         if log_dir := self.log_dir:
             log_dir = str(Path(log_dir).resolve())

nshtrainer-0.11.3/src/nshtrainer/util/_useful_types.py

@@ -0,0 +1,307 @@
+"""Credit to useful-types from https://github.com/hauntsaninja/useful_types"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Iterable, Iterator, Sequence, Sized
+from collections.abc import Set as AbstractSet
+from os import PathLike
+from typing import Any, TypeVar, overload
+
+from typing_extensions import Buffer, Literal, Protocol, SupportsIndex, TypeAlias
+
+_KT = TypeVar("_KT")
+_KT_co = TypeVar("_KT_co", covariant=True)
+_KT_contra = TypeVar("_KT_contra", contravariant=True)
+_VT = TypeVar("_VT")
+_VT_co = TypeVar("_VT_co", covariant=True)
+_T = TypeVar("_T")
+_T_co = TypeVar("_T_co", covariant=True)
+_T_contra = TypeVar("_T_contra", contravariant=True)
+
+# For partially known annotations. Usually, fields where type annotations
+# haven't been added are left unannotated, but in some situations this
+# isn't possible or a type is already partially known. In cases like these,
+# use Incomplete instead of Any as a marker. For example, use
+# "Incomplete | None" instead of "Any | None".
+Incomplete: TypeAlias = Any
+
+
+class IdentityFunction(Protocol):
+    def __call__(self, __x: _T) -> _T: ...
+
+
+# ====================
+# Comparison protocols
+# ====================
+
+
+class SupportsDunderLT(Protocol[_T_contra]):
+    def __lt__(self, __other: _T_contra) -> bool: ...
+
+
+class SupportsDunderGT(Protocol[_T_contra]):
+    def __gt__(self, __other: _T_contra) -> bool: ...
+
+
+class SupportsDunderLE(Protocol[_T_contra]):
+    def __le__(self, __other: _T_contra) -> bool: ...
+
+
+class SupportsDunderGE(Protocol[_T_contra]):
+    def __ge__(self, __other: _T_contra) -> bool: ...
+
+
+class SupportsAllComparisons(
+    SupportsDunderLT[Any],
+    SupportsDunderGT[Any],
+    SupportsDunderLE[Any],
+    SupportsDunderGE[Any],
+    Protocol,
+): ...
+
+
+SupportsRichComparison: TypeAlias = SupportsDunderLT[Any] | SupportsDunderGT[Any]
+SupportsRichComparisonT = TypeVar(
+    "SupportsRichComparisonT", bound=SupportsRichComparison
+)
+
+# ====================
+# Dunder protocols
+# ====================
+
+
+class SupportsNext(Protocol[_T_co]):
+    def __next__(self) -> _T_co: ...
+
+
+class SupportsAnext(Protocol[_T_co]):
+    def __anext__(self) -> Awaitable[_T_co]: ...
+
+
+class SupportsAdd(Protocol[_T_contra, _T_co]):
+    def __add__(self, __x: _T_contra) -> _T_co: ...
+
+
+class SupportsRAdd(Protocol[_T_contra, _T_co]):
+    def __radd__(self, __x: _T_contra) -> _T_co: ...
+
+
+class SupportsSub(Protocol[_T_contra, _T_co]):
+    def __sub__(self, __x: _T_contra) -> _T_co: ...
+
+
+class SupportsRSub(Protocol[_T_contra, _T_co]):
+    def __rsub__(self, __x: _T_contra) -> _T_co: ...
+
+
+class SupportsDivMod(Protocol[_T_contra, _T_co]):
+    def __divmod__(self, __other: _T_contra) -> _T_co: ...
+
+
+class SupportsRDivMod(Protocol[_T_contra, _T_co]):
+    def __rdivmod__(self, __other: _T_contra) -> _T_co: ...
+
+
+# This protocol is generic over the iterator type, while Iterable is
+# generic over the type that is iterated over.
+class SupportsIter(Protocol[_T_co]):
+    def __iter__(self) -> _T_co: ...
+
+
+# This protocol is generic over the iterator type, while AsyncIterable is
+# generic over the type that is iterated over.
+class SupportsAiter(Protocol[_T_co]):
+    def __aiter__(self) -> _T_co: ...
+
+
+class SupportsLenAndGetItem(Protocol[_T_co]):
+    def __len__(self) -> int: ...
+    def __getitem__(self, __k: int) -> _T_co: ...
+
+
+class SupportsTrunc(Protocol):
+    def __trunc__(self) -> int: ...
+
+
+# ====================
+# Mapping-like protocols
+# ====================
+
+
+class SupportsItems(Protocol[_KT_co, _VT_co]):
+    def items(self) -> AbstractSet[tuple[_KT_co, _VT_co]]: ...
+
+
+class SupportsKeysAndGetItem(Protocol[_KT, _VT_co]):
+    def keys(self) -> Iterable[_KT]: ...
+    def __getitem__(self, __key: _KT) -> _VT_co: ...
+
+
+class SupportsGetItem(Protocol[_KT_contra, _VT_co]):
+    def __contains__(self, __x: Any) -> bool: ...
+    def __getitem__(self, __key: _KT_contra) -> _VT_co: ...
+
+
+class SupportsItemAccess(SupportsGetItem[_KT_contra, _VT], Protocol[_KT_contra, _VT]):
+    def __setitem__(self, __key: _KT_contra, __value: _VT) -> None: ...
+    def __delitem__(self, __key: _KT_contra) -> None: ...
+
+
+# ====================
+# File handling
+# ====================
+
+StrPath: TypeAlias = str | PathLike[str]
+BytesPath: TypeAlias = bytes | PathLike[bytes]
+StrOrBytesPath: TypeAlias = str | bytes | PathLike[str] | PathLike[bytes]
+
+OpenTextModeUpdating: TypeAlias = Literal[
+    "r+",
+    "+r",
+    "rt+",
+    "r+t",
+    "+rt",
+    "tr+",
+    "t+r",
+    "+tr",
+    "w+",
+    "+w",
+    "wt+",
+    "w+t",
+    "+wt",
+    "tw+",
+    "t+w",
+    "+tw",
+    "a+",
+    "+a",
+    "at+",
+    "a+t",
+    "+at",
+    "ta+",
+    "t+a",
+    "+ta",
+    "x+",
+    "+x",
+    "xt+",
+    "x+t",
+    "+xt",
+    "tx+",
+    "t+x",
+    "+tx",
+]
+OpenTextModeWriting: TypeAlias = Literal[
+    "w", "wt", "tw", "a", "at", "ta", "x", "xt", "tx"
+]
+OpenTextModeReading: TypeAlias = Literal[
+    "r", "rt", "tr", "U", "rU", "Ur", "rtU", "rUt", "Urt", "trU", "tUr", "Utr"
+]
+OpenTextMode: TypeAlias = (
+    OpenTextModeUpdating | OpenTextModeWriting | OpenTextModeReading
+)
+OpenBinaryModeUpdating: TypeAlias = Literal[
+    "rb+",
+    "r+b",
+    "+rb",
+    "br+",
+    "b+r",
+    "+br",
+    "wb+",
+    "w+b",
+    "+wb",
+    "bw+",
+    "b+w",
+    "+bw",
+    "ab+",
+    "a+b",
+    "+ab",
+    "ba+",
+    "b+a",
+    "+ba",
+    "xb+",
+    "x+b",
+    "+xb",
+    "bx+",
+    "b+x",
+    "+bx",
+]
+OpenBinaryModeWriting: TypeAlias = Literal["wb", "bw", "ab", "ba", "xb", "bx"]
+OpenBinaryModeReading: TypeAlias = Literal[
+    "rb", "br", "rbU", "rUb", "Urb", "brU", "bUr", "Ubr"
+]
+OpenBinaryMode: TypeAlias = (
+    OpenBinaryModeUpdating | OpenBinaryModeReading | OpenBinaryModeWriting
+)
+
+
+class HasFileno(Protocol):
+    def fileno(self) -> int: ...
+
+
+FileDescriptor: TypeAlias = int
+FileDescriptorLike: TypeAlias = int | HasFileno
+FileDescriptorOrPath: TypeAlias = int | StrOrBytesPath
+
+
+class SupportsRead(Protocol[_T_co]):
+    def read(self, __length: int = ...) -> _T_co: ...
+
+
+class SupportsReadline(Protocol[_T_co]):
+    def readline(self, __length: int = ...) -> _T_co: ...
+
+
+class SupportsNoArgReadline(Protocol[_T_co]):
+    def readline(self) -> _T_co: ...
+
+
+class SupportsWrite(Protocol[_T_contra]):
+    def write(self, __s: _T_contra) -> object: ...
+
+
+# ====================
+# Buffer protocols
+# ====================
+
+# Unfortunately PEP 688 does not allow us to distinguish read-only
+# from writable buffers. We use these aliases for readability for now.
+# Perhaps a future extension of the buffer protocol will allow us to
+# distinguish these cases in the type system.
+ReadOnlyBuffer: TypeAlias = Buffer
+# Anything that implements the read-write buffer interface.
+WriteableBuffer: TypeAlias = Buffer
+# Same as WriteableBuffer, but also includes read-only buffer types (like bytes).
+ReadableBuffer: TypeAlias = Buffer
+
+
+class SliceableBuffer(Buffer, Protocol):
+    def __getitem__(self, __slice: slice) -> Sequence[int]: ...
+
+
+class IndexableBuffer(Buffer, Protocol):
+    def __getitem__(self, __i: int) -> int: ...
+
+
+class SupportsGetItemBuffer(SliceableBuffer, IndexableBuffer, Protocol):
+    def __contains__(self, __x: Any) -> bool: ...
+    @overload
+    def __getitem__(self, __slice: slice) -> Sequence[int]: ...
+    @overload
+    def __getitem__(self, __i: int) -> int: ...
+
+
+class SizedBuffer(Sized, Buffer, Protocol): ...
+
+
+# Source from https://github.com/python/typing/issues/256#issuecomment-1442633430
+# This works because str.__contains__ does not accept object (either in typeshed or at runtime)
+class SequenceNotStr(Protocol[_T_co]):
+    @overload
+    def __getitem__(self, index: SupportsIndex, /) -> _T_co: ...
+    @overload
+    def __getitem__(self, index: slice, /) -> Sequence[_T_co]: ...
+    def __contains__(self, value: object, /) -> bool: ...
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[_T_co]: ...
+    def index(self, value: Any, start: int = 0, stop: int = ..., /) -> int: ...
+    def count(self, value: Any, /) -> int: ...
+    def __reversed__(self) -> Iterator[_T_co]: ...