Set up AWS Authentication for SageMaker Deployment

This topic describes how to set up IAM roles to allow you to deploy MLflow models to AWS SageMaker. It is possible to use access keys for an AWS user with similar permissions as the IAM role specified here, but Databricks recommends using IAM roles to give a cluster permission to deploy to SageMaker.

Step 1: Create an AWS IAM role and attach SageMaker permission policy

  1. In the AWS console, go to the IAM service.

  2. Click the Roles tab in the sidebar.

  3. Click Create role.

    1. Under Select type of trusted entity, select AWS service.

    2. Under Choose the service that will use this role, click the EC2 service.

      ../../../_images/sagemaker-service.png
    3. Click Next: Permissions.

  4. In the Attach permissions policies screen, select AmazonSageMakerFullAccess.

    ../../../_images/sagemaker-permissions.png
  5. Click Next: Review.

  6. In the Role name field, enter a role name.

  7. Click Create role.

  8. In the Roles list, click the role name.

    ../../../_images/sagemaker-role-summary.png

Make note of your Role ARN, which is of the format arn:aws:iam::<account-id>:role/<role-name>.

Step 2: Add an inline policy for access to SageMaker deployment resources

Add a policy to the role.

  1. Click Inline policy.

  2. Paste in the following JSON definition:

    {
      "Statement": [
        {
          "Action": [
            "s3:PutObjectAcl",
            "s3:PutObjectTagging"
          ],
          "Resource": [
            "arn:aws:s3:::mlflow-sagemaker-*-<account-id>",
            "arn:aws:s3:::mlflow-sagemaker-*-<account-id>/*"
          ],
          "Effect": "Allow"
        },
        {
          "Action": [
            "iam:GetRole"
          ],
          "Resource": [
            "arn:aws:iam::<account-id>:role/<role-name>"
          ],
          "Effect": "Allow"
        },
        {
          "Action": [
            "ecr:DescribeRepositories"
          ],
          "Resource": [
            "arn:aws:ecr:*:<account-id>:repository/*"
          ],
          "Effect": "Allow"
        }
      ]
    }
    

These permissions are required to allow the Databricks cluster to:

  1. Obtain the new role’s canonical ARN.
  2. Upload permission-scoped objects to S3 for use by SageMaker endpoint servers.

The role’s permissions will look like:

../../../_images/sagemaker-policy.png

Step 3: Update the role’s trust policy

Add iam:AssumeRole access to sagemaker.amazonaws.com.

  1. Go to Role Summary > Trust relationships > Edit trust relationship.

  2. Paste and save the following JSON:

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

Your role’s trust relationships should resemble the following:

../../../_images/sagemaker-policy-summary.png

Step 4: Allow your Databricks workspace AWS role to pass the role

  1. Go to your Databricks workspace AWS role.

  2. Click Inline policy.

  3. Paste and save the following JSON definition:

    {
      "Version": "2012-10-17",
      "Statement": [
        {
          "Action": [
            "iam:PassRole"
          ],
          "Resource": [
            "arn:aws:iam::<account-id>:role/<role-name>"
          ],
          "Effect": "Allow"
        }
      ]
    }
    
where account-id is the ID of the account running the AWS SageMaker service and role-name is the role you defined in Step 1.

Step 5: Create a Databricks cluster IAM role

  1. In your Databricks workspace, go to: Admin Console > IAM Roles > Add IAM Role.

  2. Paste in the instance profile ARN associated with the AWS role you created. This ARN is of the form arn:aws:iam::<account-id>:instance-profile/<role-name> and can be found in the AWS console:

    ../../../_images/sagemaker-inline-policy.png
  3. Click the Add button.

    ../../../_images/sagemaker-iam-role.png

For details, see Step 5: Add the S3 IAM role to Databricks in Secure Access to S3 Buckets Using IAM Roles.