Setup env variable in aws SageMaker container (bring your own container)

Question

We are using aws sagemaker that is using ecs container, Is there a way, we can setup environment variable (e.g. stage or prod) in container when calling sagemaker api using low level python sdk

Answer 1

Even invoking the API directly (which is lower level than using the python SDK) you cannot directly set environment arbitrary variables inside the container. You can however pass arbitrary hyperparameters in as configuration for a TrainingJob, for example pass in a hyperparameter like {"mystage": "prod"}. Hyperparameters show up in the container in a file called /opt/ml/input/config/hyperparameters.json which is a simple key-value map as a JSON object. You can use this to set the environment variable in a launching script like this:

#!/bin/bash

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

# Now run your code...

You can get SageMaker to invoke this script either by making it the ENTRYPOINT in your Dockerfile, or by calling it train and making sure it's on the PATH for the shell if you're not setting an ENTRYPOINT.