Simplify entry management and auditing for Amazon SageMaker Studio utilizing trusted identification propagation

AWS helps trusted identification propagation, a characteristic that permits AWS companies to securely propagate a person’s identification throughout service boundaries. With trusted identification propagation, you may have fine-grained entry controls based mostly on a bodily person’s identification relatively than counting on IAM roles. This integration permits for the implementation of entry management by way of companies corresponding to Amazon S3 Entry Grants and maintains detailed audit logs of person actions throughout supported AWS companies corresponding to Amazon EMR. Moreover, it helps long-running person background classes for coaching jobs, so you possibly can sign off of your interactive ML utility whereas the background job continues to run.

Amazon SageMaker Studio now helps trusted identification propagation, providing a strong answer for enterprises in search of to reinforce their ML system safety. By integrating trusted identification propagation with SageMaker Studio, organizations can simplify entry administration by granting permissions to present AWS IAM Identification Heart identities.

On this submit, we discover easy methods to allow and use trusted identification propagation in SageMaker Studio, demonstrating its advantages by way of sensible use instances and implementation pointers. We stroll by way of the setup course of, focus on key concerns, and showcase how this characteristic can remodel your group’s method to safety and entry controls.

Answer overview

On this part, we overview the structure for the proposed answer and the steps to allow trusted identification propagation to your SageMaker Studio area.

The next diagram exhibits the interplay between the completely different parts that enable the person’s identification to propagate from their identification supplier and IAM Identification Heart to downstream companies corresponding to Amazon EMR and Amazon Athena.

With a trusted identification propagation-enabled SageMaker Studio area, customers can entry knowledge throughout supported AWS companies utilizing their finish person identification and group membership, along with entry allowed by their area or person execution position. As well as, API calls from SageMaker Studio notebooks and supported AWS companies and Amazon SageMaker AI options log the person identification in AWS CloudTrail. For a listing of supported AWS companies and SageMaker AI options, see Trusted identification propagation structure and compatibility. Within the following sections, we present easy methods to allow trusted identification propagation to your area.

This answer applies for SageMaker Studio domains arrange utilizing IAM Identification Heart as the strategy of authentication. In case your area is ready up utilizing IAM, see Implement user-level entry management for multi-tenant ML platforms on Amazon SageMaker AI for finest practices on managing and scaling entry management.

Conditions

To comply with together with this submit, you should have the next:

• An AWS account with a corporation occasion of IAM Identification Heart configured by way of AWS Organizations
• Administrator permissions (or elevated permissions permitting modification of IAM principals, and SageMaker administrator entry to create and replace domains)

Create or replace the SageMaker execution position

For trusted identification propagation to work, the SageMaker execution position (area and person profile execution position), ought to enable the sts:SetContext permissions, along with sts:AssumeRole, in its belief coverage. For a brand new SageMaker AI area, create a website execution position by following the directions in Create execution position. For present domains, comply with the directions in Get your execution position to search out the person or area’s execution position.

Subsequent, to replace the belief coverage for the position, full the next steps:

1. Within the navigation pane of the IAM console, select Roles.
2. Within the checklist of roles in your account, select the area or person execution position.
3. On the Belief relationships tab, select Edit belief coverage.
4. Replace the belief coverage with the next assertion:

{
  "Model": "2012-10-17",
  "Assertion": [
     .....
    {
      "Effect": "Allow",
      "Principal": {
        "Service": [
          "sagemaker.amazonaws.com",
        ]
      },
      "Motion": [
        "sts:AssumeRole",
        "sts:SetContext"
      ],
      "Situation": {
	"aws:SourceAccount": ""
         }
       }
    }
  ]
}

5. Select Replace coverage to save lots of your modifications.

Trusted identification propagation solely works for personal areas on the time of launch.

Create a SageMaker AI area with trusted identification propagation enabled

SageMaker AI domains utilizing IAM Identification Heart for authentication can solely be arrange in the identical AWS Area because the IAM Identification Heart occasion. To create a brand new SageMaker area, comply with the steps in Use customized setup for Amazon SageMaker AI. For Trusted identification propagation, choose Allow trusted identification propagation for all customers on this area, and proceed with the remainder of the setup to create a website and assign customers and teams, selecting the position you created within the earlier step.

Replace an present SageMaker AI area

You may as well replace your present SageMaker AI area to allow trusted identification propagation. You possibly can allow trusted identification propagation even whereas the area or person has lively SageMaker Studio purposes. Nevertheless, for the modifications to be utilized, the lively purposes should be restarted. You should utilize the EffectiveTrustedIdentityPropagationStatus area within the response to the DescribeApp API for working purposes to find out if the appliance has trusted identification propagation enabled.

To allow trusted identification propagation for the area utilizing the SageMaker AI console, select Edit beneath Authentication and permissions on the Area settings tab.

For Trusted identification propagation, choose Allow trusted identification propagation for all customers on this area, and select Submit to save lots of the modifications.

(Non-obligatory) Replace person background session configuration in IAM Identification Heart

IAM Identification Heart now helps working person background classes, and the session period is ready by default to 7 days. With background classes, customers can launch long-running SageMaker coaching jobs that assume the person’s identification context together with the SageMaker execution position. As an administrator, you possibly can allow or disable person background classes, and modify the session period for person background classes. As of the time of writing, the utmost session period you could set for person background classes is 90 days. The person’s session is stopped on the finish of the desired period, and consequently, the coaching job may even fail on the finish of the session period.

To disable or replace the session period, navigate to the IAM Identification Heart console, select Settings within the navigation pane, and select Configure beneath Session period.

For Consumer background classes, choose Allow person background classes and use the dropdown to alter the session period. If person background classes are disabled, the person should be logged in in the course of the coaching job; in any other case, the coaching job will fail as soon as the person logs out. Updating this configuration doesn’t have an effect on present working classes and solely applies to newly created person background classes. Select Save to save lots of your settings.

Use instances

Think about you’re an enterprise with lots of and even hundreds of customers, every requiring various ranges of entry to knowledge throughout a number of groups. You’re answerable for sustaining an AI/ML system on SageMaker AI and managing entry permissions throughout numerous knowledge sources corresponding to Amazon Easy Storage Service (Amazon S3), Amazon Redshift, and AWS Lake Formation. Historically, this has concerned sustaining advanced IAM insurance policies for customers, companies, and assets, together with bucket insurance policies the place relevant. This method will not be solely tedious but additionally makes it difficult to trace and audit knowledge entry with out sustaining a separate position for every person.

That is exactly the state of affairs that trusted identification propagation goals to handle. With trusted identification propagation assist, now you can keep service-specific roles with minimal permissions, corresponding to s3:GetDataAccess or LakeFormation:GetDataAccess, together with extra permissions to begin jobs, view job statuses, and carry out different mandatory duties. For knowledge entry, you possibly can assign fine-grained insurance policies on to particular person customers. As an example, Jane may need learn entry to buyer knowledge and full entry to gross sales and pricing knowledge, whereas Laura would possibly solely have learn entry to gross sales tendencies. Each Jane and Laura can assume the identical SageMaker AI position to entry their SageMaker Studio purposes, whereas sustaining separate knowledge entry permissions based mostly on their particular person identities.Within the following sections, we discover how this may be achieved for frequent use instances, demonstrating the facility and suppleness of trusted identification propagation in simplifying knowledge entry administration whereas sustaining sturdy safety and auditability.

State of affairs 1: Experiment with Amazon S3 knowledge in notebooks

S3 Entry Grants present a simplified method to handle knowledge entry at scale. Not like conventional IAM roles and insurance policies that require an in depth information of IAM ideas, and frequent coverage updates as new assets are added, with S3 Entry Grants, you possibly can outline entry to knowledge based mostly on acquainted database-like grants that mechanically scale along with your knowledge. This method considerably reduces the operational overhead of managing hundreds of IAM insurance policies and bucket insurance policies, and overcomes the restrictions of IAM permissions, whereas strengthening safety by way of entry patterns. If you happen to don’t have S3 Entry Grants arrange, see Create an S3 Entry Grant occasion to get began. For detailed structure and use instances, you can even consult with Scaling knowledge entry with Amazon S3 Entry Grants. After you may have arrange S3 Entry Grants, you possibly can grant entry to your datasets to customers based mostly on their identification in IAM Identification Heart.

To make use of S3 Entry Grants from SageMaker Studio, replace the next IAM roles with insurance policies and belief insurance policies.

For the area or person execution position, add the next inline coverage:

{
    "Model": "2012-10-17",
    "Assertion": [
        {
            "Sid": "AllowDataAccessAPI",
            "Effect": "Allow",
            "Action": [
                "s3:GetDataAccess"
            ],
            "Useful resource": [
                "arn:aws:s3:::access-grants/default"
            ]
        },
        {
            "Sid": "RequiredForTIP",
            "Impact": "Enable",
            "Motion": "sts:SetContext",
            "Useful resource": "arn:aws:iam:::position/"
        }
    ]
}

Be certain the S3 Entry Grants position’s belief coverage permits the sts:SetContext motion along with sts:AssumeRole. The next is a pattern belief coverage:

{
    "Model": "2012-10-17",
    "Assertion": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": [
                    "access-grants.s3.amazonaws.com"
                ]
            },
            "Motion": [
                "sts:AssumeRole",
                "sts:SetContext"
            ],
            "Situation": {
                "StringEquals": {
                    "aws:SourceArn": "arn:aws:s3:::access-grants/default"
                }
            }
        }
    ]

import boto3
from botocore.config import Config

def get_access_grant_credentials(account_id: str, goal: str, 
                                 permission: str="READ"):
    s3control = boto3.consumer('s3control')
    response = s3control.get_data_access(
        AccountId=account_id,
        Goal=goal,
        Permission=permission
    )
    return response['Credentials']

def create_s3_client_from_credentials(credentials) -> boto3.consumer:
    return boto3.consumer(
        's3',
        aws_access_key_id=credentials['AccessKeyId'],
        aws_secret_access_key=credentials['SecretAccessKey'],
        aws_session_token=credentials['SessionToken']
    )

# Create consumer
credentials = get_access_grant_credentials('',
                                        "s3:////")
s3 = create_s3_client_from_credentials(credentials)

# Will succeed
s3.list_objects(Bucket="", Prefix="")

# Will fail
s3.list_objects(Bucket="", Prefix="")

State of affairs 2: Entry Lake Formation by way of Athena

Lake Formation supplies centralized governance and fine-grained entry management administration for knowledge saved in Amazon S3 and metadata within the AWS Glue Information Catalog. The Lake Formation permission mannequin operates together with IAM permissions, providing granular controls on the database, desk, column, row, and cell ranges. This dual-layer safety mannequin supplies complete knowledge governance whereas sustaining flexibility in entry patterns.

Information ruled by way of Lake Formation might be accessed by way of numerous AWS analytics companies. On this state of affairs, we display utilizing Athena, a serverless question engine that integrates seamlessly with Lake Formation’s permission mannequin. For different companies like Amazon EMR on EC2, ensure that the useful resource is configured to assist trusted identification propagation, together with establishing safety configurations and ensuring the EMR cluster is configured with IAM roles that assist trusted identification propagation.

The next directions assume that you’ve already arrange Lake Formation. If not, see Arrange AWS Lake Formation and comply with the AWS Lake Formation tutorials to arrange Lake Formation and herald your knowledge.

Full the next steps to entry your ruled knowledge in trusted identification propagation-enabled SageMaker Studio notebooks utilizing Athena:

1. Combine Lake Formation with IAM Identification Heart by following the directions in Integrating IAM Identification Heart. At a excessive degree, this consists of creating an IAM position permitting creating and updating utility configurations in Lake Formation and IAM Identification Heart, and offering the one sign-on (SSO) occasion ID.
2. Grant permissions to the IAM Identification Heart person to the related assets (database, desk, row or column) utilizing Lake Formation. See Granting permissions on Information Catalog assets directions.
3. Create an Athena workgroup that helps trusted identification propagation by following directions in Create a workgroup and selecting IAM Identification Heart as the strategy of authentication. Be certain the person has entry to write down to the question outcomes location supplied right here utilizing S3 Entry Grants, as a result of Athena makes use of entry grants by default when selecting IAM Identification Heart because the authentication technique.
4. Replace the Athena workgroup’s IAM position with the next belief coverage (add sts:SetContext to the prevailing belief coverage). You will discover the IAM position by selecting the workgroup you created earlier and searching for Function title.

{
    "Model": "2012-10-17",
    "Assertion": [
        {
            "Sid": "AthenaTrustPolicy",
            "Effect": "Allow",
            "Principal": {
                "Service": "athena.amazonaws.com"
            },
            "Action": [
                "sts:AssumeRole",
                "sts:SetContext"
            ],
            "Situation": {
                "StringEquals": {
                    "aws:SourceAccount": ""
                },
                "ArnLike": {
                    "aws:SourceArn": "arn:aws:athena:::workgroup/"
                }
            }
        }
    ]
}

The setup is now full. Now you can launch SageMaker Studio utilizing an IAM Identification Heart person, launch a JupyterLab or Code Editor utility, and question the database. See the next instance code to get began:

import time
import boto3
import pandas as pd
athena_client = boto3.consumer("athena")

database = ""
desk = ""
question = f"SELECT * FROM {database}.{desk}"
output_location = "s3:///queries"  # bucket title and site from Step 3

response = athena_client.start_query_execution(
    QueryString=question,
    QueryExecutionContext={'Database': database},
    ResultConfiguration={'OutputLocation': output_location}
)

# Get the question execution ID
query_execution_id = response['QueryExecutionId']

# await question to finish
whereas True:
    query_status = athena_client.get_query_execution(QueryExecutionId=query_execution_id)
    standing = query_status['QueryExecution']['Status']['State']
    if standing in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
        break
    time.sleep(1)

# If the question succeeded, fetch and show outcomes
if standing == 'SUCCEEDED':
    outcomes = athena_client.get_query_results(QueryExecutionId=query_execution_id)
    
    # Extract column names and knowledge
    columns = [col['Name'] for col in outcomes['ResultSet']['ResultSetMetadata']['ColumnInfo']]
    knowledge = []
    for row in outcomes['ResultSet']['Rows'][1:]:  # Skip the header row
        knowledge.append([field.get('VarCharValue', '') for field in row['Data']])
    
    # Create a pandas DataFrame
    df = pd.DataFrame(knowledge, columns=columns)
    
    # Show the primary few rows
    print(df.head())
else:
    print(f"Question failed with standing: {standing}")

State of affairs 3: Create a coaching job supported with person background classes

For a trusted identification propagation-enabled area, a person background session is a session that continues to run even when the end-user has logged out of their interactive session corresponding to JupyterLab purposes in SageMaker Studio. For instance, the person can provoke a coaching job from their SageMaker Studio house, and the job can run within the background for days or perhaps weeks whatever the person’s exercise, and use the person’s identification to entry knowledge and log audit trails. In case your area doesn’t have trusted identification propagation enabled, you possibly can proceed to run coaching jobs and processing jobs as earlier than; nonetheless, if trusted identification propagation is enabled, ensure that your person background session time is up to date to mirror the period of your coaching jobs, as a result of the default is ready mechanically to 7 days. When you have enabled person background classes, replace your SageMaker Studio area or person’s execution position with the next permissions to supply a seamless expertise for knowledge scientists:

{
    "Model": "2012-10-17",
    "Assertion": [
        {
            "Sid": "AllowDataAccessAPI",
            "Effect": "Allow",
            "Action": [
                "s3:GetDataAccess",
                "s3:GetAccessGrantsInstanceForPrefix"
            ],
            "Useful resource": [
                "arn:aws:s3:::access-grants/default"
            ]
        },
        {
            "Sid": "RequiredForTIP",
            "Impact": "Enable",
            "Motion": "sts:SetContext",
            "Useful resource": "arn:aws:iam:::position/"
        }
    ]
}

With this setup, a knowledge scientist can use an Amazon S3 location that they’ve entry to by way of S3 Entry Grants. SageMaker mechanically appears to be like for knowledge entry utilizing S3 Entry Grants and falls again to the job’s IAM position in any other case. For instance, within the following SDK name to create the coaching job, the person supplies the S3 Amazon URI the place the information is saved, they’ve entry to it by way of S3 Entry Grants, and so they can run this job with out extra setup:

    response = sm.create_training_job(
        TrainingJobName=training_job_name,
        AlgorithmSpecification={
            'TrainingImage': '763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04',
            'TrainingInputMode': 'File',
            ...
                    RoleArn='arn:aws:iam:::position/tip-domain-role',
        InputDataConfig=[
            {
                'ChannelName': 'training',
                'DataSource': {
                    'S3DataSource': {
                        'S3DataType': 'S3Prefix',
                        'S3Uri': 's3:///',
                        'S3DataDistributionType': 'FullyReplicated'
                    }
                },
                'CompressionType': 'None',
                'RecordWrapperType': 'None'
            },
            ...
        }

(Optional) View and manage user background sessions on IAM Identity Center

When training jobs are run as user background sessions, you can view these sessions as user background sessions on IAM Identity Center. The administrator can view a list of all user background sessions and optionally stop a session if the user has left the team, for example. When the user background session is ended, the training job subsequently fails.

To view a list of all user background sessions, on the IAM Identity Center console, choose Users and choose the user you want view the user background sessions for. Choose the Active sessions tab to view a list of sessions. The user background session can be identified by the Session type column, which shows if the session is interactive or a user background session. The list also shows the job’s Amazon Resource Name (ARN) under the Used by column.

To end a session, select the session and choose End sessions.

You will be prompted to confirm the action. Enter confirm to confirm that you want to end the session and choose End sessions to stop the user background session.

Scenario 4: Auditing using CloudTrail

After trusted identity propagation is enabled for your domain, you can now track the user that performed specific actions through CloudTrail. To try this out, log in to SageMaker Studio, and create and open a JupyterLab space. Open a terminal and enter aws s3 ls to list the available buckets in your Region.

On the CloudTrail console, choose Event history in the navigation pane. Update the Lookup attributes to Event name and in the search box, enter ListBuckets. You should see a list of events, as shown in the following screenshot (it might take up to 5 minutes for the logs to be available in CloudTrail).

Choose the event to view its details (verify the user name is SageMaker if you have also listed buckets through the AWS console or APIs). In the event details, you should be able to see an additional field called onBehalfOf that has the user’s identity.

Supported services and SageMaker AI features called from a trusted identity propagation-enabled SageMaker Studio domain will have the OnBehalfOf field in CloudTrail.

Clean up

If you have created a SageMaker Studio domain for the purposes of trying out trusted identity propagation, delete the domain and its associated Amazon Elastic File System (Amazon EFS) volume to avoid incurring additional charges. Before deleting a domain, you must delete all the users and their associated spaces and applications. For detailed instructions, see Stop and delete your Studio running applications and spaces.

If you created a SageMaker training job, they are ephemeral, and the compute is shut down automatically when the job is complete.

Athena is a serverless analytics service that charges per query billing. No cleanup is necessary, but for best practices, delete the workgroup to remove unused resources.

Conclusion

In this post, we showed you how to enable trusted identity propagation for SageMaker AI domains that use IAM Identity Center as the mode of authentication. With trusted identity propagation, administrators can manage user authorization to other AWS services through the user’s physical identity in conjunction with IAM roles. Administrators can streamline permissions management by maintaining a single domain execution role and manage granular access to other AWS services and data sources through the user’s identity. In addition, trusted identity propagation supports auditing, so administrators can track user activity without the need for managing a role for each user profile.

To learn more about enabling this feature and its use cases, see Trusted identity propagation use cases and Trusted identity propagation with Studio. This post covered a subset of supported applications; we encourage you to check out the documentation and choose the services that best serve your use case and share your feedback!

About the authors

Amit Shyam Jaisinghani is a Software Engineer on the SageMaker Studio team at Amazon Web Services, and he earned his Master’s degree in Computer Science from Rochester Institute of Technology. Since joining Amazon in 2019, he has built and enhanced several AWS services, including AWS WorkSpaces and Amazon SageMaker Studio. Outside of work, he explores hiking trails, plays with his two cats, Missy and Minnie, and enjoys playing Age of Empire.

Durga Sury is a Senior Solutions Architect at Amazon SageMaker, where she helps enterprise customers build secure and scalable AI/ML systems. When she’s not architecting solutions, you can find her enjoying sunny walks with her dog, immersing herself in murder mystery books, or catching up on her favorite Netflix shows.

Khushboo Srivastava is a Senior Product Manager for Amazon SageMaker. She enjoys building products that simplify machine learning workflows for customers, and loves playing with her 1-year old daughter.

Krishnan Manivannan is a Senior Software Engineer at Amazon Web Services and a founding member of the SageMaker AI API team. He has 8 years of experience in the architecture and security of large-scale machine learning services. His specialties include API design, service scalability, identity and access management, and inventing new approaches for building and operating distributed systems. Krishnan has led multiple engineering efforts from design through global launch, delivering reliable and secure systems for customers worldwide.