Overview

The MLflow Model Registry component is a centralized model store, set of APIs, and UI, to collaboratively manage the full lifecycle of MLflow Models. It provides model lineage (which MLflow Experiment and Run produced the model), model versioning, stage transitions, annotations, and deployment management.

In this notebook, you use each of the MLflow Model Registry's components to develop and manage a production machine learning application. This notebook covers the following topics:

• Track and log models with MLflow
• Register models with the Model Registry
• Describe models and make model version stage transitions
• Integrate registered models with production applications
• Search and discover models in the Model Registry
• Archive and delete models

Requirements

• Databricks Runtime for Machine Learning.

Machine learning application: Forecasting wind power

In this notebook, you use the MLflow Model Registry to build a machine learning application that forecasts the daily power output of a wind farm. Wind farm power output depends on weather conditions: generally, more energy is produced at higher wind speeds. Accordingly, the machine learning models used in the notebook predict power output based on weather forecasts with three features: wind direction, wind speed, and air temperature.

This notebook uses altered data from the National WIND Toolkit dataset provided by NREL, which is publicly available and cited as follows:

Draxl, C., B.M. Hodge, A. Clifton, and J. McCaa. 2015. Overview and Meteorological Validation of the Wind Integration National Dataset Toolkit (Technical Report, NREL/TP-5000-61740). Golden, CO: National Renewable Energy Laboratory.

Draxl, C., B.M. Hodge, A. Clifton, and J. McCaa. 2015. "The Wind Integration National Dataset (WIND) Toolkit." Applied Energy 151: 355366.

Lieberman-Cribbin, W., C. Draxl, and A. Clifton. 2014. Guide to Using the WIND Toolkit Validation Code (Technical Report, NREL/TP-5000-62595). Golden, CO: National Renewable Energy Laboratory.

King, J., A. Clifton, and B.M. Hodge. 2014. Validation of Power Output for the WIND Toolkit (Technical Report, NREL/TP-5D00-61714). Golden, CO: National Renewable Energy Laboratory.

Load the dataset

The following cells load a dataset containing weather data and power output information for a wind farm in the United States. The dataset contains wind direction, wind speed, and air temperature features sampled every eight hours (once at 00:00, once at 08:00, and once at 16:00), as well as daily aggregate power output (power), over several years.

Display a sample of the data for reference.

Train a power forecasting model and track it with MLflow

The following cells train a neural network to predict power output based on the weather features in the dataset. MLflow is used to track the model's hyperparameters, performance metrics, source code, and artifacts.

Define a power forecasting model using TensorFlow Keras.

Train the model and use MLflow to track its parameters, metrics, artifacts, and source code.

Register the model with the MLflow Model Registry API

Now that a forecasting model has been trained and tracked with MLflow, the next step is to register it with the MLflow Model Registry. You can register and manage models using the MLflow UI or the MLflow API .

The following cells use the API to register your forecasting model, add rich model descriptions, and perform stage transitions. See the documentation for the UI workflow.

Create a new registered model using the API

The following cells use the mlflow.register_model() function to create a new registered model whose name begins with the string power-forecasting-model. This also creates a new model version (for example, Version 1 of power-forecasting-model).

After creating a model version, it may take a short period of time to become ready. Certain operations, such as model stage transitions, require the model to be in the READY state. Other operations, such as adding a description or fetching model details, can be performed before the model version is ready (for example, while it is in the PENDING_REGISTRATION state).

The following cell uses the MlflowClient.get_model_version() function to wait until the model is ready.

Add model descriptions

You can add descriptions to registered models as well as model versions:

• Model version descriptions are useful for detailing the unique attributes of a particular model version (such as the methodology and algorithm used to develop the model).
• Registered model descriptions are useful for recording information that applies to multiple model versions (such as a general overview of the modeling problem and dataset).

Add a high-level description to the registered model, including the machine learning problem and dataset.

Add a model version description with information about the model architecture and machine learning framework.

Perform a model stage transition

The MLflow Model Registry defines several model stages: None, Staging, Production, and Archived. Each stage has a unique meaning. For example, Staging is meant for model testing, while Production is for models that have completed the testing or review processes and have been deployed to applications.

Users with appropriate permissions can transition models between stages. In private preview, any user can transition a model to any stage. In the near future, administrators in your organization will be able to control these permissions on a per-user and per-model basis.

If you have permission to transition a model to a particular stage, you can make the transition directly by using the MlflowClient.update_model_version() function. If you do not have permission, you can request a stage transition using the REST API; for example:

%sh curl -i -X POST -H "X-Databricks-Org-Id: <YOUR_ORG_ID>" -H "Authorization: Bearer <YOUR_ACCESS_TOKEN>" https://<YOUR_DATABRICKS_WORKSPACE_URL>/api/2.0/preview/mlflow/transition-requests/create -d '{"comment": "Please move this model into production!", "model_version": {"version": 1, "registered_model": {"name": "power-forecasting-model"}}, "stage": "Production"}'

Now that you've learned about stage transitions, transition the model to the Production stage.

Use the MlflowClient.get_model_version() function to fetch the model's current stage.

The MLflow Model Registry allows multiple model versions to share the same stage. When referencing a model by stage, the Model Registry will use the latest model version (the model version with the largest version ID). The MlflowClient.get_latest_versions() function fetches the latest model version for a given stage or set of stages. The following cell uses this function to print the latest version of the power forecasting model that is in the Production stage.

Integrate the model with the forecasting application

Now that you have trained and registered a power forecasting model with the MLflow Model Registry, the next step is to integrate it with an application. This application fetches a weather forecast for the wind farm over the next five days and uses the model to produce power forecasts. For example purposes, the application consists of a simple forecast_power() function (defined below) that is executed within this notebook. In practice, you may want to execute this function as a recurring batch inference job using the Databricks Jobs service.

The following section demonstrates how to load model versions from the MLflow Model Registry for use in applications. The Forecast power output with the production model section uses the Production model to forecast power output for the next five days.

Load versions of the registered model

The MLflow Models component defines functions for loading models from several machine learning frameworks. For example, mlflow.tensorflow.load_model() is used to load Tensorflow Keras models that were saved in MLflow format, and mlflow.sklearn.load_model() is used to load scikit-learn models that were saved in MLflow format.

These functions can load models from the MLflow Model Registry.

You can load a model by specifying its name (for example, power-forecast-model) and version number (in this case, 1). The following cell uses the mlflow.pyfunc.load_model() API to load Version 1 of the registered power forecasting model as a generic Python function.

You can also load a specific model stage. The following cell loads the Production stage of the power forecasting model.

Forecast power output with the production model

In this section, the production model is used to evaluate weather forecast data for the wind farm. The forecast_power() application loads the latest version of the forecasting model from the specified stage and uses it to forecast power production over the next five days.

Create and deploy a new model version

The MLflow Model Registry enables you to create multiple model versions corresponding to a single registered model. By performing stage transitions, you can seamlessly integrate new model versions into your staging or production environments. Model versions can be trained in different machine learning frameworks (such as scikit-learn and tensorflow); MLflow's python_function provides a consistent inference API across machine learning frameworks, ensuring that the same application code continues to work when a new model version is introduced.

The following sections create a new version of the power forecasting model using scikit-learn, perform model testing in Staging, and update the production application by transitioning the new model version to Production.

Create a new model version

Classical machine learning techniques are also effective for power forecasting. The following cell trains a random forest model using scikit-learn and registers it with the MLflow Model Registry via the mlflow.sklearn.log_model() function.

Fetch the new model version ID using MLflow Model Registry Search

The MlflowClient.search_model_versions() function searches for model versions by model name, MLflow run ID, or artifact source location. All model versions satisfying a particular filter query are returned.

The following cell uses this search function to fetch the version ID of the new model. It searches for the maximum value of the version ID (that is, the most recent version).

Wait for the new model version to become ready.

Add a description to the new model version

Transition the new model version to Staging

Before deploying a model to a production application, it is often best practice to test it in a staging environment. The following cells transition the new model version to Staging and evaluate its performance.

Evaluate the new model's forecasting performance in Staging

Transition the new model version to Production

After verifying that the new model version performs well in staging, the following cells transition the model version to Production and use the exact same application code from the Forecast power output with the production model section to produce a power forecast.

There are now two model versions of the forecasting model in the Production stage: the model version trained in Tensorflow Keras and the version trained in scikit-learn.

When referencing a model by stage, the MLflow Model Model Registry automatically uses the latest production version. This enables you to update your production models without changing any application code.

See the documentation for how to transition the model to Production using the UI.

Transition the new model version to Production using the API

Archive and delete models

When a model version is no longer being used, you can archive it or delete it. You can also delete an entire registered model; this removes all of its associated model versions.

Archive Version 1 of the power forecasting model

Archive Version 1 of the power forecasting model because it is no longer being used. You can archive models in the MLflow Model Registry UI or via the MLflow API. See the documentation for the UI workflow.

Archive Version 1 using the MLflow API

The following cell uses the MlflowClient.update_model_version() function to archive Version 1 of the power forecasting model.

Delete Version 1 of the power forecasting model

You can also use the MLflow UI or MLflow API to delete model versions. Model version deletion is permanent and cannot be undone.

The following cells provide a reference for deleting Version 1 of the power forecasting model using the MLflow API. See the documentation for how to delete a model version using the UI.

Delete Version 1 using the MLflow API

The following cell permanently deletes Version 1 of the power forecasting model.

Delete the power forecasting model

If you want to delete an entire registered model, including all of its model versions, you can use the MlflowClient.delete_registered_model() to do so. This action cannot be undone. You must first transition all model version stages to None or Archived.

Warning: The following cell permanently deletes the power forecasting model, including all of its versions.