fix: integs fallback from p3 to p2 instance (#3168)

Author: mufaddal-rohawala

tests/conftest.py

@@ -46,6 +46,7 @@
     "ca-central-1",  # it has p3, but not enough
     "eu-central-1",  # it has p3, but not enough
     "eu-north-1",
+    "eu-west-1",  # it has p3, but not enough
     "eu-west-2",  # it has p3, but not enough
     "eu-west-3",
     "eu-south-1",

tests/integ/__init__.py

@@ -74,6 +74,7 @@
     "ca-central-1",  # it has p3, but not enough
     "eu-central-1",  # it has p3, but not enough
     "eu-north-1",
+    "eu-west-1",  # it has p3, but not enough
     "eu-west-2",  # it has p3, but not enough
     "eu-west-3",
     "eu-south-1",

tests/integ/test_horovod.py

@@ -22,6 +22,7 @@
 
 import sagemaker.utils
 import tests.integ as integ
+from tests.integ.utils import gpu_list, retry_with_instance_list
 from sagemaker.tensorflow import TensorFlow
 from tests.integ import timeout
 
@@ -51,18 +52,19 @@ def test_hvd_cpu(
     and integ.test_region() in integ.TRAINING_NO_P3_REGIONS,
     reason="no ml.p2 or ml.p3 instances in this region",
 )
+@retry_with_instance_list(gpu_list(integ.test_region()))
 def test_hvd_gpu(
     sagemaker_session,
     tensorflow_training_latest_version,
     tensorflow_training_latest_py_version,
-    gpu_instance_type,
     tmpdir,
+    **kwargs,
 ):
     _create_and_fit_estimator(
         sagemaker_session,
         tensorflow_training_latest_version,
         tensorflow_training_latest_py_version,
-        gpu_instance_type,
+        kwargs["instance_type"],
         tmpdir,
     )
 

tests/integ/test_horovod_mx.py

@@ -24,6 +24,7 @@
 import tests.integ as integ
 from sagemaker.mxnet import MXNet
 from tests.integ import timeout
+from tests.integ.utils import gpu_list, retry_with_instance_list
 
 horovod_dir = os.path.join(os.path.dirname(__file__), "..", "data", "horovod")
 
@@ -51,18 +52,19 @@ def test_hvd_cpu(
     and integ.test_region() in integ.TRAINING_NO_P3_REGIONS,
     reason="no ml.p2 or ml.p3 instances in this region",
 )
+@retry_with_instance_list(gpu_list(integ.test_region()))
 def test_hvd_gpu(
     mxnet_training_latest_version,
     mxnet_training_latest_py_version,
     sagemaker_session,
-    gpu_instance_type,
     tmpdir,
+    **kwargs,
 ):
     _create_and_fit_estimator(
         mxnet_training_latest_version,
         mxnet_training_latest_py_version,
         sagemaker_session,
-        gpu_instance_type,
+        kwargs["instance_type"],
         tmpdir,
     )
 

tests/integ/test_huggingface.py

@@ -69,12 +69,13 @@ def test_framework_processing_job_with_deps(
     and integ.test_region() in integ.TRAINING_NO_P3_REGIONS,
     reason="no ml.p2 or ml.p3 instances in this region",
 )
+@retry_with_instance_list(gpu_list(integ.test_region()))
 def test_huggingface_training(
     sagemaker_session,
-    gpu_instance_type,
     huggingface_training_latest_version,
     huggingface_training_pytorch_latest_version,
     huggingface_pytorch_latest_training_py_version,
+    **kwargs,
 ):
     with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, "huggingface")
@@ -86,7 +87,7 @@ def test_huggingface_training(
             transformers_version=huggingface_training_latest_version,
             pytorch_version=huggingface_training_pytorch_latest_version,
             instance_count=1,
-            instance_type=gpu_instance_type,
+            instance_type=kwargs["instance_type"],
             hyperparameters={
                 "model_name_or_path": "distilbert-base-cased",
                 "task_name": "wnli",

tests/integ/test_tf.py

@@ -182,21 +182,51 @@ def test_server_side_encryption(sagemaker_session, tf_full_version, tf_full_py_v
 
 
 @pytest.mark.release
-def test_mnist_distributed(
+def test_mnist_distributed_cpu(
     sagemaker_session,
-    instance_type,
+    cpu_instance_type,
     tensorflow_training_latest_version,
     tensorflow_training_latest_py_version,
 ):
+    _create_and_fit_estimator(
+        sagemaker_session,
+        tensorflow_training_latest_version,
+        tensorflow_training_latest_py_version,
+        cpu_instance_type,
+    )
+
+
+@pytest.mark.release
+@pytest.mark.skipif(
+    tests.integ.test_region() in tests.integ.TRAINING_NO_P2_REGIONS
+    and tests.integ.test_region() in tests.integ.TRAINING_NO_P3_REGIONS,
+    reason="no ml.p2 or ml.p3 instances in this region",
+)
+@retry_with_instance_list(gpu_list(tests.integ.test_region()))
+def test_mnist_distributed_gpu(
+    sagemaker_session,
+    tensorflow_training_latest_version,
+    tensorflow_training_latest_py_version,
+    **kwargs,
+):
+    _create_and_fit_estimator(
+        sagemaker_session,
+        tensorflow_training_latest_version,
+        tensorflow_training_latest_py_version,
+        kwargs["instance_type"],
+    )
+
+
+def _create_and_fit_estimator(sagemaker_session, tf_version, py_version, instance_type):
     estimator = TensorFlow(
         entry_point=SCRIPT,
         source_dir=MNIST_RESOURCE_PATH,
         role=ROLE,
         instance_count=2,
         instance_type=instance_type,
         sagemaker_session=sagemaker_session,
-        framework_version=tensorflow_training_latest_version,
-        py_version=tensorflow_training_latest_py_version,
+        framework_version=tf_version,
+        py_version=py_version,
         distribution=PARAMETER_SERVER_DISTRIBUTION,
         disable_profiler=True,
     )