nshtrainer 0.23.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

nshtrainer/_checkpoint/metadata.py

@@ -11,7 +11,7 @@ import numpy as np
 import torch
 
 from ..util._environment_info import EnvironmentConfig
-from ..util.path import get_relative_path
+from ..util.path import compute_file_checksum, get_relative_path
 
 if TYPE_CHECKING:
     from ..model import BaseConfig, LightningModuleBase
@@ -28,6 +28,7 @@ class CheckpointMetadata(C.Config):
 
     checkpoint_path: Path
     checkpoint_filename: str
+    checkpoint_checksum: str
 
     run_id: str
     name: str
@@ -81,6 +82,7 @@ def _generate_checkpoint_metadata(
         # moving the checkpoint directory
         checkpoint_path=checkpoint_path.relative_to(metadata_path.parent),
         checkpoint_filename=checkpoint_path.name,
+        checkpoint_checksum=compute_file_checksum(checkpoint_path),
         run_id=config.id,
         name=config.run_name,
         project=config.project,

nshtrainer/_hf_hub.py

@@ -1,7 +1,9 @@
-import io
+import contextlib
 import logging
 import os
 import re
+from dataclasses import dataclass
+from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
@@ -16,7 +18,6 @@ if TYPE_CHECKING:
     from huggingface_hub import HfApi  # noqa: F401
 
     from .model.base import BaseConfig
-    from .trainer.trainer import Trainer
 
 
 log = logging.getLogger(__name__)
@@ -108,36 +109,6 @@ def _api(token: str | None = None):
     return api
 
 
-def _enabled_and_valid(
-    trainer: "Trainer",
-    config: HuggingFaceHubConfig,
-    *,
-    rank_zero_only: bool,
-):
-    # Make sure this is enabled and the config is valid
-    if not config:
-        return None
-
-    # If `rank_zero_only` and this is not rank 0, stop here.
-    if rank_zero_only and not trainer.is_global_zero:
-        return
-
-    # Make sure that `huggingface_hub` is installed
-    try:
-        import huggingface_hub  # noqa: F401
-    except ImportError:
-        log.exception(
-            "Could not import `huggingface_hub`. Please install it using `pip install huggingface_hub`."
-        )
-        return None
-
-    # Create and authenticate the API instance
-    if (api := getattr(trainer, "_hf_hub_api", None)) is None:
-        api = _api(config.token)
-        setattr(trainer, "_hf_hub_api", api)
-    return cast(huggingface_hub.HfApi, api)
-
-
 def _repo_name(api: "HfApi", root_config: "BaseConfig"):
     username = None
     if (ac := root_config.trainer.hf_hub.auto_create) and ac.namespace:
@@ -173,346 +144,210 @@ def _repo_name(api: "HfApi", root_config: "BaseConfig"):
     return f"{username}/{repo_name}"
 
 
-def _init(*, 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.auto_create:
-        return
-
-    from huggingface_hub.utils import RepositoryNotFoundError
-
-    # Resolve the repository name
-    repo_name = _repo_name(api, root_config)
-
-    # Create the repository, if it doesn't exist
-    try:
-        # Check if the repository exists
-        api.repo_info(repo_id=repo_name, repo_type="model")
-        log.info(f"Repository '{repo_name}' already exists.")
-    except RepositoryNotFoundError:
-        # Repository doesn't exist, so create it
-        try:
-            api.create_repo(
-                repo_id=repo_name,
-                repo_type="model",
-                private=config.auto_create.private,
-                exist_ok=True,
-            )
-            log.info(f"Created new repository '{repo_name}'.")
-        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)
-
-    # Upload the code
-    _save_code(repo_name, config=config, trainer=trainer)
-
-
-def _save_code(
-    repo_name: str,
-    *,
-    config: HuggingFaceHubConfig,
-    trainer: "Trainer",
-):
-    if (
-        api := _enabled_and_valid(
-            trainer,
-            config,
-            rank_zero_only=True,
-        )
-    ) is None or not config.save_code:
-        return
-
-    # If a snapshot has been taken (which can be detected using the SNAPSHOT_DIR env),
-    # then upload all contents within the snapshot directory to the repository.
-    snapshot_dir = os.environ.get(SNAPSHOT_DIR)
-    if not snapshot_dir:
-        log.info("No snapshot directory found. Skipping upload.")
-        return
-
-    snapshot_path = Path(snapshot_dir)
-    if not snapshot_path.exists() or not snapshot_path.is_dir():
-        log.warning(
-            f"Snapshot directory '{snapshot_dir}' does not exist or is not a directory."
-        )
-        return
-
-    try:
-        api.upload_folder(
-            folder_path=str(snapshot_path),
-            repo_id=repo_name,
-            repo_type="model",
-            path_in_repo="code",  # Prefix with "code" folder
-            run_as_future=cast(Any, config.save_in_background),
-        )
-        log.info(
-            f"Uploaded snapshot contents to repository '{repo_name}' under 'code' folder."
-        )
-    except Exception:
-        log.exception(
-            f"Failed to upload snapshot contents to repository '{repo_name}' under 'code' folder."
-        )
-
-
-def _save_config(
-    root_config: "BaseConfig",
-    *,
-    trainer: "Trainer",
-):
-    config = root_config.trainer.hf_hub
-    if (
-        api := _enabled_and_valid(
-            trainer,
-            config,
-            rank_zero_only=True,
+@dataclass
+class _Upload:
+    local_path: Path
+    path_in_repo: Path
+
+    @classmethod
+    def from_local_path(
+        cls,
+        local_path: Path,
+        root_config: "BaseConfig",
+    ):
+        # Resolve the checkpoint directory
+        checkpoint_dir = root_config.directory.resolve_subdirectory(
+            root_config.id, "checkpoint"
         )
-    ) is None or not config.save_config:
-        return
-
-    # Convert the root config to a JSON string
-    # NOTE: This is a utf-8 string.
-    config_json = root_config.model_dump_json(indent=4)
 
-    # Resolve the repository name
-    repo_name = _repo_name(api, root_config)
-
-    # Upload the config file to the repository
-    try:
-        api.upload_file(
-            path_or_fileobj=config_json.encode("utf-8"),
-            path_in_repo="config.json",
-            repo_id=repo_name,
-            repo_type="model",
-            run_as_future=cast(Any, config.save_in_background),
-        )
-        log.info(f"Uploaded config.json to repository '{repo_name}'.")
-    except Exception:
-        log.exception(f"Failed to upload config.json to repository '{repo_name}'.")
-
-
-def _save_checkpoint_files(
-    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)
-
-    # Let's read all the files to memory right now,
-    # in case they get used/removed by other processes.
-    # 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:
-            log.exception(f"Failed to read checkpoint file {p}.")
-            file_contents.append(None)
-
-    # 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)
+            relative_path = local_path.relative_to(checkpoint_dir)
         except ValueError:
-            log.warning(
-                f"Checkpoint path {p} is not within the checkpoint directory {checkpoint_dir}."
+            raise ValueError(
+                f"Checkpoint path {local_path} is not within the checkpoint directory {checkpoint_dir}."
             )
-            continue
 
         # Prefix the path in repo with "checkpoints"
         path_in_repo = Path("checkpoints") / relative_path
 
-        # Upload the checkpoint file to the repository
-        try:
-            api.upload_file(
-                path_or_fileobj=io.BytesIO(contents),
-                path_in_repo=str(path_in_repo),
-                repo_id=repo_name,
-                repo_type="model",
-                run_as_future=cast(Any, config.save_in_background),
-            )
-            log.info(
-                f"Uploaded checkpoint file {relative_path} to repository '{repo_name}'."
-            )
-        except Exception:
-            log.exception(
-                f"Failed to upload checkpoint file {relative_path} to repository '{repo_name}'."
-            )
-
-    log.info(f"Completed uploading checkpoint files to repository '{repo_name}'.")
-
+        return cls(local_path=local_path, path_in_repo=path_in_repo)
 
-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"
-    )
+class HFHubCallback(NTCallbackBase):
+    @contextlib.contextmanager
+    def _with_error_handling(self, opeartion: str):
+        try:
+            yield
+        except Exception:
+            log.exception(f"Failed to {opeartion}, repo_id={self._repo_id}")
+        else:
+            log.debug(f"Successfully {opeartion}, repo_id={self._repo_id}")
 
-    # Resolve the repository name
-    repo_name = _repo_name(api, root_config)
+    def __init__(self, config: HuggingFaceHubConfig):
+        super().__init__()
 
-    # Create a commit for copying the files
-    from huggingface_hub.hf_api import CommitOperation, CommitOperationCopy
+        self.config = config
 
-    commits: list[CommitOperation] = []
-    for p in paths:
-        assert p.is_symlink(), f"Path {p} is not a symlink."
+        self._repo_id = None
+        self._checksum_to_path_in_repo: dict[str, Path] = {}
 
-        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
+    @override
+    def setup(self, trainer, pl_module, stage):
+        from .trainer.trainer import Trainer
 
-        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}."
+        if not isinstance(trainer, Trainer):
+            raise ValueError(
+                f"HFHubCallback requires a `nshtrainer.Trainer` instance, got {type(trainer)}."
             )
-            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
+        root_config = cast("BaseConfig", pl_module.hparams)
 
-        # 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),
+        # Create the repository, if it doesn't exist
+        self._repo_id = self.api.create_repo(
+            repo_id=_repo_name(self.api, root_config),
+            repo_type="model",
+            private=self.config.auto_create.private,
+            exist_ok=True,
         )
-        commits.append(copy_op)
 
-    log.info(f"Creating a commit with the following operations: {commits}")
+        # Upload the config and code
+        self._save_config(root_config)
+        self._save_code()
 
-    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}'"
-        )
+    @override
+    def on_checkpoint_saved(self, ckpt_path, metadata_path, trainer, pl_module):
+        root_config = cast("BaseConfig", pl_module.hparams)
 
-    log.info(f"Completed copying checkpoint symlinks to repository '{repo_name}'.")
+        # If HF Hub is enabled, then we upload
+        if self.config and trainer.is_global_zero:
+            with self._with_error_handling("save checkpoints"):
+                self._save_checkpoint(
+                    _Upload.from_local_path(ckpt_path, root_config),
+                    _Upload.from_local_path(metadata_path, root_config)
+                    if metadata_path is not None
+                    else None,
+                )
 
+    @cached_property
+    def api(self):
+        # Create and authenticate the API instance
+        if (api := _api(self.config.token)) is None:
+            raise ValueError("Failed to create Hugging Face Hub API instance.")
+        return api
+
+    @property
+    def repo_id(self):
+        if self._repo_id is None:
+            raise ValueError("Repository id has not been initialized.")
+        return self._repo_id
+
+    def _save_config(self, root_config: "BaseConfig"):
+        with self._with_error_handling("upload config"):
+            self.api.upload_file(
+                path_or_fileobj=root_config.model_dump_json(indent=4).encode("utf-8"),
+                path_in_repo="config.json",
+                repo_id=self.repo_id,
+                repo_type="model",
+                run_as_future=cast(Any, self.config.save_in_background),
+            )
 
-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
+    def _save_code(self):
+        # If a snapshot has been taken (which can be detected using the SNAPSHOT_DIR env),
+        # then upload all contents within the snapshot directory to the repository.
+        if not (snapshot_dir := os.environ.get(SNAPSHOT_DIR)):
+            log.debug("No snapshot directory found. Skipping upload.")
+            return
+
+        with self._with_error_handling("save code"):
+            snapshot_dir = Path(snapshot_dir)
+            if not snapshot_dir.exists() or not snapshot_dir.is_dir():
+                log.warning(
+                    f"Snapshot directory '{snapshot_dir}' does not exist or is not a directory."
+                )
+                return
 
-    # Resolve the checkpoint directory
-    checkpoint_dir = root_config.directory.resolve_subdirectory(
-        root_config.id, "checkpoint"
-    )
+            self.api.upload_folder(
+                folder_path=str(snapshot_dir),
+                repo_id=self.repo_id,
+                repo_type="model",
+                path_in_repo="code",  # Prefix with "code" folder
+                run_as_future=cast(Any, self.config.save_in_background),
+            )
 
-    # Resolve the repository name
-    repo_name = _repo_name(api, root_config)
+    def _save_file(self, p: _Upload):
+        with self._with_error_handling("save file"):
+            # Upload the checkpoint files to the repository
+            self.api.upload_file(
+                path_or_fileobj=p.local_path,
+                path_in_repo=str(p.path_in_repo),
+                repo_id=self.repo_id,
+                repo_type="model",
+                run_as_future=cast(Any, self.config.save_in_background),
+            )
 
-    # 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}'."
-        )
+    def _copy_file(self, source_path_in_repo: Path, dest_path_in_repo: Path):
+        # Create a commit for copying the files
+        from huggingface_hub.hf_api import CommitOperationCopy
 
-    log.info(f"Completed uploading checkpoint files to repository '{repo_name}'.")
+        with self._with_error_handling("copy file"):
+            copy_op = CommitOperationCopy(
+                src_path_in_repo=str(source_path_in_repo),
+                path_in_repo=str(dest_path_in_repo),
+            )
 
+            self.api.create_commit(
+                repo_id=self.repo_id,
+                repo_type="model",
+                commit_message="Copy checkpoint file",
+                operations=[copy_op],
+                run_as_future=cast(Any, self.config.save_in_background),
+            )
 
-class HFHubCallback(NTCallbackBase):
-    def __init__(self, config: HuggingFaceHubConfig):
-        super().__init__()
-        self.config = config
+    def _save_checkpoint(self, path: _Upload, metadata_path: _Upload | None):
+        if not self.config.save_checkpoints:
+            return
 
-    @override
-    def setup(self, trainer, pl_module, stage):
-        from .trainer.trainer import Trainer
+        # If no metadata, just save regularly.
+        if metadata_path is None:
+            self._save_file(path)
+            return
 
-        if not isinstance(trainer, Trainer):
-            raise ValueError(
-                f"HFHubCallback requires a `nshtrainer.Trainer` instance, got {type(trainer)}."
+        # Otherwise, let's check to see if we've already uploaded the metadata.
+        # If so, we can just copy the checkpoint file.
+        from ._checkpoint.metadata import CheckpointMetadata
+
+        metadata = CheckpointMetadata.from_file(metadata_path.local_path)
+        if (
+            existing_ckpt_path := self._checksum_to_path_in_repo.get(
+                metadata.checkpoint_checksum
+            )
+        ) is not None:
+            self._copy_file(existing_ckpt_path, path.path_in_repo)
+        else:
+            # Otherwise, we save the checkpoint & keep the checksum so we don't
+            # re-upload the same file again.
+            self._save_file(path)
+            self._checksum_to_path_in_repo[metadata.checkpoint_checksum] = (
+                path.path_in_repo
             )
 
-        root_config = cast("BaseConfig", pl_module.hparams)
-        _init(trainer=trainer, root_config=root_config)
+        # Save the metadata file
+        # NOTE: This file is fairly small, so we can just upload it directly.
+        # No need to copy.
+        self._save_file(metadata_path)
 
     @override
-    def teardown(self, trainer, pl_module, stage):
-        if hasattr(trainer, "_hf_hub_api"):
-            delattr(trainer, "_hf_hub_api")
+    def state_dict(self):
+        return {
+            "repo_id": self._repo_id,
+            "checksum_to_path_in_repo": {
+                k: str(v) for k, v in self._checksum_to_path_in_repo.items()
+            },
+        }
 
     @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
-                )
+    def load_state_dict(self, state_dict):
+        self._repo_id = state_dict["repo_id"]
+        self._checksum_to_path_in_repo = {
+            k: Path(v) for k, v in state_dict["checksum_to_path_in_repo"].items()
+        }

nshtrainer/callbacks/checkpoint/_base.py

@@ -11,7 +11,6 @@ from typing_extensions import TypeVar, override
 
 from ..._checkpoint.metadata import CheckpointMetadata, _metadata_path
 from ..._checkpoint.saver import _link_checkpoint, _remove_checkpoint
-from ...util.path import find_symlinks
 from ..base import CallbackConfigBase
 
 if TYPE_CHECKING:
@@ -117,47 +116,9 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
                 )
                 continue
 
-            self._remove_checkpoint_with_link_support(
-                trainer, old_ckpt_path, metadata=True
-            )
+            _remove_checkpoint(trainer, old_ckpt_path, metadata=True)
             log.debug(f"Removed old checkpoint: {old_ckpt_path}")
 
-    def _remove_checkpoint_with_link_support(
-        self,
-        trainer: Trainer,
-        path: Path,
-        metadata: bool,
-    ):
-        # Find all the symlinks to the checkpoint
-        ckpt_callbacks: list[CheckpointBase] = [
-            callback
-            for callback in trainer.checkpoint_callbacks
-            if isinstance(callback, CheckpointBase) and callback is not self
-        ]
-        symlink_paths = find_symlinks(
-            path,
-            *[callback.dirpath for callback in ckpt_callbacks],
-            glob_pattern=f"*.{self.extension()}",
-        )
-
-        # If there are no symlinks, just remove the checkpoint
-        if not symlink_paths:
-            _remove_checkpoint(trainer, path, metadata=metadata)
-            return
-
-        log.debug(
-            f"Removing checkpoint with symlinks: {path}, symlinks: {symlink_paths}"
-        )
-
-        # For the first symlink, we can just move the checkpoint file
-        # to the symlink path. For the rest, we need to make new symlinks.
-        new_target = symlink_paths.pop(0)
-        path.rename(new_target)
-        log.debug(f"New symlink target: {new_target}")
-
-        for symlink_path in symlink_paths:
-            _link_checkpoint(new_target, symlink_path, metadata=False)
-
     def current_metrics(self, trainer: Trainer) -> dict[str, Any]:
         current_metrics: dict[str, Any] = {
             "epoch": trainer.current_epoch,
@@ -195,6 +156,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
             filepath,
             self.config.save_weights_only,
             use_checkpoint_cache=None,
+            ckpt_cache_use_symlink=False,
         )
 
         if trainer.is_global_zero:

nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -70,5 +70,5 @@ class BestCheckpoint(CheckpointBase[BestCheckpointCallbackConfig]):
 
     # Events
     @override
-    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule):
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
         self.save_checkpoints(trainer)

nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -39,5 +39,5 @@ class LastCheckpoint(CheckpointBase[LastCheckpointCallbackConfig]):
         return True
 
     @override
-    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
         self.save_checkpoints(trainer)

nshtrainer/trainer/trainer.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import shutil
 from collections.abc import Sequence
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
@@ -416,11 +417,12 @@ class Trainer(LightningTrainer):
         weights_only: bool = False,
         storage_options: Any | None = None,
         use_checkpoint_cache: bool | None = None,
+        ckpt_cache_use_symlink: bool = False,
     ):
         lm = self._base_module
-        hparams = cast(BaseConfig, lm.hparams)
+        root_config = cast(BaseConfig, lm.hparams)
         if use_checkpoint_cache is None:
-            use_checkpoint_cache = hparams.trainer.use_checkpoint_cache
+            use_checkpoint_cache = root_config.trainer.use_checkpoint_cache
 
         filepath = Path(filepath)
 
@@ -440,7 +442,10 @@ class Trainer(LightningTrainer):
             # 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}.")
             if self.is_global_zero:
-                _link_checkpoint(cached_path, filepath, metadata=False)
+                if ckpt_cache_use_symlink:
+                    _link_checkpoint(cached_path, filepath, metadata=False)
+                else:
+                    shutil.copy(cached_path, filepath)
             self.strategy.barrier("Trainer.save_checkpoint")
         else:
             super().save_checkpoint(filepath, weights_only, storage_options)
@@ -454,7 +459,7 @@ class Trainer(LightningTrainer):
 
         # Save the checkpoint metadata
         metadata_path = None
-        if hparams.trainer.save_checkpoint_metadata and self.is_global_zero:
+        if root_config.trainer.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
             if (
                 metadata_path := _write_checkpoint_metadata(self, lm, filepath)

nshtrainer/util/path.py

@@ -1,3 +1,4 @@
+import hashlib
 import os
 from pathlib import Path
 from typing import TypeAlias
@@ -50,3 +51,20 @@ def find_symlinks(
                     pass
 
     return symlinks
+
+
+def compute_file_checksum(file_path: Path) -> str:
+    """
+    Calculate the SHA256 checksum of a file.
+
+    Args:
+        file_path (Path): The path to the file.
+
+    Returns:
+        str: The hexadecimal representation of the file's SHA256 checksum.
+    """
+    sha256_hash = hashlib.sha256()
+    with file_path.open("rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256_hash.update(byte_block)
+    return sha256_hash.hexdigest()

{nshtrainer-0.23.0.dist-info → nshtrainer-0.25.0.dist-info}/METADATA

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

{nshtrainer-0.23.0.dist-info → nshtrainer-0.25.0.dist-info}/RECORD

@@ -1,18 +1,18 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
 nshtrainer/_callback.py,sha256=A1zLsTy4b_wOYnInLLXGSRdHzT2yNa6mPEql-ozm0u0,1013
 nshtrainer/_checkpoint/loader.py,sha256=5vjg-OFChXJjgiOVv8vnV8nwTscfdDtEdxQRz6uPfDE,14158
-nshtrainer/_checkpoint/metadata.py,sha256=E4tfiGzhnn65X95P0Y6K2d_YfPWqvHZoF0FF1-smEJc,5221
+nshtrainer/_checkpoint/metadata.py,sha256=BpxC3VGrgSDHvZpc40A2icjBEDvRNGEsTis9YkFY8Kc,5341
 nshtrainer/_checkpoint/saver.py,sha256=fvRKGI5aeXtsHBOIO4cwGe__wmO-6DiD0-744VASYA4,1500
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
-nshtrainer/_hf_hub.py,sha256=iqhXH54RhSqmot_K3UCVcHTC_TC81_YY7cwvHGHXXlw,16782
+nshtrainer/_hf_hub.py,sha256=0K3uWa8hd2KyGuUYM7OXARcA7vuUiWWGSlP2USysY7o,12066
 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=NpjeKmonJ1Kaz5_39XSn3LlDwvbGjk6WV8BpHSNCvI4,3508
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=g-3zIthupERKqWZQw-A_busQPaPRkto6iHBV-M7nK1Y,527
-nshtrainer/callbacks/checkpoint/_base.py,sha256=zKN-6n61aGze-Hf8MBY1Surh6B-xDwNSApqQJtPcTUs,8048
-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/_base.py,sha256=pLhLj7tvfY5czGY_vT0xRfWHzGJYC4iOBRLokFVq0mE,6733
+nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=8BHgLAd3Tuzf5sup0guEAKF1jJiAwYsjdKBFYZw98ac,2171
+nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CWWv0cSwQ1VAX26N7hAyMxbNCk26Keh39oQguBEK5To,1102
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
 nshtrainer/callbacks/early_stopping.py,sha256=VWuJz0oN87b6SwBeVc32YNpeJr1wts8K45k8JJJmG9I,4617
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
@@ -78,15 +78,15 @@ nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=r0ir4xYSdf_jebM0x09qaO6nJsvsiRQDyM0fs80ppOQ,2347
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
-nshtrainer/trainer/trainer.py,sha256=Leh3ADxoYsRWlJFIW20netohLcKx0XxUrRhD9LM4jws,19201
+nshtrainer/trainer/trainer.py,sha256=Zwdcqfmrr7yuonsp4VrNOget8wkaZY9lf-_yeJ94lkk,19397
 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=WbPWXpu5LIDocQihQC3-72qxN1sa6-d1kPOmKDR-NC8,1520
+nshtrainer/util/path.py,sha256=RUkIOrlj9b8zPPXE3JLhdihBNitJSUWw1whZ33u-2Yk,2005
 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.23.0.dist-info/METADATA,sha256=wkbqsz6A4d0h1u-8CCZwfYYmqLm7YjirdnS-fTA-mkI,935
-nshtrainer-0.23.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.23.0.dist-info/RECORD,,
+nshtrainer-0.25.0.dist-info/METADATA,sha256=Rqdeh2yp2AhZ_nOHlD47v5YPDrLc2fHN6WGwqJnDv04,935
+nshtrainer-0.25.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.25.0.dist-info/RECORD,,