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

Amazon SageMaker Studio is a web-based built-in growth surroundings (IDE) for machine studying (ML) that allows you to construct, prepare, debug, deploy, and monitor your ML fashions. For provisioning Studio in your AWS account and Area, you first have to 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, software, coverage, and Amazon Virtual Private Cloud (Amazon VPC) configurations.

When creating your SageMaker area, you’ll be able to 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 instances; on this put up, we give attention to SageMaker domains with IAM Identification Middle, or single sign-on (SSO) mode, because the authentication technique.

With SSO mode, you arrange an SSO consumer and group in IAM Identification Middle after which grant entry to both the SSO group or consumer from the Studio console. At present, all SSO customers in a website inherit the area’s execution position. This may occasionally not work for all organizations. As an example, directors could need to arrange IAM permissions for a Studio SSO consumer primarily based on their Energetic 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 a whole bunch of customers.

On this put up, we offer prescriptive steerage for the answer to provision SSO customers to Studio with least privilege permissions primarily based on AD group membership. This steerage lets you shortly scale for onboarding a whole bunch of customers to Studio and obtain your safety and compliance posture.

Resolution overview

The next diagram illustrates the answer structure.

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

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

   Alternatively, you’ll be able to undertake a naming commonplace for IAM position ARNs primarily based 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 Middle:
   1. In case you’re utilizing an id supplier (IdP) that helps SCIM, use the SCIM API integration with IAM Identification Middle.
   2. In case you are utilizing self-managed AD, you might 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 Middle.
   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 Middle.
6. When the AD consumer is assigned to an AD group, an IAM Identification Middle 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 Middle 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 comparable 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 degree.

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’ll be able to alternatively hard-code the area ID or use a DynamoDB desk as effectively when you’ve got 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 in case you’re utilizing a table-driven strategy. In case you use DynamoDB, your Lambda perform’s execution position wants permissions to learn from the desk as effectively.
4. Retrieve the SSO consumer and AD group membership data from IAM Identification Middle, primarily based 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})

    shopper = boto3.shopper('identitystore')
    sm_client = boto3.shopper('sagemaker')
    
    event_detail = occasion['detail']
    group_response = shopper.describe_group(
        IdentityStoreId=event_detail['requestParameters']['identityStoreId'],
        GroupId=event_detail['requestParameters']['groupId'],
    )
    group_name = group_response['DisplayName']
    
    user_response = shopper.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 e-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 shouldn't be 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 checklist SSO customers. After you create the Lambda perform, entry the perform’s execution position on IAM and connect the next coverage as an inline coverage after scoping down as wanted primarily based in your group necessities.

{
    "Model": "2012-10-17",
    "Assertion": [
        {
            "Action": [
                "identitystore:DescribeGroup",
                "identitystore:DescribeUser"
            ],
            "Impact": "Enable",
            "Useful resource": "*"
        },
        {
            "Motion": "sagemaker:CreateUserProfile",
            "Impact": "Enable",
            "Useful resource": "*"
        },
        {
            "Motion": "iam:PassRole",
            "Impact": "Enable",
            "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 purposes 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 Evaluation and create web page.

After you’ve set the Lambda perform and the EventBridge rule, you’ll be able to 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’ll be able to confirm the Lambda perform logs to examine the occasion and likewise see the Studio consumer provisioned robotically. Moreover, you need to use the DescribeUserProfile API name to confirm that the consumer is created with applicable permissions.

Supporting a number of Studio accounts

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

1. Arrange an AD group mapped to every Studio account degree.
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, primarily based on the IAM position mapping ARN and created consumer profile.

Deprovisioning customers

When a consumer is faraway from their AD group, you must take away their entry from the Studio area as effectively. With SSO, when a consumer is eliminated, the consumer is disabled in IAM Identification Middle robotically if the AD to IAM Identification Middle sync is in place, and their Studio software 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 hear for the DeleteGroupMembership occasion. Within the Lambda perform, full the next steps:

1. Acquire the consumer profile title from the consumer and group ID.
2. Listing all working apps for the consumer profile utilizing the ListApps API name, filtering by the UserProfileNameEquals parameter. Be sure that to examine for the paginated response, to checklist all apps for the consumer.
3. Delete all working apps for the consumer and wait till all apps are deleted. You should 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 features.

The next code reveals 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 reveals 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 put up, we mentioned how directors can scale Studio onboarding for a whole bunch of customers primarily based 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. In case you’re in search of a scalable answer to automate your consumer onboarding, do that answer, and go away you suggestions under! For extra details about onboarding to Studio, see Onboard to Amazon SageMaker Domain.

Concerning 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 purposes. He’s enthusiastic about constructing safe and scalable AI/ML and massive knowledge options to assist enterprise prospects 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 crew. 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 prospects. When she isn’t working, she loves motorbike rides, thriller novels, and mountaineering together with her 5-year-old husky.