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 AI 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.

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.

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)