updates to report_model_version.py for model card report

mlflow_tools/reports/report_model_version.py

@@ -2,7 +2,7 @@
 
 from mlflow_tools.client.http_client import MlflowHttpClient
 from mlflow_tools.common import mlflow_utils, model_download_utils
-from mlflow_tools.common.timestamp_utils import fmt_ts_millis
+from mlflow_tools.common.timestamp_utils import fmt_ts_millis, ts_now_fmt_utc
 from mlflow_tools.common.click_options import (
     opt_show_system_info,
     opt_format,
@@ -29,19 +29,21 @@ def build_report(model_name, version):
     model_uri = f"models:/{model_name}/{version}"
     model_info = _dump_mlflow_model.build(model_uri)
     model_artifact_path = model_download_utils.get_relative_model_path(vr["source"], vr["run_id"])
-    model_summary = _mk_model_summary(vr, run, model_artifact_path, model_info["model_info"])
+
+    mlflow_model = _mk_mlflow_model(vr, run, model_info)
+    model_summary = _mk_model_summary(vr, run, model_artifact_path, model_info["model_info"], mlflow_model)
 
     return {
         "model_summary": model_summary,
-        "mlflow_model": _mk_mlflow_model(vr, run, model_info),
+        "mlflow_model": mlflow_model,
         "registered_model_version": _mk_version_summary(vr),
         "registered_model": _mk_registered_model(model_name),
         "run": _mk_run_summary(run),
         "experiment": _mk_experiment(run["info"]["experiment_id"])
     }
 
 
-def mk_native_flavor_summary(model_info):
+def mk_native_flavor_summary(model_info, mlflow_model):
     """ 
     Make native flavor summary
     """
@@ -71,29 +73,33 @@ def _prune(flavor):
 
     return {
         "time_created": model_info.get("_utc_time_created"),
-        "native_flavor": flavor_summary
+        "native_flavor": flavor_summary,
+        "size": mlflow_model.get("model_artifacts_size")
     }
 
 
-def _mk_model_summary(vr, run, model_artifact_path, model_info):
+def _mk_model_summary(vr, run, model_artifact_path, model_info, mlflow_model):
     """
     Make top-level model summary
     """
-    native_flavor = mk_native_flavor_summary(model_info)
+    native_flavor = mk_native_flavor_summary(model_info, mlflow_model)
     info = run["info"]
     return {
         "general": {
             "user": utils.get_user(run),
-            "time_created": model_info.get("_utc_time_created"),
+            "report_time_created": ts_now_fmt_utc,
+            "model_time_created": model_info.get("_utc_time_created"),
             "tracking_server": str(http_client)
         },
         "registered_model": {
             "name": vr["name"],
-            "version": vr["version"]
+            "version": vr["version"],
+            "stage": vr["current_stage"]
         },
         "mlflow_model": {
             "model_name": model_artifact_path,
             "run_id": vr["run_id"],
+            "run_name": _get_run_name(run["data"]["tags"]),
             "experiment": {
                 "name": info["_experiment_name"],
                 "experiment_id": info["experiment_id"]
@@ -136,6 +142,7 @@ def _mk_run_summary(run):
     run = run.copy()
 
     info = run["info"]
+    info["run_name"] = _get_run_name(run["data"]["tags"])
     info["start_time"] = info["_start_time"]
     info["end_time"] = info["_end_time"]
     info["duration_seconds"] = info["_duration"]
@@ -154,6 +161,10 @@ def _mk_run_summary(run):
     return run
 
 
+def _get_run_name(tags):
+    return tags.get("mlflow.runName")
+
+
 def _mk_mlflow_model_sources(vr):
     """
     Represents a version's two URI sources for MLflow model: the run URI and the 'cached registry' model
@@ -223,7 +234,7 @@ def _mk_experiment(experiment_id):
     exp = rsp["experiment"]
     exp["creation_time"] =  fmt_ts_millis(exp.get("creation_time"))
     exp["last_update_time"] =  fmt_ts_millis(exp.get("last_update_time"))
-    exp["tags"] = mlflow_utils.mk_tags_dict(exp["tags"])
+    exp["tags"] = mlflow_utils.mk_tags_dict(exp.get("tags"))
     return exp