如何在SageMaker的RegisterModel中设置SAGEMAKER_SUBMIT_DIRECTORY环境变量

问题解答 #

首先，RegisterModel 函数本身没有直接传入自定义环境变量的参数，但是可以通过前置构造带环境变量的Model对象的方式，实现将SAGEMAKER_SUBMIT_DIRECTORY写入注册模型配置的需求，具体操作如下：

操作步骤

1. 首先构造包含目标环境变量的Model实例
   你需要先导入SageMaker的Model类，传入推理镜像、模型数据、执行角色，同时通过env参数配置你需要的环境变量：

from sagemaker.model import Model

# 构造带自定义环境变量的模型对象
custom_model = Model(
    image_uri=estimator.training_image_uri(), # 也可直接指定推理镜像的URI
    model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,
    role=role, # 替换为你的SageMaker执行角色ARN
    env={
        "SAGEMAKER_SUBMIT_DIRECTORY": "s3://<你的S3桶名>/<推理代码存储路径>/", # 替换为实际的S3路径
        "SAGEMAKER_PROGRAM": "inference.py" # 可选，指定推理入口脚本名，和提交目录搭配使用
    }
)

2. 调整RegisterModel的传入参数
   将你原先传入的estimator、model_data参数移除，改为传入上面构造好的custom_model对象即可：

step_register = RegisterModel(
    name="RegisterCustomModel",
    model=custom_model, # 传入带环境变量的模型对象
    content_types=["text/csv"],
    response_types=["text/csv"],
    inference_instances=["ml.t2.medium", "ml.m5.large"],
    transform_instances=["ml.m5.large"],
    model_package_group_name=model_package_group_name,
    approval_status=model_approval_status,
    model_metrics=model_metrics,
)

注意事项

• 上述方式注册的模型包会自动携带你配置的环境变量，部署时不需要额外配置即可生效。
• SAGEMAKER_SUBMIT_DIRECTORY对应的S3路径需要确保SageMaker执行角色有读取权限，路径下需要包含推理脚本、依赖文件等所需资源。
• 不要同时给RegisterModel传入model、estimator、model_data三个参数，否则会触发参数冲突报错。

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