fix: xfail gated training test if capacity error

Author: evakravi

tests/integ/sagemaker/jumpstart/estimator/test_jumpstart_estimator.py

@@ -26,6 +26,7 @@
 from tests.integ.sagemaker.jumpstart.utils import (
     get_sm_session,
     get_training_dataset_for_model_and_version,
+    x_fail_if_ice,
 )
 
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
@@ -75,8 +76,7 @@ def test_jumpstart_estimator(setup):
     assert response is not None
 
 
-# instance capacity errors require retries
-@pytest.mark.flaky(reruns=5, reruns_delay=60)
+@x_fail_if_ice
 @pytest.mark.skipif(
     tests.integ.test_region() not in GATED_TRAINING_MODEL_SUPPORTED_REGIONS,
     reason=f"JumpStart gated training models unavailable in {tests.integ.test_region()}.",

tests/integ/sagemaker/jumpstart/utils.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
+import functools
 import json
 
 import uuid
@@ -19,6 +20,7 @@
 import pandas as pd
 import os
 from botocore.config import Config
+import pytest
 
 
 from tests.integ.sagemaker.jumpstart.constants import (
@@ -50,6 +52,19 @@ def get_training_dataset_for_model_and_version(model_id: str, version: str) -> d
     return TRAINING_DATASET_MODEL_DICT[(model_id, version)]
 
 
+def x_fail_if_ice(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            if "CapacityError" in str(e):
+                pytest.xfail(str(e))
+            raise
+
+    return wrapper
+
+
 def download_inference_assets():
 
     if not os.path.exists(TMP_DIRECTORY_PATH):