如何从AWS StepFunctions TuningStep配合SageMaker获取最优训练模型

结论 #

stepfunctions.steps.TuningStep 确实没有内置与SageMaker Pipelines中TuningStep对等的get_expected_model()、get_top_model_s3_uri()方法，两类TuningStep分属不同工作流框架，设计逻辑存在差异。

差异说明 #

• Step Functions Data Science SDK提供的TuningStep是对SageMaker超参数调优API的轻量封装，仅返回调优任务的原生响应字段，没有额外封装最优模型提取的语法糖能力。
• 你提到的get_expected_model()等方法是SageMaker Pipelines专属组件的能力，该框架对模型训练上下游链路做了更深的内置适配，相关方法不支持在Step Functions SDK的组件中使用。

可行解决方案 #

你可以通过以下两种方式从TuningStep中获取最优训练模型：

方式1：直接引用TuningStep输出字段

调优任务完成后，TuningStep的输出结果默认包含最优训练任务的全量信息，你可以直接通过字段路径在后续步骤中引用最优模型地址：

${TuningStep.Output.BestTrainingJob.TrainingJobOutput.S3ModelArtifacts}

将上述路径直接传入后续model_step的模型Artifact配置项即可直接使用。

方式2：新增LambdaStep提取最优模型信息

如果需要对最优模型做额外的指标校验、字段预处理，可以在tuning_step和model_step之间新增LambdaStep，调用SageMaker API拉取调优任务的详细结果，示例Lambda代码如下：

import boto3
sagemaker_client = boto3.client('sagemaker')

def lambda_handler(event, context):
    # 从入参获取TuningStep返回的调优任务名称
    tuning_job_name = event['TuningStep']['Output']['HyperParameterTuningJobName']
    tuning_detail = sagemaker_client.describe_hyper_parameter_tuning_job(
        HyperParameterTuningJobName=tuning_job_name
    )
    best_job = tuning_detail['BestTrainingJob']
    return {
        "best_model_s3_uri": best_job['TrainingJobOutput']['S3ModelArtifacts'],
        "best_training_job_name": best_job['TrainingJobName'],
        "best_job_metrics": best_job['FinalHyperParameterTuningJobObjectiveMetric']
    }

后续步骤可以直接引用LambdaStep返回的best_model_s3_uri字段创建模型。

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