feature: Add support for SageMaker lineage queries

Co-authored-by: Payton Staub <[email protected]>

Author: 2 people

src/sagemaker/lineage/_utils.py

@@ -52,3 +52,15 @@ def get_module(module_name):
         return import_module(module_name)
     except ImportError:
         raise Exception("Cannot import module {}, please try again.".format(module_name))
+
+
+def get_resource_name_from_arn(arn):
+    """Extract the resource name from an ARN string.
+
+    Args:
+        arn (str): An ARN.
+
+    Returns:
+        str: The resource name.
+    """
+    return arn.split(":", 5)[5].split("/", 1)[1]

src/sagemaker/lineage/artifact.py

@@ -17,12 +17,19 @@
 import math
 
 from datetime import datetime
-from typing import Iterator, Union, Any, Optional
+from typing import Iterator, Union, Any, Optional, List
 
 from sagemaker.apiutils import _base_types, _utils
 from sagemaker.lineage import _api_types
 from sagemaker.lineage._api_types import ArtifactSource, ArtifactSummary
-from sagemaker.lineage._utils import get_module, _disassociate
+from sagemaker.lineage.query import (
+    LineageQuery,
+    LineageFilter,
+    LineageSourceEnum,
+    LineageEntityEnum,
+    LineageQueryDirectionEnum,
+)
+from sagemaker.lineage._utils import get_module, _disassociate, get_resource_name_from_arn
 from sagemaker.lineage.association import Association
 
 LOGGER = logging.getLogger("sagemaker")
@@ -333,11 +340,13 @@ class ModelArtifact(Artifact):
     """A SageMaker lineage artifact representing a model.
 
     Common model specific lineage traversals to discover how the model is connected
-    to otherentities.
+    to other entities.
     """
 
+    from sagemaker.lineage.context import Context
+
     def endpoints(self) -> list:
-        """Given a model artifact, get all associated endpoint context.
+        """Get association summaries for endpoints deployed with this model.
 
         Returns:
             [AssociationSummary]: A list of associations repesenting the endpoints using the model.
@@ -359,6 +368,104 @@ def endpoints(self) -> list:
         ]
         return endpoint_context_list
 
+    def endpoint_contexts(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.DESCENDANTS
+    ) -> List[Context]:
+        """Get contexts representing endpoints from the models's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of Contexts: Contexts representing an endpoint.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.CONTEXT], sources=[LineageSourceEnum.ENDPOINT]
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.artifact_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+
+        endpoint_contexts = []
+        for vertex in query_result.vertices:
+            endpoint_contexts.append(vertex.to_lineage_object())
+        return endpoint_contexts
+
+    def dataset_artifacts(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[Artifact]:
+        """Get artifacts representing datasets from the model's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of Artifacts: Artifacts representing a 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,
+        )
+
+        dataset_artifacts = []
+        for vertex in query_result.vertices:
+            dataset_artifacts.append(vertex.to_lineage_object())
+        return dataset_artifacts
+
+    def training_job_arns(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[str]:
+        """Get ARNs for all training jobs that appear in the model's lineage.
+
+        Returns:
+            list of str: Training job ARNs.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.TRIAL_COMPONENT], sources=[LineageSourceEnum.TRAINING_JOB]
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.artifact_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+
+        training_job_arns = []
+        for vertex in query_result.vertices:
+            trial_component_name = get_resource_name_from_arn(vertex.arn)
+            trial_component = self.sagemaker_session.sagemaker_client.describe_trial_component(
+                TrialComponentName=trial_component_name
+            )
+            training_job_arns.append(trial_component["Source"]["SourceArn"])
+        return training_job_arns
+
+    def pipeline_execution_arn(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> str:
+        """Get the ARN for the pipeline execution associated with this model (if any).
+
+        Returns:
+            str: A pipeline execution ARN.
+        """
+        training_job_arns = self.training_job_arns(direction=direction)
+        for training_job_arn in training_job_arns:
+            tags = self.sagemaker_session.sagemaker_client.list_tags(ResourceArn=training_job_arn)[
+                "Tags"
+            ]
+            for tag in tags:
+                if tag["Key"] == "sagemaker:pipeline-execution-arn":
+                    return tag["Value"]
+
+        return None
+
 
 class DatasetArtifact(Artifact):
     """A SageMaker Lineage artifact representing a dataset.
@@ -367,7 +474,9 @@ class DatasetArtifact(Artifact):
     connect to related entities.
     """
 
-    def trained_models(self) -> list:
+    from sagemaker.lineage.context import Context
+
+    def trained_models(self) -> List[Association]:
         """Given a dataset artifact, get associated trained models.
 
         Returns:
@@ -387,3 +496,29 @@ def trained_models(self) -> list:
                 result.extend(models)
 
         return result
+
+    def endpoint_contexts(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.DESCENDANTS
+    ) -> List[Context]:
+        """Get contexts representing endpoints from the dataset's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of Contexts: Contexts representing an endpoint.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.CONTEXT], sources=[LineageSourceEnum.ENDPOINT]
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.artifact_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+
+        endpoint_contexts = []
+        for vertex in query_result.vertices:
+            endpoint_contexts.append(vertex.to_lineage_object())
+        return endpoint_contexts

src/sagemaker/lineage/context.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from datetime import datetime
-from typing import Iterator, Optional
+from typing import Iterator, Optional, List
 
 from sagemaker.apiutils import _base_types
 from sagemaker.lineage import (
@@ -23,6 +23,14 @@
     association,
 )
 from sagemaker.lineage._api_types import ContextSummary
+from sagemaker.lineage.query import (
+    LineageQuery,
+    LineageFilter,
+    LineageSourceEnum,
+    LineageEntityEnum,
+    LineageQueryDirectionEnum,
+)
+from sagemaker.lineage.artifact import Artifact
 
 
 class Context(_base_types.Record):
@@ -252,7 +260,7 @@ def list(
 class EndpointContext(Context):
     """An Amazon SageMaker endpoint context, which is part of a SageMaker lineage."""
 
-    def models(self) -> list:
+    def models(self) -> List[association.Association]:
         """Get all models deployed by all endpoint versions of the endpoint.
 
         Returns:
@@ -274,3 +282,72 @@ def models(self) -> list:
             )
         ]
         return model_list
+
+    def dataset_artifacts(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[Artifact]:
+        """Get artifacts representing datasets from the endpoint's lineage.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of Artifacts: Artifacts representing a dataset.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.ARTIFACT], sources=[LineageSourceEnum.DATASET]
+        )
+        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 training_job_arns(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[str]:
+        """Get ARNs for all training jobs that appear in the endpoint's lineage.
+
+        Returns:
+            list of str: Training job ARNs.
+        """
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.TRIAL_COMPONENT], sources=[LineageSourceEnum.TRAINING_JOB]
+        )
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.context_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+
+        training_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
+            )
+            training_job_arns.append(trial_component["Source"]["SourceArn"])
+        return training_job_arns
+
+    def pipeline_execution_arn(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> str:
+        """Get the ARN for the pipeline execution associated with this endpoint (if any).
+
+        Returns:
+            str: A pipeline execution ARN.
+        """
+        training_job_arns = self.training_job_arns(direction=direction)
+        for training_job_arn in training_job_arns:
+            tags = self.sagemaker_session.sagemaker_client.list_tags(ResourceArn=training_job_arn)[
+                "Tags"
+            ]
+            for tag in tags:
+                if tag["Key"] == "sagemaker:pipeline-execution-arn":
+                    return tag["Value"]
+
+        return None