sagemaker-core 1.0.47__py3-none-any.whl → 2.1.1__py3-none-any.whl

sagemaker/core/lineage/association.py

@@ -0,0 +1,190 @@
+# 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 create and manage SageMaker ``Artifact``."""
+from __future__ import absolute_import
+
+from typing import Optional, Iterator
+from datetime import datetime
+import logging
+
+from sagemaker.core.apiutils import _base_types
+from sagemaker.core.lineage import _api_types
+from sagemaker.core.lineage._api_types import AssociationSummary
+from sagemaker.core.common_utils import format_tags
+
+logger = logging.getLogger(__name__)
+
+
+class Association(_base_types.Record):
+    """An Amazon SageMaker artifact, which is part of a SageMaker lineage.
+
+    Examples:
+        .. code-block:: python
+
+            from sagemaker.lineage import association
+
+            my_association = association.Association.create(
+                source_arn=artifact_arn,
+                destination_arn=trial_component_arn,
+                association_type='ContributedTo')
+
+            for assoctn in association.Association.list():
+                print(assoctn)
+
+            my_association.delete()
+
+    Attributes:
+        source_arn (str): The ARN of the source entity.
+        source_type (str): The type of the source entity.
+        destination_arn (str): The ARN of the destination entity.
+        destination_type (str): The type of the destination entity.
+        association_type (str): the type of the association.
+    """
+
+    source_arn: str = None
+    destination_arn: str = None
+
+    _boto_create_method: str = "add_association"
+    _boto_delete_method: str = "delete_association"
+
+    _custom_boto_types: dict = {}
+
+    _boto_delete_members = [
+        "source_arn",
+        "destination_arn",
+    ]
+
+    def delete(self):
+        """Delete this Association from SageMaker."""
+        self._invoke_api(self._boto_delete_method, self._boto_delete_members)
+
+    def set_tag(self, tag=None):
+        """Add a tag to the object.
+
+        Args:
+            tag (obj): Key value pair to set tag.
+
+        Returns:
+            list({str:str}): a list of key value pairs
+        """
+        logger.warning(
+            "set_tag on Association is deprecated. Use set_tag on the source or destination\
+             entity instead."
+        )
+        return self._set_tags(resource_arn=self.source_arn, tags=[tag])
+
+    def set_tags(self, tags=None):
+        """Add tags to the object.
+
+        Args:
+            tags ([{key:value}]): list of key value pairs.
+
+        Returns:
+            list({str:str}): a list of key value pairs
+        """
+        logger.warning(
+            "set_tags on Association is deprecated. Use set_tags on the source or destination\
+            entity instead."
+        )
+        return self._set_tags(resource_arn=self.source_arn, tags=format_tags(tags))
+
+    @classmethod
+    def create(
+        cls,
+        source_arn: str,
+        destination_arn: str,
+        association_type: str = None,
+        sagemaker_session=None,
+    ) -> "Association":
+        """Add an association and return an ``Association`` object representing it.
+
+        Args:
+            source_arn (str): The ARN of the source.
+            destination_arn (str): The ARN of the destination.
+            association_type (str): The type of the association. ContributedTo, AssociatedWith,
+                DerivedFrom, or Produced.
+            sagemaker_session (sagemaker.session.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed. If not specified, one is created using the
+                default AWS configuration chain.
+
+        Returns:
+            association: A SageMaker ``Association`` object.
+        """
+        return super(Association, cls)._construct(
+            cls._boto_create_method,
+            source_arn=source_arn,
+            destination_arn=destination_arn,
+            association_type=association_type,
+            sagemaker_session=sagemaker_session,
+        )
+
+    @classmethod
+    def list(
+        cls,
+        source_arn: str = None,
+        destination_arn: str = None,
+        source_type: str = None,
+        destination_type: str = None,
+        association_type: str = None,
+        created_after: Optional[datetime] = None,
+        created_before: Optional[datetime] = None,
+        sort_by: Optional[str] = None,
+        sort_order: Optional[str] = None,
+        max_results: Optional[int] = None,
+        next_token: Optional[str] = None,
+        sagemaker_session=None,
+    ) -> Iterator[AssociationSummary]:
+        """Return a list of context summaries.
+
+        Args:
+            source_arn (str): The ARN of the source entity.
+            destination_arn (str): The ARN of the destination entity.
+            source_type (str): The type of the source entity.
+            destination_type (str): The type of the destination entity.
+            association_type (str): The type of the association.
+            created_after (datetime.datetime, optional): Return contexts created after this
+                instant.
+            created_before (datetime.datetime, optional): Return contexts created before this
+                instant.
+            sort_by (str, optional): Which property to sort results by.
+                One of 'SourceArn', 'CreatedBefore', 'CreatedAfter'
+            sort_order (str, optional): One of 'Ascending', or 'Descending'.
+            max_results (int, optional): maximum number of contexts to retrieve
+            next_token (str, optional): token for next page of results
+            sagemaker_session (sagemaker.session.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed. If not specified, one is created using the
+                default AWS configuration chain.
+
+        Returns:
+            collections.Iterator[AssociationSummary]: An iterator
+                over ``AssociationSummary`` objects.
+        """
+        return super(Association, cls)._list(
+            "list_associations",
+            _api_types.AssociationSummary.from_boto,
+            "AssociationSummaries",
+            source_arn=source_arn,
+            destination_arn=destination_arn,
+            source_type=source_type,
+            destination_type=destination_type,
+            association_type=association_type,
+            created_after=created_after,
+            created_before=created_before,
+            sort_by=sort_by,
+            sort_order=sort_order,
+            max_results=max_results,
+            next_token=next_token,
+            sagemaker_session=sagemaker_session,
+        )

sagemaker/core/lineage/context.py

@@ -0,0 +1,505 @@
+# 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 create and manage SageMaker ``Context``."""
+from __future__ import absolute_import
+
+from datetime import datetime
+from typing import Iterator, Optional, List
+
+from sagemaker.core.apiutils import _base_types
+from sagemaker.core.lineage import (
+    _api_types,
+    _utils,
+    association,
+)
+from sagemaker.core.lineage._api_types import ContextSummary
+from sagemaker.core.lineage.query import (
+    LineageQuery,
+    LineageFilter,
+    LineageSourceEnum,
+    LineageEntityEnum,
+    LineageQueryDirectionEnum,
+)
+from sagemaker.core.lineage.artifact import Artifact
+from sagemaker.core.lineage.action import Action
+from sagemaker.core.lineage.lineage_trial_component import LineageTrialComponent
+from sagemaker.core.common_utils import format_tags
+
+
+class Context(_base_types.Record):
+    """An Amazon SageMaker context, which is part of a SageMaker lineage.
+
+    Attributes:
+        context_arn (str): The ARN of the context.
+        context_name (str): The name of the context.
+        context_type (str): The type of the context.
+        description (str): A description of the context.
+        source (obj): The source of the context with a URI and type.
+        properties (dict): Dictionary of properties.
+        tags (List[dict[str, str]]): A list of tags to associate with the context.
+        creation_time (datetime): When the context was created.
+        created_by (obj): Contextual info on which account created the context.
+        last_modified_time (datetime): When the context was last modified.
+        last_modified_by (obj): Contextual info on which account created the context.
+    """
+
+    context_arn: str = None
+    context_name: str = None
+    context_type: str = None
+    properties: dict = None
+    tags: list = None
+    creation_time: datetime = None
+    created_by: str = None
+    last_modified_time: datetime = None
+    last_modified_by: str = None
+
+    _boto_load_method: str = "describe_context"
+    _boto_create_method: str = "create_context"
+    _boto_update_method: str = "update_context"
+    _boto_delete_method: str = "delete_context"
+
+    _custom_boto_types = {
+        "source": (_api_types.ContextSource, False),
+    }
+
+    _boto_update_members = [
+        "context_name",
+        "description",
+        "properties",
+        "properties_to_remove",
+    ]
+    _boto_delete_members = ["context_name"]
+
+    def save(self) -> "Context":
+        """Save the state of this Context to SageMaker.
+
+        Returns:
+            obj: boto API response.
+        """
+        return self._invoke_api(self._boto_update_method, self._boto_update_members)
+
+    def delete(self, disassociate: bool = False):
+        """Delete the context object.
+
+        Args:
+            disassociate (bool): When set to true, disassociate incoming and outgoing association.
+
+        Returns:
+            obj: boto API response.
+        """
+        if disassociate:
+            _utils._disassociate(
+                source_arn=self.context_arn, sagemaker_session=self.sagemaker_session
+            )
+            _utils._disassociate(
+                destination_arn=self.context_arn,
+                sagemaker_session=self.sagemaker_session,
+            )
+        return self._invoke_api(self._boto_delete_method, self._boto_delete_members)
+
+    def set_tag(self, tag=None):
+        """Add a tag to the object.
+
+        Args:
+            tag (obj): Key value pair to set tag.
+
+        Returns:
+            list({str:str}): a list of key value pairs
+        """
+        return self._set_tags(resource_arn=self.context_arn, tags=[tag])
+
+    def set_tags(self, tags=None):
+        """Add tags to the object.
+
+        Args:
+            tags ([{key:value}]): list of key value pairs.
+
+        Returns:
+            list({str:str}): a list of key value pairs
+        """
+        return self._set_tags(resource_arn=self.context_arn, tags=format_tags(tags))
+
+    @classmethod
+    def load(cls, context_name: str, sagemaker_session=None) -> "Context":
+        """Load an existing context and return an ``Context`` object representing it.
+
+        Examples:
+            .. code-block:: python
+
+                from sagemaker.lineage import context
+
+                my_context = context.Context.create(
+                    context_name='MyContext',
+                    context_type='Endpoint',
+                    source_uri='arn:aws:...')
+
+                my_context.properties["added"] = "property"
+                my_context.save()
+
+                for ctx in context.Context.list():
+                    print(ctx)
+
+                my_context.delete()
+
+            Args:
+                context_name (str): Name of the context
+                sagemaker_session (sagemaker.session.Session): Session object which
+                    manages interactions with Amazon SageMaker APIs and any other
+                    AWS services needed. If not specified, one is created using the
+                    default AWS configuration chain.
+
+            Returns:
+                Context: A SageMaker ``Context`` object
+        """
+        context = cls._construct(
+            cls._boto_load_method,
+            context_name=context_name,
+            sagemaker_session=sagemaker_session,
+        )
+        return context
+
+    @classmethod
+    def create(
+        cls,
+        context_name: str = None,
+        source_uri: str = None,
+        source_type: str = None,
+        context_type: str = None,
+        description: str = None,
+        properties: dict = None,
+        tags: dict = None,
+        sagemaker_session=None,
+    ) -> "Context":
+        """Create a context and return a ``Context`` object representing it.
+
+        Args:
+            context_name (str): The name of the context.
+            source_uri (str): The source URI of the context.
+            source_type (str): The type of the source.
+            context_type (str): The type of the context.
+            description (str): Description of the context.
+            properties (dict): Metadata associated with the context.
+            tags (dict): Tags to add to the context.
+            sagemaker_session (sagemaker.session.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed. If not specified, one is created using the
+                default AWS configuration chain.
+
+        Returns:
+            Context: A SageMaker ``Context`` object.
+        """
+        return super(Context, cls)._construct(
+            cls._boto_create_method,
+            context_name=context_name,
+            source=_api_types.ContextSource(source_uri=source_uri, source_type=source_type),
+            context_type=context_type,
+            description=description,
+            properties=properties,
+            tags=tags,
+            sagemaker_session=sagemaker_session,
+        )
+
+    @classmethod
+    def list(
+        cls,
+        source_uri: Optional[str] = None,
+        context_type: Optional[str] = None,
+        created_after: Optional[datetime] = None,
+        created_before: Optional[datetime] = None,
+        sort_by: Optional[str] = None,
+        sort_order: Optional[str] = None,
+        max_results: Optional[int] = None,
+        next_token: Optional[str] = None,
+        sagemaker_session=None,
+    ) -> Iterator[ContextSummary]:
+        """Return a list of context summaries.
+
+        Args:
+            source_uri (str, optional): A source URI.
+            context_type (str, optional): An context type.
+            created_before (datetime.datetime, optional): Return contexts created before this
+                instant.
+            created_after (datetime.datetime, optional): Return contexts created after this instant.
+            sort_by (str, optional): Which property to sort results by.
+                One of 'SourceArn', 'CreatedBefore', 'CreatedAfter'
+            sort_order (str, optional): One of 'Ascending', or 'Descending'.
+            max_results (int, optional): maximum number of contexts to retrieve
+            next_token (str, optional): token for next page of results
+            sagemaker_session (sagemaker.session.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed. If not specified, one is created using the
+                default AWS configuration chain.
+
+        Returns:
+            collections.Iterator[ContextSummary]: An iterator
+                over ``ContextSummary`` objects.
+        """
+        return super(Context, cls)._list(
+            "list_contexts",
+            _api_types.ContextSummary.from_boto,
+            "ContextSummaries",
+            source_uri=source_uri,
+            context_type=context_type,
+            created_before=created_before,
+            created_after=created_after,
+            sort_by=sort_by,
+            sort_order=sort_order,
+            max_results=max_results,
+            next_token=next_token,
+            sagemaker_session=sagemaker_session,
+        )
+
+    def actions(self, direction: LineageQueryDirectionEnum) -> List[Action]:
+        """Use the lineage query to retrieve actions that use this context.
+
+        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."""
+
+    def models(self) -> List[association.Association]:
+        """Use Lineage API to get all models deployed by this endpoint.
+
+        Returns:
+            list of Associations: Associations that destination represents an endpoint's model.
+        """
+        endpoint_actions: Iterator = association.Association.list(
+            sagemaker_session=self.sagemaker_session,
+            source_arn=self.context_arn,
+            destination_type="ModelDeployment",
+        )
+
+        model_list: list = [
+            model
+            for endpoint_action in endpoint_actions
+            for model in association.Association.list(
+                source_arn=endpoint_action.destination_arn,
+                destination_type="Model",
+                sagemaker_session=self.sagemaker_session,
+            )
+        ]
+        return model_list
+
+    def models_v2(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.DESCENDANTS
+    ) -> List[Artifact]:
+        """Use the lineage query to retrieve downstream model artifacts that use this endpoint.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of Artifacts: Artifacts representing a model.
+        """
+        # Firstly query out the model_deployment vertices
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.ACTION], sources=[LineageSourceEnum.MODEL_DEPLOYMENT]
+        )
+        model_deployment_query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.context_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        if not model_deployment_query_result:
+            return []
+
+        model_deployment_vertices: [] = model_deployment_query_result.vertices
+
+        # Secondary query model based on model deployment
+        model_vertices = []
+        for vertex in model_deployment_vertices:
+            query_result = LineageQuery(self.sagemaker_session).query(
+                start_arns=[vertex.arn],
+                query_filter=LineageFilter(
+                    entities=[LineageEntityEnum.ARTIFACT], sources=[LineageSourceEnum.MODEL]
+                ),
+                direction=LineageQueryDirectionEnum.DESCENDANTS,
+                include_edges=False,
+            )
+            model_vertices.extend(query_result.vertices)
+
+        if not model_vertices:
+            return []
+
+        model_artifacts = []
+        for vertex in model_vertices:
+            lineage_object = vertex.to_lineage_object()
+            model_artifacts.append(lineage_object)
+
+        return model_artifacts
+
+    def dataset_artifacts(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[Artifact]:
+        """Use the lineage query to retrieve datasets that use this endpoint.
+
+        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.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        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 processing_jobs(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[LineageTrialComponent]:
+        """Use the lineage query to retrieve processing jobs that use this endpoint.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of LineageTrialComponent: Lineage trial component that represent Processing jobs.
+        """
+        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,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]
+
+    def transform_jobs(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[LineageTrialComponent]:
+        """Use the lineage query to retrieve transform jobs that use this endpoint.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of LineageTrialComponent: Lineage trial component that represent 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]:
+        """Use the lineage query to retrieve trial components that use this endpoint.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of LineageTrialComponent: Lineage trial component.
+        """
+        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.
+        """
+        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 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.get("PipelineExecutionArn")