Fetch HF metadata only when explicit type is not selected (#4407)

Author: samruds

src/sagemaker/serve/builder/model_builder.py

@@ -579,19 +579,19 @@ def build(
 
         self.serve_settings = self._get_serve_setting()
 
-        hf_model_md = get_huggingface_model_metadata(
-            self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
-        )
-
         if isinstance(self.model, str):
             if self._is_jumpstart_model_id():
                 return self._build_for_jumpstart()
-            if self._is_djl():
+            if self._is_djl():  # pylint: disable=R1705
                 return self._build_for_djl()
-            if hf_model_md.get("pipeline_tag") == "text-generation":  # pylint: disable=R1705
-                return self._build_for_tgi()
             else:
-                return self._build_for_transformers()
+                hf_model_md = get_huggingface_model_metadata(
+                    self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
+                )
+                if hf_model_md.get("pipeline_tag") == "text-generation":  # pylint: disable=R1705
+                    return self._build_for_tgi()
+                else:
+                    return self._build_for_transformers()
 
         self._build_validations()
 

tests/integ/sagemaker/serve/test_serve_js_happy.py

@@ -13,7 +13,6 @@
 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 (
@@ -33,7 +32,6 @@
 ]
 JS_MODEL_ID = "huggingface-textgeneration1-gpt-neo-125m-fp16"
 ROLE_NAME = "SageMakerRole"
-MOCK_HF_MODEL_METADATA_JSON = {"mock_key": "mock_value"}
 
 
 @pytest.fixture
@@ -47,23 +45,14 @@ 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(
-    mock_urllib, mock_json, happy_model_builder, gpu_instance_type
-):
+def test_happy_tgi_sagemaker_endpoint(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,7 +19,6 @@
 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
@@ -38,7 +37,6 @@
 logger = logging.getLogger(__name__)
 
 ROLE_NAME = "SageMakerRole"
-MOCK_HF_MODEL_METADATA_JSON = {"mock_key": "mock_value"}
 
 
 @pytest.fixture
@@ -183,8 +181,6 @@ 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",
@@ -194,17 +190,12 @@ def model_builder(request):
 )
 @pytest.mark.slow_test
 def test_happy_pytorch_sagemaker_endpoint(
-    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/unit/sagemaker/serve/builder/test_model_builder.py

@@ -42,7 +42,6 @@
 mock_s3_model_data_url = "sample s3 data url"
 mock_secret_key = "mock_secret_key"
 mock_instance_type = "mock instance type"
-MOCK_HF_MODEL_METADATA_JSON = {"mock_key": "mock_value"}
 
 supported_model_server = {
     ModelServer.TORCHSERVE,
@@ -55,15 +54,7 @@
 
 class TestModelBuilder(unittest.TestCase):
     @patch("sagemaker.serve.builder.model_builder._ServeSettings")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
-    def test_validation_in_progress_mode_not_supported(
-        self, mock_serveSettings, mock_urllib, mock_json
-    ):
-        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
-
+    def test_validation_in_progress_mode_not_supported(self, mock_serveSettings):
         builder = ModelBuilder()
         self.assertRaisesRegex(
             Exception,
@@ -75,15 +66,7 @@ def test_validation_in_progress_mode_not_supported(
         )
 
     @patch("sagemaker.serve.builder.model_builder._ServeSettings")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
-    def test_validation_cannot_set_both_model_and_inference_spec(
-        self, mock_serveSettings, mock_urllib, mock_json
-    ):
-        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
-
+    def test_validation_cannot_set_both_model_and_inference_spec(self, mock_serveSettings):
         builder = ModelBuilder(inference_spec="some value", model=Mock(spec=object))
         self.assertRaisesRegex(
             Exception,
@@ -95,15 +78,7 @@ def test_validation_cannot_set_both_model_and_inference_spec(
         )
 
     @patch("sagemaker.serve.builder.model_builder._ServeSettings")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
-    def test_validation_unsupported_model_server_type(
-        self, mock_serveSettings, mock_urllib, mock_json
-    ):
-        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
-
+    def test_validation_unsupported_model_server_type(self, mock_serveSettings):
         builder = ModelBuilder(model_server="invalid_model_server")
         self.assertRaisesRegex(
             Exception,
@@ -116,15 +91,7 @@ def test_validation_unsupported_model_server_type(
         )
 
     @patch("sagemaker.serve.builder.model_builder._ServeSettings")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
-    def test_validation_model_server_not_set_with_image_uri(
-        self, mock_serveSettings, mock_urllib, mock_json
-    ):
-        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
-
+    def test_validation_model_server_not_set_with_image_uri(self, mock_serveSettings):
         builder = ModelBuilder(image_uri="image_uri")
         self.assertRaisesRegex(
             Exception,
@@ -137,15 +104,9 @@ def test_validation_model_server_not_set_with_image_uri(
         )
 
     @patch("sagemaker.serve.builder.model_builder._ServeSettings")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
     def test_save_model_throw_exception_when_none_of_model_and_inference_spec_is_set(
-        self, mock_serveSettings, mock_urllib, mock_json
+        self, mock_serveSettings
     ):
-        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
-
         builder = ModelBuilder(inference_spec=None, model=None)
         self.assertRaisesRegex(
             Exception,
@@ -165,12 +126,8 @@ def test_save_model_throw_exception_when_none_of_model_and_inference_spec_is_set
     @patch("sagemaker.serve.builder.model_builder.SageMakerEndpointMode")
     @patch("sagemaker.serve.builder.model_builder.Model")
     @patch("os.path.exists")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
     def test_build_happy_path_with_sagemaker_endpoint_mode_and_byoc(
         self,
-        mock_urllib,
-        mock_json,
         mock_path_exists,
         mock_sdk_model,
         mock_sageMakerEndpointMode,
@@ -189,10 +146,6 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_byoc(
             else 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
-
         mock_detect_fw_version.return_value = framework, version
 
         mock_prepare_for_torchserve.side_effect = (
@@ -273,12 +226,8 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_byoc(
     @patch("sagemaker.serve.builder.model_builder.SageMakerEndpointMode")
     @patch("sagemaker.serve.builder.model_builder.Model")
     @patch("os.path.exists")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
     def test_build_happy_path_with_sagemaker_endpoint_mode_and_1p_dlc_as_byoc(
         self,
-        mock_urllib,
-        mock_json,
         mock_path_exists,
         mock_sdk_model,
         mock_sageMakerEndpointMode,
@@ -296,11 +245,6 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_1p_dlc_as_byoc(
             and instance_type == "ml.c5.xlarge"
             else 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
-
         mock_detect_fw_version.return_value = framework, version
 
         mock_prepare_for_torchserve.side_effect = (
@@ -381,12 +325,8 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_1p_dlc_as_byoc(
     @patch("sagemaker.serve.builder.model_builder.SageMakerEndpointMode")
     @patch("sagemaker.serve.builder.model_builder.Model")
     @patch("os.path.exists")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
     def test_build_happy_path_with_sagemaker_endpoint_mode_and_inference_spec(
         self,
-        mock_urllib,
-        mock_json,
         mock_path_exists,
         mock_sdk_model,
         mock_sageMakerEndpointMode,
@@ -402,10 +342,6 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_inference_spec(
             lambda model_path: mock_native_model if model_path == MODEL_PATH else 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
-
         mock_detect_fw_version.return_value = framework, version
 
         mock_detect_container.side_effect = (
@@ -490,12 +426,8 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_inference_spec(
     @patch("sagemaker.serve.builder.model_builder.SageMakerEndpointMode")
     @patch("sagemaker.serve.builder.model_builder.Model")
     @patch("os.path.exists")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
     def test_build_happy_path_with_sagemakerEndpoint_mode_and_model(
         self,
-        mock_urllib,
-        mock_json,
         mock_path_exists,
         mock_sdk_model,
         mock_sageMakerEndpointMode,
@@ -514,10 +446,6 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_model(
             else 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
-
         mock_detect_fw_version.return_value = framework, version
 
         mock_prepare_for_torchserve.side_effect = (
@@ -601,12 +529,8 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_model(
     @patch("sagemaker.serve.builder.model_builder.SageMakerEndpointMode")
     @patch("sagemaker.serve.builder.model_builder.Model")
     @patch("os.path.exists")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
     def test_build_happy_path_with_sagemakerEndpoint_mode_and_xgboost_model(
         self,
-        mock_urllib,
-        mock_json,
         mock_path_exists,
         mock_sdk_model,
         mock_sageMakerEndpointMode,
@@ -626,10 +550,6 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_xgboost_model(
             else 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
-
         mock_detect_fw_version.return_value = "xgboost", version
 
         mock_prepare_for_torchserve.side_effect = (
@@ -714,12 +634,8 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_xgboost_model(
     @patch("sagemaker.serve.builder.model_builder.LocalContainerMode")
     @patch("sagemaker.serve.builder.model_builder.Model")
     @patch("os.path.exists")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
     def test_build_happy_path_with_local_container_mode(
         self,
-        mock_urllib,
-        mock_json,
         mock_path_exists,
         mock_sdk_model,
         mock_localContainerMode,
@@ -734,10 +650,6 @@ def test_build_happy_path_with_local_container_mode(
             lambda model_path: mock_native_model if model_path == MODEL_PATH else 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
-
         mock_detect_container.side_effect = (
             lambda model, region, instance_type: mock_image_uri
             if model == mock_native_model
@@ -816,12 +728,8 @@ def test_build_happy_path_with_local_container_mode(
     @patch("sagemaker.serve.builder.model_builder.LocalContainerMode")
     @patch("sagemaker.serve.builder.model_builder.Model")
     @patch("os.path.exists")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
     def test_build_happy_path_with_localContainer_mode_overwritten_with_sagemaker_mode(
         self,
-        mock_urllib,
-        mock_json,
         mock_path_exists,
         mock_sdk_model,
         mock_localContainerMode,
@@ -838,10 +746,6 @@ def test_build_happy_path_with_localContainer_mode_overwritten_with_sagemaker_mo
             lambda model_path: mock_native_model if model_path == MODEL_PATH else 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
-
         mock_detect_fw_version.return_value = framework, version
 
         mock_detect_container.side_effect = (
@@ -965,12 +869,8 @@ def test_build_happy_path_with_localContainer_mode_overwritten_with_sagemaker_mo
     @patch("sagemaker.serve.builder.model_builder.LocalContainerMode")
     @patch("sagemaker.serve.builder.model_builder.Model")
     @patch("os.path.exists")
-    @patch("sagemaker.huggingface.llm_utils.urllib")
-    @patch("sagemaker.huggingface.llm_utils.json")
     def test_build_happy_path_with_sagemaker_endpoint_mode_overwritten_with_local_container(
         self,
-        mock_urllib,
-        mock_json,
         mock_path_exists,
         mock_sdk_model,
         mock_localContainerMode,
@@ -984,10 +884,6 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_overwritten_with_local_co
         # setup mocks
         mock_detect_fw_version.return_value = framework, version
 
-        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
-
         mock_detect_container.side_effect = (
             lambda model, region, instance_type: mock_image_uri
             if model == mock_fw_model