From a322ca131a3f9e40de808bbae65b2bb81164f4bc Mon Sep 17 00:00:00 2001 From: ZealSV Date: Fri, 24 Jul 2026 12:32:27 -0700 Subject: [PATCH] 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 | 72 ++++-------------- 2 files changed, 24 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..bb35726a89 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,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"