如何在S3上传文件时自动运行AWS SageMaker Studio Notebook

背景说明 #

SageMaker Studio和传统Notebook Instance的核心差异在于：Studio基于用户空间和Kernel Gateway运行，没有独立的可直接启动的EC2实例资源。因此无法照搬Notebook Instance的触发逻辑，需要通过「S3事件→Lambda→启动Studio Kernel→执行Notebook」的链路实现。

分步实现 #

1. 预处理Studio Notebook #

• 在SageMaker Studio中创建并调试好目标Notebook（例如process_s3_upload.ipynb），确保它能读取指定S3路径的文件并完成业务逻辑。
• 两种执行方式任选：
  • 转成Python脚本：在Studio终端执行 jupyter nbconvert --to script process_s3_upload.ipynb，后续直接运行脚本。
  • 保留Notebook：用nbconvert直接执行（适合需要保留输出日志的场景）。

2. 配置S3事件通知 #

• 打开目标S3桶的属性→事件通知，新建通知规则：
  • 事件类型：选择「所有对象创建事件」或仅「Put」事件（按需筛选）。
  • 目标：指定后续创建的Lambda函数。
  • 前缀/后缀：可设置（例如仅触发.csv文件上传）。

3. 开发Lambda触发函数 #

Lambda需要具备访问S3、SageMaker Studio、SSM的权限，核心逻辑分为两步：检查/启动Studio Kernel、执行Notebook。

示例代码（Python）

import boto3
import time
import json

sagemaker = boto3.client('sagemaker')
ssm = boto3.client('ssm')

def lambda_handler(event, context):
    # 提取S3上传事件信息
    s3_rec = event['Records'][0]['s3']
    bucket = s3_rec['bucket']['name']
    file_key = s3_rec['object']['key']

    # Studio环境配置（替换为你的实际信息）
    DOMAIN_ID = 'd-xxxxxxxxxx'  # 在SageMaker控制台的Studio域名详情中获取
    USER_PROFILE = 'default'     # Studio用户名
    APP_NAME = 'default'         # Studio中运行时的名称（例如Python 3 (Data Science)）

    # 检查Kernel Gateway是否已启动，未启动则创建
    try:
        sagemaker.describe_app(
            DomainId=DOMAIN_ID,
            UserProfileName=USER_PROFILE,
            AppType='KernelGateway',
            AppName=APP_NAME
        )
    except sagemaker.exceptions.ResourceNotFound:
        # 启动Kernel，替换为你的实例类型和镜像ARN
        sagemaker.create_app(
            DomainId=DOMAIN_ID,
            UserProfileName=USER_PROFILE,
            AppType='KernelGateway',
            AppName=APP_NAME,
            ResourceSpec={
                'InstanceType': 'ml.t3.medium',
                'SageMakerImageArn': 'arn:aws:sagemaker:us-east-1:081325390199:image/data-science-1-0'
            }
        )
        # 等待Kernel启动（根据实例类型调整等待时间）
        time.sleep(120)

    # 获取Studio Kernel对应的EC2实例ID（通过SSM标签筛选）
    instance_list = ssm.describe_instance_information(
        Filters=[
            {'Key': 'tag:aws:sagemaker:domain-id', 'Values': [DOMAIN_ID]},
            {'Key': 'tag:aws:sagemaker:user-profile-name', 'Values': [USER_PROFILE]},
            {'Key': 'tag:aws:sagemaker:app-name', 'Values': [APP_NAME]}
        ]
    )['InstanceInformationList']

    if not instance_list:
        raise RuntimeError("Studio Kernel instance not registered in SSM")
    
    instance_id = instance_list[0]['InstanceId']

    # 构建执行命令：传递S3参数并运行Notebook
    notebook_path = '/home/sagemaker-user/process_s3_upload.ipynb'
    run_cmd = f"""
    export S3_BUCKET={bucket}
    export S3_FILE_KEY={file_key}
    jupyter nbconvert --execute {notebook_path} --ExecutePreprocessor.timeout=3600 --to html
    """

    # 通过SSM发送命令执行
    ssm.send_command(
        InstanceIds=[instance_id],
        DocumentName='AWS-RunShellScript',
        Parameters={'commands': [run_cmd]},
        CloudWatchOutputConfig={
            'CloudWatchLogGroupName': '/aws/lambda/studio-notebook-trigger',
            'CloudWatchOutputEnabled': True
        }
    )

    return {'statusCode': 200, 'body': json.dumps(f"Triggered notebook for s3://{bucket}/{file_key}")}

4. 配置Lambda权限 #

给Lambda的IAM角色添加以下权限（建议最小化权限，避免过度授权）：

• S3只读权限（限制到目标桶）
• SageMaker权限：sagemaker:DescribeApp、sagemaker:CreateApp
• SSM权限：ssm:DescribeInstanceInformation、ssm:SendCommand
• CloudWatch日志权限：用于记录执行日志

5. 调试与优化 #

• 测试：上传文件到目标S3桶，查看Lambda和CloudWatch日志，确认Kernel启动和Notebook执行状态。
• 优化点：
  • 避免重复启动：Lambda中可增加Kernel状态检查，已启动则跳过启动步骤。
  • 参数传递：在Notebook中通过os.environ.get('S3_BUCKET')获取环境变量，动态处理上传的文件。
  • 超时处理：根据Notebook执行时长调整ExecutePreprocessor.timeout参数。

关键注意事项 #

• Studio Kernel有空闲超时机制，若长时间未使用会自动关闭，Lambda需处理重新启动逻辑。
• 实例类型和镜像ARN必须与Studio中使用的一致，否则会出现环境不兼容问题。
• 成本控制：选择合适的实例类型，避免闲置实例持续运行。

内容的提问来源于stack exchange，提问作者Logan27