WIP: add ai-hub metrics to MLFlow

src/state.py

@@ -67,6 +67,18 @@ class CliStateStore:
     def set_latest_experiment_model_version(self, version: str) -> None:
         self.update(latest_experiment_model_version=version)
 
+    def get_aihub_model_provenance(self, model_id: str) -> dict[str, Any]:
+        provenance = self._aihub_model_provenance(self.read())
+        value = provenance.get(model_id, {})
+        return dict(value) if isinstance(value, dict) else {}
+
+    def update_aihub_model_provenance(self, model_id: str, provenance: dict[str, Any]) -> None:
+        state = self.read()
+        model_provenance = self._aihub_model_provenance(state)
+        model_provenance[model_id] = provenance
+        state["aihub_model_provenance"] = model_provenance
+        self._write(state)
+
     def _write(self, state: dict[str, Any]) -> None:
         with open(self.path, "w") as f:
             json.dump(state, f, indent=2)
@@ -75,6 +87,10 @@ class CliStateStore:
         value = state.get("training_jobs", {})
         return dict(value) if isinstance(value, dict) else {}
 
+    def _aihub_model_provenance(self, state: dict[str, Any]) -> dict[str, Any]:
+        value = state.get("aihub_model_provenance", {})
+        return dict(value) if isinstance(value, dict) else {}
+
 
 def store(config_path: str) -> CliStateStore:
     config_dir = str(Path(config_path).parent)