如何在SageMaker流水线中从模型注册表加载已训练模型执行推理？

方案可行性说明 #

你的需求完全可行，问题出在从模型注册表加载模型的具体实现方式上。以下是修正后的代码和配置步骤：

1. 从模型注册表获取模型信息 #

首先通过模型包ARN获取模型的镜像URI和模型数据S3路径：

import boto3

# 初始化SageMaker客户端
sm_client = boto3.client('sagemaker')
# 替换为你的目标模型包ARN
model_package_arn = "arn:aws:sagemaker:your-region:your-account-id:model-package/your-model-package/version"

# 获取模型包详情
model_package = sm_client.describe_model_package(ModelPackageName=model_package_arn)
# 提取推理所需的镜像和模型数据路径
model_image_uri = model_package['InferenceSpecification']['Containers'][0]['Image']
model_data_s3_uri = model_package['InferenceSpecification']['Containers'][0]['ModelDataUrl']

2. 替换原代码中的Model创建逻辑 #

用上面获取的模型信息替换原训练步骤的输出：

# 从模型注册表加载的信息创建Model对象
model = Model(
    image_uri=model_image_uri,        
    model_data=model_data_s3_uri,
    name=f"from-idea-to-prod-xgboost-model",
    sagemaker_session=session,
    role=sm_role,
)

# 定义ModelStep步骤
step_create_model = ModelStep(
    name=f"{pipeline_name}-model",
    step_args=model.create(instance_type="ml.m5.large", accelerator_type="ml.eia1.medium"),
)

3. 配置Lambda触发流水线 #

要让流水线支持Lambda触发，需完成以下配置：

• 给Lambda执行角色添加最小权限：允许调用sagemaker:StartPipelineExecution，同时确保角色有权限访问流水线相关资源
• Lambda函数示例代码：

import boto3

def lambda_handler(event, context):
    sm_client = boto3.client('sagemaker')
    pipeline_name = "your-pipeline-name"
    
    # 启动流水线执行，可按需传递参数
    execution_response = sm_client.start_pipeline_execution(
        PipelineName=pipeline_name,
        PipelineExecutionDisplayName="Lambda-Triggered-Batch-Transform",
        # 若需动态传入数据集路径等参数，可添加如下配置
        # PipelineParameters=[
        #     {"Name": "InputDataS3Uri", "Value": event["input_data_uri"]}
        # ]
    )
    
    return {
        "statusCode": 200,
        "body": f"流水线已启动，执行ARN: {execution_response['PipelineExecutionArn']}"
    }

常见问题排查 #

• 若模型包ARN创建失败，检查以下几点：
  • 模型包是否处于Approved状态（只有已批准的模型包可用于推理）
  • SageMaker会话角色是否拥有sagemaker:DescribeModelPackage权限
  • 模型包中的镜像URI和S3模型数据路径是否有效，且角色有权限访问对应S3桶

内容的提问来源于stack exchange，提问作者Luis Quiñones