我们正在使用使用 ecs 容器的 aws sagemaker,有没有一种方法,我们可以在使用低级 python sdk 调用 sagemaker api 时在容器中设置环境变量(例如 stage 或 prod)
即使直接调用 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
。
您可以为ECS任务配置环境变量,这是区分开发/生产模式的常见方法。
environment - 传递给容器的环境变量。这 参数映射到 Docker 的创建容器部分中的 Env 远程 API 和 docker run 的 --env 选项。
我的答案与 Sagemaker 无关,因为我认为该问题仅涉及 ECS。
如果您使用的是低级 Boto SageMaker 客户端,那么使用
create_model
方法为模型设置环境变量可能会有效。此方法允许您将环境变量定义为 PrimaryContainer 的一部分,这些变量将与容器实例中的模型工件一起使用。
来源:
在高级 sagemaker Python 包中,也可以设置环境变量,例如通过
Estimator.deploy()
和 Estimator.create_model()
方法(因为附加参数将传递给 Model
)。
来源:
注意:这种方法似乎只在推理时有效,在训练过程中不起作用。
@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
并且可以直接访问。
在为模型定义容器时可以使用“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],
)