nshtrainer 0.16.1__tar.gz → 0.17.1__tar.gz

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

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

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

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

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

@@ -10,6 +10,8 @@ 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
@@ -36,13 +38,13 @@ class CheckpointMetadata(C.Config):
     global_step: int
     training_time: datetime.timedelta
     metrics: dict[str, Any]
-    environment: dict[str, Any]
+    environment: EnvironmentConfig
 
-    hparams: dict[str, Any] | None
+    hparams: Any
 
     @classmethod
     def from_file(cls, path: Path):
-        return cls.model_validate_json(path.read_text())
+        return cls.model_validate_json(path.read_text(encoding="utf-8"))
 
     @classmethod
     def from_ckpt_path(cls, checkpoint_path: Path):
@@ -89,8 +91,8 @@ def _generate_checkpoint_metadata(
         global_step=trainer.global_step,
         training_time=training_time,
         metrics=metrics,
-        environment=config.environment.model_dump(mode="json"),
-        hparams=config.model_dump(mode="json"),
+        environment=config.environment,
+        hparams=config.model_dump(),
     )
 
 
@@ -107,9 +109,9 @@ def _write_checkpoint_metadata(
 
     # 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}")
 
@@ -118,8 +120,8 @@ def _remove_checkpoint_metadata(checkpoint_path: Path):
     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}")
 
@@ -142,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-0.16.1 → nshtrainer-0.17.1}/src/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()