fix: address review comments (iteration #4)

Author: sagemaker-bot

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

@@ -2702,6 +2702,52 @@ def _wait_for_endpoint(
 
         return desc
 
+    @staticmethod
+    def _apply_optional_ic_params(inference_component_spec, **kwargs):
+        """Apply optional IC-level parameters to an inference component spec dict.
+
+        Wires data_cache_config, base_inference_component_name, and container
+        into the given inference_component_spec dict.  Shared by
+        _deploy_core_endpoint and _update_inference_component to avoid
+        code duplication.
+
+        Args:
+            inference_component_spec (dict): The spec dict to mutate in-place.
+            **kwargs: May contain data_cache_config, base_inference_component_name,
+                and container.
+        """
+        from sagemaker.serve.model_builder_utils import _ModelBuilderUtils
+
+        ic_data_cache_config = kwargs.get("data_cache_config")
+        if ic_data_cache_config is not None:
+            resolved = _ModelBuilderUtils._resolve_data_cache_config(
+                None, ic_data_cache_config
+            )
+            if resolved is not None:
+                inference_component_spec["DataCacheConfig"] = {
+                    "EnableCaching": resolved.enable_caching
+                }
+
+        ic_base_component_name = kwargs.get("base_inference_component_name")
+        if ic_base_component_name is not None:
+            inference_component_spec["BaseInferenceComponentName"] = ic_base_component_name
+
+        ic_container = kwargs.get("container")
+        if ic_container is not None:
+            resolved_container = _ModelBuilderUtils._resolve_container_spec(
+                None, ic_container
+            )
+            if resolved_container is not None:
+                container_dict = {}
+                if resolved_container.image:
+                    container_dict["Image"] = resolved_container.image
+                if resolved_container.artifact_url:
+                    container_dict["ArtifactUrl"] = resolved_container.artifact_url
+                if resolved_container.environment:
+                    container_dict["Environment"] = resolved_container.environment
+                if container_dict:
+                    inference_component_spec["Container"] = container_dict
+
     def _deploy_core_endpoint(self, **kwargs):
         # Extract and update self parameters
         initial_instance_count = kwargs.get(
@@ -2930,31 +2976,7 @@ def _deploy_core_endpoint(self, **kwargs):
             }
 
             # Wire optional IC-level parameters into the specification
-            ic_data_cache_config = kwargs.get("data_cache_config")
-            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
-                    }
-
-            ic_base_component_name = kwargs.get("base_inference_component_name")
-            if ic_base_component_name is not None:
-                inference_component_spec["BaseInferenceComponentName"] = ic_base_component_name
-
-            ic_container = kwargs.get("container")
-            if ic_container is not None:
-                resolved_container = self._resolve_container_spec(ic_container)
-                if resolved_container is not None:
-                    container_dict = {}
-                    if resolved_container.image:
-                        container_dict["Image"] = resolved_container.image
-                    if resolved_container.artifact_url:
-                        container_dict["ArtifactUrl"] = resolved_container.artifact_url
-                    if resolved_container.environment:
-                        container_dict["Environment"] = resolved_container.environment
-                    if container_dict:
-                        inference_component_spec["Container"] = container_dict
+            self._apply_optional_ic_params(inference_component_spec, **kwargs)
 
             runtime_config = {"CopyCount": resources.copy_count}
             self.inference_component_name = (
@@ -3148,31 +3170,7 @@ def _update_inference_component(
         }
 
         # Wire optional IC-level parameters into the update specification
-        ic_data_cache_config = kwargs.get("data_cache_config")
-        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
-                }
-
-        ic_base_component_name = kwargs.get("base_inference_component_name")
-        if ic_base_component_name is not None:
-            inference_component_spec["BaseInferenceComponentName"] = ic_base_component_name
-
-        ic_container = kwargs.get("container")
-        if ic_container is not None:
-            resolved_container = self._resolve_container_spec(ic_container)
-            if resolved_container is not None:
-                container_dict = {}
-                if resolved_container.image:
-                    container_dict["Image"] = resolved_container.image
-                if resolved_container.artifact_url:
-                    container_dict["ArtifactUrl"] = resolved_container.artifact_url
-                if resolved_container.environment:
-                    container_dict["Environment"] = resolved_container.environment
-                if container_dict:
-                    inference_component_spec["Container"] = container_dict
+        self._apply_optional_ic_params(inference_component_spec, **kwargs)
 
         runtime_config = {"CopyCount": resource_requirements.copy_count}
 
@@ -4384,6 +4382,9 @@ def _deploy_model_customization(
         data_cache_config: Optional[Union["InferenceComponentDataCacheConfig", Dict[str, Any]]] = None,
         **kwargs,
     ) -> Endpoint:
+        # NOTE: For backward compatibility, model customization deployments
+        # default variant_name to endpoint_name (not "AllTraffic") when the
+        # caller does not provide an explicit value.
         """Deploy a model customization (fine-tuned) model to an endpoint with inference components.
 
         This method handles the special deployment flow for fine-tuned models, creating:
@@ -4423,8 +4424,9 @@ def _deploy_model_customization(
         # Fetch model package
         model_package = self._fetch_model_package()
 
-        # Resolve variant_name: use provided value or default to "AllTraffic"
-        effective_variant_name = variant_name or "AllTraffic"
+        # Resolve variant_name: preserve backward-compatible default of
+        # endpoint_name for model customization deployments.
+        effective_variant_name = variant_name or endpoint_name or "AllTraffic"
 
         # Resolve data_cache_config if provided
         resolved_data_cache_config = None

tests/unit/sagemaker/serve/test_resolve_ic_params.py

@@ -167,6 +167,82 @@ def test_invalid_type_list_raises(self, utils):
             utils._resolve_container_spec([{"image": "img"}])
 
 
+# ============================================================
+# Tests for _apply_optional_ic_params helper
+# ============================================================
+
+class TestApplyOptionalIcParams:
+    """Tests for the static helper that wires optional IC params into a spec dict."""
+
+    def test_no_params_no_mutation(self):
+        from sagemaker.serve.model_builder import ModelBuilder
+        spec = {"ModelName": "m"}
+        ModelBuilder._apply_optional_ic_params(spec)
+        assert "DataCacheConfig" not in spec
+        assert "BaseInferenceComponentName" not in spec
+        assert "Container" not in spec
+
+    def test_data_cache_config_dict(self):
+        from sagemaker.serve.model_builder import ModelBuilder
+        spec = {"ModelName": "m"}
+        ModelBuilder._apply_optional_ic_params(
+            spec, data_cache_config={"enable_caching": True}
+        )
+        assert spec["DataCacheConfig"] == {"EnableCaching": True}
+
+    def test_data_cache_config_typed(self):
+        from sagemaker.serve.model_builder import ModelBuilder
+        spec = {"ModelName": "m"}
+        cfg = InferenceComponentDataCacheConfig(enable_caching=False)
+        ModelBuilder._apply_optional_ic_params(spec, data_cache_config=cfg)
+        assert spec["DataCacheConfig"] == {"EnableCaching": False}
+
+    def test_base_inference_component_name(self):
+        from sagemaker.serve.model_builder import ModelBuilder
+        spec = {"ModelName": "m"}
+        ModelBuilder._apply_optional_ic_params(
+            spec, base_inference_component_name="base-ic"
+        )
+        assert spec["BaseInferenceComponentName"] == "base-ic"
+
+    def test_container_dict(self):
+        from sagemaker.serve.model_builder import ModelBuilder
+        spec = {"ModelName": "m"}
+        ModelBuilder._apply_optional_ic_params(
+            spec,
+            container={
+                "image": "img:latest",
+                "artifact_url": "s3://b/a",
+                "environment": {"K": "V"},
+            },
+        )
+        assert spec["Container"] == {
+            "Image": "img:latest",
+            "ArtifactUrl": "s3://b/a",
+            "Environment": {"K": "V"},
+        }
+
+    def test_container_typed(self):
+        from sagemaker.serve.model_builder import ModelBuilder
+        spec = {"ModelName": "m"}
+        c = InferenceComponentContainerSpecification(image="img")
+        ModelBuilder._apply_optional_ic_params(spec, container=c)
+        assert spec["Container"] == {"Image": "img"}
+
+    def test_all_params_together(self):
+        from sagemaker.serve.model_builder import ModelBuilder
+        spec = {"ModelName": "m"}
+        ModelBuilder._apply_optional_ic_params(
+            spec,
+            data_cache_config={"enable_caching": True},
+            base_inference_component_name="base",
+            container={"image": "img"},
+        )
+        assert spec["DataCacheConfig"] == {"EnableCaching": True}
+        assert spec["BaseInferenceComponentName"] == "base"
+        assert spec["Container"] == {"Image": "img"}
+
+
 # ============================================================
 # Tests for core wiring logic in _deploy_core_endpoint
 # ============================================================
@@ -419,5 +495,163 @@ def test_variant_name_passed_to_production_variant_on_new_endpoint(self, mock_en
             # Verify production_variant was called with variant_name="CustomVariant"
             mock_pv.assert_called_once()
             pv_kwargs = mock_pv.call_args
-            assert pv_kwargs.kwargs.get("variant_name") == "CustomVariant" or \
-                (len(pv_kwargs.args) > 0 and False)  # variant_name is always a kwarg
+            assert pv_kwargs.kwargs.get("variant_name") == "CustomVariant"
+
+
+# ============================================================
+# Tests for _update_inference_component wiring
+# ============================================================
+
+class TestUpdateInferenceComponentWiring:
+    """Tests that _update_inference_component correctly wires optional IC params."""
+
+    def _make_model_builder(self):
+        from sagemaker.serve.model_builder import ModelBuilder
+
+        mb = object.__new__(ModelBuilder)
+        mb.model_name = "test-model"
+        mb.sagemaker_session = MagicMock()
+        return mb
+
+    def test_update_ic_with_data_cache_config(self):
+        mb = self._make_model_builder()
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._update_inference_component(
+            "my-ic", resources, data_cache_config={"enable_caching": True}
+        )
+
+        call_kwargs = mb.sagemaker_session.update_inference_component.call_args
+        spec = call_kwargs.kwargs["specification"]
+        assert spec["DataCacheConfig"] == {"EnableCaching": True}
+
+    def test_update_ic_with_container(self):
+        mb = self._make_model_builder()
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._update_inference_component(
+            "my-ic", resources, container={"image": "img:v1"}
+        )
+
+        call_kwargs = mb.sagemaker_session.update_inference_component.call_args
+        spec = call_kwargs.kwargs["specification"]
+        assert spec["Container"] == {"Image": "img:v1"}
+
+    def test_update_ic_with_base_inference_component_name(self):
+        mb = self._make_model_builder()
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._update_inference_component(
+            "my-ic", resources, base_inference_component_name="base-ic"
+        )
+
+        call_kwargs = mb.sagemaker_session.update_inference_component.call_args
+        spec = call_kwargs.kwargs["specification"]
+        assert spec["BaseInferenceComponentName"] == "base-ic"
+
+    def test_update_ic_no_optional_params(self):
+        mb = self._make_model_builder()
+        from sagemaker.core.inference_config import ResourceRequirements
+        resources = ResourceRequirements(
+            requests={"memory": 8192, "num_accelerators": 1, "num_cpus": 2, "copies": 1}
+        )
+
+        mb._update_inference_component("my-ic", resources)
+
+        call_kwargs = mb.sagemaker_session.update_inference_component.call_args
+        spec = call_kwargs.kwargs["specification"]
+        assert "DataCacheConfig" not in spec
+        assert "BaseInferenceComponentName" not in spec
+        assert "Container" not in spec
+
+
+# ============================================================
+# Tests for deploy() parameter forwarding
+# ============================================================
+
+class TestDeployParameterForwarding:
+    """Tests that deploy() correctly forwards new IC params into kwargs."""
+
+    def test_deploy_forwards_variant_name_to_kwargs(self):
+        """deploy() should set kwargs['variant_name'] to the provided value."""
+        from sagemaker.serve.model_builder import ModelBuilder
+
+        mb = object.__new__(ModelBuilder)
+        mb.built_model = MagicMock()
+        mb._deployed = False
+        mb._is_sharded_model = False
+        mb.model_name = "test"
+        mb.instance_type = "ml.m5.large"
+        mb.endpoint_name = None
+        mb.mode = None
+        mb.model_server = None
+
+        # Mock _is_model_customization to return False
+        mb._is_model_customization = MagicMock(return_value=False)
+        # Mock _deploy to capture kwargs
+        captured = {}
+
+        def fake_deploy(**kw):
+            captured.update(kw)
+            return MagicMock()
+
+        mb._deploy = fake_deploy
+
+        mb.deploy(
+            endpoint_name="ep",
+            instance_type="ml.m5.large",
+            initial_instance_count=1,
+            variant_name="MyVariant",
+            data_cache_config={"enable_caching": True},
+            base_inference_component_name="base-ic",
+            container={"image": "img"},
+        )
+
+        assert captured["variant_name"] == "MyVariant"
+        assert captured["data_cache_config"] == {"enable_caching": True}
+        assert captured["base_inference_component_name"] == "base-ic"
+        assert captured["container"] == {"image": "img"}
+
+    def test_deploy_defaults_variant_name_to_all_traffic(self):
+        """deploy() should default variant_name to 'AllTraffic' when not provided."""
+        from sagemaker.serve.model_builder import ModelBuilder
+
+        mb = object.__new__(ModelBuilder)
+        mb.built_model = MagicMock()
+        mb._deployed = False
+        mb._is_sharded_model = False
+        mb.model_name = "test"
+        mb.instance_type = "ml.m5.large"
+        mb.endpoint_name = None
+        mb.mode = None
+        mb.model_server = None
+        mb._is_model_customization = MagicMock(return_value=False)
+
+        captured = {}
+
+        def fake_deploy(**kw):
+            captured.update(kw)
+            return MagicMock()
+
+        mb._deploy = fake_deploy
+
+        mb.deploy(
+            endpoint_name="ep",
+            instance_type="ml.m5.large",
+            initial_instance_count=1,
+        )
+
+        assert captured["variant_name"] == "AllTraffic"
+        # Optional params should not be in kwargs when not provided
+        assert "data_cache_config" not in captured
+        assert "base_inference_component_name" not in captured
+        assert "container" not in captured