fix: integs for training compiler in non-PDX regions

Author: mufaddal-rohawala

tests/integ/test_training_compiler.py

@@ -20,23 +20,25 @@
 from tests import integ
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout
-
-
-@pytest.fixture(scope="module")
-def gpu_instance_type(request):
-    return "ml.p3.2xlarge"
+from tests.integ.utils import gpu_list, retry_with_instance_list
 
 
 @pytest.mark.release
 @pytest.mark.skipif(
     integ.test_region() not in integ.TRAINING_COMPILER_SUPPORTED_REGIONS,
     reason="SageMaker Training Compiler is not supported in this region",
 )
+@pytest.mark.skipif(
+    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",
+)
+@retry_with_instance_list(gpu_list(integ.test_region()))
 def test_huggingface_pytorch(
     sagemaker_session,
-    gpu_instance_type,
     huggingface_training_compiler_latest_version,
     huggingface_training_compiler_pytorch_latest_version,
+    **kwargs,
 ):
     with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, "huggingface")
@@ -48,7 +50,7 @@ def test_huggingface_pytorch(
             transformers_version=huggingface_training_compiler_latest_version,
             pytorch_version=huggingface_training_compiler_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",
@@ -78,11 +80,17 @@ def test_huggingface_pytorch(
     integ.test_region() not in integ.TRAINING_COMPILER_SUPPORTED_REGIONS,
     reason="SageMaker Training Compiler is not supported in this region",
 )
+@pytest.mark.skipif(
+    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",
+)
+@retry_with_instance_list(gpu_list(integ.test_region()))
 def test_huggingface_tensorflow(
     sagemaker_session,
-    gpu_instance_type,
     huggingface_training_compiler_latest_version,
     huggingface_training_compiler_tensorflow_latest_version,
+    **kwargs,
 ):
     with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, "huggingface")
@@ -94,7 +102,7 @@ def test_huggingface_tensorflow(
             transformers_version=huggingface_training_compiler_latest_version,
             tensorflow_version=huggingface_training_compiler_tensorflow_latest_version,
             instance_count=1,
-            instance_type=gpu_instance_type,
+            instance_type=kwargs["instance_type"],
             hyperparameters={
                 "model_name_or_path": "distilbert-base-cased",
                 "per_device_train_batch_size": 128,