Setting tags during model logging mlflow

Question

I am logging the model using

mlflow.sklearn.log_model(model, "my-model")

and I want to set tags to the model during logging, I checked that this method does not allow to set tags, there is a mlflow.set_tags() method but it is tagging the run not the model.

Does anyone know how to tag the model during logging?

Thank you!

Answer 1

When using mlflow.sklearn.log_model you work with the experiment registry which is run-focused so only experiments and runs can be described and tagged.

If you want to set tags on models, you need to work with the model registry.

The solution I would recommend is to register the model when logging using registered_model_name (there are more fine-grained ways, too) and use MLFlowClient API to set custom properties (like tags) of the already registered model.

Here is a working example:

import mlflow
from mlflow.client import MlflowClient

mlflow.set_tracking_uri('http://0.0.0.0:5000')

experiment_name = 'test_mlflow'
try:
    experiment_id = mlflow.create_experiment(experiment_name)
except:
    experiment_id = mlflow.get_experiment_by_name(experiment_name).experiment_id

from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_score

with mlflow.start_run(experiment_id = experiment_id):
    # log performance and register the model
    X, y = load_iris(return_X_y=True)
    params = {"C": 0.1, "random_state": 42}
    mlflow.log_params(params)
    lr = LogisticRegression(**params).fit(X, y)
    y_pred = lr.predict(X)
    mlflow.log_metric("accuracy", accuracy_score(y, y_pred))
    mlflow.sklearn.log_model(lr, 
        artifact_path="models", 
        registered_model_name='test-model'
    )
    # set extra tags on the model
    client = MlflowClient(mlflow.get_tracking_uri())
    model_info = client.get_latest_versions('test-model')[0]
    client.set_model_version_tag(
        name='test-model',
        version=model_info.version,
        key='task',
        value='regression'
    )

Here is the illustration

See also this excellent documentation of MLFlow Client.