Onboard customers to Amazon SageMaker Studio with Lively Listing group-specific IAM roles

Amazon SageMaker Studio is a web-based built-in growth surroundings (IDE) for machine studying (ML) that permits you to construct, prepare, debug, deploy, and monitor your ML fashions. For provisioning Studio in your AWS account and Area, you first must create an Amazon SageMaker area—a assemble that encapsulates your ML surroundings. Extra concretely, a SageMaker area consists of an related Amazon Elastic File System (Amazon EFS) quantity, a listing of approved customers, and a wide range of safety, utility, coverage, and Amazon Virtual Private Cloud (Amazon VPC) configurations.

When creating your SageMaker area, you possibly can select to make use of both AWS IAM Identity Center (successor to AWS Single Signal-On) or AWS Identity and Access Management (IAM) for consumer authentication strategies. Each authentication strategies have their very own set of use circumstances; on this publish, we give attention to SageMaker domains with IAM Identification Heart, or single sign-on (SSO) mode, because the authentication technique.

With SSO mode, you arrange an SSO consumer and group in IAM Identification Heart after which grant entry to both the SSO group or consumer from the Studio console. At the moment, all SSO customers in a website inherit the area’s execution position. This will likely not work for all organizations. As an illustration, directors could wish to arrange IAM permissions for a Studio SSO consumer based mostly on their Lively Listing (AD) group membership. Moreover, as a result of directors are required to manually grant SSO customers entry to Studio, the method could not scale when onboarding tons of of customers.

On this publish, we offer prescriptive steerage for the answer to provision SSO customers to Studio with least privilege permissions based mostly on AD group membership. This steerage lets you rapidly scale for onboarding tons of of customers to Studio and obtain your safety and compliance posture.

Answer overview

The next diagram illustrates the answer structure.

The workflow to provision AD customers in Studio contains the next steps:

1. Arrange a Studio domain in SSO mode.
2. For every AD group:
   1. Arrange your Studio execution position with acceptable fine-grained IAM insurance policies
   2. Document an entry within the AD group-role mapping Amazon DynamoDB desk.

   Alternatively, you possibly can undertake a naming normal for IAM position ARNs based mostly on the AD group title and derive the IAM position ARN while not having to retailer the mapping in an exterior database.

3. Sync your AD customers and teams and memberships to AWS Identification Heart:
   1. If you happen to’re utilizing an id supplier (IdP) that helps SCIM, use the SCIM API integration with IAM Identification Heart.
   2. In case you are utilizing self-managed AD, it’s possible you’ll use AD Connector.
4. When the AD group is created in your company AD, full the next steps:
   1. Create a corresponding SSO group in IAM Identification Heart.
   2. Affiliate the SSO group to the Studio area utilizing the SageMaker console.
5. When an AD consumer is created in your company AD, a corresponding SSO consumer is created in IAM Identification Heart.
6. When the AD consumer is assigned to an AD group, an IAM Identification Heart API (CreateGroupMembership) is invoked, and SSO group membership is created.
7. The previous occasion is logged in AWS CloudTrail with the title AddMemberToGroup.
8. An Amazon EventBridge rule listens to CloudTrail occasions and matches the AddMemberToGroup rule sample.
9. The EventBridge rule triggers the goal AWS Lambda perform.
10. This Lambda perform will name again IAM Identification Heart APIs, get the SSO consumer and group data, and carry out the next steps to create the Studio consumer profile (CreateUserProfile) for the SSO consumer:
    1. Search for the DynamoDB desk to fetch the IAM position similar to the AD group.
    2. Create a consumer profile with the SSO consumer and the IAM position obtained from the lookup desk.
    3. The SSO consumer is granted entry to Studio.
11. The SSO consumer is redirected to the Studio IDE by way of the Studio area URL.

Word that, as of writing, Step 4b (affiliate the SSO group to the Studio area) must be carried out manually by an admin utilizing the SageMaker console on the SageMaker area stage.

Arrange a Lambda perform to create the consumer profiles

The answer makes use of a Lambda perform to create the Studio consumer profiles. We offer the next pattern Lambda perform that you may copy and modify to satisfy your wants for automating the creation of the Studio consumer profile. This perform performs the next actions:

1. Obtain the CloudTrail AddMemberToGroup occasion from EventBridge.
2. Retrieve the Studio DOMAIN_ID from the surroundings variable (you possibly can alternatively hard-code the area ID or use a DynamoDB desk as nicely when you have a number of domains).
3. Learn from a dummy markup desk to match AD customers to execution roles. You’ll be able to change this to fetch from the DynamoDB desk if you happen to’re utilizing a table-driven strategy. If you happen to use DynamoDB, your Lambda perform’s execution position wants permissions to learn from the desk as nicely.
4. Retrieve the SSO consumer and AD group membership data from IAM Identification Heart, based mostly on the CloudTrail occasion knowledge.
5. Create a Studio consumer profile for the SSO consumer, with the SSO particulars and the matching execution position.

import os
import json
import boto3
DOMAIN_ID = os.environ.get('DOMAIN_ID', 'd-xxxx')


def lambda_handler(occasion, context):
    
    print({"Occasion": occasion})

    consumer = boto3.consumer('identitystore')
    sm_client = boto3.consumer('sagemaker')
    
    event_detail = occasion['detail']
    group_response = consumer.describe_group(
        IdentityStoreId=event_detail['requestParameters']['identityStoreId'],
        GroupId=event_detail['requestParameters']['groupId'],
    )
    group_name = group_response['DisplayName']
    
    user_response = consumer.describe_user(
        IdentityStoreId=event_detail['requestParameters']['identityStoreId'],
        UserId=event_detail['requestParameters']['member']['memberId']
    )
    user_name = user_response['UserName']
    print(f"Occasion particulars: {user_name} has been added to {group_name}")
    
    mapping_dict = {
        "ad-group-1": "<execution-role-arn>",
        "ad-group-2": "<execution-role-arn>”
    }
    
    user_role = mapping_dict.get(group_name)
    
    if user_role:
        response = sm_client.create_user_profile(
            DomainId=DOMAIN_ID,
            SingleSignOnUserIdentifier="UserName",
            SingleSignOnUserValue=user_name,
            # if the SSO user_name worth is an electronic mail, 
	  #  add logic to deal with it since Studio consumer profiles don’t settle for @ character
            UserProfileName=user_name, 
            UserSettings={
                "ExecutionRole": user_role
            }
        )
        print(response)
    else:
        response = "Group is just not approved to make use of SageMaker. Doing nothing."
        print(response)
    return {
        'statusCode': 200,
        'physique': json.dumps(response)
    }

Word that by default, the Lambda execution position doesn’t have entry to create consumer profiles or listing SSO customers. After you create the Lambda perform, entry the perform’s execution position on IAM and fasten the next coverage as an inline coverage after scoping down as wanted based mostly in your group necessities.

{
    "Model": "2012-10-17",
    "Assertion": [
        {
            "Action": [
                "identitystore:DescribeGroup",
                "identitystore:DescribeUser"
            ],
            "Impact": "Permit",
            "Useful resource": "*"
        },
        {
            "Motion": "sagemaker:CreateUserProfile",
            "Impact": "Permit",
            "Useful resource": "*"
        },
        {
            "Motion": "iam:PassRole",
            "Impact": "Permit",
            "Useful resource": [
                "<list-of-studio-execution-roles>"
            ]
        }
    ]
}

Arrange the EventBridge rule for the CloudTrail occasion

EventBridge is a serverless occasion bus service that you need to use to attach your functions with knowledge from a wide range of sources. On this answer, we create a rule-based set off: EventBridge listens to occasions and matches towards the offered sample and triggers a Lambda perform if the sample match is profitable. As defined within the answer overview, we take heed to the AddMemberToGroup occasion. To set it up, full the next steps:

1. On the EventBridge console, select Guidelines within the navigation pane.
2. Select Create rule.
3. Present a rule title, for instance, AddUserToADGroup.
4. Optionally, enter an outline.
5. Choose default for the occasion bus.
6. Below Rule kind, select Rule with an occasion sample, then select Subsequent.
7. On the Construct occasion sample web page, select Occasion supply as AWS occasions or EventBridge companion occasions.
8. Below Occasion sample, select the Customized patterns (JSON editor) tab and enter the next sample:

   {
     "supply": ["aws.sso-directory"],
     "detail-type": ["AWS API Call via CloudTrail"],
     "element": {
       "eventSource": ["sso-directory.amazonaws.com"],
       "eventName": ["AddMemberToGroup"]
     }
   }

9. Select Subsequent.
10. On the Choose goal(s) web page, select the AWS service for the goal kind, the Lambda perform because the goal, and the perform you created earlier, then select Subsequent.
11. Select Subsequent on the Configure tags web page, then select Create rule on the Overview and create web page.

After you’ve set the Lambda perform and the EventBridge rule, you possibly can take a look at out this answer. To take action, open your IdP and add a consumer to one of many AD teams with the Studio execution position mapped. When you add the consumer, you possibly can confirm the Lambda perform logs to examine the occasion and likewise see the Studio consumer provisioned routinely. Moreover, you need to use the DescribeUserProfile API name to confirm that the consumer is created with acceptable permissions.

Supporting a number of Studio accounts

To help a number of Studio accounts with the previous structure, we suggest the next adjustments:

1. Arrange an AD group mapped to every Studio account stage.
2. Arrange a group-level IAM position in every Studio account.
3. Arrange or derive the group to IAM position mapping.
4. Arrange a Lambda perform to carry out cross-account role assumption, based mostly on the IAM position mapping ARN and created consumer profile.

Deprovisioning customers

When a consumer is faraway from their AD group, you need to take away their entry from the Studio area as nicely. With SSO, when a consumer is eliminated, the consumer is disabled in IAM Identification Heart routinely if the AD to IAM Identification Heart sync is in place, and their Studio utility entry is straight away revoked.

Nevertheless, the consumer profile on Studio nonetheless persists. You’ll be able to add the same workflow with CloudTrail and a Lambda perform to take away the consumer profile from Studio. The EventBridge set off ought to now pay attention for the DeleteGroupMembership occasion. Within the Lambda perform, full the next steps:

1. Receive the consumer profile title from the consumer and group ID.
2. Record all working apps for the consumer profile utilizing the ListApps API name, filtering by the UserProfileNameEquals parameter. Make certain to verify for the paginated response, to listing all apps for the consumer.
3. Delete all working apps for the consumer and wait till all apps are deleted. You need to use the DescribeApp API to view the app’s standing.
4. When all apps are in a Deleted state (or Failed), delete the consumer profile.

With this answer in place, ML platform directors can keep group memberships in a single central location and automate the Studio consumer profile administration by way of EventBridge and Lambda capabilities.

The next code exhibits a pattern CloudTrail occasion:

"AddMemberToGroup": 
{
    "eventVersion": "1.08",
    "userIdentity": {
        "kind": "Unknown",
        "accountId": "<account-id>",
        "accessKeyId": "30997fec-b566-4b8b-810b-60934abddaa2"
    },
    "eventTime": "2022-09-26T22:24:18Z",
    "eventSource": "sso-directory.amazonaws.com",
    "eventName": "AddMemberToGroup",
    "awsRegion": "us-east-1",
    "sourceIPAddress": "54.189.184.116",
    "userAgent": "Okta SCIM Shopper 1.0.0",
    "requestParameters": {
        "identityStoreId": "d-906716eb24",
        "groupId": "14f83478-a061-708f-8de4-a3a2b99e9d89",
        "member": {
            "memberId": "04c8e458-a021-702e-f9d1-7f430ff2c752"
        }
    },
    "responseElements": null,
    "requestID": "b24a123b-afb3-4fb6-8650-b0dc1f35ea3a",
    "eventID": "c2c0873b-5c49-404c-add7-f10d4a6bd40c",
    "readOnly": false,
    "eventType": "AwsApiCall",
    "managementEvent": true,
    "recipientAccountId": "<account-id>",
    "eventCategory": "Administration",
    "tlsDetails": {
        "tlsVersion": "TLSv1.2",
        "cipherSuite": "ECDHE-RSA-AES128-GCM-SHA256",
        "clientProvidedHostHeader": "up.sso.us-east-1.amazonaws.com"
    }
}

The next code exhibits a pattern Studio consumer profile API request:

create-user-profile 
--domain-id d-xxxxxx 
--user-profile-name ssouserid
--single-sign-on-user-identifier 'userName' 
--single-sign-on-user-value 'ssouserid‘ 
--user-settings ExecutionRole=arn:aws:iam::<account id>:position/title

Conclusion

On this publish, we mentioned how directors can scale Studio onboarding for tons of of customers based mostly on their AD group membership. We demonstrated an end-to-end answer structure that organizations can undertake to automate and scale their onboarding course of to satisfy their agility, safety, and compliance wants. If you happen to’re in search of a scalable answer to automate your consumer onboarding, do this answer, and go away you suggestions beneath! For extra details about onboarding to Studio, see Onboard to Amazon SageMaker Domain.

In regards to the authors

Ram Vittal is an ML Specialist Options Architect at AWS. He has over 20 years of expertise architecting and constructing distributed, hybrid, and cloud functions. He’s enthusiastic about constructing safe and scalable AI/ML and large knowledge options to assist enterprise clients with their cloud adoption and optimization journey to enhance their enterprise outcomes. In his spare time, he rides his motorbike and walks along with his 2-year-old sheep-a-doodle!

Durga Sury is an ML Options Architect within the Amazon SageMaker Service SA workforce. She is enthusiastic about making machine studying accessible to everybody. In her 4 years at AWS, she has helped arrange AI/ML platforms for enterprise clients. When she isn’t working, she loves motorbike rides, thriller novels, and mountaineering along with her 5-year-old husky.