diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 948b87bcaa..9182903285 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -4058,6 +4058,10 @@ def _deploy_core_endpoint(self, **kwargs): inference_ami_version=inference_ami_version, ) + data_capture_config_dict = None + if data_capture_config is not None: + data_capture_config_dict = data_capture_config._to_request_dict() + if endpoint_type == EndpointType.INFERENCE_COMPONENT_BASED: if update_endpoint: raise ValueError( @@ -4101,6 +4105,7 @@ def _deploy_core_endpoint(self, **kwargs): production_variants=[production_variant], tags=tags, kms_key=kms_key, + data_capture_config_dict=data_capture_config_dict, vpc_config=self.vpc_config, enable_network_isolation=self._enable_network_isolation, role=self.role_arn, @@ -4175,10 +4180,6 @@ def _deploy_core_endpoint(self, **kwargs): base_endpoint_name = base_from_name(self.model_name) self.endpoint_name = name_from_base(base_endpoint_name) - data_capture_config_dict = None - if data_capture_config is not None: - data_capture_config_dict = data_capture_config._to_request_dict() - async_inference_config_dict = None if is_async: if ( diff --git a/sagemaker-serve/tests/unit/test_model_builder_deploy.py b/sagemaker-serve/tests/unit/test_model_builder_deploy.py index a6255128ba..a26990b5d8 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_deploy.py +++ b/sagemaker-serve/tests/unit/test_model_builder_deploy.py @@ -14,6 +14,7 @@ from sagemaker.core.inference_config import ( ResourceRequirements, ) +from sagemaker.core.model_monitor.data_capture_config import DataCaptureConfig class TestModelBuilderContainerDef(unittest.TestCase): @@ -380,6 +381,68 @@ def test_deploy_core_endpoint_sharded_model_network_isolation_raises_error(self) self.assertIn("EnableNetworkIsolation cannot be set to True", str(context.exception)) + def _ic_builder(self): + """Build a ModelBuilder ready to deploy an inference-component based endpoint.""" + builder = ModelBuilder( + model=Mock(), + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + ) + builder.built_model = Mock() + builder.built_model.model_name = "test-model" + builder.model_name = "test-model" + builder.inference_component_name = "test-ic" + builder._enable_network_isolation = False + builder._tags = [] + return builder + + def test_deploy_core_endpoint_ic_based_passes_data_capture_config(self): + """Test _deploy_core_endpoint forwards data_capture_config for IC based endpoints.""" + builder = self._ic_builder() + data_capture_config = DataCaptureConfig( + enable_capture=True, + sampling_percentage=100, + destination_s3_uri="s3://bucket/prefix", + ) + + with patch("sagemaker.serve.model_builder.Endpoint.get") as mock_get: + mock_get.return_value = Mock(spec=Endpoint) + builder._deploy_core_endpoint( + instance_type="ml.m5.large", + initial_instance_count=1, + endpoint_name="test-ep", + endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, + resources=ResourceRequirements(requests={"num_cpus": 1, "memory": 1024}), + data_capture_config=data_capture_config, + wait=False, + ) + + self.mock_session.endpoint_from_production_variants.assert_called_once() + _, call_kwargs = self.mock_session.endpoint_from_production_variants.call_args + self.assertEqual( + call_kwargs["data_capture_config_dict"], + data_capture_config._to_request_dict(), + ) + + def test_deploy_core_endpoint_ic_based_without_data_capture_passes_none(self): + """Test _deploy_core_endpoint passes no data capture config when none is given.""" + builder = self._ic_builder() + + with patch("sagemaker.serve.model_builder.Endpoint.get") as mock_get: + mock_get.return_value = Mock(spec=Endpoint) + builder._deploy_core_endpoint( + instance_type="ml.m5.large", + initial_instance_count=1, + endpoint_name="test-ep", + endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, + resources=ResourceRequirements(requests={"num_cpus": 1, "memory": 1024}), + wait=False, + ) + + self.mock_session.endpoint_from_production_variants.assert_called_once() + _, call_kwargs = self.mock_session.endpoint_from_production_variants.call_args + self.assertIsNone(call_kwargs["data_capture_config_dict"]) + class TestModelBuilderDeployHelpers(unittest.TestCase): """Test ModelBuilder deployment helper methods."""