chore: integration test for gated jumpstart training model

Author: evakravi

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,6 +13,8 @@
 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
@@ -63,6 +65,46 @@ def test_jumpstart_estimator(setup):
     assert response is not None
 
 
+# instance capacity errors require retries
+@pytest.mark.flaky(reruns=5, reruns_delay=60)
+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"