chore: add flaky test markers & skip region with low P3 instance capacity (#4870)

* chore: add flaky test markers

* skip p2 instance tests in eu-west-1

* Fix to check p3

Author: benieric

tests/integ/__init__.py

@@ -87,6 +87,7 @@
     "ap-south-1",
     "ap-northeast-2",  # it has p3, but not enough
     "us-east-2",  # it has p3, but not enough
+    "eu-west-1",  # it has p3, but not enough
 ]
 
 # EI is currently only supported in the following regions

tests/integ/test_huggingface.py

@@ -29,6 +29,10 @@
 
 
 @pytest.mark.release
+@pytest.mark.skipif(
+    tests.integ.test_region() in tests.integ.TRAINING_NO_P3_REGIONS,
+    reason="No P3 instances or low capacity in this region",
+)
 def test_framework_processing_job_with_deps(
     sagemaker_session,
     huggingface_training_latest_version,
@@ -59,6 +63,10 @@ def test_framework_processing_job_with_deps(
 
 
 @pytest.mark.release
+@pytest.mark.skipif(
+    tests.integ.test_region() in tests.integ.TRAINING_NO_P3_REGIONS,
+    reason="No P3 instances or low capacity in this region",
+)
 def test_huggingface_training(
     sagemaker_session,
     huggingface_training_latest_version,

tests/integ/test_training_compiler.py

@@ -90,6 +90,10 @@ def skip_if_incompatible(gpu_instance_type, request):
         pytest.param("ml.p3.16xlarge", 2),
     ],
 )
+@pytest.mark.skipif(
+    integ.test_region() in integ.TRAINING_NO_P3_REGIONS,
+    reason="No P3 instances or low capacity in this region",
+)
 def test_huggingface_pytorch(
     sagemaker_session,
     gpu_instance_type,

tests/unit/sagemaker/jumpstart/test_notebook_utils.py

@@ -231,6 +231,7 @@ def test_list_jumpstart_models_simple_case(
 
     @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
     @patch("sagemaker.jumpstart.notebook_utils.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.read_s3_file")
+    @pytest.mark.flaky(reruns=5, reruns_delay=1)
     def test_list_jumpstart_models_script_filter(
         self, patched_read_s3_file: Mock, patched_get_manifest: Mock
     ):
@@ -583,6 +584,7 @@ def vulnerable_training_model_spec(bucket, key, *args, **kwargs):
 
     @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
     @patch("sagemaker.jumpstart.notebook_utils.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.read_s3_file")
+    @pytest.mark.flaky(reruns=5, reruns_delay=1)
     def test_list_jumpstart_models_deprecated_models(
         self,
         patched_read_s3_file: Mock,

tests/unit/sagemaker/workflow/test_transform_step.py

@@ -70,6 +70,7 @@
         custom_step.properties.OutputDataConfig.S3OutputPath,
     ],
 )
+@pytest.mark.flaky(reruns=5, reruns_delay=1)
 def test_transform_step_with_transformer(model_name, data, output_path, pipeline_session):
     transformer = Transformer(
         model_name=model_name,