fix: handle tags when upsert pipeine

Author: EthanShouhanCheng

src/sagemaker/workflow/pipeline.py

@@ -143,7 +143,7 @@ def describe(self) -> Dict[str, Any]:
         """
         return self.sagemaker_session.sagemaker_client.describe_pipeline(PipelineName=self.name)
 
-    def update(self, role_arn: str, description: str = None) -> Dict[str, Any]:
+    def update(self, role_arn: str, description: str = None, ) -> Dict[str, Any]:
         """Updates a Pipeline in the Workflow service.
 
         Args:
@@ -182,6 +182,20 @@ 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 old_tags]
+
+                    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
+                    )
             else:
                 raise
         return response

tests/integ/test_workflow.py

@@ -410,13 +410,13 @@ def test_three_step_definition(
 
 def test_one_step_sklearn_processing_pipeline(
     sagemaker_session,
-    role,
     sklearn_latest_version,
     cpu_instance_type,
     pipeline_name,
     region_name,
     athena_dataset_definition,
 ):
+    role = "arn:aws:iam::734680132978:role/SageMakerFullAccessRole"
     instance_count = ParameterInteger(name="InstanceCount", default_value=2)
     script_path = os.path.join(DATA_DIR, "dummy_script.py")
     input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
@@ -458,20 +458,26 @@ 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": "abc"}, {"Key": "bar", "Value": "xyz"}])
         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"}]:
+            assert tag in original_tags["Tags"]
 
         pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
-        response = pipeline.update(role)
+        response = pipeline.upsert(role, tags=[{"Key": "abc", "Value": "foo"}, {"Key": "xyz", "Value": "bar"}])
         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"}]:
+            assert tag in updated_tags["Tags"]
 
         execution = pipeline.start(parameters={})
         assert re.match(

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -102,21 +102,39 @@ 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")
+    assert sagemaker_session_mock.sagemaker_client.delete_tags(
+        ResourceArn="mock_pipeline_arn", TagKeys=["dummy"]
+    )
+    assert sagemaker_session_mock.sagemaker_client.add_tags.called_with(
+        ResourceArn="mock_pipeline_arn", Tags=tags
+    )
 
 
 def test_pipeline_delete(sagemaker_session_mock):