fix(serve): repair HuggingFace -> JumpStart redirect in ModelBuilder (#5958)

Two bugs prevented `ModelBuilder(model="<HF_id_with_JS_mirror>").build()`
from working with any bare HuggingFace id (gpt2, Qwen/Qwen3-0.6B,
microsoft/phi-2, etc.), reproducible end-to-end:

  ModelBuilder.build()
    -> _build_single_modelbuilder              (model_builder.py:2579)
    -> _use_jumpstart_equivalent               (model_builder_utils.py)
    -> _hf_schema_builder_init                 raises TaskNotFoundException
    -> redirect aborts before HF -> JS id swap

1. `sagemaker/serve/schema/task.json` was missing from the published
   wheel. setuptools needs an explicit `[tool.setuptools.package-data]`
   declaration; `include-package-data = true` alone is a no-op without a
   MANIFEST.in or version-controlled file list. Without the JSON in the
   wheel, `retrieve_local_schemas("text-generation")` raised
   FileNotFoundError, the caller wrapped it in TaskNotFoundException, and
   the redirect aborted before swapping to the JumpStart id.

2. `_use_jumpstart_equivalent` called `_build_for_jumpstart()` itself,
   then `_build_single_modelbuilder` called it again on return. The
   second call hit `_prepare_for_mode` with consumed state and raised
   FileNotFoundError. Removed the inner call; the caller already handles
   it.

Also wraps the HF schema bootstrap in a try/except so a missing local
task or flaky remote schema fetch can't abort the redirect — the
JumpStart build path supplies its own schema downstream.

Verified end-to-end: ModelBuilder(model="gpt2").build() and
ModelBuilder(model="Qwen/Qwen3-0.6B").build() both succeed (model
created with TGI / DJL_SERVING respectively). Pure-boto3 control with
the same JumpStart djl-inference image + S3 artifacts was already
accepted by CreateModel/CreateEndpointConfig/CreateEndpoint, so the
AWS contract and JumpStart artifacts were never the issue.

Author: ZealSV

sagemaker-serve/pyproject.toml

@@ -61,6 +61,9 @@ where = ["src"]
 include = ["sagemaker*"]
 namespaces = true
 
+[tool.setuptools.package-data]
+"sagemaker.serve" = ["schema/*.json"]
+
 [tool.setuptools.dynamic]
 version = { file = "VERSION"}
 

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

@@ -942,17 +942,20 @@ def _use_jumpstart_equivalent(self) -> bool:
                 if not model_task:
                     model_task = hf_model_md.get("pipeline_tag")
                 if model_task:
-                    self._hf_schema_builder_init(model_task)
+                    try:
+                        self._hf_schema_builder_init(model_task)
+                    except (TaskNotFoundException, FileNotFoundError, OSError) as e:
+                        logger.warning(
+                            "Could not initialize HF schema builder for task %r "
+                            "(%s: %s); falling back to the JumpStart-supplied schema.",
+                            model_task, type(e).__name__, e,
+                        )
 
             huggingface_model_id = self.model
             jumpstart_model_id = self._jumpstart_mapping[huggingface_model_id]["jumpstart-model-id"]
             self.model = jumpstart_model_id
             merged_date = self._jumpstart_mapping[huggingface_model_id].get("merged-at")
 
-            # Call _build_for_jumpstart if method exists
-            if hasattr(self, "_build_for_jumpstart"):
-                self._build_for_jumpstart()
-
             compare_model_diff_message = (
                 "If you want to identify the differences between the two, "
                 "please use model_uris.retrieve() to retrieve the model "

sagemaker-serve/tests/unit/test_model_builder_utils_additional_gaps.py

@@ -10,6 +10,7 @@
 
 from sagemaker.serve.model_builder_utils import _ModelBuilderUtils
 from sagemaker.serve.constants import Framework
+from sagemaker.serve.utils.exceptions import TaskNotFoundException
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.train import ModelTrainer
 
@@ -189,10 +190,56 @@ def test_auto_detect_image_uri_object_model(self, mock_detect_obj):
 class TestUseJumpStartEquivalent(unittest.TestCase):
     """Test _use_jumpstart_equivalent method."""
 
+    def _make_utils(self):
+        utils = _ModelBuilderUtils()
+        utils.model = "gpt2"
+        utils.image_uri = None
+        utils.env_vars = None
+        utils.schema_builder = None
+        utils.model_metadata = None
+        utils._is_gated_model = Mock(return_value=False)
+        utils._build_for_jumpstart = Mock()
+        return utils
+
+    @patch.object(_ModelBuilderUtils, 'get_huggingface_model_metadata')
+    @patch.object(_ModelBuilderUtils, '_hf_schema_builder_init')
+    @patch.object(_ModelBuilderUtils, '_retrieve_hugging_face_model_mapping')
+    def test_use_jumpstart_equivalent_redirects_hf_to_js(
+        self, mock_retrieve, mock_schema_init, mock_md
+    ):
+        """HF id with a JumpStart mirror is rewritten to its JS id."""
+        utils = self._make_utils()
+        mock_retrieve.return_value = {
+            "gpt2": {"jumpstart-model-id": "huggingface-textgeneration-gpt2", "merged-at": "2024-01-01"}
+        }
+        mock_md.return_value = {"pipeline_tag": "text-generation"}
+
+        result = utils._use_jumpstart_equivalent()
+
+        self.assertTrue(result)
+        self.assertEqual(utils.model, "huggingface-textgeneration-gpt2")
+        utils._build_for_jumpstart.assert_not_called()
+        mock_schema_init.assert_called_once_with("text-generation")
+
+    @patch.object(_ModelBuilderUtils, 'get_huggingface_model_metadata')
+    @patch.object(_ModelBuilderUtils, '_hf_schema_builder_init')
     @patch.object(_ModelBuilderUtils, '_retrieve_hugging_face_model_mapping')
-    def test_use_jumpstart_equivalent_no_image_uri(self, mock_retrieve):
-        """Test using JumpStart equivalent without image_uri - skipped (complex schema builder init)."""
-        pass
+    def test_use_jumpstart_equivalent_swallows_schema_failures(
+        self, mock_retrieve, mock_schema_init, mock_md
+    ):
+        """Schema bootstrap failures are tolerated and the JumpStart id swap still happens."""
+        mock_retrieve.return_value = {
+            "gpt2": {"jumpstart-model-id": "huggingface-textgeneration-gpt2"}
+        }
+        mock_md.return_value = {"pipeline_tag": "text-generation"}
+
+        for exc in (TaskNotFoundException("x"), FileNotFoundError("x"), OSError("x")):
+            with self.subTest(exception=type(exc).__name__):
+                mock_schema_init.side_effect = exc
+                utils = self._make_utils()
+                result = utils._use_jumpstart_equivalent()
+                self.assertTrue(result)
+                self.assertEqual(utils.model, "huggingface-textgeneration-gpt2")
 
     def test_use_jumpstart_equivalent_with_image_uri(self):
         """Test using JumpStart equivalent with image_uri provided."""