feature: Add support for SageMaker lineage queries in artifact, context and trial component

Author: yzhu0

src/sagemaker/lineage/action.py

@@ -16,7 +16,7 @@
 from typing import Optional, Iterator
 from datetime import datetime
 
-from sagemaker import Session
+from sagemaker.session import Session
 from sagemaker.apiutils import _base_types
 from sagemaker.lineage import _api_types, _utils
 from sagemaker.lineage._api_types import ActionSource, ActionSummary
@@ -116,7 +116,7 @@ def delete(self, disassociate: bool = False):
         self._invoke_api(self._boto_delete_method, self._boto_delete_members)
 
     @classmethod
-    def load(cls, action_name: str, sagemaker_session: Session = None) -> "Action":
+    def load(cls, action_name: str, sagemaker_session=None) -> "Action":
         """Load an existing action and return an ``Action`` object representing it.
 
         Args:

src/sagemaker/lineage/artifact.py

@@ -146,7 +146,7 @@ def downstream_trials(self, sagemaker_session=None) -> list:
         """Retrieve all trial runs which that use this artifact.
 
         Args:
-            sagemaker_session (obj): Sagemaker Sesssion to use. If not provided a default session
+            sagemaker_session (obj): Sagemaker Session to use. If not provided a default session
                 will be created.
 
         Returns:
@@ -159,6 +159,57 @@ def downstream_trials(self, sagemaker_session=None) -> list:
         )
         trial_component_arns: list = list(map(lambda x: x.destination_arn, outgoing_associations))
 
+        return self._get_trial_from_trial_component(trial_component_arns)
+
+    def downstream_trials_v2(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.DESCENDANTS
+    ) -> list:
+        """Retrieve all downstream trial runs which that use this artifact by using lineage query.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            [Trial]: A list of SageMaker `Trial` objects.
+        """
+        query_filter = LineageFilter(entities=[LineageEntityEnum.TRIAL_COMPONENT])
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.artifact_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        trial_component_arns: list = list(map(lambda x: x.arn, query_result.vertices))
+
+        return self._get_trial_from_trial_component(trial_component_arns)
+
+    def get_upstream_trials(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[str]:
+        """Retrieve all upstream trial runs which that use this artifact by using lineage query.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            [Trial]: A list of SageMaker `Trial` objects.
+        """
+        query_filter = LineageFilter(entities=[LineageEntityEnum.TRIAL_COMPONENT])
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.artifact_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        trial_component_arns: list = list(map(lambda x: x.arn, query_result.vertices))
+        return self._get_trial_from_trial_component(trial_component_arns)
+
+    def _get_trial_from_trial_component(self, trial_component_arns: list):
+        """Retrieve all upstream trial runs which that use the trial component arns.
+
+        Returns:
+            [Trial]: A list of SageMaker `Trial` objects.
+        """
         if not trial_component_arns:
             # no outgoing associations for this artifact
             return []
@@ -170,7 +221,7 @@ def downstream_trials(self, sagemaker_session=None) -> list:
         num_search_batches = math.ceil(len(trial_component_arns) % max_search_by_arn)
         trial_components: list = []
 
-        sagemaker_session = sagemaker_session or _utils.default_session()
+        sagemaker_session = self.sagemaker_session or _utils.default_session()
         sagemaker_client = sagemaker_session.sagemaker_client
 
         for i in range(num_search_batches):
@@ -335,6 +386,17 @@ def list(
             sagemaker_session=sagemaker_session,
         )
 
+    def get_artifacts(self, s3_uri: str):
+        """Return a list of artifact that use provided s3 uri.
+
+        Args:
+            s3_uri (str): A S3 URI.
+
+        Returns:
+            A list of ``Artifacts``
+        """
+        return self.sagemaker_session.sagemaker_client.list_artifacts(SourceUri=s3_uri)
+
 
 class ModelArtifact(Artifact):
     """A SageMaker lineage artifact representing a model.
@@ -349,7 +411,7 @@ def endpoints(self) -> list:
         """Get association summaries for endpoints deployed with this model.
 
         Returns:
-            [AssociationSummary]: A list of associations repesenting the endpoints using the model.
+            [AssociationSummary]: A list of associations representing the endpoints using the model.
         """
         endpoint_development_actions: Iterator = Association.list(
             source_arn=self.artifact_arn,
@@ -522,3 +584,75 @@ def endpoint_contexts(
         for vertex in query_result.vertices:
             endpoint_contexts.append(vertex.to_lineage_object())
         return endpoint_contexts
+
+    def get_upstream_datasets(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[Artifact]:
+        """Get upstream artifacts representing dataset from the dataset's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of Artifacts: Artifacts representing an dataset.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.ARTIFACT], sources=[LineageSourceEnum.DATASET]
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.artifact_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]
+
+    def get_downstream_datasets(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.DESCENDANTS
+    ) -> List[Artifact]:
+        """Get downstream artifacts representing dataset from the dataset's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of Artifacts: Artifacts representing an dataset.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.ARTIFACT], sources=[LineageSourceEnum.DATASET]
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.artifact_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]
+
+
+class ImageArtifact(Artifact):
+    """A SageMaker lineage artifact representing an image.
+
+    Common model specific lineage traversals to discover how the image is connected
+    to other entities.
+    """
+
+    def get_dataset(self, direction: LineageQueryDirectionEnum) -> List[Artifact]:
+        """Get artifacts representing dataset from the image artifact's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum): The query direction.
+
+        Returns:
+            list of Artifacts: Artifacts representing an dataset.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.ARTIFACT], sources=[LineageSourceEnum.DATASET]
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.artifact_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]

src/sagemaker/lineage/context.py

@@ -31,6 +31,8 @@
     LineageQueryDirectionEnum,
 )
 from sagemaker.lineage.artifact import Artifact
+from sagemaker.lineage.action import Action
+from sagemaker.lineage.lineage_trial_component import LineageTrialComponent
 
 
 class Context(_base_types.Record):
@@ -256,6 +258,24 @@ def list(
             sagemaker_session=sagemaker_session,
         )
 
+    def actions(self, direction: LineageQueryDirectionEnum) -> List[Action]:
+        """Get actions from the endpoint's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum): The query direction.
+
+        Returns:
+            list of Actions: Actions.
+        """
+        query_filter = LineageFilter(entities=[LineageEntityEnum.ACTION])
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.context_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]
+
 
 class EndpointContext(Context):
     """An Amazon SageMaker endpoint context, which is part of a SageMaker lineage."""
@@ -360,6 +380,9 @@ def training_job_arns(
     ) -> List[str]:
         """Get ARNs for all training jobs that appear in the endpoint's lineage.
 
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
         Returns:
             list of str: Training job ARNs.
         """
@@ -382,11 +405,86 @@ 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.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        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 transform_jobs(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[LineageTrialComponent]:
+        """Get ARNs for all transform jobs that appear in the endpoint's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of str: Transform jobs.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.TRIAL_COMPONENT], sources=[LineageSourceEnum.TRANSFORM_JOB]
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.context_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]
+
+    def trial_components(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[LineageTrialComponent]:
+        """Get ARNs for all trial components that appear in the endpoint's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        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,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]
+
     def pipeline_execution_arn(
         self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
     ) -> str:
         """Get the ARN for the pipeline execution associated with this endpoint (if any).
 
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
         Returns:
             str: A pipeline execution ARN.
         """