From d5072939bf812ea5fbbec895dfb8d601570a14b6 Mon Sep 17 00:00:00 2001 From: Ryan Steakley Date: Wed, 23 Sep 2026 05:23:12 +0000 Subject: [PATCH] fix(serve): pass data_capture_config when creating inference-component endpoints ModelBuilder.deploy(data_capture_config=...) was silently ignored for INFERENCE_COMPONENT_BASED endpoints. _deploy_core_endpoint only built data_capture_config_dict inside the model-based else branch and only passed it to create_endpoint_config()/endpoint_from_production_variants() there. On the inference-component path the dict was never built and never passed, so the endpoint was created without a DataCaptureConfig and no error was raised - data capture simply did not happen. The dict is now computed once, before the endpoint_type branch, and passed to endpoint_from_production_variants() on both branches. The model-based branch keeps passing exactly the value it passed before; no other behavior changes. Adds two unit tests to sagemaker-serve/tests/unit/test_model_builder_deploy.py: test_deploy_core_endpoint_ic_based_passes_data_capture_config asserts the data_capture_config_dict kwarg equals DataCaptureConfig._to_request_dict(), and test_deploy_core_endpoint_ic_based_without_data_capture_passes_none asserts it is None when no config is supplied. --- .../src/sagemaker/serve/model_builder.py | 9 +-- .../tests/unit/test_model_builder_deploy.py | 63 +++++++++++++++++++ 2 files changed, 68 insertions(+), 4 deletions(-) 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."""