nshtrainer 0.26.1__py3-none-any.whl → 0.27.0__py3-none-any.whl

nshtrainer/_hf_hub.py

@@ -188,27 +188,19 @@ class HFHubCallback(NTCallbackBase):
 
         self.config = config
 
-        self._repo_id = None
+        self._repo_id: str | None = None
         self._checksum_to_path_in_repo: dict[str, Path] = {}
 
     @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)
+        self._repo_id = _repo_name(self.api, root_config)
+
+        if not self.config or not trainer.is_global_zero:
+            return
 
         # 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,
-        )
+        self._create_repo_if_not_exists()
 
         # Upload the config and code
         self._save_config(root_config)
@@ -216,17 +208,22 @@ class HFHubCallback(NTCallbackBase):
 
     @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 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,
-                )
+        if (
+            not self.config
+            or not self.config.save_checkpoints
+            or not trainer.is_global_zero
+        ):
+            return
+
+        with self._with_error_handling("save checkpoints"):
+            root_config = cast("BaseConfig", pl_module.hparams)
+            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):
@@ -241,6 +238,33 @@ class HFHubCallback(NTCallbackBase):
             raise ValueError("Repository id has not been initialized.")
         return self._repo_id
 
+    def _create_repo_if_not_exists(self):
+        if not self.config or not self.config.auto_create:
+            return
+
+        # Create the repository, if it doesn't exist
+        with self._with_error_handling("create repository"):
+            from huggingface_hub.utils import RepositoryNotFoundError
+
+            try:
+                # Check if the repository exists
+                self.api.repo_info(repo_id=self.repo_id, repo_type="model")
+                log.info(f"Repository '{self.repo_id}' already exists.")
+            except RepositoryNotFoundError:
+                # Repository doesn't exist, so create it
+                try:
+                    self.api.create_repo(
+                        repo_id=self.repo_id,
+                        repo_type="model",
+                        private=self.config.auto_create.private,
+                        exist_ok=True,
+                    )
+                    log.info(f"Created new repository '{self.repo_id}'.")
+                except Exception:
+                    log.exception(f"Failed to create repository '{self.repo_id}'")
+            except Exception:
+                log.exception(f"Error checking repository '{self.repo_id}'")
+
     def _save_config(self, root_config: "BaseConfig"):
         with self._with_error_handling("upload config"):
             self.api.upload_file(

nshtrainer/callbacks/checkpoint/_base.py

@@ -153,10 +153,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         # Save the new checkpoint
         filepath = self.resolve_checkpoint_path(self.current_metrics(trainer))
         trainer._nshtrainer_save_checkpoint(
-            filepath,
-            self.config.save_weights_only,
-            use_checkpoint_cache=None,
-            ckpt_cache_use_symlink=False,
+            filepath, self.config.save_weights_only, use_checkpoint_cache=None
         )
 
         if trainer.is_global_zero:

nshtrainer/model/config.py

@@ -1012,7 +1012,7 @@ 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
+    use_checkpoint_cache: bool = False
     """If enabled, will optimize the saving of duplicate checkpoints by creating symlinks instead of copying the file."""
 
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()

nshtrainer/trainer/trainer.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import shutil
+from collections import defaultdict
 from collections.abc import Sequence
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
@@ -18,7 +19,6 @@ 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,
@@ -280,6 +280,12 @@ class Trainer(LightningTrainer):
     if TYPE_CHECKING:
         callbacks: list[Callback]
 
+    def _nshtrainer_checkpoint_cache_get(self, key: tuple[int, int]):
+        return next(
+            (ckpt for ckpt in self._nshtrainer_checkpoint_cache[key] if ckpt.exists()),
+            None,
+        )
+
     @override
     def __init__(
         self,
@@ -287,7 +293,9 @@ class Trainer(LightningTrainer):
         /,
         **kwargs: Unpack[LightningTrainerKwargs],
     ):
-        self._nshtrainer_checkpoint_cache: dict[tuple[int, int], Path] = {}
+        self._nshtrainer_checkpoint_cache = defaultdict[tuple[int, int], list[Path]](
+            lambda: []
+        )
 
         self._pre_init(config)
 
@@ -417,7 +425,6 @@ 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
         root_config = cast(BaseConfig, lm.hparams)
@@ -433,7 +440,7 @@ class Trainer(LightningTrainer):
         if (
             use_checkpoint_cache
             and (
-                cached_path := self._nshtrainer_checkpoint_cache.get(
+                cached_path := self._nshtrainer_checkpoint_cache_get(
                     (self.current_epoch, self.global_step)
                 )
             )
@@ -442,10 +449,7 @@ 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:
-                if ckpt_cache_use_symlink:
-                    _link_checkpoint(cached_path, filepath, metadata=False)
-                else:
-                    shutil.copy(cached_path, filepath)
+                shutil.copy(cached_path, filepath)
             self.strategy.barrier("Trainer.save_checkpoint")
         else:
             super().save_checkpoint(filepath, weights_only, storage_options)
@@ -454,7 +458,7 @@ class Trainer(LightningTrainer):
         if use_checkpoint_cache and cached_path is None:
             self._nshtrainer_checkpoint_cache[
                 (self.current_epoch, self.global_step)
-            ] = filepath
+            ].append(filepath)
             log.debug(f"Checkpoint saved to cache: {filepath}")
 
         # Save the checkpoint metadata

{nshtrainer-0.26.1.dist-info → nshtrainer-0.27.0.dist-info}/METADATA

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

{nshtrainer-0.26.1.dist-info → nshtrainer-0.27.0.dist-info}/RECORD

@@ -4,13 +4,13 @@ nshtrainer/_checkpoint/loader.py,sha256=5vjg-OFChXJjgiOVv8vnV8nwTscfdDtEdxQRz6uP
 nshtrainer/_checkpoint/metadata.py,sha256=hxZwwsUKVbBtt4wjqcKZbObx0PuO-qCdF3BTdnyqaQo,4711
 nshtrainer/_checkpoint/saver.py,sha256=1loCDYDy_Cay37uKs_wvxnkwvr41WMmga85qefct80Q,1271
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
-nshtrainer/_hf_hub.py,sha256=0K3uWa8hd2KyGuUYM7OXARcA7vuUiWWGSlP2USysY7o,12066
+nshtrainer/_hf_hub.py,sha256=v1DV3Vn6Pdwr2KYI7yL_Xv1dyJLaG7yZhKFWZFy3QFk,13087
 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/_base.py,sha256=iDtQ8urdSfl0uUJMvsSvvYo-fsL3ZVXCv79FGS67Jx4,6666
 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
@@ -58,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=22_xIcdEO2pJzXgrFaqGFtk3PQEiwKiMZY1cjhoyWaA,43486
+nshtrainer/model/config.py,sha256=9fPNrMkpOUIWwAM1UTQQRTVnuV7m3CDTtivya2GCyQY,43487
 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
@@ -78,7 +78,7 @@ 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=Zwdcqfmrr7yuonsp4VrNOget8wkaZY9lf-_yeJ94lkk,19397
+nshtrainer/trainer/trainer.py,sha256=Bm-kHOF_9W905SpOQcfEwHSUhF6KW8aAFdWNyIotyeM,19450
 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
@@ -87,6 +87,6 @@ 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.26.1.dist-info/METADATA,sha256=tMMpyg1BTKec5d69ziW6XBxDXaI0gSK5tDMPCmj7VCQ,916
-nshtrainer-0.26.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.26.1.dist-info/RECORD,,
+nshtrainer-0.27.0.dist-info/METADATA,sha256=fn2L1iuf7uaxT8UYzZd1lhKSKGEldmA0_026HG5f6ek,916
+nshtrainer-0.27.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.27.0.dist-info/RECORD,,