nshtrainer 0.44.0__py3-none-any.whl → 1.0.0b9__py3-none-any.whl

nshtrainer/trainer/trainer.py

@@ -2,90 +2,120 @@ from __future__ import annotations
 
 import logging
 import os
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, cast
+from typing import IO, TYPE_CHECKING, Any, cast
 
 import torch
 from lightning.fabric.plugins.environments.lsf import LSFEnvironment
 from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
+from lightning.fabric.utilities.cloud_io import _load as pl_load
+from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
 from lightning.pytorch import LightningModule
 from lightning.pytorch import Trainer as LightningTrainer
 from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.core.saving import (
+    _default_map_location,
+    load_hparams_from_tags_csv,
+    load_hparams_from_yaml,
+)
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.trainer.states import TrainerFn
+from lightning.pytorch.utilities.migration import pl_legacy_patch
+from lightning.pytorch.utilities.migration.utils import _pl_migrate_checkpoint
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
-from typing_extensions import Unpack, assert_never, override
+from typing_extensions import Never, Unpack, assert_never, deprecated, override
 
 from .._checkpoint.metadata import _write_checkpoint_metadata
 from ..callbacks.base import resolve_all_callbacks
+from ..util._environment_info import EnvironmentConfig
 from ..util.bf16 import is_bf16_supported_no_emulation
 from ._config import (
     AcceleratorConfigProtocol,
     LightningTrainerKwargs,
     StrategyConfigProtocol,
+    TrainerConfig,
 )
 from ._runtime_callback import RuntimeTrackerCallback, Stage
 from .checkpoint_connector import _CheckpointConnector
 from .signal_connector import _SignalConnector
 
-if TYPE_CHECKING:
-    from ..model.config import BaseConfig
-
 log = logging.getLogger(__name__)
 
 
 class Trainer(LightningTrainer):
+    CHECKPOINT_HYPER_PARAMS_KEY = "trainer_hyper_parameters"
+
+    @property
+    def hparams(self) -> TrainerConfig:
+        """The collection of hyperparameters saved with :meth:`save_hyperparameters`. It is mutable by the user. For
+        the frozen set of initial hyperparameters, use :attr:`hparams_initial`.
+
+        Returns:
+            Mutable hyperparameters dictionary
+
+        """
+        return self._hparams
+
+    @property
+    @deprecated("Use `hparams` instead")
+    def config(self):
+        return cast(Never, self.hparams)
+
+    @classmethod
+    def hparams_cls(cls):
+        return TrainerConfig
+
     @classmethod
-    def _pre_init(cls, config: "BaseConfig"):
-        if (precision := config.trainer.set_float32_matmul_precision) is not None:
+    def _pre_init(cls, hparams: TrainerConfig):
+        if (precision := hparams.set_float32_matmul_precision) is not None:
             torch.set_float32_matmul_precision(precision)
 
     @classmethod
     def _update_kwargs(
         cls,
-        config: "BaseConfig",
+        hparams: TrainerConfig,
         kwargs_ctor: LightningTrainerKwargs,
     ):
         kwargs: LightningTrainerKwargs = {
-            "deterministic": config.trainer.reproducibility.deterministic,
-            "fast_dev_run": config.trainer.fast_dev_run,
-            "max_epochs": config.trainer.max_epochs,
-            "min_epochs": config.trainer.min_epochs,
-            "max_steps": config.trainer.max_steps,
-            "min_steps": config.trainer.min_steps,
-            "max_time": config.trainer.max_time,
-            "limit_train_batches": config.trainer.limit_train_batches,
-            "limit_val_batches": config.trainer.limit_val_batches,
-            "limit_test_batches": config.trainer.limit_test_batches,
-            "limit_predict_batches": config.trainer.limit_predict_batches,
-            "overfit_batches": config.trainer.overfit_batches,
-            "val_check_interval": config.trainer.val_check_interval,
-            "num_sanity_val_steps": config.trainer.num_sanity_val_steps,
-            "log_every_n_steps": config.trainer.log_every_n_steps,
-            "inference_mode": config.trainer.inference_mode,
+            "deterministic": hparams.reproducibility.deterministic,
+            "fast_dev_run": hparams.fast_dev_run,
+            "max_epochs": hparams.max_epochs,
+            "min_epochs": hparams.min_epochs,
+            "max_steps": hparams.max_steps,
+            "min_steps": hparams.min_steps,
+            "max_time": hparams.max_time,
+            "limit_train_batches": hparams.limit_train_batches,
+            "limit_val_batches": hparams.limit_val_batches,
+            "limit_test_batches": hparams.limit_test_batches,
+            "limit_predict_batches": hparams.limit_predict_batches,
+            "overfit_batches": hparams.overfit_batches,
+            "val_check_interval": hparams.val_check_interval,
+            "num_sanity_val_steps": hparams.num_sanity_val_steps,
+            "log_every_n_steps": hparams.log_every_n_steps,
+            "inference_mode": hparams.inference_mode,
             "callbacks": [],
             "plugins": [],
             "logger": [],
             # Moved to `lightning_kwargs`:
-            # "enable_checkpointing": config.trainer.enable_checkpointing,
-            # "accelerator": config.trainer.accelerator,
-            # "strategy": config.trainer.strategy,
-            # "num_nodes": config.trainer.num_nodes,
-            # "precision": config.trainer.precision,
-            # "logger": config.trainer.logging.enabled,
-            # "log_every_n_steps": config.trainer.log_every_n_steps,
-            # "enable_progress_bar": config.trainer.enable_progress_bar,
-            # "enable_model_summary": config.trainer.enable_model_summary,
-            # "accumulate_grad_batches": config.trainer.accumulate_grad_batches,
-            # "benchmark": config.trainer.benchmark,
-            # "use_distributed_sampler": config.trainer.use_distributed_sampler,
-            # "detect_anomaly": config.trainer.detect_anomaly,
-            # "barebones": config.trainer.barebones,
-            # "plugins": config.trainer.plugins,
-            # "sync_batchnorm": config.trainer.sync_batchnorm,
-            # "reload_dataloaders_every_n_epochs": config.trainer.reload_dataloaders_every_n_epochs,
+            # "enable_checkpointing": hparams.enable_checkpointing,
+            # "accelerator": hparams.accelerator,
+            # "strategy": hparams.strategy,
+            # "num_nodes": hparams.num_nodes,
+            # "precision": hparams.precision,
+            # "logger": hparams.logging.enabled,
+            # "log_every_n_steps": hparams.log_every_n_steps,
+            # "enable_progress_bar": hparams.enable_progress_bar,
+            # "enable_model_summary": hparams.enable_model_summary,
+            # "accumulate_grad_batches": hparams.accumulate_grad_batches,
+            # "benchmark": hparams.benchmark,
+            # "use_distributed_sampler": hparams.use_distributed_sampler,
+            # "detect_anomaly": hparams.detect_anomaly,
+            # "barebones": hparams.barebones,
+            # "plugins": hparams.plugins,
+            # "sync_batchnorm": hparams.sync_batchnorm,
+            # "reload_dataloaders_every_n_epochs": hparams.reload_dataloaders_every_n_epochs,
         }
 
         def _update_key(key: str, new_value: Any):
@@ -115,20 +145,22 @@ class Trainer(LightningTrainer):
                 _update_key(key, value)
 
         # Set `default_root_dir` if `auto_set_default_root_dir` is enabled.
-        if config.trainer.auto_set_default_root_dir:
+        if hparams.auto_set_default_root_dir:
             if kwargs.get("default_root_dir"):
                 raise ValueError(
-                    "You have set `config.trainer.default_root_dir`. "
+                    "You have set `hparams.default_root_dir`. "
                     "But we are trying to set it automatically. "
-                    "Please use `config.directory.base` rather than `config.trainer.default_root_dir`. "
-                    "If you want to set it manually, please set `config.trainer.auto_set_default_root_dir=False`."
+                    "Please use `hparams.directory.base` rather than `hparams.default_root_dir`. "
+                    "If you want to set it manually, please set `hparams.auto_set_default_root_dir=False`."
                 )
 
             _update_kwargs(
-                default_root_dir=config.directory.resolve_run_root_directory(config.id)
+                default_root_dir=hparams.directory.resolve_run_root_directory(
+                    hparams.id
+                )
             )
 
-        if (devices_input := config.trainer.devices) is not None:
+        if (devices_input := hparams.devices) is not None:
             match devices_input:
                 case "all":
                     devices = -1
@@ -141,22 +173,20 @@ class Trainer(LightningTrainer):
 
             _update_kwargs(devices=devices)
 
-        if (
-            use_distributed_sampler := config.trainer.use_distributed_sampler
-        ) is not None:
+        if (use_distributed_sampler := hparams.use_distributed_sampler) is not None:
             _update_kwargs(use_distributed_sampler=use_distributed_sampler)
 
-        if (accelerator := config.trainer.accelerator) is not None:
+        if (accelerator := hparams.accelerator) is not None:
             if isinstance(accelerator, AcceleratorConfigProtocol):
                 accelerator = accelerator.create_accelerator()
             _update_kwargs(accelerator=accelerator)
 
-        if (strategy := config.trainer.strategy) is not None:
+        if (strategy := hparams.strategy) is not None:
             if isinstance(strategy, StrategyConfigProtocol):
                 strategy = strategy.create_strategy()
             _update_kwargs(strategy=strategy)
 
-        if (precision := config.trainer.precision) is not None:
+        if (precision := hparams.precision) is not None:
             resolved_precision: _PRECISION_INPUT
             match precision:
                 case "64-true" | "32-true" | "bf16-mixed":
@@ -184,11 +214,11 @@ class Trainer(LightningTrainer):
 
             _update_kwargs(precision=resolved_precision)
 
-        if (detect_anomaly := config.trainer.detect_anomaly) is not None:
+        if (detect_anomaly := hparams.detect_anomaly) is not None:
             _update_kwargs(detect_anomaly=detect_anomaly)
 
         if (
-            grad_clip_config := config.trainer.optimizer.gradient_clipping
+            grad_clip_config := hparams.optimizer.gradient_clipping
         ) is not None and grad_clip_config.enabled:
             # kwargs["gradient_clip_algorithm"] = grad_clip_config.algorithm
             # kwargs["gradient_clip_val"] = grad_clip_config.value
@@ -197,9 +227,9 @@ class Trainer(LightningTrainer):
                 gradient_clip_val=grad_clip_config.value,
             )
 
-        if profiler_config := config.trainer.profiler:
-            if (profiler := profiler_config.create_profiler(config)) is None:
-                log.warning(f"Profiler config {profiler_config=} returned None.")
+        if profiler_config := hparams.profiler:
+            if (profiler := profiler_config.create_profiler(hparams)) is None:
+                log.warning(f"Profiler hparams {profiler_config=} returned None.")
             # Make sure that the profiler is an instance of `Profiler`.
             elif not isinstance(profiler, Profiler):
                 raise ValueError(f"{profiler=} is not an instance of `{Profiler}`.")
@@ -208,23 +238,29 @@ class Trainer(LightningTrainer):
             else:
                 _update_kwargs(profiler=profiler)
 
-        if callbacks := resolve_all_callbacks(config):
+        if callbacks := resolve_all_callbacks(hparams):
             _update_kwargs(callbacks=callbacks)
 
-        if plugin_configs := config.trainer.plugins:
+        if plugin_configs := hparams.plugins:
             _update_kwargs(
                 plugins=[
                     plugin_config.create_plugin() for plugin_config in plugin_configs
                 ]
             )
 
-        if not config.trainer.logging.enabled:
-            log.critical(f"Disabling logger because {config.trainer.logging.enabled=}.")
+        if not hparams.logging.enabled:
+            log.critical(f"Disabling logger because {hparams.logging.enabled=}.")
             kwargs["logger"] = False
         else:
-            _update_kwargs(logger=list(config.trainer.logging.create_loggers(config)))
+            _update_kwargs(
+                logger=[
+                    logger
+                    for logger in hparams.logging.create_loggers(hparams)
+                    if logger is not None
+                ]
+            )
 
-        if config.trainer.auto_determine_num_nodes:
+        if hparams.auto_determine_num_nodes:
             # When num_nodes is auto, we need to detect the number of nodes.
             if SLURMEnvironment.detect():
                 if (num_nodes := os.environ.get("SLURM_NNODES")) is not None:
@@ -243,12 +279,12 @@ class Trainer(LightningTrainer):
                 _update_kwargs(num_nodes=num_nodes)
             else:
                 log.info(
-                    "config.trainer.auto_determine_num_nodes ignored because no SLURM or LSF detected."
+                    "hparams.auto_determine_num_nodes ignored because no SLURM or LSF detected."
                 )
 
         # Update the kwargs with the additional trainer kwargs
-        _update_kwargs(**cast(Any, config.trainer.additional_lightning_kwargs))
-        _update_kwargs(**config.trainer.lightning_kwargs)
+        _update_kwargs(**cast(Any, hparams.additional_lightning_kwargs))
+        _update_kwargs(**hparams.lightning_kwargs)
         _update_kwargs(**kwargs_ctor)
 
         return kwargs
@@ -259,15 +295,29 @@ class Trainer(LightningTrainer):
     @override
     def __init__(
         self,
-        config: "BaseConfig",
+        hparams: TrainerConfig | Mapping[str, Any],
         /,
         **kwargs: Unpack[LightningTrainerKwargs],
     ):
-        self._pre_init(config)
+        # Validate the hparams.
+        hparams_cls = Trainer.hparams_cls()
+        if isinstance(hparams, Mapping):
+            hparams = hparams_cls.model_validate(hparams)
+        elif not isinstance(hparams, hparams_cls):
+            raise ValueError(
+                f"Trainer hparams must either be an instance of {hparams_cls} or a mapping. "
+                f"Got {type(hparams)=} instead."
+            )
+        hparams = hparams.model_deep_validate()
+
+        self._pre_init(hparams)
 
-        kwargs = self._update_kwargs(config, kwargs)
+        kwargs = self._update_kwargs(hparams, kwargs)
         log.critical(f"LightningTrainer.__init__ with {kwargs=}.")
 
+        self._hparams = hparams
+        self.debug = self.hparams.debug
+
         super().__init__(**kwargs)
 
         # Add our own start time callback to measure the start time.
@@ -285,7 +335,7 @@ class Trainer(LightningTrainer):
         log.critical(f"LightningTrainer log directory: {self.log_dir}.")
 
         # Set the checkpoint
-        if (ckpt_path := config.trainer.ckpt_path) is not None:
+        if (ckpt_path := hparams.ckpt_path) is not None:
             self.ckpt_path = str(Path(ckpt_path).resolve().absolute())
 
     def __runtime_tracker(self):
@@ -372,7 +422,16 @@ class Trainer(LightningTrainer):
         We patch the `Trainer._run` method to throw if gradient clipping is enabled
         and `model.automatic_optimization` is False.
         """
+        # Save the current environment information
+        datamodule = getattr(self, "datamodule", None)
+        self.hparams.environment = EnvironmentConfig.from_current_environment(
+            self.hparams, model, datamodule
+        )
 
+        # If gradient clipping is enabled, then we need to make sure that
+        # `model.automatic_optimization` is enabled. Otherwise, gradient clipping
+        # is not actually going to do anything, as we expect the user to manually
+        # call `optimizer.step()` and `optimizer.zero_grad()`.
         if not model.automatic_optimization and (
             self.gradient_clip_val is not None
             or self.gradient_clip_algorithm is not None
@@ -401,12 +460,10 @@ class Trainer(LightningTrainer):
 
         # Save the checkpoint metadata
         metadata_path = None
-        lm = self._base_module
-        root_config = cast("BaseConfig", lm.hparams)
-        if root_config.trainer.save_checkpoint_metadata and self.is_global_zero:
+        if self.hparams.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
             if (
-                metadata_path := _write_checkpoint_metadata(self, lm, filepath)
+                metadata_path := _write_checkpoint_metadata(self, filepath)
             ) is not None:
                 written_files.append(metadata_path)
 
@@ -414,3 +471,64 @@ class Trainer(LightningTrainer):
         from .. import _callback
 
         _callback._call_on_checkpoint_saved(self, filepath, metadata_path)
+
+    @classmethod
+    def load_from_checkpoint(
+        cls,
+        checkpoint_path: _PATH | IO,
+        map_location: _MAP_LOCATION_TYPE = None,
+        hparams_file: _PATH | None = None,
+        **kwargs: Any,
+    ):
+        loaded = _load_from_checkpoint(
+            checkpoint_path,
+            map_location=map_location,
+            hparams_file=hparams_file,
+            **kwargs,
+        )
+        return loaded
+
+
+def _load_from_checkpoint(
+    checkpoint_path: _PATH | IO,
+    map_location: _MAP_LOCATION_TYPE = None,
+    hparams_file: _PATH | None = None,
+    **kwargs: Any,
+):
+    map_location = map_location or _default_map_location
+    with pl_legacy_patch():
+        checkpoint = pl_load(checkpoint_path, map_location=map_location)
+
+    # convert legacy checkpoints to the new format
+    checkpoint = _pl_migrate_checkpoint(
+        checkpoint,
+        checkpoint_path=(
+            checkpoint_path if isinstance(checkpoint_path, (str, Path)) else None
+        ),
+    )
+
+    if hparams_file is not None:
+        extension = str(hparams_file).split(".")[-1]
+        if extension.lower() == "csv":
+            hparams = load_hparams_from_tags_csv(hparams_file)
+        elif extension.lower() in ("yml", "yaml"):
+            hparams = load_hparams_from_yaml(hparams_file)
+        else:
+            raise ValueError(".csv, .yml or .yaml is required for `hparams_file`")
+
+        # overwrite hparams by the given file
+        checkpoint[Trainer.CHECKPOINT_HYPER_PARAMS_KEY] = hparams
+
+    # for past checkpoint need to add the new key
+    checkpoint.setdefault(Trainer.CHECKPOINT_HYPER_PARAMS_KEY, {})
+    # override the hparams with values that were passed in
+    checkpoint[Trainer.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
+
+    # load the hparams
+    hparams = Trainer.hparams_cls().model_validate(
+        checkpoint[Trainer.CHECKPOINT_HYPER_PARAMS_KEY]
+    )
+
+    # create the trainer
+    trainer = Trainer(hparams)
+    return trainer

nshtrainer/util/_environment_info.py

@@ -15,14 +15,14 @@ from typing import TYPE_CHECKING, Any, cast
 import nshconfig as C
 import psutil
 import torch
+from lightning.pytorch import LightningDataModule, LightningModule
 from packaging import version
 from typing_extensions import Self
 
 from .slurm import parse_slurm_node_list
 
 if TYPE_CHECKING:
-    from ..model.base import LightningModuleBase
-    from ..model.config import BaseConfig
+    from ..trainer._config import TrainerConfig
 
 
 log = logging.getLogger(__name__)
@@ -708,6 +708,9 @@ class EnvironmentConfig(C.Config):
     model: EnvironmentClassInformationConfig | None = None
     """The Lightning module class information."""
 
+    datamodule: EnvironmentClassInformationConfig | None = None
+    """The Lightning data module class information."""
+
     linux: EnvironmentLinuxEnvironmentConfig | None = None
     """The Linux environment information."""
 
@@ -768,8 +771,9 @@ class EnvironmentConfig(C.Config):
     @classmethod
     def from_current_environment(
         cls,
-        root_config: "BaseConfig",
-        model: "LightningModuleBase",
+        trainer_config: TrainerConfig,
+        model: LightningModule,
+        datamodule: LightningDataModule | None = None,
     ):
         draft = cls.draft()
         draft.cwd = Path(os.getcwd())
@@ -777,23 +781,27 @@ class EnvironmentConfig(C.Config):
         draft.python_path = [Path(path) for path in sys.path]
         draft.python_version = sys.version
         draft.python_packages = EnvironmentPackageConfig.from_current_environment()
-        draft.config = EnvironmentClassInformationConfig.from_instance(root_config)
+        draft.config = EnvironmentClassInformationConfig.from_instance(trainer_config)
         draft.model = EnvironmentClassInformationConfig.from_instance(model)
+        if datamodule is not None:
+            draft.datamodule = EnvironmentClassInformationConfig.from_instance(
+                datamodule
+            )
         draft.linux = EnvironmentLinuxEnvironmentConfig.from_current_environment()
         draft.hardware = EnvironmentHardwareConfig.from_current_environment()
         draft.slurm = EnvironmentSLURMInformationConfig.from_current_environment()
         draft.lsf = EnvironmentLSFInformationConfig.from_current_environment()
-        draft.base_dir = root_config.directory.resolve_run_root_directory(
-            root_config.id
+        draft.base_dir = trainer_config.directory.resolve_run_root_directory(
+            trainer_config.id
         )
-        draft.log_dir = root_config.directory.resolve_subdirectory(
-            root_config.id, "log"
+        draft.log_dir = trainer_config.directory.resolve_subdirectory(
+            trainer_config.id, "log"
         )
-        draft.checkpoint_dir = root_config.directory.resolve_subdirectory(
-            root_config.id, "checkpoint"
+        draft.checkpoint_dir = trainer_config.directory.resolve_subdirectory(
+            trainer_config.id, "checkpoint"
         )
-        draft.stdio_dir = root_config.directory.resolve_subdirectory(
-            root_config.id, "stdio"
+        draft.stdio_dir = trainer_config.directory.resolve_subdirectory(
+            trainer_config.id, "stdio"
         )
         draft.seed = (
             int(seed_str) if (seed_str := os.environ.get("PL_GLOBAL_SEED")) else None

nshtrainer/util/config/dtype.py

@@ -9,7 +9,7 @@ from typing_extensions import assert_never
 from ..bf16 import is_bf16_supported_no_emulation
 
 if TYPE_CHECKING:
-    from ...model.base import BaseConfig
+    from ...trainer._config import TrainerConfig
 
 DTypeName: TypeAlias = Literal[
     "float32",
@@ -59,8 +59,8 @@ class DTypeConfig(C.Config):
     """The name of the dtype."""
 
     @classmethod
-    def from_base_config(cls, config: "BaseConfig"):
-        if (precision := config.trainer.precision) is None:
+    def from_trainer_config(cls, trainer_config: TrainerConfig):
+        if (precision := trainer_config.precision) is None:
             precision = "32-true"
 
         match precision:
@@ -79,7 +79,7 @@ class DTypeConfig(C.Config):
             case "64-true":
                 return cls(name="float64")
             case _:
-                assert_never(config.trainer.precision)
+                assert_never(trainer_config.precision)
 
     @property
     def torch_dtype(self):

nshtrainer/util/typing_utils.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 from typing_extensions import TypeVar
 
-TBase = TypeVar("TBase")
+TBase = TypeVar("TBase", infer_variance=True)
 
 
 def mixin_base_type(base_class: type[TBase]) -> type[TBase]:

{nshtrainer-0.44.0.dist-info → nshtrainer-1.0.0b9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.44.0
+Version: 1.0.0b9
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -15,7 +15,7 @@ Requires-Dist: huggingface-hub ; extra == "extra"
 Requires-Dist: lightning
 Requires-Dist: nshconfig
 Requires-Dist: nshrunner
-Requires-Dist: nshutils
+Requires-Dist: nshutils ; extra == "extra"
 Requires-Dist: numpy
 Requires-Dist: packaging
 Requires-Dist: psutil