fix(serve): remove IC data-source collapse hack from recommendation deploy

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
@@ -5158,12 +5158,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
@@ -5275,65 +5272,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,

sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py

@@ -0,0 +1,251 @@
+# 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``."""
+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. 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; 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):
+    """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}"
+
+    from sagemaker.core.jumpstart.configs import JumpStartConfig
+
+    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 = 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)
+
+        # 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)
+
+        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)
+
+        # 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.
+        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:
+        # 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}",
+            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 _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."""
+    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)