Fix test builds

Author: EC2 Default User

src/sagemaker/serve/builder/model_builder.py

@@ -62,7 +62,6 @@
     ModelServer.TORCHSERVE,
     ModelServer.TRITON,
     ModelServer.DJL_SERVING,
-    ModelServer.MMS,
 }
 
 

tests/integ/sagemaker/serve/test_serve_js_happy.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import pytest
+from unittest.mock import patch, Mock
 from sagemaker.serve.builder.model_builder import ModelBuilder
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
 from tests.integ.sagemaker.serve.constants import (
@@ -26,13 +27,13 @@
 
 logger = logging.getLogger(__name__)
 
-
 SAMPLE_PROMPT = {"inputs": "Hello, I'm a language model,", "parameters": {}}
 SAMPLE_RESPONSE = [
     {"generated_text": "Hello, I'm a language model, and I'm here to help you with your English."}
 ]
 JS_MODEL_ID = "huggingface-textgeneration1-gpt-neo-125m-fp16"
 ROLE_NAME = "SageMakerRole"
+MOCK_HF_MODEL_METADATA_JSON = {"mock_key": "mock_value"}
 
 
 @pytest.fixture
@@ -46,15 +47,23 @@ def happy_model_builder(sagemaker_session):
     )
 
 
+@patch("sagemaker.huggingface.llm_utils.urllib")
+@patch("sagemaker.huggingface.llm_utils.json")
 @pytest.mark.skipif(
     PYTHON_VERSION_IS_NOT_310,
     reason="The goal of these test are to test the serving components of our feature",
 )
 @pytest.mark.slow_test
-def test_happy_tgi_sagemaker_endpoint(happy_model_builder, gpu_instance_type):
+def test_happy_tgi_sagemaker_endpoint(
+    mock_urllib, mock_json, happy_model_builder, gpu_instance_type
+):
     logger.info("Running in SAGEMAKER_ENDPOINT mode...")
     caught_ex = None
 
+    mock_json.load.return_value = MOCK_HF_MODEL_METADATA_JSON
+    mock_hf_model_metadata_url = Mock()
+    mock_urllib.request.Request.side_effect = mock_hf_model_metadata_url
+
     model = happy_model_builder.build()
 
     with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):

tests/integ/sagemaker/serve/test_serve_pt_happy.py

@@ -19,6 +19,7 @@
 import io
 import numpy as np
 
+from unittest.mock import patch, Mock
 from sagemaker.serve.builder.model_builder import ModelBuilder, Mode
 from sagemaker.serve.builder.schema_builder import SchemaBuilder, CustomPayloadTranslator
 from sagemaker.serve.spec.inference_spec import InferenceSpec
@@ -37,6 +38,7 @@
 logger = logging.getLogger(__name__)
 
 ROLE_NAME = "SageMakerRole"
+MOCK_HF_MODEL_METADATA_JSON = {"mock_key": "mock_value"}
 
 
 @pytest.fixture
@@ -181,6 +183,8 @@ def model_builder(request):
 #                 ), f"{caught_ex} was thrown when running pytorch squeezenet local container test"
 
 
+@patch("sagemaker.huggingface.llm_utils.urllib")
+@patch("sagemaker.huggingface.llm_utils.json")
 @pytest.mark.skipif(
     PYTHON_VERSION_IS_NOT_310,  # or NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE,
     reason="The goal of these test are to test the serving components of our feature",
@@ -190,9 +194,17 @@ def model_builder(request):
 )
 @pytest.mark.slow_test
 def test_happy_pytorch_sagemaker_endpoint(
-    sagemaker_session, model_builder, cpu_instance_type, test_image
+    mock_urllib,
+    mock_json,
+    sagemaker_session,
+    model_builder,
+    cpu_instance_type,
+    test_image,
 ):
     logger.info("Running in SAGEMAKER_ENDPOINT mode...")
+    mock_json.load.return_value = MOCK_HF_MODEL_METADATA_JSON
+    mock_hf_model_metadata_url = Mock()
+    mock_urllib.request.Request.side_effect = mock_hf_model_metadata_url
     caught_ex = None
 
     iam_client = sagemaker_session.boto_session.client("iam")

tests/integ/sagemaker/serve/test_serve_transformers.py

@@ -13,7 +13,7 @@
 from __future__ import absolute_import
 
 import pytest
-
+from unittest.mock import patch, Mock
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
 from sagemaker.serve.builder.model_builder import ModelBuilder, Mode
 
@@ -29,6 +29,8 @@
 
 logger = logging.getLogger(__name__)
 
+MOCK_HF_MODEL_METADATA_JSON = {"mock_key": "mock_value"}
+
 sample_input = {
     "inputs": "The man worked as a [MASK].",
 }
@@ -86,15 +88,20 @@ def model_builder(request):
     return request.getfixturevalue(request.param)
 
 
+@patch("sagemaker.huggingface.llm_utils.urllib")
+@patch("sagemaker.huggingface.llm_utils.json")
 @pytest.mark.skipif(
     PYTHON_VERSION_IS_NOT_310,
     reason="Testing feature",
 )
 @pytest.mark.parametrize("model_builder", ["model_builder_model_schema_builder"], indirect=True)
 def test_pytorch_transformers_sagemaker_endpoint(
-    sagemaker_session, model_builder, gpu_instance_type, input
+    mock_urllib, mock_json, sagemaker_session, model_builder, gpu_instance_type, input
 ):
     logger.info("Running in SAGEMAKER_ENDPOINT mode...")
+    mock_json.load.return_value = MOCK_HF_MODEL_METADATA_JSON
+    mock_hf_model_metadata_url = Mock()
+    mock_urllib.request.Request.side_effect = mock_hf_model_metadata_url
     caught_ex = None
 
     iam_client = sagemaker_session.boto_session.client("iam")