Skip to main content

Model Registry Fundamentals: Version, Stage, Deploy

The model registry is the single source of truth for all trained models in your organization. Instead of models scattered across folders, emails, and local disks, the registry centralizes them with versioning, metadata, and a workflow for promoting models from development to production. In this article, you will learn how to register models, manage versions, transition models through stages (Staging, Production, Archived), and deploy them safely.

What Is the Model Registry?​

The model registry is a database that tracks:

  • Model versions: Every trained model gets a version number (1, 2, 3, ...).
  • Stages: Each version is in a stage (None, Staging, Production, or Archived). Only one version per stage at a time.
  • Metadata: Description, creation date, last modified date, and custom tags.
  • Lineage: Which experiment and run created the model.

The workflow is simple: train a model, register it as version 1 (stage None), move it to Staging for testing, then to Production when ready. If the new version is buggy, revert to the previous production version. If a version is old, archive it.

Registering Your First Model​

To register a model, you must first log it in an experiment run (article 2 covers this). Then, register it from the MLflow UI or programmatically.

Step 1: Train and log a model

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 sklearn.metrics import accuracy_score

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

mlflow.set_experiment("Iris Classification")
mlflow.start_run(run_name="rf_v1")

# Train and log
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
accuracy = accuracy_score(y_test, model.predict(X_test))

mlflow.log_param("n_estimators", 100)
mlflow.log_metric("accuracy", accuracy)
mlflow.sklearn.log_model(model, "model")

mlflow.end_run()

print(f"Model logged. Accuracy: {accuracy:.3f}")

Step 2: Register the model

In the MLflow UI, find the run, click the logged model, and click "Register Model". Enter a name like iris_classifier. This creates version 1 of the iris_classifier model.

Or, programmatically:

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Register the model from the run
run_id = "run_id_from_above"
model_name = "iris_classifier"
model_uri = f"runs:/{run_id}/model"

model_version = mlflow.register_model(model_uri, model_name)
print(f"Registered {model_name} version {model_version.version}")

Managing Versions and Stages​

A model name can have many versions. Each version is in a stage.

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Transition version 1 to Staging
client.transition_model_version_stage(
    name="iris_classifier",
    version=1,
    stage="Staging"
)

# Later, after testing, transition to Production
client.transition_model_version_stage(
    name="iris_classifier",
    version=1,
    stage="Production"
)

# If version 2 is ready, move version 1 to Archived
client.transition_model_version_stage(
    name="iris_classifier",
    version=1,
    stage="Archived"
)

At any point, you can query which version is in which stage:

# Get all versions of a model
versions = client.search_model_versions(f"name='{model_name}'")
for v in versions:
    print(f"Version {v.version}: {v.current_stage}")

# Get the production version
prod_version = client.get_model_version(
    name="iris_classifier",
    version="production"  # You can query by stage name
)
print(f"Production version: {prod_version.version}")

Loading and Serving a Model from the Registry​

The key benefit of the model registry is easy loading. Instead of passing a file path, you reference the model by name and stage.

import mlflow.pyfunc

# Load the production version
model_uri = "models:/iris_classifier/Production"
model = mlflow.pyfunc.load_model(model_uri)

# Make predictions
predictions = model.predict(new_data)

This single line loads whichever version is currently in the Production stage. If you promote a new version to Production, the next time your application restarts, it automatically loads the new model. No code changes needed.

For serving models over an API, MLflow provides a built-in server:

mlflow models serve -m "models:/iris_classifier/Production" --port 5001

This launches an HTTP server. Send predictions via:

curl -X POST "http://localhost:5001/invocations" \
  -H "Content-Type: application/json" \
  -d '{"dataframe_split": {"columns": ["SepalLength", "SepalWidth", "PetalLength", "PetalWidth"], "data": [[5.1, 3.5, 1.4, 0.2]]}}'

Setting Model Metadata and Descriptions​

Document your models. The registry stores descriptions, aliases, and custom tags.

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Update the model description
client.update_model_version(
    name="iris_classifier",
    version=1,
    description="Random Forest trained on Iris dataset. Achieves 96% accuracy on test set."
)

# Add custom tags for organization
client.set_model_version_tag(
    name="iris_classifier",
    version=1,
    key="team",
    value="data_science"
)

client.set_model_version_tag(
    name="iris_classifier",
    version=1,
    key="approved_by",
    value="[email protected]"
)

# Retrieve the version and inspect metadata
version = client.get_model_version(name="iris_classifier", version=1)
print(f"Description: {version.description}")
print(f"Tags: {version.tags}")

These tags and descriptions appear in the MLflow UI and help teams understand which model to use, who approved it, and why.

A Complete Workflow: Train, Register, Promote, Deploy​

Here's a realistic scenario: you train a new model, stage it for testing, and promote it to production after validation.

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import accuracy_score
from mlflow.tracking import MlflowClient

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

mlflow.set_experiment("Iris Classification - Production Pipeline")
client = MlflowClient()

# 1. Train the new model
mlflow.start_run(run_name="rf_v2_candidate")

model = RandomForestClassifier(n_estimators=150, random_state=42)
model.fit(X_train, y_train)

# Evaluate thoroughly
test_accuracy = accuracy_score(y_test, model.predict(X_test))
cv_scores = cross_val_score(model, X_train, y_train, cv=5)

mlflow.log_param("n_estimators", 150)
mlflow.log_metric("test_accuracy", test_accuracy)
mlflow.log_metric("cv_mean_accuracy", cv_scores.mean())
mlflow.log_metric("cv_std_accuracy", cv_scores.std())

mlflow.sklearn.log_model(model, "model")
run_id = mlflow.active_run().info.run_id
mlflow.end_run()

print(f"Trained model. Test accuracy: {test_accuracy:.3f}")

# 2. Register the model
model_uri = f"runs:/{run_id}/model"
model_name = "iris_classifier"

# Check if model exists; if not, register it
model_versions = client.search_model_versions(f"name='{model_name}'")
if not model_versions:
    version_info = mlflow.register_model(model_uri, model_name)
    new_version = version_info.version
else:
    new_version = max([int(v.version) for v in model_versions]) + 1
    # Register as new version
    version_info = mlflow.register_model(model_uri, model_name)
    new_version = version_info.version

print(f"Registered {model_name} version {new_version}")

# 3. Transition to Staging for testing
client.transition_model_version_stage(
    name=model_name,
    version=new_version,
    stage="Staging"
)
print(f"Transitioned to Staging")

# 4. Simulate testing (in reality, run integration tests, A/B tests, etc.)
staging_model = mlflow.pyfunc.load_model(f"models:/{model_name}/Staging")
staging_predictions = staging_model.predict(X_test)
staging_accuracy = accuracy_score(y_test, staging_predictions)
print(f"Staging model accuracy (validation): {staging_accuracy:.3f}")

# 5. If validation passes, promote to Production
if staging_accuracy >= 0.95:
    # Archive the old production version (if any)
    prod_versions = client.search_model_versions(
        f"name='{model_name}' and current_stage='Production'"
    )
    for pv in prod_versions:
        client.transition_model_version_stage(
            name=model_name,
            version=pv.version,
            stage="Archived"
        )
    
    # Promote new version to Production
    client.transition_model_version_stage(
        name=model_name,
        version=new_version,
        stage="Production"
    )
    print(f"Promoted version {new_version} to Production!")
else:
    print(f"Staging accuracy {staging_accuracy:.3f} below threshold. Not promoting.")

This workflow is idempotent and safe: it clearly separates concerns (training, staging, promotion) and makes rollback trivial (just transition the old production version back to Production).

Key Takeaways​

  • The model registry centralizes all trained models with versioning, staging, and metadata.
  • Register models from experiment runs; each model name can have multiple versions.
  • Stages (None, Staging, Production, Archived) control which model is "active" for different purposes.
  • Load models by name and stage: models:/model_name/Production automatically loads the current production version.
  • Set descriptions and tags to document who trained, approved, and deployed each version.
  • A staged workflow (train -> stage -> validate -> promote -> archive) is safer than directly deploying new models.

Frequently Asked Questions​

Can I register the same run's model under multiple names?​

Yes. You can call mlflow.register_model(model_uri, different_name) multiple times with the same model_uri. This creates separate model entries, which is useful if the same model serves multiple use cases.

What happens if I transition a model to a stage that already has a model?​

The previous model is moved to stage None (archived implicitly). Only one model per stage at a time.

Can I roll back a production model if something goes wrong?​

Yes. Transition the old production version back to Production and transition the buggy version to Archived (or Staging for investigation). This takes seconds.

How do I delete a model or version?​

Use client.delete_model_version(name, version) to delete a version or client.delete_registered_model(name) to delete all versions of a model. Deletion is permanent and should be done carefully.

Can I use the model registry without MLflow's UI?​

Yes. The Python client (MlflowClient) provides full functionality. Many teams automate the entire workflow with scripts and cron jobs.

Further Reading​