train: Set model URI tag with absolute artifact path

Set the "training.output.model_uri" tag to the absolute path of the
corresponding artifact when logging a model for an MLflow run. This is
useful for downstream applications that need to fetch the model
generated by a specific run directly without having to reconstruct the
URI.

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

app/management/tracker_client.py

@@ -256,7 +256,10 @@ def save_model(self,
         else:
             model_manager.log_model(model_name, filepath)
 
-        return mlflow.get_artifact_uri(model_name)
+        artifact_uri = mlflow.get_artifact_uri(model_name)
+        mlflow.set_tag("training.output.model_uri", artifact_uri)
+
+        return artifact_uri
 
     @staticmethod
     def _get_experiment_id(experiment_name: str) -> str:

tests/app/monitoring/test_tracker_client.py

@@ -156,7 +156,13 @@ def test_save_model(mlflow_fixture):
     assert "artifacts/model_name" in artifact_uri
     model_manager.log_model.assert_called_once_with("model_name", "path/to/file.zip", "model_name")
     mlflow_client.set_model_version_tag.assert_called_once_with(name="model_name", version="1", key="validation_status", value="validation_status")
-    mlflow.set_tag.assert_called_once_with("training.output.package", "file.zip")
+    mlflow.set_tag.has_calls(
+        [
+            call("training.output.package", "file.zip"),
+            call("training.output.model_uri", artifact_uri),
+        ],
+        any_order=False,
+    )
 
 
 def test_save_model_local(mlflow_fixture):