Implementing Efficient Access Control in Amazon SageMaker AI Environments

Overview of Access Control Challenges in ML Workflows

Strategies for Efficient Permission Management

Implementing Attribute-Based Access Control (ABAC)

Key Concepts in IAM for SageMaker

Prerequisites for Effective ABAC Implementation

Overview of the Proposed Solution

Access Control in SageMaker Studio

Amazon S3 Access Control Strategies

Managing Secrets Access with AWS Secrets Manager

Controlling Amazon EMR Cluster Access

File System Access Control in SageMaker Training Jobs

Monitoring User Access with Source Identity

Tracking Access to AWS Glue Data Catalog

Best Practices for ABAC and Source Identity Implementation

Conclusion: Achieving User-Level Access Control in SageMaker AI

About the Authors

Managing Access Control in Enterprise Machine Learning Environments

Managing access control in enterprise machine learning (ML) environments can be a daunting challenge, especially in scenarios where multiple teams are accessing shared resources within a single Amazon Web Services (AWS) account. While Amazon SageMaker Studio offers user-level execution roles, this approach often becomes unwieldy as organizations scale. In this post, we’ll explore strategies using attribute-based access control (ABAC) patterns to enable granular user access management, focusing on proven best practices to enhance security and compliance without compromising operational efficiency.

Challenges with Resource Isolation Across Workloads

In organizations with centralized account structures, such as those in highly regulated industries like finance and healthcare, a single ML platform team manages a vast infrastructure serving hundreds of data science teams. This centralized approach facilitates consistent governance and resource utilization, but it introduces challenges in maintaining workload isolation between teams.

For instance, platform teams using SageMaker AI can establish dedicated SageMaker Studio domains for each business unit. While this isolates resources effectively, applying team- or domain-level roles often leads to security issues and cumbersome auditing processes. Conversely, maintaining user-level roles can result in creating excessive IAM roles, potentially hitting service quotas.

To address these challenges, we demonstrate how to implement ABAC with IAM policy variables that allow user-level access controls while retaining domain-level execution roles. This method provides a more scalable approach to IAM in SageMaker.

Key Concepts

Two critical IAM concepts underpin our solution: source identity and context keys.

• Source Identity: A custom string passed during role assumption that identifies the user or application performing actions. This identity aids in tracking and auditing actions in AWS CloudTrail.

• Context Keys: SageMaker Studio supports specific condition context keys like sagemaker:DomainId and sagemaker:UserProfileName, which allow for dynamic ABAC policies based on the user’s identity and domain.

Prerequisites

Before implementing an ABAC-based solution, ensure your SageMaker Studio domain meets the following criteria:

1. Roles used with SageMaker AI incorporate the sts:SetSourceIdentity permission in their trust policy.

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

2. Update your domain settings to use user profile names for identity configuration.

   aws sagemaker update-domain --domain-id <domain_id> --domain-settings-for-update "ExecutionRoleIdentityConfig=USER_PROFILE_NAME"

3. Grant the sagemaker:AddTags permission to roles associated with SageMaker Studio.

Solution Overview

This post details how to apply IAM policy variables and source identity for scalable and user-level access control in SageMaker AI. This approach allows you to:

• Implement user-level access control without managing multiple IAM roles.
• Enforce resource isolation between users.
• Uphold least privilege principles across various AWS resources.

Common Scenarios for Implementing Access Control

1. SageMaker AI Resource Control: Prevent users from accessing or modifying each other’s resources using the sagemaker:DomainId and sagemaker:UserProfileName keys. For example:

   {
      "Sid": "TrainingJobPermissions",
      "Effect": "Allow",
      "Action": ["sagemaker:StopTrainingJob", "sagemaker:DescribeTrainingJob"],
      "Resource": "arn:aws:sagemaker:{region}:{account_number}:training-job/*",
      "Condition": {
          "StringLike": {
              "sagemaker:ResourceTag/sagemaker:user-profile-arn": "arn:aws:sagemaker:::user-profile/${sagemaker:DomainId}/${sagemaker:UserProfileName}"
          }
      }
   }

2. Amazon S3 Access Control: Limiting access to user-specific S3 prefixes can ensure data security. An example policy is shown below:

   {
      "Version": "2012-10-17",
      "Statement": [
          {
              "Sid": "ListBucket",
              "Effect": "Allow",
              "Action": "s3:ListBucket",
              "Resource": "arn:aws:s3:::my_bucket",
              "Condition": {
                  "StringLikeIfExists": {
                      "s3:prefix": ["my_domain/users/${aws:SourceIdentity}/*"]
                  }
              }
          },
          {
              "Sid": "AccessBucketObjects",
              "Effect": "Allow",
              "Action": ["s3:GetObject", "s3:PutObject", "s3:DeleteObject"],
              "Resource": [
                  "arn:aws:s3:::my_bucket/my_domain/users/${aws:SourceIdentity}/*"
              ]
          }
      ]
   }

3. Secrets Manager Access: You can restrict Secrets Manager access to user-specific secrets:

   {
      "Version": "2012-10-17",
      "Statement": [
          {
              "Sid": "UserSpecificSecretsAccess",
              "Effect": "Allow",
              "Action": "secretsmanager:GetSecretValue",
              "Resource": "arn:aws:secretsmanager:::secret:user-secrets/${aws:SourceIdentity}/*"
          }
      ]
   }

Monitoring User Access with Source Identity

The source identity not only helps manage access control but also enhances monitoring capabilities. By propagating the user profile name to CloudTrail logs, administrators can precisely track user actions, improving visibility for compliance audits.

Best Practices for Implementing ABAC

To effectively implement ABAC in your environment:

• Use Consistent Naming Conventions: Align resource names and tags for reliable policy referencing.
• Enforce Least Privilege Access: Grant only the permissions necessary for tasks, leveraging AWS managed policies as a starting point.
• Audit User Access Regularly: Utilize CloudTrail logs to track user activity associated with source identities.
• Standardize Identity-Based Policies: Implement consistent policies using context keys for simplified management.

Conclusion

In summary, implementing user-level access control in SageMaker Studio can be achieved without the complexities of managing individual IAM roles. By combining resource tags, context keys, and source identity propagation, administrators can create scalable, dynamic IAM policies that align with best practices for security and compliance. These strategies not only enhance administration efficiency but also bolster security across shared execution roles.

About the Authors

• Durga Sury is a Senior Solutions Architect at Amazon SageMaker, focusing on building secure, scalable AI/ML platforms.
• Itziar Molina Fernandez is a Machine Learning Engineer at AWS Professional Services with expertise in large-scale ML solutions.
• Will Parr is a Machine Learning Engineer dedicated to creating scalable ML platforms and impactful generative AI solutions.

Implement these strategies today to optimize your ML access control workflows securely and efficiently!