nshtrainer 0.44.1__py3-none-any.whl → 1.0.0b10__py3-none-any.whl

nshtrainer/config/trainer/_config/__init__.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
 # Config/alias imports
 
 if TYPE_CHECKING:
+    from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
     from nshtrainer.trainer._config import (
         BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
     )
@@ -25,9 +26,11 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         DebugFlagCallbackConfig as DebugFlagCallbackConfig,
     )
+    from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
     from nshtrainer.trainer._config import (
         EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
     )
+    from nshtrainer.trainer._config import EnvironmentConfig as EnvironmentConfig
     from nshtrainer.trainer._config import (
         GradientClippingConfig as GradientClippingConfig,
     )
@@ -35,8 +38,12 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
     )
+    from nshtrainer.trainer._config import (
+        LogEpochCallbackConfig as LogEpochCallbackConfig,
+    )
     from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
     from nshtrainer.trainer._config import LoggingConfig as LoggingConfig
+    from nshtrainer.trainer._config import MetricConfig as MetricConfig
     from nshtrainer.trainer._config import (
         OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
     )
@@ -64,80 +71,96 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "SanityCheckingConfig":
+        if name == "ActSaveLoggerConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).SanityCheckingConfig
-        if name == "TrainerConfig":
-            return importlib.import_module("nshtrainer.trainer._config").TrainerConfig
-        if name == "OnExceptionCheckpointCallbackConfig":
+            ).ActSaveLoggerConfig
+        if name == "BestCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OnExceptionCheckpointCallbackConfig
-        if name == "GradientClippingConfig":
+            ).BestCheckpointCallbackConfig
+        if name == "CSVLoggerConfig":
+            return importlib.import_module("nshtrainer.trainer._config").CSVLoggerConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).GradientClippingConfig
-        if name == "WandbLoggerConfig":
+            ).CallbackConfigBase
+        if name == "CheckpointLoadingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).WandbLoggerConfig
-        if name == "LoggingConfig":
-            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
-        if name == "TensorboardLoggerConfig":
+            ).CheckpointLoadingConfig
+        if name == "CheckpointSavingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).TensorboardLoggerConfig
-        if name == "RLPSanityChecksCallbackConfig":
+            ).CheckpointSavingConfig
+        if name == "DebugFlagCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).RLPSanityChecksCallbackConfig
-        if name == "CheckpointSavingConfig":
+            ).DebugFlagCallbackConfig
+        if name == "DirectoryConfig":
+            return importlib.import_module("nshtrainer.trainer._config").DirectoryConfig
+        if name == "EarlyStoppingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).CheckpointSavingConfig
-        if name == "CSVLoggerConfig":
-            return importlib.import_module("nshtrainer.trainer._config").CSVLoggerConfig
+            ).EarlyStoppingCallbackConfig
+        if name == "EnvironmentConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).EnvironmentConfig
+        if name == "GradientClippingConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).GradientClippingConfig
         if name == "HuggingFaceHubConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).HuggingFaceHubConfig
-        if name == "CheckpointLoadingConfig":
+        if name == "LastCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).CheckpointLoadingConfig
-        if name == "DebugFlagCallbackConfig":
+            ).LastCheckpointCallbackConfig
+        if name == "LogEpochCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).DebugFlagCallbackConfig
-        if name == "CallbackConfigBase":
+            ).LogEpochCallbackConfig
+        if name == "LoggingConfig":
+            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
+        if name == "MetricConfig":
+            return importlib.import_module("nshtrainer.trainer._config").MetricConfig
+        if name == "OnExceptionCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).CallbackConfigBase
-        if name == "LastCheckpointCallbackConfig":
+            ).OnExceptionCheckpointCallbackConfig
+        if name == "OptimizationConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).LastCheckpointCallbackConfig
-        if name == "SharedParametersCallbackConfig":
+            ).OptimizationConfig
+        if name == "RLPSanityChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).SharedParametersCallbackConfig
+            ).RLPSanityChecksCallbackConfig
         if name == "ReproducibilityConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).ReproducibilityConfig
-        if name == "EarlyStoppingCallbackConfig":
+        if name == "SanityCheckingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).EarlyStoppingCallbackConfig
-        if name == "OptimizationConfig":
+            ).SanityCheckingConfig
+        if name == "SharedParametersCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OptimizationConfig
-        if name == "BestCheckpointCallbackConfig":
+            ).SharedParametersCallbackConfig
+        if name == "TensorboardLoggerConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).BestCheckpointCallbackConfig
+            ).TensorboardLoggerConfig
+        if name == "TrainerConfig":
+            return importlib.import_module("nshtrainer.trainer._config").TrainerConfig
+        if name == "WandbLoggerConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).WandbLoggerConfig
         if name == "CallbackConfig":
             return importlib.import_module("nshtrainer.trainer._config").CallbackConfig
         if name == "CheckpointCallbackConfig":

nshtrainer/config/trainer/trainer/__init__.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.trainer.trainer import EnvironmentConfig as EnvironmentConfig
+    from nshtrainer.trainer.trainer import TrainerConfig as TrainerConfig
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "EnvironmentConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer.trainer"
+            ).EnvironmentConfig
+        if name == "TrainerConfig":
+            return importlib.import_module("nshtrainer.trainer.trainer").TrainerConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+# Submodule exports

nshtrainer/config/util/__init__.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.util._environment_info import (
+        EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
+    )
+    from nshtrainer.util._environment_info import EnvironmentConfig as EnvironmentConfig
+    from nshtrainer.util._environment_info import (
+        EnvironmentCUDAConfig as EnvironmentCUDAConfig,
+    )
+    from nshtrainer.util._environment_info import (
+        EnvironmentGPUConfig as EnvironmentGPUConfig,
+    )
+    from nshtrainer.util._environment_info import (
+        EnvironmentHardwareConfig as EnvironmentHardwareConfig,
+    )
+    from nshtrainer.util._environment_info import (
+        EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
+    )
+    from nshtrainer.util._environment_info import (
+        EnvironmentLSFInformationConfig as EnvironmentLSFInformationConfig,
+    )
+    from nshtrainer.util._environment_info import (
+        EnvironmentPackageConfig as EnvironmentPackageConfig,
+    )
+    from nshtrainer.util._environment_info import (
+        EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
+    )
+    from nshtrainer.util._environment_info import (
+        EnvironmentSnapshotConfig as EnvironmentSnapshotConfig,
+    )
+    from nshtrainer.util._environment_info import (
+        GitRepositoryConfig as GitRepositoryConfig,
+    )
+    from nshtrainer.util.config import DTypeConfig as DTypeConfig
+    from nshtrainer.util.config import DurationConfig as DurationConfig
+    from nshtrainer.util.config import EpochsConfig as EpochsConfig
+    from nshtrainer.util.config import StepsConfig as StepsConfig
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "DTypeConfig":
+            return importlib.import_module("nshtrainer.util.config").DTypeConfig
+        if name == "EnvironmentCUDAConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentCUDAConfig
+        if name == "EnvironmentClassInformationConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentClassInformationConfig
+        if name == "EnvironmentConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentConfig
+        if name == "EnvironmentGPUConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentGPUConfig
+        if name == "EnvironmentHardwareConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentHardwareConfig
+        if name == "EnvironmentLSFInformationConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentLSFInformationConfig
+        if name == "EnvironmentLinuxEnvironmentConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentLinuxEnvironmentConfig
+        if name == "EnvironmentPackageConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentPackageConfig
+        if name == "EnvironmentSLURMInformationConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentSLURMInformationConfig
+        if name == "EnvironmentSnapshotConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentSnapshotConfig
+        if name == "EpochsConfig":
+            return importlib.import_module("nshtrainer.util.config").EpochsConfig
+        if name == "GitRepositoryConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).GitRepositoryConfig
+        if name == "StepsConfig":
+            return importlib.import_module("nshtrainer.util.config").StepsConfig
+        if name == "DurationConfig":
+            return importlib.import_module("nshtrainer.util.config").DurationConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+# Submodule exports
+from . import _environment_info as _environment_info
+from . import config as config

nshtrainer/config/util/_environment_info/__init__.py

@@ -45,50 +45,50 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "EnvironmentLinuxEnvironmentConfig":
+        if name == "EnvironmentCUDAConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentLinuxEnvironmentConfig
-        if name == "EnvironmentLSFInformationConfig":
+            ).EnvironmentCUDAConfig
+        if name == "EnvironmentClassInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentLSFInformationConfig
-        if name == "EnvironmentGPUConfig":
+            ).EnvironmentClassInformationConfig
+        if name == "EnvironmentConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentGPUConfig
-        if name == "EnvironmentPackageConfig":
+            ).EnvironmentConfig
+        if name == "EnvironmentGPUConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentPackageConfig
+            ).EnvironmentGPUConfig
         if name == "EnvironmentHardwareConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).EnvironmentHardwareConfig
-        if name == "EnvironmentSnapshotConfig":
+        if name == "EnvironmentLSFInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentSnapshotConfig
-        if name == "EnvironmentClassInformationConfig":
+            ).EnvironmentLSFInformationConfig
+        if name == "EnvironmentLinuxEnvironmentConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentClassInformationConfig
-        if name == "GitRepositoryConfig":
+            ).EnvironmentLinuxEnvironmentConfig
+        if name == "EnvironmentPackageConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).GitRepositoryConfig
-        if name == "EnvironmentConfig":
+            ).EnvironmentPackageConfig
+        if name == "EnvironmentSLURMInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentConfig
-        if name == "EnvironmentCUDAConfig":
+            ).EnvironmentSLURMInformationConfig
+        if name == "EnvironmentSnapshotConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentCUDAConfig
-        if name == "EnvironmentSLURMInformationConfig":
+            ).EnvironmentSnapshotConfig
+        if name == "GitRepositoryConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentSLURMInformationConfig
+            ).GitRepositoryConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/util/config/__init__.py

@@ -18,12 +18,12 @@ else:
 
         if name in globals():
             return globals()[name]
+        if name == "DTypeConfig":
+            return importlib.import_module("nshtrainer.util.config").DTypeConfig
         if name == "EpochsConfig":
             return importlib.import_module("nshtrainer.util.config").EpochsConfig
         if name == "StepsConfig":
             return importlib.import_module("nshtrainer.util.config").StepsConfig
-        if name == "DTypeConfig":
-            return importlib.import_module("nshtrainer.util.config").DTypeConfig
         if name == "DurationConfig":
             return importlib.import_module("nshtrainer.util.config").DurationConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

nshtrainer/data/datamodule.py

@@ -1,7 +1,57 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from typing import Any, Generic, cast
+
+import nshconfig as C
 from lightning.pytorch import LightningDataModule
+from typing_extensions import Never, TypeVar, deprecated, override
+
+from ..model.mixins.callback import CallbackRegistrarModuleMixin
+from ..model.mixins.debug import _DebugModuleMixin
+
+THparams = TypeVar("THparams", bound=C.Config, infer_variance=True)
+
+
+class LightningDataModuleBase(
+    _DebugModuleMixin,
+    CallbackRegistrarModuleMixin,
+    LightningDataModule,
+    ABC,
+    Generic[THparams],
+):
+    @property
+    @override
+    def hparams(self) -> THparams:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return cast(THparams, super().hparams)
+
+    @property
+    @override
+    def hparams_initial(self):  # pyright: ignore[reportIncompatibleMethodOverride]
+        hparams = cast(THparams, super().hparams_initial)
+        return cast(Never, {"datamodule": hparams.model_dump(mode="json")})
+
+    @property
+    @deprecated("Use `hparams` instead")
+    def config(self):
+        return cast(Never, self.hparams)
+
+    @classmethod
+    @abstractmethod
+    def hparams_cls(cls) -> type[THparams]: ...
 
+    @override
+    def __init__(self, hparams: THparams | Mapping[str, Any]):
+        super().__init__()
 
-class LightningDataModuleBase(LightningDataModule):
-    pass
+        # Validate and save hyperparameters
+        hparams_cls = self.hparams_cls()
+        if isinstance(hparams, Mapping):
+            hparams = hparams_cls.model_validate(hparams)
+        elif not isinstance(hparams, hparams_cls):
+            raise TypeError(
+                f"Expected hparams to be either a Mapping or an instance of {hparams_cls}, got {type(hparams)}"
+            )
+        hparams = hparams.model_deep_validate()
+        self.save_hyperparameters(hparams)

nshtrainer/loggers/__init__.py

@@ -5,11 +5,12 @@ from typing import Annotated, TypeAlias
 import nshconfig as C
 
 from ._base import BaseLoggerConfig as BaseLoggerConfig
+from .actsave import ActSaveLoggerConfig as ActSaveLoggerConfig
 from .csv import CSVLoggerConfig as CSVLoggerConfig
 from .tensorboard import TensorboardLoggerConfig as TensorboardLoggerConfig
 from .wandb import WandbLoggerConfig as WandbLoggerConfig
 
 LoggerConfig: TypeAlias = Annotated[
-    CSVLoggerConfig | TensorboardLoggerConfig | WandbLoggerConfig,
+    CSVLoggerConfig | TensorboardLoggerConfig | WandbLoggerConfig | ActSaveLoggerConfig,
     C.Field(discriminator="name"),
 ]

nshtrainer/loggers/_base.py

@@ -7,7 +7,7 @@ import nshconfig as C
 from lightning.pytorch.loggers import Logger
 
 if TYPE_CHECKING:
-    from ..model import BaseConfig
+    from ..trainer._config import TrainerConfig
 
 
 class BaseLoggerConfig(C.Config, ABC):
@@ -21,8 +21,11 @@ class BaseLoggerConfig(C.Config, ABC):
     """Directory to save the logs to. If None, will use the default log directory for the trainer."""
 
     @abstractmethod
-    def create_logger(self, root_config: "BaseConfig") -> Logger | None: ...
+    def create_logger(self, trainer_config: TrainerConfig) -> Logger | None: ...
 
     def disable_(self):
         self.enabled = False
         return self
+
+    def __bool__(self):
+        return self.enabled

nshtrainer/loggers/actsave.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from argparse import Namespace
+from typing import Any, Literal
+
+import numpy as np
+from lightning.pytorch.loggers import Logger
+
+from ._base import BaseLoggerConfig
+
+
+class ActSaveLoggerConfig(BaseLoggerConfig):
+    name: Literal["actsave"] = "actsave"
+
+    def create_logger(self, trainer_config):
+        if not self.enabled:
+            return None
+
+        return ActSaveLogger()
+
+
+class ActSaveLogger(Logger):
+    @property
+    def name(self):
+        return None
+
+    @property
+    def version(self):
+        from nshutils import ActSave
+
+        if ActSave._saver is None:
+            return None
+
+        return ActSave._saver._id
+
+    @property
+    def save_dir(self):
+        from nshutils import ActSave
+
+        if ActSave._saver is None:
+            return None
+
+        return str(ActSave._saver._save_dir)
+
+    def log_hyperparams(
+        self,
+        params: dict[str, Any] | Namespace,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        from nshutils import ActSave
+
+        # Wrap the hparams as a object-dtype np array
+        return ActSave.save({"hyperparameters": np.array(params, dtype=object)})
+
+    def log_metrics(self, metrics: dict[str, float], step: int | None = None) -> None:
+        from nshutils import ActSave
+
+        ActSave.save({**metrics})

nshtrainer/loggers/csv.py

@@ -23,20 +23,20 @@ class CSVLoggerConfig(BaseLoggerConfig):
     """How often to flush logs to disk."""
 
     @override
-    def create_logger(self, root_config):
+    def create_logger(self, trainer_config):
         if not self.enabled:
             return None
 
         from lightning.pytorch.loggers.csv_logs import CSVLogger
 
-        save_dir = root_config.directory._resolve_log_directory_for_logger(
-            root_config.id,
+        save_dir = trainer_config.directory._resolve_log_directory_for_logger(
+            trainer_config.id,
             self,
         )
         return CSVLogger(
             save_dir=save_dir,
-            name=root_config.run_name,
-            version=root_config.id,
+            name=trainer_config.full_name,
+            version=trainer_config.id,
             prefix=self.prefix,
             flush_logs_every_n_steps=self.flush_logs_every_n_steps,
         )

nshtrainer/loggers/tensorboard.py

@@ -56,20 +56,20 @@ class TensorboardLoggerConfig(BaseLoggerConfig):
     """A string to put at the beginning of metric keys."""
 
     @override
-    def create_logger(self, root_config):
+    def create_logger(self, trainer_config):
         if not self.enabled:
             return None
 
         from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
 
-        save_dir = root_config.directory._resolve_log_directory_for_logger(
-            root_config.id,
+        save_dir = trainer_config.directory._resolve_log_directory_for_logger(
+            trainer_config.id,
             self,
         )
         return TensorBoardLogger(
             save_dir=save_dir,
-            name=root_config.run_name,
-            version=root_config.id,
+            name=trainer_config.full_name,
+            version=trainer_config.id,
             log_graph=self.log_graph,
             default_hp_metric=self.default_hp_metric,
         )