MLflow’s Model Registry isn’t just a fancy version control system for your models; it’s a central nervous system for managing their lifecycle, acting as the single source of truth for what’s ready for production and what’s not.
Let’s see it in action. Imagine we’ve just trained a model and want to register it.
from mlflow import MlflowClient
from mlflow.models import infer_signature
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification
# Sample data and model training
X, y = make_classification(n_samples=100, n_features=4,
n_informative=2, n_redundant=0,
random_state=42)
data = pd.DataFrame(X, columns=[f'feature_{i}' for i in range(X.shape[1])])
data['target'] = y
model = RandomForestClassifier(random_state=42)
model.fit(data[['feature_0', 'feature_1', 'feature_2', 'feature_3']], data['target'])
# Log the model and get its run ID
with mlflow.start_run() as run:
signature = infer_signature(data[['feature_0', 'feature_1', 'feature_2', 'feature_3']], model.predict(data[['feature_0', 'feature_1', 'feature_2', 'feature_3']]))
mlflow.sklearn.log_model(
sk_model=model,
artifact_path="random_forest_model",
signature=signature,
registered_model_name="MyAwesomeClassifier"
)
run_id = run.info.run_id
model_uri = f"runs:/{run_id}/random_forest_model"
print(f"Model logged with run ID: {run_id}")
print(f"Model URI: {model_uri}")
After this runs, MLflow automatically creates a "MyAwesomeClassifier" registered model in the Model Registry, and the currently logged model becomes version 1 of it.
The Model Registry organizes models into "registered models," each with multiple "versions." A registered model is a container, and each version is an artifact snapshot of a trained model that you’ve decided to track. When you log a model with registered_model_name, MLflow handles this creation for you.
Here’s how you’d interact with it programmatically using the MlflowClient:
from mlflow import MlflowClient
client = MlflowClient()
# List all registered models
print("Registered Models:")
for rm in client.search_registered_models():
print(f"- {rm.name}")
# Get a specific registered model
model_name = "MyAwesomeClassifier"
registered_model = client.get_registered_model(model_name)
print(f"\nDetails for '{model_name}':")
print(f" Description: {registered_model.description}")
print(f" Created by: {registered_model.user_id}")
print(f" Latest Version: {registered_model.latest_version}")
# List versions of a specific model
print(f"\nVersions for '{model_name}':")
for mv in client.search_model_versions(f"name='{model_name}'"):
print(f"- Version: {mv.version}, Stage: {mv.current_stage}, Status: {mv.status}")
The core concept is stages. Models can exist in different stages, representing their readiness for production. The common stages are None (default, unassigned), Staging (ready for testing), Production (actively serving predictions), and Archived (no longer in use). You promote a model by transitioning its version to a higher stage.
# Transition a model version to Staging
client.transition_model_version_stage(
name=model_name,
version=registered_model.latest_version, # We'll transition the latest version
stage="Staging"
)
# Archive an older version (if any)
# Assuming version 1 is the latest we just staged
# You might archive version 0 if it existed and was previously Production
# For this example, let's assume we want to archive version 1 if it was already Staging
# If version 1 is the *only* version and we just put it in Staging,
# we wouldn't archive it. Let's simulate archiving an older, hypothetical version 0.
# In a real scenario, you'd check existing versions first.
try:
client.transition_model_version_stage(
name=model_name,
version="1", # Assuming version "1" exists and is not the latest
stage="Archived"
)
print("Archived version 1.")
except Exception as e:
print(f"Could not archive version 1: {e}") # Handle cases where version 1 might be the latest or doesn't exist.
# Verify the changes
print(f"\nUpdated Stages for '{model_name}':")
for mv in client.search_model_versions(f"name='{model_name}'"):
print(f"- Version: {mv.version}, Stage: {mv.current_stage}")
When you transition a model version to Production, MLflow clients (like mlflow.pyfunc.load_model) can directly load the model from that stage, abstracting away the need to know specific run IDs or artifact paths. This is crucial for CI/CD pipelines.
The "stage" is not just a label; it’s a contractual commitment. When a model version is in the Production stage, it implies that it has passed all necessary quality gates, testing, and is deemed fit to serve live traffic. MLflow enforces this by allowing only one model version to be in the Production stage at a time for a given registered model name. If you try to promote a new version to Production while another is already there, MLflow will prompt you to either archive the existing Production model or transition it to another stage. This prevents accidental overwrites of live models.
The next step in mastering the Model Registry is understanding how to deploy models from specific stages to serving endpoints.