Skip gpu tests in regions without ml.p2.xlarge (#461)

Author: yangaws

tests/integ/__init__.py

@@ -16,11 +16,14 @@
 import os
 import sys
 
+import boto3
+
 DATA_DIR = os.path.join(os.path.dirname(__file__), '..', 'data')
 TRAINING_DEFAULT_TIMEOUT_MINUTES = 20
 TUNING_DEFAULT_TIMEOUT_MINUTES = 20
 TRANSFORM_DEFAULT_TIMEOUT_MINUTES = 20
 PYTHON_VERSION = 'py' + str(sys.version_info.major)
+REGION = boto3.session.Session().region_name
 
 logging.getLogger('boto3').setLevel(logging.INFO)
 logging.getLogger('botocore').setLevel(logging.INFO)

tests/integ/test_chainer_train.py

@@ -22,7 +22,7 @@
 from sagemaker.chainer.estimator import Chainer
 from sagemaker.chainer.model import ChainerModel
 from sagemaker.utils import sagemaker_timestamp
-from tests.integ import DATA_DIR, PYTHON_VERSION, TRAINING_DEFAULT_TIMEOUT_MINUTES
+from tests.integ import DATA_DIR, PYTHON_VERSION, TRAINING_DEFAULT_TIMEOUT_MINUTES, REGION
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
@@ -35,6 +35,8 @@ def test_distributed_cpu_training(sagemaker_session, chainer_full_version):
     _run_mnist_training_job(sagemaker_session, "ml.c4.xlarge", 2, chainer_full_version)
 
 
+@pytest.mark.skipif(REGION in ['us-west-1', 'eu-west-2', 'ca-central-1'],
+                    reason='No ml.p2.xlarge supported in these regions')
 def test_distributed_gpu_training(sagemaker_session, chainer_full_version):
     _run_mnist_training_job(sagemaker_session, "ml.p2.xlarge", 2, chainer_full_version)
 

tests/integ/test_pytorch_train.py

@@ -18,7 +18,7 @@
 from sagemaker.pytorch.estimator import PyTorch
 from sagemaker.pytorch.model import PyTorchModel
 from sagemaker.utils import sagemaker_timestamp
-from tests.integ import DATA_DIR, PYTHON_VERSION, TRAINING_DEFAULT_TIMEOUT_MINUTES
+from tests.integ import DATA_DIR, PYTHON_VERSION, TRAINING_DEFAULT_TIMEOUT_MINUTES, REGION
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 MNIST_DIR = os.path.join(DATA_DIR, 'pytorch_mnist')
@@ -69,6 +69,8 @@ def test_deploy_model(pytorch_training_job, sagemaker_session):
         assert output.shape == (batch_size, 10)
 
 
+@pytest.mark.skipif(REGION in ['us-west-1', 'eu-west-2', 'ca-central-1'],
+                    reason='No ml.p2.xlarge supported in these regions')
 def test_async_fit_deploy(sagemaker_session, pytorch_full_version):
     training_job_name = ""
     # TODO: add tests against local mode when it's ready to be used

tests/integ/test_tf_cifar.py

@@ -19,7 +19,7 @@
 import pytest
 
 from sagemaker.tensorflow import TensorFlow
-from tests.integ import DATA_DIR, PYTHON_VERSION
+from tests.integ import DATA_DIR, PYTHON_VERSION, REGION
 from tests.integ.timeout import timeout_and_delete_endpoint_by_name, timeout
 
 PICKLE_CONTENT_TYPE = 'application/python-pickle'
@@ -35,6 +35,8 @@ def __call__(self, data):
 
 @pytest.mark.continuous_testing
 @pytest.mark.skipif(PYTHON_VERSION != 'py2', reason="TensorFlow image supports only python 2.")
+@pytest.mark.skipif(REGION in ['us-west-1', 'eu-west-2', 'ca-central-1'],
+                    reason='No ml.p2.xlarge supported in these regions')
 def test_cifar(sagemaker_session, tf_full_version):
     with timeout(minutes=45):
         script_path = os.path.join(DATA_DIR, 'cifar_10', 'source')

tests/integ/test_tf_keras.py

@@ -18,12 +18,14 @@
 import pytest
 
 from sagemaker.tensorflow import TensorFlow
-from tests.integ import DATA_DIR, PYTHON_VERSION
+from tests.integ import DATA_DIR, PYTHON_VERSION, REGION
 from tests.integ.timeout import timeout_and_delete_endpoint_by_name, timeout
 
 
 @pytest.mark.continuous_testing
 @pytest.mark.skipif(PYTHON_VERSION != 'py2', reason="TensorFlow image supports only python 2.")
+@pytest.mark.skipif(REGION in ['us-west-1', 'eu-west-2', 'ca-central-1'],
+                    reason='No ml.p2.xlarge supported in these regions')
 def test_keras(sagemaker_session, tf_full_version):
     script_path = os.path.join(DATA_DIR, 'cifar_10', 'source')
     dataset_path = os.path.join(DATA_DIR, 'cifar_10', 'data')