Integrate MLflow with your environment - Amazon SageMaker
Install MLflow and the AWS MLflow plugin Connect to your MLflow Tracking Server

Integrate MLflow with your environment

The following page describes how to get started with the MLflow SDK and the AWS MLflow plugin within your development environment. This can include local IDEs or a Jupyter Notebook environment within Studio or Studio Classic.

Amazon SageMaker uses an MLflow plugin to customize the behavior of the MLflow Python client and integrate AWS tooling. The AWS MLflow plugin authenticates API calls made with MLflow using AWS Signature Version 4. The AWS 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.

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 AWS MLflow plugin

Within your development environment, install both MLflow and the AWS 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)