chore: add jumpstart llama 2 tests

Author: evakravi

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

@@ -14,6 +14,10 @@
 import os
 import time
 
+import pytest
+
+import tests.integ
+
 from sagemaker.jumpstart.model import JumpStartModel
 from tests.integ.sagemaker.jumpstart.constants import (
     ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
@@ -29,6 +33,8 @@
 
 MAX_INIT_TIME_SECONDS = 5
 
+MODEL_PACKAGE_ARN_SUPPORTED_REGIONS = {"us-west-2", "us-east-1"}
+
 
 def test_non_prepacked_jumpstart_model(setup):
 
@@ -73,6 +79,35 @@ def test_prepacked_jumpstart_model(setup):
     assert response is not None
 
 
+@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()}.",
+)
+def test_model_package_arn_jumpstart_model(setup):
+
+    model_id = "meta-textgeneration-llama-2-7b"
+
+    model = JumpStartModel(
+        model_id=model_id,
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+    )
+
+    # uses ml.g5.2xlarge instance
+    predictor = model.deploy(
+        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
+
+
 def test_instatiating_model_not_too_slow(setup):
 
     model_id = "catboost-regression-model"

tests/unit/sagemaker/jumpstart/constants.py

@@ -14,6 +14,63 @@
 
 
 SPECIAL_MODEL_SPECS_DICT = {
+    "js-model-package-arn": {
+        "model_id": "meta-textgeneration-llama-2-7b-f",
+        "url": "https://ai.meta.com/resources/models-and-libraries/llama-downloads/",
+        "version": "1.0.0",
+        "min_sdk_version": "2.173.0",
+        "training_supported": False,
+        "incremental_training_supported": False,
+        "hosting_ecr_specs": {
+            "framework": "pytorch",
+            "framework_version": "1.12.0",
+            "py_version": "py38",
+        },
+        "hosting_artifact_key": "meta-infer/infer-meta-textgeneration-llama-2-7b-f.tar.gz",
+        "hosting_script_key": "source-directory-tarballs/meta/inference/textgeneration/v1.0.0/sourcedir.tar.gz",
+        "hosting_eula_key": "fmhMetadata/eula/llamaEula.txt",
+        "hosting_model_package_arns": {
+            "us-west-2": "arn:aws:sagemaker:us-west-2:594846645681:model-package/"
+            "llama2-7b-f-e46eb8a833643ed58aaccd81498972c3",
+            "us-east-1": "arn:aws:sagemaker:us-east-1:865070037744:model-package/"
+            "llama2-7b-f-e46eb8a833643ed58aaccd81498972c3",
+        },
+        "inference_vulnerable": False,
+        "inference_dependencies": [],
+        "inference_vulnerabilities": [],
+        "training_vulnerable": False,
+        "training_dependencies": [],
+        "training_vulnerabilities": [],
+        "deprecated": False,
+        "inference_environment_variables": [],
+        "metrics": [],
+        "default_inference_instance_type": "ml.g5.2xlarge",
+        "supported_inference_instance_types": [
+            "ml.g5.2xlarge",
+            "ml.g5.4xlarge",
+            "ml.g5.8xlarge",
+            "ml.g5.12xlarge",
+            "ml.g5.24xlarge",
+            "ml.g5.48xlarge",
+            "ml.p4d.24xlarge",
+        ],
+        "model_kwargs": {},
+        "deploy_kwargs": {
+            "model_data_download_timeout": 3600,
+            "container_startup_health_check_timeout": 3600,
+        },
+        "predictor_specs": {
+            "supported_content_types": ["application/json"],
+            "supported_accept_types": ["application/json"],
+            "default_content_type": "application/json",
+            "default_accept_type": "application/json",
+        },
+        "inference_volume_size": 256,
+        "inference_enable_network_isolation": True,
+        "validation_supported": False,
+        "fine_tuning_supported": False,
+        "resource_name_base": "meta-textgeneration-llama-2-7b-f",
+    },
     "js-trainable-model-prepacked": {
         "model_id": "huggingface-text2text-flan-t5-base",
         "url": "https://huggingface.co/google/flan-t5-base",

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -15,6 +15,7 @@
 from typing import Optional, Set
 from unittest import mock
 import unittest
+from mock import MagicMock
 import pytest
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.jumpstart.enums import JumpStartScriptScope
@@ -567,6 +568,79 @@ def test_model_id_not_found_refeshes_cach_inference(
             ]
         )
 
+    @mock.patch("sagemaker.jumpstart.model.is_valid_model_id")
+    @mock.patch("sagemaker.jumpstart.factory.model.Session")
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
+    def test_jumpstart_model_package_arn(
+        self,
+        mock_get_model_specs: mock.Mock,
+        mock_session: mock.Mock,
+        mock_is_valid_model_id: mock.Mock,
+    ):
+
+        mock_is_valid_model_id.return_value = True
+
+        model_id, _ = "js-model-package-arn", "*"
+
+        mock_get_model_specs.side_effect = get_special_model_spec
+
+        mock_session.return_value = MagicMock(sagemaker_config={})
+
+        model = JumpStartModel(model_id=model_id)
+
+        model.deploy()
+
+        self.assertEqual(
+            mock_session.return_value.create_model.call_args[0][2],
+            {
+                "ModelPackageName": "arn:aws:sagemaker:us-west-2:594846645681:model-package"
+                "/llama2-7b-f-e46eb8a833643ed58aaccd81498972c3"
+            },
+        )
+
+    @mock.patch("sagemaker.jumpstart.model.is_valid_model_id")
+    @mock.patch("sagemaker.jumpstart.factory.model.Session")
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
+    def test_jumpstart_model_package_arn_override(
+        self,
+        mock_get_model_specs: mock.Mock,
+        mock_session: mock.Mock,
+        mock_is_valid_model_id: mock.Mock,
+    ):
+
+        mock_is_valid_model_id.return_value = True
+
+        # arbitrary model without model packarn arn
+        model_id, _ = "js-trainable-model", "*"
+
+        mock_get_model_specs.side_effect = get_special_model_spec
+
+        mock_session.return_value = MagicMock(sagemaker_config={})
+
+        model_package_arn = (
+            "arn:aws:sagemaker:us-west-2:867530986753:model-package/"
+            "llama2-ynnej-f-e46eb8a833643ed58aaccd81498972c3"
+        )
+        model = JumpStartModel(model_id=model_id, model_package_arn=model_package_arn)
+
+        model.deploy()
+
+        self.assertEqual(
+            mock_session.return_value.create_model.call_args[0][2],
+            {
+                "ModelPackageName": model_package_arn,
+                "Environment": {
+                    "ENDPOINT_SERVER_TIMEOUT": "3600",
+                    "MODEL_CACHE_ROOT": "/opt/ml/model",
+                    "SAGEMAKER_ENV": "1",
+                    "SAGEMAKER_MODEL_SERVER_WORKERS": "1",
+                    "SAGEMAKER_PROGRAM": "inference.py",
+                },
+            },
+        )
+
 
 def test_jumpstart_model_requires_model_id():
     with pytest.raises(ValueError):