Report metrics from training stage

Author: orazve

examples/kge-distmult-nations.py

@@ -192,7 +192,7 @@ def inspect_graph(G):
 
     model_name = "dummyModelName_" + str(time.time())
 
-    gds.kge.model.train(
+    res = gds.kge.model.train(
         G_train,
         model_name=model_name,
         scoring_function="distmult",
@@ -202,6 +202,7 @@ def inspect_graph(G):
         epochs_per_val=5,
         split_ratios={"TRAIN": 0.8, "VALID": 0.1, "TEST": 0.1},
     )
+    print(res["metrics"])
 
     predict_result = gds.kge.model.predict(
         model_name=model_name,
@@ -216,7 +217,6 @@ def inspect_graph(G):
 
     print(predict_result.to_string())
 
-    print(predict_result.to_string())
     for index, row in predict_result.iterrows():
         h = row["head"]
         r = row["rel"]

graphdatascience/model/kge_runner.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import time
@@ -114,7 +115,12 @@ def train(
 
         self._wait_for_job(job_id)
 
-        return Series({"status": "finished"})
+        return Series(
+            {
+                "status": "finished",
+                "metrics": self._get_metrics(config["user_name"], config["task_config"]["modelname"], job_id),
+            }
+        )
 
     @client_only_endpoint("gds.kge.model")
     def predict(
@@ -205,6 +211,25 @@ def _stream_results(self, user_name: str, model_name: str, job_id: str) -> DataF
         os.remove(res_file_name)
         return df
 
+    def _get_metrics(self, user_name: str, model_name: str, job_id: str) -> DataFrame:
+        res = requests.get(
+            f"{self._compute_cluster_web_uri}/internal/fetch-model-metadata",
+            params={"user_name": user_name, "modelname": model_name},
+        )
+        res.raise_for_status()
+
+        res_file_name = f"metadata_{job_id}.json"
+
+        with open(res_file_name, mode="wb+") as f:
+            f.write(res.content)
+
+        with open(res_file_name, mode="r") as f:
+            metadata = json.load(f)
+
+        os.remove(res_file_name)
+
+        return metadata["metrics"]
+
     def _start_job(self, config: Dict[str, Any]) -> str:
         url = f"{self._compute_cluster_web_uri}/api/machine-learning/start"
         res = requests.post(url, json=config)