update trial_compoment pipeline execution arn

Author: yzhu0

src/sagemaker/lineage/lineage_trial_component.py

@@ -130,8 +130,14 @@ def pipeline_execution_arn(self) -> str:
         Returns:
             str: A pipeline execution ARN.
         """
+        trial_component = self.load(trial_component_name=self.trial_component_name,
+                                    sagemaker_session=self.sagemaker_session)
+
+        if trial_component.source is None or trial_component.source["SourceArn"] is None:
+            return None
+
         tags = self.sagemaker_session.sagemaker_client.list_tags(
-            ResourceArn=self.trial_component_arn
+            ResourceArn=trial_component.source["SourceArn"]
         )["Tags"]
         for tag in tags:
             if tag["Key"] == "sagemaker:pipeline-execution-arn":

tests/integ/sagemaker/lineage/conftest.py

@@ -561,14 +561,14 @@ def static_approval_action(
 
 
 @pytest.fixture
-def static_model_deployment_action(sagemaker_session, static_endpoint_context):
+def static_model_deployment_action(sagemaker_session, static_processing_job_trial_component):
     query_filter = LineageFilter(
         entities=[LineageEntityEnum.ACTION], sources=[LineageSourceEnum.MODEL_DEPLOYMENT]
     )
     query_result = LineageQuery(sagemaker_session).query(
-        start_arns=[static_endpoint_context.context_arn],
+        start_arns=[static_processing_job_trial_component.trial_component_arn],
         query_filter=query_filter,
-        direction=LineageQueryDirectionEnum.ASCENDANTS,
+        direction=LineageQueryDirectionEnum.DESCENDANTS,
         include_edges=False,
     )
     model_approval_actions = []
@@ -579,14 +579,14 @@ def static_model_deployment_action(sagemaker_session, static_endpoint_context):
 
 @pytest.fixture
 def static_processing_job_trial_component(
-    sagemaker_session, static_endpoint_context
+    sagemaker_session, static_dataset_artifact
 ) -> LineageTrialComponent:
     query_filter = LineageFilter(
         entities=[LineageEntityEnum.TRIAL_COMPONENT], sources=[LineageSourceEnum.PROCESSING_JOB]
     )
 
     query_result = LineageQuery(sagemaker_session).query(
-        start_arns=[static_endpoint_context.context_arn],
+        start_arns=[static_dataset_artifact.artifact_arn],
         query_filter=query_filter,
         direction=LineageQueryDirectionEnum.ASCENDANTS,
         include_edges=False,
@@ -600,16 +600,16 @@ def static_processing_job_trial_component(
 
 @pytest.fixture
 def static_training_job_trial_component(
-    sagemaker_session, static_endpoint_context
+    sagemaker_session, static_dataset_artifact
 ) -> LineageTrialComponent:
     query_filter = LineageFilter(
         entities=[LineageEntityEnum.TRIAL_COMPONENT], sources=[LineageSourceEnum.TRAINING_JOB]
     )
 
     query_result = LineageQuery(sagemaker_session).query(
-        start_arns=[static_endpoint_context.context_arn],
+        start_arns=[static_dataset_artifact.artifact_arn],
         query_filter=query_filter,
-        direction=LineageQueryDirectionEnum.ASCENDANTS,
+        direction=LineageQueryDirectionEnum.DESCENDANTS,
         include_edges=False,
     )
     training_jobs = []
@@ -738,12 +738,12 @@ def static_dataset_artifact(static_model_artifact, sagemaker_session):
 
 
 @pytest.fixture
-def static_image_artifact(static_model_artifact, sagemaker_session):
+def static_image_artifact(static_dataset_artifact, sagemaker_session):
     query_filter = LineageFilter(
         entities=[LineageEntityEnum.ARTIFACT], sources=[LineageSourceEnum.IMAGE]
     )
     query_result = LineageQuery(sagemaker_session).query(
-        start_arns=[static_model_artifact.artifact_arn],
+        start_arns=[static_dataset_artifact.artifact_arn],
         query_filter=query_filter,
         direction=LineageQueryDirectionEnum.ASCENDANTS,
         include_edges=False,