Manage models

Databricks provides a hosted version of the MLflow Model Registry to help you to manage the full lifecycle of MLflow Models. Model Registry provides chronological model lineage (which MLflow experiment and run produced the model at a given time), model versioning, stage transitions (for example, from staging to production or archived), and model and model version annotations and descriptions.

This article covers how to use the Model Registry as part of your machine learning workflow and includes instructions for both the Model Registry UI and the Model Registry API.

For an overview of Model Registry concepts, see MLflow guide.

In this section:

Requirements
Create or register a model
Control access to models
Transition a model stage
Annotate a model or model version
Rename a model (API only)
Search for a model
Delete a model or model version
Share models across workspaces
Example

Requirements

To use the Model Registry UI:

Databricks Runtime 6.4 or above with mlflow>=1.7.0 installed
Databricks Runtime 6.4 ML with mlflow>=1.7.0 installed
Databricks Runtime 6.5 ML or above

To use the Model Registry API:

Databricks Runtime 6.5 or above
Databricks Runtime 6.5 ML or above
Databricks Runtime 6.4 or below with mlflow>=1.7.0 installed
Databricks Runtime 6.4 ML or below with mlflow>=1.7.0 installed

Create or register a model

In this section:

Create or register a model using the UI
Register a model using the API

Create or register a model using the UI

There are two ways to register a model in the Model Registry. You can register an existing model that has been logged to MLflow, or you can create and register a new, empty model and then assign a previously logged model to it.

Register an existing logged model from a notebook

In the Workspace, identify the MLflow run containing the model you want to register.
1. Click the Experiment icon in the notebook tool bar.
2. In the Experiment Runs sidebar, click the icon next to the date of the run. The MLflow Run page displays. This page shows details of the run including parameters, metrics, tags, and list of artifacts.
In the Artifacts section, click the directory named xxx-model.
Click the Register Model button at the far right.
In the Register Model dialog, select Create New Model from the Model drop-down menu, and in the Model Name field, enter a model name, for example scikit-learn-power-forecasting.
Click Register. This registers a model named scikit-learn-power-forecasting, copies the model into a secure location managed by the MLflow Model Registry, and creates a new version of the model.

After a few moments, the MLflow Run UI replaces the Register Model button with a link to the new registered model version.
Click the link to open the new model version in the Model Registry UI. You can also find the model in the Model Registry by clicking the icon in the left navigation bar.

Create a new registered model and assign a logged model to it

You can use the Create Model button on the registered models page to create a new, empty model and then assign a logged model to it. Follow these steps:

On the registered models page, click Create Model. Enter a name for the model and click Create.
Follow Steps 1 through 3 in Register an existing logged model from a notebook.
In the Register Model dialog, select the name of the model you created in Step 1 and click Register. This registers a model with the name you created, copies the model into a secure location managed by the MLflow Model Registry, and creates a model version: Version 1.

After a few moments, the MLflow Run UI replaces the Register Model button with a link to the new registered model version. You can now select the model from the Model drop-down list in the Register Model dialog on the Experiment Runs page. You can also register new versions of the model by specifying its name in API commands like Create ModelVersion.

Register a model using the API

There are three programmatic ways to register a model in the Model Registry. All methods copy the model into a secure location managed by the MLflow Model Registry.

To log a model and register it with the specified name during an MLflow experiment, use the mlflow.<model-flavor>.log_model(...) method. If a registered model with the name doesn’t exist, the method registers a new model, creates Version 1, and returns a ModelVersion MLflow object. If a registered model with the name exists already, the method creates a new model version and returns the version object.
```
with mlflow.start_run(run_name=<run-name>) as run:
  ...
  mlflow.<model-flavor>.log_model(<model-flavor>=<model>,
    artifact_path="<model-path>",
    registered_model_name="<model-name>"
  )
```
To register a model with the specified name after all your experiment runs complete and you have decided which model is most suitable to add to the registry, use the mlflow.register_model() method. For this method, you need the run ID for the mlruns:URI argument. If a registered model with the name doesn’t exist, the method registers a new model, creates Version 1, and returns a ModelVersion MLflow object. If a registered model with the name exists already, the method creates a new model version and returns the version object.
```
result=mlflow.register_model("runs:<model-path>", "<model-name>")
```
To create a new registered model with the specified name, use the MLflow Client API create_registered_model() method. If the model name exists, this method throws an MLflowException.
```
client = MlflowClient()
result = client.create_registered_model("<model-name>")
```

Control access to models

To learn how to control access to models registered in Model Registry, see MLflow Model permissions.

Transition a model stage

A model version has one of the following stages: None, Staging, Production, or Archived. The Staging stage is meant for model testing and validating, while the Production stage is for model versions that have completed the testing or review processes and have been deployed to applications for live scoring. An Archived model version is assumed to be inactive, at which point you can consider deleting it. Different versions of a model can be in different stages.

A user with appropriate permission can transition a model version between stages. If you have permission to transition a model version to a particular stage, you can make the transition directly. If you do not have permission, you can request a stage transition and a user that has permission to transition model versions can approve, reject, or cancel the request.

In this section:

Transition a model stage using the UI
Transition a model stage using the API

Transition a model stage using the UI

Follow these instructions to transition a model’s stage.

To display the list of available model stages and your available options, in a model version page, click the Stage: <Stage> button and request or select a transition to another stage.
Enter an optional comment and click OK.

Transition a model version to the Production stage

After testing and validation, you can transition or request a transition to the Production stage.

Model Registry allows more than one version of the registered model in each stage. If you want to have only one version in Production, you can transition all versions of the model currently in Production to Archived by checking Transition existing Production model versions to Archived.

Approve, reject, or cancel a model version stage transition request

A user without stage transition permission can request a stage transition. The request appears in the Pending Requests section in the model version page:

To approve, reject, or cancel a stage transition request, click the Approve, Reject, or Cancel link.

The creator of a transition request can also cancel the request.

View model version activities

To view all the transitions requested, approved, pending, and applied to a model version, go to the Activities section. This record of activities provides a lineage of the model’s lifecycle for auditing or inspection.

Transition a model stage using the API

Users with appropriate permissions can transition a model version to a new stage.

To update a model version stage to a new stage, use the MLflow Client API transition_model_version_stage() method:

  client = MlflowClient()
  client.transition_model_version_stage(
    name="<model-name>",
    version=<model-version>,
    stage="<stage>",
    description="<description>"
  )

The accepted values for <stage> are: "Staging"|"staging", "Archived"|"archived", "Production"|"production", "None"|"none".

Annotate a model or model version

You can provide information about a model or model version by annotating it. For example, you may want to include an overview of the problem or information about the methodology and algorithm used.

Annotate a model or model version using the UI

To add or update a model or model version description:

From the registered model or model version page, click the Description icon. An edit window displays.
Enter or edit the description in the edit window.
Click Save.

If you entered a description of a model version, the description appears in the Description column in the table on the registered model page. The column displays a maximum of 32 characters or one line of text, whichever is shorter.

Annotate a model version using the API

To update a model version description, use the MLflow Client API update_model_version() method:

client = MlflowClient()
client.update_model_version(
  name="<model-name>",
  version=<model-version>,
  description="<description>"
)

Rename a model (API only)

To rename a registered model, use the MLflow Client API rename_registered_model() method:

client=MlflowClient()
client.rename_registered_model("<model-name>", "<new-model-name>")

Note

You can rename a registered model only if it has no versions, or all versions are in the None or Archived stage.

Search for a model

All registered models live in the MLflow Model Registry. You can search for models using the UI or the API.

Search for a model using the UI

To display all registered models, click the icon in the sidebar.

To search for a specific model, type in the model name in the search box.

Search for a model using the API

To retrieve a list of all registered models, use the MLflow Client API list_model_versions() method:

from pprint import pprint

client = MlflowClient()
for rm in client.list_registered_models():
  pprint(dict(rm), indent=4)

This outputs:

{   'creation_timestamp': 1582671933216,
    'description': None,
    'last_updated_timestamp': 1582671960712,
    'latest_versions': [<ModelVersion: creation_timestamp=1582671933246, current_stage='Production', description='A random forest model containing 100 decision trees trained in scikit-learn', last_updated_timestamp=1582671960712, name='sk-learn-random-forest-reg-model', run_id='ae2cc01346de45f79a44a320aab1797b', source='./mlruns/0/ae2cc01346de45f79a44a320aab1797b/artifacts/sklearn-model', status='READY', status_message=None, user_id=None, version=1>,
    <ModelVersion: creation_timestamp=1582671960628, current_stage='None', description=None, last_updated_timestamp=1582671960628, name='sk-learn-random-forest-reg-model', run_id='d994f18d09c64c148e62a785052e6723', source='./mlruns/0/d994f18d09c64c148e62a785052e6723/artifacts/sklearn-model', status='READY', status_message=None, user_id=None, version=2>],
    'name': 'sk-learn-random-forest-reg-model'}

You can also search for a specific model name and list its version details using MLflow Client API search_model_versions() method:

from pprint import pprint

client=MlflowClient()
[pprint(mv) for mv in client.search_model_versions("name=<model-name>")]

This outputs:

{   'creation_timestamp': 1582671933246,
    'current_stage': 'Production',
    'description': 'A random forest model containing 100 decision trees '
                   'trained in scikit-learn',
    'last_updated_timestamp': 1582671960712,
    'name': 'sk-learn-random-forest-reg-model',
    'run_id': 'ae2cc01346de45f79a44a320aab1797b',
    'source': './mlruns/0/ae2cc01346de45f79a44a320aab1797b/artifacts/sklearn-model',
    'status': 'READY',
    'status_message': None,
    'user_id': None,
    'version': 1 }

{   'creation_timestamp': 1582671960628,
    'current_stage': 'None',
    'description': None,
    'last_updated_timestamp': 1582671960628,
    'name': 'sk-learn-random-forest-reg-model',
    'run_id': 'd994f18d09c64c148e62a785052e6723',
    'source': './mlruns/0/d994f18d09c64c148e62a785052e6723/artifacts/sklearn-model',
    'status': 'READY',
    'status_message': None,
    'user_id': None,
    'version': 2 }

Delete a model or model version

You can delete a model using the UI or the API.

Delete a model version or model using the UI

Warning

You cannot undo this action. You can transition a model version to the Archived stage rather than deleting it from the registry. When you delete a model, all model artifacts stored by the Model Registry and all the metadata associated with the registered model are deleted.

Note

You can only delete models and model versions in the None or Archived stage. If a registered model has versions in the Staging or Production stage, you must transition them to either the None or Archived stage before deleting the model.

To delete a model version:

Click the icon in the sidebar.
Click a model name.
Click a model version.
Select the next to the version number.
Click the Delete button.

To delete a model:

Click the icon in the sidebar.
Click a model name.
Select the next to the model.
Click the Delete button.

Delete a model version or model using the API

Warning

You cannot undo this action. You can transition a model version to the Archived stage rather than deleting it from the registry. When you delete a model, all model artifacts stored by the Model Registry and all the metadata associated with the registered model are deleted.

Note

You can only delete models and model versions in the None or Archived stage. If a registered model has versions in the Staging or Production stage, you must transition them to either the None or Archived stage before deleting the model.

Delete a model version

To delete a model version, use the MLflow Client API delete_model_version() method:

# Delete versions 1,2, and 3 of the model
client = MlflowClient()
versions=[1, 2, 3]
for version in versions:
  client.delete_model_version(name="<model-name>", version=version)

Delete a model

To delete a model, use the MLflow Client API delete_registered_model() method:

client = MlflowClient()
client.delete_registered_model(name="<model-name>")

Example

This example illustrates how to use the Model Registry to build a machine learning application.

MLflow Model Registry example