nshtrainer 0.16.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

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/_checkpoint/metadata.py

@@ -40,7 +40,7 @@ class CheckpointMetadata(C.Config):
     metrics: dict[str, Any]
     environment: EnvironmentConfig
 
-    hparams: dict[str, Any] | None
+    hparams: Any
 
     @classmethod
     def from_file(cls, path: Path):
@@ -48,9 +48,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}"
             )
@@ -94,7 +92,7 @@ def _generate_checkpoint_metadata(
         training_time=training_time,
         metrics=metrics,
         environment=config.environment,
-        hparams=config.model_dump(mode="json"),
+        hparams=config.model_dump(),
     )
 
 
@@ -104,26 +102,26 @@ 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
     )
 
     # Write the metadata to the checkpoint directory
     try:
-        metadata_path.write_text(metadata.model_dump_json(indent=4))
-    except Exception as e:
-        log.warning(f"Failed to write metadata to {checkpoint_path}: {e}")
+        metadata_path.write_text(metadata.model_dump_json(indent=4), encoding="utf-8")
+    except Exception:
+        log.exception(f"Failed to write metadata to {checkpoint_path}")
     else:
         log.debug(f"Checkpoint metadata written to {checkpoint_path}")
 
 
 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:
-        log.warning(f"Failed to remove {path}: {e}")
+    except Exception:
+        log.exception(f"Failed to remove {path}")
     else:
         log.debug(f"Removed {path}")
 
@@ -133,8 +131,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)
@@ -146,8 +144,8 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
             # on Windows, special permissions are required to create symbolic links as a regular user
             # fall back to copying the file
             shutil.copy(path, linked_path)
-    except Exception as e:
-        log.warning(f"Failed to link {path} to {linked_path}: {e}")
+    except Exception:
+        log.exception(f"Failed to link {path} to {linked_path}")
     else:
         log.debug(f"Linked {path} to {linked_path}")
 

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/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/util/_environment_info.py

@@ -429,12 +429,12 @@ class EnvironmentPackageConfig(C.Config):
                         version=clean_version,
                         path=Path(str(f)) if (f := dist.locate_file("")) else None,
                         summary=metadata["Summary"] if "Summary" in metadata else None,
-                        author=metadata["Author"] if "Summary" in metadata else None,
-                        license=metadata["License"] if "Summary" in metadata else None,
+                        author=metadata["Author"] if "Author" in metadata else None,
+                        license=metadata["License"] if "License" in metadata else None,
                         requires=requires,
                     )
-                except Exception as e:
-                    log.warning(f"Error processing package {dist.name}: {str(e)}")
+                except Exception:
+                    log.exception(f"Error processing package {dist.name}")
 
         except ImportError:
             log.warning(
@@ -672,8 +672,8 @@ class GitRepositoryConfig(C.Config):
             draft.is_dirty = repo.is_dirty()
         except git.InvalidGitRepositoryError:
             draft.is_git_repo = False
-        except Exception as e:
-            log.warning(f"Failed to get Git repository information: {e}")
+        except Exception:
+            log.exception("Failed to get Git repository information")
             draft.is_git_repo = None
 
         return draft.finalize()

{nshtrainer-0.16.0.dist-info → nshtrainer-0.17.0.dist-info}/METADATA

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

{nshtrainer-0.16.0.dist-info → nshtrainer-0.17.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
-nshtrainer/_checkpoint/loader.py,sha256=DSaNR8194kWon4O1svslNsCcN_8vlyLbF0LNCPfUpzI,13789
-nshtrainer/_checkpoint/metadata.py,sha256=onmetLp5eKbA86abq1PTkwAOO7bWj7Pa1EGUjl2TEjQ,5153
+nshtrainer/_checkpoint/loader.py,sha256=myFObRsPdb8jBncMK73vjr5FDJIfKhF86Ec_kSjXtwg,13837
+nshtrainer/_checkpoint/metadata.py,sha256=ilq3Wz9QmILxwpy5qYm7kphhtqPUYAOh8D-OXc5SqSc,5131
 nshtrainer/_checkpoint/saver.py,sha256=DkbCH0YeOJ71m32vAARiQdGBf0hvwwdoAV8LOFGy-0Y,1428
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
 nshtrainer/callbacks/__init__.py,sha256=4qocBDzQbLLhhbIEfvbA3SQB_Dy9ZJH7keMwPay-ZS8,2359
@@ -8,7 +8,7 @@ nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHm
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,3555
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=g-3zIthupERKqWZQw-A_busQPaPRkto6iHBV-M7nK1Y,527
-nshtrainer/callbacks/checkpoint/_base.py,sha256=yBUZ63P4a3m5jviIUGxBd2WXq_TigpKt3w4KdzqrzLs,6094
+nshtrainer/callbacks/checkpoint/_base.py,sha256=YT_V-oihO9iB4ETl46CGYTCQjIYl-CpV7TMViTn07Lk,6144
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=DJiLo7NDzd-lp-O3v7Cv8WejyjXPV_6_RmfltKO9fvE,2165
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CqB_8Xv32rtpLCaEEPi6DbRZm4ph5TWS-LfqIHXUIUA,1097
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
@@ -76,14 +76,14 @@ nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=F2tkHogbMAa5U7335sm77sZBkjEDa5v46XbJCH9Mg6c,2167
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
-nshtrainer/trainer/trainer.py,sha256=M97phnALfG18VxkMLoDr5AKFf4UaPBdc6S2BghdBtas,17103
-nshtrainer/util/_environment_info.py,sha256=Nmhls6u5rtMWbeDLGCjEk58efUutc_ONSUg3fs59TSI,24210
+nshtrainer/trainer/trainer.py,sha256=jIqiNrq1I0f5pQP7lHshtgjCAYfpoWPoqwS74LHU9iM,17148
+nshtrainer/util/_environment_info.py,sha256=gIdq9TJgzGCdcVzZxjHcwYasJ_HmEGVHbvE-KJVVtWs,24187
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
 nshtrainer/util/environment.py,sha256=AeW_kLl-N70wmb6L_JLz1wRj0kA70xs6RCmc9iUqczE,4159
 nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.16.0.dist-info/METADATA,sha256=O8wSotKFVUpCPLCLibdBmC6kUICBBcVMXsThhwuMmKI,885
-nshtrainer-0.16.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.16.0.dist-info/RECORD,,
+nshtrainer-0.17.0.dist-info/METADATA,sha256=6qrwuGiaXKIaFBpYwvuqBjyDhkLX4bNjU1ojJeTLsQE,885
+nshtrainer-0.17.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.17.0.dist-info/RECORD,,