Set up IAM permissions for MLflow

You must configure the necessary IAM service roles to get started with MLflow in Amazon SageMaker AI.

If you create a new Amazon SageMaker AI domain to access your experiments in Studio, you can configure the necessary IAM permissions during domain setup. For more information, see Set up MLflow IAM permissions when creating a new domain.

To set up permissions using the IAM console, see Create necessary IAM service roles in the IAM console.

You must configure authorization controls for sagemaker-mlflow actions. You can optionally define more granular authorization controls to govern action-specific MLflow permissions. For more information, see Create action-specific authorization controls.

Set up MLflow IAM permissions when creating a new domain

When setting up a new Amazon SageMaker AI domain for your organization, you can configure IAM permissions for your domain service role through the Users and ML Activities settings.

The following MLflow ML activities are available in Amazon SageMaker Role Manager:

For more information about ML activities in Role Manager, see ML activity reference.

Create necessary IAM service roles in the IAM console

If you did not create or update your domain service role, you must instead create the following service roles in the IAM console in order to create and use an MLflow Tracking Server:

IAM policies for the tracking server IAM service role

The tracking server IAM service role is used by the tracking server to access the resources it needs such as Amazon S3 and the SageMaker Model Registry.

When creating the tracking server IAM service role, use the following IAM trust policy:

{
     "Version": "2012-10-17",
     "Statement": [
         {
             "Effect": "Allow",
             "Principal": {
                 "Service": [                     
                      "sagemaker.amazonaws.com"
                 ]
             },
             "Action": "sts:AssumeRole"
         }
     ]
 }

In the IAM console, add the following permissions policy to your tracking server service role:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "s3:Get*",
                "s3:Put*",
                "s3:List*",
                "sagemaker:AddTags",
                "sagemaker:CreateModelPackageGroup",
                "sagemaker:CreateModelPackage",
                "sagemaker:UpdateModelPackage",
                "sagemaker:DescribeModelPackageGroup"
            ],
            "Resource": "*"
        }
    ]
}

IAM policy for the SageMaker AI IAM service role

The SageMaker AI service role is used by the client accessing the MLflow Tracking Server and needs permissions to call MLflow REST APIs. The SageMaker AI service role also needs SageMaker API permissions to create, view update, start, stop, and delete tracking servers.

You can create a new role or update an existing role. The SageMaker AI service role needs the following policy:

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

Create action-specific authorization controls

You must set up authorization controls for sagemaker-mlflow, and can optionally configure action-specific authorization controls to govern more granular MLflow permissions that your users have on an MLflow Tracking Server.

The following command creates a file called mlflow-policy.json that provides your tracking server with IAM permissions for all available SageMaker AI MLflow actions. You can optionally limit the permissions a user has by choosing the specific actions you want that user to perform. For a list of available actions, see IAM actions supported for MLflow.

# Replace "Resource":"*" with "Resource":"TrackingServerArn" 
# Replace "sagemaker-mlflow:*" with specific actions

printf '{
    "Version": "2012-10-17",    
    "Statement": [        
        {            
            "Effect": "Allow",            
            "Action": "sagemaker-mlflow:*",            
            "Resource": "*"        
        }        
    ]
}' > mlflow-policy.json

Use the mlflow-policy.json file to create an IAM policy using the AWS CLI.

aws iam create-policy \
  --policy-name MLflowPolicy \
  --policy-document file://mlflow-policy.json

Retrieve your account ID and attach the policy to your IAM role.

# Get your account ID
aws sts get-caller-identity

# Attach the IAM policy using your exported role and account ID
aws iam attach-role-policy \
  --role-name $role_name \
  --policy-arn arn:aws:iam::123456789012:policy/MLflowPolicy

IAM actions supported for MLflow

The following SageMaker AI MLflow actions are supported for authorization access control: