Use the test argement framework_version in all tests (#158)

* Use the test argement framework_version in all tests

* Make flake8 happy

Author: icywang86rui

test/integration/local/test_horovod.py

@@ -32,7 +32,8 @@ def test_distributed_training_horovod_basic(instances,
                                             processes,
                                             sagemaker_local_session,
                                             docker_image,
-                                            tmpdir):
+                                            tmpdir,
+                                            framework_version):
     output_path = 'file://%s' % tmpdir
     estimator = TensorFlow(
         entry_point=os.path.join(RESOURCE_PATH, 'hvdbasic', 'train_hvd_basic.py'),
@@ -42,6 +43,7 @@ def test_distributed_training_horovod_basic(instances,
         train_instance_count=instances,
         image_name=docker_image,
         output_path=output_path,
+        framework_version=framework_version,
         hyperparameters={'sagemaker_mpi_enabled': True,
                          'sagemaker_network_interface_name': 'eth0',
                          'sagemaker_mpi_num_of_processes_per_host': processes})

test/integration/local/test_keras.py

@@ -26,7 +26,7 @@
 
 
 @pytest.mark.skip_gpu
-def test_keras_training(sagemaker_local_session, docker_image, tmpdir):
+def test_keras_training(sagemaker_local_session, docker_image, tmpdir, framework_version):
     entry_point = os.path.join(RESOURCE_PATH, 'keras_inception.py')
     output_path = 'file://{}'.format(tmpdir)
 
@@ -39,14 +39,14 @@ def test_keras_training(sagemaker_local_session, docker_image, tmpdir):
         sagemaker_session=sagemaker_local_session,
         model_dir='/opt/ml/model',
         output_path=output_path,
-        framework_version='1.11.0',
+        framework_version=framework_version,
         py_version='py3')
 
     estimator.fit()
 
     model = serving.Model(model_data=output_path,
                           role='SageMakerRole',
-                          framework_version='1.11.0',
+                          framework_version=framework_version,
                           sagemaker_session=sagemaker_local_session)
 
     predictor = model.deploy(initial_instance_count=1, instance_type='local')

test/integration/local/test_training.py

@@ -39,52 +39,62 @@ def test_py_versions(docker_image, processor, py_full_version):
 
 
 @pytest.mark.skip_gpu
-def test_mnist_cpu(sagemaker_local_session, docker_image, tmpdir):
+def test_mnist_cpu(sagemaker_local_session, docker_image, tmpdir, framework_version):
     output_path = 'file://{}'.format(tmpdir)
     run_tf_training(script=os.path.join(RESOURCE_PATH, 'mnist', 'mnist.py'),
                     instance_type='local',
                     instance_count=1,
                     sagemaker_local_session=sagemaker_local_session,
                     docker_image=docker_image,
+                    framework_version=framework_version,
                     output_path=output_path,
                     training_data_path='file://{}'.format(
                         os.path.join(RESOURCE_PATH, 'mnist', 'data')))
     _assert_files_exist_in_tar(output_path, ['my_model.h5'])
 
 
 @pytest.mark.skip_cpu
-def test_gpu(sagemaker_local_session, docker_image):
+def test_gpu(sagemaker_local_session, docker_image, framework_version):
     run_tf_training(script=os.path.join(RESOURCE_PATH, 'gpu_device_placement.py'),
                     instance_type='local_gpu',
                     instance_count=1,
                     sagemaker_local_session=sagemaker_local_session,
                     docker_image=docker_image,
+                    framework_version=framework_version,
                     training_data_path='file://{}'.format(
                         os.path.join(RESOURCE_PATH, 'mnist', 'data')))
 
 
 @pytest.mark.skip_gpu
-def test_distributed_training_cpu_no_ps(sagemaker_local_session, docker_image, tmpdir):
+def test_distributed_training_cpu_no_ps(sagemaker_local_session,
+                                        docker_image,
+                                        tmpdir,
+                                        framework_version):
     output_path = 'file://{}'.format(tmpdir)
     run_tf_training(script=os.path.join(RESOURCE_PATH, 'mnist', 'mnist_estimator.py'),
                     instance_type='local',
                     instance_count=2,
                     sagemaker_local_session=sagemaker_local_session,
                     docker_image=docker_image,
+                    framework_version=framework_version,
                     output_path=output_path,
                     training_data_path='file://{}'.format(
                         os.path.join(RESOURCE_PATH, 'mnist', 'data-distributed')))
     _assert_files_exist_in_tar(output_path, TF_CHECKPOINT_FILES)
 
 
 @pytest.mark.skip_gpu
-def test_distributed_training_cpu_ps(sagemaker_local_session, docker_image, tmpdir):
+def test_distributed_training_cpu_ps(sagemaker_local_session,
+                                     docker_image,
+                                     tmpdir,
+                                     framework_version):
     output_path = 'file://{}'.format(tmpdir)
     run_tf_training(script=os.path.join(RESOURCE_PATH, 'mnist', 'mnist_estimator.py'),
                     instance_type='local',
                     instance_count=2,
                     sagemaker_local_session=sagemaker_local_session,
                     docker_image=docker_image,
+                    framework_version=framework_version,
                     output_path=output_path,
                     hyperparameters={'sagemaker_parameter_server_enabled': True},
                     training_data_path='file://{}'.format(
@@ -96,7 +106,10 @@ def run_tf_training(script,
                     instance_type,
                     instance_count,
                     sagemaker_local_session,
-                    docker_image, training_data_path, output_path=None,
+                    docker_image,
+                    framework_version,
+                    training_data_path,
+                    output_path=None,
                     hyperparameters=None):
 
     hyperparameters = hyperparameters or {}
@@ -111,7 +124,7 @@ def run_tf_training(script,
                            output_path=output_path,
                            hyperparameters=hyperparameters,
                            base_job_name='test-tf',
-                           framework_version='1.11.0',
+                           framework_version=framework_version,
                            py_version='py3')
 
     estimator.fit(training_data_path)

test/integration/sagemaker/test_horovod.py

@@ -24,7 +24,8 @@ def test_distributed_training_horovod(sagemaker_session,
                                       sagemaker_local_session,
                                       instance_type,
                                       ecr_image,
-                                      tmpdir):
+                                      tmpdir,
+                                      framework_version):
 
     mpi_options = '-verbose -x orte_base_help_aggregate=0'
     estimator = TensorFlow(
@@ -33,7 +34,7 @@ def test_distributed_training_horovod(sagemaker_session,
         train_instance_type=instance_type,
         train_instance_count=2,
         image_name=ecr_image,
-        framework_version='1.12',
+        framework_version=framework_version,
         py_version='py3',
         script_mode=True,
         hyperparameters={'sagemaker_mpi_enabled': True,

test/integration/sagemaker/test_mnist.py

@@ -21,7 +21,7 @@
 from sagemaker_tensorflow_container.training import SAGEMAKER_PARAMETER_SERVER_ENABLED
 
 
-def test_mnist(sagemaker_session, ecr_image, instance_type):
+def test_mnist(sagemaker_session, ecr_image, instance_type, framework_version):
     resource_path = os.path.join(os.path.dirname(__file__), '../..', 'resources')
     script = os.path.join(resource_path, 'mnist', 'mnist.py')
     estimator = TensorFlow(entry_point=script,
@@ -30,7 +30,7 @@ def test_mnist(sagemaker_session, ecr_image, instance_type):
                            train_instance_count=1,
                            sagemaker_session=sagemaker_session,
                            image_name=ecr_image,
-                           framework_version='1.12.0',
+                           framework_version=framework_version,
                            py_version='py3',
                            base_job_name='test-sagemaker-mnist')
     inputs = estimator.sagemaker_session.upload_data(
@@ -40,7 +40,7 @@ def test_mnist(sagemaker_session, ecr_image, instance_type):
     _assert_s3_file_exists(estimator.model_data)
 
 
-def test_distributed_mnist_no_ps(sagemaker_session, ecr_image, instance_type):
+def test_distributed_mnist_no_ps(sagemaker_session, ecr_image, instance_type, framework_version):
     resource_path = os.path.join(os.path.dirname(__file__), '../..', 'resources')
     script = os.path.join(resource_path, 'mnist', 'mnist.py')
     estimator = TensorFlow(entry_point=script,
@@ -49,7 +49,7 @@ def test_distributed_mnist_no_ps(sagemaker_session, ecr_image, instance_type):
                            train_instance_type=instance_type,
                            sagemaker_session=sagemaker_session,
                            image_name=ecr_image,
-                           framework_version='1.12.0',
+                           framework_version=framework_version,
                            py_version='py3',
                            base_job_name='test-tf-sm-distributed-mnist')
     inputs = estimator.sagemaker_session.upload_data(
@@ -59,7 +59,7 @@ def test_distributed_mnist_no_ps(sagemaker_session, ecr_image, instance_type):
     _assert_s3_file_exists(estimator.model_data)
 
 
-def test_distributed_mnist_ps(sagemaker_session, ecr_image, instance_type):
+def test_distributed_mnist_ps(sagemaker_session, ecr_image, instance_type, framework_version):
     resource_path = os.path.join(os.path.dirname(__file__), '..', '..', 'resources')
     script = os.path.join(resource_path, 'mnist', 'mnist_estimator.py')
     estimator = TensorFlow(entry_point=script,
@@ -69,7 +69,7 @@ def test_distributed_mnist_ps(sagemaker_session, ecr_image, instance_type):
                            train_instance_type=instance_type,
                            sagemaker_session=sagemaker_session,
                            image_name=ecr_image,
-                           framework_version='1.12.0',
+                           framework_version=framework_version,
                            py_version='py3',
                            base_job_name='test-tf-sm-distributed-mnist')
     inputs = estimator.sagemaker_session.upload_data(
@@ -80,7 +80,7 @@ def test_distributed_mnist_ps(sagemaker_session, ecr_image, instance_type):
     _assert_s3_file_exists(estimator.model_data)
 
 
-def test_s3_plugin(sagemaker_session, ecr_image, instance_type, region):
+def test_s3_plugin(sagemaker_session, ecr_image, instance_type, region, framework_version):
     resource_path = os.path.join(os.path.dirname(__file__), '..', '..', 'resources')
     script = os.path.join(resource_path, 'mnist', 'mnist_estimator.py')
     estimator = TensorFlow(entry_point=script,
@@ -91,7 +91,7 @@ def test_s3_plugin(sagemaker_session, ecr_image, instance_type, region):
                                # Disable throttling for checkpoint and model saving
                                'throttle-secs': 0,
                                # Without the patch training jobs would fail around 100th to
-                               # 150th steps
+                               # 150th step
                                'max-steps': 200,
                                # Large batch size would result in a larger checkpoint file
                                'batch-size': 2048,
@@ -103,7 +103,7 @@ def test_s3_plugin(sagemaker_session, ecr_image, instance_type, region):
                            train_instance_type=instance_type,
                            sagemaker_session=sagemaker_session,
                            image_name=ecr_image,
-                           framework_version='1.12.0',
+                           framework_version=framework_version,
                            py_version='py3',
                            base_job_name='test-tf-sm-s3-mnist')
     estimator.fit('s3://sagemaker-sample-data-{}/tensorflow/mnist'.format(region))