fix: integ test only in supported regions

Author: evakravi

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

@@ -18,6 +18,7 @@
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 
 from sagemaker.jumpstart.estimator import JumpStartEstimator
+import tests
 from tests.integ.sagemaker.jumpstart.constants import (
     ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
     JUMPSTART_TAG,
@@ -32,6 +33,15 @@
 
 MAX_INIT_TIME_SECONDS = 5
 
+GATED_TRAINING_MODEL_SUPPORTED_REGIONS = {
+    "us-west-2",
+    "us-east-1",
+    "eu-west-1",
+    "ap-southeast-1",
+    "us-east-2",
+    "ap-southeast-2",
+}
+
 
 def test_jumpstart_estimator(setup):
 
@@ -67,6 +77,10 @@ def test_jumpstart_estimator(setup):
 
 # instance capacity errors require retries
 @pytest.mark.flaky(reruns=5, reruns_delay=60)
+@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()}.",
+)
 def test_gated_model_training(setup):
 
     model_id, model_version = "meta-textgeneration-llama-2-7b", "*"

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -33,7 +33,14 @@
 
 MAX_INIT_TIME_SECONDS = 5
 
-MODEL_PACKAGE_ARN_SUPPORTED_REGIONS = {"us-west-2", "us-east-1"}
+GATED_INFERENCE_MODEL_SUPPORTED_REGIONS = {
+    "us-west-2",
+    "us-east-1",
+    "eu-west-1",
+    "ap-southeast-1",
+    "us-east-2",
+    "ap-southeast-2",
+}
 
 
 def test_non_prepacked_jumpstart_model(setup):
@@ -80,8 +87,8 @@ def test_prepacked_jumpstart_model(setup):
 
 
 @pytest.mark.skipif(
-    tests.integ.test_region() not in MODEL_PACKAGE_ARN_SUPPORTED_REGIONS,
-    reason=f"JumpStart Model Package models unavailable in {tests.integ.test_region()}.",
+    tests.integ.test_region() not in GATED_INFERENCE_MODEL_SUPPORTED_REGIONS,
+    reason=f"JumpStart gated inference models unavailable in {tests.integ.test_region()}.",
 )
 def test_model_package_arn_jumpstart_model(setup):