fix: address review comments (iteration #1)

Author: sagemaker-bot

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

@@ -4211,8 +4211,12 @@ def deploy(
         if not hasattr(self, "built_model") and not hasattr(self, "_deployables"):
             raise ValueError("Model needs to be built before deploying")
 
-        # Centralize variant_name defaulting and always forward IC-level params
-        kwargs["variant_name"] = variant_name or "AllTraffic"
+        # Only forward variant_name when explicitly provided by the caller.
+        # Each downstream path has its own default:
+        #   - _deploy_core_endpoint defaults to "AllTraffic"
+        #   - _deploy_model_customization defaults to endpoint_name
+        if variant_name is not None:
+            kwargs["variant_name"] = variant_name
         if inference_component_name is not None:
             kwargs["inference_component_name"] = inference_component_name
         if data_cache_config is not None:

sagemaker-serve/tests/integ/test_ic_deploy_params_integration.py

@@ -37,9 +37,6 @@
 # Use the same JumpStart model as test_jumpstart_integration.py
 MODEL_ID = "huggingface-llm-falcon-7b-bf16"
 
-# Training job for model customization path (same as test_model_customization_deployment.py)
-TRAINING_JOB_NAME = "meta-textgeneration-llama-3-2-1b-instruct-sft-20251201172445"
-
 
 def _cleanup_endpoint(endpoint_name, sagemaker_client):
     """Delete endpoint, endpoint config, and all inference components."""
@@ -158,85 +155,4 @@ def test_deploy_with_data_cache_config_and_variant_name_via_ic_path():
         _cleanup_model(model_name, sagemaker_client)
 
 
-@pytest.mark.slow_test
-def test_deploy_with_data_cache_config_via_model_customization_path():
-    """Deploy a fine-tuned model via _deploy_model_customization with data_cache_config.
-
-    Verifies:
-    - The IC was created with DataCacheConfig.EnableCaching == True
-    - The variant_name defaults to endpoint_name (backward compat) when not explicitly provided
-    """
-    from sagemaker.core.resources import TrainingJob
-
-    unique_id = uuid.uuid4().hex[:8]
-    model_name = f"ic-mc-test-model-{unique_id}"
-    endpoint_name = f"ic-mc-test-ep-{unique_id}"
-
-    sagemaker_client = boto3.client("sagemaker")
-
-    try:
-        training_job = TrainingJob.get(training_job_name=TRAINING_JOB_NAME)
-        model_builder = ModelBuilder(
-            model=training_job, instance_type="ml.g5.4xlarge"
-        )
-        model_builder.accept_eula = True
-        core_model = model_builder.build(model_name=model_name)
-        logger.info("Model created: %s", core_model.model_name)
-
-        # Deploy with data_cache_config but WITHOUT explicit variant_name
-        # so it should default to endpoint_name for model customization path
-        endpoint = model_builder.deploy(
-            endpoint_name=endpoint_name,
-            initial_instance_count=1,
-            data_cache_config={"enable_caching": True},
-        )
-        logger.info("Endpoint created: %s", endpoint.endpoint_name)
-
-        # Find inference components on this endpoint
-        paginator = sagemaker_client.get_paginator("list_inference_components")
-        ic_names = []
-        for page in paginator.paginate(EndpointNameEquals=endpoint_name):
-            for ic in page.get("InferenceComponents", []):
-                ic_names.append(ic["InferenceComponentName"])
-
-        assert len(ic_names) > 0, (
-            f"Expected at least one inference component on endpoint '{endpoint_name}'"
-        )
-
-        # Check the first (or base) IC for DataCacheConfig
-        # For LORA, the base IC should have data_cache_config; for non-LORA, the single IC.
-        peft_type = model_builder._fetch_peft()
-        if peft_type == "LORA":
-            # Base IC is named <endpoint_name>-inference-component
-            base_ic_name = f"{endpoint_name}-inference-component"
-        else:
-            base_ic_name = f"{endpoint_name}-inference-component"
-
-        ic_desc = sagemaker_client.describe_inference_component(
-            InferenceComponentName=base_ic_name
-        )
 
-        # Verify DataCacheConfig.EnableCaching == True
-        spec = ic_desc.get("Specification", {})
-        data_cache = spec.get("DataCacheConfig", {})
-        assert data_cache.get("EnableCaching") is True, (
-            f"Expected DataCacheConfig.EnableCaching=True, got {data_cache}"
-        )
-
-        # Verify variant_name defaults to endpoint_name (backward compat)
-        actual_variant = ic_desc.get("VariantName")
-        assert actual_variant == endpoint_name, (
-            f"Expected VariantName='{endpoint_name}' (backward compat default), "
-            f"got '{actual_variant}'"
-        )
-
-        logger.info(
-            "Test passed: IC '%s' has DataCacheConfig.EnableCaching=True "
-            "and VariantName='%s' (backward compat default)",
-            base_ic_name,
-            endpoint_name,
-        )
-
-    finally:
-        _cleanup_endpoint(endpoint_name, sagemaker_client)
-        _cleanup_model(model_name, sagemaker_client)

tests/unit/sagemaker/serve/test_resolve_ic_params.py

@@ -621,8 +621,13 @@ def fake_deploy(**kw):
         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."""
+    def test_deploy_does_not_set_variant_name_when_not_provided(self):
+        """deploy() should NOT set variant_name in kwargs when not provided.
+
+        This allows downstream methods to use their own defaults:
+        - _deploy_core_endpoint defaults to 'AllTraffic'
+        - _deploy_model_customization defaults to endpoint_name
+        """
         from sagemaker.serve.model_builder import ModelBuilder
 
         mb = object.__new__(ModelBuilder)
@@ -650,8 +655,45 @@ def fake_deploy(**kw):
             initial_instance_count=1,
         )
 
-        assert captured["variant_name"] == "AllTraffic"
+        # variant_name should NOT be in kwargs when not explicitly provided
+        assert "variant_name" not in captured
         # 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
+
+    def test_deploy_forwards_variant_name_none_is_not_forwarded(self):
+        """deploy(variant_name=None) should NOT forward variant_name.
+
+        None is the default, so it should behave the same as not providing it.
+        """
+        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,
+            variant_name=None,
+        )
+
+        # variant_name=None should not be forwarded
+        assert "variant_name" not in captured