fix: handle unrecognized JumpStart container images in ModelBuilder

Collapse endpoint mode routing into a single code path since all
JumpStart builders (DJL, TGI, MMS) perform identical operations
for endpoint deployment. Unknown container images (e.g. vLLM) now
pass through instead of raising ValueError.

For local container mode, unrecognized images raise ValueError with
a message directing users to use SAGEMAKER_ENDPOINT mode.

Author: lnealrei

sagemaker-serve/src/sagemaker/serve/model_builder_servers.py

@@ -930,25 +930,28 @@ def _build_for_jumpstart(self) -> Model:
                     )
                 return self._build_for_mms_jumpstart(init_kwargs)
             else:
-                raise ValueError(f"Unsupported JumpStart image URI: {self.image_uri}")
+                raise ValueError(
+                    f"Local container mode is not yet supported for JumpStart image: {self.image_uri}. "
+                    f"Use Mode.SAGEMAKER_ENDPOINT for deployment."
+                )
 
         else:
-            # SAGEMAKER_ENDPOINT mode - prepare artifacts if needed
+            # SAGEMAKER_ENDPOINT mode — all JumpStart containers follow the same
+            # pattern: JumpStart provides the full image URI, env vars, and S3
+            # model artifacts. No framework-specific prep is needed.
             if not self._optimizing:
                 self._prepare_for_mode()
-            
-            # Route to appropriate model server based on image URI
+
             if "djl-inference" in self.image_uri:
                 self.model_server = ModelServer.DJL_SERVING
-                return self._build_for_djl_jumpstart(init_kwargs)
             elif "tgi-inference" in self.image_uri:
                 self.model_server = ModelServer.TGI
-                return self._build_for_tgi_jumpstart(init_kwargs)
             elif "huggingface-pytorch-inference" in self.image_uri:
                 self.model_server = ModelServer.MMS
-                return self._build_for_mms_jumpstart(init_kwargs)
-            else:
-                raise ValueError(f"Unsupported JumpStart image URI: {self.image_uri}")
+
+            if not hasattr(self, "s3_upload_path") or not self.s3_upload_path:
+                self.s3_upload_path = init_kwargs.model_data
+            return self._create_model()
 
 
 

sagemaker-serve/tests/integ/test_jumpstart_vllm_integration.py

@@ -0,0 +1,58 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import uuid
+import pytest
+import logging
+
+from sagemaker.serve.model_builder import ModelBuilder
+from sagemaker.core.jumpstart.configs import JumpStartConfig
+from sagemaker.train.configs import Compute
+
+logger = logging.getLogger(__name__)
+
+MODEL_ID = "huggingface-vlm-qwen3-5-4b"
+INSTANCE_TYPE = "ml.g6.12xlarge"
+MODEL_NAME_PREFIX = "js-vllm-test-model"
+
+
+@pytest.mark.slow_test
+def test_jumpstart_vllm_build():
+    """Integration test for JumpStart model using vLLM container image.
+
+    Validates that ModelBuilder correctly handles vLLM container images
+    which do not match the legacy djl-inference/tgi-inference/huggingface-pytorch-inference
+    patterns in the endpoint mode routing.
+    """
+    logger.info("Starting JumpStart vLLM build test...")
+
+    compute = Compute(instance_type=INSTANCE_TYPE)
+    jumpstart_config = JumpStartConfig(model_id=MODEL_ID)
+    model_builder = ModelBuilder.from_jumpstart_config(
+        jumpstart_config=jumpstart_config, compute=compute
+    )
+    unique_id = str(uuid.uuid4())[:8]
+
+    core_model = model_builder.build(model_name=f"{MODEL_NAME_PREFIX}-{unique_id}")
+    logger.info(f"Model Successfully Created: {core_model.model_name}")
+
+    try:
+        assert "vllm" in model_builder.image_uri, (
+            f"Expected vLLM image URI, got: {model_builder.image_uri}"
+        )
+        assert model_builder.s3_upload_path is not None, "s3_upload_path should be set"
+        logger.info("JumpStart vLLM build test completed successfully")
+    finally:
+        core_model.delete()
+        logger.info("Model Successfully Deleted!")

sagemaker-serve/tests/unit/test_model_builder_servers_coverage.py

@@ -319,7 +319,7 @@ def setUp(self):
     @patch('sagemaker.core.jumpstart.factory.utils.get_init_kwargs')
     @patch('sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode')
     def test_build_for_jumpstart_unsupported_image_raises(self, mock_prepare, mock_get_kwargs):
-        """Test that unsupported JumpStart image raises error."""
+        """Test that unsupported JumpStart image in LOCAL_CONTAINER mode raises error."""
         mock_init_kwargs = Mock()
         mock_init_kwargs.image_uri = "unsupported-image:latest"
         mock_init_kwargs.env = {}
@@ -329,27 +329,58 @@ def test_build_for_jumpstart_unsupported_image_raises(self, mock_prepare, mock_g
             model="some-model-id",
             role_arn=MOCK_ROLE_ARN,
             sagemaker_session=self.mock_session,
-            mode=Mode.SAGEMAKER_ENDPOINT
+            mode=Mode.LOCAL_CONTAINER
         )
         builder._optimizing = False
 
         with self.assertRaises(ValueError) as context:
             builder._build_for_jumpstart()
 
-        self.assertIn("Unsupported JumpStart image URI", str(context.exception))
+        self.assertIn("Local container mode is not yet supported", str(context.exception))
+        self.assertIn("Mode.SAGEMAKER_ENDPOINT", str(context.exception))
+
+    @patch('sagemaker.core.jumpstart.factory.utils.get_init_kwargs')
+    @patch('sagemaker.serve.model_builder.ModelBuilder._create_model')
+    @patch('sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode')
+    def test_build_for_jumpstart_unknown_image_endpoint_uses_default(self, mock_prepare, mock_create, mock_get_kwargs):
+        """Test that unknown image in endpoint mode passes through without error."""
+        mock_init_kwargs = Mock()
+        mock_init_kwargs.image_uri = "763104351884.dkr.ecr.us-east-1.amazonaws.com/vllm:0.17-gpu-py312-cu129-ubuntu22.04-sagemaker-v1"
+        mock_init_kwargs.env = {}
+        mock_init_kwargs.model_data = "s3://jumpstart-cache/models/vllm/model.tar.gz"
+        mock_get_kwargs.return_value = mock_init_kwargs
+
+        mock_model = Mock(spec=Model)
+        mock_create.return_value = mock_model
+
+        builder = ModelBuilder(
+            model="huggingface-vlm-qwen3-5-4b",
+            role_arn=MOCK_ROLE_ARN,
+            sagemaker_session=self.mock_session,
+            mode=Mode.SAGEMAKER_ENDPOINT
+        )
+        builder._optimizing = False
+
+        result = builder._build_for_jumpstart()
+
+        self.assertEqual(result, mock_model)
+        self.assertIsNone(builder.model_server)
+        self.assertEqual(builder.s3_upload_path, "s3://jumpstart-cache/models/vllm/model.tar.gz")
+        mock_create.assert_called_once()
 
     @patch('sagemaker.core.jumpstart.factory.utils.get_init_kwargs')
-    @patch('sagemaker.serve.model_builder.ModelBuilder._build_for_djl_jumpstart')
+    @patch('sagemaker.serve.model_builder.ModelBuilder._create_model')
     @patch('sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode')
-    def test_build_for_jumpstart_routes_to_djl(self, mock_prepare, mock_build_djl, mock_get_kwargs):
+    def test_build_for_jumpstart_routes_to_djl(self, mock_prepare, mock_create, mock_get_kwargs):
         """Test JumpStart routing to DJL builder."""
         mock_init_kwargs = Mock()
         mock_init_kwargs.image_uri = "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.21.0-deepspeed0.8.0-cu117"
         mock_init_kwargs.env = {}
+        mock_init_kwargs.model_data = "s3://jumpstart-cache/models/djl/model.tar.gz"
         mock_get_kwargs.return_value = mock_init_kwargs
 
         mock_model = Mock(spec=Model)
-        mock_build_djl.return_value = mock_model
+        mock_create.return_value = mock_model
 
         builder = ModelBuilder(
             model="huggingface-llm-falcon-7b",
@@ -363,20 +394,21 @@ def test_build_for_jumpstart_routes_to_djl(self, mock_prepare, mock_build_djl, m
 
         self.assertEqual(result, mock_model)
         self.assertEqual(builder.model_server, ModelServer.DJL_SERVING)
-        mock_build_djl.assert_called_once()
+        mock_create.assert_called_once()
 
     @patch('sagemaker.core.jumpstart.factory.utils.get_init_kwargs')
-    @patch('sagemaker.serve.model_builder.ModelBuilder._build_for_djl_jumpstart')
+    @patch('sagemaker.serve.model_builder.ModelBuilder._create_model')
     @patch('sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode')
-    def test_build_for_jumpstart_passes_config_name(self, mock_prepare, mock_build_djl, mock_get_kwargs):
+    def test_build_for_jumpstart_passes_config_name(self, mock_prepare, mock_create, mock_get_kwargs):
         """Test that config_name is forwarded to get_init_kwargs."""
         mock_init_kwargs = Mock()
         mock_init_kwargs.image_uri = "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.21.0-deepspeed0.8.0-cu117"
         mock_init_kwargs.env = {}
+        mock_init_kwargs.model_data = "s3://jumpstart-cache/models/djl/model.tar.gz"
         mock_get_kwargs.return_value = mock_init_kwargs
 
         mock_model = Mock(spec=Model)
-        mock_build_djl.return_value = mock_model
+        mock_create.return_value = mock_model
 
         builder = ModelBuilder(
             model="meta-textgeneration-llama-3-3-70b-instruct",
@@ -394,17 +426,18 @@ def test_build_for_jumpstart_passes_config_name(self, mock_prepare, mock_build_d
         self.assertEqual(call_kwargs.kwargs.get("config_name") or call_kwargs[1].get("config_name"), "lmi-optimized")
 
     @patch('sagemaker.core.jumpstart.factory.utils.get_init_kwargs')
-    @patch('sagemaker.serve.model_builder.ModelBuilder._build_for_tgi_jumpstart')
+    @patch('sagemaker.serve.model_builder.ModelBuilder._create_model')
     @patch('sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode')
-    def test_build_for_jumpstart_routes_to_tgi(self, mock_prepare, mock_build_tgi, mock_get_kwargs):
+    def test_build_for_jumpstart_routes_to_tgi(self, mock_prepare, mock_create, mock_get_kwargs):
         """Test JumpStart routing to TGI builder."""
         mock_init_kwargs = Mock()
         mock_init_kwargs.image_uri = "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.0.1-tgi0.9.3-gpu-py39-cu118-ubuntu20.04"
         mock_init_kwargs.env = {}
+        mock_init_kwargs.model_data = "s3://jumpstart-cache/models/tgi/model.tar.gz"
         mock_get_kwargs.return_value = mock_init_kwargs
 
         mock_model = Mock(spec=Model)
-        mock_build_tgi.return_value = mock_model
+        mock_create.return_value = mock_model
 
         builder = ModelBuilder(
             model="huggingface-llm-mistral-7b",
@@ -418,20 +451,21 @@ def test_build_for_jumpstart_routes_to_tgi(self, mock_prepare, mock_build_tgi, m
 
         self.assertEqual(result, mock_model)
         self.assertEqual(builder.model_server, ModelServer.TGI)
-        mock_build_tgi.assert_called_once()
+        mock_create.assert_called_once()
 
     @patch('sagemaker.core.jumpstart.factory.utils.get_init_kwargs')
-    @patch('sagemaker.serve.model_builder.ModelBuilder._build_for_mms_jumpstart')
+    @patch('sagemaker.serve.model_builder.ModelBuilder._create_model')
     @patch('sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode')
-    def test_build_for_jumpstart_routes_to_mms(self, mock_prepare, mock_build_mms, mock_get_kwargs):
+    def test_build_for_jumpstart_routes_to_mms(self, mock_prepare, mock_create, mock_get_kwargs):
         """Test JumpStart routing to MMS builder."""
         mock_init_kwargs = Mock()
         mock_init_kwargs.image_uri = "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04"
         mock_init_kwargs.env = {}
+        mock_init_kwargs.model_data = "s3://jumpstart-cache/models/mms/model.tar.gz"
         mock_get_kwargs.return_value = mock_init_kwargs
 
         mock_model = Mock(spec=Model)
-        mock_build_mms.return_value = mock_model
+        mock_create.return_value = mock_model
 
         builder = ModelBuilder(
             model="pytorch-ic-mobilenet-v2",
@@ -445,7 +479,7 @@ def test_build_for_jumpstart_routes_to_mms(self, mock_prepare, mock_build_mms, m
 
         self.assertEqual(result, mock_model)
         self.assertEqual(builder.model_server, ModelServer.MMS)
-        mock_build_mms.assert_called_once()
+        mock_create.assert_called_once()
 
 
 class TestDeployWrappers(unittest.TestCase):