train: Set model URI tag to absolute artifact path (#10)

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 <phoevos.kalemkeris@ucl.ac.uk>

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):