nshtrainer 0.11.7__py3-none-any.whl → 0.11.9__py3-none-any.whl

nshtrainer/_checkpoint/loader.py

@@ -236,7 +236,7 @@ def _load_ckpt_meta(
         error_msg = f"Skipping checkpoint {path} because it belongs to a different run"
         match on_error:
             case "warn":
-                log.warn(error_msg)
+                log.warning(error_msg)
             case "raise":
                 raise ValueError(error_msg)
             case _:
@@ -325,13 +325,13 @@ def _resolve_checkpoint(
                     ),
                 ]
                 if not candidates:
-                    log.warn(
+                    log.warning(
                         "No checkpoint candidates found for `best` checkpoint strategy."
                     )
                     continue
 
                 if (metric := strategy.metric or root_config.primary_metric) is None:
-                    log.warn(
+                    log.warning(
                         "No metric specified for `best` checkpoint strategy, "
                         "and no primary metric is set in the configuration. "
                         "Skipping strategy."
@@ -360,7 +360,7 @@ def _resolve_checkpoint(
                     ),
                 ]
                 if not candidates:
-                    log.warn(
+                    log.warning(
                         "No checkpoint candidates found for `last` checkpoint strategy."
                     )
                     continue

nshtrainer/_checkpoint/metadata.py

@@ -4,7 +4,7 @@ import logging
 import shutil
 from collections.abc import Callable
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, ClassVar, cast
 
 import nshconfig as C
 import numpy as np
@@ -20,10 +20,11 @@ log = logging.getLogger(__name__)
 
 
 METADATA_PATH_SUFFIX = ".metadata.json"
-HPARAMS_PATH_SUFFIX = ".hparams.json"
 
 
 class CheckpointMetadata(C.Config):
+    PATH_SUFFIX: ClassVar[str] = METADATA_PATH_SUFFIX
+
     checkpoint_path: Path
     checkpoint_filename: str
 
@@ -39,6 +40,8 @@ class CheckpointMetadata(C.Config):
     metrics: dict[str, Any]
     environment: EnvironmentConfig
 
+    hparams: dict[str, Any] | None
+
     @classmethod
     def from_file(cls, path: Path):
         return cls.model_validate_json(path.read_text())
@@ -55,7 +58,10 @@ class CheckpointMetadata(C.Config):
 
 
 def _generate_checkpoint_metadata(
-    config: "BaseConfig", trainer: "Trainer", checkpoint_path: Path
+    config: "BaseConfig",
+    trainer: "Trainer",
+    checkpoint_path: Path,
+    metadata_path: Path,
 ):
     checkpoint_timestamp = datetime.datetime.now()
     start_timestamp = trainer.start_time()
@@ -70,7 +76,11 @@ def _generate_checkpoint_metadata(
                 metrics[name] = metric
 
     return CheckpointMetadata(
-        checkpoint_path=checkpoint_path,
+        # checkpoint_path=checkpoint_path,
+        # We should store the path as a relative path
+        # to the metadata file to avoid issues with
+        # moving the checkpoint directory
+        checkpoint_path=checkpoint_path.relative_to(metadata_path.parent),
         checkpoint_filename=checkpoint_path.name,
         run_id=config.id,
         name=config.run_name,
@@ -84,6 +94,7 @@ def _generate_checkpoint_metadata(
         training_time=training_time,
         metrics=metrics,
         environment=config.environment,
+        hparams=config.model_dump(mode="json"),
     )
 
 
@@ -93,36 +104,28 @@ def _write_checkpoint_metadata(
     checkpoint_path: Path,
 ):
     config = cast("BaseConfig", model.config)
-    metadata = _generate_checkpoint_metadata(config, trainer, checkpoint_path)
+    metadata_path = checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
+    metadata = _generate_checkpoint_metadata(
+        config, trainer, checkpoint_path, metadata_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.debug(f"Checkpoint metadata written to {checkpoint_path}")
 
-    # Write the hparams to the checkpoint directory
+
+def _remove_checkpoint_metadata(checkpoint_path: Path):
+    path = checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
     try:
-        hparams_path = checkpoint_path.with_suffix(HPARAMS_PATH_SUFFIX)
-        hparams_path.write_text(config.model_dump_json(indent=4))
+        path.unlink(missing_ok=True)
     except Exception as e:
-        log.warning(f"Failed to write hparams to {checkpoint_path}: {e}")
+        log.warning(f"Failed to remove {path}: {e}")
     else:
-        log.debug(f"Checkpoint metadata written to {checkpoint_path}")
-
-
-def _remove_checkpoint_metadata(checkpoint_path: Path):
-    for suffix in (METADATA_PATH_SUFFIX, HPARAMS_PATH_SUFFIX):
-        path = checkpoint_path.with_suffix(suffix)
-        try:
-            path.unlink(missing_ok=True)
-        except Exception as e:
-            log.warning(f"Failed to remove {path}: {e}")
-        else:
-            log.debug(f"Removed {path}")
+        log.debug(f"Removed {path}")
 
 
 def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Path):
@@ -130,20 +133,19 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
     _remove_checkpoint_metadata(linked_checkpoint_path)
 
     # Link the metadata files to the new checkpoint
-    for suffix in (METADATA_PATH_SUFFIX, HPARAMS_PATH_SUFFIX):
-        path = checkpoint_path.with_suffix(suffix)
-        linked_path = linked_checkpoint_path.with_suffix(suffix)
+    path = checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
+    linked_path = linked_checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
+    try:
         try:
-            try:
-                linked_path.symlink_to(path)
-            except OSError:
-                # on Windows, special permissions are required to create symbolic links as a regular user
-                # fall back to copying the file
-                shutil.copy(path, linked_path)
-        except Exception as e:
-            log.warning(f"Failed to link {path} to {linked_path}: {e}")
-        else:
-            log.debug(f"Linked {path} to {linked_path}")
+            linked_path.symlink_to(path)
+        except OSError:
+            # on Windows, special permissions are required to create symbolic links as a regular user
+            # fall back to copying the file
+            shutil.copy(path, linked_path)
+    except Exception as e:
+        log.warning(f"Failed to link {path} to {linked_path}: {e}")
+    else:
+        log.debug(f"Linked {path} to {linked_path}")
 
 
 def _sort_ckpts_by_metadata(

nshtrainer/callbacks/__init__.py

@@ -6,12 +6,8 @@ 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 LastCheckpoint as LastCheckpoint
+from .checkpoint import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
 from .checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
 from .checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
@@ -46,8 +42,7 @@ CallbackConfig = Annotated[
     | GradientSkippingConfig
     | EMAConfig
     | BestCheckpointCallbackConfig
-    | ModelCheckpointCallbackConfig
-    | LatestEpochCheckpointCallbackConfig
+    | LastCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig
     | WandbWatchConfig,
     C.Field(discriminator="name"),

nshtrainer/callbacks/checkpoint/__init__.py

@@ -2,13 +2,9 @@ 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 .last_checkpoint import LastCheckpoint as LastCheckpoint
+from .last_checkpoint import (
+    LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
 )
 from .on_exception_checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
 from .on_exception_checkpoint import (

nshtrainer/callbacks/checkpoint/_base.py

@@ -0,0 +1,175 @@
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generic, Literal
+
+import numpy as np
+import torch
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import Checkpoint
+from typing_extensions import TypeVar, override
+
+from ..._checkpoint.metadata import CheckpointMetadata
+from ..._checkpoint.saver import _link_checkpoint, _remove_checkpoint
+from ..base import CallbackConfigBase
+
+if TYPE_CHECKING:
+    from ...model.config import BaseConfig
+
+log = logging.getLogger(__name__)
+
+
+class BaseCheckpointCallbackConfig(CallbackConfigBase, ABC):
+    dirpath: str | Path | None = None
+    """Directory path to save the checkpoint file."""
+
+    filename: str | None = None
+    """Checkpoint filename. This must not include the extension.
+    If None, the default filename will be used."""
+
+    save_weights_only: bool = False
+    """Whether to save only the model's weights or the entire model object."""
+
+    save_symlink: bool = True
+    """Whether to create a symlink to the saved checkpoint."""
+
+    topk: int | Literal["all"] = 1
+    """The number of checkpoints to keep."""
+
+    @abstractmethod
+    def create_checkpoint(
+        self,
+        root_config: "BaseConfig",
+        dirpath: Path,
+    ) -> "CheckpointBase": ...
+
+    @override
+    def create_callbacks(self, root_config):
+        dirpath = Path(
+            self.dirpath
+            or root_config.directory.resolve_subdirectory(root_config.id, "checkpoint")
+        )
+
+        yield self.create_checkpoint(root_config, dirpath)
+
+
+TConfig = TypeVar("TConfig", bound=BaseCheckpointCallbackConfig, infer_variance=True)
+
+
+class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
+    def __init__(self, config: TConfig, dirpath: Path):
+        super().__init__()
+
+        self.config = config
+        self.dirpath = dirpath / self.name()
+        self.symlink_dirpath = dirpath
+
+        self._last_global_step_saved = 0
+
+    @abstractmethod
+    def default_filename(self) -> str: ...
+
+    @abstractmethod
+    def name(self) -> str: ...
+
+    def extension(self) -> str:
+        return ".ckpt"
+
+    @abstractmethod
+    def topk_sort_key(self, metadata: CheckpointMetadata) -> Any: ...
+
+    def symlink_path(self):
+        if not self.config.save_symlink:
+            return None
+
+        return self.symlink_dirpath / f"{self.name()}{self.extension()}"
+
+    def resolve_checkpoint_path(self, current_metrics: dict[str, Any]) -> Path:
+        if (filename := self.config.filename) is None:
+            filename = self.default_filename()
+        filename = filename.format(**current_metrics)
+        return self.dirpath / f"{filename}{self.extension()}"
+
+    def remove_old_checkpoints(self, trainer: Trainer):
+        if (topk := self.config.topk) == "all":
+            return
+
+        # Get all the checkpoint metadata
+        metas = [
+            CheckpointMetadata.from_file(p)
+            for p in self.dirpath.glob(f"*{CheckpointMetadata.PATH_SUFFIX}")
+        ]
+
+        # Sort by the topk sort key
+        metas = sorted(metas, key=self.topk_sort_key)
+
+        # Now, the metas are sorted from the best to the worst,
+        # so we can remove the worst checkpoints
+        for meta in metas[topk:]:
+            if not (old_ckpt_path := self.dirpath / meta.checkpoint_filename).exists():
+                log.warning(
+                    f"Checkpoint file not found: {old_ckpt_path}\n"
+                    "Skipping removal of the checkpoint metadata."
+                )
+                continue
+
+            _remove_checkpoint(trainer, old_ckpt_path, metadata=True)
+            log.debug(f"Removed old checkpoint: {old_ckpt_path}")
+
+    def current_metrics(self, trainer: Trainer) -> dict[str, Any]:
+        current_metrics: dict[str, Any] = {
+            "epoch": trainer.current_epoch,
+            "step": trainer.global_step,
+        }
+
+        for name, value in trainer.callback_metrics.items():
+            match value:
+                case torch.Tensor() if value.numel() == 1:
+                    value = value.detach().cpu().item()
+                case np.ndarray() if value.size == 1:
+                    value = value.item()
+                case _:
+                    pass
+
+            current_metrics[name] = value
+
+        return current_metrics
+
+    def save_checkpoints(self, trainer: Trainer):
+        if self._should_skip_saving_checkpoint(trainer):
+            return
+
+        # Save the new checkpoint
+        filepath = self.resolve_checkpoint_path(self.current_metrics(trainer))
+        trainer.save_checkpoint(filepath, self.config.save_weights_only)
+
+        if trainer.is_global_zero:
+            # Remove old checkpoints
+            self.remove_old_checkpoints(trainer)
+
+            # Create the latest symlink
+            if (symlink_filename := self.symlink_path()) is not None:
+                symlink_path = self.dirpath / symlink_filename
+                _link_checkpoint(filepath, symlink_path, metadata=True)
+                log.debug(f"Created latest symlink: {symlink_path}")
+
+        # Barrier to ensure all processes have saved the checkpoint,
+        # deleted the old checkpoints, and created the symlink before continuing
+        trainer.strategy.barrier()
+
+        # Set the last global step saved
+        self._last_global_step_saved = trainer.global_step
+
+    def _should_skip_saving_checkpoint(self, trainer: Trainer) -> bool:
+        from lightning.pytorch.trainer.states import TrainerFn
+
+        return (
+            bool(
+                getattr(trainer, "fast_dev_run", False)
+            )  # disable checkpointing with fast_dev_run
+            or trainer.state.fn
+            != TrainerFn.FITTING  # don't save anything during non-fit
+            or trainer.sanity_checking  # don't save anything during sanity check
+            or self._last_global_step_saved
+            == trainer.global_step  # already saved at the last step
+        )

nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -1,47 +1,27 @@
 import logging
 from pathlib import Path
-from typing import Any, Literal
+from typing import Literal
 
 from lightning.pytorch import LightningModule, Trainer
-from lightning.pytorch.callbacks import Checkpoint
-from typing_extensions import override
+from typing_extensions import final, override
+
+from nshtrainer._checkpoint.metadata import CheckpointMetadata
 
-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
+from ._base import BaseCheckpointCallbackConfig, CheckpointBase
 
 log = logging.getLogger(__name__)
 
 
-class BestCheckpointCallbackConfig(CallbackConfigBase):
+@final
+class BestCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
     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")
-        )
-
+    def create_checkpoint(self, root_config, dirpath):
         # Resolve metric
         if (metric := self.metric) is None and (
             metric := root_config.primary_metric
@@ -50,143 +30,41 @@ class BestCheckpointCallbackConfig(CallbackConfigBase):
                 "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)
+        return BestCheckpoint(self, dirpath, metric)
 
 
-class BestCheckpoint(Checkpoint):
-    PREFIX = "best_"
-    EXTENSION = ".ckpt"
+@final
+class BestCheckpoint(CheckpointBase[BestCheckpointCallbackConfig]):
+    @property
+    def _metric_name_normalized(self):
+        return self.metric.name.replace("/", "_").replace(" ", "_").replace(".", "_")
 
+    @override
     def __init__(
         self,
         config: BestCheckpointCallbackConfig,
-        metric: MetricConfig,
         dirpath: Path,
+        metric: MetricConfig,
     ):
-        super().__init__()
-        self.config = config
+        super().__init__(config, dirpath)
         self.metric = metric
-        self.dirpath = dirpath
-
-        self._last_global_step_saved = 0  # no need to save when no steps were taken
 
     @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 name(self):
+        return f"best_{self._metric_name_normalized}"
 
-    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
+    @override
+    def default_filename(self):
+        return f"epoch{{epoch:03d}}-{self._metric_name_normalized}{{{self.metric.validation_monitor}}}"
 
-    def _get_metric_value(self, metrics: dict[str, Any]):
-        return metrics.get(
+    @override
+    def topk_sort_key(self, metadata: CheckpointMetadata):
+        return metadata.metrics.get(
             self.metric.validation_monitor,
             float("-inf" if self.metric.mode == "max" else "inf"),
         )
 
-    def _sorted_ckpts(self):
-        """
-        Get sorted checkpoints by the metric value.
-
-        Sort order: best -> worst
-        """
-        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 == "max"),
-        )
-
-    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(best_ckpt_path, symlink_path, metadata=True)
-        log.debug(f"Created best symlink: {symlink_path}")
-
-    def _save_best_checkpoint(self, trainer: Trainer):
-        # Skip saving the checkpoint if we're not in the fitting state
-        if self._should_skip_saving_checkpoint(trainer):
-            return
-
-        # Get the current metric value
-        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)
-        log.debug(f"Saved best checkpoint: {filepath}")
-
-        if trainer.is_global_zero:
-            # 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()
-
-            # Remove worst checkpoint if we've reached save_top_k
-            if (topk := self.config.save_top_k) != "all" and len(sorted_ckpts) > topk:
-                # NOTE: Sort order is best -> worst. Let's get the worst checkpoints.
-                for _, ckpt_path in sorted_ckpts[topk:]:
-                    _remove_checkpoint(trainer, ckpt_path, metadata=True)
-
-            # Create symlink to best model
-            if sorted_ckpts:
-                _, best_ckpt_path = sorted_ckpts[0]
-                self._create_symlink(trainer, best_ckpt_path)
-
-        # Update the last global step saved
-        self._last_global_step_saved = trainer.global_step
-
-        # Barrier to ensure all processes have saved the checkpoint before continuing
-        trainer.strategy.barrier()
-
-    def _should_skip_saving_checkpoint(self, trainer: Trainer) -> bool:
-        from lightning.pytorch.trainer.states import TrainerFn
-
-        return (
-            bool(
-                getattr(trainer, "fast_dev_run", False)
-            )  # disable checkpointing with fast_dev_run
-            or trainer.state.fn
-            != TrainerFn.FITTING  # don't save anything during non-fit
-            or trainer.sanity_checking  # don't save anything during sanity check
-            or self._last_global_step_saved
-            == trainer.global_step  # already saved at the last step
-        )
+    # Events
+    @override
+    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule):
+        self.save_checkpoints(trainer)

nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -0,0 +1,39 @@
+import logging
+from typing import Literal
+
+from lightning.pytorch import LightningModule, Trainer
+from typing_extensions import final, override
+
+from nshtrainer._checkpoint.metadata import CheckpointMetadata
+
+from ._base import BaseCheckpointCallbackConfig, CheckpointBase
+
+log = logging.getLogger(__name__)
+
+
+@final
+class LastCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
+    name: Literal["last_checkpoint"] = "last_checkpoint"
+
+    @override
+    def create_checkpoint(self, root_config, dirpath):
+        return LastCheckpoint(self, dirpath)
+
+
+@final
+class LastCheckpoint(CheckpointBase[LastCheckpointCallbackConfig]):
+    @override
+    def name(self):
+        return "last"
+
+    @override
+    def default_filename(self):
+        return "epoch{epoch:03d}-step{step:07d}"
+
+    @override
+    def topk_sort_key(self, metadata: CheckpointMetadata):
+        return metadata.checkpoint_timestamp
+
+    @override
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        self.save_checkpoints(trainer)

nshtrainer/model/__init__.py

@@ -5,17 +5,15 @@ from .base import LightningModuleBase as LightningModuleBase
 from .config import BaseConfig as BaseConfig
 from .config import BaseLoggerConfig as BaseLoggerConfig
 from .config import BaseProfilerConfig as BaseProfilerConfig
+from .config import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
 from .config import CheckpointLoadingConfig as CheckpointLoadingConfig
 from .config import CheckpointSavingConfig as CheckpointSavingConfig
 from .config import DirectoryConfig as DirectoryConfig
 from .config import EarlyStoppingConfig as EarlyStoppingConfig
 from .config import GradientClippingConfig as GradientClippingConfig
-from .config import (
-    LatestEpochCheckpointCallbackConfig as LatestEpochCheckpointCallbackConfig,
-)
+from .config import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
 from .config import LoggingConfig as LoggingConfig
 from .config import MetricConfig as MetricConfig
-from .config import ModelCheckpointCallbackConfig as ModelCheckpointCallbackConfig
 from .config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )

nshtrainer/model/config.py

@@ -39,8 +39,7 @@ from .._checkpoint.loader import CheckpointLoadingConfig
 from ..callbacks import (
     BestCheckpointCallbackConfig,
     CallbackConfig,
-    LatestEpochCheckpointCallbackConfig,
-    ModelCheckpointCallbackConfig,
+    LastCheckpointCallbackConfig,
     OnExceptionCheckpointCallbackConfig,
     WandbWatchConfig,
 )
@@ -771,9 +770,8 @@ class ReproducibilityConfig(C.Config):
 
 
 CheckpointCallbackConfig: TypeAlias = Annotated[
-    ModelCheckpointCallbackConfig
-    | BestCheckpointCallbackConfig
-    | LatestEpochCheckpointCallbackConfig
+    BestCheckpointCallbackConfig
+    | LastCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig,
     C.Field(discriminator="name"),
 ]
@@ -784,9 +782,8 @@ class CheckpointSavingConfig(CallbackConfigBase):
     """Enable checkpoint saving."""
 
     checkpoint_callbacks: Sequence[CheckpointCallbackConfig] = [
-        # ModelCheckpointCallbackConfig(),
         BestCheckpointCallbackConfig(),
-        LatestEpochCheckpointCallbackConfig(),
+        LastCheckpointCallbackConfig(),
         OnExceptionCheckpointCallbackConfig(),
     ]
     """Checkpoint callback configurations."""
@@ -804,36 +801,6 @@ class CheckpointSavingConfig(CallbackConfigBase):
 
         return True
 
-    @property
-    def model_checkpoint(self) -> ModelCheckpointCallbackConfig | None:
-        return next(
-            (
-                callback
-                for callback in self.checkpoint_callbacks
-                if isinstance(callback, ModelCheckpointCallbackConfig)
-            ),
-        )
-
-    @property
-    def latest_epoch_checkpoint(self) -> LatestEpochCheckpointCallbackConfig | None:
-        return next(
-            (
-                callback
-                for callback in self.checkpoint_callbacks
-                if isinstance(callback, LatestEpochCheckpointCallbackConfig)
-            ),
-        )
-
-    @property
-    def on_exception_checkpoint(self) -> OnExceptionCheckpointCallbackConfig | None:
-        return next(
-            (
-                callback
-                for callback in self.checkpoint_callbacks
-                if isinstance(callback, OnExceptionCheckpointCallbackConfig)
-            ),
-        )
-
     @override
     def create_callbacks(self, root_config: "BaseConfig"):
         if not self.should_save_checkpoints(root_config):

{nshtrainer-0.11.7.dist-info → nshtrainer-0.11.9.dist-info}/METADATA

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

{nshtrainer-0.11.7.dist-info → nshtrainer-0.11.9.dist-info}/RECORD

@@ -1,19 +1,19 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
-nshtrainer/_checkpoint/loader.py,sha256=_3jBf-k-fJCFfmU8wjDwbnE9rb4WoKYEyQiKGsBOCi4,13777
-nshtrainer/_checkpoint/metadata.py,sha256=M9eAZ2xMs36Z1G1xULu9MHZhsHxN8_9mNt3Iv7wuq-I,5069
+nshtrainer/_checkpoint/loader.py,sha256=DSaNR8194kWon4O1svslNsCcN_8vlyLbF0LNCPfUpzI,13789
+nshtrainer/_checkpoint/metadata.py,sha256=n2PMGdA3Jn3BuoyDXVCF9dUnNjuuru5CL9jqMD-X4Vk,4918
 nshtrainer/_checkpoint/saver.py,sha256=TuSAP39DOOVvSnSukQ9RitMV60JnDg6L27fMRc2uVJc,1358
 nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ4b-PoWM,120
 nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
 nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
 nshtrainer/_experimental/flops/module_tracker.py,sha256=bUL-IRTd0aF_DwmXkZjHZAA31p4ZEhyqhc26XWKQUUY,4922
-nshtrainer/callbacks/__init__.py,sha256=5lFHe7bNdKxqvw8VRHG18W3R5l34n02Z1wVuQ38_PTg,2488
+nshtrainer/callbacks/__init__.py,sha256=k-DbpIlH2t5-oR3gHGHr8KiyCd_Twers4PcIUM1noqQ,2262
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,3555
-nshtrainer/callbacks/checkpoint/__init__.py,sha256=zrEVCGFikfkt0iOMceOFzXsZG2-6QrqY79RKBCS7bu4,738
-nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=BUO0sWqlwfyxD1UeII5DZ-01SGLiawJAEsL8HjGX4XA,7018
-nshtrainer/callbacks/checkpoint/latest_epoch_checkpoint.py,sha256=CQ0IqhuPI7zAxpQLy48kK8qqfVfwXEJoHGRqI4h8xNk,4819
-nshtrainer/callbacks/checkpoint/model_checkpoint.py,sha256=JS1z2YuEiQxk61HgZU1jySzF_pzdfXYO54_qHo-q3CQ,6776
+nshtrainer/callbacks/checkpoint/__init__.py,sha256=g-3zIthupERKqWZQw-A_busQPaPRkto6iHBV-M7nK1Y,527
+nshtrainer/callbacks/checkpoint/_base.py,sha256=9HQSa-toOyjtuDldQ71gaDVBRdryAaB_nRv5Y554tIk,5938
+nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=O6d4SqIYsxpnRj_IYX8A9VLgOBwxTdz-j2FV_nn3BT8,2067
+nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CM8f37dwaYHkjQFfJNTZTzSoF45zEjFRm-Fg1CzYmP4,1037
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=s8tOHrnb_uVqLVeV2K38ZszXrXPTEGdDVfXuXgo_KDQ,3277
 nshtrainer/callbacks/early_stopping.py,sha256=LGn3rdbvkFfUo9kwMzK4eMGlPAqD9uFdowDx6VdfozQ,3761
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
@@ -52,9 +52,9 @@ nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=mn6cyizyI_stkXtg6zxIEGF9b
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-WKArwxxf7ZQizhmbCw,2782
 nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
 nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ,1288
-nshtrainer/model/__init__.py,sha256=NpvyQHmGaHB8xdraHmm8l7kDHLmvJSgBNQKkfYqtgyI,1454
+nshtrainer/model/__init__.py,sha256=RlGW5a46DZcqK6cYICYxDaKpZIEj-8zLxoMrl432tno,1429
 nshtrainer/model/base.py,sha256=AXRfEsFAT0Ln7zjYVPU5NgtHS_c8FZM-M4pyLamO7OA,17516
-nshtrainer/model/config.py,sha256=48Vx4RiPUGjqEhoHEX4ukOGy6KlI7RmhnShJQaRQ3io,54885
+nshtrainer/model/config.py,sha256=F-doUiqPVw4yepT6FRqhflEiiMEWr_PuFC6lKzFWktA,53809
 nshtrainer/model/modules/callback.py,sha256=K0-cyEtBcQhI7Q2e-AGTE8T-GghUPY9DYmneU6ULV6g,6401
 nshtrainer/model/modules/debug.py,sha256=Yy7XEdPou9BkCsD5hJchwJGmCVGrfUru5g9VjPM4uAw,1120
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
@@ -82,6 +82,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.11.7.dist-info/METADATA,sha256=htPbfKNDbqr1taf0bEvSbl-hQOaPfjKkJhvrWoMz2r0,860
-nshtrainer-0.11.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.11.7.dist-info/RECORD,,
+nshtrainer-0.11.9.dist-info/METADATA,sha256=rFmi4wYXJz8srZhqFl_5ROxfmgFru7jhYUpUp7ZZjMg,860
+nshtrainer-0.11.9.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.11.9.dist-info/RECORD,,

nshtrainer/callbacks/checkpoint/latest_epoch_checkpoint.py

@@ -1,131 +0,0 @@
-import logging
-from pathlib import Path
-from typing import 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 ..base import CallbackConfigBase
-
-log = logging.getLogger(__name__)
-
-
-class LatestEpochCheckpointCallbackConfig(CallbackConfigBase):
-    name: Literal["latest_epoch_checkpoint"] = "latest_epoch_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."""
-
-    latest_symlink_filename: str | None = "latest"
-    """Filename for the latest symlink. If None, no symlink will be created."""
-
-    latest_k: int | Literal["all"] = 1
-    """Number of latest checkpoints to keep. If "all", all checkpoints are kept."""
-
-    @override
-    def create_callbacks(self, root_config):
-        dirpath = self.dirpath or root_config.directory.resolve_subdirectory(
-            root_config.id, "checkpoint"
-        )
-        dirpath = Path(dirpath)
-
-        yield LatestEpochCheckpoint(self, dirpath)
-
-
-class LatestEpochCheckpoint(Checkpoint):
-    PREFIX = "latest_"
-    EXTENSION = ".ckpt"
-
-    def __init__(self, config: LatestEpochCheckpointCallbackConfig, dirpath: Path):
-        super().__init__()
-
-        self.config = config
-        self.dirpath = dirpath
-
-        self._last_global_step_saved = 0
-
-    @override
-    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
-        self._save_new_checkpoint(trainer)
-
-    def _latest_symlink_filename(self):
-        if (filename := self.config.latest_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_old_checkpoints(self, trainer: Trainer):
-        if (latest_k := self.config.latest_k) == "all":
-            return
-
-        # Get all configs, ignoring the latest symlink
-        ckpts = list(self.dirpath.glob(f"{self.PREFIX}*{self.EXTENSION}"))
-        # Ignore the latest symlink
-        if (latest_symlink_filename := self._latest_symlink_filename()) is not None:
-            ckpts = [p for p in ckpts if p.name != latest_symlink_filename]
-
-        # Sort by epoch, then step, then last modified
-        ckpts = _sort_ckpts_by_metadata(
-            ckpts,
-            key=lambda meta, p: (meta.epoch, meta.global_step, p.stat().st_mtime),
-            reverse=True,
-        )
-
-        # Remove all but the latest k checkpoints
-        # NOTE: We add 1 to the latest_k here because
-        # we're about to save a new checkpoint.
-        for _, ckpt_path in ckpts[latest_k:]:
-            _remove_checkpoint(trainer, ckpt_path, metadata=True)
-
-    def _save_new_checkpoint(self, trainer: Trainer):
-        if self._should_skip_saving_checkpoint(trainer):
-            return
-
-        # Save the new checkpoint
-        filepath = self._ckpt_path(trainer)
-        trainer.save_checkpoint(filepath, self.config.save_weights_only)
-
-        if trainer.is_global_zero:
-            # Remove old checkpoints
-            self._remove_old_checkpoints(trainer)
-
-            # Create the latest symlink
-            if (symlink_filename := self._latest_symlink_filename()) is not None:
-                symlink_path = self.dirpath / symlink_filename
-                _link_checkpoint(filepath, symlink_path, metadata=True)
-                log.debug(f"Created latest symlink: {symlink_path}")
-
-        # Set the last global step saved
-        self._last_global_step_saved = trainer.global_step
-
-        # Barrier to ensure all processes have saved the checkpoint before continuing
-        trainer.strategy.barrier()
-
-    def _should_skip_saving_checkpoint(self, trainer: Trainer) -> bool:
-        from lightning.pytorch.trainer.states import TrainerFn
-
-        return (
-            bool(
-                getattr(trainer, "fast_dev_run", False)
-            )  # disable checkpointing with fast_dev_run
-            or trainer.state.fn
-            != TrainerFn.FITTING  # don't save anything during non-fit
-            or trainer.sanity_checking  # don't save anything during sanity check
-            or self._last_global_step_saved
-            == trainer.global_step  # already saved at the last step
-        )

nshtrainer/callbacks/checkpoint/model_checkpoint.py

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