Set S3 environment variables (aws#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

@@ -20,6 +20,7 @@
 
 def configure(model_dir, job_region):
 
+<<<<<<< HEAD
     os.environ['S3_REGION'] = _s3_region(job_region, model_dir)
 
     # setting log level to WARNING
@@ -41,3 +42,24 @@ def _s3_region(job_region, model_dir):
         return bucket_location or job_region
     else:
         return 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'
+>>>>>>> Set S3 environment variables (#112)

src/sagemaker_tensorflow_container/training.py

@@ -15,10 +15,14 @@
 import json
 import logging
 <<<<<<< HEAD
+<<<<<<< HEAD
 import multiprocessing
 import os
 =======
 >>>>>>> Add distributed training support (#98)
+=======
+import os
+>>>>>>> Set S3 environment variables (#112)
 import subprocess
 import time
 
@@ -27,6 +31,8 @@
 
 from sagemaker_tensorflow_container import s3_utils
 
+import sagemaker_tensorflow_container.s3_utils as s3_utils
+
 
 logger = logging.getLogger(__name__)
 
@@ -305,6 +311,7 @@ def main():
     """
     hyperparameters = framework.env.read_hyperparameters()
     env = framework.training_env(hyperparameters=hyperparameters)
+<<<<<<< HEAD
 
     user_hyperparameters = env.hyperparameters
 
@@ -318,3 +325,8 @@ def main():
     s3_utils.configure(user_hyperparameters.get('model_dir'), os.environ.get('SAGEMAKER_REGION'))
     train(env, framework.mapping.to_cmd_args(user_hyperparameters))
     _log_model_missing_warning(MODEL_DIR)
+=======
+    s3_utils.configure(env.hyperparameters.get('model_dir'), os.environ.get('SAGEMAKER_REGION'))
+    logger.setLevel(env.log_level)
+    train(env)
+>>>>>>> Set S3 environment variables (#112)

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

@@ -32,6 +32,7 @@ def test_configure(client):
     client.return_value = s3
     loc = {'LocationConstraint': BUCKET_REGION}
     s3.get_bucket_location.return_value = loc
+<<<<<<< HEAD
 
     s3_utils.configure(MODEL_DIR, JOB_REGION)
 
@@ -46,3 +47,9 @@ def test_configure_local_dir():
     assert os.environ['S3_REGION'] == JOB_REGION
     assert os.environ['TF_CPP_MIN_LOG_LEVEL'] == '1'
     assert os.environ['S3_USE_HTTPS'] == '1'
+=======
+    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'
+>>>>>>> Set S3 environment variables (#112)

test/unit/test_training.py

@@ -319,6 +319,7 @@ def test_main_tuning_mpi_model_dir(configure_s3_env, read_hyperparameters, train
 >>>>>>> Add tox.ini and configure coverage and flake runs (#80)
 =======
 import json
+import os
 
 >>>>>>> Add distributed training support (#98)
 from mock import MagicMock, patch
@@ -343,6 +344,8 @@ def test_main_tuning_mpi_model_dir(configure_s3_env, read_hyperparameters, train
 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
@@ -368,7 +371,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
@@ -505,11 +508,18 @@ 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)
+<<<<<<< HEAD
 >>>>>>> Add tox.ini and configure coverage and flake runs (#80)
+=======
+    configure_s3_env.assert_called_once()
+>>>>>>> Set S3 environment variables (#112)