Add extra unit tests for invalid inputs

Author: Xiong Zeng

src/sagemaker/huggingface/llm_utils.py

@@ -81,7 +81,8 @@ def get_huggingface_model_metadata(model_id: str, hf_hub_token: Optional[str] =
     Returns:
         dict: The model metadata retrieved with the HuggingFace API
     """
-
+    if len(model_id) == 0:
+        raise ValueError("Model ID is empty. Please provide a valid Model ID.")
     hf_model_metadata_url = f"https://huggingface.co/api/models/{model_id}"
     hf_model_metadata_json = None
     try:

src/sagemaker/serve/builder/model_builder.py

@@ -627,9 +627,8 @@ def build(
 
                 if model_task is None:
                     model_task = hf_model_md.get("pipeline_tag")
-                if self.schema_builder is None and model_task:
+                if self.schema_builder is None and model_task is not None:
                     self._schema_builder_init(model_task)
-
                 if model_task == "text-generation":  # pylint: disable=R1705
                     return self._build_for_tgi()
                 else:

tests/unit/sagemaker/serve/builder/test_model_builder.py

@@ -1072,7 +1072,7 @@ def test_build_negative_path_when_schema_builder_not_present(
 
         model_builder = ModelBuilder(model="CompVis/stable-diffusion-v1-4")
 
-        self.assertRaisesRegexp(
+        self.assertRaisesRegex(
             TaskNotFoundException,
             "Error Message: Schema builder for text-to-image could not be found.",
             lambda: model_builder.build(sagemaker_session=mock_session),
@@ -1131,7 +1131,7 @@ def test_build_happy_path_override_with_task_provided(
     @patch("sagemaker.huggingface.llm_utils.json")
     @patch("sagemaker.model_uris.retrieve")
     @patch("sagemaker.serve.builder.model_builder._ServeSettings")
-    def test_build_negative_path_override_with_task_provided(
+    def test_build_task_override_with_invalid_task_provided(
         self,
         mock_serveSettings,
         mock_model_uris_retrieve,
@@ -1157,11 +1157,39 @@ def test_build_negative_path_override_with_task_provided(
         mock_model_urllib.request.Request.side_effect = Mock()
 
         mock_image_uris_retrieve.return_value = "https://some-image-uri"
+        model_ids_with_invalid_task = ["bert-base-uncased:invalid-task", "bert-base-uncased:"]
+        for model_id in model_ids_with_invalid_task:
+            model_builder = ModelBuilder(model=model_id)
+
+            provided_task = model_id.split(":")[1]
+            self.assertRaisesRegex(
+                TaskNotFoundException,
+                f"Error Message: Schema builder for {provided_task} could not be found.",
+                lambda: model_builder.build(sagemaker_session=mock_session),
+            )
 
-        model_builder = ModelBuilder(model="bert-base-uncased:invalid-task")
+    @patch("sagemaker.image_uris.retrieve")
+    @patch("sagemaker.model_uris.retrieve")
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    def test_build_task_override_with_invalid_model_provided(
+        self,
+        mock_serveSettings,
+        mock_model_uris_retrieve,
+        mock_image_uris_retrieve,
+    ):
+        # Setup mocks
 
-        self.assertRaisesRegexp(
-            TaskNotFoundException,
-            "Error Message: Schema builder for invalid-task could not be found.",
-            lambda: model_builder.build(sagemaker_session=mock_session),
-        )
+        mock_setting_object = mock_serveSettings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        # HF Pipeline Tag
+        mock_model_uris_retrieve.side_effect = KeyError
+
+        mock_image_uris_retrieve.return_value = "https://some-image-uri"
+        invalid_model_ids_with_task = [":fill-mask", "bert-base-uncased;fill-mask"]
+
+        for model_id in invalid_model_ids_with_task:
+            model_builder = ModelBuilder(model=model_id)
+            with self.assertRaises(Exception):
+                model_builder.build(sagemaker_session=mock_session)