nshtrainer 0.19.3__py3-none-any.whl → 0.21.0__py3-none-any.whl

nshtrainer/_callback.py

@@ -0,0 +1,40 @@
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from lightning.pytorch.callbacks import Callback as _LightningCallback
+
+if TYPE_CHECKING:
+    from .model import LightningModuleBase
+    from .trainer import Trainer
+
+
+class NTCallbackBase(_LightningCallback):
+    def on_checkpoint_saved(
+        self,
+        ckpt_path: Path,
+        metadata_path: Path | None,
+        trainer: "Trainer",
+        pl_module: "LightningModuleBase",
+    ) -> None:
+        """Called after a checkpoint is saved."""
+        pass
+
+
+def _call_on_checkpoint_saved(
+    trainer: "Trainer",
+    ckpt_path: str | Path,
+    metadata_path: str | Path | None,
+):
+    ckpt_path = Path(ckpt_path)
+    metadata_path = Path(metadata_path) if metadata_path else None
+
+    for callback in trainer.callbacks:
+        if not isinstance(callback, NTCallbackBase):
+            continue
+
+        callback.on_checkpoint_saved(
+            ckpt_path,
+            metadata_path,
+            trainer,
+            trainer._base_module,
+        )

nshtrainer/_checkpoint/loader.py

@@ -76,7 +76,11 @@ class CheckpointLoadingConfig(C.Config):
     """Whether to include checkpoints from HPC pre-emption."""
 
     @classmethod
-    def _auto_train(cls, ckpt: Literal["best", "last"] | str | Path | None):
+    def none(cls, include_hpc: bool = False):
+        return cls(strategies=[], include_hpc=include_hpc)
+
+    @classmethod
+    def _auto_train(cls, ckpt: Literal["best", "last", "none"] | str | Path | None):
         if ckpt is None:
             ckpt = "last"
         match ckpt:
@@ -90,6 +94,8 @@ class CheckpointLoadingConfig(C.Config):
                     strategies=[LastCheckpointStrategyConfig()],
                     include_hpc=True,
                 )
+            case "none":
+                return cls.none()
             case Path() | str():
                 ckpt = Path(ckpt)
                 return cls(
@@ -103,7 +109,7 @@ class CheckpointLoadingConfig(C.Config):
                 assert_never(ckpt)
 
     @classmethod
-    def _auto_eval(cls, ckpt: Literal["best", "last"] | str | Path | None):
+    def _auto_eval(cls, ckpt: Literal["best", "last", "none"] | str | Path | None):
         if ckpt is None:
             log.warn("No checkpoint specified for evaluation. Defaulting to `last`.")
             ckpt = "last"
@@ -119,6 +125,8 @@ class CheckpointLoadingConfig(C.Config):
                     strategies=[LastCheckpointStrategyConfig()],
                     include_hpc=False,
                 )
+            case "none":
+                return cls.none(include_hpc=False)
             case Path() | str():
                 ckpt = Path(ckpt)
                 return cls(
@@ -131,7 +139,7 @@ class CheckpointLoadingConfig(C.Config):
     @classmethod
     def auto(
         cls,
-        ckpt: Literal["best", "last"] | str | Path | None,
+        ckpt: Literal["best", "last", "none"] | str | Path | None,
         trainer_mode: TrainerFn,
     ):
         """
@@ -142,7 +150,7 @@ class CheckpointLoadingConfig(C.Config):
 
         Parameters:
         -----------
-        ckpt : Literal["best", "last"] | str | Path | None
+        ckpt : Literal["best", "last", "none"] | str | Path | None
             Specifies the checkpoint loading preference:
             - "best": Use the best checkpoint based on the primary metric.
             - "last": Use the most recent checkpoint.

nshtrainer/_checkpoint/metadata.py

@@ -96,13 +96,17 @@ def _generate_checkpoint_metadata(
     )
 
 
+def _metadata_path(checkpoint_path: Path):
+    return checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
+
+
 def _write_checkpoint_metadata(
     trainer: "Trainer",
     model: "LightningModuleBase",
     checkpoint_path: Path,
 ):
     config = cast("BaseConfig", model.config)
-    metadata_path = checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
+    metadata_path = _metadata_path(checkpoint_path)
     metadata = _generate_checkpoint_metadata(
         config, trainer, checkpoint_path, metadata_path
     )
@@ -119,7 +123,7 @@ def _write_checkpoint_metadata(
 
 
 def _remove_checkpoint_metadata(checkpoint_path: Path):
-    path = checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
+    path = _metadata_path(checkpoint_path)
     try:
         path.unlink(missing_ok=True)
     except Exception:
@@ -133,8 +137,8 @@ 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
-    path = checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
-    linked_path = linked_checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
+    path = _metadata_path(checkpoint_path)
+    linked_path = _metadata_path(linked_checkpoint_path)
     try:
         try:
             # linked_path.symlink_to(path)

nshtrainer/_checkpoint/saver.py

@@ -4,6 +4,7 @@ from pathlib import Path
 
 from lightning.pytorch import Trainer
 
+from ..util.path import get_relative_path
 from .metadata import _link_checkpoint_metadata, _remove_checkpoint_metadata
 
 
@@ -14,10 +15,8 @@ def _link_checkpoint(
     metadata: bool,
     remove_existing: bool = True,
 ):
-    if not isinstance(filepath, Path):
-        filepath = Path(filepath)
-    if not isinstance(linkpath, Path):
-        linkpath = Path(linkpath)
+    filepath = Path(filepath)
+    linkpath = Path(linkpath)
 
     if remove_existing:
         if linkpath.exists():
@@ -30,7 +29,7 @@ def _link_checkpoint(
             _remove_checkpoint_metadata(linkpath)
 
     try:
-        linkpath.symlink_to(filepath.relative_to(linkpath.parent))
+        linkpath.symlink_to(get_relative_path(linkpath, filepath))
     except OSError:
         # on Windows, special permissions are required to create symbolic links as a regular user
         # fall back to copying the file
@@ -46,9 +45,9 @@ def _remove_checkpoint(
     *,
     metadata: bool,
 ):
-    if not isinstance(filepath, Path):
-        filepath = Path(filepath)
+    filepath = Path(filepath)
 
     trainer.strategy.remove_checkpoint(filepath)
+
     if metadata:
         _remove_checkpoint_metadata(filepath)

nshtrainer/_hf_hub.py

@@ -6,11 +6,10 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
 import nshconfig as C
-from lightning.pytorch import Callback
-from lightning.pytorch.trainer import Trainer
 from nshrunner._env import SNAPSHOT_DIR
 from typing_extensions import override
 
+from ._callback import NTCallbackBase
 from .callbacks.base import (
     CallbackConfigBase,
     CallbackMetadataConfig,
@@ -22,6 +21,8 @@ if TYPE_CHECKING:
 
     from .model.base import BaseConfig
     from .trainer.trainer import Trainer
+
+
 log = logging.getLogger(__name__)
 
 
@@ -102,9 +103,9 @@ def _api(token: str | None = None):
 
         # Verify authentication
         api.whoami()
-    except Exception as e:
+    except Exception:
         log.exception(
-            f"Authentication failed for Hugging Face Hub: {str(e)}. "
+            "Authentication failed for Hugging Face Hub. "
             "Please make sure you are logged in using `huggingface-cli login`, "
             "by setting the HUGGING_FACE_HUB_TOKEN environment variable, "
             "or by providing a valid token in the configuration."
@@ -210,10 +211,10 @@ def _init(*, trainer: "Trainer", root_config: "BaseConfig"):
                 exist_ok=True,
             )
             log.info(f"Created new repository '{repo_name}'.")
-        except Exception as e:
-            log.exception(f"Failed to create repository '{repo_name}': {str(e)}")
-    except Exception as e:
-        log.exception(f"Error checking repository '{repo_name}': {str(e)}")
+        except Exception:
+            log.exception(f"Failed to create repository '{repo_name}'")
+    except Exception:
+        log.exception(f"Error checking repository '{repo_name}'")
 
     # Upload the config
     _save_config(root_config, trainer=trainer)
@@ -262,9 +263,9 @@ def _save_code(
         log.info(
             f"Uploaded snapshot contents to repository '{repo_name}' under 'code' folder."
         )
-    except Exception as e:
+    except Exception:
         log.exception(
-            f"Failed to upload snapshot contents to repository '{repo_name}' under 'code' folder: {str(e)}"
+            f"Failed to upload snapshot contents to repository '{repo_name}' under 'code' folder."
         )
 
 
@@ -300,10 +301,8 @@ def _save_config(
             run_as_future=cast(Any, config.save_in_background),
         )
         log.info(f"Uploaded config.json to repository '{repo_name}'.")
-    except Exception as e:
-        log.exception(
-            f"Failed to upload config.json to repository '{repo_name}': {str(e)}"
-        )
+    except Exception:
+        log.exception(f"Failed to upload config.json to repository '{repo_name}'.")
 
 
 def _save_checkpoint_files(
@@ -331,17 +330,24 @@ def _save_checkpoint_files(
     # Read all the files to memory
     file_contents: list[bytes | None] = []
     for p in paths:
+        assert not p.is_symlink(), f"Path {p} is a symlink."
+        assert p.is_file(), f"Path {p} is not a file."
         try:
             with open(p, "rb") as f:
                 file_contents.append(f.read())
-        except IOError as e:
-            log.warning(f"Failed to read checkpoint file {p}: {str(e)}")
+        except IOError:
+            log.exception(f"Failed to read checkpoint file {p}.")
             file_contents.append(None)
 
-    for p, contents in zip(paths, file_contents):
-        if contents is None:
-            continue
+    # Remove the paths that failed to read
+    file_contents_and_paths = [
+        (contents, p)
+        for contents, p in zip(file_contents, paths)
+        if contents is not None
+    ]
 
+    # Upload the checkpoint files to the repository
+    for contents, p in file_contents_and_paths:
         try:
             relative_path = p.relative_to(checkpoint_dir)
         except ValueError:
@@ -365,21 +371,136 @@ def _save_checkpoint_files(
             log.info(
                 f"Uploaded checkpoint file {relative_path} to repository '{repo_name}'."
             )
-        except Exception as e:
+        except Exception:
             log.exception(
-                f"Failed to upload checkpoint file {relative_path} to repository '{repo_name}': {str(e)}"
+                f"Failed to upload checkpoint file {relative_path} to repository '{repo_name}'."
             )
 
     log.info(f"Completed uploading checkpoint files to repository '{repo_name}'.")
 
 
-class HFHubCallback(Callback):
+def _save_checkpoint_symlinks(
+    trainer: "Trainer",
+    paths: list[Path],
+    *,
+    root_config: "BaseConfig",
+):
+    config = root_config.trainer.hf_hub
+    if (
+        api := _enabled_and_valid(trainer, config, rank_zero_only=True)
+    ) is None or not config.save_checkpoints:
+        return
+
+    # Resolve the checkpoint directory
+    checkpoint_dir = root_config.directory.resolve_subdirectory(
+        root_config.id, "checkpoint"
+    )
+
+    # Resolve the repository name
+    repo_name = _repo_name(api, root_config)
+
+    # Create a commit for copying the files
+    from huggingface_hub.hf_api import CommitOperation, CommitOperationCopy
+
+    commits: list[CommitOperation] = []
+    for p in paths:
+        assert p.is_symlink(), f"Path {p} is not a symlink."
+
+        try:
+            dest_relative_path = p.relative_to(checkpoint_dir)
+        except ValueError:
+            log.warning(
+                f"Checkpoint path {p} is not within the checkpoint directory {checkpoint_dir}."
+            )
+            continue
+
+        try:
+            source_relative_path = p.resolve().relative_to(checkpoint_dir)
+        except ValueError:
+            log.warning(
+                f"Checkpoint symlink target {p.resolve()} is not within the checkpoint directory {checkpoint_dir}."
+            )
+            continue
+
+        # Prefix the path in repo with "checkpoints"
+        dest_path_in_repo = Path("checkpoints") / dest_relative_path
+        source_path_in_repo = Path("checkpoints") / source_relative_path
+
+        # Create and append a CommitOperationCopy for copying the symlink
+        copy_op = CommitOperationCopy(
+            src_path_in_repo=str(source_path_in_repo),
+            path_in_repo=str(dest_path_in_repo),
+        )
+        commits.append(copy_op)
+
+    log.info(f"Creating a commit with the following operations: {commits}")
+
+    try:
+        api.create_commit(
+            repo_id=repo_name,
+            repo_type="model",
+            commit_message="Copy checkpoint symlinks",
+            operations=commits,
+            run_as_future=cast(Any, config.save_in_background),
+        )
+        log.info(
+            f"Created commit to copy checkpoint symlinks to repository '{repo_name}'."
+        )
+    except Exception:
+        log.exception(
+            f"Failed to create commit to copy checkpoint symlinks to repository '{repo_name}'"
+        )
+
+    log.info(f"Completed copying checkpoint symlinks to repository '{repo_name}'.")
+
+
+def _save_checkpoint_directory(trainer: "Trainer", *, root_config: "BaseConfig"):
+    config = root_config.trainer.hf_hub
+    if (
+        api := _enabled_and_valid(trainer, config, rank_zero_only=True)
+    ) is None or not config.save_checkpoints:
+        return
+
+    # Resolve the checkpoint directory
+    checkpoint_dir = root_config.directory.resolve_subdirectory(
+        root_config.id, "checkpoint"
+    )
+
+    # Resolve the repository name
+    repo_name = _repo_name(api, root_config)
+
+    # Upload the checkpoint directory to the repository
+    try:
+        api.upload_folder(
+            folder_path=str(checkpoint_dir),
+            repo_id=repo_name,
+            repo_type="model",
+            path_in_repo="checkpoints",
+            run_as_future=cast(Any, config.save_in_background),
+        )
+        log.info(f"Uploaded checkpoint directory to repository '{repo_name}'.")
+    except Exception:
+        log.exception(
+            f"Failed to upload checkpoint directory to repository '{repo_name}'."
+        )
+
+    log.info(f"Completed uploading checkpoint files to repository '{repo_name}'.")
+
+
+class HFHubCallback(NTCallbackBase):
     def __init__(self, config: HuggingFaceHubConfig):
         super().__init__()
         self.config = config
 
     @override
     def setup(self, trainer, pl_module, stage):
+        from .trainer.trainer import Trainer
+
+        if not isinstance(trainer, Trainer):
+            raise ValueError(
+                f"HFHubCallback requires a `nshtrainer.Trainer` instance, got {type(trainer)}."
+            )
+
         root_config = cast("BaseConfig", pl_module.hparams)
         _init(trainer=trainer, root_config=root_config)
 
@@ -387,3 +508,18 @@ class HFHubCallback(Callback):
     def teardown(self, trainer, pl_module, stage):
         if hasattr(trainer, "_hf_hub_api"):
             delattr(trainer, "_hf_hub_api")
+
+    @override
+    def on_checkpoint_saved(self, ckpt_path, metadata_path, trainer, pl_module):
+        root_config = cast("BaseConfig", pl_module.hparams)
+
+        # If HF Hub is enabled, then we upload
+        if root_config.trainer.hf_hub and trainer.is_global_zero:
+            # Upload the regular files first, then the symlinks
+            all_paths = [p for p in (ckpt_path, metadata_path) if p is not None]
+            if regular_paths := [p for p in all_paths if not p.is_symlink()]:
+                _save_checkpoint_files(trainer, regular_paths, root_config=root_config)
+            if symlink_paths := [p for p in all_paths if p.is_symlink()]:
+                _save_checkpoint_symlinks(
+                    trainer, symlink_paths, root_config=root_config
+                )

nshtrainer/callbacks/checkpoint/_base.py

@@ -9,7 +9,7 @@ from lightning.pytorch import Trainer
 from lightning.pytorch.callbacks import Checkpoint
 from typing_extensions import TypeVar, override
 
-from ..._checkpoint.metadata import CheckpointMetadata
+from ..._checkpoint.metadata import CheckpointMetadata, _metadata_path
 from ..._checkpoint.saver import _link_checkpoint, _remove_checkpoint
 from ..base import CallbackConfigBase
 
@@ -65,8 +65,6 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         self.dirpath.mkdir(parents=True, exist_ok=True)
         self.symlink_dirpath = dirpath
 
-        self._last_global_step_saved = 0
-
     @abstractmethod
     def default_filename(self) -> str: ...
 
@@ -144,9 +142,21 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         if self._should_skip_saving_checkpoint(trainer):
             return
 
+        from ...trainer import Trainer as NTTrainer
+
+        if not isinstance(trainer, NTTrainer):
+            raise TypeError(
+                f"Trainer must be an instance of {NTTrainer.__name__}, "
+                f"but got {type(trainer).__name__}"
+            )
+
         # Save the new checkpoint
         filepath = self.resolve_checkpoint_path(self.current_metrics(trainer))
-        trainer.save_checkpoint(filepath, self.config.save_weights_only)
+        trainer._nshtrainer_save_checkpoint(
+            filepath,
+            self.config.save_weights_only,
+            use_checkpoint_cache=None,
+        )
 
         if trainer.is_global_zero:
             # Create the latest symlink
@@ -162,8 +172,15 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         # 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
+        # Call the on save checkpoint callback for the symlink (if it exists)
+        if (symlink_filename := self.symlink_path()) is not None:
+            from ... import _callback
+
+            symlink_path = self.dirpath / symlink_filename
+            symlink_metadata_path = _metadata_path(symlink_path)
+            _callback._call_on_checkpoint_saved(
+                trainer, symlink_path, symlink_metadata_path
+            )
 
     def _should_skip_saving_checkpoint(self, trainer: Trainer) -> bool:
         from lightning.pytorch.trainer.states import TrainerFn
@@ -175,6 +192,4 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
             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/model/config.py

@@ -811,11 +811,14 @@ class SanityCheckingConfig(C.Config):
 
 
 class TrainerConfig(C.Config):
-    ckpt_path: str | Path | None = None
-    """Path to a checkpoint to load and resume training from."""
+    ckpt_path: Literal["none"] | str | Path | None = None
+    """Path to a checkpoint to load and resume training from. If ``"none"``, will not load a checkpoint."""
 
-    checkpoint_loading: CheckpointLoadingConfig | Literal["auto"] = "auto"
-    """Checkpoint loading configuration options."""
+    checkpoint_loading: CheckpointLoadingConfig | Literal["auto", "none"] = "auto"
+    """Checkpoint loading configuration options.
+    `"auto"` will automatically determine the best checkpoint loading strategy based on the provided.
+    `"none"` will disable checkpoint loading.
+    """
 
     checkpoint_saving: CheckpointSavingConfig = CheckpointSavingConfig()
     """Checkpoint saving configuration options."""
@@ -1009,6 +1012,8 @@ class TrainerConfig(C.Config):
     """If enabled, the model supports scaling the gradients of shared parameters that are registered using `LightningModuleBase.register_shared_parameters(...)`"""
     save_checkpoint_metadata: bool = True
     """If enabled, will save additional metadata whenever a checkpoint is saved."""
+    use_checkpoint_cache: bool = True
+    """If enabled, will optimize the saving of duplicate checkpoints by creating symlinks instead of copying the file."""
 
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()
     """

nshtrainer/trainer/checkpoint_connector.py

@@ -31,8 +31,14 @@ class _CheckpointConnector(_LightningCheckpointConnector):
 
         # Now, resolve the checkpoint loader config.
         root_config = cast("BaseConfig", trainer._base_module.config)
-        if (ckpt_loader_config := root_config.trainer.checkpoint_loading) == "auto":
-            ckpt_loader_config = CheckpointLoadingConfig.auto(ckpt_path, state_fn)
+        ckpt_loader_config = root_config.trainer.checkpoint_loading
+        match ckpt_loader_config:
+            case "auto":
+                ckpt_loader_config = CheckpointLoadingConfig.auto(ckpt_path, state_fn)
+            case "none":
+                ckpt_loader_config = CheckpointLoadingConfig.none()
+            case _:
+                pass
         log.debug(f"Checkpoint loader config: {ckpt_loader_config}")
 
         # Use the config to resolve the checkpoint.

nshtrainer/trainer/trainer.py

@@ -17,6 +17,7 @@ from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
 from typing_extensions import Unpack, assert_never, override
 
 from .._checkpoint.metadata import _write_checkpoint_metadata
+from .._checkpoint.saver import _link_checkpoint
 from ..callbacks.base import resolve_all_callbacks
 from ..model.config import (
     AcceleratorConfigProtocol,
@@ -285,6 +286,8 @@ class Trainer(LightningTrainer):
         /,
         **kwargs: Unpack[LightningTrainerKwargs],
     ):
+        self._nshtrainer_checkpoint_cache: dict[tuple[int, int], Path] = {}
+
         self._pre_init(config)
 
         kwargs = self._update_kwargs(config, kwargs)
@@ -407,29 +410,70 @@ class Trainer(LightningTrainer):
 
         return super()._run(model, ckpt_path)
 
-    @override
-    def save_checkpoint(
+    def _nshtrainer_save_checkpoint(
         self,
         filepath: str | Path,
         weights_only: bool = False,
         storage_options: Any | None = None,
+        use_checkpoint_cache: bool | None = None,
     ):
+        lm = self._base_module
+        hparams = cast(BaseConfig, lm.hparams)
+        if use_checkpoint_cache is None:
+            use_checkpoint_cache = hparams.trainer.use_checkpoint_cache
+
         filepath = Path(filepath)
-        ret_val = super().save_checkpoint(filepath, weights_only, storage_options)
+
+        # List of files that we should upload to HF
+        written_files: list[Path] = [filepath]
+
+        cached_path = None
+        if (
+            use_checkpoint_cache
+            and (
+                cached_path := self._nshtrainer_checkpoint_cache.get(
+                    (self.current_epoch, self.global_step)
+                )
+            )
+            is not None
+        ):
+            # If we have a cached path, then we symlink it to the new path.
+            log.info(f"Re-using cached checkpoint {cached_path} for {filepath}.")
+            _link_checkpoint(cached_path, filepath, metadata=False)
+        else:
+            super().save_checkpoint(filepath, weights_only, storage_options)
+
+        # If we are using the cache but we don't have a cached path, then we save the checkpoint to the cache.
+        if use_checkpoint_cache and cached_path is None:
+            self._nshtrainer_checkpoint_cache[
+                (self.current_epoch, self.global_step)
+            ] = filepath
+            log.debug(f"Checkpoint saved to cache: {filepath}")
 
         # Save the checkpoint metadata
-        lm = self._base_module
-        hparams = cast(BaseConfig, lm.hparams)
         metadata_path = None
         if hparams.trainer.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
-            metadata_path = _write_checkpoint_metadata(self, lm, filepath)
+            if (
+                metadata_path := _write_checkpoint_metadata(self, lm, filepath)
+            ) is not None:
+                written_files.append(metadata_path)
 
-        # If HF Hub is enabled, then we upload
-        if hparams.trainer.hf_hub and self.is_global_zero:
-            from .._hf_hub import _save_checkpoint_files
+        # Call the `on_checkpoint_saved` method on all callbacks
+        from .. import _callback
 
-            files = [f for f in (filepath, metadata_path) if f is not None]
-            _save_checkpoint_files(self, files, root_config=hparams)
+        _callback._call_on_checkpoint_saved(self, filepath, metadata_path)
 
-        return ret_val
+    @override
+    def save_checkpoint(
+        self,
+        filepath: str | Path,
+        weights_only: bool = False,
+        storage_options: Any | None = None,
+    ):
+        return self._nshtrainer_save_checkpoint(
+            filepath=filepath,
+            weights_only=weights_only,
+            storage_options=storage_options,
+            use_checkpoint_cache=False,
+        )

nshtrainer/util/path.py

@@ -0,0 +1,29 @@
+import os
+from pathlib import Path
+from typing import TypeAlias
+
+_Path: TypeAlias = str | Path | os.PathLike
+
+
+def get_relative_path(source: _Path, destination: _Path):
+    # Get the absolute paths
+    source = os.path.abspath(source)
+    destination = os.path.abspath(destination)
+
+    # Split the paths into components
+    source_parts = source.split(os.sep)
+    destination_parts = destination.split(os.sep)
+
+    # Find the point where the paths diverge
+    i = 0
+    for i in range(min(len(source_parts), len(destination_parts))):
+        if source_parts[i] != destination_parts[i]:
+            break
+    else:
+        i += 1
+
+    # Build the relative path
+    up = os.sep.join([".." for _ in range(len(source_parts) - i - 1)])
+    down = os.sep.join(destination_parts[i:])
+
+    return Path(os.path.normpath(os.path.join(up, down)))

{nshtrainer-0.19.3.dist-info → nshtrainer-0.21.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.19.3
+Version: 0.21.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -26,6 +26,7 @@ Requires-Dist: torchmetrics ; extra == "extra"
 Requires-Dist: typing-extensions
 Requires-Dist: wandb ; extra == "extra"
 Requires-Dist: wrapt ; extra == "extra"
+Requires-Dist: zstandard ; extra == "extra"
 Description-Content-Type: text/markdown
 
 

{nshtrainer-0.19.3.dist-info → nshtrainer-0.21.0.dist-info}/RECORD

@@ -1,15 +1,16 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
-nshtrainer/_checkpoint/loader.py,sha256=vc9f22qEDVw-y-Clpy71jVeI2EPxWNqRy-cAslMTb8c,13868
-nshtrainer/_checkpoint/metadata.py,sha256=p5e7dhVPpOGrXeuesq_7Y_RHi5lguzDAR_UXtMJXzWU,5175
-nshtrainer/_checkpoint/saver.py,sha256=DkbCH0YeOJ71m32vAARiQdGBf0hvwwdoAV8LOFGy-0Y,1428
+nshtrainer/_callback.py,sha256=A1zLsTy4b_wOYnInLLXGSRdHzT2yNa6mPEql-ozm0u0,1013
+nshtrainer/_checkpoint/loader.py,sha256=5vjg-OFChXJjgiOVv8vnV8nwTscfdDtEdxQRz6uPfDE,14158
+nshtrainer/_checkpoint/metadata.py,sha256=TLAt7yR3KhSRbXCtomLMxcMvOiAju873A1ZRo8VWNwA,5179
+nshtrainer/_checkpoint/saver.py,sha256=6W-Rbc3QGuhcF_mcwN8v31uEjLQCsZvt8CPuqPs4m5g,1342
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
-nshtrainer/_hf_hub.py,sha256=To3BnnGWbMNNMBdzVtgrNOcNU2fi1dQpwwuclusFAbI,12169
+nshtrainer/_hf_hub.py,sha256=0bOhJNyIjQGJsMRaW7qQJc1oTnUMHj08auuztzTQvZ0,16906
 nshtrainer/callbacks/__init__.py,sha256=4qocBDzQbLLhhbIEfvbA3SQB_Dy9ZJH7keMwPay-ZS8,2359
 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=g-3zIthupERKqWZQw-A_busQPaPRkto6iHBV-M7nK1Y,527
-nshtrainer/callbacks/checkpoint/_base.py,sha256=YT_V-oihO9iB4ETl46CGYTCQjIYl-CpV7TMViTn07Lk,6144
+nshtrainer/callbacks/checkpoint/_base.py,sha256=r6IPpl3sGUmxBNv80y9r326lTrPAIVSU3Fu-3LrYH2s,6691
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=DJiLo7NDzd-lp-O3v7Cv8WejyjXPV_6_RmfltKO9fvE,2165
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CqB_8Xv32rtpLCaEEPi6DbRZm4ph5TWS-LfqIHXUIUA,1097
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
@@ -57,7 +58,7 @@ nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy0
 nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ,1288
 nshtrainer/model/__init__.py,sha256=VyRziPT3YilP6xjLi_StsSqtlvn7N4LOMzgukRsOnF8,1380
 nshtrainer/model/base.py,sha256=oQVolDk81acy4OlckwQEBHuX2gCaVSYiIA0JaDIfhQ4,17517
-nshtrainer/model/config.py,sha256=147uV7IukvuYE4G_ZuQNxVjnlog1BdCrAVbcj_sx9Vs,43104
+nshtrainer/model/config.py,sha256=22_xIcdEO2pJzXgrFaqGFtk3PQEiwKiMZY1cjhoyWaA,43486
 nshtrainer/model/modules/callback.py,sha256=1z6gUDBd35KG3phGzRekgZM6SIk-wj5Uo6APN4YhRR0,8549
 nshtrainer/model/modules/debug.py,sha256=Yy7XEdPou9BkCsD5hJchwJGmCVGrfUru5g9VjPM4uAw,1120
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
@@ -75,16 +76,17 @@ nshtrainer/runner.py,sha256=USAjrExHkN5oVNVunsoPnLxfQrEHSaa54S3RipOe544,3605
 nshtrainer/scripts/find_packages.py,sha256=ixYivZobumyyGsf2B9oYMLyLTRcBzY_vUv-u3bNW-hs,1424
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
-nshtrainer/trainer/checkpoint_connector.py,sha256=F2tkHogbMAa5U7335sm77sZBkjEDa5v46XbJCH9Mg6c,2167
+nshtrainer/trainer/checkpoint_connector.py,sha256=r0ir4xYSdf_jebM0x09qaO6nJsvsiRQDyM0fs80ppOQ,2347
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
-nshtrainer/trainer/trainer.py,sha256=TTtVkgSB_ekgDlHg24d58Vzddtkpp6ZHOTVprXdXMH0,17503
+nshtrainer/trainer/trainer.py,sha256=DNKA4mcW083i7qLk0fi3j5-Qj4KNBtlLuyIsNxykebw,19100
 nshtrainer/util/_environment_info.py,sha256=gIdq9TJgzGCdcVzZxjHcwYasJ_HmEGVHbvE-KJVVtWs,24187
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
 nshtrainer/util/environment.py,sha256=AeW_kLl-N70wmb6L_JLz1wRj0kA70xs6RCmc9iUqczE,4159
+nshtrainer/util/path.py,sha256=A_Ocag3_hbwns_zAxFDlH-5eVHWFlcy2DKxHQ7jddvk,837
 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.19.3.dist-info/METADATA,sha256=0LXA6hNdn7QjIJEGK-tvPQHiTEuxwsDzifvEyBMCYmo,935
-nshtrainer-0.19.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.19.3.dist-info/RECORD,,
+nshtrainer-0.21.0.dist-info/METADATA,sha256=7QfSX_yXi-Up6uxOVFfDPn4ieGK5b3UgQfO_KFsNzXk,979
+nshtrainer-0.21.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.21.0.dist-info/RECORD,,