nshtrainer 0.10.8__tar.gz → 0.10.10__tar.gz

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/PKG-INFO

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

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.10.8"
+version = "0.10.10"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -11,14 +11,14 @@ nshrunner = "*"
 nshconfig = "*"
 nshutils = "*"
 psutil = "*"
+numpy = "*"
 torch = "*"
 typing-extensions = "*"
 lightning = "*"
 pytorch-lightning = "*"
-torchmetrics = "*"
-numpy = "*"
-wrapt = "*"
-GitPython = "*"
+torchmetrics = { version = "*", optional = true }
+wrapt = { version = "*", optional = true }
+GitPython = { version = "*", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 pyright = "^1.1.372"

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/callbacks/early_stopping.py

@@ -1,5 +1,5 @@
+import logging
 import math
-from logging import getLogger
 
 from lightning.fabric.utilities.rank_zero import _get_rank
 from lightning.pytorch import Trainer
@@ -7,7 +7,7 @@ from lightning.pytorch.callbacks import EarlyStopping as _EarlyStopping
 from lightning.pytorch.utilities.rank_zero import rank_prefixed_message
 from typing_extensions import override
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 class EarlyStopping(_EarlyStopping):

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/callbacks/finite_checks.py

@@ -1,4 +1,4 @@
-from logging import getLogger
+import logging
 from typing import Literal
 
 import torch
@@ -7,7 +7,7 @@ from typing_extensions import override
 
 from .base import CallbackConfigBase
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def finite_checks(

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/callbacks/gradient_skipping.py

@@ -1,8 +1,8 @@
-from logging import getLogger
-from typing import Literal, Protocol, runtime_checkable
+import importlib.util
+import logging
+from typing import Any, Literal, Protocol, runtime_checkable
 
 import torch
-import torchmetrics
 from lightning.pytorch import Callback, LightningModule, Trainer
 from torch.optim import Optimizer
 from typing_extensions import override
@@ -10,23 +10,29 @@ from typing_extensions import override
 from .base import CallbackConfigBase
 from .norm_logging import compute_norm
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 @runtime_checkable
 class HasGradSkippedSteps(Protocol):
-    grad_skipped_steps: torchmetrics.SumMetric
+    grad_skipped_steps: Any
 
 
 class GradientSkipping(Callback):
     def __init__(self, config: "GradientSkippingConfig"):
-        super().__init__()
+        if importlib.util.find_spec("torchmetrics") is not None:
+            raise ImportError(
+                "To use the GradientSkipping callback, please install torchmetrics: pip install torchmetrics"
+            )
 
+        super().__init__()
         self.config = config
 
     @override
     def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
         if not isinstance(pl_module, HasGradSkippedSteps):
+            import torchmetrics  # type: ignore
+
             pl_module.grad_skipped_steps = torchmetrics.SumMetric()
 
     @override
@@ -47,12 +53,7 @@ class GradientSkipping(Callback):
         ):
             return
 
-        norm = compute_norm(
-            pl_module,
-            optimizer,
-            self.config.norm_type,
-            grad=True,
-        )
+        norm = compute_norm(pl_module, optimizer, self.config.norm_type, grad=True)
 
         # If the norm is NaN/Inf, we don't want to skip the step
         # beacuse AMP checks for NaN/Inf grads to adjust the loss scale.

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/callbacks/model_checkpoint.py

@@ -1,6 +1,6 @@
+import logging
 import re
 from datetime import timedelta
-from logging import getLogger
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal
 
@@ -15,7 +15,7 @@ from .base import CallbackConfigBase
 if TYPE_CHECKING:
     from ..model.config import BaseConfig
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def _convert_string(input_string: str):

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/callbacks/norm_logging.py

@@ -1,4 +1,4 @@
-from logging import getLogger
+import logging
 from typing import Literal, cast
 
 import torch
@@ -9,7 +9,7 @@ from typing_extensions import override
 
 from .base import CallbackConfigBase
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def grad_norm(

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/callbacks/throughput_monitor.py

@@ -1,4 +1,4 @@
-from logging import getLogger
+import logging
 from typing import Any, Literal, Protocol, TypedDict, cast, runtime_checkable
 
 from typing_extensions import NotRequired, override
@@ -6,7 +6,7 @@ from typing_extensions import NotRequired, override
 from ._throughput_monitor_callback import ThroughputMonitor as _ThroughputMonitor
 from .base import CallbackConfigBase
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 class ThroughputMonitorBatchStats(TypedDict):

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/callbacks/wandb_watch.py

@@ -1,4 +1,4 @@
-from logging import getLogger
+import logging
 from typing import Literal, Protocol, cast, runtime_checkable
 
 import torch.nn as nn
@@ -9,7 +9,7 @@ from typing_extensions import override
 
 from .base import CallbackConfigBase
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 @runtime_checkable

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/data/balanced_batch_sampler.py

@@ -1,6 +1,6 @@
 import heapq
+import logging
 from functools import cached_property
-from logging import getLogger
 from typing import Any, Protocol, runtime_checkable
 
 import numpy as np
@@ -10,7 +10,7 @@ from lightning_fabric.utilities.distributed import _DatasetSamplerWrapper
 from torch.utils.data import BatchSampler, Dataset, DistributedSampler
 from typing_extensions import override
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def _all_gather(tensor: torch.Tensor, device: torch.device | None = None):

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/data/transform.py

@@ -22,7 +22,13 @@ def transform(
         deepcopy: Whether to deep copy each item before applying the transform.
     """
 
-    import wrapt
+    try:
+        import wrapt
+    except ImportError:
+        raise ImportError(
+            "wrapt is not installed. wrapt is required for the transform function."
+            "Please install it using 'pip install wrapt'"
+        )
 
     class _TransformedDataset(wrapt.ObjectProxy):
         def __getitem__(self, idx):
@@ -52,7 +58,13 @@ def transform_with_index(
         deepcopy: Whether to deep copy each item before applying the transform.
     """
 
-    import wrapt
+    try:
+        import wrapt
+    except ImportError:
+        raise ImportError(
+            "wrapt is not installed. wrapt is required for the transform function."
+            "Please install it using 'pip install wrapt'"
+        )
 
     class _TransformedWithIndexDataset(wrapt.ObjectProxy):
         def __getitem__(self, idx: int):

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/model/_environment.py

@@ -9,7 +9,6 @@ from datetime import timedelta
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
-import git
 import nshconfig as C
 import psutil
 import torch
@@ -618,6 +617,11 @@ class GitRepositoryConfig(C.Config):
 
     @classmethod
     def from_current_directory(cls):
+        try:
+            import git
+        except ImportError:
+            return cls()
+
         draft = cls.draft()
         try:
             repo = git.Repo(os.getcwd(), search_parent_directories=True)

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/model/base.py

@@ -1,7 +1,7 @@
 import inspect
+import logging
 from abc import ABC, abstractmethod
 from collections.abc import MutableMapping
-from logging import getLogger
 from typing import IO, TYPE_CHECKING, Any, Generic, cast
 
 import torch
@@ -21,7 +21,7 @@ from .modules.profiler import ProfilerMixin
 from .modules.rlp_sanity_checks import RLPSanityCheckModuleMixin
 from .modules.shared_parameters import SharedParametersModuleMixin
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 THparams = TypeVar("THparams", bound=BaseConfig, infer_variance=True)
 

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/model/config.py

@@ -1,4 +1,5 @@
 import copy
+import logging
 import os
 import string
 import time
@@ -6,7 +7,6 @@ import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
 from datetime import timedelta
-from logging import getLogger
 from pathlib import Path
 from typing import (
     Annotated,
@@ -46,7 +46,7 @@ from ..callbacks.base import CallbackConfigBase
 from ..metrics import MetricConfig
 from ._environment import EnvironmentConfig
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 class IdSeedWarning(Warning):

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/model/modules/callback.py

@@ -1,6 +1,6 @@
+import logging
 from collections import abc
 from collections.abc import Callable, Iterable
-from logging import getLogger
 from typing import Any, TypeAlias, cast, final
 
 from lightning.pytorch import Callback, LightningModule
@@ -9,7 +9,7 @@ from typing_extensions import override
 
 from ...util.typing_utils import mixin_base_type
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 CallbackFn: TypeAlias = Callable[[], Callback | Iterable[Callback] | None]
 

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/model/modules/debug.py

@@ -1,9 +1,9 @@
-from logging import getLogger
+import logging
 
 import torch
 import torch.distributed
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 class DebugModuleMixin:

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/model/modules/rlp_sanity_checks.py

@@ -1,5 +1,5 @@
+import logging
 from collections.abc import Mapping
-from logging import getLogger
 from typing import cast
 
 import torch
@@ -14,7 +14,7 @@ from ...util.typing_utils import mixin_base_type
 from ..config import BaseConfig
 from .callback import CallbackModuleMixin
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def _on_train_start_callback(trainer: Trainer, pl_module: LightningModule):

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/model/modules/shared_parameters.py

@@ -1,5 +1,5 @@
+import logging
 from collections.abc import Sequence
-from logging import getLogger
 from typing import cast
 
 import torch.nn as nn
@@ -10,7 +10,7 @@ from ...util.typing_utils import mixin_base_type
 from ..config import BaseConfig
 from .callback import CallbackRegistrarModuleMixin
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def _parameters_to_names(parameters: Sequence[nn.Parameter], model: nn.Module):

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/util/environment.py

@@ -1,8 +1,8 @@
+import logging
 import os
 from contextlib import contextmanager
-from logging import getLogger
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 @contextmanager

{nshtrainer-0.10.8 → nshtrainer-0.10.10}/src/nshtrainer/util/seed.py

@@ -1,8 +1,8 @@
-from logging import getLogger
+import logging
 
 import lightning.fabric.utilities.seed as LS
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 def seed_everything(seed: int | None, *, workers: bool = False):

nshtrainer-0.10.8/src/nshtrainer/scripts/check_env.py

@@ -1,41 +0,0 @@
-REQUIRED_PACKAGES = [
-    "beartype",
-    "cloudpickle",
-    "jaxtyping",
-    "lightning",
-    "lightning_fabric",
-    "lightning_utilities",
-    "lovely_numpy",
-    "lovely_tensors",
-    "numpy",
-    "psutil",
-    "pydantic",
-    "pydantic_core",
-    "pysnooper",
-    "rich",
-    "tabulate",
-    "torch",
-    "torchmetrics",
-    "tqdm",
-    "typing_extensions",
-    "wrapt",
-    "yaml",
-]
-
-
-def main():
-    import importlib.util
-    import sys
-
-    missing_packages: list[str] = []
-    for package_name in REQUIRED_PACKAGES:
-        spec = importlib.util.find_spec(package_name)
-        if spec is None:
-            missing_packages.append(package_name)
-
-    if missing_packages:
-        sys.exit(f"Error: Missing required packages: {', '.join(missing_packages)}")
-
-
-if __name__ == "__main__":
-    main()