Use 'AWS_REGION' again

Author: laurenyu

src/tf_container/experiment_trainer.py

@@ -1,26 +1,27 @@
 #  Copyright 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://www.apache.org/licenses/LICENSE-2.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
+#  
+#  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.
-import inspect
 
+import boto3
+import inspect
+import os
 import tensorflow as tf
+import tf_container.s3_fs as s3_fs
+from tf_container.run import logger
 from tensorflow.contrib.learn import RunConfig, Experiment
 from tensorflow.contrib.learn.python.learn import learn_runner
 from tensorflow.contrib.learn.python.learn.utils import saved_model_export_utils
 from tensorflow.contrib.training import HParams
 
-import tf_container.s3_fs as s3_fs
-from tf_container.run import logger
-
 
 class Trainer(object):
     DEFAULT_TRAINING_CHANNEL = 'training'
@@ -72,7 +73,7 @@ def __init__(self,
         self.customer_params = customer_params
 
         if model_path.startswith('s3://'):
-            s3_fs.configure_s3_fs(model_path, region_name=customer_params.get('sagemaker_region'))
+            s3_fs.configure_s3_fs(model_path)
 
     def _get_task_type(self, masters):
         if self.current_host in masters:

src/tf_container/s3_fs.py

@@ -3,7 +3,9 @@
 import container_support as cs
 
 
-def configure_s3_fs(checkpoint_path, region_name=None):
+def configure_s3_fs(checkpoint_path):
+    # If env variable is not set, defaults to None, which will use the global endpoint.
+    region_name = os.environ.get('AWS_REGION')
     s3 = boto3.client('s3', region_name=region_name)
 
     # We get the AWS region of the checkpoint bucket, which may be different from

src/tf_container/serve.py

@@ -10,26 +10,27 @@
 #  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.
-import csv
+
 import json
-import os
 import shutil
 import subprocess
-import time
-
 import boto3
+import container_support as cs
 import google.protobuf.json_format as json_format
+import os
+
 from grpc import StatusCode
 from grpc.framework.interfaces.face.face import AbortionError
-from six import StringIO
 from tensorflow.core.framework import tensor_pb2
-
-import container_support as cs
+from tf_container import proxy_client
+from six import StringIO
+import csv
 from container_support.serving import UnsupportedContentTypeError, UnsupportedAcceptTypeError, \
                                       JSON_CONTENT_TYPE, CSV_CONTENT_TYPE, \
                                       OCTET_STREAM_CONTENT_TYPE, ANY_CONTENT_TYPE
-from tf_container import proxy_client
 from tf_container.run import logger
+import time
+
 
 TF_SERVING_PORT = 9000
 GENERIC_MODEL_NAME = "generic_model"

src/tf_container/trainer.py

@@ -1,19 +1,20 @@
 #  Copyright 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://www.apache.org/licenses/LICENSE-2.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
+#  
+#  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.
-import inspect
 
+import boto3
+import inspect
+import os
 import tensorflow as tf
-
 from tf_container.run import logger
 import tf_container.s3_fs as s3_fs
 
@@ -61,7 +62,7 @@ def __init__(self,
         self.customer_params = customer_params
 
         if model_path.startswith('s3://'):
-            s3_fs.configure_s3_fs(model_path, region_name=customer_params.get('sagemaker_region'))
+            s3_fs.configure_s3_fs(model_path)
 
     def train(self):
         run_config = self._build_run_config()

test/unit/test_experiment_trainer.py

@@ -10,9 +10,11 @@
 #  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.
+import os
 
 import pytest
 from mock import patch, call, MagicMock, ANY
+
 from test.unit.utils import mock_import_modules
 
 mock_script = {}
@@ -113,12 +115,12 @@ def test_build_tf_config_with_multiple_hosts(trainer):
 @patch('os.environ')
 def test_configure_s3_file_system(os_env, botocore, boto_client, trainer):
     region = 'my-region'
+    os_env.get.return_value = region
 
     trainer.Trainer(customer_script=mock_script,
                     current_host=current_host,
                     hosts=hosts,
-                    model_path='s3://my/s3/path',
-                    customer_params={'sagemaker_region': region})
+                    model_path='s3://my/s3/path')
 
     boto_client.assert_called_once_with('s3', region_name=region)
     boto_client('s3', region_name=region).get_bucket_location.assert_called_once_with(Bucket='my')
@@ -129,6 +131,7 @@ def test_configure_s3_file_system(os_env, botocore, boto_client, trainer):
     ]
 
     os_env.__setitem__.assert_has_calls(calls, any_order=True)
+    os_env.get.assert_called_with('AWS_REGION')
 
 
 @patch('boto3.client')

test/unit/test_trainer.py

@@ -10,9 +10,11 @@
 #  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.
+import os
 
 import pytest
 from mock import patch, call, MagicMock, ANY
+
 from test.unit.utils import mock_import_modules
 
 
@@ -393,12 +395,12 @@ def test_build_tf_config_with_multiple_hosts(trainer):
 @patch('os.environ')
 def test_configure_s3_file_system(os_env, botocore, boto_client, trainer_module):
     region = 'my-region'
+    os_env.get.return_value = region
 
     trainer_module.Trainer(customer_script=MOCK_SCRIPT,
                            current_host=CURRENT_HOST,
                            hosts=HOSTS,
-                           model_path='s3://my/s3/path',
-                           customer_params={'sagemaker_region': region})
+                           model_path='s3://my/s3/path')
 
     boto_client.assert_called_once_with('s3', region_name=region)
     boto_client('s3', region_name=region).get_bucket_location.assert_called_once_with(Bucket='my')
@@ -409,6 +411,7 @@ def test_configure_s3_file_system(os_env, botocore, boto_client, trainer_module)
     ]
 
     os_env.__setitem__.assert_has_calls(calls, any_order=False)
+    os_env.get.assert_called_with('AWS_REGION')
 
 
 CUSTOMER_PARAMS = HYPERPARAMETERS.copy()