fix: handle tags when upsert pipeine (#2488)

* fix: handle tags when upsert pipeine

* fix: typo and remove local test changes

* fix: typo

* fix: black check failure

* fix: merge tags instead of delete

Co-authored-by: icywang86rui <[email protected]>

Author: EthanShouhanCheng

src/sagemaker/workflow/pipeline.py

@@ -182,6 +182,19 @@ def upsert(
                 and "Pipeline names must be unique within" in error["Message"]
             ):
                 response = self.update(role_arn, description)
+                if tags is not None:
+                    old_tags = self.sagemaker_session.sagemaker_client.list_tags(
+                        ResourceArn=response["PipelineArn"]
+                    )["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.add_tags(
+                        ResourceArn=response["PipelineArn"], Tags=tags
+                    )
             else:
                 raise
         return response

tests/integ/test_workflow.py

@@ -458,20 +458,34 @@ def test_one_step_sklearn_processing_pipeline(
         # sagemaker entities. However, the jobs created in the steps themselves execute
         # 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)
+        response = pipeline.create(
+            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": "123"}, {"Key": "bar", "Value": "456"}]:
+            assert tag in original_tags["Tags"]
 
         pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
-        response = pipeline.update(role)
+        response = pipeline.upsert(
+            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": "foo", "Value": "abc"},
+            {"Key": "bar", "Value": "456"},
+            {"Key": "baz", "Value": "789"},
+        ]:
+            assert tag in updated_tags["Tags"]
 
         execution = pipeline.start(parameters={})
         assert re.match(

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -102,21 +102,37 @@ def test_pipeline_upsert(sagemaker_session_mock, role_arn):
                 }
             },
         ),
+        {"PipelineArn": "mock_pipeline_arn"},
+        [{"Key": "dummy", "Value": "dummy_tag"}],
         {},
     ]
+
     pipeline = Pipeline(
         name="MyPipeline",
         parameters=[],
         steps=[],
         sagemaker_session=sagemaker_session_mock,
     )
-    pipeline.update(role_arn=role_arn)
+
+    tags = [
+        {"Key": "foo", "Value": "abc"},
+        {"Key": "bar", "Value": "xyz"},
+    ]
+    pipeline.upsert(role_arn=role_arn, tags=tags)
     assert sagemaker_session_mock.sagemaker_client.create_pipeline.called_with(
         PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn
     )
     assert sagemaker_session_mock.sagemaker_client.update_pipeline.called_with(
         PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn
     )
+    assert sagemaker_session_mock.sagemaker_client.list_tags.called_with(
+        ResourceArn="mock_pipeline_arn"
+    )
+
+    tags.append({"Key": "dummy", "Value": "dummy_tag"})
+    assert sagemaker_session_mock.sagemaker_client.add_tags.called_with(
+        ResourceArn="mock_pipeline_arn", Tags=tags
+    )
 
 
 def test_pipeline_delete(sagemaker_session_mock):