mlflow: Add model type to training run tags

phoevos · phoevos · commit 6aeaecb51407 · 2025-12-02T11:36:41.000Z
Signed-off-by: Phoevos Kalemkeris &lt;phoevos.kalemkeris@ucl.ac.uk&gt;
diff --git a/app/management/tracker_client.py b/app/management/tracker_client.py
@@ -574,6 +574,7 @@ def save_model(
 
         artifact_uri = mlflow.get_artifact_uri(model_name)
         mlflow.set_tag("training.output.model_uri", artifact_uri)
+        mlflow.set_tag("training.output.model_type", model_type)
 
         return artifact_uri
 
diff --git a/tests/app/monitoring/test_tracker_client.py b/tests/app/monitoring/test_tracker_client.py
@@ -182,6 +182,7 @@ def test_save_model(mlflow_fixture):
         [
             call("training.output.package", "file.zip"),
             call("training.output.model_uri", artifact_uri),
+            call("training.output.model_type", "model_type"),
         ],
         any_order=False,
     )

Original file line number	Diff line number	Diff line change
`@@ -182,6 +182,7 @@ def test_save_model(mlflow_fixture):`
`182`	`182`	`[`
`183`	`183`	`call("training.output.package", "file.zip"),`
`184`	`184`	`call("training.output.model_uri", artifact_uri),`
	`185`	`+ call("training.output.model_type", "model_type"),`
`185`	`186`	`],`
`186`	`187`	`any_order=False,`
`187`	`188`	`)`