MLflow tutorial: inference

This notebook shows how to load a model previously logged to MLflow and use it to make predictions on data in different formats. The notebook includes two examples of applying the model:

• as a scikit-learn model to a pandas DataFrame
• as a PySpark UDF to a Spark DataFrame

Requirements

• If you are using a cluster running Databricks Runtime, you must install MLflow. See "Install a library on a cluster" (AWS|Azure|GCP). Select Library Source PyPI and enter mlflow in the Package field.
• If you are using a cluster running Databricks Runtime ML, MLflow is already installed.

Prerequsite

• This notebook uses the ElasticNet models from MLflow tutorial part 1: training and logging (AWS|Azure|GCP).

Find and copy the run ID of the run that created the model

Find and copy a run ID associated with an ElasticNet training run from the MLflow tutorial part 1: training and logging notebook. The run ID appears on the run details page; it is a 32-character alphanumeric string shown after the label "Run".

To navigate to the run details page for the MLflow tutorial part 1: training and logging notebook, open that notebook and click Experiment in the upper right corner. The Experiments sidebar displays. Do one of the following:

• In the Experiments sidebar, click the icon at the far right of the date and time of the run. The run details page appears in a new tab.

• Click the square icon with the arrow to the right of Experiment Runs. The Experiment page displays in a new tab. This page lists all of the runs associated with this notebook. To display the run details page for a particular run, click the link in the Start Time column for that run.

For more information, see "View notebook experiment" (AWS|Azure|GCP).

Load the model as a scikit-learn model

Use the MLflow API to load the model from the MLflow server that was created by the run. After loading the model, you can use just like you would any scikit-learn model.

Create a PySpark UDF and use it for batch inference

In this section, you use the MLflow API to create a PySpark UDF from the model you saved to MLflow. For more information, see Export a python_function model as an Apache Spark UDF.

Saving the model as a PySpark UDF allows you to run the model to make predictions on a Spark DataFrame.

Use the Spark function withColumn() to apply the PySpark UDF to the DataFrame and return a new DataFrame with a prediction column.