nshtrainer 0.10.9__tar.gz → 0.10.11__tar.gz

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/PKG-INFO

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

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/pyproject.toml

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

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/_checkpoint/metadata.py

@@ -1,6 +1,8 @@
 import copy
 import datetime
 import logging
+import shutil
+from collections.abc import Callable
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
@@ -100,3 +102,74 @@ def _write_checkpoint_metadata(
         log.warning(f"Failed to write hparams to {checkpoint_path}: {e}")
     else:
         log.info(f"Checkpoint metadata written to {checkpoint_path}")
+
+
+def _remove_checkpoint_metadata(checkpoint_path: Path):
+    for path in (
+        checkpoint_path.with_suffix(METADATA_PATH_SUFFIX),
+        checkpoint_path.with_suffix(HPARAMS_PATH_SUFFIX),
+    ):
+        try:
+            path.unlink(missing_ok=True)
+        except Exception as e:
+            log.warning(f"Failed to remove {path}: {e}")
+        else:
+            log.info(f"Removed {path}")
+
+
+def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Path):
+    # First, remove any existing metadata files
+    _remove_checkpoint_metadata(linked_checkpoint_path)
+
+    # Link the metadata files to the new checkpoint
+    for path in (
+        checkpoint_path.with_suffix(METADATA_PATH_SUFFIX),
+        checkpoint_path.with_suffix(HPARAMS_PATH_SUFFIX),
+    ):
+        linked_path = linked_checkpoint_path.with_suffix(path.suffix)
+        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.info(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],
+):
+    # 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))

nshtrainer-0.10.11/src/nshtrainer/_checkpoint/saver.py

@@ -0,0 +1,52 @@
+import os
+import shutil
+from pathlib import Path
+
+from lightning.pytorch import Trainer
+
+from .metadata import _link_checkpoint_metadata, _remove_checkpoint_metadata
+
+
+def _link_checkpoint(
+    trainer: Trainer,
+    filepath: str | Path | os.PathLike,
+    linkpath: str | Path | os.PathLike,
+    *,
+    barrier: bool,
+    metadata: bool,
+):
+    if not isinstance(filepath, Path):
+        filepath = Path(filepath)
+    if not isinstance(linkpath, Path):
+        linkpath = Path(linkpath)
+
+    if trainer.is_global_zero:
+        if linkpath.exists():
+            if linkpath.is_symlink() or linkpath.is_file():
+                linkpath.unlink()
+            elif linkpath.is_dir():
+                shutil.rmtree(linkpath)
+            _remove_checkpoint_metadata(linkpath)
+
+        try:
+            target_path = filepath.relative_to(linkpath.parent)
+            linkpath.symlink_to(target_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(filepath, linkpath)
+
+        _link_checkpoint_metadata(filepath, linkpath)
+    if barrier:
+        trainer.strategy.barrier()
+
+
+def _remove_checkpoint(
+    trainer: Trainer,
+    filepath: str | Path | os.PathLike,
+    remove_metadata: bool = True,
+):
+    if not isinstance(filepath, Path):
+        filepath = Path(filepath)
+    trainer.strategy.remove_checkpoint(filepath)
+    _remove_checkpoint_metadata(filepath)

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/callbacks/early_stopping.py

@@ -1,5 +1,5 @@
+import logging
 import math
-from logging import getLogger
 
 from lightning.fabric.utilities.rank_zero import _get_rank
 from lightning.pytorch import Trainer
@@ -7,7 +7,7 @@ from lightning.pytorch.callbacks import EarlyStopping as _EarlyStopping
 from lightning.pytorch.utilities.rank_zero import rank_prefixed_message
 from typing_extensions import override
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 class EarlyStopping(_EarlyStopping):

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/callbacks/finite_checks.py

@@ -1,4 +1,4 @@
-from logging import getLogger
+import logging
 from typing import Literal
 
 import torch
@@ -7,7 +7,7 @@ from typing_extensions import override
 
 from .base import CallbackConfigBase
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def finite_checks(

nshtrainer-0.10.11/src/nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -0,0 +1,114 @@
+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
+
+    @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_checkpoints(self, trainer: Trainer, ckpt_paths: list[Path]):
+        for ckpt_path in ckpt_paths:
+            _remove_checkpoint(trainer, ckpt_path, remove_metadata=True)
+
+    def _remove_old_checkpoints(self, trainer: Trainer):
+        if (latest_k := self.config.latest_k) == "all":
+            return
+
+        # Get all configs, ignoring the latest symlink
+        ckpt_paths = list(self.dirpath.glob(f"{self.PREFIX}*{self.EXTENSION}"))
+        # Ignore the latest symlink
+        if (latest_symlink_filename := self._latest_symlink_filename()) is not None:
+            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(
+            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
+        )
+
+        # Remove all but the latest k checkpoints
+        ckpts_to_remove = ckpt_paths[:-latest_k]
+        self._remove_checkpoints(trainer, ckpts_to_remove)
+
+    def _save_new_checkpoint(self, trainer: Trainer):
+        # Remove old checkpoints
+        self._remove_old_checkpoints(trainer)
+
+        # Save the new checkpoint
+        filepath = self._ckpt_path(trainer)
+        trainer.save_checkpoint(filepath, self.config.save_weights_only)
+
+        # Create the latest symlink
+        if (symlink_filename := self._latest_symlink_filename()) is not None:
+            symlink_path = self.dirpath / symlink_filename
+            _link_checkpoint(
+                trainer,
+                filepath,
+                symlink_path,
+                barrier=True,
+                metadata=True,
+            )
+            log.info(f"Created latest symlink: {symlink_path}")

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/callbacks/model_checkpoint.py

@@ -1,21 +1,23 @@
+import logging
 import re
 from datetime import timedelta
-from logging import getLogger
 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, _remove_checkpoint
 from ..metrics import MetricConfig
 from .base import CallbackConfigBase
 
 if TYPE_CHECKING:
     from ..model.config import BaseConfig
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def _convert_string(input_string: str):
@@ -158,6 +160,8 @@ class ModelCheckpointCallbackConfig(CallbackConfigBase):
 
 
 class ModelCheckpoint(_ModelCheckpoint):
+    CHECKPOINT_NAME_LAST = "best"
+
     @override
     def __init__(
         self,
@@ -185,3 +189,17 @@ class ModelCheckpoint(_ModelCheckpoint):
             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]
+        return _link_checkpoint(
+            trainer,
+            filepath,
+            linkpath,
+            barrier=True,
+            metadata=True,
+        )
+
+    @override
+    def _remove_checkpoint(self, trainer: Trainer, filepath: str):
+        return _remove_checkpoint(trainer, filepath, remove_metadata=True)

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/callbacks/norm_logging.py

@@ -1,4 +1,4 @@
-from logging import getLogger
+import logging
 from typing import Literal, cast
 
 import torch
@@ -9,7 +9,7 @@ from typing_extensions import override
 
 from .base import CallbackConfigBase
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def grad_norm(

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/callbacks/throughput_monitor.py

@@ -1,4 +1,4 @@
-from logging import getLogger
+import logging
 from typing import Any, Literal, Protocol, TypedDict, cast, runtime_checkable
 
 from typing_extensions import NotRequired, override
@@ -6,7 +6,7 @@ from typing_extensions import NotRequired, override
 from ._throughput_monitor_callback import ThroughputMonitor as _ThroughputMonitor
 from .base import CallbackConfigBase
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 class ThroughputMonitorBatchStats(TypedDict):

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/callbacks/wandb_watch.py

@@ -1,4 +1,4 @@
-from logging import getLogger
+import logging
 from typing import Literal, Protocol, cast, runtime_checkable
 
 import torch.nn as nn
@@ -9,7 +9,7 @@ from typing_extensions import override
 
 from .base import CallbackConfigBase
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 @runtime_checkable

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/data/balanced_batch_sampler.py

@@ -1,6 +1,6 @@
 import heapq
+import logging
 from functools import cached_property
-from logging import getLogger
 from typing import Any, Protocol, runtime_checkable
 
 import numpy as np
@@ -10,7 +10,7 @@ from lightning_fabric.utilities.distributed import _DatasetSamplerWrapper
 from torch.utils.data import BatchSampler, Dataset, DistributedSampler
 from typing_extensions import override
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def _all_gather(tensor: torch.Tensor, device: torch.device | None = None):

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/model/base.py

@@ -1,7 +1,7 @@
 import inspect
+import logging
 from abc import ABC, abstractmethod
 from collections.abc import MutableMapping
-from logging import getLogger
 from typing import IO, TYPE_CHECKING, Any, Generic, cast
 
 import torch
@@ -21,7 +21,7 @@ from .modules.profiler import ProfilerMixin
 from .modules.rlp_sanity_checks import RLPSanityCheckModuleMixin
 from .modules.shared_parameters import SharedParametersModuleMixin
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 THparams = TypeVar("THparams", bound=BaseConfig, infer_variance=True)
 

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/model/config.py

@@ -1,4 +1,5 @@
 import copy
+import logging
 import os
 import string
 import time
@@ -6,7 +7,6 @@ import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
 from datetime import timedelta
-from logging import getLogger
 from pathlib import Path
 from typing import (
     Annotated,
@@ -46,7 +46,7 @@ from ..callbacks.base import CallbackConfigBase
 from ..metrics import MetricConfig
 from ._environment import EnvironmentConfig
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 class IdSeedWarning(Warning):

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/model/modules/callback.py

@@ -1,6 +1,6 @@
+import logging
 from collections import abc
 from collections.abc import Callable, Iterable
-from logging import getLogger
 from typing import Any, TypeAlias, cast, final
 
 from lightning.pytorch import Callback, LightningModule
@@ -9,7 +9,7 @@ from typing_extensions import override
 
 from ...util.typing_utils import mixin_base_type
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 CallbackFn: TypeAlias = Callable[[], Callback | Iterable[Callback] | None]
 

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/model/modules/debug.py

@@ -1,9 +1,9 @@
-from logging import getLogger
+import logging
 
 import torch
 import torch.distributed
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 class DebugModuleMixin:

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/model/modules/rlp_sanity_checks.py

@@ -1,5 +1,5 @@
+import logging
 from collections.abc import Mapping
-from logging import getLogger
 from typing import cast
 
 import torch
@@ -14,7 +14,7 @@ from ...util.typing_utils import mixin_base_type
 from ..config import BaseConfig
 from .callback import CallbackModuleMixin
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def _on_train_start_callback(trainer: Trainer, pl_module: LightningModule):

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/model/modules/shared_parameters.py

@@ -1,5 +1,5 @@
+import logging
 from collections.abc import Sequence
-from logging import getLogger
 from typing import cast
 
 import torch.nn as nn
@@ -10,7 +10,7 @@ from ...util.typing_utils import mixin_base_type
 from ..config import BaseConfig
 from .callback import CallbackRegistrarModuleMixin
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def _parameters_to_names(parameters: Sequence[nn.Parameter], model: nn.Module):

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/util/environment.py

@@ -1,8 +1,8 @@
+import logging
 import os
 from contextlib import contextmanager
-from logging import getLogger
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 @contextmanager

{nshtrainer-0.10.9 → nshtrainer-0.10.11}/src/nshtrainer/util/seed.py

@@ -1,8 +1,8 @@
-from logging import getLogger
+import logging
 
 import lightning.fabric.utilities.seed as LS
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def seed_everything(seed: int | None, *, workers: bool = False):

nshtrainer-0.10.9/src/nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -1,74 +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 .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 = "latest_epoch{epoch:02d}_step{step:04d}.ckpt"
-    """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.ckpt"
-    """Filename for the latest symlink. If None, no symlink will be created."""
-
-    @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):
-    def __init__(self, config: LatestEpochCheckpointCallbackConfig, dirpath: Path):
-        super().__init__()
-
-        self.config = config
-        self.dirpath = dirpath
-
-    def _ckpt_path(self, trainer: Trainer):
-        return self.dirpath / self.config.filename.format(
-            epoch=trainer.current_epoch, step=trainer.global_step
-        )
-
-    @override
-    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
-        # Save the new checkpoint
-        filepath = self._ckpt_path(trainer)
-        trainer.save_checkpoint(filepath, self.config.save_weights_only)
-
-        # Create the latest symlink
-        if (
-            trainer.is_global_zero
-            and (symlink_filename := self.config.latest_symlink_filename) is not None
-        ):
-            symlink_path = self.dirpath / symlink_filename
-            symlink_path.unlink(missing_ok=True)
-            symlink_path.symlink_to(filepath.name)
-            log.info(f"Created latest symlink: {symlink_path}")
-
-    def latest_checkpoint(self):
-        if (symlink_filename := self.config.latest_symlink_filename) is None:
-            return None
-
-        if not (symlink_path := self.dirpath / symlink_filename).exists():
-            return None
-
-        return symlink_path