Model Registry

Professional+

The Calabi ML Model Registry is a centralised store for managing the lifecycle of ML models from experimentation through production deployment. When a run produces a model you want to promote beyond the experiment, you register it in the registry. The registry maintains a full version history and supports a stage-based promotion workflow — Staging, Production, and Archived — to govern which model version is active in production at any given time.

---

Registry Concepts

Concept	Description
Registered Model	A named model in the registry (e.g., churn-predictor). Persists across runs and experiments.
Model Version	One specific model artifact promoted to the registry from a run. Each version is immutable once created.
Stage	The lifecycle phase of a version: None, Staging, Production, or Archived
Alias	A mutable pointer to a version (e.g., champion pointing to version 7)
Tags	Key-value metadata on registered models or specific versions

---

Registering a Model

From a Run (Python SDK)

import mlflow
import mlflow.sklearn

mlflow.set_tracking_uri("https://calabi.<your-domain>/mlflow")
mlflow.set_experiment("churn-prediction-v2")

with mlflow.start_run(run_name="gb_best_v1") as run:
    # ... train model ...
    mlflow.sklearn.log_model(
        sk_model=pipeline,
        artifact_path="model",
        registered_model_name="churn-predictor",   # Register immediately
    )

If churn-predictor does not exist in the registry, it is created automatically. If it already exists, a new version is appended.

From an Existing Run (After the Fact)

from mlflow.tracking import MlflowClient

client = MlflowClient(tracking_uri="https://calabi.<your-domain>/mlflow")

# Register a model from a completed run's artifact URI
run_id = "3f8c2a1b9e4d"
model_uri = f"runs:/{run_id}/model"

result = mlflow.register_model(
    model_uri=model_uri,
    name="churn-predictor",
)
print(f"Registered version: {result.version}")

From the UI

1. Navigate to Calabi ML → Experiments and open the run you want to register.
2. In the Artifacts panel, find the model artifact folder.
3. Click Register Model.
4. Select an existing registered model or enter a new name.
5. Click Register.

---

Model Versions

Every time a model is registered, a new immutable version is created. Versions are numbered sequentially (1, 2, 3, …) in the order they were registered, regardless of which run produced them.

Viewing Versions

Navigate to Calabi ML → Models → churn-predictor. The versions table shows:

Column	Description
Version	Sequential version number
Stage	Current lifecycle stage
Created	Timestamp when this version was registered
Run ID	The experiment run that produced this model
Description	Optional notes about this version

---

Model Stages

Calabi ML uses stages to communicate the deployment status of each model version.

Stage	Meaning
None	Freshly registered — not yet reviewed or deployed
Staging	Under evaluation — deployed to a staging/test environment
Production	Live in production — used by serving infrastructure
Archived	Retired — no longer active; retained for audit/rollback

Only one version can be in Production at a time

When you transition a version to Production, the previously active Production version is automatically moved to Archived. This ensures there is always a single, clear production model.

Transitioning a Version's Stage

Via the Python SDK:

client = MlflowClient(tracking_uri="https://calabi.<your-domain>/mlflow")

# Promote version 3 to Staging
client.transition_model_version_stage(
    name="churn-predictor",
    version="3",
    stage="Staging",
    archive_existing_versions=False,
)

# Promote version 3 to Production (archives current Production version)
client.transition_model_version_stage(
    name="churn-predictor",
    version="3",
    stage="Production",
    archive_existing_versions=True,
)

Via the UI:

1. Open Calabi ML → Models → churn-predictor.
2. Click the version number to open its detail page.
3. Click Stage → Transition to Staging (or Production / Archived).
4. Confirm the transition in the dialog.

---

Model Version Descriptions

Add descriptions to document why a version was promoted, what changed, and any known limitations:

client.update_model_version(
    name="churn-predictor",
    version="3",
    description=(
        "XGBoost pipeline with behavioral_v2 features. "
        "Val AUC 0.91 vs. 0.87 for v2 (production). "
        "Trained on 18 months of data. Validated by data science team 2026-04-01."
    ),
)

---

Aliases

Aliases are mutable named pointers to a specific model version. They decouple your serving code from version numbers — you can update the alias to point to a new version without changing any serving code.

Setting an Alias

client.set_registered_model_alias(
    name="churn-predictor",
    alias="champion",
    version="3",
)

Loading a Model by Alias

model = mlflow.sklearn.load_model("models:/churn-predictor@champion")
predictions = model.predict(X_test)

Updating an Alias During Promotion

# Move champion alias to the new version
client.set_registered_model_alias(name="churn-predictor", alias="champion", version="4")

# Keep a challenger alias for A/B testing
client.set_registered_model_alias(name="churn-predictor", alias="challenger", version="5")

---

Loading Models for Inference

Load a model from the registry for batch scoring or online inference:

import mlflow.sklearn

# Load by stage (always gets the current Production version)
model = mlflow.sklearn.load_model("models:/churn-predictor/Production")

# Load by specific version number
model = mlflow.sklearn.load_model("models:/churn-predictor/3")

# Load by alias
model = mlflow.sklearn.load_model("models:/churn-predictor@champion")

# Run inference
import pandas as pd
features = pd.read_parquet("s3://my-bucket/scoring/batch_2026_04_06.parquet")
predictions = model.predict_proba(features)[:, 1]

Batch Scoring in Calabi Pipelines

from airflow.decorators import dag, task
from datetime import datetime

@dag(
    dag_id="churn_batch_scoring",
    schedule="@daily",
    start_date=datetime(2024, 1, 1),
    catchup=False,
)
def churn_batch_scoring():

    @task()
    def score_customers():
        import mlflow.sklearn
        import pandas as pd

        model = mlflow.sklearn.load_model("models:/churn-predictor/Production")
        features = pd.read_parquet("s3://my-bucket/scoring/latest.parquet")
        scores = model.predict_proba(features)[:, 1]
        features["churn_probability"] = scores
        features.to_parquet("s3://my-bucket/scores/churn_scores_{{ ds }}.parquet", index=False)

    score_customers()

dag_instance = churn_batch_scoring()

---

Promotion Workflow

A typical model promotion workflow for a production-grade model:

Recommended Checklist Before Promoting to Production

• Model version passes all automated unit tests
• Val/test metrics exceed the current production model
• Feature importance reviewed — no data leakage
• Model signature validates against live serving schema
• Model description updated with training data date range and key metrics
• Signed off by data science lead in model version description

---

Model Tags

Tags can be applied to registered models or specific versions for filtering and governance:

# Tag the registered model
client.set_registered_model_tag("churn-predictor", "owner", "alice@example.com")
client.set_registered_model_tag("churn-predictor", "use_case", "retention_scoring")

# Tag a specific version
client.set_model_version_tag("churn-predictor", "3", "validated_by", "bob@example.com")
client.set_model_version_tag("churn-predictor", "3", "validation_date", "2026-04-01")

---

Deleting Models and Versions

# Archive a version (preferred — retains history)
client.transition_model_version_stage(
    name="churn-predictor", version="1", stage="Archived"
)

# Delete a specific version permanently
client.delete_model_version(name="churn-predictor", version="1")

# Delete the entire registered model (all versions)
client.delete_registered_model(name="churn-predictor")

Deletion is permanent

Deleting a model version removes the registry entry and the associated artifact from S3. This action cannot be undone. Prefer archiving over deleting for auditability.

---

Next Steps

• Comparing Runs — Find the best run before registering
• Logging Runs — Record the metrics that inform promotion decisions
• Experiments — Organise runs before promoting to the registry