How to beginner to intermediate · 4 min read

How to use MLflow model registry

Q: How to use MLflow model registry

Use the MLflow Model Registry to centrally manage your machine learning models' lifecycle by registering models, tracking versions, and transitioning stages like Staging or Production. Interact with the registry via the mlflow Python API to register models, update metadata, and deploy models efficiently.

Quick answer

Use the MLflow Model Registry to centrally manage your machine learning models' lifecycle by registering models, tracking versions, and transitioning stages like Staging or Production. Interact with the registry via the mlflow Python API to register models, update metadata, and deploy models efficiently.

PREREQUISITES

Python 3.8+
pip install mlflow>=2.0
Access to an MLflow tracking server (local or remote)

Setup

Install MLflow and set up your environment to use the model registry. You need a running MLflow tracking server, which can be local or remote. Set the MLFLOW_TRACKING_URI environment variable to point to your server.

bash

pip install mlflow

# Example to set tracking URI in bash
export MLFLOW_TRACKING_URI=http://localhost:5000

Step by step

This example shows how to train a simple model, log it to MLflow, register it in the model registry, and transition its stage.

python

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

# Load data
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, random_state=42)

# Train model
model = RandomForestClassifier(n_estimators=10)
model.fit(X_train, y_train)

# Start MLflow run
with mlflow.start_run() as run:
    # Log model
    mlflow.sklearn.log_model(model, "model")
    run_id = run.info.run_id

# Register the model
model_uri = f"runs:/{run_id}/model"
model_name = "IrisRandomForest"

# Create or update model in registry
result = mlflow.register_model(model_uri, model_name)

# Transition model stage to 'Staging'
client = mlflow.tracking.MlflowClient()
client.transition_model_version_stage(
    name=model_name,
    version=result.version,
    stage="Staging"
)

print(f"Model {model_name} version {result.version} registered and moved to Staging stage.")

output

Model IrisRandomForest version 1 registered and moved to Staging stage.

Common variations

You can automate stage transitions, archive old versions, or use async APIs. Also, MLflow supports different model flavors like TensorFlow, PyTorch, and custom models. Use mlflow..log_model() accordingly.

python

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Archive an old version
client.transition_model_version_stage(
    name="IrisRandomForest",
    version=1,
    stage="Archived"
)

# List all versions
versions = client.get_latest_versions("IrisRandomForest")
for v in versions:
    print(f"Version: {v.version}, Stage: {v.current_stage}")

output

Version: 1, Stage: Archived
Version: 2, Stage: Staging

Troubleshooting

If you get ModelNotFound, ensure the model name is correct and the tracking server is reachable.
If stage transition fails, check you have permission to modify the registry or that the version exists.
For connection issues, verify MLFLOW_TRACKING_URI is set properly.

✅

Key Takeaways

Use MLflow Model Registry to track and manage model versions centrally.
Register models after logging them to MLflow runs for lifecycle control.
Transition model stages (e.g., Staging, Production) to manage deployment readiness.
Use MlflowClient API for advanced registry operations like archiving and listing versions.

Verified 2026-04 · mlflow

Verify ↗