Model stages (staging, production, archived) in MLOps - Commands & Configuration

Introduction

When you train machine learning models, you need a way to organize them by their readiness. Model stages help you separate models that are being tested, models ready for use, and old models you want to keep but not use. This keeps your work clear and safe.

When you want to test a new model version without affecting the current one used in your app.

When you have a model ready and want to mark it as the official one for predictions.

When you want to keep old models for reference or rollback but not use them actively.

When you want to automate deployment steps based on model readiness.

When you want to track model progress clearly in your machine learning workflow.

Commands

This command starts a local server to serve the model from a specific run. It helps you test the model in staging before production.

Terminal

mlflow models serve -m runs:/1234567890abcdef/model -p 1234

Expected Output

2024/06/01 12:00:00 INFO mlflow.models.cli: Starting MLflow model server...

2024/06/01 12:00:01 INFO mlflow.models.cli: Model server listening on port 1234

→

-m - Specifies the model path to serve

→

-p - Sets the port number for the server

This command moves the model version 1 to the Production stage, marking it as ready for real use.

Terminal

mlflow models transition-stage --model-name my-model --version 1 --stage Production

Expected Output

Model version '1' of 'my-model' transitioned to stage 'Production'.

→

--model-name - Specifies the registered model name

→

--version - Specifies the model version number

→

--stage - Specifies the target stage (e.g., Production, Staging, Archived)

This command archives the model version 1, meaning it is kept but not used for serving or training.

Terminal

mlflow models transition-stage --model-name my-model --version 1 --stage Archived

Expected Output

Model version '1' of 'my-model' transitioned to stage 'Archived'.

→

--model-name - Specifies the registered model name

→

--version - Specifies the model version number

→

--stage - Sets the model stage to Archived

This command lists all versions of the model with their current stages, so you can see which are in staging, production, or archived.

Terminal

mlflow models list-versions --model-name my-model

Expected Output

Version  Stage      Run ID

1        Production 1234567890abcdef

2        Staging    fedcba0987654321

3        Archived   1122334455667788

→

--model-name - Specifies the model to list versions for

Key Concept

If you remember nothing else from this pattern, remember: model stages help you safely manage and deploy machine learning models by clearly marking their readiness and use.

Code Example

MLOps

import mlflow
import mlflow.sklearn

# Log a model and set its stage
with mlflow.start_run() as run:
    mlflow.log_param("param1", 5)
    mlflow.sklearn.log_model(sk_model=None, artifact_path="model")
    model_uri = f"runs:/{run.info.run_id}/model"

# Register the model
model_name = "my-model"
result = mlflow.register_model(model_uri, model_name)

# Transition model to staging
client = mlflow.tracking.MlflowClient()
client.transition_model_version_stage(name=model_name, version=result.version, stage="Staging")

print(f"Model version {result.version} is now in Staging stage.")

Common Mistakes

Trying to serve a model without setting it to Production stage first.

The model might not be stable or ready, causing errors or bad predictions in real use.

Test the model in Staging stage first, then transition it to Production when ready.

Archiving a model version that is still in use by the application.

Archived models are not served or updated, so the app will fail to find the model.

Ensure no active deployments use the model before archiving it.

Not listing model versions to check their stages before deployment.

You might deploy the wrong model version or miss important updates.

Always list and verify model versions and stages before deploying.

Summary

Use model stages to organize models by readiness: Staging for testing, Production for live use, Archived for old models.

Transition models between stages using mlflow CLI or Python client to control deployment safely.

List model versions and their stages to keep track of which models are active or archived.