fix: only run DataParallel and EdgePackaging tests in supported regions (#2023)

Co-authored-by: Rui Wang Napieralski <[email protected]>

Author: icywang86rui

tests/integ/__init__.py

@@ -118,6 +118,18 @@
     "eu-south-1",
 ]
 NO_MODEL_MONITORING_REGIONS = ["me-south-1", "af-south-1", "eu-south-1"]
+EDGE_PACKAGING_SUPPORTED_REGIONS = [
+    "us-east-2",
+    "us-west-2",
+    "us-east-1",
+    "eu-west-1",
+    "ap-northeast-1",
+    "eu-central-1",
+]
+# Data parallelism need to be tested with p3.16xlarge.
+# The instance type is expensive and not supported in all the regions.
+# Limiting the test to run in IAD and CMH
+DATA_PARALLEL_TESTING_REGIONS = ["us-east-2", "us-east-1"]
 
 
 EFS_TEST_ENABLED_REGION = []

tests/integ/test_edge.py

@@ -17,7 +17,12 @@
 import pytest
 
 from sagemaker.mxnet.estimator import MXNet
-from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
+from tests.integ import (
+    DATA_DIR,
+    TRAINING_DEFAULT_TIMEOUT_MINUTES,
+    EDGE_PACKAGING_SUPPORTED_REGIONS,
+    test_region,
+)
 from tests.integ.timeout import timeout
 
 
@@ -53,6 +58,10 @@ def mxnet_training_job(
         return mx.latest_training_job.name
 
 
+@pytest.mark.skipif(
+    test_region() not in EDGE_PACKAGING_SUPPORTED_REGIONS,
+    reason="Edge packaging isn't supported in that specific region.",
+)
 def test_edge_packaging_job(mxnet_training_job, sagemaker_session):
     estimator = MXNet.attach(mxnet_training_job, sagemaker_session=sagemaker_session)
     model = estimator.compile_model(

tests/integ/test_smdataparallel_pt.py

@@ -14,6 +14,8 @@
 
 import os
 
+import pytest
+
 import sagemaker.utils
 import tests.integ as integ
 
@@ -26,6 +28,10 @@
 )
 
 
+@pytest.mark.skipif(
+    integ.test_region() not in integ.DATA_PARALLEL_TESTING_REGIONS,
+    reason="Only allow this test to run in IAD and CMH to limit usage of p3.16xlarge",
+)
 def test_smdataparallel_pt_mnist(
     sagemaker_session,
     pytorch_training_latest_version,

tests/integ/test_smdataparallel_tf.py

@@ -14,6 +14,8 @@
 
 import os
 
+import pytest
+
 import sagemaker.utils
 import tests.integ as integ
 
@@ -25,6 +27,10 @@
 )
 
 
+@pytest.mark.skipif(
+    integ.test_region() not in integ.DATA_PARALLEL_TESTING_REGIONS,
+    reason="Only allow this test to run in IAD and CMH to limit usage of p3.16xlarge",
+)
 def test_smdataparallel_tf_mnist(
     sagemaker_session,
     tensorflow_training_latest_version,