fix: gpu use p3/p2 per avail for region (#2902)

Author: mufaddal-rohawala

tests/conftest.py

@@ -37,6 +37,25 @@
     "me-south-1",
 ]
 
+NO_P3_REGIONS = [
+    "af-south-1",
+    "ap-east-1",
+    "ap-southeast-1",  # it has p3, but not enough
+    "ap-southeast-2",  # it has p3, but not enough
+    "ca-central-1",  # it has p3, but not enough
+    "eu-central-1",  # it has p3, but not enough
+    "eu-north-1",
+    "eu-west-2",  # it has p3, but not enough
+    "eu-west-3",
+    "eu-south-1",
+    "me-south-1",
+    "sa-east-1",
+    "us-west-1",
+    "ap-northeast-1",  # it has p3, but not enough
+    "ap-south-1",
+    "ap-northeast-2",  # it has p3, but not enough
+]
+
 NO_T2_REGIONS = ["eu-north-1", "ap-east-1", "me-south-1"]
 
 FRAMEWORKS_FOR_GENERATED_VERSION_FIXTURES = (
@@ -361,9 +380,13 @@ def cpu_instance_type(sagemaker_session, request):
         return "ml.m4.xlarge"
 
 
-@pytest.fixture(scope="module")
-def gpu_instance_type(request):
-    return "ml.p3.2xlarge"
+@pytest.fixture(scope="session")
+def gpu_instance_type(sagemaker_session, request):
+    region = sagemaker_session.boto_session.region_name
+    if region in NO_P3_REGIONS:
+        return "ml.p2.xlarge"
+    else:
+        return "ml.p3.2xlarge"
 
 
 @pytest.fixture(scope="session")
@@ -405,10 +428,16 @@ def pytest_generate_tests(metafunc):
 
         params = [cpu_instance_type]
         if not (
+            region in tests.integ.HOSTING_NO_P3_REGIONS
+            or region in tests.integ.TRAINING_NO_P3_REGIONS
+        ):
+            params.append("ml.p3.2xlarge")
+        elif not (
             region in tests.integ.HOSTING_NO_P2_REGIONS
             or region in tests.integ.TRAINING_NO_P2_REGIONS
         ):
-            params.append("ml.p3.2xlarge")
+            params.append("ml.p2.xlarge")
+
         metafunc.parametrize("instance_type", params, scope="session")
 
     _generate_all_framework_version_fixtures(metafunc)

tests/integ/__init__.py

@@ -66,6 +66,24 @@
     "sa-east-1",
     "us-west-1",
 ]
+TRAINING_NO_P3_REGIONS = [
+    "af-south-1",
+    "ap-east-1",
+    "ap-southeast-1",  # it has p3, but not enough
+    "ap-southeast-2",  # it has p3, but not enough
+    "ca-central-1",  # it has p3, but not enough
+    "eu-central-1",  # it has p3, but not enough
+    "eu-north-1",
+    "eu-west-2",  # it has p3, but not enough
+    "eu-west-3",
+    "eu-south-1",
+    "me-south-1",
+    "sa-east-1",
+    "us-west-1",
+    "ap-northeast-1",  # it has p3, but not enough
+    "ap-south-1",
+    "ap-northeast-2",  # it has p3, but not enough
+]
 
 # EI is currently only supported in the following regions
 # regions were derived from https://aws.amazon.com/machine-learning/elastic-inference/pricing/

tests/integ/test_horovod.py

@@ -47,19 +47,22 @@ def test_hvd_cpu(
 
 @pytest.mark.release
 @pytest.mark.skipif(
-    integ.test_region() in integ.TRAINING_NO_P2_REGIONS, reason="no ml.p2 instances in this region"
+    integ.test_region() in integ.TRAINING_NO_P2_REGIONS
+    and integ.test_region() in integ.TRAINING_NO_P3_REGIONS,
+    reason="no ml.p2 or ml.p3 instances in this region",
 )
 def test_hvd_gpu(
     sagemaker_session,
     tensorflow_training_latest_version,
     tensorflow_training_latest_py_version,
+    gpu_instance_type,
     tmpdir,
 ):
     _create_and_fit_estimator(
         sagemaker_session,
         tensorflow_training_latest_version,
         tensorflow_training_latest_py_version,
-        "ml.p3.2xlarge",
+        gpu_instance_type,
         tmpdir,
     )
 

tests/integ/test_horovod_mx.py

@@ -47,7 +47,9 @@ def test_hvd_cpu(
 
 @pytest.mark.release
 @pytest.mark.skipif(
-    integ.test_region() in integ.TRAINING_NO_P2_REGIONS, reason="no ml.p2 instances in this region"
+    integ.test_region() in integ.TRAINING_NO_P2_REGIONS
+    and integ.test_region() in integ.TRAINING_NO_P3_REGIONS,
+    reason="no ml.p2 or ml.p3 instances in this region",
 )
 def test_hvd_gpu(
     mxnet_training_latest_version,

tests/integ/test_huggingface.py

@@ -28,8 +28,9 @@
 
 @pytest.mark.release
 @pytest.mark.skipif(
-    integ.test_region() in integ.TRAINING_NO_P2_REGIONS,
-    reason="no ml.p2 instances in this region",
+    integ.test_region() in integ.TRAINING_NO_P2_REGIONS
+    and integ.test_region() in integ.TRAINING_NO_P3_REGIONS,
+    reason="no ml.p2 or ml.p3 instances in this region",
 )
 def test_framework_processing_job_with_deps(
     sagemaker_session,
@@ -63,8 +64,9 @@ def test_framework_processing_job_with_deps(
 
 @pytest.mark.release
 @pytest.mark.skipif(
-    integ.test_region() in integ.TRAINING_NO_P2_REGIONS,
-    reason="no ml.p2 instances in this region",
+    integ.test_region() in integ.TRAINING_NO_P2_REGIONS
+    and integ.test_region() in integ.TRAINING_NO_P3_REGIONS,
+    reason="no ml.p2 or ml.p3 instances in this region",
 )
 def test_huggingface_training(
     sagemaker_session,
@@ -108,7 +110,9 @@ def test_huggingface_training(
 
 @pytest.mark.release
 @pytest.mark.skipif(
-    integ.test_region() in integ.TRAINING_NO_P2_REGIONS, reason="no ml.p2 instances in this region"
+    integ.test_region() in integ.TRAINING_NO_P2_REGIONS
+    and integ.test_region() in integ.TRAINING_NO_P3_REGIONS,
+    reason="no ml.p2 or ml.p3 instances in this region",
 )
 def test_huggingface_training_tf(
     sagemaker_session,