How can I pass environment variables to a custom training script in Amazon SageMaker using the Python SDK?

Question

I'm training a custom model using a script in Amazon SageMaker and launching the job with the Python SDK. I want to pass some environment variables (like API keys or config flags) to the training job so they’re accessible inside the script via os.environ.

Here’s a simplified version of my code:

from sagemaker.estimator import Estimator

estimator = Estimator(
    image_uri='123456789012.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:latest',
    role=role,
    instance_count=1,
    instance_type='ml.g5.xlarge',
    entry_point='train.py',
    source_dir='src',
    environment={
        'MY_API_KEY': 'abcdef123456',
        'DEBUG_MODE': 'true'
    }
)

In my training script, I try to read the variable:

import os

api_key = os.environ.get('MY_API_KEY')
print("API Key:", api_key)

Is this the correct way to pass environment variables to a SageMaker training job using the Python SDK? Are there any limitations or best practices I should be aware of, especially for sensitive information like API keys?

os.environ.get() is standard method used in Python - and it seems OK. Someone may say that only problem is that you can see it directly in system using linux command env. Other method is to keep keys in file env and use special module to read it - python-dotenv - this way you may have many projects with different keys. But if you send code to GitHub or backup then you may have to remeber to remove this file because someone could get your keys. — furas
– furas, Commented Mar 11 at 16:42

Janosch Woschitz · Accepted Answer · 2025-03-12 15:17:40Z

Yes, your approach is correct. Using the environment parameter in the Estimator and accessing variables with os.environ.get() in your script is the standard way to pass environment variables in SageMaker. As @furas pointed out in their comment, os.environ.get() is the common approach in Python.

That said, for handling secrets like API keys, it's better to avoid hardcoding them in your code or environment. A more secure approach is to store them in AWS Secrets Manager and fetch them inside your training script at runtime. You can pass the secret's name as an environment variable and retrieve the value securely using boto3:

import boto3  
import os  

secret_name = os.environ.get('API_KEY_SECRET_NAME')  
region = os.environ.get('AWS_REGION', 'us-west-2')  

client = boto3.client('secretsmanager', region_name=region)  
secret_value = client.get_secret_value(SecretId=secret_name)  
api_key = secret_value['SecretString']

print("API Key:", api_key)

This keeps the actual secret out of your environment config and allows for better access control via IAM.

Collectives™ on Stack Overflow

How can I pass environment variables to a custom training script in Amazon SageMaker using the Python SDK?

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related