nshtrainer 0.44.0__tar.gz → 1.0.0b9__tar.gz

{nshtrainer-0.44.0 → nshtrainer-1.0.0b9}/PKG-INFO

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

{nshtrainer-0.44.0 → nshtrainer-1.0.0b9}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.44.0"
+version = "1.0.0-beta9"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -9,7 +9,7 @@ readme = "README.md"
 python = "^3.10"
 nshrunner = "*"
 nshconfig = "*"
-nshutils = "*"
+nshutils = { version = "*", optional = true }
 psutil = "*"
 numpy = "*"
 torch = "*"
@@ -50,4 +50,11 @@ ignore = ["F722", "F821", "E731", "E741"]
 required-imports = ["from __future__ import annotations"]
 
 [tool.poetry.extras]
-extra = ["wrapt", "GitPython", "wandb", "tensorboard", "huggingface-hub"]
+extra = [
+    "wrapt",
+    "GitPython",
+    "wandb",
+    "tensorboard",
+    "huggingface-hub",
+    "nshutils",
+]

{nshtrainer-0.44.0 → nshtrainer-1.0.0b9}/src/nshtrainer/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from . import _experimental as _experimental
 from . import callbacks as callbacks
-from . import config as config
 from . import data as data
 from . import lr_scheduler as lr_scheduler
 from . import metrics as metrics
@@ -12,7 +11,11 @@ from . import optimizer as optimizer
 from . import profiler as profiler
 from .data import LightningDataModuleBase as LightningDataModuleBase
 from .metrics import MetricConfig as MetricConfig
-from .model import BaseConfig as BaseConfig
 from .model import LightningModuleBase as LightningModuleBase
-from .runner import Runner as Runner
 from .trainer import Trainer as Trainer
+from .trainer import TrainerConfig as TrainerConfig
+
+try:
+    from . import config as config
+except BaseException:
+    pass

nshtrainer-1.0.0b9/src/nshtrainer/_callback.py

@@ -0,0 +1,337 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+from lightning.pytorch.callbacks import Callback as _LightningCallback
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from torch.optim import Optimizer
+
+if TYPE_CHECKING:
+    from .model import LightningModuleBase
+    from .trainer import Trainer
+
+
+class NTCallbackBase(_LightningCallback):
+    def setup(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase, stage: str
+    ) -> None:
+        """Called when fit, validate, test, predict, or tune begins."""
+
+    def teardown(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase, stage: str
+    ) -> None:
+        """Called when fit, validate, test, predict, or tune ends."""
+
+    def on_fit_start(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Called when fit begins."""
+
+    def on_fit_end(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Called when fit ends."""
+
+    def on_sanity_check_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the validation sanity check starts."""
+
+    def on_sanity_check_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the validation sanity check ends."""
+
+    def on_train_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        """Called when the train batch begins."""
+
+    def on_train_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        """Called when the train batch ends.
+
+        Note:
+            The value ``outputs["loss"]`` here will be the normalized value w.r.t ``accumulate_grad_batches`` of the
+            loss returned from ``training_step``.
+
+        """
+
+    def on_train_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the train epoch begins."""
+
+    def on_train_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the train epoch ends.
+
+        To access all batch outputs at the end of the epoch, you can cache step outputs as an attribute of the
+        :class:`lightning.pytorch.core.LightningModule` and access them in this hook:
+
+        .. code-block:: python
+
+            class MyLightningModule(L.LightningModule):
+                 def __init__(self):
+                    super().__init__() # pyright: ignore[reportIncompatibleMethodOverride]
+                    self.training_step_outputs = []
+
+                 def training_step(self):
+                    loss = ... # pyright: ignore[reportIncompatibleMethodOverride]
+                    self.training_step_outputs.append(loss)
+                    return loss
+
+
+            class MyCallback(L.Callback):
+                 def on_train_epoch_end(self, trainer, pl_module):
+                    # do something with all training_step outputs, for example: # pyright: ignore[reportIncompatibleMethodOverride]
+                    epoch_mean = torch.stack(pl_module.training_step_outputs).mean()
+                    pl_module.log("training_epoch_mean", epoch_mean)
+                    # free up the memory
+                    pl_module.training_step_outputs.clear()
+
+        """
+
+    def on_validation_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the val epoch begins."""
+
+    def on_validation_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the val epoch ends."""
+
+    def on_test_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the test epoch begins."""
+
+    def on_test_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the test epoch ends."""
+
+    def on_predict_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the predict epoch begins."""
+
+    def on_predict_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the predict epoch ends."""
+
+    def on_validation_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Called when the validation batch begins."""
+
+    def on_validation_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Called when the validation batch ends."""
+
+    def on_test_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Called when the test batch begins."""
+
+    def on_test_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Called when the test batch ends."""
+
+    def on_predict_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Called when the predict batch begins."""
+
+    def on_predict_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        outputs: Any,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Called when the predict batch ends."""
+
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Called when the train begins."""
+
+    def on_train_end(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Called when the train ends."""
+
+    def on_validation_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the validation loop begins."""
+
+    def on_validation_end(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the validation loop ends."""
+
+    def on_test_start(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Called when the test begins."""
+
+    def on_test_end(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Called when the test ends."""
+
+    def on_predict_start(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called when the predict begins."""
+
+    def on_predict_end(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Called when predict ends."""
+
+    def on_exception(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        exception: BaseException,
+    ) -> None:
+        """Called when any trainer execution is interrupted by an exception."""
+
+    def state_dict(self) -> dict[str, Any]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Called when saving a checkpoint, implement to generate callback's ``state_dict``.
+
+        Returns:
+            A dictionary containing callback state.
+
+        """
+        return {}
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Called when loading a checkpoint, implement to reload callback state given callback's ``state_dict``.
+
+        Args:
+            state_dict: the callback state returned by ``state_dict``.
+
+        """
+        pass
+
+    def on_save_checkpoint(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        checkpoint: dict[str, Any],
+    ) -> None:
+        r"""Called when saving a checkpoint to give you a chance to store anything else you might want to save.
+
+        Args:
+            trainer: the current :class:`~lightning.pytorch.trainer.trainer.Trainer` instance.
+            pl_module: the current :class:`~lightning.pytorch.core.LightningModule` instance.
+            checkpoint: the checkpoint dictionary that will be saved.
+
+        """
+
+    def on_load_checkpoint(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        checkpoint: dict[str, Any],
+    ) -> None:
+        r"""Called when loading a model checkpoint, use to reload state.
+
+        Args:
+            trainer: the current :class:`~lightning.pytorch.trainer.trainer.Trainer` instance.
+            pl_module: the current :class:`~lightning.pytorch.core.LightningModule` instance.
+            checkpoint: the full checkpoint dictionary that got loaded by the Trainer.
+
+        """
+
+    def on_before_backward(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase, loss: torch.Tensor
+    ) -> None:
+        """Called before ``loss.backward()``."""
+
+    def on_after_backward(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, trainer: Trainer, pl_module: LightningModuleBase
+    ) -> None:
+        """Called after ``loss.backward()`` and before optimizers are stepped."""
+
+    def on_before_optimizer_step(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        optimizer: Optimizer,
+    ) -> None:
+        """Called before ``optimizer.step()``."""
+
+    def on_before_zero_grad(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        trainer: Trainer,
+        pl_module: LightningModuleBase,
+        optimizer: Optimizer,
+    ) -> None:
+        """Called before ``optimizer.zero_grad()``."""
+
+    def on_checkpoint_saved(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        ckpt_path: Path,
+        metadata_path: Path | None,
+        trainer: "Trainer",
+        pl_module: "LightningModuleBase",
+    ) -> None:
+        """Called after a checkpoint is saved."""
+        pass
+
+
+def _call_on_checkpoint_saved(
+    trainer: "Trainer",
+    ckpt_path: str | Path,
+    metadata_path: str | Path | None,
+):
+    ckpt_path = Path(ckpt_path)
+    metadata_path = Path(metadata_path) if metadata_path else None
+
+    for callback in trainer.callbacks:
+        if not isinstance(callback, NTCallbackBase):
+            continue
+
+        callback.on_checkpoint_saved(
+            ckpt_path,
+            metadata_path,
+            trainer,
+            trainer._base_module,
+        )

{nshtrainer-0.44.0 → nshtrainer-1.0.0b9}/src/nshtrainer/_checkpoint/loader.py

@@ -7,7 +7,6 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias, overload
 
 import nshconfig as C
-from lightning.pytorch import Trainer as LightningTrainer
 from lightning.pytorch.trainer.states import TrainerFn
 from typing_extensions import assert_never
 
@@ -15,7 +14,8 @@ from ..metrics._config import MetricConfig
 from .metadata import CheckpointMetadata
 
 if TYPE_CHECKING:
-    from ..model.config import BaseConfig
+    from ..trainer import Trainer
+    from ..trainer._config import TrainerConfig
 
 log = logging.getLogger(__name__)
 
@@ -228,22 +228,22 @@ class _CkptCandidate:
 @overload
 def _load_ckpt_meta(
     path: Path,
-    root_config: "BaseConfig",
+    trainer_config: TrainerConfig,
     on_error: Literal["warn"] = "warn",
 ) -> _CkptCandidate | None: ...
 @overload
 def _load_ckpt_meta(
     path: Path,
-    root_config: "BaseConfig",
+    trainer_config: TrainerConfig,
     on_error: Literal["raise"],
 ) -> _CkptCandidate: ...
 def _load_ckpt_meta(
     path: Path,
-    root_config: "BaseConfig",
+    trainer_config: TrainerConfig,
     on_error: Literal["warn", "raise"] = "warn",
 ):
     meta = CheckpointMetadata.from_file(path)
-    if root_config.id != meta.run_id:
+    if trainer_config.id != meta.run_id:
         error_msg = f"Skipping checkpoint {path} because it belongs to a different run"
         match on_error:
             case "warn":
@@ -256,16 +256,13 @@ def _load_ckpt_meta(
     return _CkptCandidate(meta, path)
 
 
-def _checkpoint_candidates(
-    root_config: "BaseConfig",
-    trainer: LightningTrainer,
-    *,
-    include_hpc: bool = True,
-):
+def _checkpoint_candidates(trainer: Trainer, *, include_hpc: bool = True):
     # Load the checkpoint directory, and throw if it doesn't exist.
     # This indicates a non-standard setup, and we don't want to guess
     # where the checkpoints are.
-    ckpt_dir = root_config.directory.resolve_subdirectory(root_config.id, "checkpoint")
+    ckpt_dir = trainer.hparams.directory.resolve_subdirectory(
+        trainer.hparams.id, "checkpoint"
+    )
     if not ckpt_dir.is_dir():
         raise FileNotFoundError(
             f"Checkpoint directory {ckpt_dir} not found. "
@@ -275,46 +272,40 @@ def _checkpoint_candidates(
     # Load all checkpoints in the directory.
     # We can do this by looking for metadata files.
     for path in ckpt_dir.glob(f"*{CheckpointMetadata.PATH_SUFFIX}"):
-        if (meta := _load_ckpt_meta(path, root_config)) is not None:
+        if (meta := _load_ckpt_meta(path, trainer.hparams)) is not None:
             yield meta
 
     # If we have a pre-empted checkpoint, load it
     if include_hpc and (hpc_path := trainer._checkpoint_connector._hpc_resume_path):
         hpc_meta_path = Path(hpc_path).with_suffix(CheckpointMetadata.PATH_SUFFIX)
-        if (meta := _load_ckpt_meta(hpc_meta_path, root_config)) is not None:
+        if (meta := _load_ckpt_meta(hpc_meta_path, trainer.hparams)) is not None:
             yield meta
 
 
 def _additional_candidates(
-    additional_candidates: Iterable[Path], root_config: "BaseConfig"
+    additional_candidates: Iterable[Path], trainer_config: TrainerConfig
 ):
     for path in additional_candidates:
         if (
             meta := _load_ckpt_meta(
-                path.with_suffix(CheckpointMetadata.PATH_SUFFIX), root_config
+                path.with_suffix(CheckpointMetadata.PATH_SUFFIX), trainer_config
             )
         ) is None:
             continue
         yield meta
 
 
-def _resolve_checkpoint(
-    config: CheckpointLoadingConfig,
-    root_config: "BaseConfig",
-    trainer: LightningTrainer,
-):
+def _resolve_checkpoint(config: CheckpointLoadingConfig, trainer: Trainer):
     # We lazily load the checkpoint candidates to avoid loading them
     # if they are not needed.
     _ckpt_candidates: list[_CkptCandidate] | None = None
 
     def ckpt_candidates():
-        nonlocal _ckpt_candidates, root_config, trainer
+        nonlocal _ckpt_candidates, trainer
 
         if _ckpt_candidates is None:
             _ckpt_candidates = list(
-                _checkpoint_candidates(
-                    root_config, trainer, include_hpc=config.include_hpc
-                )
+                _checkpoint_candidates(trainer, include_hpc=config.include_hpc)
             )
         return _ckpt_candidates
 
@@ -324,7 +315,7 @@ def _resolve_checkpoint(
             case UserProvidedPathCheckpointStrategyConfig():
                 meta = _load_ckpt_meta(
                     strategy.path.with_suffix(CheckpointMetadata.PATH_SUFFIX),
-                    root_config,
+                    trainer.hparams,
                     on_error=strategy.on_error,
                 )
                 if meta is None:
@@ -334,7 +325,7 @@ def _resolve_checkpoint(
                 candidates = [
                     *ckpt_candidates(),
                     *_additional_candidates(
-                        strategy.additional_candidates, root_config
+                        strategy.additional_candidates, trainer.hparams
                     ),
                 ]
                 if not candidates:
@@ -343,7 +334,9 @@ def _resolve_checkpoint(
                     )
                     continue
 
-                if (metric := strategy.metric or root_config.primary_metric) is None:
+                if (
+                    metric := strategy.metric or trainer.hparams.primary_metric
+                ) is None:
                     log.warning(
                         "No metric specified for `best` checkpoint strategy, "
                         "and no primary metric is set in the configuration. "
@@ -369,7 +362,7 @@ def _resolve_checkpoint(
                 candidates = [
                     *ckpt_candidates(),
                     *_additional_candidates(
-                        strategy.additional_candidates, root_config
+                        strategy.additional_candidates, trainer.hparams
                     ),
                 ]
                 if not candidates:

{nshtrainer-0.44.0 → nshtrainer-1.0.0b9}/src/nshtrainer/_checkpoint/metadata.py

@@ -5,7 +5,7 @@ import datetime
 import logging
 from collections.abc import Callable
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, cast
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import nshconfig as C
 import numpy as np
@@ -15,7 +15,6 @@ from ..util._environment_info import EnvironmentConfig
 from ..util.path import compute_file_checksum, try_symlink_or_copy
 
 if TYPE_CHECKING:
-    from ..model import BaseConfig, LightningModuleBase
     from ..trainer.trainer import Trainer
 
 log = logging.getLogger(__name__)
@@ -24,6 +23,19 @@ log = logging.getLogger(__name__)
 METADATA_PATH_SUFFIX = ".metadata.json"
 
 
+def _full_hparams_dict(trainer: Trainer):
+    hparams = {}
+    hparams["trainer"] = trainer.hparams.model_dump(mode="json")
+
+    if trainer.lightning_module is not None:
+        from ..model import LightningModuleBase
+
+        if isinstance(trainer.lightning_module, LightningModuleBase):
+            hparams["model"] = trainer.lightning_module.hparams.model_dump(mode="json")
+
+    return hparams
+
+
 class CheckpointMetadata(C.Config):
     PATH_SUFFIX: ClassVar[str] = METADATA_PATH_SUFFIX
 
@@ -59,8 +71,7 @@ class CheckpointMetadata(C.Config):
 
 
 def _generate_checkpoint_metadata(
-    config: "BaseConfig",
-    trainer: "Trainer",
+    trainer: Trainer,
     checkpoint_path: Path,
     metadata_path: Path,
 ):
@@ -84,9 +95,9 @@ def _generate_checkpoint_metadata(
         checkpoint_path=checkpoint_path.relative_to(metadata_path.parent),
         checkpoint_filename=checkpoint_path.name,
         checkpoint_checksum=compute_file_checksum(checkpoint_path),
-        run_id=config.id,
-        name=config.run_name,
-        project=config.project,
+        run_id=trainer.hparams.id,
+        name=trainer.hparams.full_name,
+        project=trainer.hparams.project,
         checkpoint_timestamp=checkpoint_timestamp,
         start_timestamp=start_timestamp.datetime
         if start_timestamp is not None
@@ -95,8 +106,8 @@ def _generate_checkpoint_metadata(
         global_step=trainer.global_step,
         training_time=training_time,
         metrics=metrics,
-        environment=config.environment,
-        hparams=config.model_dump(),
+        environment=trainer.hparams.environment,
+        hparams=_full_hparams_dict(trainer),
     )
 
 
@@ -104,16 +115,9 @@ def _metadata_path(checkpoint_path: Path):
     return checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
 
 
-def _write_checkpoint_metadata(
-    trainer: "Trainer",
-    model: "LightningModuleBase",
-    checkpoint_path: Path,
-):
-    config = cast("BaseConfig", model.config)
+def _write_checkpoint_metadata(trainer: Trainer, checkpoint_path: Path):
     metadata_path = _metadata_path(checkpoint_path)
-    metadata = _generate_checkpoint_metadata(
-        config, trainer, checkpoint_path, metadata_path
-    )
+    metadata = _generate_checkpoint_metadata(trainer, checkpoint_path, metadata_path)
 
     # Write the metadata to the checkpoint directory
     try:

nshtrainer-1.0.0b9/src/nshtrainer/_experimental/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations