[Feature] Propagate tags to lineage resources (#4027)

Author: jiapinw

src/sagemaker/feature_store/feature_processor/feature_scheduler.py

@@ -129,8 +129,8 @@ def to_pipeline(
             code for Lineage tracking. This code is not used for actual transformation.
         max_retries (Optional[int]): The number of times to retry sagemaker pipeline step.
             If not specified, sagemaker pipline step will not retry.
-        tags (List[Tuple[str, str]): A list of tags attached to the pipeline. If not specified,
-            no custom tags will be attached to the pipeline.
+        tags (List[Tuple[str, str]): A list of tags attached to the pipeline and all corresponding
+            lineage resources that support tags. If not specified, no custom tags will be attached.
         sagemaker_session (Optional[Session]): Session object which manages interactions
             with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
             function creates one using the default AWS configuration chain.
@@ -220,6 +220,9 @@ def to_pipeline(
 
     describe_pipeline_response = pipeline.describe()
     pipeline_arn = describe_pipeline_response["PipelineArn"]
+    tags_propagate_to_lineage_resources = _get_tags_from_pipeline_to_propagate_to_lineage_resources(
+        pipeline_arn, _sagemaker_session
+    )
 
     lineage_handler = FeatureProcessorLineageHandler(
         pipeline_name=pipeline_name,
@@ -230,7 +233,8 @@ def to_pipeline(
         transformation_code=transformation_code,
         sagemaker_session=_sagemaker_session,
     )
-    lineage_handler.create_lineage()
+    lineage_handler.create_lineage(tags_propagate_to_lineage_resources)
+    lineage_handler.upsert_tags_for_lineage_resources(tags_propagate_to_lineage_resources)
 
     pipeline_lineage_names: Dict[str, str] = lineage_handler.get_pipeline_lineage_names()
 
@@ -299,15 +303,19 @@ def schedule(
     describe_pipeline_response = _sagemaker_session.sagemaker_client.describe_pipeline(
         PipelineName=pipeline_name
     )
+    pipeline_arn = describe_pipeline_response["PipelineArn"]
+    tags_propagate_to_lineage_resources = _get_tags_from_pipeline_to_propagate_to_lineage_resources(
+        pipeline_arn, _sagemaker_session
+    )
 
     logger.info("Creating/Updating EventBridge Schedule for pipeline %s.", pipeline_name)
     event_bridge_schedule_arn = event_bridge_scheduler_helper.upsert_schedule(
-        pipeline_name,
-        describe_pipeline_response["PipelineArn"],
-        schedule_expression,
-        state,
-        _start_date,
-        _role_arn,
+        schedule_name=pipeline_name,
+        pipeline_arn=pipeline_arn,
+        schedule_expression=schedule_expression,
+        state=state,
+        start_date=_start_date,
+        role=_role_arn,
     )
     logger.info("Created/Updated EventBridge Schedule for pipeline %s.", pipeline_name)
     lineage_handler = FeatureProcessorLineageHandler(
@@ -322,6 +330,7 @@ def schedule(
         schedule_expression=schedule_expression,
         state=state,
         start_date=_start_date,
+        tags=tags_propagate_to_lineage_resources,
     )
     return event_bridge_schedule_arn["ScheduleArn"]
 
@@ -805,7 +814,7 @@ def _get_feature_processor_outputs(
 def _parse_name_from_arn(fg_uri: str) -> str:
     """Parse the name from a string, if it's an ARN. Otherwise, return the string.
 
-    Attributes:
+    Args:
         fg_uri (str): The Feature Group Name or ARN.
 
     Returns:
@@ -816,3 +825,23 @@ def _parse_name_from_arn(fg_uri: str) -> str:
         feature_group_name = match.group(4)
         return feature_group_name
     return fg_uri
+
+
+def _get_tags_from_pipeline_to_propagate_to_lineage_resources(
+    pipeline_arn: str, sagemaker_session: Session
+) -> List[Dict[str, str]]:
+    """Retrieve custom tags attached to sagemakre pipeline
+
+    Args:
+        pipeline_arn (str): SageMaker Pipeline Arn.
+        sagemaker_session (Session): Session object which manages interactions
+            with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
+            function creates one using the default AWS configuration chain.
+
+    Returns:
+        List[Dict[str, str]]: List of custom tags to be propagated to lineage resources.
+    """
+    tags_in_pipeline = sagemaker_session.sagemaker_client.list_tags(ResourceArn=pipeline_arn)[
+        "Tags"
+    ]
+    return [d for d in tags_in_pipeline if d["Key"] not in TO_PIPELINE_RESERVED_TAG_KEYS]

src/sagemaker/feature_store/feature_processor/lineage/_feature_processor_lineage.py

@@ -18,6 +18,9 @@
 import attr
 from botocore.exceptions import ClientError
 
+from sagemaker.feature_store.feature_processor._event_bridge_scheduler_helper import (
+    EventBridgeSchedulerHelper,
+)
 from sagemaker.feature_store.feature_processor.lineage._lineage_association_handler import (
     LineageAssociationHandler,
 )
@@ -102,7 +105,7 @@ class FeatureProcessorLineageHandler:
     output: str = attr.ib(default=None)
     transformation_code: TransformationCode = attr.ib(default=None)
 
-    def create_lineage(self) -> None:
+    def create_lineage(self, tags: Optional[List[Dict[str, str]]] = None) -> None:
         """Create and Update Feature Processor Lineage"""
         input_feature_group_contexts: List[
             FeatureGroupContexts
@@ -120,6 +123,8 @@ def create_lineage(self) -> None:
         )
         if transformation_code_artifact is not None:
             logger.info("Created Transformation Code Artifact: %s", transformation_code_artifact)
+            if tags:
+                transformation_code_artifact.set_tags(tags)  # pylint: disable=E1101
         # Create the Pipeline Lineage for the first time
         if not self._check_if_pipeline_lineage_exists():
             self._create_new_pipeline_lineage(
@@ -160,6 +165,7 @@ def create_schedule_lineage(
         schedule_expression,
         state,
         start_date: datetime,
+        tags: Optional[List[Dict[str, str]]] = None,
     ) -> None:
         """Class to Create and Update FeatureProcessor Lineage Entities.
 
@@ -171,10 +177,12 @@ def create_schedule_lineage(
             state (str):Specifies whether the schedule is enabled or disabled. Valid values are
                 ENABLED and DISABLED. See https://docs.aws.amazon.com/scheduler/latest/APIReference/
                 API_CreateSchedule.html#scheduler-CreateSchedule-request-State for more details.
-                If not specified, it will default to DISABLED.
+                If not specified, it will default to ENABLED.
             start_date (Optional[datetime]): The date, in UTC, after which the schedule can begin
                 invoking its target. Depending on the schedule’s recurrence expression, invocations
                 might occur on, or after, the StartDate you specify.
+            tags (Optional[List[Dict[str, str]]]): Custom tags to be attached to schedule
+                lineage resource.
         """
         pipeline_context: Context = self._get_pipeline_context()
         pipeline_version_context: Context = self._get_pipeline_version_context(
@@ -192,13 +200,55 @@ def create_schedule_lineage(
             pipeline_schedule=pipeline_schedule,
             sagemaker_session=self.sagemaker_session,
         )
+        if tags:
+            schedule_artifact.set_tags(tags)
 
         LineageAssociationHandler.add_upstream_schedule_associations(
             schedule_artifact=schedule_artifact,
             pipeline_version_context_arn=pipeline_version_context.context_arn,
             sagemaker_session=self.sagemaker_session,
         )
 
+    def upsert_tags_for_lineage_resources(self, tags: List[Dict[str, str]]) -> None:
+        """Add or update tags for lineage resources using tags attached to sagemaker pipeline as
+
+        source of truth.
+
+        Args:
+            tags (List[Dict[str, str]]): Custom tags to be attached to lineage resources.
+        """
+        if not tags:
+            return
+        pipeline_context: Context = self._get_pipeline_context()
+        current_pipeline_version_context: Context = self._get_pipeline_version_context(
+            last_update_time=pipeline_context.properties[LAST_UPDATE_TIME]
+        )
+        input_raw_data_artifacts: List[Artifact] = self._retrieve_input_raw_data_artifacts()
+        pipeline_context.set_tags(tags)
+        current_pipeline_version_context.set_tags(tags)
+        for input_raw_data_artifact in input_raw_data_artifacts:
+            input_raw_data_artifact.set_tags(tags)
+
+        event_bridge_scheduler_helper = EventBridgeSchedulerHelper(
+            self.sagemaker_session,
+            self.sagemaker_session.boto_session.client("scheduler"),
+        )
+        event_bridge_schedule = event_bridge_scheduler_helper.describe_schedule(self.pipeline_name)
+
+        if event_bridge_schedule:
+            schedule_artifact_summary = S3LineageEntityHandler._load_artifact_from_s3_uri(
+                s3_uri=event_bridge_schedule["Arn"],
+                sagemaker_session=self.sagemaker_session,
+            )
+            if schedule_artifact_summary is not None:
+                pipeline_schedule_artifact: Artifact = (
+                    S3LineageEntityHandler.load_artifact_from_arn(
+                        artifact_arn=schedule_artifact_summary.artifact_arn,
+                        sagemaker_session=self.sagemaker_session,
+                    )
+                )
+                pipeline_schedule_artifact.set_tags(tags)
+
     def _create_new_pipeline_lineage(
         self,
         input_feature_group_contexts: List[FeatureGroupContexts],

tests/unit/sagemaker/feature_store/feature_processor/lineage/test_constants.py

@@ -313,7 +313,6 @@
     },
 }
 
-
 PIPELINE_CONTEXT: Context = Context(
     context_arn=f"{PIPELINE_NAME}-context-arn",
     context_name=f"sm-fs-fe-{PIPELINE_NAME}-{CREATION_TIME}-fep",
@@ -341,7 +340,6 @@
     name="test-name",
 )
 
-
 TRANSFORMATION_CODE_INPUT_2: TransformationCode = TransformationCode(
     s3_uri="s3://sagemaker-us-west-2-789975069016/transform-2023-04-28-21-50-14-616/"
     "transform-2023-04-28-21-50-14-616/output/model.tar.gz/2",