nshtrainer 0.1.0__tar.gz → 0.2.0__tar.gz

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.1.0
+Version: 0.2.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -9,10 +9,21 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: beartype (>=0.18.5,<0.19.0)
+Requires-Dist: jaxtyping (>=0.2.33,<0.3.0)
+Requires-Dist: lightning
+Requires-Dist: lovely-numpy (>=0.2.13,<0.3.0)
+Requires-Dist: lovely-tensors (>=0.1.16,<0.2.0)
 Requires-Dist: nshconfig (>=0.2.0,<0.3.0)
-Requires-Dist: nshrunner (>=0.1.0,<0.2.0)
+Requires-Dist: nshrunner (>=0.5.4,<0.6.0)
+Requires-Dist: numpy
+Requires-Dist: pysnooper
+Requires-Dist: pytorch-lightning
+Requires-Dist: rich
 Requires-Dist: torch
+Requires-Dist: torchmetrics
 Requires-Dist: typing-extensions
+Requires-Dist: wrapt
 Description-Content-Type: text/markdown
 
 

nshtrainer-0.2.0/pyproject.toml

@@ -0,0 +1,47 @@
+[tool.poetry]
+name = "nshtrainer"
+version = "0.2.0"
+description = ""
+authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.10"
+nshrunner = "^0.5.4"
+nshconfig = "^0.2.0"
+torch = "*"
+typing-extensions = "*"
+lightning = "*"
+pytorch-lightning = "*"
+torchmetrics = "*"
+numpy = "*"
+jaxtyping = "^0.2.33"
+beartype = "^0.18.5"
+lovely-numpy = "^0.2.13"
+lovely-tensors = "^0.1.16"
+pysnooper = "*"
+wrapt = "*"
+rich = "*"
+
+
+[tool.poetry.group.dev.dependencies]
+pyright = "^1.1.372"
+ruff = "^0.5.4"
+ipykernel = "^6.29.5"
+ipywidgets = "^8.1.3"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.pyright]
+typeCheckingMode = "standard"
+deprecateTypingAliases = true
+strictListInference = true
+strictDictionaryInference = true
+strictSetInference = true
+reportPrivateImportUsage = false
+ignore = ["./build/"]
+
+[tool.ruff.lint]
+ignore = ["F722", "F821", "E731", "E741"]

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/src/nshtrainer/__init__.py

@@ -1,25 +1,16 @@
 from . import _experimental as _experimental
 from . import actsave as actsave
 from . import callbacks as callbacks
+from . import config as config
 from . import lr_scheduler as lr_scheduler
 from . import model as model
 from . import nn as nn
 from . import optimizer as optimizer
-from . import snapshot as snapshot
 from . import typecheck as typecheck
 from ._snoop import snoop as snoop
 from .actsave import ActLoad as ActLoad
 from .actsave import ActSave as ActSave
-from .config import MISSING as MISSING
-from .config import AllowMissing as AllowMissing
-from .config import Field as Field
-from .config import MissingField as MissingField
-from .config import PrivateAttr as PrivateAttr
-from .config import TypedConfig as TypedConfig
 from .data import dataset_transform as dataset_transform
-from .log import init_python_logging as init_python_logging
-from .log import lovely as lovely
-from .log import pretty as pretty
 from .lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from .model import ActSaveConfig as ActSaveConfig
 from .model import Base as Base
@@ -41,24 +32,17 @@ from .model import (
     EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
 )
 from .model import GradientClippingConfig as GradientClippingConfig
-from .model import LightningDataModuleBase as LightningDataModuleBase
 from .model import LightningModuleBase as LightningModuleBase
 from .model import LoggingConfig as LoggingConfig
 from .model import MetricConfig as MetricConfig
 from .model import OptimizationConfig as OptimizationConfig
 from .model import PrimaryMetricConfig as PrimaryMetricConfig
-from .model import PythonLogging as PythonLogging
 from .model import ReproducibilityConfig as ReproducibilityConfig
-from .model import RunnerConfig as RunnerConfig
 from .model import SanityCheckingConfig as SanityCheckingConfig
-from .model import SeedConfig as SeedConfig
 from .model import TrainerConfig as TrainerConfig
 from .model import WandbWatchConfig as WandbWatchConfig
 from .nn import TypedModuleDict as TypedModuleDict
 from .nn import TypedModuleList as TypedModuleList
 from .optimizer import OptimizerConfig as OptimizerConfig
 from .runner import Runner as Runner
-from .runner import SnapshotConfig as SnapshotConfig
 from .trainer import Trainer as Trainer
-from .util.singleton import Registry as Registry
-from .util.singleton import Singleton as Singleton

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/src/nshtrainer/callbacks/__init__.py

@@ -1,6 +1,7 @@
 from typing import Annotated
 
-from ..config import Field
+import nshconfig as C
+
 from .base import CallbackConfigBase as CallbackConfigBase
 from .early_stopping import EarlyStopping as EarlyStopping
 from .ema import EMA as EMA
@@ -31,5 +32,5 @@ CallbackConfig = Annotated[
     | NormLoggingConfig
     | GradientSkippingConfig
     | EMAConfig,
-    Field(discriminator="name"),
+    C.Field(discriminator="name"),
 ]

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/src/nshtrainer/callbacks/base.py

@@ -4,10 +4,9 @@ from collections.abc import Iterable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, TypeAlias, TypedDict
 
+import nshconfig as C
 from lightning.pytorch import Callback
 
-from ..config import TypedConfig
-
 if TYPE_CHECKING:
     from ..model.config import BaseConfig
 
@@ -20,7 +19,7 @@ class CallbackMetadataDict(TypedDict, total=False):
     """Priority of the callback. Callbacks with higher priority will be loaded first."""
 
 
-class CallbackMetadataConfig(TypedConfig):
+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."""
 
@@ -37,7 +36,7 @@ class CallbackWithMetadata:
 ConstructedCallback: TypeAlias = Callback | CallbackWithMetadata
 
 
-class CallbackConfigBase(TypedConfig, ABC):
+class CallbackConfigBase(C.Config, ABC):
     metadata: CallbackMetadataConfig = CallbackMetadataConfig()
     """Metadata for the callback."""
 

nshtrainer-0.2.0/src/nshtrainer/config.py

@@ -0,0 +1,4 @@
+from nshconfig import *  # type: ignore # noqa: F403
+from nshconfig import Config
+
+TypedConfig = Config

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/src/nshtrainer/lr_scheduler/__init__.py

@@ -1,6 +1,7 @@
 from typing import Annotated, TypeAlias
 
-from ..config import Field
+import nshconfig as C
+
 from ._base import LRSchedulerConfigBase as LRSchedulerConfigBase
 from ._base import LRSchedulerMetadata as LRSchedulerMetadata
 from .linear_warmup_cosine import (
@@ -14,5 +15,5 @@ from .reduce_lr_on_plateau import ReduceLROnPlateauConfig as ReduceLROnPlateauCo
 
 LRSchedulerConfig: TypeAlias = Annotated[
     LinearWarmupCosineDecayLRSchedulerConfig | ReduceLROnPlateauConfig,
-    Field(discriminator="name"),
+    C.Field(discriminator="name"),
 ]

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/src/nshtrainer/lr_scheduler/_base.py

@@ -1,8 +1,9 @@
 import math
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, ClassVar, Literal, TypeAlias
+from typing import TYPE_CHECKING, Literal
 
+import nshconfig as C
 from lightning.pytorch.utilities.types import (
     LRSchedulerConfigType,
     LRSchedulerTypeUnion,
@@ -10,8 +11,6 @@ from lightning.pytorch.utilities.types import (
 from torch.optim import Optimizer
 from typing_extensions import NotRequired, TypedDict
 
-from ..config import TypedConfig
-
 if TYPE_CHECKING:
     from ..model.base import LightningModuleBase
 
@@ -37,9 +36,7 @@ class LRSchedulerMetadata(TypedDict):
     """Whether to enforce that the monitor exists for reducing the learning rate on plateau. Default is `True`."""
 
 
-class LRSchedulerConfigBase(TypedConfig, ABC):
-    Metadata: ClassVar[TypeAlias] = LRSchedulerMetadata
-
+class LRSchedulerConfigBase(C.Config, ABC):
     @abstractmethod
     def metadata(self) -> LRSchedulerMetadata: ...
 

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/src/nshtrainer/lr_scheduler/linear_warmup_cosine.py

@@ -2,12 +2,12 @@ import math
 import warnings
 from typing import Literal
 
+import nshconfig as C
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from typing_extensions import override
 
-from ..config import Field
-from ._base import LRSchedulerConfigBase
+from ._base import LRSchedulerConfigBase, LRSchedulerMetadata
 
 
 class LinearWarmupCosineAnnealingLR(LRScheduler):
@@ -91,11 +91,11 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
 class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
     name: Literal["linear_warmup_cosine_decay"] = "linear_warmup_cosine_decay"
 
-    warmup_epochs: int = Field(ge=0)
+    warmup_epochs: int = C.Field(ge=0)
     r"""The number of epochs for the linear warmup phase.
     The learning rate is linearly increased from `warmup_start_lr` to the initial learning rate over this number of epochs."""
 
-    max_epochs: int = Field(gt=0)
+    max_epochs: int = C.Field(gt=0)
     r"""The total number of epochs.
     The learning rate is decayed to `min_lr` over this number of epochs."""
 
@@ -113,7 +113,7 @@ class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
     If `True`, the learning rate will be decayed to `min_lr` over `max_epochs` epochs, and then the learning rate will be increased back to the initial learning rate over `max_epochs` epochs, and so on (this is called a cosine annealing schedule)."""
 
     @override
-    def metadata(self) -> LRSchedulerConfigBase.Metadata:
+    def metadata(self) -> LRSchedulerMetadata:
         return {
             "interval": "step",
         }

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/src/nshtrainer/lr_scheduler/reduce_lr_on_plateau.py

@@ -1,12 +1,11 @@
 from typing import TYPE_CHECKING, Literal, cast
 
+from lightning.pytorch.utilities.types import LRSchedulerConfigType
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 from typing_extensions import override
 
-from ll.lr_scheduler._base import LRSchedulerMetadata
-
 from ..model.config import MetricConfig
-from ._base import LRSchedulerConfigBase
+from ._base import LRSchedulerConfigBase, LRSchedulerMetadata
 
 if TYPE_CHECKING:
     from ..model.base import BaseConfig
@@ -43,7 +42,9 @@ class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
     r"""One of `rel`, `abs`. In `rel` mode, dynamic_threshold = best * (1 + threshold) in 'max' mode or best * (1 - threshold) in `min` mode. In `abs` mode, dynamic_threshold = best + threshold in `max` mode or best - threshold in `min` mode. Default: 'rel'."""
 
     @override
-    def create_scheduler_impl(self, optimizer, lightning_module, lr):
+    def create_scheduler_impl(
+        self, optimizer, lightning_module, lr
+    ) -> LRSchedulerConfigType:
         if (metric := self.metric) is None:
             lm_config = cast("BaseConfig", lightning_module.config)
             assert (

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/src/nshtrainer/model/__init__.py

@@ -1,7 +1,6 @@
 from typing_extensions import TypeAlias
 
 from .base import Base as Base
-from .base import LightningDataModuleBase as LightningDataModuleBase
 from .base import LightningModuleBase as LightningModuleBase
 from .config import ActSaveConfig as ActSaveConfig
 from .config import BaseConfig as BaseConfig
@@ -33,11 +32,8 @@ from .config import (
 )
 from .config import OptimizationConfig as OptimizationConfig
 from .config import PrimaryMetricConfig as PrimaryMetricConfig
-from .config import PythonLogging as PythonLogging
 from .config import ReproducibilityConfig as ReproducibilityConfig
-from .config import RunnerConfig as RunnerConfig
 from .config import SanityCheckingConfig as SanityCheckingConfig
-from .config import SeedConfig as SeedConfig
 from .config import TrainerConfig as TrainerConfig
 from .config import WandbWatchConfig as WandbWatchConfig
 

{nshtrainer-0.1.0 → nshtrainer-0.2.0}/src/nshtrainer/model/base.py

@@ -23,11 +23,12 @@ from .config import (
     EnvironmentLinuxEnvironmentConfig,
     EnvironmentLSFInformationConfig,
     EnvironmentSLURMInformationConfig,
+    EnvironmentSnapshotConfig,
 )
-from .modules.callback import CallbackModuleMixin, CallbackRegistrarModuleMixin
+from .modules.callback import CallbackModuleMixin
 from .modules.debug import DebugModuleMixin
 from .modules.distributed import DistributedMixin
-from .modules.logger import LoggerLightningModuleMixin, LoggerModuleMixin
+from .modules.logger import LoggerLightningModuleMixin
 from .modules.profiler import ProfilerMixin
 from .modules.rlp_sanity_checks import RLPSanityCheckModuleMixin
 from .modules.shared_parameters import SharedParametersModuleMixin
@@ -265,6 +266,9 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
             boot_time=_try_get(lambda: _psutil().boot_time()),
             load_avg=_try_get(lambda: os.getloadavg()),
         )
+        hparams.environment.snapshot = (
+            EnvironmentSnapshotConfig.from_current_environment()
+        )
 
     def pre_init_update_hparams_dict(self, hparams: MutableMapping[str, Any]):
         """
@@ -309,15 +313,12 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
     @property
     def datamodule(self):
         datamodule = getattr(self.trainer, "datamodule", None)
-        if datamodule is None:
+        if (datamodule := getattr(self.trainer, "datamodule", None)) is None:
             return None
-
-        if not isinstance(datamodule, LightningDataModuleBase):
+        if not isinstance(datamodule, LightningDataModule):
             raise TypeError(
-                f"datamodule must be a LightningDataModuleBase: {type(datamodule)}"
+                f"datamodule must be a LightningDataModule: {type(datamodule)}"
             )
-
-        datamodule = cast(LightningDataModuleBase[THparams], datamodule)
         return datamodule
 
     if TYPE_CHECKING:
@@ -576,66 +577,3 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
             "batch": batch,
             "batch_idx": batch_idx,
         }
-
-
-class LightningDataModuleBase(
-    LoggerModuleMixin,
-    CallbackRegistrarModuleMixin,
-    Base[THparams],
-    LightningDataModule,
-    ABC,
-    Generic[THparams],
-):
-    hparams: THparams  # pyright: ignore[reportIncompatibleMethodOverride]
-    hparams_initial: THparams  # pyright: ignore[reportIncompatibleMethodOverride]
-
-    def pre_init_update_hparams_dict(self, hparams: MutableMapping[str, Any]):
-        """
-        Override this method to update the hparams dictionary before it is used to create the hparams object.
-        Mapping-based parameters are passed to the constructor of the hparams object when we're loading the model from a checkpoint.
-        """
-        return hparams
-
-    def pre_init_update_hparams(self, hparams: THparams):
-        """
-        Override this method to update the hparams object before it is used to create the hparams_initial object.
-        """
-        return hparams
-
-    @classmethod
-    def _update_environment(cls, hparams: THparams):
-        hparams.environment.data = _cls_info(cls)
-
-    @override
-    def __init__(self, hparams: THparams):
-        if not isinstance(hparams, BaseConfig):
-            if not isinstance(hparams, MutableMapping):
-                raise TypeError(
-                    f"hparams must be a BaseConfig or a MutableMapping: {type(hparams)}"
-                )
-
-            hparams = self.pre_init_update_hparams_dict(hparams)
-            hparams = self.config_cls().from_dict(hparams)
-        self._update_environment(hparams)
-        hparams = self.pre_init_update_hparams(hparams)
-        super().__init__(hparams)
-
-        self.save_hyperparameters(hparams)
-
-    @property
-    def lightning_module(self):
-        if not self.trainer:
-            raise ValueError("Trainer has not been set.")
-
-        module = self.trainer.lightning_module
-        if not isinstance(module, LightningModuleBase):
-            raise ValueError(
-                f"Trainer's lightning_module is not a LightningModuleBase: {type(module)}"
-            )
-
-        module = cast(LightningModuleBase[THparams], module)
-        return module
-
-    @property
-    def device(self):
-        return self.lightning_module.device