Track machine learning experiments using InfinStor MLflow with Amazon SageMaker Studio

MLflow is an open-source platform to manage the machine learning (ML) lifecycle, including experimentation, reproducibility, deployment, and a central model registry. Amazon SageMaker Studio is a fully integrated development environment (IDE) for machine learning. By integrating SageMaker Studio with MLFlow, you can enjoy the benefits of open-standards-based MLflow experiment tracking and model management in SageMaker Studio.

Popular ML frameworks such as Pytorch, Tensorflow, and Xgboost include built-in support for MLflow Experiment Tracking, also known as auto-logging, and MLflow model management. With this solution, you do not need to write special code for experiment tracking or model management. You can also share experiments and models with authorized colleagues.

In this post, I show how to use InfinStor MLflow with Amazon SageMaker Studio to experiment, collaborate, train, and run inferences using this ML platform. SageMaker Studio provides the Notebook and remote IPython kernel portion of the solution, and InfinStor MLflow provides the experiment tracking and model management.

Prerequisites

• AWS account
• Familiarity with Amazon SageMaker Studio
• Familiarity with ML frameworks
• The trial version of InfinStor MLflow, available in AWS Marketplace

Solution walkthrough: Use InfinStor MLflow with Amazon SageMaker Studio

A. Sign up for InfinStor MLflow service in AWS Marketplace

To subscribe to InfinStor Mlflow in AWS Marketplace, do the following:

1. Sign in to your AWS Management Console and navigate to AWS Marketplace or follow this InfinStor MLflow link. In the upper right, choose Try for free.
2. On the Configure your software contract page, review the offer and in the lower right, choose Create contract. When the subscription confirmation dialogue box appears, choose Set up your account.
3. You will be redirected to the InfinStor Service Dashboard. Below the Log in button, choose the Sign up! link.
4. On the registration page, to create your InfinStor administrator account, enter a user name, email, and a password. Choose InfinStor MLflow is a multi-user system with authentication and authorization, and the user you are creating in this page is the InfinStor administrator account. You can add additional data scientist users later on.

B. Configure your AWS account for InfinStor MLflow

InfinStor stores all MLflow artifacts in an artifacts bucket in your own AWS account, and the MLflow Tracking Database is created in your AWS account. InfinStor provides an AWS CloudFormation template to create these resources. To configure your AWS account for InfinStor Mlflow, do the following:

1. Go to the InfinStor service dashboard. Log in using the InfinStor administrator account you created step A.4.
2. In the left navigation, choose Configuration, then Configure AWS.
3. To go to the CloudFormation template, choose Start. You can now start setting up your InfinStor customer account. For more information, refer to Prepare AWS Account for InfinStor.

C. Create data scientist users in InfinStor MLflow

You must create InfinStor MLflow data scientist user accounts for each data scientist in your organization. To do that, follow these steps:

1. Log in to the InfinStor service dashboard using the administrator account you created in step A.4.
2. In the left navigation, choose Users and then Manage Users.
3. To add new users, choose Add Data Scientist user. For more information, refer to InfinStor MLflow SaaS Service Administration.

D. Use SageMaker Studio with InfinStor MLflow

To use InfinStor MLflow with SageMaker Studio, you use the lifecycle configuration feature of Amazon SageMaker. This automatically initializes your IPython kernel for use with InfinStor MLflow. This section has step-by-step instructions for creating this lifecycle configuration. For general information about SageMaker Studio Lifecycle Configuration, refer to Create a Lifecycle Configuration from the SageMaker Console.

InfinStor MLflow enhances your open-source MLflow with authentication and authorization. So users must log in either using an interactive user name and password or using a token downloaded from the InfinStor service dashboard.

Furthermore, you must set the environment variable MLFLOW_TRACKING_URI to infinstor://mlflow.infinstor.com/, so that all MLflow Python calls are sent to the InfinStor MLflow service. To do that and have InfinStor MLflow automatically configured in each Studio notebook that you create, do the following:

1. Create a lifecycle configuration for your user

• 1. In a browser, navigate to the InfinStor MLflow service dashboard and sign in with the user name and password for the InfinStor MLflow data scientist account you created in step C.
  2. In the left navigation, choose Configuration and then Manage Token.
  3. In the section Show AWS SageMaker Studio Lifecycle Script, choose Show Script. Copy the script.

The browser window will display a shell script similar to the following:

#!/bin/bash

set -eux

PACKAGE=infinstor_mlflow_plugin

pip install --upgrade $PACKAGE

mkdir -p /root/.infinstor

set +x

echo 'Token=Custom aaaaaaaa-bbbb-1111-2222-aaaaaaaaaaaa:abcdefghijklmnopqestuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz' > /root/.infinstor/token

set -x

mkdir -p /root/.ipython/profile_default/startup

echo 'import os' > /root/.ipython/profile_default/startup/00-set-mlflow-env.py

echo "os.environ['MLFLOW_TRACKING_URI'] = 'infinstor://mlflow.infinstor.com/'" >> /root/.ipython/profile_default/startup/00-set-mlflow-env.py

1. 4. On the SageMaker console navigation pane, choose Lifecycle configurations, and then choose Create configuration. Choose Jupyter kernel gateway app (not Jupyter server app). Choose Next.
   5. Enter the text that you copied from the InfinStor service dashboard in step D.1.c into the Scripts box.
   6. Enter a name for your lifecycle configuration. I strongly recommend a name that includes your user name because the script is specific to each InfinStor MLflow user. For example, I used jack-infinstor-mlflow.
   7. On the SageMaker console navigation pane, choose Domains, and then choose your SageMaker domain. Choose the user that you want to associate this lifecycle configuration with.
2. Attach the authentication token to your user name

InfinStor lifecycle configuration scripts contain the authentication token for a specific user. Therefore, that token must be attached to the SageMaker user and not to the SageMaker domain. To attach the token to your user name, do the following:

1. 1. In the SageMaker page for the user, choose Edit. You’re redirected to the main page for this specific user in this specific SageMaker domain.
   2. Choose Next.
   3. In the Lifecycle configurations attached to user section, choose Attach.
   4. Select the lifecycle configuration you created in step D.1.f (for this example, jack-infinstor-mlflow) and choose Attach to user.
   5. Select the lifecycle configuration. To ensure that the script is called for each notebook that this user creates, choose Set as default.
   6. Choose Next, Next again, and then choose Submit.

This completes the one-time configuration for using InfinStor MLflow with SageMaker Studio.

E. Create Studio notebooks with InfinStor MLflow enabled

To create SageMaker Studio notebooks with InfinStor MLflow enabled, complete the following steps:

1. On the SageMaker Studio console, choose Open Launcher.
2. Choose Change Environment, then choose the image Data Science 3.0 and kernel Python 3. An incompatibility exists between MLflow and the default Data Science image with Anaconda Individual Edition. Therefore, the best image for use with MLflow is Data Science 3.0.
3. Choose your preferred instance, and make sure that the startup script created and attached to this user name is listed.
4. Choose Create notebook.
5. After the notebook is created, in the notebook cell, enter !mlflow experiments create -n test-exper and run the cell.

You should get a response such as the following:

Created experiment 'test-exper' with id 2

6. The MLflow UI is now available at MLflow UI, and you must log in using your user name and password. The following screenshot shows the MLflow UI. It shows the Experiments tab for test-exper and column headings of Run Name, Created, Duration, Source, and Models. The list is empty, with the notation No runs logged.

You now have a fully functional 24×7 MLflow service that is automatically configured into your SageMaker Studio notebooks.

Cleanup

To stop using the InfinStor service and clean up all resources used by the InfinStor service in your AWS account, follow these steps:

1. Empty the S3 bucket used for MLflow artifact storage by going to your AWS Management Console, choosing the S3 service, and choosing the MLflow artifacts bucket named infinstor-mlflow-artifacts-<your_infinstor_admin_name>. Choose the Empty button.
2. Delete the AWS CloudFormation stack named ForInfinStorMLflow by going to the AWS Management Console’s CloudFormation page, choosing this stack, and choosing the Delete button.

Conclusion

In this post, I showed how to use InfinStor MLflow with SageMaker Studio to automatically track experiments, create models, and share both with authorized colleagues. You can use MLflow with your SageMaker Studio notebooks to create models from Pytorch, Tensorflow, and Xgboost. Get started with a free trial in AWS Marketplace today, and share any feedback and questions in the comments section.

About the author

Jagane Sundar is co-founder of InfinStor, which develops and markets InfinStor MLflow. As part of this platform, Jagane invented InfinSnap snapshots for Cloud Object Stores. His expertise includes Machine Learning Infrastructure, Storage, Cloud, Kubernetes, Hadoop, Containers, Distributed Coordination Algorithms, and old-school Linux software engineering.