How to Setup Amazon SageMaker with MLflow

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.

Updates and Imports

Imports

Session variables

MLflow Permissions

IAM Role for the MLflow Tracking Server

To run the next cell, make sure the IAM role used while running this notebook has permission to create an IAM Role. The iam:CreateRole, iam:CreatePolicy, iam:ListPolicies, and iam:AttachRolePolicy action must be allowed by the notebook execution role's policy.

If you are running this notebook from SageMaker Studio, you can update your notebook execution role through the following steps:

1. Navigate to the AWS Console and select the Domain you are using
2. Under the Domain, select the User Profile you are using. You will see the Execution Role listed there.
3. Navigate to the IAM Console, search for the Execution Role under "Roles", and update your role with a policy that allows the iam:CreateRole, iam:CreatePolicy, iam:ListPolicies, and iam:AttachRolePolicy actions.

If you are not using a SageMaker Studio Notebook, confirm that the role you have used to configure your AWS CLI has appropriate permissions to create an IAM role and attach a policy to it.

Here is an example of an inline policy you can add to your role -

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "Statement1",
            "Effect": "Allow",
            "Action": [
                "iam:ListPolicies",
                "iam:CreatePolicy",
                "iam:CreateRole",
                "iam:AttachRolePolicy"
            ],
            "Resource": [
                "*"
            ]
        }
    ]
}

Note that your SageMaker execution role should have the following permissions to call Mlflow REST APIs:

{
    "Version": "2012-10-17",    
    "Statement": [        
        {            
            "Effect": "Allow",            
            "Action": [
                "sagemaker-mlflow:*",
                "sagemaker:CreateMlflowTrackingServer",
                "sagemaker:UpdateMlflowTrackingServer",
                "sagemaker:DeleteMlflowTrackingServer",
                "sagemaker:StartMlflowTrackingServer",
                "sagemaker:StopMlflowTrackingServer",
                "sagemaker:CreatePresignedMlflowTrackingServerUrl"
            ],            
            "Resource": "*"        
        }        
    ]
}

Create MLflow Tracking Server

Install the MLflow SDK and our MLflow AWS Plugin

MLflow tracking test

Connect to tracking server

Log a metric

See results in MLflow UI. You can either launch the MLflow UI from within SageMaker Studio, or generate a pre-signed URL like this:

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.