From ee08d99fd8defb6c1ac91f0fb3ae7237aacafbc4 Mon Sep 17 00:00:00 2001 From: DWarez Date: Thu, 9 Jul 2026 13:33:36 +0200 Subject: [PATCH] fix: vllm and vllm-omni tasks in routing logic Signed-off-by: DWarez --- .../src/sagemaker/serve/constants.py | 13 ++++++++----- .../src/sagemaker/serve/model_builder.py | 3 ++- .../sagemaker/serve/model_builder_utils.py | 3 ++- sagemaker-serve/tests/unit/test_constants.py | 12 +++++++++++- ...t_model_builder_utils_extended_coverage.py | 18 ++++++++++++++++++ .../unit/test_model_builder_workflows.py | 19 ++++++++++++++----- 6 files changed, 55 insertions(+), 13 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/constants.py b/sagemaker-serve/src/sagemaker/serve/constants.py index 591f851a71..d8b3f08dba 100644 --- a/sagemaker-serve/src/sagemaker/serve/constants.py +++ b/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()), } - diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index b7be863d9b..30461eb91e 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -139,6 +139,7 @@ LOCAL_MODES, SUPPORTED_MODEL_SERVERS, OMNI_TASKS, + VLLM_TASKS, Framework, ) from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline @@ -2810,7 +2811,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: diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py index e58ea4d7ad..d237a04d0a 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py +++ b/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 diff --git a/sagemaker-serve/tests/unit/test_constants.py b/sagemaker-serve/tests/unit/test_constants.py index 43187a0f75..507cf29fb7 100644 --- a/sagemaker-serve/tests/unit/test_constants.py +++ b/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: diff --git a/sagemaker-serve/tests/unit/test_model_builder_utils_extended_coverage.py b/sagemaker-serve/tests/unit/test_model_builder_utils_extended_coverage.py index 142bd66234..10dc2f228a 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_utils_extended_coverage.py +++ b/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): diff --git a/sagemaker-serve/tests/unit/test_model_builder_workflows.py b/sagemaker-serve/tests/unit/test_model_builder_workflows.py index 695948eb69..31b74d0b58 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_workflows.py +++ b/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')