Skip to content

Commit ce8354f

Browse files
authored
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>
1 parent 5b5976b commit ce8354f

File tree

2 files changed

+11
-2
lines changed

2 files changed

+11
-2
lines changed

app/management/tracker_client.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -256,7 +256,10 @@ def save_model(self,
256256
else:
257257
model_manager.log_model(model_name, filepath)
258258

259-
return mlflow.get_artifact_uri(model_name)
259+
artifact_uri = mlflow.get_artifact_uri(model_name)
260+
mlflow.set_tag("training.output.model_uri", artifact_uri)
261+
262+
return artifact_uri
260263

261264
@staticmethod
262265
def _get_experiment_id(experiment_name: str) -> str:

tests/app/monitoring/test_tracker_client.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,13 @@ def test_save_model(mlflow_fixture):
156156
assert "artifacts/model_name" in artifact_uri
157157
model_manager.log_model.assert_called_once_with("model_name", "path/to/file.zip", "model_name")
158158
mlflow_client.set_model_version_tag.assert_called_once_with(name="model_name", version="1", key="validation_status", value="validation_status")
159-
mlflow.set_tag.assert_called_once_with("training.output.package", "file.zip")
159+
mlflow.set_tag.has_calls(
160+
[
161+
call("training.output.package", "file.zip"),
162+
call("training.output.model_uri", artifact_uri),
163+
],
164+
any_order=False,
165+
)
160166

161167

162168
def test_save_model_local(mlflow_fixture):

0 commit comments

Comments
 (0)