nshtrainer 0.21.0__tar.gz → 0.22.1__tar.gz

{nshtrainer-0.21.0 → nshtrainer-0.22.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.21.0
+Version: 0.22.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -26,7 +26,6 @@ Requires-Dist: torchmetrics ; extra == "extra"
 Requires-Dist: typing-extensions
 Requires-Dist: wandb ; extra == "extra"
 Requires-Dist: wrapt ; extra == "extra"
-Requires-Dist: zstandard ; extra == "extra"
 Description-Content-Type: text/markdown
 
 

{nshtrainer-0.21.0 → nshtrainer-0.22.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.21.0"
+version = "0.22.1"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -23,7 +23,6 @@ GitPython = { version = "*", optional = true }
 wandb = { version = "*", optional = true }
 tensorboard = { version = "*", optional = true }
 huggingface-hub = { version = "*", optional = true }
-zstandard = { version = "*", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 pyright = "^1.1.372"
@@ -54,5 +53,4 @@ extra = [
     "wandb",
     "tensorboard",
     "huggingface-hub",
-    "zstandard",
 ]

{nshtrainer-0.21.0 → nshtrainer-0.22.1}/src/nshtrainer/_checkpoint/metadata.py

@@ -11,6 +11,7 @@ import numpy as np
 import torch
 
 from ..util._environment_info import EnvironmentConfig
+from ..util.path import get_relative_path
 
 if TYPE_CHECKING:
     from ..model import BaseConfig, LightningModuleBase
@@ -145,7 +146,7 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
             # 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(path.relative_to(linked_path.parent))
+            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

{nshtrainer-0.21.0 → nshtrainer-0.22.1}/src/nshtrainer/_hf_hub.py

@@ -10,11 +10,7 @@ from nshrunner._env import SNAPSHOT_DIR
 from typing_extensions import override
 
 from ._callback import NTCallbackBase
-from .callbacks.base import (
-    CallbackConfigBase,
-    CallbackMetadataConfig,
-    CallbackWithMetadata,
-)
+from .callbacks.base import CallbackConfigBase
 
 if TYPE_CHECKING:
     from huggingface_hub import HfApi  # noqa: F401
@@ -81,10 +77,7 @@ class HuggingFaceHubConfig(CallbackConfigBase):
 
     @override
     def create_callbacks(self, root_config):
-        yield CallbackWithMetadata(
-            HFHubCallback(self),
-            CallbackMetadataConfig(ignore_if_exists=True),
-        )
+        yield self.with_metadata(HFHubCallback(self), ignore_if_exists=True)
 
 
 def _api(token: str | None = None):

{nshtrainer-0.21.0 → nshtrainer-0.22.1}/src/nshtrainer/callbacks/base.py

@@ -2,29 +2,24 @@ from abc import ABC, abstractmethod
 from collections import Counter
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, TypeAlias, TypedDict
+from typing import TYPE_CHECKING, ClassVar, TypeAlias
 
 import nshconfig as C
 from lightning.pytorch import Callback
+from typing_extensions import TypedDict, Unpack
 
 if TYPE_CHECKING:
     from ..model.config import BaseConfig
 
 
-class CallbackMetadataDict(TypedDict, total=False):
+class CallbackMetadataConfig(TypedDict, total=False):
     ignore_if_exists: bool
-    """If `True`, the callback will not be added if another callback with the same class already exists."""
+    """If `True`, the callback will not be added if another callback with the same class already exists.
+    Default is `False`."""
 
     priority: int
-    """Priority of the callback. Callbacks with higher priority will be loaded first."""
-
-
-class CallbackMetadataConfig(C.Config):
-    ignore_if_exists: bool = False
-    """If `True`, the callback will not be added if another callback with the same class already exists."""
-
-    priority: int = 0
-    """Priority of the callback. Callbacks with higher priority will be loaded first."""
+    """Priority of the callback. Callbacks with higher priority will be loaded first.
+    Default is `0`."""
 
 
 @dataclass(frozen=True)
@@ -37,13 +32,18 @@ ConstructedCallback: TypeAlias = Callback | CallbackWithMetadata
 
 
 class CallbackConfigBase(C.Config, ABC):
-    metadata: CallbackMetadataConfig = CallbackMetadataConfig()
+    metadata: ClassVar[CallbackMetadataConfig] = CallbackMetadataConfig()
     """Metadata for the callback."""
 
-    def with_metadata(self, callback: Callback, **metadata: CallbackMetadataDict):
-        return CallbackWithMetadata(
-            callback=callback, metadata=self.metadata.model_copy(update=metadata)
-        )
+    @classmethod
+    def with_metadata(
+        cls, callback: Callback, **kwargs: Unpack[CallbackMetadataConfig]
+    ):
+        metadata: CallbackMetadataConfig = {}
+        metadata.update(cls.metadata)
+        metadata.update(kwargs)
+
+        return CallbackWithMetadata(callback=callback, metadata=metadata)
 
     @abstractmethod
     def create_callbacks(
@@ -73,7 +73,7 @@ def _filter_ignore_if_exists(callbacks: list[CallbackWithMetadata]):
     for callback in callbacks:
         # If `ignore_if_exists` is `True` and there is already a callback of the same class, skip this callback
         if (
-            callback.metadata.ignore_if_exists
+            callback.metadata.get("ignore_if_exists", False)
             and callback_classes[callback.callback.__class__] > 1
         ):
             continue
@@ -89,7 +89,10 @@ def _process_and_filter_callbacks(
     callbacks = list(callbacks)
 
     # Sort by priority (higher priority first)
-    callbacks.sort(key=lambda callback: callback.metadata.priority, reverse=True)
+    callbacks.sort(
+        key=lambda callback: callback.metadata.get("priority", 0),
+        reverse=True,
+    )
 
     # Process `ignore_if_exists`
     callbacks = _filter_ignore_if_exists(callbacks)

{nshtrainer-0.21.0 → nshtrainer-0.22.1}/src/nshtrainer/trainer/trainer.py

@@ -439,7 +439,8 @@ 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}.")
-            _link_checkpoint(cached_path, filepath, metadata=False)
+            if self.is_global_zero:
+                _link_checkpoint(cached_path, filepath, metadata=False)
         else:
             super().save_checkpoint(filepath, weights_only, storage_options)