使用Boto3 API传递SageMaker训练容器参数失败求助

问题描述 #

将自定义Docker镜像发布至AWS ECR后，通过Boto3调用create_training_job启动SageMaker训练作业，在AlgorithmSpecification.ContainerArguments中传入'--mode training'和'--region_id 1'参数，训练作业启动后报错：

entry_point.py: error: the following arguments are required: --mode

错误原因 #

1. ContainerArguments参数格式错误：SageMaker的ContainerArguments要求每个命令行参数（包括参数名和对应的值）作为独立的列表元素传入。如果将--mode training放在同一个字符串中，Docker会将其当作单个参数传递给entry_point.py，argparse无法识别这是--mode参数及其值training，导致认为--mode参数缺失。
2. 参数名不匹配：entry_point.py中定义的区域参数是--region，但Boto3中传递的是--region_id，同时代码中引用的是args.region_id，与定义的args.region不匹配，后续也会引发错误。

解决方案 #

1. 修正Boto3的ContainerArguments格式 #

将参数名和对应的值拆分为独立的列表元素：

session = boto3.Session(profile_name='algoprod')
client = session.client('sagemaker', region_name='us-east-1')
training_job_name = 'sagemaker-training-demo'
resp = client.create_training_job(
                    TrainingJobName=training_job_name,
                    RoleArn="xxxx",
                    AlgorithmSpecification={
                            'TrainingImage': "image:latest",
                            'TrainingInputMode': "File",
                            'ContainerArguments': [
                                    '--mode', 'training',
                                    '--region', '1'
                             ]
                    }
)

print(resp)

2. 修正entry_point.py的参数引用 #

确保参数名和代码中的引用一致：

parser = argparse.ArgumentParser()
parser.add_argument("--mode", type=str, required=True)
parser.add_argument("--region", type=int)

args = parser.parse_args()

if args.mode == "inference":
    run_inference(args.region)
elif args.mode == "training":
    run_training(args.region)
else:
    raise ValueError(f"Unknown mode: {args.mode}")

如果希望保留region_id作为参数名，可修改entry_point.py如下：

parser = argparse.ArgumentParser()
parser.add_argument("--mode", type=str, required=True)
parser.add_argument("--region_id", type=int)

args = parser.parse_args()

if args.mode == "inference":
    run_inference(args.region_id)
elif args.mode == "training":
    run_training(args.region_id)
else:
    raise ValueError(f"Unknown mode: {args.mode}")

对应的Boto3参数需调整为：

'ContainerArguments': [
        '--mode', 'training',
        '--region_id', '1'
]

3. 本地验证（可选） #

在推送镜像到ECR之前，可本地运行Docker镜像测试参数传递是否正常：

docker run image:latest --mode training --region 1

确认脚本能正常执行后，再重新发布镜像并启动SageMaker训练作业。

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