Fixed run metadata file not being saved in the W&B artifact

CHANGELOG.md

@@ -55,6 +55,8 @@ to communicate reliably.
   of model it is.
 - Made sure that LoRA layers are trainable.
 - Limiting W&B job names to 128 characters.
+- Fixed run metadata file not being saved in the W&B artifact, meaning the run package
+  version wasn't being determined correctly.
 
 
 ## [2.0.0] - 2025-06-16

nip/parameters/version.py

@@ -131,6 +131,9 @@ def convert_hyper_param_dict(hyper_param_dict: dict) -> dict:
         hyper_param_dict = conversion_func(hyper_param_dict)
         dict_version = next_version
 
+    # Finally, set the _package_version field to the current version
+    hyper_param_dict["_package_version"] = version_tuple_to_string(current_version)
+
     return hyper_param_dict
 
 

nip/trainers/trainer_base.py

@@ -176,6 +176,7 @@ def save_checkpoint(self, log: bool = True):
                 CHECKPOINT_STATE_ARTIFACT_TYPE,
             )
             artifact.add_file(self.checkpoint_state_path)
+            artifact.add_file(self.checkpoint_metadata_path)
             artifact.add_file(self.checkpoint_params_path)
             self.settings.wandb_run.log_artifact(artifact)
 
@@ -328,6 +329,15 @@ def _load_state_dict_from_active_run_checkpoint(self) -> dict:
                 metadata = json.load(f)
                 package_name = metadata.get("package_name", None)
                 package_version = metadata.get("package_version", None)
+
+        # Fall back to reading the hyper parameters file, which should contain the
+        # package name and version
+        elif self.checkpoint_params_path.exists():
+            with open(self.checkpoint_params_path, "r") as f:
+                hyper_params = json.load(f)
+                package_name = hyper_params.get("_package_name", None)
+                package_version = hyper_params.get("_package_version", None)
+
         else:
             package_name = None
             package_version = None