在 aws SageMaker 容器中设置环境变量（自带容器）

Question

我们正在使用使用 ecs 容器的 aws sagemaker，有没有一种方法，我们可以在使用低级 python sdk 调用 sagemaker api 时在容器中设置环境变量（例如 stage 或 prod）

Answer 1

即使直接调用 API（比使用 python SDK 级别更低），您也无法直接在容器内设置环境任意变量。 但是，您可以传递任意超参数作为 TrainingJob 的配置，例如传递像

{"mystage": "prod"}

这样的超参数。 超参数显示在容器中名为

/opt/ml/input/config/hyperparameters.json

的文件中，这是一个作为 JSON 对象的简单键值映射。 您可以使用它在启动脚本中设置环境变量，如下所示：

#!/bin/bash

export STAGE=$(jq -r ".mystage" /opt/ml/input/config/hyperparameters.json)

# Now run your code...

您可以通过将其设为

ENTRYPOINT

中的

Dockerfile

来让 SageMaker 调用此脚本，或者通过调用它

train

并确保它位于 shell 的

PATH

上（如果您未设置

） ENTRYPOINT

。

Answer 2

您可以为ECS任务配置环境变量，这是区分开发/生产模式的常见方法。

environment - 传递给容器的环境变量。这 参数映射到 Docker 的创建容器部分中的 Env 远程 API 和 docker run 的 --env 选项。

我的答案与 Sagemaker 无关，因为我认为该问题仅涉及 ECS。

Answer 3

如果您使用的是低级 Boto SageMaker 客户端，那么使用

create_model

方法为模型设置环境变量可能会有效。此方法允许您将环境变量定义为 PrimaryContainer 的一部分，这些变量将与容器实例中的模型工件一起使用。

来源：

• https://boto3.readthedocs.io/en/latest/reference/services/sagemaker.html#SageMaker.Client.create_model

在高级 sagemaker Python 包中，也可以设置环境变量，例如通过

Estimator.deploy()

和

Estimator.create_model()

方法（因为附加参数将传递给

Model

）。

来源：

• http://sagemaker.readthedocs.io/en/latest/estimators.html#sagemaker.estimator.Estimator.deploy
• http://sagemaker.readthedocs.io/en/latest/estimators.html#sagemaker.estimator.Estimator.create_model
• http://sagemaker.readthedocs.io/en/latest/model.html

注意：这种方法似乎只在推理时有效，在训练过程中不起作用。

Answer 4

@leopd 答案的扩展是解析 Sagemaker 设置的 EnvVar

SM_TRAINING_ENV

并直接从 Python 代码中使用它 (

train.py

):

import json
import os

if __name__ == '__main__':
    envs = dict(os.environ)
    sm_training_env = envs.get('SM_TRAINING_ENV')
    sm_training_env = json.loads(sm_training_env)
    hyperparameters = sm_training_env.get('hyperparameters')
    #do_train(hyperparameters)

此外，刚刚意识到传递给估计器的超参数被设置为

SM_HP_NAMEOFPARAMETER

并且可以直接访问。

Answer 5

在为模型定义容器时可以使用“Environment”属性。请参阅boto3 create_model 文档了解更多信息。 @anthnyprschka 提到了这一点，但这里有一个例子。

def create_model(model_name: str, docker_img_url: str, account_id: str):
    container = {
        "Image": docker_img_url,
        "Mode": "SingleModel",
        # Pass through to model container running the inference endpoint
        "Environment": {
            "S3_BUCKET": os.environ["S3_BUCKET"],            
        },
    }
    role = os.environ["AWS_EXEC_ROLE"]

    sm_client = boto3.client(service_name="sagemaker")
    sm_client.create_model(
        ModelName=model_name,
        ExecutionRoleArn=f"arn:aws:iam::{account_id}:role/{role}",
        Containers=[container],
    )