Set up IAM permissions for MLflow

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

If you create a new Amazon SageMaker 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 AuthZ controls for sagemaker-mlflow actions. You can optionally define more granular AuthZ controls to govern action-specific MLflow permissions. For more information, see Action-specific AuthZ controls.

Set up MLflow IAM permissions when creating a new domain

When setting up a new Amazon SageMaker 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:

Use the following steps to add the MLflow ML activities to your domain service role:

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:

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

To create the tracking server IAM service role, create 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 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": "*"
        }
    ]
}

Create the SageMaker IAM service role

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

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

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

Action-specific AuthZ controls

You must set up AuthZ controls for sagemaker-mlflow, and can optionally configure action-specific AuthZ 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 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 Amazon 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 MLflow actions are supported for AuthZ access control: