Skip to content

Commit d1fd74f

Browse files
committed
fix(serve): repair HuggingFace -> JumpStart redirect in ModelBuilder
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.
1 parent 9101cef commit d1fd74f

2 files changed

Lines changed: 11 additions & 5 deletions

File tree

sagemaker-serve/pyproject.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,9 @@ where = ["src"]
6161
include = ["sagemaker*"]
6262
namespaces = true
6363

64+
[tool.setuptools.package-data]
65+
"sagemaker.serve" = ["schema/*.json"]
66+
6467
[tool.setuptools.dynamic]
6568
version = { file = "VERSION"}
6669

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

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -942,17 +942,20 @@ def _use_jumpstart_equivalent(self) -> bool:
942942
if not model_task:
943943
model_task = hf_model_md.get("pipeline_tag")
944944
if model_task:
945-
self._hf_schema_builder_init(model_task)
945+
try:
946+
self._hf_schema_builder_init(model_task)
947+
except TaskNotFoundException:
948+
logger.warning(
949+
"HF schema builder unavailable for task %r; "
950+
"JumpStart build path will supply its own schema.",
951+
model_task,
952+
)
946953

947954
huggingface_model_id = self.model
948955
jumpstart_model_id = self._jumpstart_mapping[huggingface_model_id]["jumpstart-model-id"]
949956
self.model = jumpstart_model_id
950957
merged_date = self._jumpstart_mapping[huggingface_model_id].get("merged-at")
951958

952-
# Call _build_for_jumpstart if method exists
953-
if hasattr(self, "_build_for_jumpstart"):
954-
self._build_for_jumpstart()
955-
956959
compare_model_diff_message = (
957960
"If you want to identify the differences between the two, "
958961
"please use model_uris.retrieve() to retrieve the model "

0 commit comments

Comments
 (0)