nshtrainer 0.22.1__py3-none-any.whl → 0.24.0__py3-none-any.whl

nshtrainer/_checkpoint/saver.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import shutil
 from pathlib import Path
@@ -7,6 +8,8 @@ from lightning.pytorch import Trainer
 from ..util.path import get_relative_path
 from .metadata import _link_checkpoint_metadata, _remove_checkpoint_metadata
 
+log = logging.getLogger(__name__)
+
 
 def _link_checkpoint(
     filepath: str | Path | os.PathLike,
@@ -19,11 +22,14 @@ def _link_checkpoint(
     linkpath = Path(linkpath)
 
     if remove_existing:
-        if linkpath.exists():
-            if linkpath.is_symlink() or linkpath.is_file():
-                linkpath.unlink()
-            elif linkpath.is_dir():
-                shutil.rmtree(linkpath)
+        try:
+            if linkpath.exists():
+                if linkpath.is_dir():
+                    shutil.rmtree(linkpath, ignore_errors=True)
+                else:
+                    linkpath.unlink(missing_ok=True)
+        except Exception:
+            log.exception(f"Failed to remove {linkpath}")
 
         if metadata:
             _remove_checkpoint_metadata(linkpath)

nshtrainer/_hf_hub.py

@@ -2,6 +2,7 @@ import io
 import logging
 import os
 import re
+from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
@@ -110,6 +111,7 @@ def _api(token: str | None = None):
 
 def _enabled_and_valid(
     trainer: "Trainer",
+    callback: "HFHubCallback",
     config: HuggingFaceHubConfig,
     *,
     rank_zero_only: bool,
@@ -120,22 +122,9 @@ def _enabled_and_valid(
 
     # 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)
+    return callback._hf_hub_api
 
 
 def _repo_name(api: "HfApi", root_config: "BaseConfig"):
@@ -173,11 +162,12 @@ def _repo_name(api: "HfApi", root_config: "BaseConfig"):
     return f"{username}/{repo_name}"
 
 
-def _init(*, trainer: "Trainer", root_config: "BaseConfig"):
+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,
         )
@@ -210,10 +200,10 @@ def _init(*, trainer: "Trainer", root_config: "BaseConfig"):
         log.exception(f"Error checking repository '{repo_name}'")
 
     # Upload the config
-    _save_config(root_config, trainer=trainer)
+    _save_config(root_config, trainer=trainer, callback=callback)
 
     # Upload the code
-    _save_code(repo_name, config=config, trainer=trainer)
+    _save_code(repo_name, config=config, trainer=trainer, callback=callback)
 
 
 def _save_code(
@@ -221,10 +211,12 @@ def _save_code(
     *,
     config: HuggingFaceHubConfig,
     trainer: "Trainer",
+    callback: "HFHubCallback",
 ):
     if (
         api := _enabled_and_valid(
             trainer,
+            callback,
             config,
             rank_zero_only=True,
         )
@@ -266,11 +258,13 @@ 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,
         )
@@ -298,15 +292,45 @@ def _save_config(
         log.exception(f"Failed to upload config.json to repository '{repo_name}'.")
 
 
+def _is_link(p: Path, trainer: "Trainer"):
+    if p.is_symlink():
+        return True
+
+    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)
+
+    return False
+
+
+def _resolve_link(p: Path, trainer: "Trainer"):
+    if p.is_symlink():
+        return p.resolve()
+
+    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, config, rank_zero_only=True)
+        api := _enabled_and_valid(trainer, callback, config, rank_zero_only=True)
     ) is None or not config.save_checkpoints:
         return
 
@@ -323,7 +347,7 @@ 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 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:
@@ -374,13 +398,14 @@ def _save_checkpoint_files(
 
 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, config, rank_zero_only=True)
+        api := _enabled_and_valid(trainer, callback, config, rank_zero_only=True)
     ) is None or not config.save_checkpoints:
         return
 
@@ -397,7 +422,7 @@ def _save_checkpoint_symlinks(
 
     commits: list[CommitOperation] = []
     for p in paths:
-        assert p.is_symlink(), f"Path {p} is not a symlink."
+        assert _is_link(p, trainer=trainer), f"Path {p} is not a symlink."
 
         try:
             dest_relative_path = p.relative_to(checkpoint_dir)
@@ -407,11 +432,15 @@ def _save_checkpoint_symlinks(
             )
             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.resolve().relative_to(checkpoint_dir)
+            source_relative_path = p_resolved.relative_to(checkpoint_dir)
         except ValueError:
             log.warning(
-                f"Checkpoint symlink target {p.resolve()} is not within the checkpoint directory {checkpoint_dir}."
+                f"Checkpoint symlink target {p_resolved} is not within the checkpoint directory {checkpoint_dir}."
             )
             continue
 
@@ -447,39 +476,6 @@ def _save_checkpoint_symlinks(
     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__()
@@ -495,12 +491,7 @@ class HFHubCallback(NTCallbackBase):
             )
 
         root_config = cast("BaseConfig", pl_module.hparams)
-        _init(trainer=trainer, root_config=root_config)
-
-    @override
-    def teardown(self, trainer, pl_module, stage):
-        if hasattr(trainer, "_hf_hub_api"):
-            delattr(trainer, "_hf_hub_api")
+        _init(trainer=trainer, callback=self, root_config=root_config)
 
     @override
     def on_checkpoint_saved(self, ckpt_path, metadata_path, trainer, pl_module):
@@ -510,9 +501,18 @@ class HFHubCallback(NTCallbackBase):
         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()]:
+            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, symlink_paths, root_config=root_config
+                    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)

nshtrainer/callbacks/checkpoint/_base.py

@@ -156,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/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
@@ -279,6 +280,13 @@ 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,
@@ -287,6 +295,7 @@ 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)
 
@@ -416,11 +425,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 +450,14 @@ 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._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)
 
@@ -453,7 +470,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

@@ -27,3 +27,26 @@ def get_relative_path(source: _Path, destination: _Path):
     down = os.sep.join(destination_parts[i:])
 
     return Path(os.path.normpath(os.path.join(up, down)))
+
+
+def find_symlinks(
+    target_file: _Path,
+    *search_directories: _Path,
+    glob_pattern: str = "*",
+):
+    target_file = Path(target_file).resolve()
+    symlinks: list[Path] = []
+
+    for search_directory in search_directories:
+        search_directory = Path(search_directory)
+        for path in search_directory.rglob(glob_pattern):
+            if path.is_symlink():
+                try:
+                    link_target = path.resolve()
+                    if link_target.samefile(target_file):
+                        symlinks.append(path)
+                except FileNotFoundError:
+                    # Handle broken symlinks
+                    pass
+
+    return symlinks

{nshtrainer-0.22.1.dist-info → nshtrainer-0.24.0.dist-info}/METADATA

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

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

@@ -2,15 +2,15 @@ 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/saver.py,sha256=6W-Rbc3QGuhcF_mcwN8v31uEjLQCsZvt8CPuqPs4m5g,1342
+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=Ac4y7jmuAMEQOJPJgoYmiaIGlZvgyUcqpipb6fPuHSE,16587
 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=r6IPpl3sGUmxBNv80y9r326lTrPAIVSU3Fu-3LrYH2s,6691
+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/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
@@ -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=KXsvAhgVgYjmYfoqzH_qoQXqd6nVx7-vs9ObQJpwbIk,19140
+nshtrainer/trainer/trainer.py,sha256=vswxAhyLqTL99kRJvU4Q3uEyQT80eM3mN74yMhsyn_I,19905
 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/path.py,sha256=WbPWXpu5LIDocQihQC3-72qxN1sa6-d1kPOmKDR-NC8,1520
 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.22.1.dist-info/METADATA,sha256=XF3QXKbeAN7I5vYHNbjExlV_6CF8QgPqPYFsCxs52rA,935
-nshtrainer-0.22.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.22.1.dist-info/RECORD,,
+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,,