fix: update lineage_trial_compoment get pipeline execution arn (#2944)

Co-authored-by: Shreya Pandit <[email protected]>

Author: 2 people

src/sagemaker/lineage/lineage_trial_component.py

@@ -130,8 +130,15 @@ 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

@@ -233,7 +233,7 @@ def upstream_trial_associated_artifact(
         sagemaker_session=sagemaker_session,
     )
     trial_obj.add_trial_component(trial_component_obj)
-    time.sleep(3)
+    time.sleep(4)
     yield artifact_obj
     trial_obj.remove_trial_component(trial_component_obj)
     assntn.delete()
@@ -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,14 +600,14 @@ def static_processing_job_trial_component(
 
 @pytest.fixture
 def static_training_job_trial_component(
-    sagemaker_session, static_endpoint_context
+    sagemaker_session, static_model_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_model_artifact.artifact_arn],
         query_filter=query_filter,
         direction=LineageQueryDirectionEnum.ASCENDANTS,
         include_edges=False,
@@ -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,

tests/unit/sagemaker/lineage/test_lineage_trial_component.py

@@ -114,19 +114,39 @@ def test_pipeline_execution_arn(sagemaker_session):
     trial_component_arn = (
         "arn:aws:sagemaker:us-west-2:123456789012:trial_component/lineage-unit-3b05f017-0d87-4c37"
     )
-    obj = lineage_trial_component.LineageTrialComponent(
-        sagemaker_session, trial_component_name="foo", trial_component_arn=trial_component_arn
+    training_job_arn = (
+        "arn:aws:sagemaker:us-west-2:123456789012:training-job/pipelines-bs6gaeln463r-abalonetrain"
     )
+    context = lineage_trial_component.LineageTrialComponent(
+        sagemaker_session,
+        trial_component_name="foo",
+        trial_component_arn=trial_component_arn,
+        source={
+            "SourceArn": training_job_arn,
+            "SourceType": "SageMakerTrainingJob",
+        },
+    )
+    obj = {
+        "TrialComponentName": "pipelines-bs6gaeln463r-AbaloneTrain-A0QiDGuY6z-aws-training-job",
+        "TrialComponentArn": trial_component_arn,
+        "DisplayName": "pipelines-bs6gaeln463r-AbaloneTrain-A0QiDGuY6z-aws-training-job",
+        "Source": {
+            "SourceArn": training_job_arn,
+            "SourceType": "SageMakerTrainingJob",
+        },
+    }
+    sagemaker_session.sagemaker_client.describe_trial_component.return_value = obj
 
     sagemaker_session.sagemaker_client.list_tags.return_value = {
         "Tags": [
             {"Key": "sagemaker:pipeline-execution-arn", "Value": "tag1"},
         ],
     }
     expected_calls = [
-        unittest.mock.call(ResourceArn=trial_component_arn),
+        unittest.mock.call(ResourceArn=training_job_arn),
     ]
-    pipeline_execution_arn_result = obj.pipeline_execution_arn()
+    pipeline_execution_arn_result = context.pipeline_execution_arn()
+
     assert pipeline_execution_arn_result == "tag1"
     assert expected_calls == sagemaker_session.sagemaker_client.list_tags.mock_calls
 
@@ -135,19 +155,77 @@ def test_no_pipeline_execution_arn(sagemaker_session):
     trial_component_arn = (
         "arn:aws:sagemaker:us-west-2:123456789012:trial_component/lineage-unit-3b05f017-0d87-4c37"
     )
-    obj = lineage_trial_component.LineageTrialComponent(
-        sagemaker_session, trial_component_name="foo", trial_component_arn=trial_component_arn
+    training_job_arn = (
+        "arn:aws:sagemaker:us-west-2:123456789012:training-job/pipelines-bs6gaeln463r-abalonetrain"
     )
+    context = lineage_trial_component.LineageTrialComponent(
+        sagemaker_session,
+        trial_component_name="foo",
+        trial_component_arn=trial_component_arn,
+        source={
+            "SourceArn": training_job_arn,
+            "SourceType": "SageMakerTrainingJob",
+        },
+    )
+    obj = {
+        "TrialComponentName": "pipelines-bs6gaeln463r-AbaloneTrain-A0QiDGuY6z-aws-training-job",
+        "TrialComponentArn": trial_component_arn,
+        "DisplayName": "pipelines-bs6gaeln463r-AbaloneTrain-A0QiDGuY6z-aws-training-job",
+        "Source": {
+            "SourceArn": training_job_arn,
+            "SourceType": "SageMakerTrainingJob",
+        },
+    }
+    sagemaker_session.sagemaker_client.describe_trial_component.return_value = obj
 
     sagemaker_session.sagemaker_client.list_tags.return_value = {
         "Tags": [
             {"Key": "abcd", "Value": "efg"},
         ],
     }
     expected_calls = [
-        unittest.mock.call(ResourceArn=trial_component_arn),
+        unittest.mock.call(ResourceArn=training_job_arn),
     ]
-    pipeline_execution_arn_result = obj.pipeline_execution_arn()
+    pipeline_execution_arn_result = context.pipeline_execution_arn()
+    expected_result = None
+    assert pipeline_execution_arn_result == expected_result
+    assert expected_calls == sagemaker_session.sagemaker_client.list_tags.mock_calls
+
+
+def test_no_source_arn_pipeline_execution_arn(sagemaker_session):
+    trial_component_arn = (
+        "arn:aws:sagemaker:us-west-2:123456789012:trial_component/lineage-unit-3b05f017-0d87-4c37"
+    )
+    training_job_arn = (
+        "arn:aws:sagemaker:us-west-2:123456789012:training-job/pipelines-bs6gaeln463r-abalonetrain"
+    )
+    context = lineage_trial_component.LineageTrialComponent(
+        sagemaker_session,
+        trial_component_name="foo",
+        trial_component_arn=trial_component_arn,
+        source={
+            "SourceArn": training_job_arn,
+            "SourceType": "SageMakerTrainingJob",
+        },
+    )
+    obj = {
+        "TrialComponentName": "pipelines-bs6gaeln463r-AbaloneTrain-A0QiDGuY6z-aws-training-job",
+        "TrialComponentArn": trial_component_arn,
+        "DisplayName": "pipelines-bs6gaeln463r-AbaloneTrain-A0QiDGuY6z-aws-training-job",
+        "Source": {
+            "SourceArn": None,
+            "SourceType": None,
+        },
+    }
+    sagemaker_session.sagemaker_client.describe_trial_component.return_value = obj
+
+    sagemaker_session.sagemaker_client.list_tags.return_value = {
+        "Tags": [
+            {"Key": "abcd", "Value": "efg"},
+        ],
+    }
+    expected_calls = []
+    pipeline_execution_arn_result = context.pipeline_execution_arn()
     expected_result = None
     assert pipeline_execution_arn_result == expected_result
     assert expected_calls == sagemaker_session.sagemaker_client.list_tags.mock_calls