Integration tests

Author: Jonathan Makunga

src/sagemaker/jumpstart/enums.py

@@ -93,8 +93,10 @@ class JumpStartTag(str, Enum):
     MODEL_ID = "sagemaker-sdk:jumpstart-model-id"
     MODEL_VERSION = "sagemaker-sdk:jumpstart-model-version"
     MODEL_TYPE = "sagemaker-sdk:jumpstart-model-type"
+
     INFERENCE_CONFIG_NAME = "sagemaker-sdk:jumpstart-inference-config-name"
     TRAINING_CONFIG_NAME = "sagemaker-sdk:jumpstart-training-config-name"
+
     HUB_CONTENT_ARN = "sagemaker-sdk:hub-content-arn"
 
 

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

@@ -11,7 +11,10 @@
 # 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 io
 import os
+import sys
 import time
 from unittest import mock
 
@@ -349,3 +352,43 @@ def test_register_gated_jumpstart_model(setup):
     predictor.delete_predictor()
 
     assert response is not None
+
+
+def test_jumpstart_model_with_deployment_configs(setup):
+    model_id = "meta-textgeneration-llama-2-7b-f"
+
+    model = JumpStartModel(
+        model_id=model_id,
+        model_version="*",
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+    )
+
+    captured_output = io.StringIO()
+    sys.stdout = captured_output
+    model.display_benchmark_metrics()
+    sys.stdout = sys.__stdout__
+    assert captured_output.getvalue() is not None
+
+    configs = model.list_deployment_configs()
+    assert len(configs) > 0
+
+    model.set_deployment_config(
+        configs[0]["ConfigName"],
+        "ml.g5.2xlarge",
+    )
+    assert model.config_name == configs[0]["ConfigName"]
+
+    predictor = model.deploy(
+        accept_eula=True,
+        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+    )
+
+    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

tests/integ/sagemaker/serve/test_serve_js_happy.py

@@ -12,6 +12,9 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import io
+import sys
+
 import pytest
 
 from sagemaker.serve.builder.model_builder import ModelBuilder
@@ -54,6 +57,19 @@ def happy_model_builder(sagemaker_session):
     )
 
 
+@pytest.fixture
+def meta_textgeneration_llama_2_7b_f_schema():
+    prompt = "Hello, I'm a language model,"
+    response = "Hello, I'm a language model, and I'm here to help you with your English."
+    sample_input = {"inputs": prompt}
+    sample_output = [{"generated_text": response}]
+
+    return SchemaBuilder(
+        sample_input=sample_input,
+        sample_output=sample_output,
+    )
+
+
 @pytest.fixture
 def happy_mms_model_builder(sagemaker_session):
     iam_client = sagemaker_session.boto_session.client("iam")
@@ -125,3 +141,59 @@ def test_happy_mms_sagemaker_endpoint(happy_mms_model_builder, gpu_instance_type
             )
             if caught_ex:
                 raise caught_ex
+
+
+@pytest.mark.skipif(
+    PYTHON_VERSION_IS_NOT_310,
+    reason="The goal of these test are to test the serving components of our feature",
+)
+def test_js_model_with_deployment_configs(
+    meta_textgeneration_llama_2_7b_f_schema,
+    sagemaker_session,
+):
+    logger.info("Running in SAGEMAKER_ENDPOINT mode...")
+    caught_ex = None
+    iam_client = sagemaker_session.boto_session.client("iam")
+    role_arn = iam_client.get_role(RoleName=ROLE_NAME)["Role"]["Arn"]
+
+    model_builder = ModelBuilder(
+        model="meta-textgeneration-llama-2-7b-f",
+        schema_builder=meta_textgeneration_llama_2_7b_f_schema,
+    )
+    configs = model_builder.list_deployment_configs()
+
+    assert len(configs) > 0
+
+    captured_output = io.StringIO()
+    sys.stdout = captured_output
+    model_builder.display_benchmark_metrics()
+    sys.stdout = sys.__stdout__
+    assert captured_output.getvalue() is not None
+
+    model_builder.set_deployment_config(
+        configs[0]["ConfigName"],
+        "ml.g5.2xlarge",
+    )
+    model = model_builder.build(role_arn=role_arn, sagemaker_session=sagemaker_session)
+    assert model.config_name == configs[0]["ConfigName"]
+    assert model_builder.get_deployment_config() is not None
+
+    with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
+        try:
+            logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
+            predictor = model.deploy(accept_eula=True)
+            logger.info("Endpoint successfully deployed.")
+
+            updated_sample_input = model_builder.schema_builder.sample_input
+
+            predictor.predict(updated_sample_input)
+        except Exception as e:
+            caught_ex = e
+        finally:
+            cleanup_model_resources(
+                sagemaker_session=sagemaker_session,
+                model_name=model.name,
+                endpoint_name=model.endpoint_name,
+            )
+            if caught_ex:
+                raise caught_ex