Set S3 environment variables (#112)

* Setting S3 environment variables before training starts

* Remove S3 environment variable setting in test training script

* Add unit tests

Author: icywang86rui

src/sagemaker_tensorflow_container/s3_utils.py

@@ -0,0 +1,41 @@
+# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the 'license' file accompanying this file. This file is
+# distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import os
+
+import boto3
+from six.moves.urllib.parse import urlparse
+
+
+def configure(model_dir, job_region):
+
+    if not model_dir:
+        return
+
+    s3 = boto3.client('s3', region_name=job_region)
+
+    # We get the AWS region of the checkpoint bucket, which may be different from
+    # the region this container is currently running in.
+    parsed_url = urlparse(model_dir)
+    bucket_name = parsed_url.netloc
+
+    bucket_location = s3.get_bucket_location(Bucket=bucket_name)['LocationConstraint']
+
+    # Configure environment variables used by TensorFlow S3 file system
+    if bucket_location:
+        os.environ['S3_REGION'] = bucket_location
+
+    # setting log level to WARNING
+    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
+    os.environ['S3_USE_HTTPS'] = '1'

src/sagemaker_tensorflow_container/training.py

@@ -15,11 +15,14 @@
 
 import json
 import logging
+import os
 import subprocess
 import time
 
 import sagemaker_containers.beta.framework as framework
 
+import sagemaker_tensorflow_container.s3_utils as s3_utils
+
 
 logger = logging.getLogger(__name__)
 
@@ -151,5 +154,6 @@ def main():
     """
     hyperparameters = framework.env.read_hyperparameters()
     env = framework.training_env(hyperparameters=hyperparameters)
+    s3_utils.configure(env.hyperparameters.get('model_dir'), os.environ.get('SAGEMAKER_REGION'))
     logger.setLevel(env.log_level)
     train(env)

test/resources/mnist/distributed_mnist.py

@@ -136,11 +136,6 @@ def _parse_args():
     tf_logger = tf_logging._get_logger()
     tf_logger.handlers = [_handler]
 
-    if args.checkpoint_path.startswith('s3://'):
-        os.environ['S3_REGION'] = 'us-west-2'
-        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'
-        os.environ['S3_USE_HTTPS'] = '1'
-
     train_data, train_labels = _load_training_data(args.train)
     eval_data, eval_labels = _load_testing_data(args.train)
 

test/unit/test_s3_utils.py

@@ -0,0 +1,38 @@
+# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the 'license' file accompanying this file. This file is
+# distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import os
+
+from mock import MagicMock, patch
+
+from sagemaker_tensorflow_container import s3_utils
+
+
+BUCKET_REGION = 'us-west-2'
+JOB_REGION = 'us-west-1'
+JOB_BUKCET = 'sagemaker-us-west-2-000-00-1'
+PREFIX = 'sagemaker/something'
+MODEL_DIR = 's3://{}/{}'.format(JOB_BUKCET, PREFIX)
+
+
+@patch('boto3.client')
+def test_configure(client):
+    s3 = MagicMock()
+    client.return_value = s3
+    loc = {'LocationConstraint': BUCKET_REGION}
+    s3.get_bucket_location.return_value = loc
+    s3_utils.configure(MODEL_DIR, JOB_REGION)
+    assert os.environ['S3_REGION'] == BUCKET_REGION
+    assert os.environ['TF_CPP_MIN_LOG_LEVEL'] == '1'
+    assert os.environ['S3_USE_HTTPS'] == '1'

test/unit/test_training.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import json
+import os
 
 from mock import MagicMock, patch
 import pytest
@@ -36,6 +37,8 @@
 WORKER_TASK = {'index': 0, 'type': 'worker'}
 PS_TASK_1 = {'index': 0, 'type': 'ps'}
 PS_TASK_2 = {'index': 1, 'type': 'ps'}
+MODEL_DIR = 's3://bucket/prefix'
+REGION = 'us-west-2'
 
 
 @pytest.fixture
@@ -61,7 +64,7 @@ def single_machine_training_env():
 
     env.module_dir = MODULE_DIR
     env.module_name = MODULE_NAME
-    env.hyperparameters = {}
+    env.hyperparameters = {'model_dir': MODEL_DIR}
     env.log_level = LOG_LEVEL
 
     return env
@@ -195,10 +198,14 @@ def test_build_tf_config_error():
 @patch('logging.Logger.setLevel')
 @patch('sagemaker_containers.beta.framework.training_env')
 @patch('sagemaker_containers.beta.framework.env.read_hyperparameters', return_value={})
-def test_main(read_hyperparameters, training_env, set_level, train, single_machine_training_env):
+@patch('sagemaker_tensorflow_container.s3_utils.configure')
+def test_main(configure_s3_env, read_hyperparameters, training_env,
+              set_level, train, single_machine_training_env):
     training_env.return_value = single_machine_training_env
+    os.environ['SAGEMAKER_REGION'] = REGION
     training.main()
     read_hyperparameters.assert_called_once_with()
     training_env.assert_called_once_with(hyperparameters={})
     set_level.assert_called_once_with(LOG_LEVEL)
     train.assert_called_once_with(single_machine_training_env)
+    configure_s3_env.assert_called_once()