fix: merge tags instead of delete

Author: EthanShouhanCheng

src/sagemaker/workflow/pipeline.py

@@ -187,11 +187,11 @@ def upsert(
                         ResourceArn=response["PipelineArn"]
                     )["Tags"]
 
-                    tag_keys = [tag["Key"] for tag in old_tags]
+                    tag_keys = [tag["Key"] for tag in tags]
+                    for old_tag in old_tags:
+                        if old_tag["Key"] not in tag_keys:
+                            tags.append(old_tag)
 
-                    self.sagemaker_session.sagemaker_client.delete_tags(
-                        ResourceArn=response["PipelineArn"], TagKeys=tag_keys
-                    )
                     self.sagemaker_session.sagemaker_client.add_tags(
                         ResourceArn=response["PipelineArn"], Tags=tags
                     )

tests/integ/test_workflow.py

@@ -459,28 +459,32 @@ def test_one_step_sklearn_processing_pipeline(
         # under a potentially different role, often requiring access to S3 and other
         # artifacts not required to during creation of the jobs in the pipeline steps.
         response = pipeline.create(
-            role, tags=[{"Key": "foo", "Value": "abc"}, {"Key": "bar", "Value": "xyz"}]
+            role, tags=[{"Key": "foo", "Value": "123"}, {"Key": "bar", "Value": "456"}]
         )
         create_arn = response["PipelineArn"]
         assert re.match(
             fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
             create_arn,
         )
         original_tags = sagemaker_session.sagemaker_client.list_tags(ResourceArn=create_arn)
-        for tag in [{"Key": "foo", "Value": "abc"}, {"Key": "bar", "Value": "xyz"}]:
+        for tag in [{"Key": "foo", "Value": "123"}, {"Key": "bar", "Value": "456"}]:
             assert tag in original_tags["Tags"]
 
         pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
         response = pipeline.upsert(
-            role, tags=[{"Key": "abc", "Value": "foo"}, {"Key": "xyz", "Value": "bar"}]
+            role, tags=[{"Key": "foo", "Value": "abc"}, {"Key": "baz", "Value": "789"}]
         )
         update_arn = response["PipelineArn"]
         assert re.match(
             fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
             update_arn,
         )
         updated_tags = sagemaker_session.sagemaker_client.list_tags(ResourceArn=create_arn)
-        for tag in [{"Key": "abc", "Value": "foo"}, {"Key": "xyz", "Value": "bar"}]:
+        for tag in [
+            {"Key": "foo", "Value": "abc"},
+            {"Key": "bar", "Value": "456"},
+            {"Key": "baz", "Value": "789"},
+        ]:
             assert tag in updated_tags["Tags"]
 
         execution = pipeline.start(parameters={})

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -105,7 +105,6 @@ def test_pipeline_upsert(sagemaker_session_mock, role_arn):
         {"PipelineArn": "mock_pipeline_arn"},
         [{"Key": "dummy", "Value": "dummy_tag"}],
         {},
-        {},
     ]
 
     pipeline = Pipeline(
@@ -129,9 +128,8 @@ def test_pipeline_upsert(sagemaker_session_mock, role_arn):
     assert sagemaker_session_mock.sagemaker_client.list_tags.called_with(
         ResourceArn="mock_pipeline_arn"
     )
-    assert sagemaker_session_mock.sagemaker_client.delete_tags(
-        ResourceArn="mock_pipeline_arn", TagKeys=["dummy"]
-    )
+
+    tags.append({"Key": "dummy", "Value": "dummy_tag"})
     assert sagemaker_session_mock.sagemaker_client.add_tags.called_with(
         ResourceArn="mock_pipeline_arn", Tags=tags
     )