feature: Add support for SageMaker lineage queries context

Author: yzhu0

src/sagemaker/lineage/context.py

@@ -333,6 +333,57 @@ def training_job_arns(
             training_job_arns.append(trial_component["Source"]["SourceArn"])
         return training_job_arns
 
+    def processing_job_arns(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[str]:
+        """Get ARNs for all processing jobs that appear in the endpoint's lineage.
+
+        Returns:
+            list of str: Processing job ARNs.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.TRIAL_COMPONENT],
+            sources=[LineageSourceEnum.PROCESSING_JOB],
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.context_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        processing_job_arns = []
+        for vertex in query_result.vertices:
+            trial_component_name = _utils.get_resource_name_from_arn(vertex.arn)
+            trial_component = self.sagemaker_session.sagemaker_client.describe_trial_component(
+                TrialComponentName=trial_component_name
+            )
+            processing_job_arns.append(trial_component["Source"]["SourceArn"])
+        return processing_job_arns
+
+    def trial_components_arns(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[str]:
+        """Get ARNs for all trial components that appear in the endpoint's lineage.
+
+        Returns:
+            list of str: Trial components ARNs.
+        """
+        query_filter = LineageFilter(entities=[LineageEntityEnum.TRIAL_COMPONENT])
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.context_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        transform_job_arns = []
+        for vertex in query_result.vertices:
+            trial_component_name = _utils.get_resource_name_from_arn(vertex.arn)
+            trial_component = self.sagemaker_session.sagemaker_client.describe_trial_component(
+                TrialComponentName=trial_component_name
+            )
+            transform_job_arns.append(trial_component["Source"]["SourceArn"])
+        return transform_job_arns
+
     def pipeline_execution_arn(
         self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
     ) -> str:
@@ -351,3 +402,16 @@ def pipeline_execution_arn(
                     return tag["Value"]
 
         return None
+
+
+class ModelPackageGroup(Context):
+    """An Amazon SageMaker model package group context, which is part of a SageMaker lineage."""
+
+    def pipeline_execution_arn(self) -> str:
+        """Get the ARN for the pipeline execution associated with this model package group (if any).
+
+        Returns:
+            str: A pipeline execution ARN.
+        """
+        # return self.properties["PipelineExecutionArn"]
+        return self.properties.get("PipelineExecutionArn")

src/sagemaker/lineage/query.py

@@ -41,6 +41,8 @@ class LineageSourceEnum(Enum):
     MODEL_REPLACE = "ModelReplaced"
     TENSORBOARD = "TensorBoard"
     TRAINING_JOB = "TrainingJob"
+    PROCESSING_JOB = "ProcessingJob"
+    TRANSFORM_JOB = "TransformJob"
 
 
 class LineageQueryDirectionEnum(Enum):
@@ -87,6 +89,7 @@ def to_lineage_object(self):
         from sagemaker.lineage.artifact import Artifact, ModelArtifact
         from sagemaker.lineage.context import Context, EndpointContext
         from sagemaker.lineage.artifact import DatasetArtifact
+        from sagemaker.lineage.action import Action
 
         if self.lineage_entity == LineageEntityEnum.CONTEXT.value:
             resource_name = get_resource_name_from_arn(self.arn)
@@ -103,6 +106,8 @@ def to_lineage_object(self):
                 return DatasetArtifact.load(artifact_arn=self.arn, sagemaker_session=self._session)
             return Artifact.load(artifact_arn=self.arn, sagemaker_session=self._session)
 
+        if self.lineage_entity == LineageEntityEnum.ACTION.value:
+            return Action.load(action_name=self.arn, sagemaker_session=self._session)
         raise ValueError("Vertex cannot be converted to a lineage object.")
 
 
@@ -155,11 +160,11 @@ def __init__(
     def _to_request_dict(self):
         """Convert the lineage filter to its API representation."""
         filter_request = {}
-        if self.entities:
+        if self.sources:
             filter_request["Types"] = list(
                 map(lambda x: x.value if isinstance(x, LineageSourceEnum) else x, self.sources)
             )
-        if self.sources:
+        if self.entities:
             filter_request["LineageTypes"] = list(
                 map(lambda x: x.value if isinstance(x, LineageEntityEnum) else x, self.entities)
             )

tests/integ/sagemaker/lineage/conftest.py

@@ -38,6 +38,7 @@
 SLEEP_TIME_SECONDS = 1
 STATIC_PIPELINE_NAME = "SdkIntegTestStaticPipeline17"
 STATIC_ENDPOINT_NAME = "SdkIntegTestStaticEndpoint17"
+STATIC_MODEL_PACKAGE_GROUP_NAME = "SdkIntegTestStaticPipeline17ModelPackageGroup"
 
 
 @pytest.fixture
@@ -543,6 +544,29 @@ def static_endpoint_context(sagemaker_session, static_pipeline_execution_arn):
     )
 
 
+@pytest.fixture
+def static_model_package_group_context(sagemaker_session, static_pipeline_execution_arn):
+
+    model_package_group_arn = get_model_package_group_arn_from_static_pipeline(sagemaker_session)
+
+    contexts = sagemaker_session.sagemaker_client.list_contexts(SourceUri=model_package_group_arn)[
+        "ContextSummaries"
+    ]
+    if len(contexts) != 1:
+        raise (
+            Exception(
+                f"Got an unexpected number of Contexts for \
+                model package group {STATIC_MODEL_PACKAGE_GROUP_NAME} from pipeline \
+                execution {static_pipeline_execution_arn}. \
+                Expected 1 but got {len(contexts)}"
+            )
+        )
+
+    yield context.ModelPackageGroup.load(
+        contexts[0]["ContextName"], sagemaker_session=sagemaker_session
+    )
+
+
 @pytest.fixture
 def static_model_artifact(sagemaker_session, static_pipeline_execution_arn):
     model_package_arn = get_model_package_arn_from_static_pipeline(
@@ -590,6 +614,31 @@ def static_dataset_artifact(static_model_artifact, sagemaker_session):
     )
 
 
+@pytest.fixture
+def static_image_artifact(static_model_artifact, sagemaker_session):
+    dataset_associations = sagemaker_session.sagemaker_client.list_associations(
+        DestinationArn=static_model_artifact.artifact_arn, SourceType="Image"
+    )
+    if len(dataset_associations["AssociationSummaries"]) == 0:
+        # no directly associated dataset. work backwards from the model
+        model_associations = sagemaker_session.sagemaker_client.list_associations(
+            DestinationArn=static_model_artifact.artifact_arn, SourceType="Model"
+        )
+        training_job_associations = sagemaker_session.sagemaker_client.list_associations(
+            DestinationArn=model_associations["AssociationSummaries"][0]["SourceArn"],
+            SourceType="SageMakerTrainingJob",
+        )
+        dataset_associations = sagemaker_session.sagemaker_client.list_associations(
+            DestinationArn=training_job_associations["AssociationSummaries"][0]["SourceArn"],
+            SourceType="Image",
+        )
+
+    yield artifact.ImageArtifact.load(
+        dataset_associations["AssociationSummaries"][0]["SourceArn"],
+        sagemaker_session=sagemaker_session,
+    )
+
+
 def get_endpoint_arn_from_static_pipeline(sagemaker_session):
     try:
         endpoint_arn = sagemaker_session.sagemaker_client.describe_endpoint(
@@ -604,6 +653,15 @@ def get_endpoint_arn_from_static_pipeline(sagemaker_session):
         raise e
 
 
+def get_model_package_group_arn_from_static_pipeline(sagemaker_session):
+    static_model_package_group_arn = (
+        sagemaker_session.sagemaker_client.describe_model_package_group(
+            ModelPackageGroupName=STATIC_MODEL_PACKAGE_GROUP_NAME
+        )["ModelPackageGroupArn"]
+    )
+    return static_model_package_group_arn
+
+
 def get_model_package_arn_from_static_pipeline(pipeline_execution_arn, sagemaker_session):
     # get the model package ARN from the pipeline
     pipeline_execution_steps = sagemaker_session.sagemaker_client.list_pipeline_execution_steps(

tests/integ/sagemaker/lineage/test_endpoint_context.py

@@ -41,6 +41,22 @@ def test_training_job_arns(
         assert "training-job" in arn
 
 
+def test_processing_job_arns(static_endpoint_context):
+    processing_job_arns = static_endpoint_context.processing_job_arns()
+
+    assert len(processing_job_arns) > 0
+    for arn in processing_job_arns:
+        assert "processing-job" in arn
+
+
+def test_trial_components_arns(static_endpoint_context):
+    trial_components_arns = static_endpoint_context.trial_components_arns()
+
+    assert len(trial_components_arns) > 0
+    for arn in trial_components_arns:
+        assert "job" in arn
+
+
 def test_pipeline_execution_arn(static_endpoint_context, static_pipeline_execution_arn):
     pipeline_execution_arn = static_endpoint_context.pipeline_execution_arn()
 

tests/integ/sagemaker/lineage/test_model_package_group_context.py

@@ -0,0 +1,20 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module contains code to test SageMaker ``ModelPackageGroup``"""
+from __future__ import absolute_import
+
+
+def test_pipeline_execution_arn(static_model_package_group_context, static_pipeline_execution_arn):
+    pipeline_execution_arn = static_model_package_group_context.pipeline_execution_arn()
+
+    assert pipeline_execution_arn == static_pipeline_execution_arn

tests/unit/sagemaker/lineage/test_query.py

@@ -243,7 +243,7 @@ def test_vertex_to_object_artifact(sagemaker_session):
 def test_vertex_to_object_unconvertable(sagemaker_session):
     vertex = Vertex(
         arn="arn:aws:sagemaker:us-west-2:0123456789012:artifact/e66eef7f19c05e75284089183491bd4f",
-        lineage_entity=LineageEntityEnum.ACTION.value,
+        lineage_entity=LineageEntityEnum.TRIAL_COMPONENT.value,
         lineage_source=LineageSourceEnum.TENSORBOARD.value,
         sagemaker_session=sagemaker_session,
     )