Address PR comments

Author: Chia-Eng

src/sagemaker/estimator.py

@@ -264,8 +264,7 @@ 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]) : A string to string map contains environment
-                variables to set in the Docker container.
+            environment (dict[str, str]) : Environment variables to be set for use during training job (default: ``None``)
 
         """
         instance_count = renamed_kwargs(
@@ -1808,8 +1807,7 @@ 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]) : A string to string map contains environment
-                variables to set in the Docker container.
+            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 {}

src/sagemaker/session.py

@@ -526,8 +526,7 @@ def train(  # 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]) : A string to string map contains environment
-                variables to set in the Docker container.
+            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.
@@ -662,8 +661,7 @@ 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]) : A string to string map contains environment
-                variables to set in the Docker container.
+            environment (dict[str, str]) : Environment variables to be set for use during training job (default: ``None``)
 
         Returns:
             Dict: a training request dict

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"
@@ -83,6 +85,7 @@ def test_mnist_with_checkpoint_config(
         TrainingJobName=training_job_name
     )["CheckpointConfig"]
     assert actual_training_checkpoint_config == expected_training_checkpoint_config
+    assert actual_training_checkpoint_config['Environment'] == ENV_INPUT
 
 
 def test_server_side_encryption(sagemaker_session, tf_full_version, tf_full_py_version):