MLflow Model Registry
The MLflow Model Registry is a centralized model store, set of APIs and a UI designed to collaboratively manage the full lifecycle of a model. It provides lineage (i.e., which MLflow experiment and run produced the model), versioning, aliasing, metadata tagging and annotation support to ensure that you have the full spectrum of information at every stage from development to production deployment.
Why Model Registry?β
As machine learning projects grow in complexity and scale, managing models manually across different environments, teams, and iterations becomes increasingly error-prone and inefficient. The MLflow Model Registry addresses this challenge by providing a centralized, structured system for organizing and governing ML models throughout their lifecycle.
Using the Model Registry offers the following benefits:
- ποΈ Version Control: The registry automatically tracks versions of each model, allowing teams to compare iterations, roll back to previous states, and manage multiple versions in parallel (e.g., staging vs. production).
- 𧬠Model Lineage and Traceability: Each registered model version is linked to the MLflow run, logged model or notebook that produced it, enabling full reproducibility. You can trace back exactly how a model was trained, with what data and parameters.
- π Production-Ready Workflows: Features like model aliases (e.g., @champion) and tags make it easier to manage deployment workflows, promoting models to experimental, staging, or production environments in a controlled and auditable way.
- π‘οΈ Governance and Compliance: With structured metadata, tagging, and role-based access controls (when used with a backend like Databricks or a managed MLflow service), the Model Registry supports governance requirements critical for enterprise-grade ML operations.
Whether youβre a solo data scientist or part of a large ML platform team, the Model Registry is a foundational component for scaling reliable and maintainable machine learning systems.
Conceptsβ
The Model Registry introduces a few concepts that describe and facilitate the full lifecycle of an MLflow Model.
Concept | Description |
---|---|
Model | An MLflow Model is created with one of the model flavorβs |
Registered Model | An MLflow Model can be registered with the Model Registry. A registered model has a unique name, contains versions, aliases, tags, and other metadata. |
Model Version | Each registered model can have one or many versions. When a new model is added to the Model Registry, it is added as version 1. Each new model registered to
the same model name increments the version number. Model versions have tags, which can be useful for tracking attributes of the model version (e.g. |
Model URI | You can refer to the registered model by using a URI of this format: |
Model Alias | Model aliases allow you to assign a mutable, named reference to a particular version of a registered model. By assigning an alias to a specific model version,
you can use the alias to refer to that model version via a model URI or the model registry API. For example, you can create an alias named Aliases are especially useful for deploying models. For example, you could assign a |
Tags | Tags are key-value pairs that you associate with registered models and model versions, allowing you to label and categorize them by function or status. For example, you
could apply a tag with key |
Annotations and Descriptions | You can annotate the top-level model and each version individually using Markdown, including the description and any relevant information useful for the team such as algorithm descriptions, datasets employed or the overall methodology involved in a given version's modeling approach. |
Model Registry in practiceβ
The MLflow Model Registry is available in both open-source (OSS) MLflow and managed platforms like Databricks. Depending on the environment, the registry offers different levels of integration, governance, and collaboration features.
Model Registry in OSS MLflowβ
In the open-source version of MLflow, the Model Registry provides both a UI and API for managing the lifecycle of machine learning models. You can register models, track versions, add tags and descriptions, and transition models between stages such as Staging and Production.
Register a model in MLflow
- Python APIs
- MLflow UI
Register a model with MLflow Python APIsβ
MLflow provides several ways to register a model version
# Option 1: specify `registered_model_name` parameter when logging a model
mlflow.<flavor>.log_model(..., registered_model_name="<YOUR_MODEL_NAME>")
# Option 2: register a logged model
mlflow.register_model(model_uri="<YOUR_MODEL_URI>", name="<YOUR_MODEL_NAME>")
After registering the model, you can load it back with the model name and version
mlflow.<flavor>.load_model("models:/<YOUR_MODEL_NAME>/<YOUR_MODEL_VERSION>")
Register a model on MLflow UIβ
- Open the details page for the MLflow Run containing the logged MLflow model you'd like to register. Select the model folder containing the intended MLflow model in the Artifacts section.
-
Click the Register Model button, which will trigger a modal form to pop up.
-
In the Model dropdown menu on the form, you can either select "Create New Model" (which creates a new registered model with your MLflow model as its initial version) or select an existing registered model (which registers your model under it as a new version). The screenshot below demonstrates registering the MLflow model to a new registered model named
"iris_model_testing"
.
To learn more about the OSS Model Registry, refer to the tutorial on the model registry.
Model Resgitry in Databricksβ
Databricks extends MLflowβs capabilities by integrating the Model Registry with Unity Catalog, enabling centralized governance, fine-grained access control, and cross-workspace collaboration.
Key benefits of Unity Catalog integration include:
- π‘οΈ Enhanced governance: Apply access policies and permission controls to model assets.
- π Cross-workspace access: Register models once and access them across multiple Databricks workspaces.
- π Model lineage: Track which notebooks, datasets, and experiments were used to create each model.
- π Discovery and reuse: Browse and reuse production-grade models from a shared catalog.
Register a model in Databricks UC
- Python APIs
- Databricks UI
Register a model to Databricks UC with MLflow Python APIsβ
Prerequisite: Set tracking uri to Databricks
import mlflow
mlflow.set_registry_uri("databricks-uc")
Use MLflow APIs to register the model
# Option 1: specify `registered_model_name` parameter when logging a model
mlflow.<flavor>.log_model(..., registered_model_name="<YOUR_MODEL_NAME>")
# Option 2: register a logged model
mlflow.register_model(model_uri="<YOUR_MODEL_URI>", name="<YOUR_MODEL_NAME>")
ML model versions in UC must have a model signature. If you want to set a signature on a model that's
already logged or saved, the mlflow.models.set_signature()
API is available for this purpose.
After registering the model, you can load it back with the model name and version
mlflow.<flavor>.load_model("models:/<YOUR_MODEL_NAME>/<YOUR_MODEL_VERSION>")
Register a model on Databricks UIβ
-
From the experiment run page or models page, click Register model in the upper-right corner of the UI.
-
In the dialog, select Unity Catalog, and select a destination model from the drop down list.
- Click Register.
Registering a model can take time. To monitor progress, navigate to the destination model in Unity Catalog and refresh periodically.
For more information, refer to the Databricks documentation on managing the model lifecycle.