nshtrainer 0.24.0__py3-none-any.whl → 0.26.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,8 @@
-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
@@ -17,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__)
@@ -109,24 +109,6 @@ def _api(token: str | None = None):
     return api
 
 
-def _enabled_and_valid(
-    trainer: "Trainer",
-    callback: "HFHubCallback",
-    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 None
-
-    return callback._hf_hub_api
-
-
 def _repo_name(api: "HfApi", root_config: "BaseConfig"):
     username = None
     if (ac := root_config.trainer.hf_hub.auto_create) and ac.namespace:
@@ -162,357 +144,210 @@ def _repo_name(api: "HfApi", root_config: "BaseConfig"):
     return f"{username}/{repo_name}"
 
 
-def _init(*, trainer: "Trainer", callback: "HFHubCallback", root_config: "BaseConfig"):
-    config = root_config.trainer.hf_hub
-    if (
-        api := _enabled_and_valid(
-            trainer,
-            callback,
-            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.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,
+            relative_path = local_path.relative_to(checkpoint_dir)
+        except ValueError:
+            raise ValueError(
+                f"Checkpoint path {local_path} is not within the checkpoint directory {checkpoint_dir}."
             )
-            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, callback=callback)
-
-    # Upload the code
-    _save_code(repo_name, config=config, trainer=trainer, callback=callback)
-
-
-def _save_code(
-    repo_name: str,
-    *,
-    config: HuggingFaceHubConfig,
-    trainer: "Trainer",
-    callback: "HFHubCallback",
-):
-    if (
-        api := _enabled_and_valid(
-            trainer,
-            callback,
-            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."
-        )
+        # Prefix the path in repo with "checkpoints"
+        path_in_repo = Path("checkpoints") / relative_path
 
+        return cls(local_path=local_path, path_in_repo=path_in_repo)
 
-def _save_config(
-    root_config: "BaseConfig",
-    *,
-    trainer: "Trainer",
-    callback: "HFHubCallback",
-):
-    config = root_config.trainer.hf_hub
-    if (
-        api := _enabled_and_valid(
-            trainer,
-            callback,
-            config,
-            rank_zero_only=True,
-        )
-    ) 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)
+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__()
 
-    # 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}'.")
+        self.config = config
 
+        self._repo_id = None
+        self._checksum_to_path_in_repo: dict[str, Path] = {}
 
-def _is_link(p: Path, trainer: "Trainer"):
-    if p.is_symlink():
-        return True
+    @override
+    def setup(self, trainer, pl_module, stage):
+        from .trainer.trainer import Trainer
 
-    try:
-        link_path = trainer._nshtrainer_ckpt_link(p)
-        return link_path in trainer._nshtrainer_checkpoint_link_dict
-    except Exception:
-        log.info(f"Failed to check if path {p} is a symlink.", exc_info=True)
+        if not isinstance(trainer, Trainer):
+            raise ValueError(
+                f"HFHubCallback requires a `nshtrainer.Trainer` instance, got {type(trainer)}."
+            )
 
-    return False
+        root_config = cast("BaseConfig", pl_module.hparams)
 
+        # 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,
+        )
 
-def _resolve_link(p: Path, trainer: "Trainer"):
-    if p.is_symlink():
-        return p.resolve()
+        # Upload the config and code
+        self._save_config(root_config)
+        self._save_code()
 
-    try:
-        link_path = trainer._nshtrainer_ckpt_link(p)
-        if (
-            resolved := trainer._nshtrainer_checkpoint_link_dict.get(link_path)
-        ) is not None:
-            return resolved
-    except Exception:
-        log.info(f"Failed to resolve symlink for path {p}.", exc_info=True)
-
-    return None
-
-
-def _save_checkpoint_files(
-    trainer: "Trainer",
-    callback: "HFHubCallback",
-    paths: list[Path],
-    *,
-    root_config: "BaseConfig",
-):
-    config = root_config.trainer.hf_hub
-    if (
-        api := _enabled_and_valid(trainer, callback, 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 _is_link(p, trainer=trainer), 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)
-        except ValueError:
-            log.warning(
-                f"Checkpoint path {p} is not within the checkpoint directory {checkpoint_dir}."
-            )
-            continue
+    @override
+    def on_checkpoint_saved(self, ckpt_path, metadata_path, trainer, pl_module):
+        root_config = cast("BaseConfig", pl_module.hparams)
 
-        # Prefix the path in repo with "checkpoints"
-        path_in_repo = Path("checkpoints") / relative_path
+        # 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,
+                )
 
-        # 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,
+    @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, config.save_in_background),
+                run_as_future=cast(Any, self.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}'.")
-
-
-def _save_checkpoint_symlinks(
-    trainer: "Trainer",
-    callback: "HFHubCallback",
-    paths: list[Path],
-    *,
-    root_config: "BaseConfig",
-):
-    config = root_config.trainer.hf_hub
-    if (
-        api := _enabled_and_valid(trainer, callback, 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 _is_link(p, trainer=trainer), f"Path {p} is not a symlink."
+    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
 
-        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}."
+            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),
             )
-            continue
 
-        if (p_resolved := _resolve_link(p, trainer=trainer)) is None:
-            log.warning(f"Failed to resolve symlink for path {p}.")
-            continue
-
-        try:
-            source_relative_path = p_resolved.relative_to(checkpoint_dir)
-        except ValueError:
-            log.warning(
-                f"Checkpoint symlink target {p_resolved} is not within the checkpoint directory {checkpoint_dir}."
+    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),
             )
-            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}")
+    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
 
-    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}'"
-        )
+        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),
+            )
 
-    log.info(f"Completed copying checkpoint symlinks to repository '{repo_name}'.")
+            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),
+            )
 
+    def _save_checkpoint(self, path: _Upload, metadata_path: _Upload | None):
+        if not self.config.save_checkpoints:
+            return
 
-class HFHubCallback(NTCallbackBase):
-    def __init__(self, config: HuggingFaceHubConfig):
-        super().__init__()
-        self.config = config
+        # If no metadata, just save regularly.
+        if metadata_path is None:
+            self._save_file(path)
+            return
 
-    @override
-    def setup(self, trainer, pl_module, stage):
-        from .trainer.trainer import 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
 
-        if not isinstance(trainer, Trainer):
-            raise ValueError(
-                f"HFHubCallback requires a `nshtrainer.Trainer` instance, got {type(trainer)}."
+        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, callback=self, 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 on_checkpoint_saved(self, ckpt_path, metadata_path, trainer, pl_module):
-        root_config = cast("BaseConfig", pl_module.hparams)
+    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()
+            },
+        }
 
-        # 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 _is_link(p, trainer=trainer)
-            ]:
-                _save_checkpoint_files(
-                    trainer, self, regular_paths, root_config=root_config
-                )
-            if symlink_paths := [p for p in all_paths if _is_link(p, trainer=trainer)]:
-                _save_checkpoint_symlinks(
-                    trainer, self, symlink_paths, root_config=root_config
-                )
-
-    @cached_property
-    def _hf_hub_api(self):
-        # Create and authenticate the API instance
-        return _api(self.config.token)
+    @override
+    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/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/callbacks/gradient_skipping.py

@@ -1,8 +1,8 @@
-import importlib.util
 import logging
 from typing import Any, Literal, Protocol, runtime_checkable
 
 import torch
+import torchmetrics
 from lightning.pytorch import Callback, LightningModule, Trainer
 from torch.optim import Optimizer
 from typing_extensions import override
@@ -20,19 +20,12 @@ class HasGradSkippedSteps(Protocol):
 
 class GradientSkipping(Callback):
     def __init__(self, config: "GradientSkippingConfig"):
-        if importlib.util.find_spec("torchmetrics") is not None:
-            raise ImportError(
-                "To use the GradientSkipping callback, please install torchmetrics: pip install torchmetrics"
-            )
-
         super().__init__()
         self.config = config
 
     @override
     def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
         if not isinstance(pl_module, HasGradSkippedSteps):
-            import torchmetrics  # type: ignore
-
             pl_module.grad_skipped_steps = torchmetrics.SumMetric()
 
     @override

nshtrainer/trainer/trainer.py

@@ -280,13 +280,6 @@ class Trainer(LightningTrainer):
     if TYPE_CHECKING:
         callbacks: list[Callback]
 
-    def _nshtrainer_ckpt_link(self, ckpt_path: Path):
-        root_config = cast(BaseConfig, self._base_module.hparams)
-        ckpt_dir = root_config.directory.resolve_subdirectory(
-            root_config.id, "checkpoint"
-        )
-        return str(ckpt_path.absolute().relative_to(ckpt_dir))
-
     @override
     def __init__(
         self,
@@ -295,7 +288,6 @@ class Trainer(LightningTrainer):
         **kwargs: Unpack[LightningTrainerKwargs],
     ):
         self._nshtrainer_checkpoint_cache: dict[tuple[int, int], Path] = {}
-        self._nshtrainer_checkpoint_link_dict = dict[str, Path]()
 
         self._pre_init(config)
 
@@ -454,9 +446,6 @@ class Trainer(LightningTrainer):
                     _link_checkpoint(cached_path, filepath, metadata=False)
                 else:
                     shutil.copy(cached_path, filepath)
-                self._nshtrainer_checkpoint_link_dict[
-                    self._nshtrainer_ckpt_link(filepath)
-                ] = cached_path
             self.strategy.barrier("Trainer.save_checkpoint")
         else:
             super().save_checkpoint(filepath, weights_only, storage_options)

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.24.0.dist-info → nshtrainer-0.26.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.24.0
+Version: 0.26.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -22,7 +22,7 @@ Requires-Dist: psutil
 Requires-Dist: pytorch-lightning
 Requires-Dist: tensorboard ; extra == "extra"
 Requires-Dist: torch
-Requires-Dist: torchmetrics ; extra == "extra"
+Requires-Dist: torchmetrics
 Requires-Dist: typing-extensions
 Requires-Dist: wandb ; extra == "extra"
 Requires-Dist: wrapt ; extra == "extra"

{nshtrainer-0.24.0.dist-info → nshtrainer-0.26.0.dist-info}/RECORD

@@ -1,23 +1,23 @@
 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=Ac4y7jmuAMEQOJPJgoYmiaIGlZvgyUcqpipb6fPuHSE,16587
+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=pLhLj7tvfY5czGY_vT0xRfWHzGJYC4iOBRLokFVq0mE,6733
-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/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
 nshtrainer/callbacks/finite_checks.py,sha256=gJC_RUr3ais3FJI0uB6wUZnDdE3WRwCix3ppA3PwQXA,2077
-nshtrainer/callbacks/gradient_skipping.py,sha256=pqu5AELx4ctJxR2Y7YSSiGd5oGauVCTZFCEIIS6s88w,3665
+nshtrainer/callbacks/gradient_skipping.py,sha256=EBNkANDnD3BTszWjnG-jwY8FEj-iRqhE3e1x5LQF6M8,3393
 nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvLfaw,8080
 nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
 nshtrainer/callbacks/norm_logging.py,sha256=T2psu8mYsw9iahPKT6aUPjkGrZ4TIzm6_UUUmE09GJs,6274
@@ -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=vswxAhyLqTL99kRJvU4Q3uEyQT80eM3mN74yMhsyn_I,19905
+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.24.0.dist-info/METADATA,sha256=sYEAgF1cJ2a3Mjp137ienL4cvHswQHH76lm2GTj-Nb8,935
-nshtrainer-0.24.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.24.0.dist-info/RECORD,,
+nshtrainer-0.26.0.dist-info/METADATA,sha256=YBlbpalQ3BX8UBF_5SHk_F7v9Nq3JMsqVf6MoqH8KzU,916
+nshtrainer-0.26.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.26.0.dist-info/RECORD,,