feature: adding support for DataCaptureConfig in endpoint config (#26)

Author: brightsparc

src/stepfunctions/steps/sagemaker.py

@@ -19,7 +19,7 @@
 
 from sagemaker.workflow.airflow import training_config, transform_config, model_config, tuning_config
 from sagemaker.model import Model, FrameworkModel
-
+from sagemaker.model_monitor import DataCaptureConfig
 
 class TrainingStep(Task):
 
@@ -209,14 +209,17 @@ class EndpointConfigStep(Task):
     Creates a Task State to `create an endpoint configuration in SageMaker <https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateEndpointConfig.html>`_.
     """
 
-    def __init__(self, state_id, endpoint_config_name, model_name, initial_instance_count, instance_type, tags=None, **kwargs):
+    def __init__(self, state_id, endpoint_config_name, model_name, initial_instance_count, instance_type, data_capture_config=None, tags=None, **kwargs):
         """
         Args:
             state_id (str): State name whose length **must be** less than or equal to 128 unicode characters. State names **must be** unique within the scope of the whole state machine.
             endpoint_config_name (str or Placeholder): The name of the endpoint configuration to create. We recommend to use :py:class:`~stepfunctions.inputs.ExecutionInput` placeholder collection to pass the value dynamically in each execution.
             model_name (str or Placeholder): The name of the SageMaker model to attach to the endpoint configuration. We recommend to use :py:class:`~stepfunctions.inputs.ExecutionInput` placeholder collection to pass the value dynamically in each execution.
             initial_instance_count (int or Placeholder): The initial number of instances to run in the ``Endpoint`` created from this ``Model``.
             instance_type (str or Placeholder): The EC2 instance type to deploy this Model to. For example, 'ml.p2.xlarge'.
+            data_capture_config (sagemaker.model_monitor.DataCaptureConfig, optional): Specifies
+                configuration related to Endpoint data capture for use with
+                Amazon SageMaker Model Monitoring. Default: None.
             tags (list[dict], optional): `List to tags <https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html>`_ to associate with the resource.
         """
         parameters = {
@@ -229,6 +232,9 @@ def __init__(self, state_id, endpoint_config_name, model_name, initial_instance_
             }]
         }
 
+        if isinstance(data_capture_config, DataCaptureConfig):
+            parameters['DataCaptureConfig'] = data_capture_config._to_request_dict()
+            
         if tags:
             parameters['Tags'] = tags_dict_to_kv_list(tags)
 

tests/unit/test_sagemaker_steps.py

@@ -20,6 +20,7 @@
 from sagemaker.model import Model
 from sagemaker.tensorflow import TensorFlow
 from sagemaker.pipeline import PipelineModel
+from sagemaker.model_monitor import DataCaptureConfig
 
 from unittest.mock import MagicMock, patch
 from stepfunctions.steps.sagemaker import TrainingStep, TransformStep, ModelStep, EndpointStep, EndpointConfigStep
@@ -366,7 +367,16 @@ def test_model_step_creation(pca_model):
     }
 
 def test_endpoint_config_step_creation(pca_model):
-    step = EndpointConfigStep('Endpoint Config', endpoint_config_name='MyEndpointConfig', model_name='pca-model', initial_instance_count=1, instance_type='ml.p2.xlarge')
+    data_capture_config = DataCaptureConfig(
+        enable_capture=True,
+        sampling_percentage=100,
+        destination_s3_uri='s3://sagemaker/datacapture')
+    step = EndpointConfigStep('Endpoint Config', 
+        endpoint_config_name='MyEndpointConfig', 
+        model_name='pca-model', 
+        initial_instance_count=1, 
+        instance_type='ml.p2.xlarge',
+        data_capture_config=data_capture_config)
     assert step.to_dict() == {
         'Type': 'Task',
         'Parameters': {
@@ -376,7 +386,20 @@ def test_endpoint_config_step_creation(pca_model):
                 'InstanceType': 'ml.p2.xlarge',
                 'ModelName': 'pca-model',
                 'VariantName': 'AllTraffic'
-            }]
+            }],
+            'DataCaptureConfig': {
+                'EnableCapture': True,
+                'InitialSamplingPercentage': 100,
+                'DestinationS3Uri': 's3://sagemaker/datacapture',
+                'CaptureOptions': [
+                    {'CaptureMode': 'Input'}, 
+                    {'CaptureMode': 'Output'}
+                ],
+                'CaptureContentTypeHeader': {
+                    'CsvContentTypes': ['text/csv'],
+                    'JsonContentTypes': ['application/json']
+                }
+            }
         },
         'Resource': 'arn:aws:states:::sagemaker:createEndpointConfig',
         'End': True