nshtrainer 0.11.7__tar.gz → 0.11.9__tar.gz

{nshtrainer-0.11.7 → nshtrainer-0.11.9}/PKG-INFO

@@ -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 → nshtrainer-0.11.9}/pyproject.toml

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

{nshtrainer-0.11.7 → nshtrainer-0.11.9}/src/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-0.11.7 → nshtrainer-0.11.9}/src/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-0.11.7 → nshtrainer-0.11.9}/src/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-0.11.9/src/nshtrainer/callbacks/checkpoint/__init__.py

@@ -0,0 +1,12 @@
+from .best_checkpoint import BestCheckpoint as BestCheckpoint
+from .best_checkpoint import (
+    BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
+)
+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 (
+    OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+)

nshtrainer-0.11.9/src/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-0.11.9/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -0,0 +1,70 @@
+import logging
+from pathlib import Path
+from typing import Literal
+
+from lightning.pytorch import LightningModule, Trainer
+from typing_extensions import final, override
+
+from nshtrainer._checkpoint.metadata import CheckpointMetadata
+
+from ...metrics._config import MetricConfig
+from ._base import BaseCheckpointCallbackConfig, CheckpointBase
+
+log = logging.getLogger(__name__)
+
+
+@final
+class BestCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
+    name: Literal["best_checkpoint"] = "best_checkpoint"
+
+    metric: MetricConfig | None = None
+    """Metric to monitor, or `None` to use the default metric."""
+
+    @override
+    def create_checkpoint(self, root_config, dirpath):
+        # 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"
+            )
+
+        return BestCheckpoint(self, dirpath, metric)
+
+
+@final
+class BestCheckpoint(CheckpointBase[BestCheckpointCallbackConfig]):
+    @property
+    def _metric_name_normalized(self):
+        return self.metric.name.replace("/", "_").replace(" ", "_").replace(".", "_")
+
+    @override
+    def __init__(
+        self,
+        config: BestCheckpointCallbackConfig,
+        dirpath: Path,
+        metric: MetricConfig,
+    ):
+        super().__init__(config, dirpath)
+        self.metric = metric
+
+    @override
+    def name(self):
+        return f"best_{self._metric_name_normalized}"
+
+    @override
+    def default_filename(self):
+        return f"epoch{{epoch:03d}}-{self._metric_name_normalized}{{{self.metric.validation_monitor}}}"
+
+    @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"),
+        )
+
+    # Events
+    @override
+    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule):
+        self.save_checkpoints(trainer)

nshtrainer-0.11.9/src/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-0.11.7 → nshtrainer-0.11.9}/src/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-0.11.7 → nshtrainer-0.11.9}/src/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/src/nshtrainer/callbacks/checkpoint/__init__.py

@@ -1,16 +0,0 @@
-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,
-)