fix integ tests

Author: laurenyu

tests/integ/test_local_mode.py

@@ -57,13 +57,13 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
 
 
 @pytest.fixture(scope='module')
-def mxnet_model(sagemaker_local_session):
+def mxnet_model(sagemaker_local_session, mxnet_full_version):
     script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
     data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
 
     mx = MXNet(entry_point=script_path, role='SageMakerRole',
-               train_instance_count=1, train_instance_type='local',
-               sagemaker_session=sagemaker_local_session)
+               train_instance_count=1, train_instance_type='local', launch_parameter_server=True,
+               sagemaker_session=sagemaker_local_session, framework_version=mxnet_full_version)
 
     train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
                                                    key_prefix='integ-test-data/mxnet_mnist/train')
@@ -306,16 +306,16 @@ def test_local_mode_serving_from_local_model(sagemaker_local_session, mxnet_mode
         fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)
 
 
-def test_mxnet_local_mode(sagemaker_local_session):
+def test_mxnet_local_mode(sagemaker_local_session, mxnet_full_version):
     local_mode_lock_fd = open(LOCK_PATH, 'w')
     local_mode_lock = local_mode_lock_fd.fileno()
 
     script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
     data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
 
     mx = MXNet(entry_point=script_path, role='SageMakerRole', py_version=PYTHON_VERSION,
-               train_instance_count=1, train_instance_type='local',
-               sagemaker_session=sagemaker_local_session)
+               train_instance_count=1, train_instance_type='local', launch_parameter_server=True,
+               sagemaker_session=sagemaker_local_session, framework_version=mxnet_full_version)
 
     train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
                                                    key_prefix='integ-test-data/mxnet_mnist/train')
@@ -338,15 +338,15 @@ def test_mxnet_local_mode(sagemaker_local_session):
         fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)
 
 
-def test_mxnet_local_data_local_script():
+def test_mxnet_local_data_local_script(mxnet_full_version):
     local_mode_lock_fd = open(LOCK_PATH, 'w')
     local_mode_lock = local_mode_lock_fd.fileno()
 
     script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
     data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
 
-    mx = MXNet(entry_point=script_path, role='SageMakerRole',
-               train_instance_count=1, train_instance_type='local',
+    mx = MXNet(entry_point=script_path, role='SageMakerRole', framework_version=mxnet_full_version,
+               train_instance_count=1, train_instance_type='local', launch_parameter_server=True,
                sagemaker_session=LocalNoS3Session())
 
     train_input = 'file://' + os.path.join(data_path, 'train')
@@ -368,14 +368,15 @@ def test_mxnet_local_data_local_script():
         fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)
 
 
-def test_local_transform_mxnet(sagemaker_local_session, tmpdir):
+def test_local_transform_mxnet(sagemaker_local_session, tmpdir, mxnet_full_version):
     local_mode_lock_fd = open(LOCK_PATH, 'w')
     local_mode_lock = local_mode_lock_fd.fileno()
     data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
     script_path = os.path.join(data_path, 'mnist.py')
 
     mx = MXNet(entry_point=script_path, role='SageMakerRole', train_instance_count=1,
-               train_instance_type='ml.c4.xlarge', sagemaker_session=sagemaker_local_session)
+               train_instance_type='ml.c4.xlarge', framework_version=mxnet_full_version,
+               sagemaker_session=sagemaker_local_session, launch_parameter_server=True)
 
     train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
                                                    key_prefix='integ-test-data/mxnet_mnist/train')

tests/integ/test_transformer.py

@@ -29,12 +29,13 @@
 
 
 @pytest.mark.continuous_testing
-def test_transform_mxnet(sagemaker_session):
+def test_transform_mxnet(sagemaker_session, mxnet_full_version):
     data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
     script_path = os.path.join(data_path, 'mnist.py')
 
     mx = MXNet(entry_point=script_path, role='SageMakerRole', train_instance_count=1,
-               train_instance_type='ml.c4.xlarge', sagemaker_session=sagemaker_session)
+               train_instance_type='ml.c4.xlarge', launch_parameter_server=True,
+               sagemaker_session=sagemaker_session, framework_version=mxnet_full_version)
 
     train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
                                                    key_prefix='integ-test-data/mxnet_mnist/train')
@@ -103,7 +104,7 @@ def test_attach_transform_kmeans(sagemaker_session):
         attached_transformer.wait()
 
 
-def test_transform_mxnet_vpc(sagemaker_session):
+def test_transform_mxnet_vpc(sagemaker_session, mxnet_full_version):
     data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
     script_path = os.path.join(data_path, 'mnist.py')
 
@@ -113,6 +114,7 @@ def test_transform_mxnet_vpc(sagemaker_session):
 
     mx = MXNet(entry_point=script_path, role='SageMakerRole', train_instance_count=1,
                train_instance_type='ml.c4.xlarge', sagemaker_session=sagemaker_session,
+               framework_version=mxnet_full_version, launch_parameter_server=True,
                subnets=subnet_ids, security_group_ids=[security_group_id])
 
     train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),