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.

Author: ZealSV

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,

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,16 @@ 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.
         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 +313,10 @@ 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"
+        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 +326,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"