nshtrainer 0.11.2__tar.gz → 0.11.4__tar.gz

{nshtrainer-0.11.2 → nshtrainer-0.11.4}/PKG-INFO

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

{nshtrainer-0.11.2 → nshtrainer-0.11.4}/pyproject.toml

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

{nshtrainer-0.11.2 → nshtrainer-0.11.4}/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.2 → nshtrainer-0.11.4}/src/nshtrainer/callbacks/__init__.py

@@ -4,6 +4,8 @@ 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,
@@ -43,6 +45,7 @@ CallbackConfig = Annotated[
     | NormLoggingConfig
     | GradientSkippingConfig
     | EMAConfig
+    | BestCheckpointCallbackConfig
     | ModelCheckpointCallbackConfig
     | LatestEpochCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig

{nshtrainer-0.11.2 → nshtrainer-0.11.4}/src/nshtrainer/callbacks/checkpoint/__init__.py

@@ -1,3 +1,7 @@
+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,

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

@@ -0,0 +1,171 @@
+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 _create_symlink(self, trainer: Trainer, best_ckpt_path: Path):
+        # Resolve the symlink filename
+        if (symlink_filename := self._best_symlink_filename()) is None:
+            return
+
+        # If the symlink already exists and points to the best checkpoint,
+        # then we don't need to create a new symlink.
+        symlink_path = self.dirpath / symlink_filename
+        if symlink_path.exists() and symlink_path.resolve() == best_ckpt_path:
+            return
+
+        _link_checkpoint(
+            trainer,
+            best_ckpt_path,
+            symlink_path,
+            metadata=True,
+            barrier=False,
+        )
+        log.debug(f"Created best symlink: {symlink_path}")
+
+    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
+        _, best_ckpt_path = sorted_ckpts[0]
+        self._create_symlink(trainer, best_ckpt_path)
+        log.debug(f"Saved best checkpoint: {filepath}")
+
+        # Barrier to ensure all processes have saved the checkpoint before continuing
+        trainer.strategy.barrier()

{nshtrainer-0.11.2 → nshtrainer-0.11.4}/src/nshtrainer/callbacks/checkpoint/latest_epoch_checkpoint.py

@@ -86,15 +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
-        self._remove_checkpoints(trainer, ckpt_paths[:-latest_k])
+        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

{nshtrainer-0.11.2 → nshtrainer-0.11.4}/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.2 → nshtrainer-0.11.4}/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.4/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]: ...