nshtrainer 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

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/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/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/config.py

@@ -1,289 +1,4 @@
-from collections.abc import Mapping, MutableMapping
-from typing import TYPE_CHECKING, Any, ClassVar
+from nshconfig import *  # type: ignore # noqa: F403
+from nshconfig import Config
 
-from pydantic import BaseModel, ConfigDict
-from pydantic import Field as Field
-from pydantic import PrivateAttr as PrivateAttr
-from typing_extensions import deprecated, override
-
-from ._config.missing import MISSING, validate_no_missing_values
-from ._config.missing import AllowMissing as AllowMissing
-from ._config.missing import MissingField as MissingField
-
-_MutableMappingBase = MutableMapping[str, Any]
-if TYPE_CHECKING:
-    _MutableMappingBase = object
-
-
-_DraftConfigContextSentinel = object()
-
-
-class TypedConfig(BaseModel, _MutableMappingBase):
-    _is_draft_config: bool = PrivateAttr(default=False)
-    """
-    Whether this config is a draft config or not.
-
-    Draft configs are configs that are not yet fully validated.
-    They allow for a nicer API when creating configs, e.g.:
-
-        ```python
-        config = MyConfig.draft()
-
-        # Set some values
-        config.a = 10
-        config.b = "hello"
-
-        # Finalize the config
-        config = config.finalize()
-        ```
-    """
-
-    repr_diff_only: ClassVar[bool] = True
-    """
-    If `True`, the repr methods will only show values for fields that are different from the default.
-    """
-
-    MISSING: ClassVar[Any] = MISSING
-    """
-    Alias for the `MISSING` constant.
-    """
-
-    model_config: ClassVar[ConfigDict] = ConfigDict(
-        # By default, Pydantic will throw a warning if a field starts with "model_",
-        # so we need to disable that warning (beacuse "model_" is a popular prefix for ML).
-        protected_namespaces=(),
-        validate_assignment=True,
-        validate_return=True,
-        validate_default=True,
-        strict=True,
-        revalidate_instances="always",
-        arbitrary_types_allowed=True,
-        extra="ignore",
-        validation_error_cause=True,
-        use_attribute_docstrings=True,
-    )
-
-    def __draft_pre_init__(self):
-        """Called right before a draft config is finalized."""
-        pass
-
-    def __post_init__(self):
-        """Called after the final config is validated."""
-        pass
-
-    @classmethod
-    @deprecated("Use `model_validate` instead.")
-    def from_dict(cls, model_dict: Mapping[str, Any]):
-        return cls.model_validate(model_dict)
-
-    def model_deep_validate(self, strict: bool = True):
-        """
-        Validate the config and all of its sub-configs.
-
-        Args:
-            config: The config to validate.
-            strict: Whether to validate the config strictly.
-        """
-        config_dict = self.model_dump(round_trip=True)
-        config = self.model_validate(config_dict, strict=strict)
-
-        # Make sure that this is not a draft config
-        if config._is_draft_config:
-            raise ValueError("Draft configs are not valid. Call `finalize` first.")
-
-        return config
-
-    @classmethod
-    def draft(cls, **kwargs):
-        config = cls.model_construct_draft(**kwargs)
-        return config
-
-    def finalize(self, strict: bool = True):
-        # This must be a draft config, otherwise we raise an error
-        if not self._is_draft_config:
-            raise ValueError("Finalize can only be called on drafts.")
-
-        # First, we call `__draft_pre_init__` to allow the config to modify itself a final time
-        self.__draft_pre_init__()
-
-        # Then, we dump the config to a dict and then re-validate it
-        return self.model_deep_validate(strict=strict)
-
-    @override
-    def model_post_init(self, __context: Any) -> None:
-        super().model_post_init(__context)
-
-        # Call the `__post_init__` method if this is not a draft config
-        if __context is _DraftConfigContextSentinel:
-            return
-
-        self.__post_init__()
-
-        # After `_post_init__` is called, we perform the final round of validation
-        self.model_post_init_validate()
-
-    def model_post_init_validate(self):
-        validate_no_missing_values(self)
-
-    @classmethod
-    def model_construct_draft(cls, _fields_set: set[str] | None = None, **values: Any):
-        """
-        NOTE: This is a copy of the `model_construct` method from Pydantic's `Model` class,
-            with the following changes:
-            - The `model_post_init` method is called with the `_DraftConfigContext` context.
-            - The `_is_draft_config` attribute is set to `True` in the `values` dict.
-
-        Creates a new instance of the `Model` class with validated data.
-
-        Creates a new model setting `__dict__` and `__pydantic_fields_set__` from trusted or pre-validated data.
-        Default values are respected, but no other validation is performed.
-
-        !!! note
-            `model_construct()` generally respects the `model_config.extra` setting on the provided model.
-            That is, if `model_config.extra == 'allow'`, then all extra passed values are added to the model instance's `__dict__`
-            and `__pydantic_extra__` fields. If `model_config.extra == 'ignore'` (the default), then all extra passed values are ignored.
-            Because no validation is performed with a call to `model_construct()`, having `model_config.extra == 'forbid'` does not result in
-            an error if extra values are passed, but they will be ignored.
-
-        Args:
-            _fields_set: The set of field names accepted for the Model instance.
-            values: Trusted or pre-validated data dictionary.
-
-        Returns:
-            A new instance of the `Model` class with validated data.
-        """
-
-        values["_is_draft_config"] = True
-
-        m = cls.__new__(cls)
-        fields_values: dict[str, Any] = {}
-        fields_set = set()
-
-        for name, field in cls.model_fields.items():
-            if field.alias and field.alias in values:
-                fields_values[name] = values.pop(field.alias)
-                fields_set.add(name)
-            elif name in values:
-                fields_values[name] = values.pop(name)
-                fields_set.add(name)
-            elif not field.is_required():
-                fields_values[name] = field.get_default(call_default_factory=True)
-        if _fields_set is None:
-            _fields_set = fields_set
-
-        _extra: dict[str, Any] | None = None
-        if cls.model_config.get("extra") == "allow":
-            _extra = {}
-            for k, v in values.items():
-                _extra[k] = v
-        object.__setattr__(m, "__dict__", fields_values)
-        object.__setattr__(m, "__pydantic_fields_set__", _fields_set)
-        if not cls.__pydantic_root_model__:
-            object.__setattr__(m, "__pydantic_extra__", _extra)
-
-        if cls.__pydantic_post_init__:
-            m.model_post_init(_DraftConfigContextSentinel)
-            # update private attributes with values set
-            if (
-                hasattr(m, "__pydantic_private__")
-                and m.__pydantic_private__ is not None
-            ):
-                for k, v in values.items():
-                    if k in m.__private_attributes__:
-                        m.__pydantic_private__[k] = v
-
-        elif not cls.__pydantic_root_model__:
-            # Note: if there are any private attributes, cls.__pydantic_post_init__ would exist
-            # Since it doesn't, that means that `__pydantic_private__` should be set to None
-            object.__setattr__(m, "__pydantic_private__", None)
-
-        return m
-
-    @override
-    def __repr_args__(self):
-        # If `repr_diff_only` is `True`, we only show the fields that are different from the default.
-        if not self.repr_diff_only:
-            yield from super().__repr_args__()
-            return
-
-        # First, we get the default values for all fields.
-        default_values = self.model_construct_draft()
-
-        # Then, we compare the default values with the current values.
-        for k, v in super().__repr_args__():
-            if k is None:
-                yield k, v
-                continue
-
-            # If there is no default value or the value is different from the default, we yield it.
-            if not hasattr(default_values, k) or getattr(default_values, k) != v:
-                yield k, v
-                continue
-
-            # Otherwise, we can skip this field.
-
-    # region MutableMapping implementation
-    if not TYPE_CHECKING:
-        # This is mainly so the config can be used with lightning's hparams
-        #   transparently and without any issues.
-
-        @property
-        def _ll_dict(self):
-            return self.model_dump()
-
-        # We need to make sure every config class
-        #   is a MutableMapping[str, Any] so that it can be used
-        #   with lightning's hparams.
-        @override
-        def __getitem__(self, key: str):
-            # Key can be of the format "a.b.c"
-            #   so we need to split it into a list of keys.
-            [first_key, *rest_keys] = key.split(".")
-            value = self._ll_dict[first_key]
-
-            for key in rest_keys:
-                if isinstance(value, Mapping):
-                    value = value[key]
-                else:
-                    value = getattr(value, key)
-
-            return value
-
-        @override
-        def __setitem__(self, key: str, value: Any):
-            # Key can be of the format "a.b.c"
-            #   so we need to split it into a list of keys.
-            [first_key, *rest_keys] = key.split(".")
-            if len(rest_keys) == 0:
-                self._ll_dict[first_key] = value
-                return
-
-            # We need to traverse the keys until we reach the last key
-            #   and then set the value
-            current_value = self._ll_dict[first_key]
-            for key in rest_keys[:-1]:
-                if isinstance(current_value, Mapping):
-                    current_value = current_value[key]
-                else:
-                    current_value = getattr(current_value, key)
-
-            # Set the value
-            if isinstance(current_value, MutableMapping):
-                current_value[rest_keys[-1]] = value
-            else:
-                setattr(current_value, rest_keys[-1], value)
-
-        @override
-        def __delitem__(self, key: str):
-            # This is unsupported for this class
-            raise NotImplementedError
-
-        @override
-        def __iter__(self):
-            return iter(self._ll_dict)
-
-        @override
-        def __len__(self):
-            return len(self._ll_dict)
-
-    # endregion
+TypedConfig = Config

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/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/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/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/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/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