From 81bee88a1128d545acf8c88124e31bf8e5d13900 Mon Sep 17 00:00:00 2001 From: ZealSV Date: Fri, 24 Jul 2026 12:32:27 -0700 Subject: [PATCH 1/4] 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. --- .../src/sagemaker/serve/model_builder.py | 73 +++------------- .../test_model_builder_recommendations.py | 83 ++++++------------- 2 files changed, 35 insertions(+), 121 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 58761105e1..1d60d3e6be 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -82,7 +82,7 @@ from sagemaker.serve.validations.check_image_uri import is_1p_image_uri from sagemaker.core.inference_config import ResourceRequirements from sagemaker.serve.inference_recommendation_mixin import _InferenceRecommenderMixin -from sagemaker.serve.model_builder_utils import _ModelBuilderUtils, SPECULATIVE_DRAFT_MODEL +from sagemaker.serve.model_builder_utils import _ModelBuilderUtils from sagemaker.serve.model_builder_servers import _ModelBuilderServers from sagemaker.serve.validations.optimization import _validate_optimization_configuration from sagemaker.core.enums import Tag @@ -5120,12 +5120,9 @@ def _deploy_recommendation( import time as _time import uuid as _uuid from sagemaker.core.shapes.shapes import ( - AdditionalModelDataSource as _AdditionalModelDataSource, ContainerDefinition as _ContainerDefinition, - ModelDataSource as _ModelDataSource, ProductionVariant as _ProductionVariant, ProductionVariantRoutingConfig as _ProductionVariantRoutingConfig, - S3ModelDataSource as _S3ModelDataSource, ) role = role or self.role_arn @@ -5237,65 +5234,15 @@ def _deploy_recommendation( ) resolved_endpoint_name = endpoint_name or f"sm-rec-endpoint-{ts}-{suffix}" - # Optimized recommendations put the base weights (and any draft model) - # in additional model data sources with an empty primary source. Promote - # base_model to the primary source, and keep any draft channel attached - # so OPTION_SPECULATIVE_DRAFT_MODEL points at a populated path. - pkg_container = ( - described.get("InferenceSpecification", {}).get("Containers", [{}]) or [{}] - )[0] - additional_sources = pkg_container.get("AdditionalModelDataSources") or [] - by_channel = {s.get("ChannelName"): s for s in additional_sources} - base_source = by_channel.get("base_model") - - primary_model_dir = "/opt/ml/model" - if base_source: - base_s3 = base_source.get("S3DataSource", {}) - base_channel_path = f"{SPECULATIVE_DRAFT_MODEL}/base_model" - # Repoint env vars (e.g. HF_MODEL_ID) that referenced the old - # channel path to the primary mount now holding the base weights. - env = { - k: (primary_model_dir if v == base_channel_path else v) - for k, v in (pkg_container.get("Environment") or {}).items() - } - # Keep any draft channel attached and point the env var at its mount. - draft_sources = [] - for channel_name, source in by_channel.items(): - if channel_name == "base_model": - continue - draft_s3 = source.get("S3DataSource", {}) - draft_sources.append( - _AdditionalModelDataSource( - channel_name=channel_name, - s3_data_source=_S3ModelDataSource( - s3_uri=draft_s3.get("S3Uri"), - s3_data_type=draft_s3.get("S3DataType", "S3Prefix"), - compression_type=draft_s3.get("CompressionType", "None"), - ), - ) - ) - env["OPTION_SPECULATIVE_DRAFT_MODEL"] = ( - f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}/" - ) - primary_container = _ContainerDefinition( - image=pkg_container.get("Image"), - model_data_source=_ModelDataSource( - s3_data_source=_S3ModelDataSource( - s3_uri=base_s3.get("S3Uri"), - s3_data_type=base_s3.get("S3DataType", "S3Prefix"), - compression_type=base_s3.get("CompressionType", "None"), - ) - ), - additional_model_data_sources=draft_sources or None, - environment=env, - ) - else: - container_def_kwargs = {"model_package_name": model_package_arn} - if inference_specification_name: - container_def_kwargs[ - "inference_specification_name" - ] = inference_specification_name - primary_container = _ContainerDefinition(**container_def_kwargs) + # Deploy directly from the recommendation's ModelPackage. Optimized + # recommendations (kernel tuning / speculative decoding) carry the base + # weights and any draft model as AdditionalModelDataSources on the + # package; the hosting stack resolves those channels itself, so no + # client-side collapsing is needed. + container_def_kwargs = {"model_package_name": model_package_arn} + if inference_specification_name: + container_def_kwargs["inference_specification_name"] = inference_specification_name + primary_container = _ContainerDefinition(**container_def_kwargs) Model.create( model_name=resolved_model_name, diff --git a/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_model_builder_recommendations.py b/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_model_builder_recommendations.py index 9354b01c68..8f8375f3fa 100644 --- a/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_model_builder_recommendations.py +++ b/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_model_builder_recommendations.py @@ -247,10 +247,11 @@ def test_builder_role_used_when_role_omitted(self, mb_class): class TestDeployRecommendationSpeculativeDecoding: - """Speculative-decoding / kernel-tuning recommendations (ModelPackage carries - AdditionalModelDataSources) are collapsed to a single ModelDataSource + - OPTION_SPECULATIVE_DRAFT_MODEL env so the model can be hosted (Inference - Components reject AdditionalModelDataSources).""" + """Optimized recommendations (kernel tuning / speculative decoding) carry + the base weights and any draft model as AdditionalModelDataSources on the + ModelPackage. The hosting stack (including Inference Components) now resolves + those channels itself, so deploy just passes the ModelPackage through — no + client-side collapsing of the additional sources into a primary source.""" def _make_builder(self, mb_class, described): mb = mb_class(sagemaker_session=MagicMock()) @@ -291,52 +292,22 @@ def _channel(name, uri): }, } - def test_base_model_channel_promoted_to_primary_source(self, mb_class): - # Optimized recs put the weights in a base_model channel with an empty - # primary ModelDataSource; the base_model S3 uri becomes the primary source, - # and env vars that pointed at the old channel path are repointed to - # /opt/ml/model (where the primary source mounts). + def test_optimized_recommendation_deploys_via_model_package(self, mb_class): + # An optimized rec exposes base_model + draft_model as additional data + # sources; deploy no longer collapses them — it passes the ModelPackage + # through and lets the hosting stack resolve the channels. + # + # The AdditionalModelDataSources below are intentionally present as a + # regression guard: the current path reads only ModelApprovalStatus and + # ignores them, so the two negative assertions (no model_data_source / + # no additional_model_data_sources) only fail if the client-side + # collapse logic is ever re-introduced. described = { "ModelApprovalStatus": "Approved", "InferenceSpecification": { "Containers": [ { "Image": "123.dkr.ecr.us-west-2.amazonaws.com/djl-inference:latest", - "Environment": { - "OPTION_TENSOR_PARALLEL_DEGREE": "1", - "HF_MODEL_ID": "/opt/ml/additional-model-data-sources/base_model", - "option.model_id": "/opt/ml/additional-model-data-sources/base_model", - }, - "AdditionalModelDataSources": [ - self._channel("base_model", "s3://base/weights/") - ], - } - ] - }, - } - mb = self._make_builder(mb_class, described) - container = self._deploy(mb).call_args.kwargs["primary_container"] - - # No AdditionalModelDataSources on the created model (IC would reject it). - assert not getattr(container, "additional_model_data_sources", None) - assert not getattr(container, "model_package_name", None) - # base_model weights promoted to the primary source. - assert container.model_data_source.s3_data_source.s3_uri == "s3://base/weights/" - # Env vars that referenced the old channel path now point at the primary mount. - assert container.environment["HF_MODEL_ID"] == "/opt/ml/model" - assert container.environment["option.model_id"] == "/opt/ml/model" - # No draft channel here, so no draft env. - assert "OPTION_SPECULATIVE_DRAFT_MODEL" not in container.environment - assert container.environment["OPTION_TENSOR_PARALLEL_DEGREE"] == "1" - - def test_draft_model_channel_mounted_via_env(self, mb_class): - described = { - "ModelApprovalStatus": "Approved", - "InferenceSpecification": { - "Containers": [ - { - "Image": "123.dkr.ecr.us-west-2.amazonaws.com/djl-inference:latest", - "Environment": {}, "AdditionalModelDataSources": [ self._channel("base_model", "s3://base/weights/"), self._channel("draft_model", "s3://draft/model/"), @@ -348,19 +319,15 @@ def test_draft_model_channel_mounted_via_env(self, mb_class): mb = self._make_builder(mb_class, described) container = self._deploy(mb).call_args.kwargs["primary_container"] - # base_model is the primary source. - assert container.model_data_source.s3_data_source.s3_uri == "s3://base/weights/" - # The draft channel stays attached as an additional model data source so - # the speculative-decoding env path is actually populated at runtime - # (a model-based endpoint supports additional model data sources). - assert len(container.additional_model_data_sources) == 1 - draft = container.additional_model_data_sources[0] - assert draft.channel_name == "draft_model" - assert draft.s3_data_source.s3_uri == "s3://draft/model/" - assert ( - container.environment["OPTION_SPECULATIVE_DRAFT_MODEL"] - == "/opt/ml/additional-model-data-sources/draft_model/" - ) + # Deployed straight from the ModelPackage — no client-side collapsing. + assert container.model_package_name == "arn:.../p/A" + # The row's inference_specification_name is threaded onto the container + # (guards the `if inference_specification_name:` branch from regressing). + assert container.inference_specification_name == "A" + # No collapse: the additional sources above must NOT become a primary + # model_data_source / additional_model_data_sources on the container. + assert not getattr(container, "model_data_source", None) + assert not getattr(container, "additional_model_data_sources", None) def test_no_additional_sources_uses_model_package(self, mb_class): described = { @@ -370,7 +337,7 @@ def test_no_additional_sources_uses_model_package(self, mb_class): mb = self._make_builder(mb_class, described) model_create = self._deploy(mb) container = model_create.call_args.kwargs["primary_container"] - # Plain (non-speculative-decoding) recommendation still deploys via the ModelPackage. + # Plain (non-optimized) recommendation also deploys via the ModelPackage. assert container.model_package_name == "arn:.../p/A" From 39365fd55ee1e7d09954086500e6df6bdc1366de Mon Sep 17 00:00:00 2001 From: ZealSV Date: Wed, 29 Jul 2026 13:16:25 -0700 Subject: [PATCH 2/4] 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. --- ...ference_recommender_sdkt_ic_integration.py | 245 ++++++++++++++++++ 1 file changed, 245 insertions(+) create mode 100644 sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py diff --git a/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py b/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py new file mode 100644 index 0000000000..5b49639949 --- /dev/null +++ b/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py @@ -0,0 +1,245 @@ +# 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. +"""End-to-end: deploy a speculative-decoding / kernel-tuning model as an +Inference Component via ``ModelBuilder``. + +Why this test exists +-------------------- +Speculative-decoding (SD) and kernel-tuning (KT) optimized models host **two or +more model artifacts on one endpoint** — the base weights plus a draft model +(SD) or tuned-kernel artifacts (KT). On the model container these arrive as +``AdditionalModelDataSources``. Inference Components are the hosting surface +customers use to pack such multi-artifact models onto an endpoint, so +"deploy the optimized model as an Inference Component" is a first-class path. + +``ModelBuilder`` deploys as an Inference Component when ``deploy()`` is given an +``inference_config=ResourceRequirements(...)``. This test builds a model +carrying the SD/KT ``AdditionalModelDataSources`` shape and deploys it that way, +asserting the Inference Component reaches ``InService`` — i.e. hosting accepts +and loads the multi-source model. It guards two failure modes: an Inference +Component that rejects the multi-source model at deploy time, and a deploy path +that silently collapses the additional sources into a single primary source +(the post-deploy shape assertion catches the latter). + +Cost / scope +------------ +Producing a genuine optimized ModelPackage inline is infeasible in CI (SD/KT +optimization requires substantial multi-GPU training capacity). So this test +builds a model that carries the SD/KT ``AdditionalModelDataSources`` shape +directly (base + draft channels) — the shape is the deploy-time contract under +test, not the specific optimized weights. Marked ``slow_test`` + +``gpu_intensive`` (one GPU Inference Component endpoint). +""" +from __future__ import absolute_import + +import logging +import time +import uuid + +import pytest + +from sagemaker.core.enums import EndpointType +from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.inference_config import ResourceRequirements +from sagemaker.core.resources import Endpoint, EndpointConfig, InferenceComponent, Model +from sagemaker.serve.model_builder import ModelBuilder +from sagemaker.core.training.configs import Compute + +logger = logging.getLogger(__name__) + +# Small, ungated, chat-templated model. Its readable S3 artifact stands in for +# the optimized recommendation's base + draft channels, and it fits a +# single-GPU g6.2xlarge Inference Component. +MODEL_ID = "huggingface-reasoning-qwen3-06b" +INSTANCE_TYPE = "ml.g6.2xlarge" +# Right-sized for a 0.6B on a single L4; a too-large request overflows the host +# and the IC never leaves Creating. +IC_MIN_MEMORY_MB = 4096 +IC_NUM_ACCELERATORS = 1 + + +def _additional_model_data_sources(s3_uri): + """The AdditionalModelDataSources shape an SD/KT optimized model carries: + base weights + a draft/tuned channel. The same readable artifact backs both + — the IC deploy path's acceptance of the *shape* is what is under test, not + the specific weights.""" + + def _src(uri): + return { + "S3DataSource": { + "S3Uri": uri, + "S3DataType": "S3Prefix", + "CompressionType": "None", + } + } + + return [ + {"ChannelName": "base_model", **_src(s3_uri)}, + {"ChannelName": "draft_model", **_src(s3_uri)}, + ] + + +@pytest.mark.slow_test +@pytest.mark.gpu_intensive +def test_deploy_sdkt_model_as_inference_component(): + """A model carrying SD/KT AdditionalModelDataSources deploys as an + Inference Component and reaches InService via ``ModelBuilder.deploy``.""" + logger.info("Starting SD/KT deploy-as-Inference-Component integration test...") + + unique_id = f"{int(time.time())}-{uuid.uuid4().hex[:8]}" + role = get_execution_role(sagemaker_session=Session()) + src_model_name = f"air-sdkt-src-{unique_id}" + ic_model_name = f"air-sdkt-icmodel-{unique_id}" + endpoint_name = f"air-sdkt-ic-ep-{unique_id}" + ic_name = f"air-sdkt-ic-{unique_id}" + + source_model = None + endpoint = None + + try: + # Build a source model just to obtain a readable image + S3 artifact for + # this account/region (JumpStart resolves the container + weights). + from sagemaker.core.jumpstart.configs import JumpStartConfig + + source_mb = ModelBuilder.from_jumpstart_config( + jumpstart_config=JumpStartConfig(model_id=MODEL_ID), + compute=Compute(instance_type=INSTANCE_TYPE), + role_arn=role, + ) + source_model = source_mb.build(model_name=src_model_name) + src_container = Model.get(model_name=src_model_name) + primary = getattr(src_container, "primary_container", None) or getattr( + src_container, "containers", [None] + )[0] + base_s3 = _extract_s3_uri(primary) + assert base_s3, f"Could not resolve a readable S3 artifact for {MODEL_ID}" + logger.info("Resolved base artifact: %s", base_s3) + + # Build a model that carries the SD/KT additional sources, then deploy it + # as an Inference Component (inference_config=ResourceRequirements routes + # deploy() to the INFERENCE_COMPONENT_BASED path). + ic_mb = ModelBuilder(model_path=base_s3, role_arn=role) + ic_mb.additional_model_data_sources = _additional_model_data_sources(base_s3) + ic_mb.build(model_name=ic_model_name) + + endpoint = ic_mb.deploy( + endpoint_name=endpoint_name, + inference_config=ResourceRequirements( + requests={ + "num_accelerators": IC_NUM_ACCELERATORS, + "memory": IC_MIN_MEMORY_MB, + "copies": 1, + }, + ), + inference_component_name=ic_name, + instance_type=INSTANCE_TYPE, + initial_instance_count=1, + wait=True, + ) + logger.info("Deploy returned; endpoint=%s ic=%s", endpoint_name, ic_name) + + # The Inference Component — referencing a model with base + draft + # AdditionalModelDataSources — must be InService. This is the assertion + # that turns red if hosting rejects the multi-source shape at + # create/deploy time. + ic = InferenceComponent.get(inference_component_name=ic_name) + assert ic.inference_component_status == "InService", ( + f"Inference Component {ic_name} did not reach InService: " + f"{ic.inference_component_status} / " + f"{getattr(ic, 'failure_reason', None)}" + ) + + # And the model the IC references still carries the SD/KT additional + # sources — i.e. we validated the multi-source path, not a silently + # collapsed single-source one. + referenced = ic.specification.model_name + deployed_model = Model.get(model_name=referenced) + channels = _additional_channel_names(deployed_model) + assert {"base_model", "draft_model"}.issubset(channels), ( + f"IC model {referenced} lost its SD/KT additional sources " + f"(channels={channels}); the deploy path must not collapse them." + ) + logger.info("SD/KT model InService on IC with channels: %s", channels) + + finally: + _delete_quietly( + lambda: InferenceComponent.get(inference_component_name=ic_name), + f"InferenceComponent {ic_name}", + wait_gone=True, + ) + _delete_quietly( + lambda: Endpoint.get(endpoint_name=endpoint_name), + f"Endpoint {endpoint_name}", + ) + _delete_quietly( + lambda: EndpointConfig.get(endpoint_config_name=endpoint_name), + f"EndpointConfig {endpoint_name}", + ) + _delete_quietly( + lambda: Model.get(model_name=ic_model_name), + f"Model {ic_model_name}", + ) + if source_model: + _delete_quietly(lambda: source_model, f"Model {src_model_name}") + + +def _extract_s3_uri(container): + """Pull the S3 artifact URI off a resolved container (ModelDataSource or the + legacy ModelDataUrl), tolerating the resource-object attribute shape.""" + if container is None: + return None + mds = getattr(container, "model_data_source", None) + if mds is not None: + s3 = getattr(mds, "s3_data_source", None) + if s3 is not None and getattr(s3, "s3_uri", None): + return s3.s3_uri + return getattr(container, "model_data_url", None) + + +def _additional_channel_names(model): + """Return the set of AdditionalModelDataSources channel names on a model.""" + primary = getattr(model, "primary_container", None) or getattr( + model, "containers", [None] + )[0] + if primary is None: + return set() + sources = getattr(primary, "additional_model_data_sources", None) or [] + names = set() + for s in sources: + name = getattr(s, "channel_name", None) + if name is None and isinstance(s, dict): + name = s.get("ChannelName") + if name: + names.add(name) + return names + + +def _delete_quietly(resource_factory, label, wait_gone=False): + """Best-effort delete; log and continue on any failure. When ``wait_gone``, + block until the resource is gone (an endpoint can't be deleted while it + still hosts an Inference Component).""" + try: + resource = resource_factory() + resource.delete() + if wait_gone: + waited = 0 + while waited < 10 * 60: + try: + resource_factory() + time.sleep(20) + waited += 20 + except Exception: + break + logger.info("Deleted %s", label) + except Exception as exc: + logger.warning("Failed to delete %s: %s", label, exc) From 3786b9bfa650af7252ed943e5cd1a6d7a3e139cd Mon Sep 17 00:00:00 2001 From: ZealSV Date: Wed, 29 Jul 2026 15:07:15 -0700 Subject: [PATCH 3/4] test(serve): fix SD/KT IC integ test build path and trim docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The test built the IC model via ModelBuilder(model_path=), 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. --- ...ference_recommender_sdkt_ic_integration.py | 78 +++++++------------ 1 file changed, 27 insertions(+), 51 deletions(-) diff --git a/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py b/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py index 5b49639949..a9e74d1156 100644 --- a/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py +++ b/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py @@ -11,35 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """End-to-end: deploy a speculative-decoding / kernel-tuning model as an -Inference Component via ``ModelBuilder``. - -Why this test exists --------------------- -Speculative-decoding (SD) and kernel-tuning (KT) optimized models host **two or -more model artifacts on one endpoint** — the base weights plus a draft model -(SD) or tuned-kernel artifacts (KT). On the model container these arrive as -``AdditionalModelDataSources``. Inference Components are the hosting surface -customers use to pack such multi-artifact models onto an endpoint, so -"deploy the optimized model as an Inference Component" is a first-class path. - -``ModelBuilder`` deploys as an Inference Component when ``deploy()`` is given an -``inference_config=ResourceRequirements(...)``. This test builds a model -carrying the SD/KT ``AdditionalModelDataSources`` shape and deploys it that way, -asserting the Inference Component reaches ``InService`` — i.e. hosting accepts -and loads the multi-source model. It guards two failure modes: an Inference -Component that rejects the multi-source model at deploy time, and a deploy path -that silently collapses the additional sources into a single primary source -(the post-deploy shape assertion catches the latter). - -Cost / scope ------------- -Producing a genuine optimized ModelPackage inline is infeasible in CI (SD/KT -optimization requires substantial multi-GPU training capacity). So this test -builds a model that carries the SD/KT ``AdditionalModelDataSources`` shape -directly (base + draft channels) — the shape is the deploy-time contract under -test, not the specific optimized weights. Marked ``slow_test`` + -``gpu_intensive`` (one GPU Inference Component endpoint). -""" +Inference Component via ``ModelBuilder``.""" from __future__ import absolute_import import logging @@ -58,11 +30,12 @@ logger = logging.getLogger(__name__) # Small, ungated, chat-templated model. Its readable S3 artifact stands in for -# the optimized recommendation's base + draft channels, and it fits a -# single-GPU g6.2xlarge Inference Component. +# the base + draft channels. A GPU instance is required by the JumpStart +# vLLM/LMI container (not by the model size); a single-GPU g4dn.xlarge (T4) is +# ample for a 0.6B model. MODEL_ID = "huggingface-reasoning-qwen3-06b" -INSTANCE_TYPE = "ml.g6.2xlarge" -# Right-sized for a 0.6B on a single L4; a too-large request overflows the host +INSTANCE_TYPE = "ml.g4dn.xlarge" +# Right-sized for a 0.6B on a single GPU; a too-large request overflows the host # and the IC never leaves Creating. IC_MIN_MEMORY_MB = 4096 IC_NUM_ACCELERATORS = 1 @@ -103,32 +76,35 @@ def test_deploy_sdkt_model_as_inference_component(): endpoint_name = f"air-sdkt-ic-ep-{unique_id}" ic_name = f"air-sdkt-ic-{unique_id}" - source_model = None - endpoint = None + from sagemaker.core.jumpstart.configs import JumpStartConfig - try: - # Build a source model just to obtain a readable image + S3 artifact for - # this account/region (JumpStart resolves the container + weights). - from sagemaker.core.jumpstart.configs import JumpStartConfig - - source_mb = ModelBuilder.from_jumpstart_config( + def _jumpstart_builder(): + return ModelBuilder.from_jumpstart_config( jumpstart_config=JumpStartConfig(model_id=MODEL_ID), compute=Compute(instance_type=INSTANCE_TYPE), role_arn=role, ) - source_model = source_mb.build(model_name=src_model_name) - src_container = Model.get(model_name=src_model_name) - primary = getattr(src_container, "primary_container", None) or getattr( - src_container, "containers", [None] - )[0] - base_s3 = _extract_s3_uri(primary) + + source_model = None + endpoint = None + + try: + # First build resolves the JumpStart container + a readable weights + # artifact for this account/region; read its S3 URI to back the SD/KT + # channels (JumpStart populates the artifact during build, not before). + source_model = _jumpstart_builder().build(model_name=src_model_name) + src_primary = getattr(Model.get(model_name=src_model_name), "primary_container", None) + base_s3 = _extract_s3_uri(src_primary) assert base_s3, f"Could not resolve a readable S3 artifact for {MODEL_ID}" logger.info("Resolved base artifact: %s", base_s3) - # Build a model that carries the SD/KT additional sources, then deploy it - # as an Inference Component (inference_config=ResourceRequirements routes - # deploy() to the INFERENCE_COMPONENT_BASED path). - ic_mb = ModelBuilder(model_path=base_s3, role_arn=role) + # Second build carries the SD/KT additional data sources (base + draft + # channels) onto the model, then deploys it as an Inference Component: + # inference_config=ResourceRequirements routes deploy() to the + # INFERENCE_COMPONENT_BASED path. Both channels point at the same + # readable artifact — the deploy-time contract under test is the + # multi-source SHAPE, not the specific optimized weights. + ic_mb = _jumpstart_builder() ic_mb.additional_model_data_sources = _additional_model_data_sources(base_s3) ic_mb.build(model_name=ic_model_name) From 8c98238d7e0a4a99319be5bb3205ed88c8742dcb Mon Sep 17 00:00:00 2001 From: ZealSV Date: Wed, 29 Jul 2026 15:29:53 -0700 Subject: [PATCH 4/4] 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). --- ...ference_recommender_sdkt_ic_integration.py | 44 ++++++++++++++++--- 1 file changed, 37 insertions(+), 7 deletions(-) diff --git a/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py b/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py index a9e74d1156..7c72b9e915 100644 --- a/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py +++ b/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py @@ -29,16 +29,17 @@ logger = logging.getLogger(__name__) -# Small, ungated, chat-templated model. Its readable S3 artifact stands in for -# the base + draft channels. A GPU instance is required by the JumpStart -# vLLM/LMI container (not by the model size); a single-GPU g4dn.xlarge (T4) is -# ample for a 0.6B model. +# Small, ungated, chat-templated model. A GPU instance is required by the +# JumpStart vLLM/LMI container (not by the model size); a single-GPU g4dn.xlarge +# (T4) is ample for a 0.6B model. MODEL_ID = "huggingface-reasoning-qwen3-06b" INSTANCE_TYPE = "ml.g4dn.xlarge" -# Right-sized for a 0.6B on a single GPU; a too-large request overflows the host -# and the IC never leaves Creating. +# Right-sized for a 0.6B on a single GPU; an oversized request overflows the +# host and the IC never leaves Creating. IC_MIN_MEMORY_MB = 4096 IC_NUM_ACCELERATORS = 1 +# Upper bound for the Inference Component to reach a terminal state. +_IC_TIMEOUT_S = 40 * 60 def _additional_model_data_sources(s3_uri): @@ -124,11 +125,15 @@ def _jumpstart_builder(): ) logger.info("Deploy returned; endpoint=%s ic=%s", endpoint_name, ic_name) + # deploy(wait=True) waits for the endpoint, but the Inference Component + # is created with wait=False and finishes shortly after — poll it to a + # terminal state before asserting. + ic = _wait_for_ic_terminal(ic_name) + # The Inference Component — referencing a model with base + draft # AdditionalModelDataSources — must be InService. This is the assertion # that turns red if hosting rejects the multi-source shape at # create/deploy time. - ic = InferenceComponent.get(inference_component_name=ic_name) assert ic.inference_component_status == "InService", ( f"Inference Component {ic_name} did not reach InService: " f"{ic.inference_component_status} / " @@ -148,6 +153,13 @@ def _jumpstart_builder(): logger.info("SD/KT model InService on IC with channels: %s", channels) finally: + # An IC still in Creating cannot be deleted (the API rejects it), which + # would strand the IC and its GPU endpoint. Wait for it to leave + # Creating first, then delete and confirm it is gone before the endpoint. + try: + _wait_for_ic_terminal(ic_name) + except Exception as exc: + logger.warning("Could not resolve IC %s before teardown: %s", ic_name, exc) _delete_quietly( lambda: InferenceComponent.get(inference_component_name=ic_name), f"InferenceComponent {ic_name}", @@ -169,6 +181,24 @@ def _jumpstart_builder(): _delete_quietly(lambda: source_model, f"Model {src_model_name}") +def _wait_for_ic_terminal(ic_name, timeout_s=_IC_TIMEOUT_S): + """Poll an Inference Component until it leaves ``Creating``/``Updating``. + + Returns the resource in its terminal state (``InService`` / ``Failed``); + on timeout, returns the last-observed resource so the caller's assertion + reports the stuck status rather than this helper masking it. + """ + waited = 0 + ic = InferenceComponent.get(inference_component_name=ic_name) + while getattr(ic, "inference_component_status", None) in ("Creating", "Updating"): + if waited >= timeout_s: + break + time.sleep(20) + waited += 20 + ic = InferenceComponent.get(inference_component_name=ic_name) + return ic + + def _extract_s3_uri(container): """Pull the S3 artifact URI off a resolved container (ModelDataSource or the legacy ModelDataUrl), tolerating the resource-object attribute shape."""