Add model_metadata field to ModelBuilder

Author: Xiong Zeng

src/sagemaker/serve/builder/model_builder.py

@@ -118,9 +118,8 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
             into a stream. All translations between the server and the client are handled
             automatically with the specified input and output.
         model (Optional[Union[object, str]): Model object (with ``predict`` method to perform
-            inference) or a HuggingFace/JumpStart Model ID (followed by ``:task`` if you need
-            to override the task, e.g. bert-base-uncased:fill-mask). Either ``model`` or
-            ``inference_spec`` is required for the model builder to build the artifact.
+            inference) or a HuggingFace/JumpStart Model ID. Either ``model`` or ``inference_spec``
+            is required for the model builder to build the artifact.
         inference_spec (InferenceSpec): The inference spec file with your customized
             ``invoke`` and ``load`` functions.
         image_uri (Optional[str]): The container image uri (which is derived from a
@@ -140,6 +139,8 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
             to the model server). Possible values for this argument are
             ``TORCHSERVE``, ``MMS``, ``TENSORFLOW_SERVING``, ``DJL_SERVING``,
             ``TRITON``, and``TGI``.
+        model_metadata (Optional[Dict[str, str]): Dictionary used to override the HuggingFace
+            model metadata.
     """
 
     model_path: Optional[str] = field(
@@ -206,7 +207,6 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
             "help": (
                 'Model object with "predict" method to perform inference '
                 "or HuggingFace/JumpStart Model ID"
-                "or if you need to override task, provide input as ModelID:Task"
             )
         },
     )
@@ -237,6 +237,9 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
     model_server: Optional[ModelServer] = field(
         default=None, metadata={"help": "Define the model server to deploy to."}
     )
+    model_metadata: Optional[Dict[str, str]] = field(
+        default=None, metadata={"help": "Define the model metadata to override"}
+    )
 
     def _build_validations(self):
         """Placeholder docstring"""
@@ -613,9 +616,8 @@ def build(
 
         if isinstance(self.model, str):
             model_task = None
-            if ":" in self.model:
-                model_task = self.model.split(":")[1]
-                self.model = self.model.split(":")[0]
+            if self.model_metadata:
+                model_task = self.model_metadata.get("HF_TASK")
             if self._is_jumpstart_model_id():
                 return self._build_for_jumpstart()
             if self._is_djl():  # pylint: disable=R1705

tests/integ/sagemaker/serve/test_schema_builder.py

@@ -115,7 +115,7 @@ def test_model_builder_negative_path(sagemaker_session):
 def test_model_builder_happy_path_with_task_provided(
     model_id, task_provided, sagemaker_session, gpu_instance_type
 ):
-    model_builder = ModelBuilder(model=f"{model_id}:{task_provided}")
+    model_builder = ModelBuilder(model=model_id, model_metadata={"HF_TASK": task_provided})
 
     model = model_builder.build(sagemaker_session=sagemaker_session)
 
@@ -156,7 +156,9 @@ def test_model_builder_happy_path_with_task_provided(
 
 
 def test_model_builder_negative_path_with_invalid_task(sagemaker_session):
-    model_builder = ModelBuilder(model="bert-base-uncased:invalid-task")
+    model_builder = ModelBuilder(
+        model="bert-base-uncased", model_metadata={"HF_TASK": "invalid-task"}
+    )
 
     with pytest.raises(
         TaskNotFoundException,

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

@@ -1114,7 +1114,9 @@ def test_build_happy_path_override_with_task_provided(
 
         mock_image_uris_retrieve.return_value = "https://some-image-uri"
 
-        model_builder = ModelBuilder(model="bert-base-uncased:text-generation")
+        model_builder = ModelBuilder(
+            model="bert-base-uncased", model_metadata={"HF_TASK": "text-generation"}
+        )
         model_builder.build(sagemaker_session=mock_session)
 
         self.assertIsNotNone(model_builder.schema_builder)
@@ -1157,11 +1159,14 @@ def test_build_task_override_with_invalid_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:"]
+        model_ids_with_invalid_task = {
+            "bert-base-uncased": "invalid-task",
+            "bert-large-uncased-whole-word-masking-finetuned-squad": "",
+        }
         for model_id in model_ids_with_invalid_task:
-            model_builder = ModelBuilder(model=model_id)
+            provided_task = model_ids_with_invalid_task[model_id]
+            model_builder = ModelBuilder(model=model_id, model_metadata={"HF_TASK": provided_task})
 
-            provided_task = model_id.split(":")[1]
             self.assertRaisesRegex(
                 TaskNotFoundException,
                 f"Error Message: Schema builder for {provided_task} could not be found.",
@@ -1187,9 +1192,11 @@ def test_build_task_override_with_invalid_model_provided(
         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"]
+        invalid_model_id = ""
+        provided_task = "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)
+        model_builder = ModelBuilder(
+            model=invalid_model_id, model_metadata={"HF_TASK": provided_task}
+        )
+        with self.assertRaises(Exception):
+            model_builder.build(sagemaker_session=mock_session)