From a3a9e06ddc3c05297a78351066926f7a81bd17dc Mon Sep 17 00:00:00 2001
From: Sam Adam-Day <me@samadamday.com>
Date: Tue, 24 Jun 2025 09:48:04 +0100
Subject: [PATCH] Fixed run metadata file not being saved in the W&B artifact

---
 CHANGELOG.md                 |  2 ++
 nip/parameters/version.py    |  3 +++
 nip/trainers/trainer_base.py | 10 ++++++++++
 3 files changed, 15 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a7c4caa7..ba4e10fa 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/nip/parameters/version.py b/nip/parameters/version.py
index 7a401678..50e6bf39 100644
--- a/nip/parameters/version.py
+++ b/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
 
 
diff --git a/nip/trainers/trainer_base.py b/nip/trainers/trainer_base.py
index df3f3d42..b22a4e22 100644
--- a/nip/trainers/trainer_base.py
+++ b/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