[BUG] Tags are not added to sagemaker endpoints during deployment.

[clarkh-ncino]

Issues Policy acknowledgement

• I have read and agree to submit bug reports in accordance with the issues policy

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 and _update_sagemaker_endpoint functions are not added to endpoints as expected. Tags for endpoints are being set to an empty array here.

Tracking information

n/a

Code to reproduce issue

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

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

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
• 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
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

[harupy]

@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?

[harupy]

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.

Both model and endpoint are taggable.

@clarkh-ncino Would it makes sense to apply the same tag on both model and endpoint?

[clarkh-ncino]

@harupy Yes, I believe that is the expected behavior.

[harupy]

@clarkh-ncino Sounds good. Can you file a RP?

[clarkh-ncino]

@harupy Just to clarify, do you mean pull request (PR)? If so, absolutely. Can share it out in the near future.

[harupy]

yes, I meant pull request.

[mlflow-automation]

@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.

[clarkh-ncino]

@harupy #9310 is ready for review

[clarkh-ncino]

@B-Step62 The associated PR is not merged yet so I'm not sure if we would want to mark this as completed.

[BenWilson2]

@B-Step62 The associated PR is not merged yet so I'm not sure if we would want to mark this as completed.

Sorry for the delay on that. I just merged it!