Repository: mlflow/mlflow

Issue:
[BUG] Tags are not added to sagemaker endpoints during deployment.
### 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

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

### MLflow version

- Client: 2.5.0
- Tracking server: n/a


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: n/a
- **Python version**: 3.10.9



### Describe the problem

When deploying a model to sagemaker, tags passed to the [_create_sagemaker_endpoint](https://github.com/mlflow/mlflow/blob/c9397e6ac007a306b2c99f7c5dc82ebf98651780/mlflow/sagemaker/__init__.py#L1516) and [_update_sagemaker_endpoint](https://github.com/mlflow/mlflow/blob/c9397e6ac007a306b2c99f7c5dc82ebf98651780/mlflow/sagemaker/__init__.py#L1641) functions are not added to endpoints as expected. Tags for endpoints are being set to an empty array [here](https://github.com/mlflow/mlflow/blob/master/mlflow/sagemaker/__init__.py#L1559-L1562).

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
n/a
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
from mlflow.deployments import get_deploy_client

sagemaker_client = get_deploy_client('sagemaker')
model_uri = 'a_unique_uri'
tags = {
    "test": "testing-tags-for-sagemaker"
}
bucket_name = "test_bucket"
config = {
    'region_name': "us-east-1",
    'execution_role_arn': role_arn,
    'bucket': bucket_name,
    'instance_type': 'ml.t2.large',
    'tags': tags
}
model_name = "testing-sagemaker"
endpoint = None
flavor = "python_function"
sagemaker_client.create_deployment(
    model_name,
    model_uri,
    flavor,
    config,
    endpoint
)
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
2023/07/25 16:25:23 INFO mlflow.sagemaker: Using the python_function flavor for deployment!
2023-07-25 16:25:24,584.584 [INFO]: Found credentials in environment variables.
2023/07/25 16:25:25 INFO mlflow.sagemaker: Creating new endpoint with name: testing-sagemaker ...
2023/07/25 16:25:26 INFO mlflow.sagemaker: Created model with arn: arn:aws:sagemaker:us-east-1:-:model/testing-sagemaker-model-hxqure4htgcz1c3titoacg
2023/07/25 16:25:27 INFO mlflow.sagemaker: Created endpoint configuration with arn: arn:aws:sagemaker:us-east-1:-:endpoint-config/testing-sagemaker-config-vo1r7uqgt6ihxbx1q6lahw
2023/07/25 16:25:27 INFO mlflow.sagemaker: Created endpoint with arn: arn:aws:sagemaker:us-east-1:-:endpoint/testing-sagemaker
2023/07/25 16:25:27 INFO mlflow.sagemaker: Waiting for the deployment operation to complete...
2023/07/25 16:25:28 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:25:48 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:26:08 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:26:29 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:26:49 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:27:10 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:27:30 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:27:50 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:28:11 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:28:31 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:28:52 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:29:12 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:29:33 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:29:53 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:30:13 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:30:34 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:30:54 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:31:15 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:31:35 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:31:55 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:32:16 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:32:36 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:32:57 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:33:17 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:33:38 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:33:58 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:34:18 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:34:39 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:34:59 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:35:20 INFO mlflow.sagemaker: Waiting for endpoint to reach the "InService" state. Current endpoint status: "Creating"
2023/07/25 16:35:35 INFO mlflow.sagemaker: The deployment operation completed successfully with message: "The SageMaker endpoint was created successfully."
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
n/a
```



### 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
- [ ] `area/gateway`: AI Gateway service, Gateway client APIs, third-party Gateway integrations
- [ ] `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
- [X] `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:
@clarkh-ncino Thanks for reporting the issue! It looks like we use `tags` when creating a sagemaker model. I'm not familiar with sagemaker. What can be tagged?
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_Tag.html

> You can add tags to notebook instances, training jobs, hyperparameter tuning jobs, batch transform jobs, models, labeling jobs, work teams, endpoint configurations, and endpoints. For more information on adding tags to SageMaker resources, see [AddTags](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_AddTags.html).

Both model and endpoint are taggable.


@clarkh-ncino Would it makes sense to apply the same tag on both model and endpoint?
@harupy Yes, I believe that is the expected behavior. 
@clarkh-ncino Sounds good. Can you file a RP?
@harupy Just to clarify, do you mean pull request (PR)? If so, absolutely. Can share it out in the near future.
yes, I meant pull request.

Resolving patch:
diff --git a/mlflow/sagemaker/__init__.py b/mlflow/sagemaker/__init__.py
--- a/mlflow/sagemaker/__init__.py
+++ b/mlflow/sagemaker/__init__.py
@@ -11,7 +11,7 @@
 import time
 import urllib.parse
 from subprocess import Popen
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 import mlflow
 import mlflow.version
@@ -50,6 +50,7 @@
 
 DEFAULT_REGION_NAME = "us-west-2"
 SAGEMAKER_SERVING_ENVIRONMENT = "SageMaker"
+SAGEMAKER_APP_NAME_TAG_KEY = "app_name"
 
 _logger = logging.getLogger(__name__)
 
@@ -1346,6 +1347,26 @@ def _get_sagemaker_config_name(endpoint_name):
     return f"{endpoint_name}-config-{get_unique_resource_id()}"
 
 
+def _get_sagemaker_config_tags(endpoint_name):
+    return [{"Key": SAGEMAKER_APP_NAME_TAG_KEY, "Value": endpoint_name}]
+
+
+def _prepare_sagemaker_tags(
+    config_tags: List[Dict[str, str]],
+    sagemaker_tags: Optional[Dict[str, str]] = None,
+):
+    if not sagemaker_tags:
+        return config_tags
+
+    if SAGEMAKER_APP_NAME_TAG_KEY in sagemaker_tags:
+        raise MlflowException.invalid_parameter_value(
+            f"Duplicate tag provided for '{SAGEMAKER_APP_NAME_TAG_KEY}'"
+        )
+    parsed = [{"Key": key, "Value": str(value)} for key, value in sagemaker_tags.items()]
+
+    return config_tags + parsed
+
+
 def _create_sagemaker_transform_job(
     job_name,
     model_name,
@@ -1559,10 +1580,12 @@ def _create_sagemaker_endpoint(
         production_variant["InitialInstanceCount"] = instance_count
 
     config_name = _get_sagemaker_config_name(endpoint_name)
+    config_tags = _get_sagemaker_config_tags(endpoint_name)
+    tags_list = _prepare_sagemaker_tags(config_tags, tags)
     endpoint_config_kwargs = {
         "EndpointConfigName": config_name,
         "ProductionVariants": [production_variant],
-        "Tags": [{"Key": "app_name", "Value": endpoint_name}],
+        "Tags": config_tags,
     }
     if async_inference_config:
         endpoint_config_kwargs["AsyncInferenceConfig"] = async_inference_config
@@ -1576,7 +1599,7 @@ def _create_sagemaker_endpoint(
     endpoint_response = sage_client.create_endpoint(
         EndpointName=endpoint_name,
         EndpointConfigName=config_name,
-        Tags=[],
+        Tags=tags_list or [],
     )
     _logger.info("Created endpoint with arn: %s", endpoint_response["EndpointArn"])
 
@@ -1656,7 +1679,7 @@ def _update_sagemaker_endpoint(
                                  For more information, see https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DataCaptureConfig.html.
                                  Defaults to ``None``.
     :param env: A dictionary of environment variables to set for the model.
-    :param tags: A dictionary of tags to apply to the endpoint.
+    :param tags: A dictionary of tags to apply to the endpoint configuration.
     """
     if mode not in [DEPLOYMENT_MODE_ADD, DEPLOYMENT_MODE_REPLACE]:
         msg = f"Invalid mode `{mode}` for deployment to a pre-existing application"
@@ -1714,11 +1737,11 @@ 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)
-    # This is the hardcoded config for endpoint
+    config_tags = _get_sagemaker_config_tags(endpoint_name)
     endpoint_config_kwargs = {
         "EndpointConfigName": new_config_name,
         "ProductionVariants": production_variants,
-        "Tags": [{"Key": "app_name", "Value": endpoint_name}],
+        "Tags": config_tags,
     }
     if async_inference_config:
         endpoint_config_kwargs["AsyncInferenceConfig"] = async_inference_config


Test patch:
diff --git a/tests/sagemaker/mock/__init__.py b/tests/sagemaker/mock/__init__.py
--- a/tests/sagemaker/mock/__init__.py
+++ b/tests/sagemaker/mock/__init__.py
@@ -187,11 +187,12 @@ def list_tags(self):
         Handler for the SageMaker "ListTags" API call documented here:
         https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ListTags.html
         """
-        model_arn = self.request_params["ResourceArn"]
-
+        arn = self.request_params["ResourceArn"]
+        sagemaker_resource = (
+            "models" if "model" in arn else "endpoints" if "endpoint" in arn else None
+        )
         results = self.sagemaker_backend.list_tags(
-            resource_arn=model_arn,
-            region_name=self.region,
+            resource_arn=arn, region_name=self.region, resource_type=sagemaker_resource
         )
 
         return json.dumps({"Tags": results, "NextToken": None})
@@ -511,13 +512,17 @@ def list_models(self):
             summaries.append(summary)
         return summaries
 
-    def list_tags(self, resource_arn, region_name):  # pylint: disable=unused-argument
+    def list_tags(
+        self, resource_arn, region_name, resource_type
+    ):  # pylint: disable=unused-argument
         """
         Modifies backend state during calls to the SageMaker "ListTags" API
         https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ListTags.html
         """
-        model = next(model for model in self.models.values() if model.arn == resource_arn)
-        return model.resource.tags
+        resource_values = getattr(self, resource_type).values()
+        for sagemaker_resource in resource_values:
+            if sagemaker_resource.arn == resource_arn:
+                return sagemaker_resource.resource.tags
 
     def create_model(
         self, model_name, primary_container, execution_role_arn, tags, region_name, vpc_config=None
diff --git a/tests/sagemaker/test_sagemaker_deployment_client.py b/tests/sagemaker/test_sagemaker_deployment_client.py
--- a/tests/sagemaker/test_sagemaker_deployment_client.py
+++ b/tests/sagemaker/test_sagemaker_deployment_client.py
@@ -508,10 +508,39 @@ def test_create_deployment_create_sagemaker_and_s3_resources_with_expected_tags_
     model_name = endpoint_production_variants[0]["VariantName"]
     description = sagemaker_client.describe_model(ModelName=model_name)
 
-    tags = sagemaker_client.list_tags(ResourceArn=description["ModelArn"])
+    model_tags = sagemaker_client.list_tags(ResourceArn=description["ModelArn"])
+    endpoint_tags = sagemaker_client.list_tags(ResourceArn=endpoint_description["EndpointArn"])
 
     # Extra tags exist besides the ones we set, so avoid strict equality
-    assert all(tag in tags["Tags"] for tag in expected_tags)
+    assert all(tag in model_tags["Tags"] for tag in expected_tags)
+    assert all(tag in endpoint_tags["Tags"] for tag in expected_tags)
+
+
+def test_prepare_sagemaker_tags_without_custom_tags():
+    config_tags = [{"Key": "tag1", "Value": "value1"}]
+    tags = mfs._prepare_sagemaker_tags(config_tags, None)
+    assert tags == config_tags
+
+
+def test_prepare_sagemaker_tags_when_custom_tags_are_added():
+    config_tags = [{"Key": "tag1", "Value": "value1"}]
+    sagemaker_tags = {"tag2": "value2", "tag3": "123"}
+    expected_tags = [
+        {"Key": "tag1", "Value": "value1"},
+        {"Key": "tag2", "Value": "value2"},
+        {"Key": "tag3", "Value": "123"},
+    ]
+    tags = mfs._prepare_sagemaker_tags(config_tags, sagemaker_tags)
+    assert tags == expected_tags
+
+
+def test_prepare_sagemaker_tags_duplicate_key_raises_exception():
+    config_tags = [{"Key": "app_name", "Value": "a_cool_name"}]
+    sagemaker_tags = {"app_name": "a_cooler_name", "tag2": "value2", "tag3": "123"}
+    match = "Duplicate tag provided for 'app_name'"
+    with pytest.raises(MlflowException, match=match) as exc:
+        mfs._prepare_sagemaker_tags(config_tags, sagemaker_tags)
+    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
 
 
 @pytest.mark.parametrize("proxies_enabled", [True, False])
@@ -673,10 +702,12 @@ def test_deploy_cli_creates_sagemaker_and_s3_resources_with_expected_tags_from_l
     model_name = endpoint_production_variants[0]["VariantName"]
     description = sagemaker_client.describe_model(ModelName=model_name)
 
-    tags = sagemaker_client.list_tags(ResourceArn=description["ModelArn"])
+    model_tags = sagemaker_client.list_tags(ResourceArn=description["ModelArn"])
+    endpoint_tags = sagemaker_client.list_tags(ResourceArn=endpoint_description["EndpointArn"])
 
     # Extra tags exist besides the ones we set, so avoid strict equality
-    assert all(tag in tags["Tags"] for tag in expected_tags)
+    assert all(tag in model_tags["Tags"] for tag in expected_tags)
+    assert all(tag in endpoint_tags["Tags"] for tag in expected_tags)
 
 
 @pytest.mark.parametrize("proxies_enabled", [True, False])