5.6. Model Registries

What is a model registry?

A model registry is a central storage and management system for machine learning models. It serves as a repository for models at various stages of their lifecycle, from development to deployment. This system supports version control, model tracking, and collaboration across teams, making it a crucial tool for efficient model management.

Why do you need a model registry?

Storing and versioning AI/ML models: Model registries provide a structured way to store and manage multiple versions of models, allowing users to easily roll back to earlier versions if necessary.
Track the origin of AI/ML models: They help track the development history of models, including who trained them, on what data, and the changes made over time.
Make your model easily deployable: With a model registry, models can be moved smoothly from development to production, supporting continuous integration and delivery workflows.
Promote models ready for production: Model registries facilitate the staging and promotion of models, ensuring only verified and tested models make it to production.

Using a model registry acknowledges the unique nature of machine learning models, which are not just code, but include data, configuration, and dependencies.

Which model registry solution should you use?

There are various solutions available, ranging from cloud-based platforms like Google Vertex AI, AWS SageMaker, Azure ML. Third-party solutions such as Weights & Biases and Neptune AI offer ready-to-use model management features. For those looking for an open-source option, MLflow Model Registry is popular and accessible, supporting a range of ML frameworks.

To get started with MLflow, install it via uv:

uv add mlflow

Check its installation and start the server with:

uv run mlflow doctor
uv run mlflow server

What is an MLflow model and a registered model?

An MLflow Model is a model logged during an MLflow experiment, using methods such as mlflow.<model_flavor>.log_model(). After logging, these models can be registered in the MLflow Model Registry. A registered model in MLflow has a unique name and contains various versions, each with its own set of aliases, tags, and metadata.

How should you integrate MLflow Registry in your project?

The integration of MLflow Model Registry happens in 4 steps: Initializing, Saving, registry, and loading.

Initializing

Configure your MLflow environment before any operations:

import mlflow
mlflow.set_tracking_uri("file://./mlruns")
mlflow.set_registry_uri("file://./mlruns")
client.create_registered_model("bikes")

This will store models and their metadata in the local ./mlruns folder.

Saving

Log your model using MLflow, either automatically with autologging or manually:

import mlflow
with mlflow.start_run(run_name="training"):
    model = ...  # your model training logic here
    mlflow.sklearn.log_model(model, "models") # optional if you use autologging

Registring

Register your model to manage its lifecycle with this command following the saving step:

model_version = mlflow.register_model(name="bikes", model_uri=model_info.model_uri)

Loading

Load a model for use in production or further testing:

model_uri = "models:/bikes/1"
model = mlflow.sklearn.load_model(model_uri)
predictions = model.predict(data)

How can you sign the model to provide its inputs and outputs?

Create a model signature to clearly define what input the model expects and what output it produces:

import mlflow
from mlflow.models.signature import infer_signature
inputs, outputs = ...  # define your model inputs and outputs here
signature = infer_signature(inputs, outputs)
mlflow.sklearn.log_model(model, artifact_path="models", signature=signature, input_example=inputs)

How should you access the models in the registry?

Access models either through MLflow's UI for a visual interface or programmatically for automation and integration into other applications:

import mlflow
client = mlflow.tracking.MlflowClient()
model_versions = client.search_model_versions("name='bikes'")
for model_version in model_versions:
    print(model_version)

How should you mark your model as ready to production?

To effectively manage MLflow models, especially when marking a model as ready for production, using aliases is a highly practical approach. This method allows for a flexible handling of model versions, which can change frequently with new training runs or during model rollbacks.

Model aliases in MLflow are mutable, named references that can be assigned to specific versions of a registered model. This allows you to reference a model version using a model URI or through the model registry API without constantly updating the version number. For example, you can create an alias named "champion" that points to version 1 of a model called "MyModel." Subsequently, this model version can be referred to with the URI models:/MyModel@champion.

Assigning an Alias

You can assign an alias through the MLflow UI or programmatically. Here’s how you can do it programmatically:

from mlflow.tracking import MlflowClient

client = MlflowClient(tracking_uri="./mlruns", registry_uri="./mlruns")
client.set_registered_model_alias(name="bikes", alias="Champion", version=1)

Retrieving a Model Using an Alias

To use the model assigned to an alias in your applications, you can retrieve it as follows:

import mlflow

model_uri = "models:/bikes@Champion"
model = mlflow.pyfunc.load_model(model_uri=model_uri)
predictions = model.predict(inputs)

Automating Alias Assignment

To automate the assignment of an alias to the latest model version, you can use:

versions = client.search_model_versions("name='bikes'", max_results=1, order_by=["version_number DESC"])
last_version = versions[0].version
client.set_registered_model_alias(name="bikes", alias="Champion", version=last_version)

How can you rollback your model if something goes wrong?

To rollback a model version in MLflow, you can reassign the alias to a previous stable version either through the UI or programmatically. Remember, after rolling back, you'll need to refresh the model loading in your application either manually or set it up to automatically detect changes.

How can you define custom logic associated with your model?

MLflow supports a customizable model format called "PyFunc" (Python Function), which allows you to define custom logic that executes in conjunction with your model. Here's how you can utilize PyFunc to integrate custom logic with your model:

import mlflow.pyfunc

class CustomModel(mlflow.pyfunc.PythonModel):

    def load_context(self, context):
        # Load artifacts or dependencies here
        pass

    def predict(self, context, model_input):
        # Apply custom logic here
        return model_input.apply(some_transformation)

# Example of saving the model
mlflow.pyfunc.save_model(path="path/to/save", python_model=CustomModel())

This example defines a custom model by subclassing mlflow.pyfunc.PythonModel and implementing the necessary methods to integrate custom behavior during model predictions.