nshtrainer 1.3.4__tar.gz → 1.3.6__tar.gz

{nshtrainer-1.3.4 → nshtrainer-1.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.3.4
+Version: 1.3.6
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -15,7 +15,7 @@ Requires-Dist: GitPython ; extra == "extra"
 Requires-Dist: huggingface-hub ; extra == "extra"
 Requires-Dist: lightning
 Requires-Dist: nshconfig (>0.39)
-Requires-Dist: nshrunner
+Requires-Dist: nshrunner ; extra == "extra"
 Requires-Dist: nshutils ; extra == "extra"
 Requires-Dist: numpy
 Requires-Dist: packaging

{nshtrainer-1.3.4 → nshtrainer-1.3.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nshtrainer"
-version = "1.3.4"
+version = "1.3.6"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.10,<4.0"
@@ -8,7 +8,6 @@ readme = "README.md"
 
 dependencies = [
     "nshconfig>0.39",
-    "nshrunner",
     "psutil",
     "numpy",
     "torch",
@@ -20,6 +19,7 @@ dependencies = [
 
 [project.optional-dependencies]
 extra = [
+    "nshrunner",
     "wrapt",
     "GitPython",
     "wandb",

{nshtrainer-1.3.4 → nshtrainer-1.3.6}/src/nshtrainer/__init__.py

@@ -19,3 +19,17 @@ try:
     from . import configs as configs
 except BaseException:
     pass
+
+try:
+    from importlib.metadata import PackageNotFoundError, version
+except ImportError:
+    # For Python <3.8
+    from importlib_metadata import (  # pyright: ignore[reportMissingImports]
+        PackageNotFoundError,
+        version,
+    )
+
+try:
+    __version__ = version(__name__)
+except PackageNotFoundError:
+    __version__ = "unknown"

{nshtrainer-1.3.4 → nshtrainer-1.3.6}/src/nshtrainer/_hf_hub.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import contextlib
 import logging
-import os
 import re
 from dataclasses import dataclass
 from functools import cached_property
@@ -10,7 +9,6 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
 
 import nshconfig as C
-from nshrunner._env import SNAPSHOT_DIR
 from typing_extensions import assert_never, override
 
 from ._callback import NTCallbackBase
@@ -19,6 +17,7 @@ from .callbacks.base import (
     CallbackMetadataConfig,
     callback_registry,
 )
+from .util.code_upload import get_code_dir
 
 if TYPE_CHECKING:
     from huggingface_hub import HfApi  # noqa: F401
@@ -319,20 +318,13 @@ class HFHubCallback(NTCallbackBase):
     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)):
+        if (snapshot_dir := get_code_dir()) is None:
             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
-
             self.api.upload_folder(
-                folder_path=str(snapshot_dir),
+                folder_path=str(snapshot_dir.absolute()),
                 repo_id=self.repo_id,
                 repo_type="model",
                 path_in_repo="code",  # Prefix with "code" folder

{nshtrainer-1.3.4 → nshtrainer-1.3.6}/src/nshtrainer/callbacks/wandb_upload_code.py

@@ -1,16 +1,14 @@
 from __future__ import annotations
 
 import logging
-import os
-from pathlib import Path
 from typing import Literal, cast
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks.callback import Callback
 from lightning.pytorch.loggers import WandbLogger
-from nshrunner._env import SNAPSHOT_DIR
 from typing_extensions import final, override
 
+from ..util.code_upload import get_code_dir
 from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
@@ -62,22 +60,12 @@ class WandbUploadCodeCallback(Callback):
             log.warning("Wandb logger not found. Skipping code upload.")
             return
 
-        from wandb.wandb_run import Run
-
-        run = cast(Run, logger.experiment)
-
-        # 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.")
+        if (snapshot_dir := get_code_dir()) is None:
+            log.info("No nshrunner snapshot found. Skipping code upload.")
             return
 
-        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
+        from wandb.wandb_run import Run
 
+        run = cast(Run, logger.experiment)
         log.info(f"Uploading code from snapshot directory '{snapshot_dir}'")
         run.log_code(str(snapshot_dir.absolute()))

{nshtrainer-1.3.4 → nshtrainer-1.3.6}/src/nshtrainer/configs/__init__.py

@@ -5,7 +5,6 @@ __codegen__ = True
 from nshtrainer import MetricConfig as MetricConfig
 from nshtrainer import TrainerConfig as TrainerConfig
 from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
-from nshtrainer._directory import DirectoryConfig as DirectoryConfig
 from nshtrainer._hf_hub import CallbackConfigBase as CallbackConfigBase
 from nshtrainer._hf_hub import (
     HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig,
@@ -126,9 +125,9 @@ from nshtrainer.trainer._config import (
     CheckpointCallbackConfig as CheckpointCallbackConfig,
 )
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
+from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
 from nshtrainer.trainer._config import EnvironmentConfig as EnvironmentConfig
 from nshtrainer.trainer._config import GradientClippingConfig as GradientClippingConfig
-from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
 from nshtrainer.trainer._config import StrategyConfig as StrategyConfig
 from nshtrainer.trainer.accelerator import CPUAcceleratorConfig as CPUAcceleratorConfig
 from nshtrainer.trainer.accelerator import (
@@ -227,7 +226,6 @@ from nshtrainer.util.config import EpochsConfig as EpochsConfig
 from nshtrainer.util.config import StepsConfig as StepsConfig
 
 from . import _checkpoint as _checkpoint
-from . import _directory as _directory
 from . import _hf_hub as _hf_hub
 from . import callbacks as callbacks
 from . import loggers as loggers
@@ -338,7 +336,6 @@ __all__ = [
     "RpropConfig",
     "SGDConfig",
     "SLURMEnvironmentPlugin",
-    "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "SiLUNonlinearityConfig",
     "SigmoidNonlinearityConfig",
@@ -367,7 +364,6 @@ __all__ = [
     "XLAEnvironmentPlugin",
     "XLAPluginConfig",
     "_checkpoint",
-    "_directory",
     "_hf_hub",
     "accelerator_registry",
     "callback_registry",

{nshtrainer-1.3.4 → nshtrainer-1.3.6}/src/nshtrainer/configs/trainer/__init__.py

@@ -22,6 +22,9 @@ from nshtrainer.trainer._config import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
 from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
+from nshtrainer.trainer._config import (
+    DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
+)
 from nshtrainer.trainer._config import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -51,7 +54,6 @@ from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
-from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
@@ -152,6 +154,7 @@ __all__ = [
     "DebugFlagCallbackConfig",
     "DeepSpeedPluginConfig",
     "DirectoryConfig",
+    "DirectorySetupCallbackConfig",
     "DistributedPredictionWriterConfig",
     "DoublePrecisionPluginConfig",
     "EarlyStoppingCallbackConfig",
@@ -180,7 +183,6 @@ __all__ = [
     "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
     "SLURMEnvironmentPlugin",
-    "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "StrategyConfig",
     "StrategyConfigBase",

{nshtrainer-1.3.4 → nshtrainer-1.3.6}/src/nshtrainer/configs/trainer/_config/__init__.py

@@ -18,6 +18,9 @@ from nshtrainer.trainer._config import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
 from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
+from nshtrainer.trainer._config import (
+    DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
+)
 from nshtrainer.trainer._config import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -48,7 +51,6 @@ from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
-from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
@@ -70,6 +72,7 @@ __all__ = [
     "CheckpointSavingConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
+    "DirectorySetupCallbackConfig",
     "EarlyStoppingCallbackConfig",
     "EnvironmentConfig",
     "GradientClippingConfig",
@@ -86,7 +89,6 @@ __all__ = [
     "PluginConfig",
     "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
-    "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "StrategyConfig",
     "TensorboardLoggerConfig",