Track experiments with MLflow - Amazon SageMaker
Install MLflow and the Amazon MLflow plugin Connect to your MLflow Tracking Server
Services or capabilities described in Amazon Web Services documentation might vary by Region. To see the differences applicable to the China Regions, see Getting Started with Amazon Web Services in China (PDF).

Track experiments with MLflow

Amazon SageMaker uses an MLflow plugin to customize the behavior of the MLflow Python client and integrate Amazon tooling. The Amazon MLflow plugin authenticates API calls made with MLflow using Amazon Signature Version 4. The Amazon MLflow plugin allows you to connect to your MLflow tracking server using the tracking server ARN. For more information about plugins, see MLflow Plugins in the MLflow documentation.

Get started with the MLflow SDK and the Amazon MLflow plugin within your development environment. This can include local IDEs or a Jupyter Notebook environment within Studio or Studio Classic.

Important

Your user IAM permissions within your development environment must have access to any relevant MLflow API actions to successfully run provided examples. For more information, see Set up IAM permissions for MLflow.

For more information about using the MLflow SDK, see Python API in the MLflow documentation.

Install MLflow and the Amazon MLflow plugin

Within your development environment, install both MLflow and the Amazon MLflow plugin.

Note

To see which versions of MLflow are available to use with SageMaker, see Tracking server versions.

pip install mlflow==2.13.2 sagemaker-mlflow==0.1.0

Connect to your MLflow Tracking Server

Use mlflow.set_tracking_uri to connect to a your tracking server from your development environment using its ARN:

import mlflow

arn = "YOUR-TRACKING-SERVER-ARN"

mlflow.set_tracking_uri(arn)