如何通过Lambda创建并执行含模型注册步骤的SageMaker Pipeline？

问题翻译 #

我希望依据官方文档创建一个包含模型注册步骤的Amazon SageMaker Pipeline（示例代码如下），并通过AWS Lambda执行该流水线。另外我需要下载已有的model.tar.gz文件，重新打包后再完成模型注册操作，请问有没有相关示例参考？

from sagemaker.model import Model
from sagemaker.inputs import CreateModelInput
from sagemaker.workflow.steps import CreateModelStep
from sagemaker.model_metrics import MetricsSource, ModelMetrics
from sagemaker.workflow.step_collections import RegisterModel

model = Model(
    image_uri=image_uri, # XGBoost镜像地址
    model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts, # 训练生成的model.tar.gz的S3路径
    sagemaker_session=sagemaker_session,
    role=role,
)

inputs = CreateModelInput(
    instance_type="ml.m5.large",
    accelerator_type="ml.eia1.medium",
)

step_create_model = CreateModelStep(
    name="adultCreateModel",
    model=model,
    inputs=inputs,
)

step_register = RegisterModel(
    name="adultRegisterModel",
    estimator=xgb_train,
    model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts, # model.tar.gz的S3路径
    content_types=["text/csv"],
    response_types=["text/csv"],
    inference_instances=["ml.t2.medium", "ml.m5.xlarge"],
    transform_instances=["ml.m5.xlarge"],
)

1. 在Pipeline中添加模型重打包步骤 #

要实现下载已有model.tar.gz、修改后重新打包，需在Pipeline中加入ProcessingStep，通过SageMaker Processing完成文件处理。

1.1 编写重打包脚本（repackage_model.py） #

该脚本负责从S3下载原始模型包、解压修改、重新压缩后上传回S3：

import os
import tarfile
import boto3
from pathlib import Path

s3 = boto3.client('s3')

def repackage_model(input_s3_uri, output_s3_uri):
    # 解析S3路径
    input_bucket, input_key = input_s3_uri.replace("s3://", "").split("/", 1)
    output_bucket, output_key = output_s3_uri.replace("s3://", "").split("/", 1)
    
    # 创建临时工作目录
    temp_dir = Path("/tmp/model_repackage")
    temp_dir.mkdir(exist_ok=True, parents=True)
    
    # 下载原始模型包
    local_input_path = temp_dir / "model.tar.gz"
    s3.download_file(input_bucket, input_key, str(local_input_path))
    
    # 解压模型文件
    extract_dir = temp_dir / "extracted"
    extract_dir.mkdir(exist_ok=True)
    with tarfile.open(local_input_path, 'r:gz') as tar:
        tar.extractall(path=extract_dir)
    
    # 这里添加自定义修改逻辑，比如添加推理脚本、调整模型配置
    (extract_dir / "inference.py").write_text("""
def model_fn(model_dir):
    import xgboost as xgb
    return xgb.Booster()
""")
    
    # 重新打包为model.tar.gz
    local_output_path = temp_dir / "new_model.tar.gz"
    with tarfile.open(local_output_path, 'w:gz') as tar:
        for file_path in extract_dir.rglob("*"):
            tar.add(file_path, arcname=file_path.relative_to(extract_dir))
    
    # 上传到目标S3路径
    s3.upload_file(str(local_output_path), output_bucket, output_key)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-s3-uri", type=str, required=True)
    parser.add_argument("--output-s3-uri", type=str, required=True)
    args = parser.parse_args()
    repackage_model(args.input_s3_uri, args.output_s3_uri)

1.2 修改Pipeline代码，集成重打包步骤 #

扩展原始Pipeline，加入重打包步骤后使用新模型路径完成注册：

from sagemaker.model import Model
from sagemaker.inputs import CreateModelInput
from sagemaker.workflow.steps import CreateModelStep, ProcessingStep
from sagemaker.processing import ScriptProcessor, ProcessingInput, ProcessingOutput
from sagemaker.workflow.step_collections import RegisterModel
from sagemaker.workflow.parameters import ParameterString
from sagemaker.workflow.pipeline import Pipeline

# 定义可配置参数
original_model_s3_uri = ParameterString(name="OriginalModelS3Uri", default_value="s3://your-bucket/original/model.tar.gz")
repacked_model_s3_uri = ParameterString(name="RepackedModelS3Uri", default_value="s3://your-bucket/repacked/model.tar.gz")

# 创建脚本处理器，用于执行重打包逻辑
script_processor = ScriptProcessor(
    image_uri="763104351884.dkr.ecr.us-east-1.amazonaws.com/sagemaker-scikit-learn:0.23-1-cpu-py3",
    command=["python3"],
    role=role,
    instance_type="ml.t2.medium",
    instance_count=1,
    sagemaker_session=sagemaker_session
)

# 定义重打包步骤
step_repackage_model = ProcessingStep(
    name="RepackageModelStep",
    processor=script_processor,
    inputs=[ProcessingInput(source=original_model_s3_uri, destination="/opt/ml/processing/input")],
    outputs=[ProcessingOutput(source="/opt/ml/processing/output", destination=repacked_model_s3_uri)],
    code="repackage_model.py",
    arguments=["--input-s3-uri", original_model_s3_uri, "--output-s3-uri", repacked_model_s3_uri]
)

# 创建模型（使用重打包后的模型路径）
model = Model(
    image_uri=image_uri,
    model_data=repacked_model_s3_uri,
    sagemaker_session=sagemaker_session,
    role=role,
)

inputs = CreateModelInput(instance_type="ml.m5.large", accelerator_type="ml.eia1.medium")

step_create_model = CreateModelStep(name="adultCreateModel", model=model, inputs=inputs)

step_register = RegisterModel(
    name="adultRegisterModel",
    estimator=xgb_train,
    model_data=repacked_model_s3_uri,
    content_types=["text/csv"],
    response_types=["text/csv"],
    inference_instances=["ml.t2.medium", "ml.m5.xlarge"],
    transform_instances=["ml.m5.xlarge"],
)

# 组装Pipeline并提交定义
pipeline = Pipeline(
    name="AdultModelPipeline",
    parameters=[original_model_s3_uri, repacked_model_s3_uri],
    steps=[step_repackage_model, step_create_model, step_register]
)
pipeline.upsert(role_arn=role)

2. 通过AWS Lambda触发Pipeline执行 #

编写Lambda函数，调用SageMaker API启动Pipeline执行（需给Lambda配置AmazonSageMakerFullAccess或细粒度权限）：

import boto3

sagemaker = boto3.client('sagemaker')

def lambda_handler(event, context):
    pipeline_name = "AdultModelPipeline"
    execution_params = {
        "OriginalModelS3Uri": event.get("original_model_uri", "s3://your-bucket/original/model.tar.gz"),
        "RepackedModelS3Uri": event.get("repacked_model_uri", "s3://your-bucket/repacked/model.tar.gz")
    }
    
    # 启动流水线执行
    response = sagemaker.start_pipeline_execution(
        PipelineName=pipeline_name,
        PipelineParameters=[{"Name": k, "Value": v} for k, v in execution_params.items()]
    )
    
    return {
        'statusCode': 200,
        'body': f"流水线已启动，执行ARN：{response['PipelineExecutionArn']}"
    }

3. 关键注意事项 #

• 权限配置：确保SageMaker角色拥有S3读写权限，Lambda角色拥有SagePipeline执行权限和必要的S3权限。
• 镜像选择：ProcessingStep使用的镜像需包含boto3、tarfile等依赖，推荐使用SageMaker官方Python镜像。
• 版本管理：注册模型时可添加model_package_group_name参数，实现模型版本的统一管理。

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