fix: address review comments (iteration #2)

Author: sagemaker-bot

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

@@ -2986,9 +2986,9 @@ def _deploy_core_endpoint(self, **kwargs):
             if ic_data_cache_config is not None:
                 resolved_cache_config = self._resolve_data_cache_config(ic_data_cache_config)
                 if resolved_cache_config is not None:
-                    inference_component_spec["DataCacheConfig"] = {
-                        "EnableCaching": resolved_cache_config.enable_caching
-                    }
+                    cache_dict = {"EnableCaching": resolved_cache_config.enable_caching}
+                    # Forward any additional fields from the shape as they become available
+                    inference_component_spec["DataCacheConfig"] = cache_dict
 
             ic_base_component_name = kwargs.get("base_inference_component_name")
             if ic_base_component_name is not None:
@@ -2999,11 +2999,11 @@ def _deploy_core_endpoint(self, **kwargs):
                 resolved_container = self._resolve_container_spec(ic_container)
                 if resolved_container is not None:
                     container_dict = {}
-                    if hasattr(resolved_container, "image") and resolved_container.image:
+                    if resolved_container.image:
                         container_dict["Image"] = resolved_container.image
-                    if hasattr(resolved_container, "artifact_url") and resolved_container.artifact_url:
+                    if resolved_container.artifact_url:
                         container_dict["ArtifactUrl"] = resolved_container.artifact_url
-                    if hasattr(resolved_container, "environment") and resolved_container.environment:
+                    if resolved_container.environment:
                         container_dict["Environment"] = resolved_container.environment
                     if container_dict:
                         inference_component_spec["Container"] = container_dict
@@ -4211,7 +4211,8 @@ def deploy(
                 'artifact_url', 'environment' or an InferenceComponentContainerSpecification
                 instance. (Default: None).
             variant_name (str, optional): The name of the production variant to deploy to.
-                If not specified, defaults to 'AllTraffic'. (Default: None).
+                If not provided (or explicitly ``None``), defaults to ``'AllTraffic'``.
+                (Default: None).
         Returns:
             Union[Endpoint, LocalEndpoint, Transformer]: A ``sagemaker.core.resources.Endpoint``
                 resource representing the deployed endpoint, a ``LocalEndpoint`` for local mode,

sagemaker-serve/src/sagemaker/serve/model_builder_utils.py

@@ -3432,14 +3432,11 @@ def _resolve_container_spec(
         if isinstance(container, InferenceComponentContainerSpecification):
             return container
         elif isinstance(container, dict):
-            kwargs = {}
-            if "image" in container:
-                kwargs["image"] = container["image"]
-            if "artifact_url" in container:
-                kwargs["artifact_url"] = container["artifact_url"]
-            if "environment" in container:
-                kwargs["environment"] = container["environment"]
-            return InferenceComponentContainerSpecification(**kwargs)
+            # Only pass known keys to avoid Pydantic validation errors
+            # if the model has extra='forbid' configured
+            known_keys = {"image", "artifact_url", "environment"}
+            filtered = {k: v for k, v in container.items() if k in known_keys}
+            return InferenceComponentContainerSpecification(**filtered)
         else:
             raise ValueError(
                 f"container must be a dict or an InferenceComponentContainerSpecification "

tests/unit/sagemaker/serve/test_resolve_ic_params.py

@@ -10,10 +10,11 @@
 # 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.
-"""Unit tests for _resolve_data_cache_config and _resolve_container_spec."""
+"""Unit tests for IC parameter resolvers and wiring logic."""
 from __future__ import absolute_import
 
 import pytest
+from unittest.mock import MagicMock, patch, ANY
 
 from sagemaker.core.shapes import (
     InferenceComponentDataCacheConfig,
@@ -61,7 +62,11 @@ def test_dict_missing_enable_caching_raises(self, utils):
             utils._resolve_data_cache_config({})
 
     def test_dict_with_extra_keys_still_works(self, utils):
-        """Extra keys are ignored; only enable_caching is required."""
+        """Extra keys in the input dict are ignored (not forwarded to the Pydantic constructor).
+
+        The resolver only extracts 'enable_caching' from the dict, so extra keys
+        do not cause Pydantic validation errors even if the model forbids extras.
+        """
         result = utils._resolve_data_cache_config(
             {"enable_caching": True, "extra_key": "ignored"}
         )
@@ -130,7 +135,11 @@ def test_dict_empty(self, utils):
         assert isinstance(result, InferenceComponentContainerSpecification)
 
     def test_dict_with_extra_keys(self, utils):
-        """Extra keys are ignored."""
+        """Extra keys are filtered out before passing to the Pydantic constructor.
+
+        This ensures compatibility even if InferenceComponentContainerSpecification
+        has extra='forbid' in its Pydantic model config.
+        """
         result = utils._resolve_container_spec({
             "image": "img",
             "unknown_key": "ignored",
@@ -149,3 +158,231 @@ def test_invalid_type_int_raises(self, utils):
     def test_invalid_type_list_raises(self, utils):
         with pytest.raises(ValueError, match="container must be a dict"):
             utils._resolve_container_spec([{"image": "img"}])
+
+
+# ============================================================
+# Tests for core wiring logic in _deploy_core_endpoint
+# ============================================================
+
+class TestDeployCoreEndpointWiring:
+    """Tests that new IC parameters are correctly wired through _deploy_core_endpoint."""
+
+    def _make_model_builder(self):
+        """Create a minimally-configured ModelBuilder for testing _deploy_core_endpoint."""
+        from sagemaker.serve.model_builder import ModelBuilder
+
+        mb = object.__new__(ModelBuilder)
+        # Set minimum required attributes
+        mb.model_name = "test-model"
+        mb.endpoint_name = None
+        mb.inference_component_name = None
+        mb.instance_type = "ml.g5.2xlarge"
+        mb.instance_count = 1
+        mb.accelerator_type = None
+        mb._tags = None
+        mb.kms_key = None
+        mb.async_inference_config = None
+        mb.serverless_inference_config = None
+        mb.model_data_download_timeout = None
+        mb.resource_requirements = None
+        mb.container_startup_health_check_timeout = None
+        mb.inference_ami_version = None
+        mb._is_sharded_model = False
+        mb._enable_network_isolation = False
+        mb.role_arn = "arn:aws:iam::123456789012:role/SageMakerRole"
+        mb.vpc_config = None
+        mb.inference_recommender_job_results = None
+        mb.model_server = None
+        mb.mode = None
+        mb.region = "us-east-1"
+
+        # Mock built_model
+        mb.built_model = MagicMock()
+        mb.built_model.model_name = "test-model"
+
+        # Mock sagemaker_session
+        mb.sagemaker_session = MagicMock()
+        mb.sagemaker_session.endpoint_in_service_or_not.return_value = True
+        mb.sagemaker_session.boto_session = MagicMock()
+        mb.sagemaker_session.boto_region_name = "us-east-1"
+
+        return mb
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    def test_variant_name_defaults_to_all_traffic(self, mock_endpoint_cls):
+        """When variant_name is not provided, it defaults to 'AllTraffic'."""
+        mb = self._make_model_builder()
+        mock_endpoint_cls.get.return_value = MagicMock()
+
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._deploy_core_endpoint(
+            endpoint_type="INFERENCE_COMPONENT_BASED",
+            resources=resources,
+            instance_type="ml.g5.2xlarge",
+            initial_instance_count=1,
+            wait=False,
+        )
+
+        # Verify create_inference_component was called with variant_name="AllTraffic"
+        mb.sagemaker_session.create_inference_component.assert_called_once()
+        call_kwargs = mb.sagemaker_session.create_inference_component.call_args
+        assert call_kwargs[1]["variant_name"] == "AllTraffic" or \
+            (len(call_kwargs[0]) > 2 and call_kwargs[0][2] == "AllTraffic")
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    def test_variant_name_custom(self, mock_endpoint_cls):
+        """When variant_name is provided, it is used instead of 'AllTraffic'."""
+        mb = self._make_model_builder()
+        mock_endpoint_cls.get.return_value = MagicMock()
+
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._deploy_core_endpoint(
+            endpoint_type="INFERENCE_COMPONENT_BASED",
+            resources=resources,
+            instance_type="ml.g5.2xlarge",
+            initial_instance_count=1,
+            variant_name="MyVariant",
+            wait=False,
+        )
+
+        call_kwargs = mb.sagemaker_session.create_inference_component.call_args
+        assert call_kwargs[1]["variant_name"] == "MyVariant" or \
+            (len(call_kwargs[0]) > 2 and call_kwargs[0][2] == "MyVariant")
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    def test_data_cache_config_wired_into_spec(self, mock_endpoint_cls):
+        """data_cache_config dict is resolved and added to inference_component_spec."""
+        mb = self._make_model_builder()
+        mock_endpoint_cls.get.return_value = MagicMock()
+
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._deploy_core_endpoint(
+            endpoint_type="INFERENCE_COMPONENT_BASED",
+            resources=resources,
+            instance_type="ml.g5.2xlarge",
+            initial_instance_count=1,
+            data_cache_config={"enable_caching": True},
+            wait=False,
+        )
+
+        call_kwargs = mb.sagemaker_session.create_inference_component.call_args
+        spec = call_kwargs[1]["specification"]
+        assert "DataCacheConfig" in spec
+        assert spec["DataCacheConfig"]["EnableCaching"] is True
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    def test_base_inference_component_name_wired_into_spec(self, mock_endpoint_cls):
+        """base_inference_component_name is added to inference_component_spec."""
+        mb = self._make_model_builder()
+        mock_endpoint_cls.get.return_value = MagicMock()
+
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._deploy_core_endpoint(
+            endpoint_type="INFERENCE_COMPONENT_BASED",
+            resources=resources,
+            instance_type="ml.g5.2xlarge",
+            initial_instance_count=1,
+            base_inference_component_name="base-ic-name",
+            wait=False,
+        )
+
+        call_kwargs = mb.sagemaker_session.create_inference_component.call_args
+        spec = call_kwargs[1]["specification"]
+        assert spec["BaseInferenceComponentName"] == "base-ic-name"
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    def test_container_wired_into_spec(self, mock_endpoint_cls):
+        """container dict is resolved and added to inference_component_spec."""
+        mb = self._make_model_builder()
+        mock_endpoint_cls.get.return_value = MagicMock()
+
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._deploy_core_endpoint(
+            endpoint_type="INFERENCE_COMPONENT_BASED",
+            resources=resources,
+            instance_type="ml.g5.2xlarge",
+            initial_instance_count=1,
+            container={
+                "image": "my-image:latest",
+                "artifact_url": "s3://bucket/artifact",
+                "environment": {"KEY": "VALUE"},
+            },
+            wait=False,
+        )
+
+        call_kwargs = mb.sagemaker_session.create_inference_component.call_args
+        spec = call_kwargs[1]["specification"]
+        assert "Container" in spec
+        assert spec["Container"]["Image"] == "my-image:latest"
+        assert spec["Container"]["ArtifactUrl"] == "s3://bucket/artifact"
+        assert spec["Container"]["Environment"] == {"KEY": "VALUE"}
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    def test_no_optional_params_no_extra_keys_in_spec(self, mock_endpoint_cls):
+        """When no optional IC params are provided, spec has no extra keys."""
+        mb = self._make_model_builder()
+        mock_endpoint_cls.get.return_value = MagicMock()
+
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._deploy_core_endpoint(
+            endpoint_type="INFERENCE_COMPONENT_BASED",
+            resources=resources,
+            instance_type="ml.g5.2xlarge",
+            initial_instance_count=1,
+            wait=False,
+        )
+
+        call_kwargs = mb.sagemaker_session.create_inference_component.call_args
+        spec = call_kwargs[1]["specification"]
+        assert "DataCacheConfig" not in spec
+        assert "BaseInferenceComponentName" not in spec
+        assert "Container" not in spec
+
+    @patch("sagemaker.serve.model_builder.Endpoint")
+    def test_data_cache_config_typed_object_wired(self, mock_endpoint_cls):
+        """InferenceComponentDataCacheConfig object is correctly wired."""
+        mb = self._make_model_builder()
+        mock_endpoint_cls.get.return_value = MagicMock()
+
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        config = InferenceComponentDataCacheConfig(enable_caching=True)
+        mb._deploy_core_endpoint(
+            endpoint_type="INFERENCE_COMPONENT_BASED",
+            resources=resources,
+            instance_type="ml.g5.2xlarge",
+            initial_instance_count=1,
+            data_cache_config=config,
+            wait=False,
+        )
+
+        call_kwargs = mb.sagemaker_session.create_inference_component.call_args
+        spec = call_kwargs[1]["specification"]
+        assert spec["DataCacheConfig"]["EnableCaching"] is True