Manage models

Databricks provides a hosted version of 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 email notifications of model events. You can also create and view model descriptions and leave comments.

This article covers how to use 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.

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

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

  1. In the Workspace, identify the MLflow run containing the model you want to register.

    1. Click the Experiment icon Experiment icon in the notebook tool bar.

      Notebook toolbar
    2. In the Experiment Runs sidebar, click the External Link 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.

  2. In the Artifacts section, click the directory named xxx-model.

    Register model
  3. Click the Register Model button at the far right.

  4. In the dialog, click in the Model box and do one of the following:

    • Select Create New Model from the drop-down menu. The Model Name field appears. Enter a model name, for example scikit-learn-power-forecasting.
    • Select an existing model from the drop-down menu.
    Create new model
  5. Click Register.

    • If you selected Create New Model, 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.
    • If you selected an existing model, this registers a new version of the selected model.

    After a few moments, the Register Model button changes to a link to the new registered model version.

    Select newly created model
  6. 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 Models Icon Models in the sidebar.

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:

  1. On the registered models page, click Create Model. Enter a name for the model and click Create.

  2. Follow Steps 1 through 3 in Register an existing logged model from a notebook.

  3. 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.

Transition a model stage using the UI

Follow these instructions to transition a model’s stage.

  1. 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.

    Stage transition options
  2. 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:

Transition to production

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".

Use model for inference

Preview

This feature is in Public Preview.

After a model is registered in Model Registry, you can automatically generate a notebook to use the model for batch inference or create an endpoint to use the model for real-time serving.

In the upper right corner of the registered model page or the model version page, click use model button. The Configure model inference dialog displays, allowing you to configure batch or real-time inference.

Configure model inference dialog

Configure batch inference

When you follow these steps to create a batch inference notebook, the notebook is saved in your user folder under the Batch-Inference folder in a folder with the model’s name. You can edit the notebook as needed.

  1. Click the Batch inference tab.

  2. From the Model version drop-down, select the model version to use. The first two items in the drop-down are the current Production and Staging version of the model (if they exist). When you select one of these options, the notebook automatically uses the Production or Staging version as of the time it is run. You do not need to update the notebook as you continue to develop the model.

  3. Click the Browse button next to Input table. The Select input data dialog displays. If necessary, you can change the cluster in the Compute drop-down.

  4. Select the database and table containing the input data for the model, and click Select. The generated notebook automatically imports this data and sends it to the model. You can edit the generated notebook if the data requires any transformations before it is input to the model.

  5. Predictions are saved in a folder in the directory dbfs:/FileStore/batch-inference. By default, predictions are saved in a folder with the same name as the model. Each run of the generated notebook writes a new file to this directory with the timestamp appended to the name. You can also choose not to include the timestamp and to overwrite the file with subsequent runs of the notebook; instructions are provided in the generated notebook.

    You can change the folder where the predictions are saved by typing a new folder name into the Output table location field or by clicking the folder icon to browse the directory and select a different folder.

Configure real-time inference

  1. In the Configure model inference dialog, click the Real-time tab.

  2. If serving is not enabled for the model, click Enable serving. The Serving tab of the registered model page appears, with Status shown as Pending. After a few minutes, Status changes to Ready.

    If serving is already enabled, click View existing real-time inference to display the Serving tab.

Provide feedback

This feature is in preview, and we would love to get your feedback. To provide feedback, click Provide Feedback in the Configure model inference dialog.

Control notification preferences

You can configure Model Registry to notify you by email about activity on registered models and model versions that you specify.

On the registered model page, the Notify me about menu shows three options:

Email notifications menu
  • All new activity: Send email notifications about all activity on all model versions of this model. If you created the registered model, this setting is the default.
  • Activity on versions I follow: Send email notifications only about model versions you follow. With this selection, you receive notifications for all model versions that you follow; you cannot turn off notifications for a specific model version.
  • Mute notifications: Do not send email notifications about activity on this registered model.

The following events trigger an email notification:

  • Creation of a new model version
  • Request for a stage transition
  • Stage transition
  • New comments

You are automatically subscribed to model notifications when you do any of the following:

  • Comment on that model version
  • Transition a model version’s stage
  • Make a transition request for the model’s stage

To see if you are following a model version, look at the Follow Status field on the model version page, or at the table of model versions on the registered model page.

Turn off all email notifications

You can turn off email notifications in the Model Registry Settings tab of the User Settings menu:

  1. Click User Settings Icon Settings in the lower left corner of your Databricks workspace.
  2. Click User Settings.
  3. Go to the Model Registry Settings tab.
  4. Deselect Turn on model registry email notifications.

An admin can turn off email notifications for the entire organization in the Admin Console.

Maximum number of emails sent

Model Registry limits the number of emails sent to each user per day per activity. For example, if you receive 20 emails in one day about new model versions created for a registered model, Model Registry sends an email noting that the daily limit has been reached, and no additional emails about that event are sent until the next day.

To increase the limit of the number of emails allowed, contact your Databricks representative.

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

You can annotate a model in two ways: using a description or using comments. Descriptions are available for models and model versions; comments are only available for model versions. Descriptions are intended to provide information about the model. Comments provide a way to maintain an ongoing discussion about activities on a model version.

Add or update the description for a model or model version

  1. From the registered model or model version page, click the Description Edit Icon icon. An edit window displays.

  2. Enter or edit the description in the edit window.

  3. 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.

Add comments for a model version

  1. Scroll down the model version page and click the down arrow next to Activities.
  2. Type your comment in the edit window and click Add Comment.

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 Models Icon Models in the sidebar.

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

Registered models search

You can also filter models based on tags. Click Filter, and enter tags in the format: tags.<key>=<value>. To filter based on multiple tags, use the AND operator.

Tag-based search

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:

  1. Click Models Icon Models in the sidebar.
  2. Click a model name.
  3. Click a model version.
  4. Click Delete model version at the upper right corner of the screen and select Delete from the drop-down menu.

To delete a model:

  1. Click Models Icon Models in the sidebar.
  2. Click a model name.
  3. Click Delete model version at the upper right corner of the screen and select Delete from the drop-down menu.

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>")

Share models across workspaces

Databricks supports sharing models across multiple workspaces. For example, you can develop and log a model in your own workspace and then register it in a centralized model registry. This is useful when multiple teams share access to models. You can create multiple workspaces and use and manage models across these environments.

Example

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