fix: vllm and vllm-omni tasks in routing logic (#6007)

Signed-off-by: DWarez <dario.salvati@huggingface.co>

Author: dwarez

sagemaker-serve/src/sagemaker/serve/constants.py

@@ -71,14 +71,18 @@
     ModelServer.LLAMACPP,
 }
 
-# HuggingFace pipeline tags that should route to the vLLM-omni (multimodal) server.
-# Note: visual-question-answering is folded into image-text-to-text as of transformers v5.
-OMNI_TASKS: Set[str] = {
+# HuggingFace pipeline tags for models that generate text and should prefer vLLM.
+VLLM_TASKS: Set[str] = {
+    "text-generation",
     "image-text-to-text",
-    "any-to-any",
     "audio-text-to-text",
 }
 
+# HuggingFace pipeline tags for models that can generate multiple output modalities.
+OMNI_TASKS: Set[str] = {
+    "any-to-any",
+}
+
 # ========================================
 # Framework Enum
 # ========================================
@@ -130,4 +134,3 @@ class Framework(Enum):
     Framework.NTM: (RecordSerializer(), JSONDeserializer()),
     Framework.SMD: (JSONSerializer(), JSONDeserializer()),
 }
-

sagemaker-serve/src/sagemaker/serve/model_builder.py

@@ -143,6 +143,7 @@
     LOCAL_MODES,
     SUPPORTED_MODEL_SERVERS,
     OMNI_TASKS,
+    VLLM_TASKS,
     Framework,
 )
 from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline
@@ -3608,7 +3609,7 @@ def _build_single_modelbuilder(
                     # Task-based auto-selection. SGLang is not auto-selected by task; it is
                     # opt-in only via model_server=ModelServer.SGLANG, which is handled earlier
                     # by the _build_for_model_server() short-circuit above.
-                    if model_task == "text-generation":
+                    if model_task in VLLM_TASKS:
                         self.built_model = self._build_for_vllm()
                         return self.built_model
                     elif model_task in OMNI_TASKS:

sagemaker-serve/src/sagemaker/serve/model_builder_utils.py

@@ -63,6 +63,7 @@ def build(self):
 from sagemaker.serve.constants import (
     DEFAULT_SERIALIZERS_BY_FRAMEWORK,
     OMNI_TASKS,
+    VLLM_TASKS,
     Framework,
 )
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
@@ -680,7 +681,7 @@ def _detect_huggingface_image(self) -> None:
                 )
                 model_task = hf_model_md.get("pipeline_tag")
 
-                if model_task == "text-generation":
+                if model_task in VLLM_TASKS:
                     effective_model_server = ModelServer.VLLM
                 elif model_task in OMNI_TASKS:
                     effective_model_server = ModelServer.VLLM_OMNI

sagemaker-serve/tests/unit/test_constants.py

@@ -2,8 +2,10 @@
 from sagemaker.serve.constants import (
     Framework,
     LOCAL_MODES,
+    OMNI_TASKS,
     SUPPORTED_MODEL_SERVERS,
-    DEFAULT_SERIALIZERS_BY_FRAMEWORK
+    DEFAULT_SERIALIZERS_BY_FRAMEWORK,
+    VLLM_TASKS,
 )
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.types import ModelServer
@@ -52,6 +54,14 @@ def test_supported_model_servers(self):
         self.assertEqual(SUPPORTED_MODEL_SERVERS, expected_servers)
 
 
+class TestHuggingFaceRoutingTasks(unittest.TestCase):
+    def test_text_output_multimodal_tasks_route_to_vllm(self):
+        self.assertIn("image-text-to-text", VLLM_TASKS)
+        self.assertIn("audio-text-to-text", VLLM_TASKS)
+        self.assertNotIn("image-text-to-text", OMNI_TASKS)
+        self.assertNotIn("audio-text-to-text", OMNI_TASKS)
+
+
 class TestDefaultSerializersByFramework(unittest.TestCase):
     def test_all_frameworks_have_serializers(self):
         for framework in Framework:

sagemaker-serve/tests/unit/test_model_builder_utils_extended_coverage.py

@@ -171,6 +171,24 @@ def test_detect_hf_image_vllm(self, mock_metadata, mock_retrieve):
         self.assertEqual(utils.framework, Framework.HUGGINGFACE)
         self.assertEqual(mock_retrieve.call_args.args[0], "huggingface-vllm")
 
+    @patch('sagemaker.core.image_uris.retrieve')
+    @patch.object(_ModelBuilderUtils, 'get_huggingface_model_metadata')
+    def test_detect_hf_image_multimodal_text_routes_to_vllm(self, mock_metadata, mock_retrieve):
+        """Test image-text-to-text models resolve the vLLM framework."""
+        utils = _ModelBuilderUtils()
+        utils.model = "google/gemma-4-E2B-it"
+        utils.region = "us-west-2"
+        mock_metadata.return_value = {"pipeline_tag": "image-text-to-text"}
+        mock_retrieve.return_value = (
+            "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-vllm:latest"
+        )
+
+        utils._detect_huggingface_image()
+
+        self.assertIsNotNone(utils.image_uri)
+        self.assertEqual(utils.framework, Framework.HUGGINGFACE)
+        self.assertEqual(mock_retrieve.call_args.args[0], "huggingface-vllm")
+
     @patch('sagemaker.core.image_uris.retrieve')
     @patch.object(_ModelBuilderUtils, 'get_huggingface_model_metadata')
     def test_detect_hf_image_sglang(self, mock_metadata, mock_retrieve):

sagemaker-serve/tests/unit/test_model_builder_workflows.py

@@ -415,19 +415,28 @@ def test_build_single_with_hf_text_generation(self, mock_schema_init, mock_use_j
 
     @patch('sagemaker.serve.model_builder.ModelBuilder._is_huggingface_model')
     @patch('sagemaker.serve.model_builder.ModelBuilder.get_huggingface_model_metadata')
-    @patch('sagemaker.serve.model_builder.ModelBuilder._build_for_vllm_omni')
+    @patch('sagemaker.serve.model_builder.ModelBuilder._build_for_vllm')
     @patch('sagemaker.serve.model_builder.ModelBuilder._get_client_translators')
     @patch('sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id')
     @patch('sagemaker.serve.model_builder.ModelBuilder._use_jumpstart_equivalent')
     @patch('sagemaker.serve.model_builder.ModelBuilder._hf_schema_builder_init')
-    def test_build_single_with_hf_multimodal(self, mock_schema_init, mock_use_js, mock_is_js, mock_get_trans, mock_build_omni, mock_get_md, mock_is_hf):
-        """Test _build_single_modelbuilder routes HF multimodal models to vLLM-omni."""
+    def test_build_single_with_hf_multimodal(
+        self,
+        mock_schema_init,
+        mock_use_js,
+        mock_is_js,
+        mock_get_trans,
+        mock_build_vllm,
+        mock_get_md,
+        mock_is_hf,
+    ):
+        """Test _build_single_modelbuilder routes HF multimodal text models to vLLM."""
         mock_is_hf.return_value = True
         mock_is_js.return_value = False
         mock_use_js.return_value = False
         mock_get_md.return_value = {"pipeline_tag": "image-text-to-text"}
         mock_model = Mock(spec=Model)
-        mock_build_omni.return_value = mock_model
+        mock_build_vllm.return_value = mock_model
         mock_get_trans.return_value = (None, None)
 
         builder = ModelBuilder(
@@ -441,7 +450,7 @@ def test_build_single_with_hf_multimodal(self, mock_schema_init, mock_use_js, mo
         result = builder._build_single_modelbuilder()
 
         self.assertEqual(result, mock_model)
-        mock_build_omni.assert_called_once()
+        mock_build_vllm.assert_called_once()
 
     @patch('sagemaker.serve.model_builder.ModelBuilder._is_huggingface_model')
     @patch('sagemaker.serve.model_builder.ModelBuilder._build_for_sglang')