Repository: mlflow/mlflow

Issue:
[BUG] When performing Sagemaker update the Data Capture setting are gone from a Sagemaker Endpoint
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

- Client: 1.29.0
- Tracking server: 1.x.y


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **Python version**:
- **yarn version, if running the dev UI**:


### Describe the problem

When I create an endpoint using create_deployment call and supply the following data capture config `data_capture_config = {
    'EnableCapture': True,
    'InitialSamplingPercentage': data_capture_InitialSamplingPercentage,
    'DestinationS3Uri': data_capture_bucket,
    'CaptureOptions': [
        {'CaptureMode': 'Input'},
        {'CaptureMode': 'Output'}
    ]
}`
I do see the data capture setting in AWS Console. But when I execute the update_deployment call against the same endpoint using the same config, the data capture settings are gone from the console. 


### Tracking information

_No response_

### Code to reproduce issue

You can use any test configuration
```
if endpoint_exists:
    config['mode']=mlflow.sagemaker.DEPLOYMENT_MODE_REPLACE
    config['data_capture_config']=data_capture_config
    client.update_deployment(
    deployment_name,
    flavor="python_function",
    model_uri=model_uri,
    config=config)
else:
    client.create_deployment(
    deployment_name,
    flavor="python_function",
    model_uri=model_uri,
    config=config
)
```

### Stack trace

none

### Other info / logs

_No response_

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [X] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [X] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
Hi @edvorkin thank you for reporting. We're currently investigating this functionality and will look into this. Thank you for reporting!

Resolving patch:
diff --git a/mlflow/sagemaker/__init__.py b/mlflow/sagemaker/__init__.py
--- a/mlflow/sagemaker/__init__.py
+++ b/mlflow/sagemaker/__init__.py
@@ -403,6 +403,7 @@ def _deploy(
             sage_client=sage_client,
             s3_client=s3_client,
             variant_name=variant_name,
+            data_capture_config=data_capture_config,
         )
     else:
         deployment_operation = _create_sagemaker_endpoint(
@@ -1578,6 +1579,7 @@ def _update_sagemaker_endpoint(
     sage_client,
     s3_client,
     variant_name=None,
+    data_capture_config=None,
 ):
     """
     :param endpoint_name: The name of the SageMaker endpoint to update.
@@ -1597,6 +1599,9 @@ def _update_sagemaker_endpoint(
     :param sage_client: A boto3 client for SageMaker.
     :param s3_client: A boto3 client for S3.
     :variant_name: The name to assign to the new production variant if it doesn't already exist.
+    :param: data_capture_config: A dictionary specifying the data capture configuration to use.
+                                 For more information, see https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DataCaptureConfig.html.
+                                 Defaults to ``None``.
     """
     if mode not in [DEPLOYMENT_MODE_ADD, DEPLOYMENT_MODE_REPLACE]:
         msg = "Invalid mode `{md}` for deployment to a pre-existing application".format(md=mode)
@@ -1647,11 +1652,14 @@ def _update_sagemaker_endpoint(
     # Create the new endpoint configuration and update the endpoint
     # to adopt the new configuration
     new_config_name = _get_sagemaker_config_name(endpoint_name)
-    endpoint_config_response = sage_client.create_endpoint_config(
-        EndpointConfigName=new_config_name,
-        ProductionVariants=production_variants,
-        Tags=[{"Key": "app_name", "Value": endpoint_name}],
-    )
+    endpoint_config_kwargs = {
+        "EndpointConfigName": new_config_name,
+        "ProductionVariants": production_variants,
+        "Tags": [{"Key": "app_name", "Value": endpoint_name}],
+    }
+    if data_capture_config is not None:
+        endpoint_config_kwargs["DataCaptureConfig"] = data_capture_config
+    endpoint_config_response = sage_client.create_endpoint_config(**endpoint_config_kwargs)
     _logger.info(
         "Created new endpoint configuration with arn: %s",
         endpoint_config_response["EndpointConfigArn"],


Test patch: