使用屬性型存取控制 (ABAC) 進行多租用戶訓練

啟用 ABAC 的工作階段標籤鏈結

1. 匯入 boto3和 SageMaker Python SDK。已啟用 ABAC 的訓練任務隔離僅適用於 SageMaker AI Python SDK 的 2.217 版或更新版本。

   import boto3
   import sagemaker
   
   from sagemaker.estimator import Estimator
   from sagemaker.inputs import TrainingInput

2. 設定 AWS STS 和 SageMaker AI 用戶端以使用租戶標記的工作階段標籤。您可以變更標籤值來指定不同的租用戶。

   # Start an AWS STS client
   sts_client = boto3.client('sts')
   
   # Define your tenants using tags
   # The session tag key must match the principal tag key in your execution role policy
   tags = []
   tag = {}
   tag['Key'] = "tenant-id"
   tag['Value'] = "example-tenant"
   tags.append(tag)
   
   # Have AWS STS assume your ABAC-enabled job creation role
   response = sts_client.assume_role(
       RoleArn="arn:aws:iam::<account-id>:role/<your-training-job-creation-role>",
       RoleSessionName="SessionName",
       Tags=tags)
   credentials = response['Credentials']
   
   # Create a client with your job creation role (which was assumed with tags)
   sagemaker_client = boto3.client(
       'sagemaker',
       aws_access_key_id=credentials['AccessKeyId'],
       aws_secret_access_key=credentials['SecretAccessKey'],
       aws_session_token=credentials['SessionToken']
   )
   sagemaker_session = sagemaker.Session(sagemaker_client=sagemaker_client)

   將標籤附加"tenant-id=example-tenant"至任務建立角色時，執行角色會擷取這些標籤，以使用下列政策：

   {
       "Version": "2012-10-17",
       "Statement": [
           {
               "Action": [
                   "s3:GetObject",
                   "s3:PutObject"
               ],
               "Resource": [
                   "arn:aws:s3:::<your-input-s3-bucket>/example-tenant/*"
               ],
               "Effect": "Allow"
           },
               "Action": [
                   "s3:PutObject"
               ],
               "Resource": "arn:aws:s3:::<your-output-s3-bucket>/example-tenant/*"
           },
           {
               "Action": "s3:ListBucket",
               "Resource": "*",
               "Effect": "Allow"
           }
       ]
   }

3. 定義估算器以使用 SageMaker Python SDK 建立訓練任務。將 enable_session_tag_chaining設定為 以True允許 SageMaker AI 訓練執行角色從您的任務建立角色擷取標籤。

   # Specify your training input
   trainingInput = TrainingInput(
       s3_data='s3://<your-input-bucket>/example-tenant',
       distribution='ShardedByS3Key',
       s3_data_type='S3Prefix'
   )
   
   # Specify your training job execution role 
   execution_role_arn = "arn:aws:iam::<account-id>:role/<your-training-job-execution-role>"
   
   # Define your esimator with session tag chaining enabled
   estimator = Estimator(
       image_uri="<your-training-image-uri>",
       role=execution_role_arn,
       instance_count=1,
       instance_type='ml.m4.xlarge',
       volume_size=20,
       max_run=3600,
       sagemaker_session=sagemaker_session,
       output_path="s3://<your-output-bucket>/example-tenant",
       enable_session_tag_chaining=True
   )
   
   estimator.fit(inputs=trainingInput, job_name="abac-demo")