Add tox.ini and configure coverage and flake runs

.coveragerc_py27

@@ -0,0 +1,20 @@
+[run]
+branch = True
+timid = True
+
+[report]
+exclude_lines =
+    pragma: no cover
+    pragma: py2 no cover
+    if six.PY3
+    elif six.PY3
+
+partial_branches =
+    pragma: no cover
+    pragma: py2 no cover
+    if six.PY3
+    elif six.PY3
+
+show_missing = True
+
+fail_under = 90

.coveragerc_py36

@@ -0,0 +1,20 @@
+[run]
+branch = True
+timid = True
+
+[report]
+exclude_lines =
+    pragma: no cover
+    pragma: py3 no cover
+    if six.PY2
+    elif six.PY2
+
+partial_branches =
+    pragma: no cover
+    pragma: py3 no cover
+    if six.PY3
+    elif six.PY3
+
+show_missing = True
+
+fail_under = 90

.flake8

@@ -0,0 +1,3 @@
+[flake8]
+application_import_names = sagemaker_tensorflow_container, test
+import-order-style = google

docker/1.10.0/Dockerfile.cpu

@@ -70,4 +70,4 @@ RUN framework_installable_local=$(basename $framework_installable) \
     && rm $framework_installable_local \
     && rm $framework_support_installable_local
 
-ENV SAGEMAKER_TRAINING_MODULE tf_container.training:main
+ENV SAGEMAKER_TRAINING_MODULE sagemaker_tensorflow_container.training:main

docker/1.10.0/Dockerfile.gpu

@@ -69,4 +69,4 @@ RUN framework_installable_local=$(basename $framework_installable) && \
     rm $framework_installable_local && \
     rm $framework_support_installable_local
 
-ENV SAGEMAKER_TRAINING_MODULE tf_container.training:main
+ENV SAGEMAKER_TRAINING_MODULE sagemaker_tensorflow_container.training:main

setup.py

@@ -1,9 +1,23 @@
-import os
+# 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
+
 from glob import glob
+import os
 from os.path import basename
 from os.path import splitext
 
-from setuptools import setup, find_packages
+from setuptools import find_packages, setup
 
 
 def read(fname):
@@ -13,7 +27,8 @@ def read(fname):
 setup(
     name='sagemaker_tensorflow_container',
     version='2.0.0',
-    description='Open source library for creating TensorFlow containers to run on Amazon SageMaker.',
+    description='Open source library for creating '
+                'TensorFlow containers to run on Amazon SageMaker.',
 
     packages=find_packages(where='src', exclude=('test',)),
     package_dir={'': 'src'},
@@ -34,7 +49,8 @@ def read(fname):
         'Programming Language :: Python :: 3.6',
     ],
 
-    install_requires=['sagemaker-containers==2.1', 'numpy', 'scipy', 'sklearn', 'pandas', 'Pillow', 'h5py'],
+    install_requires=['sagemaker-containers==2.1', 'numpy', 'scipy', 'sklearn',
+                      'pandas', 'Pillow', 'h5py'],
     extras_require={
         'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'pytest-xdist', 'mock',
                  'sagemaker', 'tensorflow', 'docker-compose']

src/tf_container/training.py → src/sagemaker_tensorflow_container/training.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 import logging
+
 import sagemaker_containers.beta.framework as framework
 
 logger = logging.getLogger(__name__)

test/__init__.py

@@ -1,13 +1,13 @@
 #  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.
 from __future__ import absolute_import

test/conftest.py

@@ -1,14 +1,14 @@
 #  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.
 from __future__ import absolute_import
 
@@ -17,7 +17,6 @@
 
 import boto3
 import pytest
-
 from sagemaker import LocalSession, Session
 
 logger = logging.getLogger(__name__)
@@ -92,4 +91,3 @@ def skip_by_device_type(request, processor):
 @pytest.fixture(scope='session')
 def docker_image(docker_base_name, tag):
     return '{}:{}'.format(docker_base_name, tag)
-

test/integration/__init__.py

@@ -1,18 +1,18 @@
 #  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.
 from __future__ import absolute_import
 
 import logging
 
 logging.getLogger('boto3').setLevel(logging.INFO)
-logging.getLogger('botocore').setLevel(logging.INFO)
+logging.getLogger('botocore').setLevel(logging.INFO)

test/integration/docker_utils.py

@@ -1,14 +1,14 @@
 #  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.
 from __future__ import absolute_import
 
@@ -92,13 +92,13 @@ def execute_command(self, cmd):
         lines = []
         process = subprocess.Popen(docker_cmd, stdout=subprocess.PIPE)
         print(
-        '{}============================= container output ============================='.format(
-            CYAN_COLOR))
+            '{}========================= container output ========================='.format(
+                CYAN_COLOR))
         for line in iter(process.stdout.readline, b''):
             sys.stdout.write(line.decode('utf-8'))
             sys.stdout.flush()
             lines.append(line.decode('utf-8'))
-        msg = '\n{}============================= end of container output ============================='
+        msg = '\n{}========================= end of container output ========================='
         print(msg.format(CYAN_COLOR))
 
         process.wait()

test/integration/local/test_training.py

@@ -13,9 +13,10 @@
 from __future__ import absolute_import
 
 import os
-import pytest
 
+import pytest
 from sagemaker.tensorflow import TensorFlow
+
 from test.integration.docker_utils import Container
 
 
@@ -40,7 +41,8 @@ def test_mnist_cpu(sagemaker_local_session, docker_image):
                                          instance_type='local',
                                          sagemaker_local_session=sagemaker_local_session,
                                          docker_image=docker_image,
-                                         training_data_path='file://{}'.format(os.path.join(resource_path, 'data')))
+                                         training_data_path='file://{}'.format(
+                                             os.path.join(resource_path, 'data')))
     assert os.path.exists(os.path.join(output_path, 'my_model.h5')), 'model file not found'
 
 
@@ -51,10 +53,12 @@ def test_gpu(sagemaker_local_session, docker_image):
                            instance_type='local_gpu',
                            sagemaker_local_session=sagemaker_local_session,
                            docker_image=docker_image,
-                           training_data_path='file://{}'.format(os.path.join(resource_path, 'mnist', 'data')))
+                           training_data_path='file://{}'.format(
+                               os.path.join(resource_path, 'mnist', 'data')))
 
 
-def run_tf_single_training(script, instance_type, sagemaker_local_session, docker_image, training_data_path):
+def run_tf_single_training(script, instance_type, sagemaker_local_session,
+                           docker_image, training_data_path):
     estimator = TensorFlow(entry_point=script,
                            role='SageMakerRole',
                            training_steps=1,
@@ -66,4 +70,5 @@ def run_tf_single_training(script, instance_type, sagemaker_local_session, docke
                            base_job_name='test-tf-single')
 
     estimator.fit(training_data_path)
-    return estimator.sagemaker_session.sagemaker_client.s3_model_artifacts
+    model = estimator.create_model()
+    return model.model_data

test/resources/__init__.py

@@ -1,13 +1,13 @@
 #  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.
 from __future__ import absolute_import

test/resources/mnist/__init__.py

@@ -1,13 +1,13 @@
 #  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.
 from __future__ import absolute_import

test/unit/__init__.py

@@ -1,14 +1,13 @@
 #  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.
 from __future__ import absolute_import
-

test/unit/test_training.py

@@ -12,13 +12,15 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import pytest
 from mock import MagicMock, patch
+import pytest
 
-from tf_container import training
+from sagemaker_tensorflow_container import training
 
 MODULE_DIR = 's3://my/bucket'
 MODULE_NAME = 'script_name'
+LOG_LEVEL = 'Debug'
+
 
 @pytest.fixture
 def single_machine_training_env():
@@ -27,11 +29,26 @@ def single_machine_training_env():
     env.module_dir = MODULE_DIR
     env.module_name = MODULE_NAME
     env.hyperparameters = {}
+    env.log_level = LOG_LEVEL
 
     return env
 
+
 @patch('sagemaker_containers.beta.framework.modules.run_module')
 def test_single_machine(run_module, single_machine_training_env):
     training.train(single_machine_training_env)
     run_module.assert_called_with(MODULE_DIR, single_machine_training_env.to_cmd_args(),
                                   single_machine_training_env.to_env_vars(), MODULE_NAME)
+
+
+@patch('sagemaker_tensorflow_container.training.train')
+@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):
+    training_env.return_value = single_machine_training_env
+    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)