Local mode pass training env var (#411)

SageMaker Training sets various environment variables for every
training job, while Local Mode replicates none of these. This change adds
the environment variables for the AWS region and training job name.

Author: laurenyu

CHANGELOG.rst

@@ -2,6 +2,11 @@
 CHANGELOG
 =========
 
+1.11.1dev
+=========
+
+* enhancement: Local Mode: add training environment variables for AWS region and job name
+
 1.11.0
 ======
 

src/sagemaker/local/image.py

@@ -34,9 +34,13 @@
 import sagemaker
 from sagemaker.utils import get_config_value
 
-CONTAINER_PREFIX = "algo"
+CONTAINER_PREFIX = 'algo'
 DOCKER_COMPOSE_FILENAME = 'docker-compose.yaml'
 
+# Environment variables to be set during training
+REGION_ENV_NAME = 'AWS_REGION'
+TRAINING_JOB_NAME_ENV_NAME = 'TRAINING_JOB_NAME'
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.WARNING)
 
@@ -102,7 +106,12 @@ def train(self, input_data_config, hyperparameters):
             self.write_config_files(host, hyperparameters, input_data_config)
             shutil.copytree(data_dir, os.path.join(self.container_root, host, 'input', 'data'))
 
-        compose_data = self._generate_compose_file('train', additional_volumes=volumes)
+        training_env_vars = {
+            REGION_ENV_NAME: self.sagemaker_session.boto_region_name,
+            TRAINING_JOB_NAME_ENV_NAME: json.loads(hyperparameters.get(sagemaker.model.JOB_NAME_PARAM_NAME)),
+        }
+        compose_data = self._generate_compose_file('train', additional_volumes=volumes,
+                                                   additional_env_vars=training_env_vars)
         compose_command = self._compose()
 
         _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image)
@@ -149,7 +158,6 @@ def serve(self, model_dir, environment):
         logger.info('creating hosting dir in {}'.format(self.container_root))
 
         volumes = self._prepare_serving_volumes(model_dir)
-        env_vars = ['{}={}'.format(k, v) for k, v in environment.items()]
 
         # If the user script was passed as a file:// mount it to the container.
         if sagemaker.estimator.DIR_PARAM_NAME.upper() in environment:
@@ -161,7 +169,7 @@ def serve(self, model_dir, environment):
         _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image)
 
         self._generate_compose_file('serve',
-                                    additional_env_vars=env_vars,
+                                    additional_env_vars=environment,
                                     additional_volumes=volumes)
         compose_command = self._compose()
         self.container = _HostingContainer(compose_command)
@@ -384,7 +392,8 @@ def _generate_compose_file(self, command, additional_volumes=None, additional_en
         if aws_creds is not None:
             environment.extend(aws_creds)
 
-        environment.extend(additional_env_vars)
+        additional_env_var_list = ['{}={}'.format(k, v) for k, v in additional_env_vars.items()]
+        environment.extend(additional_env_var_list)
 
         if command == 'train':
             optml_dirs = {'output', 'output/data', 'input'}

tests/unit/test_image.py

@@ -49,12 +49,14 @@
     }
 ]
 HYPERPARAMETERS = {'a': 1,
-                   'b': 'bee',
-                   'sagemaker_submit_directory': json.dumps('s3://my_bucket/code')}
+                   'b': json.dumps('bee'),
+                   'sagemaker_submit_directory': json.dumps('s3://my_bucket/code'),
+                   'sagemaker_job_name': json.dumps('my-job')}
 
 LOCAL_CODE_HYPERPARAMETERS = {'a': 1,
                               'b': 2,
-                              'sagemaker_submit_directory': json.dumps('file:///tmp/code')}
+                              'sagemaker_submit_directory': json.dumps('file:///tmp/code'),
+                              'sagemaker_job_name': json.dumps('my-job')}
 
 
 @pytest.fixture()
@@ -244,6 +246,8 @@ def test_train(_download_folder, _cleanup, popen, _stream_output, LocalSession,
             for h in sagemaker_container.hosts:
                 assert config['services'][h]['image'] == image
                 assert config['services'][h]['command'] == 'train'
+                assert 'AWS_REGION={}'.format(REGION) in config['services'][h]['environment']
+                assert 'TRAINING_JOB_NAME=my-job' in config['services'][h]['environment']
 
         # assert that expected by sagemaker container output directories exist
         assert os.path.exists(os.path.join(sagemaker_container.container_root, 'output'))