Create regional S3 client for model exporting

src/tf_container/train_entry_point.py

@@ -10,7 +10,6 @@
 #  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 argparse
 import json
 import os
@@ -22,7 +21,6 @@
 
 import container_support as cs
 import tf_container.run
-import tf_container.s3_fs as s3_fs
 import tf_container.serve as serve
 
 _logger = tf_container.run.get_logger()

test/unit/test_experiment_trainer.py

@@ -1,18 +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 os
 
 import pytest
 from mock import patch, call, MagicMock, ANY
+
 from test.unit.utils import mock_import_modules
 
 mock_script = {}
@@ -112,7 +114,8 @@ def test_build_tf_config_with_multiple_hosts(trainer):
 @patch('botocore.session.get_session')
 @patch('os.environ')
 def test_configure_s3_file_system(os_env, botocore, boto_client, trainer):
-    region = os_env.get('AWS_REGION')
+    region = 'my-region'
+    os_env.get.return_value = region
 
     trainer.Trainer(customer_script=mock_script,
                     current_host=current_host,
@@ -128,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
 
 
@@ -186,7 +188,7 @@ def test_user_model_fn(modules, trainer):
     estimator = trainer._build_estimator(fake_run_config)
 
     estimator_mock = modules.estimator.Estimator
-    # Verify that _model_fn passed to Estimator correctly passes args through to user script model_fn 
+    # Verify that _model_fn passed to Estimator correctly passes args through to user script model_fn
     estimator_mock.assert_called_with(model_fn=ANY, params=expected_hps, config=fake_run_config)
     _, kwargs, = estimator_mock.call_args
     kwargs['model_fn'](1, 2, 3, 4)
@@ -392,7 +394,8 @@ def test_build_tf_config_with_multiple_hosts(trainer):
 @patch('botocore.session.get_session')
 @patch('os.environ')
 def test_configure_s3_file_system(os_env, botocore, boto_client, trainer_module):
-    region = os_env.get('AWS_REGION')
+    region = 'my-region'
+    os_env.get.return_value = region
 
     trainer_module.Trainer(customer_script=MOCK_SCRIPT,
                            current_host=CURRENT_HOST,
@@ -408,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()