Deploy a MLflow Model to SageMaker

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

Setup environment

Upgrade SageMaker Python SDK

Install MLflow and AWS MLflow plugin

Import necessary libraries

Declare some variables used later

Get some training data

Let's download the save the Iris dataset

And now let's upload that data to S3

Write your training script

Let's write the code to train a Decision Tree model using the scikit-learn framework

Since we're using MLflow in our training script, let's make sure the container installs mlflow along with our MLflow plugin before running our training script. We can do this by creating a requirements.txt file and putting it in the same directory as our training script.

SageMaker Training and MLflow

Train your Decision tree model by launching a SageMaker Training job.

Access the model in MLflow UI and SageMaker Studio UI

After the execution completes, you can find the trained model in both the MLflow UI and SageMaker Studio UI.

To view the model in the MLflow UI, select the "Models" tab:

To view the model in SageMaker Studio UI, you will need to navigate to SageMaker Studio:

1. Choose a domain and launch Studio from one of the user profiles associated with it
2. Select "Models" in the menu to see the SageMaker Model Registry. From here you will see your sm-job-experiment-model model

Deploy MLflow Model to SageMaker

Define of the Schema of the sklearn model

Model Builder requires the definition of the model schema, this is the input and output of the model. In this case it is a [4x1] vector for the input and an integer for the output.

Build and deploy the model

Inference on Deployed Model

Cleanup Resources

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.