feature: Add environment variable support for SageMaker training job (#2218)

Co-authored-by: Chia-Eng <[email protected]>
Co-authored-by: icywang86rui <[email protected]>

Author: 3 people

src/sagemaker/estimator.py

@@ -123,6 +123,7 @@ def __init__(
         enable_network_isolation=False,
         profiler_config=None,
         disable_profiler=False,
+        environment=None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -266,6 +267,8 @@ def __init__(
                 ``disable_profiler`` parameter to ``True``.
             disable_profiler (bool): Specifies whether Debugger monitoring and profiling
                 will be disabled (default: ``False``).
+            environment (dict[str, str]) : Environment variables to be set for
+                use during training job (default: ``None``)
 
         """
         instance_count = renamed_kwargs(
@@ -352,6 +355,8 @@ def __init__(
         self.profiler_config = profiler_config
         self.disable_profiler = disable_profiler
 
+        self.environment = environment
+
         if not _region_supports_profiler(self.sagemaker_session.boto_region_name):
             self.disable_profiler = True
 
@@ -1471,6 +1476,7 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
         train_args["tags"] = estimator.tags
         train_args["metric_definitions"] = estimator.metric_definitions
         train_args["experiment_config"] = experiment_config
+        train_args["environment"] = estimator.environment
 
         if isinstance(inputs, TrainingInput):
             if "InputMode" in inputs.config:
@@ -1659,6 +1665,7 @@ def __init__(
         enable_sagemaker_metrics=None,
         profiler_config=None,
         disable_profiler=False,
+        environment=None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -1807,6 +1814,8 @@ def __init__(
                 ``disable_profiler`` parameter to ``True``.
             disable_profiler (bool): Specifies whether Debugger monitoring and profiling
                 will be disabled (default: ``False``).
+            environment (dict[str, str]) : Environment variables to be set for
+                use during training job (default: ``None``)
         """
         self.image_uri = image_uri
         self.hyperparam_dict = hyperparameters.copy() if hyperparameters else {}
@@ -1840,6 +1849,7 @@ def __init__(
             enable_network_isolation=enable_network_isolation,
             profiler_config=profiler_config,
             disable_profiler=disable_profiler,
+            environment=environment,
             **kwargs,
         )
 

src/sagemaker/session.py

@@ -456,6 +456,7 @@ def train(  # noqa: C901
         enable_sagemaker_metrics=None,
         profiler_rule_configs=None,
         profiler_config=None,
+        environment=None,
     ):
         """Create an Amazon SageMaker training job.
 
@@ -522,9 +523,12 @@ def train(  # noqa: C901
                 Series. For more information see:
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_AlgorithmSpecification.html#SageMaker-Type-AlgorithmSpecification-EnableSageMakerMetricsTimeSeries
                 (default: ``None``).
-            profiler_rule_configs (list[dict]): A list of profiler rule configurations.
+            profiler_rule_configs (list[dict]): A list of profiler rule
+                configurations.src/sagemaker/lineage/artifact.py:285
             profiler_config (dict): Configuration for how profiling information is emitted
                 with SageMaker Profiler. (default: ``None``).
+            environment (dict[str, str]) : Environment variables to be set for
+                use during training job (default: ``None``)
 
         Returns:
             str: ARN of the training job, if it is created.
@@ -556,6 +560,7 @@ def train(  # noqa: C901
             enable_sagemaker_metrics=enable_sagemaker_metrics,
             profiler_rule_configs=profiler_rule_configs,
             profiler_config=profiler_config,
+            environment=environment,
         )
         LOGGER.info("Creating training-job with name: %s", job_name)
         LOGGER.debug("train request: %s", json.dumps(train_request, indent=4))
@@ -588,6 +593,7 @@ def _get_train_request(  # noqa: C901
         enable_sagemaker_metrics=None,
         profiler_rule_configs=None,
         profiler_config=None,
+        environment=None,
     ):
         """Constructs a request compatible for creating an Amazon SageMaker training job.
 
@@ -657,6 +663,8 @@ def _get_train_request(  # noqa: C901
             profiler_rule_configs (list[dict]): A list of profiler rule configurations.
             profiler_config(dict): Configuration for how profiling information is emitted with
                 SageMaker Profiler. (default: ``None``).
+            environment (dict[str, str]) : Environment variables to be set for
+                use during training job (default: ``None``)
 
         Returns:
             Dict: a training request dict
@@ -699,6 +707,9 @@ def _get_train_request(  # noqa: C901
         if hyperparameters and len(hyperparameters) > 0:
             train_request["HyperParameters"] = hyperparameters
 
+        if environment is not None:
+            train_request["Environment"] = environment
+
         if tags is not None:
             train_request["Tags"] = tags
 

tests/integ/test_tf.py

@@ -36,6 +36,7 @@
 PARAMETER_SERVER_DISTRIBUTION = {"parameter_server": {"enabled": True}}
 MPI_DISTRIBUTION = {"mpi": {"enabled": True}}
 TAGS = [{"Key": "some-key", "Value": "some-value"}]
+ENV_INPUT = {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"}
 
 
 def test_mnist_with_checkpoint_config(
@@ -59,6 +60,7 @@ def test_mnist_with_checkpoint_config(
         metric_definitions=[{"Name": "train:global_steps", "Regex": r"global_step\/sec:\s(.*)"}],
         checkpoint_s3_uri=checkpoint_s3_uri,
         checkpoint_local_path=checkpoint_local_path,
+        environment=ENV_INPUT,
     )
     inputs = estimator.sagemaker_session.upload_data(
         path=os.path.join(MNIST_RESOURCE_PATH, "data"), key_prefix="scriptmode/mnist"
@@ -82,7 +84,13 @@ def test_mnist_with_checkpoint_config(
     actual_training_checkpoint_config = sagemaker_session.sagemaker_client.describe_training_job(
         TrainingJobName=training_job_name
     )["CheckpointConfig"]
+    actual_training_environment_variable_config = (
+        sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=training_job_name)[
+            "Environment"
+        ]
+    )
     assert actual_training_checkpoint_config == expected_training_checkpoint_config
+    assert actual_training_environment_variable_config == ENV_INPUT
 
 
 def test_server_side_encryption(sagemaker_session, tf_full_version, tf_full_py_version):

tests/unit/sagemaker/huggingface/test_estimator.py

@@ -149,6 +149,7 @@ def _create_train_job(version, base_framework_version):
         "tags": None,
         "vpc_config": None,
         "metric_definitions": None,
+        "environment": None,
         "experiment_config": None,
         "debugger_hook_config": {
             "CollectionConfigurations": [],

tests/unit/sagemaker/tensorflow/test_estimator.py

@@ -130,6 +130,7 @@ def _create_train_job(tf_version, horovod=False, ps=False, py_version="py2", smd
         "tags": None,
         "vpc_config": None,
         "metric_definitions": None,
+        "environment": None,
         "experiment_config": None,
         "profiler_rule_configs": [
             {

tests/unit/sagemaker/tensorflow/test_estimator_init.py

@@ -20,6 +20,8 @@
 
 REGION = "us-west-2"
 
+ENV_INPUT = {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"}
+
 
 @pytest.fixture()
 def sagemaker_session():
@@ -68,6 +70,26 @@ def test_framework_name(sagemaker_session):
     assert tf._framework_name == "tensorflow"
 
 
+def test_tf_add_environment_variables(sagemaker_session):
+    tf = _build_tf(
+        sagemaker_session,
+        framework_version="1.15.2",
+        py_version="py3",
+        environment=ENV_INPUT,
+    )
+    assert tf.environment == ENV_INPUT
+
+
+def test_tf_miss_environment_variables(sagemaker_session):
+    tf = _build_tf(
+        sagemaker_session,
+        framework_version="1.15.2",
+        py_version="py3",
+        environment=None,
+    )
+    assert not tf.environment
+
+
 def test_enable_sm_metrics(sagemaker_session):
     tf = _build_tf(
         sagemaker_session,

tests/unit/test_chainer.py

@@ -143,6 +143,7 @@ def _create_train_job(version, py_version):
         "tags": None,
         "vpc_config": None,
         "metric_definitions": None,
+        "environment": None,
         "experiment_config": None,
         "debugger_hook_config": {
             "CollectionConfigurations": [],

tests/unit/test_estimator.py

@@ -71,6 +71,7 @@
 CODECOMMIT_REPO_SSH = "ssh://git-codecommit.us-west-2.amazonaws.com/v1/repos/test-repo/"
 CODECOMMIT_BRANCH = "master"
 REPO_DIR = "/tmp/repo_dir"
+ENV_INPUT = {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"}
 
 DESCRIBE_TRAINING_JOB_RESULT = {"ModelArtifacts": {"S3ModelArtifacts": MODEL_DATA}}
 
@@ -241,6 +242,7 @@ def test_framework_all_init_args(sagemaker_session):
         checkpoint_local_path="file://local/checkpoint",
         enable_sagemaker_metrics=True,
         enable_network_isolation=True,
+        environment=ENV_INPUT,
     )
     _TrainingJob.start_new(f, "s3://mydata", None)
     sagemaker_session.train.assert_called_once()
@@ -275,6 +277,7 @@ def test_framework_all_init_args(sagemaker_session):
         },
         "metric_definitions": [{"Name": "validation-rmse", "Regex": "validation-rmse=(\\d+)"}],
         "encrypt_inter_container_traffic": True,
+        "environment": {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"},
         "experiment_config": None,
         "checkpoint_s3_uri": "s3://bucket/checkpoint",
         "checkpoint_local_path": "file://local/checkpoint",
@@ -1085,6 +1088,7 @@ def test_framework_with_spot_and_checkpoints(sagemaker_session):
         "use_spot_instances": True,
         "checkpoint_s3_uri": "s3://mybucket/checkpoints/",
         "checkpoint_local_path": "/tmp/checkpoints",
+        "environment": None,
         "experiment_config": None,
     }
 
@@ -2389,6 +2393,7 @@ def test_unsupported_type_in_dict():
     "tags": None,
     "vpc_config": None,
     "metric_definitions": None,
+    "environment": None,
     "experiment_config": None,
 }
 
@@ -2678,6 +2683,24 @@ def test_generic_to_fit_with_sagemaker_metrics_missing(sagemaker_session):
     assert "enable_sagemaker_metrics" not in args
 
 
+def test_add_environment_variables_to_train_args(sagemaker_session):
+    e = Estimator(
+        IMAGE_URI,
+        ROLE,
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        output_path=OUTPUT_PATH,
+        sagemaker_session=sagemaker_session,
+        environment=ENV_INPUT,
+    )
+
+    e.fit()
+
+    sagemaker_session.train.assert_called_once()
+    args = sagemaker_session.train.call_args[1]
+    assert args["environment"] == ENV_INPUT
+
+
 def test_generic_to_fit_with_sagemaker_metrics_enabled(sagemaker_session):
     e = Estimator(
         IMAGE_URI,

tests/unit/test_mxnet.py

@@ -65,6 +65,8 @@
 
 MODEL_PKG_RESPONSE = {"ModelPackageArn": "arn:model-pkg-arn"}
 
+ENV_INPUT = {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"}
+
 
 @pytest.fixture()
 def sagemaker_session():
@@ -144,6 +146,7 @@ def _get_train_args(job_name):
         "tags": None,
         "vpc_config": None,
         "metric_definitions": None,
+        "environment": None,
         "experiment_config": None,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
@@ -959,6 +962,38 @@ def test_create_model_with_custom_hosting_image(sagemaker_session):
     assert model.image_uri == custom_hosting_image
 
 
+def test_mx_add_environment_variables(
+    sagemaker_session, mxnet_training_version, mxnet_training_py_version
+):
+    mx = MXNet(
+        entry_point=SCRIPT_PATH,
+        framework_version=mxnet_training_version,
+        py_version=mxnet_training_py_version,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        environment=ENV_INPUT,
+    )
+    assert mx.environment == ENV_INPUT
+
+
+def test_mx_missing_environment_variables(
+    sagemaker_session, mxnet_training_version, mxnet_training_py_version
+):
+    mx = MXNet(
+        entry_point=SCRIPT_PATH,
+        framework_version=mxnet_training_version,
+        py_version=mxnet_training_py_version,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        environment=None,
+    )
+    assert not mx.environment
+
+
 def test_mx_enable_sm_metrics(sagemaker_session, mxnet_training_version, mxnet_training_py_version):
     mx = MXNet(
         entry_point=SCRIPT_PATH,

tests/unit/test_pytorch.py

@@ -46,6 +46,8 @@
 
 ENDPOINT_CONFIG_DESC = {"ProductionVariants": [{"ModelName": "model-1"}, {"ModelName": "model-2"}]}
 
+ENV_INPUT = {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"}
+
 LIST_TAGS_RESULT = {"Tags": [{"Key": "TagtestKey", "Value": "TagtestValue"}]}
 
 EXPERIMENT_CONFIG = {
@@ -146,6 +148,7 @@ def _create_train_job(version, py_version):
         "tags": None,
         "vpc_config": None,
         "metric_definitions": None,
+        "environment": None,
         "experiment_config": None,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
@@ -637,6 +640,30 @@ def test_pt_disable_sm_metrics(
     assert not pytorch.enable_sagemaker_metrics
 
 
+def test_pt_add_environment_variables(
+    sagemaker_session, pytorch_training_version, pytorch_training_py_version
+):
+    pytorch = _pytorch_estimator(
+        sagemaker_session,
+        framework_version=pytorch_training_version,
+        py_version=pytorch_training_py_version,
+        environment=ENV_INPUT,
+    )
+    assert pytorch.environment
+
+
+def test_pt_miss_environment_variables(
+    sagemaker_session, pytorch_training_version, pytorch_training_py_version
+):
+    pytorch = _pytorch_estimator(
+        sagemaker_session,
+        framework_version=pytorch_training_version,
+        py_version=pytorch_training_py_version,
+        environment=None,
+    )
+    assert not pytorch.environment
+
+
 def test_pt_default_sm_metrics(
     sagemaker_session, pytorch_training_version, pytorch_training_py_version
 ):