chore: integration test for gated jumpstart training model

tests/integ/sagemaker/jumpstart/constants.py

@@ -45,6 +45,7 @@ def _to_s3_path(filename: str, s3_prefix: Optional[str]) -> str:
     ("huggingface-spc-bert-base-cased", "1.2.3"): ("training-datasets/QNLI-tiny/"),
     ("huggingface-spc-bert-base-cased", "*"): ("training-datasets/QNLI-tiny/"),
     ("js-trainable-model", "*"): ("training-datasets/QNLI-tiny/"),
+    ("meta-textgeneration-llama-2-7b", "*"): ("training-datasets/sec_amazon/"),
 }
 
 

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

@@ -13,23 +13,36 @@
 from __future__ import absolute_import
 import os
 import time
+
+import pytest
 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,
 )
 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
 
 
 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):
 
@@ -63,6 +76,49 @@ def test_jumpstart_estimator(setup):
     assert response is not None
 
 
+@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()}.",
+)
+def test_gated_model_training(setup):
+
+    model_id, model_version = "meta-textgeneration-llama-2-7b", "*"
+
+    estimator = JumpStartEstimator(
+        model_id=model_id,
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        environment={"accept_eula": "true"},
+        max_run=259200,  # avoid exceeding resource limits
+    )
+
+    # uses ml.g5.12xlarge instance
+    estimator.fit(
+        {
+            "training": f"s3://{get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME)}/"
+            f"{get_training_dataset_for_model_and_version(model_id, model_version)}",
+        }
+    )
+
+    # uses ml.g5.2xlarge instance
+    predictor = estimator.deploy(
+        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+    )
+
+    payload = {
+        "inputs": "some-payload",
+        "parameters": {"max_new_tokens": 256, "top_p": 0.9, "temperature": 0.6},
+    }
+
+    response = predictor.predict(payload, custom_attributes="accept_eula=true")
+
+    assert response is not None
+
+
 def test_instatiating_estimator_not_too_slow(setup):
 
     model_id = "xgboost-classification-model"

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):
 

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):