From 34d12edbb44d4c972549615310cc101e5d8ea88d Mon Sep 17 00:00:00 2001 From: Lauren Yu <6631887+laurenyu@users.noreply.github.com> Date: Thu, 27 Sep 2018 15:46:34 -0700 Subject: [PATCH 1/3] Pass some SageMaker Training environment variables in Local Mode 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. --- src/sagemaker/local/image.py | 19 ++++++++++++++----- tests/unit/test_image.py | 10 +++++++--- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/src/sagemaker/local/image.py b/src/sagemaker/local/image.py index edd0543b08..4008254e44 100644 --- a/src/sagemaker/local/image.py +++ b/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_session.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'} diff --git a/tests/unit/test_image.py b/tests/unit/test_image.py index 2aed2bff17..be1f8f6ac6 100644 --- a/tests/unit/test_image.py +++ b/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')) From bf9f56bc37a3586157d0c9ed42c67207f263eb57 Mon Sep 17 00:00:00 2001 From: Lauren Yu <6631887+laurenyu@users.noreply.github.com> Date: Thu, 27 Sep 2018 16:34:09 -0700 Subject: [PATCH 2/3] update changelog --- CHANGELOG.rst | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e0e0ec2d76..b98b2adcf2 100644 --- a/CHANGELOG.rst +++ b/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 ====== From d96b960054443925d33fe8f2452782beac570166 Mon Sep 17 00:00:00 2001 From: Lauren Yu <6631887+laurenyu@users.noreply.github.com> Date: Thu, 27 Sep 2018 16:43:54 -0700 Subject: [PATCH 3/3] Use boto_region_name instead of boto_session.region_name --- src/sagemaker/local/image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sagemaker/local/image.py b/src/sagemaker/local/image.py index 4008254e44..74a9e50c1b 100644 --- a/src/sagemaker/local/image.py +++ b/src/sagemaker/local/image.py @@ -107,7 +107,7 @@ def train(self, input_data_config, hyperparameters): shutil.copytree(data_dir, os.path.join(self.container_root, host, 'input', 'data')) training_env_vars = { - REGION_ENV_NAME: self.sagemaker_session.boto_session.region_name, + 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,