Skip to content

Commit 9ff3e5f

Browse files
authored
fix(serve): remove IC data-source collapse hack from recommendation deploy (#6101)
* fix(serve): remove IC data-source collapse hack from recommendation deploy Inference Components now support AdditionalModelDataSources natively (kernel tuning / speculative decoding channels), so _deploy_recommendation no longer needs to collapse an optimized recommendation's base_model + draft channels into a single primary ModelDataSource. Deploy the recommendation's ModelPackage directly and let the hosting stack resolve the channels. Removes the base_model-promotion / OPTION_SPECULATIVE_DRAFT_MODEL rewiring block and its now-unused shape imports (AdditionalModelDataSource, ModelDataSource, S3ModelDataSource, SPECULATIVE_DRAFT_MODEL). Updates the speculative-decoding unit tests to assert the ModelPackage pass-through. * test(serve): integ test deploying an SD/KT model as an Inference Component Adds an end-to-end integration test that builds a model carrying speculative-decoding / kernel-tuning AdditionalModelDataSources (base + draft channels) and deploys it as an Inference Component via ModelBuilder.deploy(inference_config=ResourceRequirements(...)). Asserts the Inference Component reaches InService and that the deployed model still carries the additional sources (guards against a silent client-side collapse). Marked slow_test + gpu_intensive. * test(serve): fix SD/KT IC integ test build path and trim docstring The test built the IC model via ModelBuilder(model_path=<s3_uri>), which raises 'Cannot detect required model or inference spec' — a raw S3 path is not a buildable model spec. Build via ModelBuilder.from_jumpstart_config so the container/framework resolve, then attach additional_model_data_sources (carried onto the model through _prepare_container_def_base) before deploying as an Inference Component. Also switch the IC instance to g4dn.xlarge (a GPU is needed for the vLLM/LMI container, not for the 0.6B size) and trim the module docstring. * test(serve): poll IC to terminal state before asserting + harden teardown deploy(wait=True) waits for the endpoint, but the Inference Component is created with wait=False, so the IC can still be Creating when deploy() returns. Add _wait_for_ic_terminal to poll the IC to InService/Failed before the status assertion, and wait it out of Creating before teardown (an IC in Creating cannot be deleted, which would strand its GPU endpoint).
1 parent 136706e commit 9ff3e5f

3 files changed

Lines changed: 286 additions & 121 deletions

File tree

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

Lines changed: 10 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@
8282
from sagemaker.serve.validations.check_image_uri import is_1p_image_uri
8383
from sagemaker.core.inference_config import ResourceRequirements
8484
from sagemaker.serve.inference_recommendation_mixin import _InferenceRecommenderMixin
85-
from sagemaker.serve.model_builder_utils import _ModelBuilderUtils, SPECULATIVE_DRAFT_MODEL
85+
from sagemaker.serve.model_builder_utils import _ModelBuilderUtils
8686
from sagemaker.serve.model_builder_servers import _ModelBuilderServers
8787
from sagemaker.serve.validations.optimization import _validate_optimization_configuration
8888
from sagemaker.core.enums import Tag
@@ -5158,12 +5158,9 @@ def _deploy_recommendation(
51585158
import time as _time
51595159
import uuid as _uuid
51605160
from sagemaker.core.shapes.shapes import (
5161-
AdditionalModelDataSource as _AdditionalModelDataSource,
51625161
ContainerDefinition as _ContainerDefinition,
5163-
ModelDataSource as _ModelDataSource,
51645162
ProductionVariant as _ProductionVariant,
51655163
ProductionVariantRoutingConfig as _ProductionVariantRoutingConfig,
5166-
S3ModelDataSource as _S3ModelDataSource,
51675164
)
51685165

51695166
role = role or self.role_arn
@@ -5275,65 +5272,15 @@ def _deploy_recommendation(
52755272
)
52765273
resolved_endpoint_name = endpoint_name or f"sm-rec-endpoint-{ts}-{suffix}"
52775274

5278-
# Optimized recommendations put the base weights (and any draft model)
5279-
# in additional model data sources with an empty primary source. Promote
5280-
# base_model to the primary source, and keep any draft channel attached
5281-
# so OPTION_SPECULATIVE_DRAFT_MODEL points at a populated path.
5282-
pkg_container = (
5283-
described.get("InferenceSpecification", {}).get("Containers", [{}]) or [{}]
5284-
)[0]
5285-
additional_sources = pkg_container.get("AdditionalModelDataSources") or []
5286-
by_channel = {s.get("ChannelName"): s for s in additional_sources}
5287-
base_source = by_channel.get("base_model")
5288-
5289-
primary_model_dir = "/opt/ml/model"
5290-
if base_source:
5291-
base_s3 = base_source.get("S3DataSource", {})
5292-
base_channel_path = f"{SPECULATIVE_DRAFT_MODEL}/base_model"
5293-
# Repoint env vars (e.g. HF_MODEL_ID) that referenced the old
5294-
# channel path to the primary mount now holding the base weights.
5295-
env = {
5296-
k: (primary_model_dir if v == base_channel_path else v)
5297-
for k, v in (pkg_container.get("Environment") or {}).items()
5298-
}
5299-
# Keep any draft channel attached and point the env var at its mount.
5300-
draft_sources = []
5301-
for channel_name, source in by_channel.items():
5302-
if channel_name == "base_model":
5303-
continue
5304-
draft_s3 = source.get("S3DataSource", {})
5305-
draft_sources.append(
5306-
_AdditionalModelDataSource(
5307-
channel_name=channel_name,
5308-
s3_data_source=_S3ModelDataSource(
5309-
s3_uri=draft_s3.get("S3Uri"),
5310-
s3_data_type=draft_s3.get("S3DataType", "S3Prefix"),
5311-
compression_type=draft_s3.get("CompressionType", "None"),
5312-
),
5313-
)
5314-
)
5315-
env["OPTION_SPECULATIVE_DRAFT_MODEL"] = (
5316-
f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}/"
5317-
)
5318-
primary_container = _ContainerDefinition(
5319-
image=pkg_container.get("Image"),
5320-
model_data_source=_ModelDataSource(
5321-
s3_data_source=_S3ModelDataSource(
5322-
s3_uri=base_s3.get("S3Uri"),
5323-
s3_data_type=base_s3.get("S3DataType", "S3Prefix"),
5324-
compression_type=base_s3.get("CompressionType", "None"),
5325-
)
5326-
),
5327-
additional_model_data_sources=draft_sources or None,
5328-
environment=env,
5329-
)
5330-
else:
5331-
container_def_kwargs = {"model_package_name": model_package_arn}
5332-
if inference_specification_name:
5333-
container_def_kwargs[
5334-
"inference_specification_name"
5335-
] = inference_specification_name
5336-
primary_container = _ContainerDefinition(**container_def_kwargs)
5275+
# Deploy directly from the recommendation's ModelPackage. Optimized
5276+
# recommendations (kernel tuning / speculative decoding) carry the base
5277+
# weights and any draft model as AdditionalModelDataSources on the
5278+
# package; the hosting stack resolves those channels itself, so no
5279+
# client-side collapsing is needed.
5280+
container_def_kwargs = {"model_package_name": model_package_arn}
5281+
if inference_specification_name:
5282+
container_def_kwargs["inference_specification_name"] = inference_specification_name
5283+
primary_container = _ContainerDefinition(**container_def_kwargs)
53375284

53385285
Model.create(
53395286
model_name=resolved_model_name,
Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,251 @@
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+
"""End-to-end: deploy a speculative-decoding / kernel-tuning model as an
14+
Inference Component via ``ModelBuilder``."""
15+
from __future__ import absolute_import
16+
17+
import logging
18+
import time
19+
import uuid
20+
21+
import pytest
22+
23+
from sagemaker.core.enums import EndpointType
24+
from sagemaker.core.helper.session_helper import Session, get_execution_role
25+
from sagemaker.core.inference_config import ResourceRequirements
26+
from sagemaker.core.resources import Endpoint, EndpointConfig, InferenceComponent, Model
27+
from sagemaker.serve.model_builder import ModelBuilder
28+
from sagemaker.core.training.configs import Compute
29+
30+
logger = logging.getLogger(__name__)
31+
32+
# Small, ungated, chat-templated model. A GPU instance is required by the
33+
# JumpStart vLLM/LMI container (not by the model size); a single-GPU g4dn.xlarge
34+
# (T4) is ample for a 0.6B model.
35+
MODEL_ID = "huggingface-reasoning-qwen3-06b"
36+
INSTANCE_TYPE = "ml.g4dn.xlarge"
37+
# Right-sized for a 0.6B on a single GPU; an oversized request overflows the
38+
# host and the IC never leaves Creating.
39+
IC_MIN_MEMORY_MB = 4096
40+
IC_NUM_ACCELERATORS = 1
41+
# Upper bound for the Inference Component to reach a terminal state.
42+
_IC_TIMEOUT_S = 40 * 60
43+
44+
45+
def _additional_model_data_sources(s3_uri):
46+
"""The AdditionalModelDataSources shape an SD/KT optimized model carries:
47+
base weights + a draft/tuned channel. The same readable artifact backs both
48+
— the IC deploy path's acceptance of the *shape* is what is under test, not
49+
the specific weights."""
50+
51+
def _src(uri):
52+
return {
53+
"S3DataSource": {
54+
"S3Uri": uri,
55+
"S3DataType": "S3Prefix",
56+
"CompressionType": "None",
57+
}
58+
}
59+
60+
return [
61+
{"ChannelName": "base_model", **_src(s3_uri)},
62+
{"ChannelName": "draft_model", **_src(s3_uri)},
63+
]
64+
65+
66+
@pytest.mark.slow_test
67+
@pytest.mark.gpu_intensive
68+
def test_deploy_sdkt_model_as_inference_component():
69+
"""A model carrying SD/KT AdditionalModelDataSources deploys as an
70+
Inference Component and reaches InService via ``ModelBuilder.deploy``."""
71+
logger.info("Starting SD/KT deploy-as-Inference-Component integration test...")
72+
73+
unique_id = f"{int(time.time())}-{uuid.uuid4().hex[:8]}"
74+
role = get_execution_role(sagemaker_session=Session())
75+
src_model_name = f"air-sdkt-src-{unique_id}"
76+
ic_model_name = f"air-sdkt-icmodel-{unique_id}"
77+
endpoint_name = f"air-sdkt-ic-ep-{unique_id}"
78+
ic_name = f"air-sdkt-ic-{unique_id}"
79+
80+
from sagemaker.core.jumpstart.configs import JumpStartConfig
81+
82+
def _jumpstart_builder():
83+
return ModelBuilder.from_jumpstart_config(
84+
jumpstart_config=JumpStartConfig(model_id=MODEL_ID),
85+
compute=Compute(instance_type=INSTANCE_TYPE),
86+
role_arn=role,
87+
)
88+
89+
source_model = None
90+
endpoint = None
91+
92+
try:
93+
# First build resolves the JumpStart container + a readable weights
94+
# artifact for this account/region; read its S3 URI to back the SD/KT
95+
# channels (JumpStart populates the artifact during build, not before).
96+
source_model = _jumpstart_builder().build(model_name=src_model_name)
97+
src_primary = getattr(Model.get(model_name=src_model_name), "primary_container", None)
98+
base_s3 = _extract_s3_uri(src_primary)
99+
assert base_s3, f"Could not resolve a readable S3 artifact for {MODEL_ID}"
100+
logger.info("Resolved base artifact: %s", base_s3)
101+
102+
# Second build carries the SD/KT additional data sources (base + draft
103+
# channels) onto the model, then deploys it as an Inference Component:
104+
# inference_config=ResourceRequirements routes deploy() to the
105+
# INFERENCE_COMPONENT_BASED path. Both channels point at the same
106+
# readable artifact — the deploy-time contract under test is the
107+
# multi-source SHAPE, not the specific optimized weights.
108+
ic_mb = _jumpstart_builder()
109+
ic_mb.additional_model_data_sources = _additional_model_data_sources(base_s3)
110+
ic_mb.build(model_name=ic_model_name)
111+
112+
endpoint = ic_mb.deploy(
113+
endpoint_name=endpoint_name,
114+
inference_config=ResourceRequirements(
115+
requests={
116+
"num_accelerators": IC_NUM_ACCELERATORS,
117+
"memory": IC_MIN_MEMORY_MB,
118+
"copies": 1,
119+
},
120+
),
121+
inference_component_name=ic_name,
122+
instance_type=INSTANCE_TYPE,
123+
initial_instance_count=1,
124+
wait=True,
125+
)
126+
logger.info("Deploy returned; endpoint=%s ic=%s", endpoint_name, ic_name)
127+
128+
# deploy(wait=True) waits for the endpoint, but the Inference Component
129+
# is created with wait=False and finishes shortly after — poll it to a
130+
# terminal state before asserting.
131+
ic = _wait_for_ic_terminal(ic_name)
132+
133+
# The Inference Component — referencing a model with base + draft
134+
# AdditionalModelDataSources — must be InService. This is the assertion
135+
# that turns red if hosting rejects the multi-source shape at
136+
# create/deploy time.
137+
assert ic.inference_component_status == "InService", (
138+
f"Inference Component {ic_name} did not reach InService: "
139+
f"{ic.inference_component_status} / "
140+
f"{getattr(ic, 'failure_reason', None)}"
141+
)
142+
143+
# And the model the IC references still carries the SD/KT additional
144+
# sources — i.e. we validated the multi-source path, not a silently
145+
# collapsed single-source one.
146+
referenced = ic.specification.model_name
147+
deployed_model = Model.get(model_name=referenced)
148+
channels = _additional_channel_names(deployed_model)
149+
assert {"base_model", "draft_model"}.issubset(channels), (
150+
f"IC model {referenced} lost its SD/KT additional sources "
151+
f"(channels={channels}); the deploy path must not collapse them."
152+
)
153+
logger.info("SD/KT model InService on IC with channels: %s", channels)
154+
155+
finally:
156+
# An IC still in Creating cannot be deleted (the API rejects it), which
157+
# would strand the IC and its GPU endpoint. Wait for it to leave
158+
# Creating first, then delete and confirm it is gone before the endpoint.
159+
try:
160+
_wait_for_ic_terminal(ic_name)
161+
except Exception as exc:
162+
logger.warning("Could not resolve IC %s before teardown: %s", ic_name, exc)
163+
_delete_quietly(
164+
lambda: InferenceComponent.get(inference_component_name=ic_name),
165+
f"InferenceComponent {ic_name}",
166+
wait_gone=True,
167+
)
168+
_delete_quietly(
169+
lambda: Endpoint.get(endpoint_name=endpoint_name),
170+
f"Endpoint {endpoint_name}",
171+
)
172+
_delete_quietly(
173+
lambda: EndpointConfig.get(endpoint_config_name=endpoint_name),
174+
f"EndpointConfig {endpoint_name}",
175+
)
176+
_delete_quietly(
177+
lambda: Model.get(model_name=ic_model_name),
178+
f"Model {ic_model_name}",
179+
)
180+
if source_model:
181+
_delete_quietly(lambda: source_model, f"Model {src_model_name}")
182+
183+
184+
def _wait_for_ic_terminal(ic_name, timeout_s=_IC_TIMEOUT_S):
185+
"""Poll an Inference Component until it leaves ``Creating``/``Updating``.
186+
187+
Returns the resource in its terminal state (``InService`` / ``Failed``);
188+
on timeout, returns the last-observed resource so the caller's assertion
189+
reports the stuck status rather than this helper masking it.
190+
"""
191+
waited = 0
192+
ic = InferenceComponent.get(inference_component_name=ic_name)
193+
while getattr(ic, "inference_component_status", None) in ("Creating", "Updating"):
194+
if waited >= timeout_s:
195+
break
196+
time.sleep(20)
197+
waited += 20
198+
ic = InferenceComponent.get(inference_component_name=ic_name)
199+
return ic
200+
201+
202+
def _extract_s3_uri(container):
203+
"""Pull the S3 artifact URI off a resolved container (ModelDataSource or the
204+
legacy ModelDataUrl), tolerating the resource-object attribute shape."""
205+
if container is None:
206+
return None
207+
mds = getattr(container, "model_data_source", None)
208+
if mds is not None:
209+
s3 = getattr(mds, "s3_data_source", None)
210+
if s3 is not None and getattr(s3, "s3_uri", None):
211+
return s3.s3_uri
212+
return getattr(container, "model_data_url", None)
213+
214+
215+
def _additional_channel_names(model):
216+
"""Return the set of AdditionalModelDataSources channel names on a model."""
217+
primary = getattr(model, "primary_container", None) or getattr(
218+
model, "containers", [None]
219+
)[0]
220+
if primary is None:
221+
return set()
222+
sources = getattr(primary, "additional_model_data_sources", None) or []
223+
names = set()
224+
for s in sources:
225+
name = getattr(s, "channel_name", None)
226+
if name is None and isinstance(s, dict):
227+
name = s.get("ChannelName")
228+
if name:
229+
names.add(name)
230+
return names
231+
232+
233+
def _delete_quietly(resource_factory, label, wait_gone=False):
234+
"""Best-effort delete; log and continue on any failure. When ``wait_gone``,
235+
block until the resource is gone (an endpoint can't be deleted while it
236+
still hosts an Inference Component)."""
237+
try:
238+
resource = resource_factory()
239+
resource.delete()
240+
if wait_gone:
241+
waited = 0
242+
while waited < 10 * 60:
243+
try:
244+
resource_factory()
245+
time.sleep(20)
246+
waited += 20
247+
except Exception:
248+
break
249+
logger.info("Deleted %s", label)
250+
except Exception as exc:
251+
logger.warning("Failed to delete %s: %s", label, exc)

0 commit comments

Comments
 (0)