Skip to content

Commit 5cc330b

Browse files
committed
fix: ModelBuilder resolves private hub artifacts correctly
Fix two defects causing ModelBuilder to ignore private hub when resolving model artifacts, forcing the execution role to access the public JumpStart S3 cache bucket. Defect 1: from_jumpstart_config sets hub_name AFTER __init__ has already called _initialize_jumpstart_config(), which takes the else branch and sets hub_arn = None. Fix: call _initialize_jumpstart_config() again after setting hub_name so hub_arn is correctly derived. Defect 2: _build_for_jumpstart does not forward hub_arn to get_init_kwargs, causing model data to resolve from the public catalog. Fix: pass hub_arn=getattr(self, 'hub_arn', None) to all get_init_kwargs call sites in the build path (_build_for_jumpstart, _detect_jumpstart_image). Impact: Customers deploying from private hubs via ModelBuilder no longer need to grant their execution role s3:GetObject on the public JumpStart cache bucket. Testing: - 8 unit tests covering both defects and end-to-end flow - Integration test for private hub deployment (requires env config)
1 parent 78bb11e commit 5cc330b

5 files changed

Lines changed: 597 additions & 2 deletions

File tree

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3895,6 +3895,16 @@ def from_jumpstart_config(
38953895
mb_instance.resource_requirements = resource_requirements
38963896
mb_instance.model_kms_key = model_kms_key
38973897
mb_instance.hub_name = jumpstart_config.hub_name
3898+
if mb_instance.hub_name and not getattr(mb_instance, "hub_arn", None):
3899+
from sagemaker.core.jumpstart.hub.utils import (
3900+
generate_hub_arn_for_init_kwargs,
3901+
)
3902+
3903+
mb_instance.hub_arn = generate_hub_arn_for_init_kwargs(
3904+
hub_name=mb_instance.hub_name,
3905+
region=mb_instance.region,
3906+
session=mb_instance.sagemaker_session,
3907+
)
38983908
mb_instance.config_name = jumpstart_config.inference_config_name
38993909
mb_instance.accept_eula = jumpstart_config.accept_eula
39003910
mb_instance.tolerate_vulnerable_model = tolerate_vulnerable_model

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -975,7 +975,7 @@ def _build_for_jumpstart(self) -> Model:
975975
self.secret_key = ""
976976

977977
# Get JumpStart model configuration
978-
init_kwargs = get_init_kwargs(
978+
init_kwargs_params = dict(
979979
model_id=self.model,
980980
model_version=self.model_version or "*",
981981
region=self.region,
@@ -984,6 +984,10 @@ def _build_for_jumpstart(self) -> Model:
984984
tolerate_deprecated_model=getattr(self, "tolerate_deprecated_model", None),
985985
config_name=getattr(self, "config_name", None),
986986
)
987+
hub_arn = getattr(self, "hub_arn", None)
988+
if hub_arn:
989+
init_kwargs_params["hub_arn"] = hub_arn
990+
init_kwargs = get_init_kwargs(**init_kwargs_params)
987991

988992
# Configure image URI and environment variables
989993
self.image_uri = self.image_uri or init_kwargs.image_uri

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -629,14 +629,18 @@ def _detect_jumpstart_image(self) -> None:
629629
ValueError: If image URI cannot be determined or JumpStart lookup fails.
630630
"""
631631
try:
632-
init_kwargs = get_init_kwargs(
632+
detect_kwargs = dict(
633633
model_id=self.model,
634634
model_version=getattr(self, "model_version", None) or "*",
635635
region=self.region,
636636
instance_type=getattr(self, "instance_type", None),
637637
tolerate_vulnerable_model=getattr(self, "tolerate_vulnerable_model", None),
638638
tolerate_deprecated_model=getattr(self, "tolerate_deprecated_model", None),
639639
)
640+
hub_arn = getattr(self, "hub_arn", None)
641+
if hub_arn:
642+
detect_kwargs["hub_arn"] = hub_arn
643+
init_kwargs = get_init_kwargs(**detect_kwargs)
640644

641645
self.image_uri = init_kwargs.get("image_uri")
642646
if not self.image_uri:
Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Integration test for private hub artifact resolution.
14+
15+
Verifies that ModelBuilder.from_jumpstart_config with hub_name resolves
16+
model artifacts through the private hub content reference, rather than
17+
falling back to the public JumpStart S3 cache.
18+
19+
This test creates its own private hub, adds a ModelReference, runs the
20+
build flow, and tears everything down afterward. Skips gracefully if
21+
the test environment lacks permissions to create hubs.
22+
"""
23+
from __future__ import absolute_import
24+
25+
import os
26+
import uuid
27+
import time
28+
import logging
29+
30+
import boto3
31+
import pytest
32+
from botocore.exceptions import ClientError
33+
34+
from sagemaker.serve.model_builder import ModelBuilder
35+
from sagemaker.core.jumpstart.configs import JumpStartConfig
36+
from sagemaker.train.configs import Compute
37+
38+
logger = logging.getLogger(__name__)
39+
40+
# Use a small/fast JumpStart model for the test
41+
TEST_MODEL_ID = "huggingface-llm-gemma-2b"
42+
TEST_MODEL_VERSION = "*"
43+
TEST_INSTANCE_TYPE = "ml.g5.xlarge"
44+
TEST_REGION = os.environ.get("TEST_REGION", "us-east-1")
45+
46+
HUB_NAME_PREFIX = "sdk-integ-test-private-hub"
47+
48+
49+
@pytest.fixture(scope="module")
50+
def private_hub():
51+
"""Create a private hub with a ModelReference, tear down after tests."""
52+
sm = boto3.client("sagemaker", region_name=TEST_REGION)
53+
hub_name = f"{HUB_NAME_PREFIX}-{uuid.uuid4().hex[:8]}"
54+
55+
# --- Setup ---
56+
logger.info("Creating private hub: %s", hub_name)
57+
try:
58+
sm.create_hub(
59+
HubName=hub_name,
60+
HubDescription="SDK integration test for private hub artifact resolution",
61+
Tags=[
62+
{"Key": "Purpose", "Value": "sdk-integ-test"},
63+
{"Key": "AutoCleanup", "Value": "true"},
64+
],
65+
)
66+
except ClientError as e:
67+
pytest.skip(f"Cannot create hub (likely missing permissions): {e}")
68+
69+
# Wait for hub to be ready
70+
for _ in range(30):
71+
resp = sm.describe_hub(HubName=hub_name)
72+
if resp["HubStatus"] == "InService":
73+
break
74+
time.sleep(2)
75+
else:
76+
pytest.skip(f"Hub {hub_name} did not reach InService state")
77+
78+
# Add a ModelReference to the public JumpStart content
79+
public_arn = (
80+
f"arn:aws:sagemaker:{TEST_REGION}:aws:hub-content/"
81+
f"SageMakerPublicHub/Model/{TEST_MODEL_ID}"
82+
)
83+
logger.info("Creating hub content reference to: %s", public_arn)
84+
try:
85+
sm.create_hub_content_reference(
86+
HubName=hub_name,
87+
SageMakerPublicHubContentArn=public_arn,
88+
)
89+
except ClientError as e:
90+
pytest.skip(f"Cannot create hub content reference: {e}")
91+
92+
# Wait for content reference to be available
93+
for _ in range(60):
94+
try:
95+
contents = sm.list_hub_contents(
96+
HubName=hub_name,
97+
HubContentType="ModelReference",
98+
)
99+
summaries = contents.get("HubContentSummaries", [])
100+
if any(
101+
s["HubContentName"] == TEST_MODEL_ID and s.get("HubContentStatus") == "Available"
102+
for s in summaries
103+
):
104+
break
105+
except Exception:
106+
pass
107+
time.sleep(3)
108+
else:
109+
pytest.skip(
110+
f"ModelReference for {TEST_MODEL_ID} not available in "
111+
f"hub {hub_name} after 3 minutes"
112+
)
113+
114+
yield hub_name
115+
116+
# --- Teardown ---
117+
logger.info("Cleaning up hub: %s", hub_name)
118+
try:
119+
contents = sm.list_hub_contents(
120+
HubName=hub_name,
121+
HubContentType="ModelReference",
122+
)
123+
for content in contents.get("HubContentSummaries", []):
124+
try:
125+
sm.delete_hub_content_reference(
126+
HubName=hub_name,
127+
HubContentName=content["HubContentName"],
128+
HubContentType="ModelReference",
129+
)
130+
except Exception as e:
131+
logger.warning("Failed to delete content ref: %s", e)
132+
133+
sm.delete_hub(HubName=hub_name)
134+
logger.info("Hub %s deleted", hub_name)
135+
except Exception as e:
136+
logger.warning("Hub cleanup failed: %s", e)
137+
138+
139+
@pytest.fixture(scope="module")
140+
def sagemaker_session():
141+
"""Create a SageMaker session using default credentials."""
142+
from sagemaker.core.helper.session_helper import Session
143+
144+
return Session(boto_session=boto3.Session(region_name=TEST_REGION))
145+
146+
147+
@pytest.fixture(scope="module")
148+
def execution_role():
149+
"""Get a SageMaker execution role from the caller's identity."""
150+
sts = boto3.client("sts", region_name=TEST_REGION)
151+
identity = sts.get_caller_identity()
152+
account_id = identity["Account"]
153+
return f"arn:aws:iam::{account_id}:role/Admin"
154+
155+
156+
@pytest.mark.slow_test
157+
def test_from_jumpstart_config_derives_hub_arn(private_hub, sagemaker_session):
158+
"""Verify from_jumpstart_config correctly derives hub_arn from hub_name."""
159+
js_config = JumpStartConfig(
160+
model_id=TEST_MODEL_ID,
161+
model_version=TEST_MODEL_VERSION,
162+
hub_name=private_hub,
163+
)
164+
165+
mb = ModelBuilder.from_jumpstart_config(
166+
jumpstart_config=js_config,
167+
compute=Compute(instance_type=TEST_INSTANCE_TYPE),
168+
sagemaker_session=sagemaker_session,
169+
)
170+
171+
assert (
172+
mb.hub_arn is not None
173+
), f"hub_arn is None after from_jumpstart_config with hub_name={private_hub}"
174+
assert private_hub in mb.hub_arn
175+
logger.info("hub_arn correctly derived: %s", mb.hub_arn)
176+
177+
178+
@pytest.mark.slow_test
179+
def test_build_resolves_artifacts_via_private_hub(private_hub, execution_role, sagemaker_session):
180+
"""Verify build() resolves model data through the private hub."""
181+
js_config = JumpStartConfig(
182+
model_id=TEST_MODEL_ID,
183+
model_version=TEST_MODEL_VERSION,
184+
hub_name=private_hub,
185+
accept_eula=True,
186+
)
187+
188+
mb = ModelBuilder.from_jumpstart_config(
189+
jumpstart_config=js_config,
190+
role_arn=execution_role,
191+
compute=Compute(instance_type=TEST_INSTANCE_TYPE),
192+
sagemaker_session=sagemaker_session,
193+
)
194+
195+
mb.build()
196+
197+
# After build, model data URI must NOT reference the public cache
198+
model_data = getattr(mb, "s3_model_data_url", None)
199+
assert model_data is not None, "No model_data found after build()"
200+
201+
model_data_str = str(model_data)
202+
assert "jumpstart-cache-prod" not in model_data_str, (
203+
f"Model data still points to public JumpStart cache: "
204+
f"{model_data_str}. Expected private hub artifact resolution."
205+
)
206+
logger.info("Model data resolved to: %s", model_data_str)

0 commit comments

Comments
 (0)