How to Use MLflow Model Registry for Managing Models

Use mlflow.register_model() to add a model to the MLflow Model Registry, then manage versions and stages with mlflow_client. The registry helps track model versions, stage transitions, and deployment status in a central place.

Syntax

The MLflow Model Registry uses these main commands:

• mlflow.register_model(model_uri, name): Registers a model under a given name.
• mlflow_client.transition_model_version_stage(name, version, stage): Changes the stage of a model version (e.g., 'Staging', 'Production').
• mlflow_client.get_latest_versions(name, stages): Retrieves the latest model versions for specified stages.

Here, model_uri is the path or run URI of the saved model, name is the registry model name, version is the model version number, and stage is the lifecycle stage.

python

from mlflow import register_model
from mlflow.tracking import MlflowClient

# Register a model
model_uri = "runs:/<run_id>/model"
model_name = "MyModel"
model_version = register_model(model_uri, model_name)

# Create client to manage registry
client = MlflowClient()

# Transition model version stage
client.transition_model_version_stage(
    name=model_name,
    version=model_version.version,
    stage="Production"
)

# Get latest production model
latest_models = client.get_latest_versions(model_name, stages=["Production"])
print(latest_models)

Example

This example shows how to register a simple sklearn model, transition its stage, and fetch the production version.

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
from mlflow.tracking import MlflowClient

# 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 model
model_uri = f"runs:/{run_id}/model"
model_name = "IrisRFModel"
client = MlflowClient()
model_version = mlflow.register_model(model_uri, model_name)

# Transition to Production
client.transition_model_version_stage(
    name=model_name,
    version=model_version.version,
    stage="Production"
)

# Fetch production model info
prod_models = client.get_latest_versions(model_name, stages=["Production"])
for m in prod_models:
    print(f"Model version: {m.version}, stage: {m.current_stage}")

Output

Model version: 1, stage: Production

Common Pitfalls

• Not specifying a unique model name: This causes conflicts in the registry.
• Forgetting to transition model stages: Models stay in None stage and are not ready for deployment.
• Using incorrect model_uri: The URI must point to a valid logged model path.
• Not waiting for model registration to complete: Registration is asynchronous; querying immediately may fail.

python

from mlflow import register_model
from mlflow.tracking import MlflowClient

# Wrong: Using a wrong model URI
wrong_model_uri = "runs:/invalid_run/model"
try:
    register_model(wrong_model_uri, "TestModel")
except Exception as e:
    print(f"Error: {e}")

# Right: Use a valid run ID and model path
# (Assuming valid run_id and model logged)
valid_model_uri = "runs:/<valid_run_id>/model"
model_version = register_model(valid_model_uri, "TestModel")

client = MlflowClient()
client.transition_model_version_stage(
    name="TestModel",
    version=model_version.version,
    stage="Staging"
)

Output

Error: Model URI runs:/invalid_run/model does not exist

Quick Reference

MLflow Model Registry key commands:

Command	Description
mlflow.register_model(model_uri, name)	Register a new model version from a saved model URI.
MlflowClient.transition_model_version_stage(name, version, stage)	Change the stage of a model version (e.g., 'Staging', 'Production').
MlflowClient.get_latest_versions(name, stages)	Get latest model versions filtered by stages.
MlflowClient.delete_model_version(name, version)	Delete a specific model version.
MlflowClient.rename_model(name, new_name)	Rename a registered model.

Key Takeaways

Register models with mlflow.register_model using a valid model URI and unique name.

Use MlflowClient to manage model version stages like 'Staging' and 'Production'.

Always verify model URI correctness to avoid registration errors.

Model stage transitions enable clear deployment workflows and tracking.

Use get_latest_versions to fetch models ready for deployment.