How to compare model versions and roll back to a previous model if the new one underperforms

Summary

• Store all model versions in MLflow with aliases (Champion, Challenger, Staging) for quick reference and comparison
• Test new model versions on production traffic using A/B testing or canary deployments; measure accuracy, latency, and business metrics
• Rollback instantly by updating the model alias: change the alias from the new version back to the previous version in Model Serving

How to compare model versions and roll back to a previous model if the new one underperforms

Comparing model versions means testing them on identical data and business metrics. Keep old versions registered in MLflow with clear aliases so you can instantly promote the better version. If a new model underperforms, rolling back takes seconds by updating an alias, avoiding extended outages.

How to organize model versions

Use MLflow aliases to mark each version's role. The Champion is production. The Challenger is a candidate replacement. Older versions stay accessible for reference or emergency rollback. This structure makes comparison and rollback straightforward.

Key capabilities to evaluate

Capability	Why it matters
Version history	Keeping all versions enables comparison and prevents data loss if a new model fails
Aliases	Champion/Challenger/Staging labels make version management intuitive without memorizing IDs
A/B testing	Comparing on real production traffic with real metrics (accuracy, revenue) is increased reliable vs offline testing
Fast rollback	Aliasing enables single-command rollbacks instead of redeployment or debugging
Metrics tracking	Logging metrics per version enables data-driven promotion decisions

How Databricks approaches model version comparison and rollback

1. Register and tag model versions
When training new models, register them with descriptive tags:

import mlflow
from datetime import datetime

with mlflow.start_run():
    # Train model
    model = train_fraud_detector()
    
    # Log with metadata
    mlflow.sklearn.log_model(model, "fraud_model")
    model_version = mlflow.register_model("runs:/<run_id>/fraud_model", "fraud_detector")
    
    # Tag with training date and status
    client = mlflow.tracking.MlflowClient()
    client.set_model_version_tag(
        name="fraud_detector",
        version=model_version.version,
        key="training_date",
        value=datetime.now().isoformat()
    )
    client.set_model_version_tag(
        name="fraud_detector",
        version=model_version.version,
        key="status",
        value="candidate"
    )
    client.set_model_version_tag(
        name="fraud_detector",
        version=model_version.version,
        key="notes",
        value="Retrained on last 6 months of data"
    )

2. Set aliases for quick reference
Use aliases to mark which version is which:

client = mlflow.tracking.MlflowClient()

    #Current production model
client.set_registered_model_alias(
    name="fraud_detector",
    alias="champion",
    version=42  # The version currently in production
)

    #Candidate for A/B testing
client.set_registered_model_alias(
    name="fraud_detector",
    alias="challenger",
    version=43  # The new model to test
)

    #Staging for internal testing
client.set_registered_model_alias(
    name="fraud_detector",
    alias="staging",
    version=44
)

3. Compare models on a test set
Evaluate all candidates on identical data:

from sklearn.metrics import accuracy_score, precision_score, recall_score, roc_auc_score

def evaluate_model(model_alias, X_test, y_test):
    model = mlflow.pyfunc.load_model(f"models:/fraud_detector@{model_alias}")
    predictions = model.predict(X_test)
    
    return {
        "alias": model_alias,
        "accuracy": accuracy_score(y_test, predictions.round()),
        "precision": precision_score(y_test, predictions.round()),
        "recall": recall_score(y_test, predictions.round()),
        "auc": roc_auc_score(y_test, predictions)
    }

    #Load test set
X_test, y_test = load_test_set()

    #Compare
results = []
for alias in ["champion", "challenger"]:
    metrics = evaluate_model(alias, X_test, y_test)
    results.append(metrics)
    print(f"{alias}: {metrics}")

    #Decide winner
challenger_metrics = results[1]
champion_metrics = results[0]

if challenger_metrics["auc"] > champion_metrics["auc"] + 0.01:
    print("Challenger wins! Promoting to champion.")
else:
    print("Champion remains superior. Keeping current.")

4. A/B test on production traffic
For high-stakes decisions, test on real traffic (covered in the A/B testing page). Route a small portion to Challenger and the majority to Champion:

from databricks.sdk import WorkspaceClient
import random

    #In your serving logic:
if random.random() < 0.1:
    model_version = "challenger"
else:
    model_version = "champion"

model = mlflow.pyfunc.load_model(f"models:/fraud_detector@{model_version}")
prediction = model.predict(features)

After sufficient traffic accumulates of production traffic, compare metrics between versions.
5. Rollback if needed
If the new model underperforms, rollback instantly by updating the alias:

client = mlflow.tracking.MlflowClient()

    #Rollback: send all traffic back to champion
client.set_registered_model_alias(
    name="fraud_detector",
    alias="champion",
    version=42  # Revert to the old version
)

print("Rollback complete. All traffic now uses version 42.")

This change takes effect immediately in Model Serving. No redeployment needed.
6. Investigate why new model failed
After rollback, debug the Challenger:

-- Compare predictions between champion and challenger
SELECT
    champion_pred,
    challenger_pred,
    true_label,
    COUNT(*) as count
FROM comparison_logs
WHERE champion_pred != challenger_pred
GROUP BY champion_pred, challenger_pred, true_label
ORDER BY count DESC

Were there systematic differences? Did the new model over-predict fraud? Under-predict? Use this analysis to inform retraining.

Use cases

• Production model upgrade: After A/B testing a new model, promote it to champion if it wins. Rollback if metrics degrade.
• Emergency rollback: A deployed model has a bug. Quickly revert to the previous version while investigating.
• Quarterly model refresh: Compare the incumbent with candidates retrained on recent data; promote the winner.

FAQs

How often should I compare versions?

Formally A/B test before any production promotion. Informally, benchmark new models against the current champion during development.

How long should a rollback take?

With aliases, instantly. Model Serving picks up alias changes instantly and routes new requests to the new version.

Can I run multiple challengers simultaneously?

Yes, but split traffic carefully. Three models (Champion + Challenger A + Challenger B) with equal splits need significantly more traffic for the same statistical power.

What if I don't have ground truth labels yet?

Use proxy labels (user feedback, clicks) or delayed labels. If unavailable, use domain experts to manually evaluate a sample or compare feature importance between versions.