使用Active Directory组特定的IAM角色将用户引入Amazon SageMaker Studio

Amazon SageMaker Studio是一个基于web的集成开发环境(IDE),专门用于机器学习(ML),可以让您构建、训练、调试、部署和监视您的ML模型。要在您的AWS帐户和区域中提供Studio,您首先需要创建一个Amazon SageMaker域——这是一个封装了您的ML环境的构造物。更具体地说,SageMaker域由一个关联的Amazon Elastic File System(Amazon EFS)卷、一个授权用户列表以及各种安全、应用程序、策略和Amazon Virtual Private Cloud(Amazon VPC)配置组成。

在创建SageMaker域时,您可以选择使用AWS IAM Identity Center(AWS Single Sign-On的后继者)或AWS身份和访问管理(IAM)作为用户认证方法。这两种身份验证方法都有自己的一套用例;在本篇文章中,我们将重点关注使用IAM Identity Center的SageMaker域,即单点登录(SSO)模式作为身份验证方法。

使用SSO模式,您可以在IAM Identity Center中设置SSO用户和组,然后从Studio控制台授予对SSO组或用户的访问权限。目前,域中的所有SSO用户都继承域的执行角色。这可能不适用于所有组织。例如,管理员可能希望基于其Active Directory(AD)组成员身份为Studio SSO用户设置IAM权限。此外,由于管理员需要手动授予SSO用户访问Studio的权限,当有数百个用户加入时,该过程可能无法扩展。

在本篇文章中,我们提供了一个解决方案的指南,以根据AD组成员身份为Studio提供最少特权权限的SSO用户。这些指导方针可以让您快速扩展到数百个用户,并实现安全和合规性。

解决方案概述

以下图表说明了解决方案架构。

将AD用户提供给Studio的工作流程包括以下步骤:

  1. 在SSO模式下设置Studio域。

  2. 对于每个AD组:

    1. 使用适当的细粒度IAM策略设置您的Studio执行角色
    2. 记录一个AD组-角色映射Amazon DynamoDB表中的条目。

    或者,您可以采用基于AD组名称的IAM角色ARN命名标准,并且无需存储映射在外部数据库中即可派生IAM角色ARN。

  3. 将您的AD用户和组以及成员身份同步到AWS Identity Center:

    1. 如果您使用支持SCIM的身份提供程序(IdP),请使用SCIM API与IAM Identity Center进行集成。
    2. 如果您使用自管理的AD,则可以使用AD Connector。
  4. 在公司AD中创建AD组时,完成以下步骤:

    1. 在IAM Identity Center中创建相应的SSO组。
    2. 使用SageMaker控制台将SSO组关联到Studio域。
  5. 当在公司AD中创建AD用户时,会创建相应的SSO用户在IAM Identity Center中。

  6. 当将AD用户分配到AD组时,将调用IAM Identity Center API(CreateGroupMembership),并创建SSO组成员身份。

  7. 前面的事件在AWS CloudTrail中记录,并以名称AddMemberToGroup出现。

  8. Amazon EventBridge规则监听CloudTrail事件,并匹配AddMemberToGroup规则模式。

  9. EventBridge规则触发目标AWS Lambda函数。

  10. 此Lambda函数将回调IAM Identity Center API,并获取SSO用户和组信息,并执行以下步骤为SSO用户创建Studio用户配置文件(CreateUserProfile):

    1. 查找DynamoDB表以获取与AD组对应的IAM角色。
    2. 使用从查找表中获取的IAM角色创建用户配置文件。
    3. SSO用户被授予访问Studio的权限。
  11. SSO用户通过Studio域URL被重定向到Studio IDE。

请注意,截至撰写本文,第4b步(将SSO组与Studio域关联)需要管理员在SageMaker控制台上手动执行。

设置Lambda函数以创建用户配置文件

该解决方案使用Lambda函数创建Studio用户配置文件。我们提供以下示例Lambda函数,您可以复制并修改以满足自动化创建Studio用户配置文件的需求。该函数执行以下操作:

  1. 从EventBridge接收CloudTrail AddMemberToGroup事件。
  2. 从环境变量中检索Studio DOMAIN_ID(如果有多个域,您也可以硬编码域ID或使用DynamoDB表)。
  3. 从虚拟标记表中读取以将AD用户与执行角色匹配。如果使用基于表的方法,可以将其更改为从DynamoDB表中获取。如果使用DynamoDB,则Lambda函数的执行角色也需要读取表的权限。
  4. 基于CloudTrail事件数据,从IAM身份中心检索SSO用户和AD组成员身份信息。
  5. 为SSO用户创建Studio用户配置文件,包括SSO详细信息和匹配的执行角色。
import os
import json
import boto3
DOMAIN_ID = os.environ.get('DOMAIN_ID', 'd-xxxx')


def lambda_handler(event, context):
    
    print({"Event": event})

    client = boto3.client('identitystore')
    sm_client = boto3.client('sagemaker')
    
    event_detail = event['detail']
    group_response = client.describe_group(
        IdentityStoreId=event_detail['requestParameters']['identityStoreId'],
        GroupId=event_detail['requestParameters']['groupId'],
    )
    group_name = group_response['DisplayName']
    
    user_response = client.describe_user(
        IdentityStoreId=event_detail['requestParameters']['identityStoreId'],
        UserId=event_detail['requestParameters']['member']['memberId']
    )
    user_name = user_response['UserName']
    print(f"Event details: {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 value is an email, 
      #  add logic to handle it since Studio user profiles don’t accept @ character
            UserProfileName=user_name, 
            UserSettings={
                "ExecutionRole": user_role
            }
        )
        print(response)
    else:
        response = "Group is not authorized to use SageMaker. Doing nothing."
        print(response)
    return {
        'statusCode': 200,
        'body': json.dumps(response)
    }

请注意,默认情况下,Lambda执行角色无法访问创建用户配置文件或列出SSO用户。创建Lambda函数后,请在IAM中访问函数的执行角色,并根据组织要求进行适当的范围限制,将以下策略作为内联策略附加。

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "identitystore:DescribeGroup",
                "identitystore:DescribeUser"
            ],
            "Effect": "Allow",
            "Resource": "*"
        },
        {
            "Action": "sagemaker:CreateUserProfile",
            "Effect": "Allow",
            "Resource": "*"
        },
        {
            "Action": "iam:PassRole",
            "Effect": "Allow",
            "Resource": [
                "<list-of-studio-execution-roles>"
            ]
        }
    ]
}

设置CloudTrail事件的EventBridge规则

EventBridge是一项无服务器事件总线服务,可将您的应用程序连接到各种来源的数据。在此解决方案中,我们创建了基于规则的触发器:EventBridge侦听事件并根据提供的模式进行匹配,如果模式匹配成功,则触发Lambda函数。如解决方案概述中所述,我们侦听AddMemberToGroup事件。设置方式请完成以下步骤:

  1. 在 EventBridge 控制台中,选择导航栏中的 规则

  2. 选择 创建规则

  3. 提供规则名称,例如 AddUserToADGroup

  4. 可选地,输入描述。

  5. 为事件总线选择 默认

  6. 规则类型 下,选择 带有事件模式的规则,然后选择 下一步

  7. 构建事件模式 页面中,将 事件源 选择为 AWS 事件或 EventBridge 合作伙伴事件

  8. 事件模式 下,选择 自定义模式(JSON 编辑器) 选项卡,并输入以下模式:

    {
      "source": ["aws.sso-directory"],
      "detail-type": ["AWS API Call via CloudTrail"],
      "detail": {
        "eventSource": ["sso-directory.amazonaws.com"],
        "eventName": ["AddMemberToGroup"]
      }
    }
  9. 选择 下一步

  10. 选择目标 页面中,选择 AWS 服务作为目标类型,Lambda 函数作为目标,并选择之前创建的函数,然后选择 下一步

  11. 配置标签 页面上选择 下一步,然后在 审阅和创建 页面上选择 创建规则

设置 Lambda 函数和 EventBridge 规则后,您可以测试此解决方案。要执行此操作,请打开您的 IdP 并将用户添加到其中一个已映射 Studio 执行角色的 AD 组中。添加用户后,您可以验证 Lambda 函数日志以检查事件,并自动查看 Studio 用户的提供。此外,您可以使用 DescribeUserProfile API 调用来验证用户是否创建并具有适当的权限。

支持多个 Studio 账户

为了支持前面的架构中的多个 Studio 账户,我们建议进行以下更改:

  1. 为每个 Studio 账户级别设置映射的 AD 组。
  2. 在每个 Studio 账户中设置组级 IAM 角色。
  3. 设置或派生组到 IAM 角色映射。
  4. 设置 Lambda 函数以执行跨帐户角色假设,基于 IAM 角色映射 ARN 和创建的用户配置文件。

取消用户资格

当用户从其 AD 组中移除时,您应该同时从 Studio 域中移除他们的访问权限。使用 SSO 时,当用户被删除时,如果已进行 AD 到 IAM Identity Center 同步,则 IAM Identity Center 会自动禁用该用户,并立即吊销其 Studio 应用程序访问权限。

但是,Studio 上的用户配置文件仍然存在。您可以使用 CloudTrail 和 Lambda 函数添加类似的工作流,以从 Studio 中删除用户配置文件。现在,EventBridge 触发器应该侦听 DeleteGroupMembership 事件。在 Lambda 函数中,请执行以下步骤:

  1. 从用户和组 ID 获取用户配置文件名称。
  2. 使用 ListApps API 调用列出用户配置文件的所有运行应用程序,并按 UserProfileNameEquals 参数进行过滤。确保检查分页响应,以列出用户的所有应用程序。
  3. 删除用户的所有运行应用程序,并等待所有应用程序被删除。您可以使用 DescribeApp API 查看应用程序的状态。
  4. 当所有应用程序处于 已删除 状态(或 已失败)时,删除用户配置文件。

通过此解决方案,ML 平台管理员可以在一个中央位置维护组成员身份,并通过 EventBridge 和 Lambda 函数自动化 Studio 用户配置文件管理。

以下代码显示了示例 CloudTrail 事件:

"AddMemberToGroup": 
{
    "eventVersion": "1.08",
    "userIdentity": {
        "type": "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 Client 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": "Management",
    "tlsDetails": {
        "tlsVersion": "TLSv1.2",
        "cipherSuite": "ECDHE-RSA-AES128-GCM-SHA256",
        "clientProvidedHostHeader": "up.sso.us-east-1.amazonaws.com"
    }
}

以下代码展示了一个样例 Studio 用户配置文件的 API 请求:

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>:role/name

结论

在本文中,我们讨论了管理员如何根据 AD 组成员身份为数百个用户扩展 Studio 入门体验的方法。我们演示了一种端到端的解决方案架构,组织可以采用该架构来自动化和扩展其入门流程,以满足其敏捷性、安全性和合规性需求。如果您正在寻找一种可扩展的解决方案来自动化您的用户入门流程,请尝试此解决方案,并在下方留下您的反馈!有关入门到 Studio 的更多信息,请参见 Onboard to Amazon SageMaker Domain。