nshtrainer 0.15.1__tar.gz → 0.16.1__tar.gz

{nshtrainer-0.15.1 → nshtrainer-0.16.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.15.1
+Version: 0.16.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.15.1 → nshtrainer-0.16.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.15.1"
+version = "0.16.1"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-0.15.1 → nshtrainer-0.16.1}/src/nshtrainer/_checkpoint/loader.py

@@ -10,7 +10,7 @@ from lightning.pytorch.trainer.states import TrainerFn
 from typing_extensions import assert_never
 
 from ..metrics._config import MetricConfig
-from .metadata import METADATA_PATH_SUFFIX, CheckpointMetadata
+from .metadata import CheckpointMetadata
 
 if TYPE_CHECKING:
     from ..model.config import BaseConfig
@@ -263,13 +263,13 @@ 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"*{METADATA_PATH_SUFFIX}"):
+    for path in ckpt_dir.glob(f"*{CheckpointMetadata.PATH_SUFFIX}"):
         if (meta := _load_ckpt_meta(path, root_config)) 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(METADATA_PATH_SUFFIX)
+        hpc_meta_path = Path(hpc_path).with_suffix(CheckpointMetadata.PATH_SUFFIX)
         if (meta := _load_ckpt_meta(hpc_meta_path, root_config)) is not None:
             yield meta
 
@@ -279,7 +279,9 @@ def _additional_candidates(
 ):
     for path in additional_candidates:
         if (
-            meta := _load_ckpt_meta(path.with_suffix(METADATA_PATH_SUFFIX), root_config)
+            meta := _load_ckpt_meta(
+                path.with_suffix(CheckpointMetadata.PATH_SUFFIX), root_config
+            )
         ) is None:
             continue
         yield meta
@@ -310,7 +312,7 @@ def _resolve_checkpoint(
         match strategy:
             case UserProvidedPathCheckpointStrategyConfig():
                 meta = _load_ckpt_meta(
-                    strategy.path.with_suffix(METADATA_PATH_SUFFIX),
+                    strategy.path.with_suffix(CheckpointMetadata.PATH_SUFFIX),
                     root_config,
                     on_error=strategy.on_error,
                 )

{nshtrainer-0.15.1 → nshtrainer-0.16.1}/src/nshtrainer/_checkpoint/metadata.py

@@ -10,8 +10,6 @@ import nshconfig as C
 import numpy as np
 import torch
 
-from ..util._environment_info import EnvironmentConfig
-
 if TYPE_CHECKING:
     from ..model import BaseConfig, LightningModuleBase
     from ..trainer.trainer import Trainer
@@ -38,7 +36,7 @@ class CheckpointMetadata(C.Config):
     global_step: int
     training_time: datetime.timedelta
     metrics: dict[str, Any]
-    environment: EnvironmentConfig
+    environment: dict[str, Any]
 
     hparams: dict[str, Any] | None
 
@@ -48,9 +46,7 @@ class CheckpointMetadata(C.Config):
 
     @classmethod
     def from_ckpt_path(cls, checkpoint_path: Path):
-        if not (
-            metadata_path := checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
-        ).exists():
+        if not (metadata_path := checkpoint_path.with_suffix(cls.PATH_SUFFIX)).exists():
             raise FileNotFoundError(
                 f"Metadata file not found for checkpoint: {checkpoint_path}"
             )
@@ -93,7 +89,7 @@ def _generate_checkpoint_metadata(
         global_step=trainer.global_step,
         training_time=training_time,
         metrics=metrics,
-        environment=config.environment,
+        environment=config.environment.model_dump(mode="json"),
         hparams=config.model_dump(mode="json"),
     )
 
@@ -104,7 +100,7 @@ def _write_checkpoint_metadata(
     checkpoint_path: Path,
 ):
     config = cast("BaseConfig", model.config)
-    metadata_path = checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
+    metadata_path = checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
     metadata = _generate_checkpoint_metadata(
         config, trainer, checkpoint_path, metadata_path
     )
@@ -119,7 +115,7 @@ def _write_checkpoint_metadata(
 
 
 def _remove_checkpoint_metadata(checkpoint_path: Path):
-    path = checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
+    path = checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
     try:
         path.unlink(missing_ok=True)
     except Exception as e:
@@ -133,8 +129,8 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
     _remove_checkpoint_metadata(linked_checkpoint_path)
 
     # Link the metadata files to the new checkpoint
-    path = checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
-    linked_path = linked_checkpoint_path.with_suffix(METADATA_PATH_SUFFIX)
+    path = checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
+    linked_path = linked_checkpoint_path.with_suffix(CheckpointMetadata.PATH_SUFFIX)
     try:
         try:
             # linked_path.symlink_to(path)

{nshtrainer-0.15.1 → nshtrainer-0.16.1}/src/nshtrainer/callbacks/checkpoint/_base.py

@@ -102,6 +102,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         metas = [
             CheckpointMetadata.from_file(p)
             for p in self.dirpath.glob(f"*{CheckpointMetadata.PATH_SUFFIX}")
+            if p.is_file() and not p.is_symlink()
         ]
 
         # Sort by the topk sort key

nshtrainer-0.16.1/src/nshtrainer/ll/snapshot.py

@@ -0,0 +1 @@
+from nshsnap import *  # pyright: ignore[reportWildcardImportFromLibrary] # noqa: F403

{nshtrainer-0.15.1 → nshtrainer-0.16.1}/src/nshtrainer/trainer/trainer.py

@@ -419,7 +419,8 @@ class Trainer(LightningTrainer):
 
         # Save the checkpoint metadata
         lm = self._base_module
-        if lm.config.trainer.save_checkpoint_metadata and self.is_global_zero:
+        hparams = cast(BaseConfig, lm.hparams)
+        if hparams.trainer.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
             _write_checkpoint_metadata(self, lm, filepath)
 

nshtrainer-0.15.1/src/nshtrainer/ll/snapshot.py

@@ -1 +0,0 @@
-from nshrunner.snapshot import *  # type: ignore # noqa: F403