nshtrainer 0.26.0__py3-none-any.whl → 0.26.2__py3-none-any.whl

nshtrainer/_checkpoint/metadata.py

@@ -1,7 +1,6 @@
 import copy
 import datetime
 import logging
-import shutil
 from collections.abc import Callable
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, cast
@@ -11,7 +10,7 @@ import numpy as np
 import torch
 
 from ..util._environment_info import EnvironmentConfig
-from ..util.path import compute_file_checksum, get_relative_path
+from ..util.path import compute_file_checksum, try_symlink_or_copy
 
 if TYPE_CHECKING:
     from ..model import BaseConfig, LightningModuleBase
@@ -142,21 +141,7 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
     # Link the metadata files to the new checkpoint
     path = _metadata_path(checkpoint_path)
     linked_path = _metadata_path(linked_checkpoint_path)
-    try:
-        try:
-            # linked_path.symlink_to(path)
-            # We should store the path as a relative path
-            # to the metadata file to avoid issues with
-            # moving the checkpoint directory
-            linked_path.symlink_to(get_relative_path(linked_path, path))
-        except OSError:
-            # on Windows, special permissions are required to create symbolic links as a regular user
-            # fall back to copying the file
-            shutil.copy(path, linked_path)
-    except Exception:
-        log.exception(f"Failed to link {path} to {linked_path}")
-    else:
-        log.debug(f"Linked {path} to {linked_path}")
+    try_symlink_or_copy(path, linked_path)
 
 
 def _sort_ckpts_by_metadata(

nshtrainer/_checkpoint/saver.py

@@ -5,7 +5,7 @@ from pathlib import Path
 
 from lightning.pytorch import Trainer
 
-from ..util.path import get_relative_path
+from ..util.path import try_symlink_or_copy
 from .metadata import _link_checkpoint_metadata, _remove_checkpoint_metadata
 
 log = logging.getLogger(__name__)
@@ -34,13 +34,7 @@ def _link_checkpoint(
         if metadata:
             _remove_checkpoint_metadata(linkpath)
 
-    try:
-        linkpath.symlink_to(get_relative_path(linkpath, filepath))
-    except OSError:
-        # on Windows, special permissions are required to create symbolic links as a regular user
-        # fall back to copying the file
-        shutil.copy(filepath, linkpath)
-
+    try_symlink_or_copy(filepath, linkpath)
     if metadata:
         _link_checkpoint_metadata(filepath, linkpath)
 

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/util/path.py

@@ -1,8 +1,13 @@
 import hashlib
+import logging
 import os
+import platform
+import shutil
 from pathlib import Path
 from typing import TypeAlias
 
+log = logging.getLogger(__name__)
+
 _Path: TypeAlias = str | Path | os.PathLike
 
 
@@ -68,3 +73,32 @@ def compute_file_checksum(file_path: Path) -> str:
         for byte_block in iter(lambda: f.read(4096), b""):
             sha256_hash.update(byte_block)
     return sha256_hash.hexdigest()
+
+
+def try_symlink_or_copy(
+    file_path: Path,
+    link_path: Path,
+    target_is_directory: bool = False,
+    relative: bool = True,
+):
+    """
+    Symlinks on Unix, copies on Windows.
+    """
+
+    symlink_target = get_relative_path(link_path, file_path) if relative else file_path
+    try:
+        if platform.system() == "Windows":
+            if target_is_directory:
+                shutil.copytree(file_path, link_path)
+            else:
+                shutil.copy(file_path, link_path)
+        else:
+            link_path.symlink_to(
+                symlink_target, target_is_directory=target_is_directory
+            )
+    except Exception:
+        log.exception(f"Failed to create symlink or copy {file_path} to {link_path}")
+        return False
+    else:
+        log.debug(f"Created symlink or copied {file_path} to {link_path}")
+        return True

{nshtrainer-0.26.0.dist-info → nshtrainer-0.26.2.dist-info}/METADATA

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

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

@@ -1,10 +1,10 @@
 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=BpxC3VGrgSDHvZpc40A2icjBEDvRNGEsTis9YkFY8Kc,5341
-nshtrainer/_checkpoint/saver.py,sha256=fvRKGI5aeXtsHBOIO4cwGe__wmO-6DiD0-744VASYA4,1500
+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
@@ -82,11 +82,11 @@ nshtrainer/trainer/trainer.py,sha256=Zwdcqfmrr7yuonsp4VrNOget8wkaZY9lf-_yeJ94lkk
 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=RUkIOrlj9b8zPPXE3JLhdihBNitJSUWw1whZ33u-2Yk,2005
+nshtrainer/util/path.py,sha256=jAEjF1qp8Aii32L5lWG4UFgVyQAFkHOMYEc_TC2hDx8,2947
 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.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,,
+nshtrainer-0.26.2.dist-info/METADATA,sha256=d3vWjdB9FT6fbWJPkyBI-4M18ekg2WcJmJiKasExchM,916
+nshtrainer-0.26.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.26.2.dist-info/RECORD,,