如何从Lambda函数向Sagemaker Notebook实例传递文件名

这里提供三种可行的实现方式，根据你的场景选择合适的方案：

方案1：通过实例标签传递文件名 #

Lambda从S3触发事件中解析出文件名，启动Notebook实例时为实例添加包含该文件名的标签；Notebook启动后，调用Sagemaker API读取自身标签获取文件名。

修改后的Lambda代码 #

import boto3

def lambda_handler(event, context):
    # 解析S3事件中的文件名和存储桶名
    s3_record = event['Records'][0]['s3']
    target_file = s3_record['object']['key']
    source_bucket = s3_record['bucket']['name']

    print(f"Initiating notebook for file: {target_file}")
    sagemaker_client = boto3.client('sagemaker')
    
    # 启动Notebook实例并附加标签
    sagemaker_client.start_notebook_instance(
        NotebookInstanceName='preprocess-dataset',
        Tags=[
            {'Key': 'TargetFile', 'Value': target_file},
            {'Key': 'SourceBucket', 'Value': source_bucket}
        ]
    )
    print("Notebook instance started")
    return 0

Notebook中获取文件名的代码 #

import boto3

sagemaker_client = boto3.client('sagemaker')
instance_name = 'preprocess-dataset'  # 也可通过环境变量获取：os.environ['NOTEBOOK_INSTANCE_NAME']

# 查询实例标签
instance_info = sagemaker_client.describe_notebook_instance(NotebookInstanceName=instance_name)
tags = instance_info['Tags']

# 提取目标文件名和存储桶
target_file = next(tag['Value'] for tag in tags if tag['Key'] == 'TargetFile')
source_bucket = next(tag['Value'] for tag in tags if tag['Key'] == 'SourceBucket')

# 读取S3文件示例
s3_client = boto3.client('s3')
file_content = s3_client.get_object(Bucket=source_bucket, Key=target_file)['Body'].read().decode('utf-8')

方案2：通过生命周期配置+环境变量传递 #

利用Sagemaker生命周期配置脚本，在实例启动时读取标签并设置系统环境变量，Notebook直接读取环境变量即可获取文件名，无需额外API调用。

步骤1：创建生命周期配置脚本 #

编写启动脚本on-start.sh，内容如下：

#!/bin/bash
set -e

# 获取当前实例名称
INSTANCE_NAME=$(curl -s http://169.254.169.254/latest/meta-data/instance-id)

# 调用API读取实例标签，需确保实例角色拥有`sagemaker:DescribeNotebookInstance`权限
TARGET_FILE=$(aws sagemaker describe-notebook-instance --notebook-instance-name "$INSTANCE_NAME" --query 'Tags[?Key==`TargetFile`].Value' --output text)
SOURCE_BUCKET=$(aws sagemaker describe-notebook-instance --notebook-instance-name "$INSTANCE_NAME" --query 'Tags[?Key==`SourceBucket`].Value' --output text)

# 将变量写入系统环境配置，确保Notebook会话能读取到
echo "export TARGET_FILE=\"$TARGET_FILE\"" >> /etc/profile.d/sagemaker-env.sh
echo "export SOURCE_BUCKET=\"$SOURCE_BUCKET\"" >> /etc/profile.d/sagemaker-env.sh

将该脚本上传到S3存储桶，然后创建生命周期配置：

import boto3

sagemaker_client = boto3.client('sagemaker')
sagemaker_client.create_notebook_instance_lifecycle_config(
    NotebookInstanceLifecycleConfigName='preprocess-lifecycle',
    OnStart=[
        {'Content': {'S3Uri': 's3://your-bucket-path/lifecycle-scripts/on-start.sh'}}
    ]
)

步骤2：修改Lambda代码关联生命周期配置 #

import boto3

def lambda_handler(event, context):
    s3_record = event['Records'][0]['s3']
    target_file = s3_record['object']['key']
    source_bucket = s3_record['bucket']['name']

    print(f"Initiating notebook for file: {target_file}")
    sagemaker_client = boto3.client('sagemaker')
    
    sagemaker_client.start_notebook_instance(
        NotebookInstanceName='preprocess-dataset',
        NotebookInstanceLifecycleConfigName='preprocess-lifecycle',
        Tags=[
            {'Key': 'TargetFile', 'Value': target_file},
            {'Key': 'SourceBucket', 'Value': source_bucket}
        ]
    )
    print("Notebook instance started")
    return 0

步骤3：Notebook中读取环境变量 #

import os
import boto3

target_file = os.environ['TARGET_FILE']
source_bucket = os.environ['SOURCE_BUCKET']

# 读取S3文件示例
s3_client = boto3.client('s3')
file_content = s3_client.get_object(Bucket=source_bucket, Key=target_file)['Body'].read().decode('utf-8')

方案3：将文件名写入S3元数据文件 #

Lambda触发时，把目标文件名写入S3的一个固定元数据文件；Notebook启动后，读取该元数据文件获取文件名。

修改后的Lambda代码 #

import boto3

def lambda_handler(event, context):
    s3_record = event['Records'][0]['s3']
    target_file = s3_record['object']['key']
    source_bucket = s3_record['bucket']['name']

    print(f"Initiating notebook for file: {target_file}")
    
    # 将文件名写入S3元数据文件
    s3_client = boto3.client('s3')
    s3_client.put_object(
        Bucket=source_bucket,
        Key='processing-metadata/current-target.txt',
        Body=target_file.encode('utf-8')
    )
    
    # 启动Notebook实例
    sagemaker_client = boto3.client('sagemaker')
    sagemaker_client.start_notebook_instance(NotebookInstanceName='preprocess-dataset')
    print("Notebook instance started")
    return 0

Notebook中读取元数据文件 #

import boto3

source_bucket = 'your-bucket-name'
s3_client = boto3.client('s3')

# 读取元数据文件获取目标文件名
metadata_response = s3_client.get_object(Bucket=source_bucket, Key='processing-metadata/current-target.txt')
target_file = metadata_response['Body'].read().decode('utf-8')

# 读取S3文件示例
file_content = s3_client.get_object(Bucket=source_bucket, Key=target_file)['Body'].read().decode('utf-8')

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