• MLOps
  • MlflowClient().transition_model_version_stage() Unity Catalog

Thank you for sharing the error β€” this clears everything up. πŸ™Œ

You're using MLflow in Databricks with Unity Catalog, which introduces new model management practices.


⚠️ What’s Happening:

You're getting this error:

❌ MlflowClient().transition_model_version_stage() is not supported in Unity Catalog

Because:

  • Unity Catalog does not use "Stages" (Production/Staging/etc.) anymore.
  • Instead, it uses Aliases (like tags or pointers to versions) for more flexible model deployment.

βœ… What You Should Use Instead: Model Aliases

Instead of setting the model to "Production" stage, assign an alias like "production" to a specific version.


🧠 Why Aliases Are Better (Unity Catalog)

| Old (Stages) | New (Unity Catalog Aliases) |
|---------------------------|-------------------------------------|
| Staging, Production | Aliases like @staging, @prod |
| Limited to one stage | Can assign multiple aliases |
| Deprecated | Future-proof with aliases |


βœ… Step-by-Step Fix Using Aliases

🎯 Step 1: Register the model (you already did this)

model_uri = f"runs:/{run_id}/iris_rf_model"

model_details = mlflow.register_model(
    model_uri=model_uri,
    name="iris_catalog.default.IrisClassifierModel"  # Update with your catalog path
)

πŸ“Œ Note: In Unity Catalog, model name includes:
catalog.schema.model_name, like:

"iris_catalog.default.IrisClassifierModel"

🎯 Step 2: Assign an alias (like @production) to the version

from mlflow.tracking import MlflowClient

client = MlflowClient()

client.set_registered_model_alias(
    name="iris_catalog.default.IrisClassifierModel",
    alias="production",
    version=1
)

🎯 Step 3: Load the model using the alias

from mlflow.pyfunc import load_model

model = load_model("models:/iris_catalog.default.IrisClassifierModel@production")

# Make predictions
preds = model.predict(X_test)
print(preds)

βœ… Summary of Fix

| Action | Old Way | βœ… New Way (Unity Catalog) |
|----------------------|--------------------------------|--------------------------------------------------|
| Transition stage | transition_model_version_stage() | ❌ Deprecated / Not allowed |
| Assign deployment tag| ❌ Not flexible | βœ… set_registered_model_alias(..., alias="production") |
| Load model | "models:/ModelName/Production" | βœ… "models:/ModelName@production" |


Let me know your exact model name path in Unity Catalog if you're unsure, and I’ll help adjust the commands. Or I can give you a complete updated notebook for Unity Catalog use.