The W&B Model Registry isn’t just a versioned file store; it’s a living catalog where models transition through distinct stages, and "Production" is the ultimate destination.
Let’s see what that looks like. Imagine you’ve just trained a fantastic new image classification model.
import wandb
import os
# Assume you have a trained model saved locally
# For demonstration, we'll create a dummy file
dummy_model_path = "my_awesome_model.h5"
with open(dummy_model_path, "w") as f:
f.write("This is my model data.")
# Log the model artifact
# This creates a new "model" artifact in W&B
# The 'aliases' argument is key for managing stages
with wandb.init(project="model-registry-demo", job_type="model-training") as run:
model_artifact = wandb.Artifact(
name="image-classifier",
type="model",
description="My latest and greatest image classifier.",
aliases=["staging"] # Initially tag it as staging
)
model_artifact.add_file(dummy_model_path)
run.log_artifact(model_artifact)
print(f"Model artifact logged with alias 'staging'.")
# Clean up the dummy file
os.remove(dummy_model_path)
Now, you’ll see image-classifier in your W&B project’s "Artifacts" section, with a version tagged staging. This staging tag is your first hint that we’re not just talking about files, but about a model’s lifecycle.
The Model Registry allows you to define specific stages for your models, like staging, production, archived, etc. When you "promote" a model, you’re not copying files; you’re changing the metadata associated with a specific artifact version, moving its designated tag.
The Stages of a Model
Think of these stages as checkpoints in your MLOps pipeline.
- Development/Staging: The model is ready for testing, integration, or shadow deployment. It’s not yet serving live traffic.
- Production: The model is actively being used to make predictions on live data.
- Archived: A previous production model that has been superseded but kept for historical reference or rollback purposes.
You can define custom stages, but these three are common starting points.
Promoting a Model: The Actual Mechanism
Promotion happens via the W&B UI or the API. Let’s look at the API. Suppose you’ve tested the staging model and are happy with it. You want to make it available for inference.
First, you need to identify the specific artifact version you want to promote. When you logged the artifact with aliases=["staging"], W&B created a version like image-classifier:staging. If you want to be more precise, you’d fetch the exact version ID.
# Assuming you know the artifact name and its current alias
artifact_name = "image-classifier"
current_alias = "staging"
production_alias = "production"
# Initialize W&B client
api = wandb.Api()
# Get the artifact and its specific version
try:
artifact_version = api.artifact(f"{wandb.run.entity}/{wandb.run.project}/{artifact_name}:{current_alias}")
print(f"Found artifact version: {artifact_version.version}")
# Promote it by adding the 'production' alias
# This is the core promotion action.
# W&B handles the metadata update.
artifact_version.aliases.append(production_alias)
artifact_version.save() # Save the changes to the artifact's metadata
print(f"Successfully promoted {artifact_name}:{artifact_version.version} to '{production_alias}'.")
# Optionally, remove the old alias if you want a clean promotion
# This is often done when a new model *replaces* an existing production model
if current_alias in artifact_version.aliases:
artifact_version.aliases.remove(current_alias)
artifact_version.save()
print(f"Removed '{current_alias}' alias from {artifact_name}:{artifact_version.version}.")
except Exception as e:
print(f"Error promoting artifact: {e}")
When artifact_version.save() is called, W&B updates the metadata for that specific artifact version. It adds the production alias and potentially removes the staging alias. This is a lightweight operation; no files are copied. Your inference service or deployment tool can then be configured to pull the artifact tagged production.
The "One Thing"
The most surprising aspect is how W&B treats aliases. They aren’t just labels; they are pointers to specific artifact versions. When you promote a model, you’re essentially updating which artifact version the production alias points to. This means you can have multiple versions of your model in the registry, but only one (or a few, if you manage multiple production slots) can bear the production tag at any given time. This de-coupling of the "production" label from the actual files is what makes rollbacks and A/B testing conceptually simple within the registry framework.
What happens if your production inference service tries to pull a model artifact that doesn’t exist or has no production alias?