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 ====== diff --git a/src/sagemaker/local/image.py b/src/sagemaker/local/image.py index edd0543b08..74a9e50c1b 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_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'))