1. What is model registry?

Docs

Model

An MLflow Model is created from an experiment or run that is logged with one of the model flavor's mlflow.<model_flavor>.log_model() methods.

Registered Model

A model that has been registered with the Model Registry, having a unique name, versions, aliases, tags, and other metadata.

2. How to Register?

There are three programmatic ways(3):

mlflow.<model_flavor>.log_model()
mlflow.register_model()
mlflow.MlflowClient.create_registered_model()

# 1 way
from sklearn.datasets import make_regression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split

import mlflow
import mlflow.sklearn
from mlflow.models import infer_signature

with mlflow.start_run() as run:
    X, y = make_regression(n_features=4, n_informative=2, random_state=0, shuffle=False)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    params = {"max_depth": 2, "random_state": 42}
    model = RandomForestRegressor(**params)
    model.fit(X_train, y_train)

    # Infer the model signature
    y_pred = model.predict(X_test)
    signature = infer_signature(X_test, y_pred)

    # Log parameters and metrics using the MLflow APIs
    mlflow.log_params(params)
    mlflow.log_metrics({"mse": mean_squared_error(y_test, y_pred)})

    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="sklearn-model",
        signature=signature,
        registered_model_name="sk-learn-random-forest-reg-model",
    )

# 2 way
result = mlflow.register_model(
    "runs:/d16076a3ec534311817565e6527539c0/sklearn-model", "sk-learn-random-forest-reg"
)


# 3 way
model_name = "sk-learn-random-forest-reg-model"
client = MlflowClient()
client.create_registered_model(model_name)
result = client.create_model_version(
    name=model_name, source="mlruns/0/d16076a3ec534311817565e6527539c0/artifacts/sklearn-model",
    run_id="d16076a3ec534311817565e6527539c0",
)

3. How to fetch the registered model?

You can fetch that model using mlflow.<model_flavor>.load_model(), or more generally, mlflow.pyfunc.load_model()

by specific model version
by model alias

By model version

import mlflow.pyfunc

model_name = "sk-learn-random-forest-reg-model"
model_version = 1

model = mlflow.pyfunc.load_model(model_uri=f"models:/{model_name}/{model_version}")

model.predict(data)

By model alias

import mlflow.pyfunc

model_name = "sk-learn-random-forest-reg-model"
alias = "champion"

champion_version = mlflow.pyfunc.load_model(f"models:/{model_name}@{alias}")

champion_version.predict(data)

4. Model Serving (Inference Server)

#!/usr/bin/env sh

# Set environment variable for the tracking URL where the Model Registry resides
export MLFLOW_TRACKING_URI=http://localhost:5000

# Serve the production model from the model registry
mlflow models serve -m "models:/sk-learn-random-forest-reg-model@champion"

제민욱 @minwook

MLflow::Model Registry