MLflow stages vs aliases: the wrong model version stayed in production for three weeks

Our churn prediction model is retrained monthly. In September, the data science team trained version 14, validated it, and promoted it. Three weeks later, a business review noticed that churn scores in the CRM had not changed character at all, while the new model was supposed to handle a new subscription plan. The serving logs showed:

2024-10-02 06:00:12 INFO loading model uri=models:/churn/Production -> version 12

Version 12. In the MLflow UI, version 14 had the alias champion. Version 12 still had stage Production.

Reproducing it

The promotion step had changed without anyone telling the platform team. The data scientists had upgraded MLflow in their notebooks and followed the current documentation, which recommends aliases. The training pipeline logs also showed a warning nobody read:

FutureWarning: ``mlflow.tracking.client.MlflowClient.transition_model_version_stage`` is deprecated
since 2.9.0. This method will be removed in a future release. Model registry stages will be removed
in a future major release.

The deploy script, written two years earlier, resolved models by stage.

What we ruled out

• A caching issue in the serving layer. It reloaded the model daily and logged the resolved version.
• A failed registration. Version 14 was registered with all artifacts.

Root cause

MLflow 2.9 deprecated model registry stages (None, Staging, Production, Archived) in favour of aliases and tags. Aliases are named, mutable pointers to versions, such as champion. Our organisation ended up with two promotion mechanisms at once: data scientists moved champion, the deploy script read Production. Neither side was wrong according to its own documentation, and nothing checked that they agreed.

What is happening under the hood

Stages were a fixed set of lifecycle labels on each model version, and a version could only be in one stage. Aliases are free-form names that point to exactly one version each, and one version can have several aliases. Model URIs support both:

models:/churn/Production     # stage (deprecated)
models:/churn@champion       # alias
models:/churn/14             # explicit version

Aliases fit real workflows better: champion and challenger for A/B tests, shadow for shadow deployments, per-region aliases. Tags on versions record validation results and approvals. Some registries, such as Unity Catalog-backed registries, support only aliases.

Environment

Minimal reproduction

client = MlflowClient()
client.set_registered_model_alias("churn", "champion", "14")     # data science promotion
model = mlflow.pyfunc.load_model("models:/churn/Production")     # deploy still reads stage -> v12

The fix

1. One promotion mechanism: aliases. The deploy script now loads by alias:

MODEL_URI = "models:/churn@champion"

def load_champion():
    client = MlflowClient()
    mv = client.get_model_version_by_alias("churn", "champion")
    log.info("loading churn version=%s run_id=%s", mv.version, mv.run_id)
    return mlflow.pyfunc.load_model(MODEL_URI), mv

2. Promotion through a pipeline, not by hand. A small promotion job checks that the candidate has passing validation tags before moving the alias:

REQUIRED_TAGS = {"validation_status": "passed", "fairness_check": "passed"}

def promote(name: str, version: str, approver: str) -> None:
    mv = client.get_model_version(name, version)
    missing = {k: v for k, v in REQUIRED_TAGS.items() if mv.tags.get(k) != v}
    if missing:
        raise RuntimeError(f"cannot promote {name} v{version}: {missing}")
    previous = client.get_model_version_by_alias(name, "champion").version
    client.set_registered_model_alias(name, "previous_champion", previous)
    client.set_registered_model_alias(name, "champion", version)
    client.set_model_version_tag(name, version, "promoted_by", approver)

Keeping previous_champion makes rollback a single alias move.

3. Migrate existing models. A one-off script set champion on every model's current Production version, and we removed stage usage from all code.

4. Unpin the deploy client from MLflow 2.3, which predates aliases, and align client versions across training and serving.

Alternatives and trade-offs

• Keep using stages until removal. Works for now, keeps the deprecation warnings, and blocks moving to registries without stages.
• Deploy explicit versions from config in Git. Very auditable, more manual. Some teams prefer this; aliases with an audit tag gave us similar traceability with less friction.

How to test the fix

An integration test against a local MLflow server registers two versions, promotes the second with the promotion job, runs the deploy loader and asserts it loads the second. A negative test tries to promote a version without validation tags and expects an error.

Confirming in production

• Version 14 went live the day after the fix; CRM churn scores shifted as expected for customers on the new plan.
• The serving log line with version and run id is now part of the daily job summary posted to the team channel.

Regression guard

• A daily check compares the version each deployment is serving with the champion alias in the registry and alerts on mismatch for more than 24 hours.
• Warnings from MLflow clients are collected in CI logs; deprecation warnings fail the training pipeline's lint stage.

Security implications

Promotion decides which model affects customers. Previously, anyone with registry write access could move a stage from the UI. Now only the promotion job's service identity can set the champion alias, enforced by registry permissions, and each promotion records who approved it. Model artifacts in S3 are write-once for the training role, so a promoted version cannot be swapped underneath its alias.

Prevention checklist

• Use one promotion mechanism, aliases, everywhere.
• Promote through a job that checks validation tags.
• Log the resolved model version and run id at load time.
• Compare served versions with registry aliases daily.
• Keep MLflow client versions aligned between training and serving.

References

• MLflow model registry, aliases and deprecation of stages — https://mlflow.org/docs/latest/model-registry.html
• MLflow 2.9.0 release notes — https://github.com/mlflow/mlflow/releases/tag/v2.9.0
• MlflowClient API (set_registered_model_alias) — https://mlflow.org/docs/latest/python_api/mlflow.client.html