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

sagemaker/core/lineage/_api_types.py

@@ -0,0 +1,239 @@
+# 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.
+"""Contains API objects for SageMaker Lineage."""
+from __future__ import absolute_import
+
+from sagemaker.core.apiutils import _base_types
+
+
+class ArtifactSource(_base_types.ApiObject):
+    """ArtifactSource.
+
+    Attributes:
+        source_uri (str): The URI of the source.
+        source_types(list[ArtifactSourceType]): List of source types
+    """
+
+    source_uri = None
+    source_types = None
+
+    def __init__(self, source_uri=None, source_types=None, **kwargs):
+        """Initialize ArtifactSource.
+
+        Args:
+            source_uri (str): Source S3 URI of the artifact.
+            source_types (array): Array of artifact source types.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super(ArtifactSource, self).__init__(
+            source_uri=source_uri, source_types=source_types, **kwargs
+        )
+
+
+class ArtifactSourceType(_base_types.ApiObject):
+    """ArtifactSourceType.
+
+    Attributes:
+        source_id_type (str): The source id type of artifact source.
+        value(str): The value of source
+    """
+
+    source_id_type = None
+    value = None
+
+    def __init__(self, source_id_type=None, value=None, **kwargs):
+        """Initialize ArtifactSourceType.
+
+        Args:
+            source_id_type (str): The type of the source id.
+            value (str): The source id.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super(ArtifactSourceType, self).__init__(
+            source_id_type=source_id_type, value=value, **kwargs
+        )
+
+
+class ActionSource(_base_types.ApiObject):
+    """ActionSource.
+
+    Attributes:
+        source_uri (str): The URI of the source.
+        source_type (str):  The type of the source URI.
+    """
+
+    source_uri = None
+    source_type = None
+
+    def __init__(self, source_uri=None, source_type=None, **kwargs):
+        """Initialize ActionSource.
+
+        Args:
+            source_uri (str): The URI of the source.
+            source_type (str): The type of the source.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super(ActionSource, self).__init__(source_uri=source_uri, source_type=source_type, **kwargs)
+
+
+class ContextSource(_base_types.ApiObject):
+    """ContextSource.
+
+    Attributes:
+        source_uri (str): The URI of the source.
+        source_type (str): The type of the source.
+    """
+
+    source_uri = None
+    source_type = None
+
+    def __init__(self, source_uri=None, source_type=None, **kwargs):
+        """Initialize ContextSource.
+
+        Args:
+           source_uri (str): The URI of the source.
+           source_type (str): The type of the source.
+           **kwargs: Arbitrary keyword arguments.
+        """
+        super(ContextSource, self).__init__(
+            source_uri=source_uri, source_type=source_type, **kwargs
+        )
+
+
+class ArtifactSummary(_base_types.ApiObject):
+    """Summary model of an Artifact.
+
+    Attributes:
+        artifact_arn (str): ARN of artifact.
+        artifact_name (str): Name of artifact.
+        source (obj): Source of artifact.
+        artifact_type (str): Type of artifact.
+        creation_time (datetime): Creation time.
+        last_modified_time (datetime): Date last modified.
+    """
+
+    _custom_boto_types = {
+        "source": (ArtifactSource, False),
+    }
+    artifact_arn = None
+    artifact_name = None
+    source = None
+    artifact_type = None
+    creation_time = None
+    last_modified_time = None
+
+
+class ActionSummary(_base_types.ApiObject):
+    """Summary model of an action.
+
+    Attributes:
+        action_arn (str): ARN of action.
+        action_name (str): Name of action.
+        source (obj): Source of action.
+        action_type (str): Type of action.
+        status (str): The status of the action.
+        creation_time (datetime): Creation time.
+        last_modified_time (datetime): Date last modified.
+    """
+
+    _custom_boto_types = {
+        "source": (ActionSource, False),
+    }
+    action_arn = None
+    action_name = None
+    source = None
+    action_type = None
+    status = None
+    creation_time = None
+    last_modified_time = None
+
+
+class ContextSummary(_base_types.ApiObject):
+    """Summary model of an context.
+
+    Attributes:
+        context_arn (str): ARN of context.
+        context_name (str): Name of context.
+        source (obj): Source of context.
+        context_type (str): Type of context.
+        creation_time (datetime): Creation time.
+        last_modified_time (datetime): Date last modified.
+    """
+
+    _custom_boto_types = {
+        "source": (ContextSource, False),
+    }
+    context_arn = None
+    context_name = None
+    source = None
+    context_type = None
+    creation_time = None
+    last_modified_time = None
+
+
+class UserContext(_base_types.ApiObject):
+    """Summary model of a user context.
+
+    Attributes:
+        user_profile_arn (str): User profile ARN.
+        user_profile_name (str): User profile name.
+        domain_id (str): DomainId.
+    """
+
+    user_profile_arn = None
+    user_profile_name = None
+    domain_id = None
+
+    def __init__(self, user_profile_arn=None, user_profile_name=None, domain_id=None, **kwargs):
+        """Initialize UserContext.
+
+        Args:
+            user_profile_arn (str): User profile ARN.
+            user_profile_name (str): User profile name.
+            domain_id (str): DomainId.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super(UserContext, self).__init__(
+            user_profile_arn=user_profile_arn,
+            user_profile_name=user_profile_name,
+            domain_id=domain_id,
+            **kwargs,
+        )
+
+
+class AssociationSummary(_base_types.ApiObject):
+    """Summary model of an association.
+
+    Attributes:
+        source_arn (str): ARN of source entity.
+        source_name (str): Name of the source entity.
+        destination_arn (str): ARN of the destination entity.
+        destination_name (str): Name of the destination entity.
+        source_type (obj): Type of the source entity.
+        destination_type (str): Type of destination entity.
+        association_type (str): The type of the association.
+        creation_time (datetime): Creation time.
+        created_by (obj): Context on creator.
+    """
+
+    _custom_boto_types = {
+        "created_by": (UserContext, False),
+    }
+    source_arn = None
+    source_name = None
+    destination_arn = None
+    destination_name = None
+    source_type = None
+    destination_type = None
+    association_type = None
+    creation_time = None

sagemaker/core/lineage/_utils.py

@@ -0,0 +1,49 @@
+# 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.
+"""SageMaker lineage utility methods."""
+from __future__ import absolute_import
+from sagemaker.core.lineage import association
+
+
+def _disassociate(source_arn=None, destination_arn=None, sagemaker_session=None):
+    """Remove the association.
+
+    Remove incoming association when source_arn is provided, remove outgoing association when
+    destination_arn is provided.
+    """
+    association_summaries = association.Association.list(
+        source_arn=source_arn,
+        destination_arn=destination_arn,
+        sagemaker_session=sagemaker_session,
+    )
+
+    for association_summary in association_summaries:
+
+        curr_association = association.Association(
+            sagemaker_session=sagemaker_session,
+            source_arn=association_summary.source_arn,
+            destination_arn=association_summary.destination_arn,
+        )
+        curr_association.delete()
+
+
+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]

sagemaker/core/lineage/action.py

@@ -0,0 +1,345 @@
+# 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 ``Actions``."""
+from __future__ import absolute_import
+
+from typing import Optional, Iterator, List
+from datetime import datetime
+
+from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.apiutils import _base_types
+from sagemaker.core.lineage import _api_types, _utils
+from sagemaker.core.lineage._api_types import ActionSource, ActionSummary
+from sagemaker.core.lineage.artifact import Artifact
+from sagemaker.core.common_utils import format_tags
+
+from sagemaker.core.lineage.query import (
+    LineageQuery,
+    LineageFilter,
+    LineageSourceEnum,
+    LineageEntityEnum,
+    LineageQueryDirectionEnum,
+)
+
+
+class Action(_base_types.Record):
+    """An Amazon SageMaker action, which is part of a SageMaker lineage.
+
+    Examples:
+        .. code-block:: python
+
+            from sagemaker.lineage import action
+
+            my_action = action.Action.create(
+                action_name='MyAction',
+                action_type='EndpointDeployment',
+                source_uri='s3://...')
+
+            my_action.properties["added"] = "property"
+            my_action.save()
+
+            for actn in action.Action.list():
+                print(actn)
+
+            my_action.delete()
+
+    Attributes:
+        action_arn (str): The ARN of the action.
+        action_name (str): The name of the action.
+        action_type (str): The type of the action.
+        description (str): A description of the action.
+        status (str): The status of the action.
+        source (obj): The source of the action with a URI and type.
+        properties (dict): Dictionary of properties.
+        tags (List[dict[str, str]]): A list of tags to associate with the action.
+        creation_time (datetime): When the action was created.
+        created_by (obj): Contextual info on which account created the action.
+        last_modified_time (datetime): When the action was last modified.
+        last_modified_by (obj): Contextual info on which account created the action.
+    """
+
+    action_arn: str = None
+    action_name: str = None
+    action_type: str = None
+    description: str = None
+    status: str = None
+    source: ActionSource = None
+    properties: dict = None
+    properties_to_remove: list = None
+    tags: list = None
+    creation_time: datetime = None
+    created_by: str = None
+    last_modified_time: datetime = None
+    last_modified_by: str = None
+
+    _boto_create_method: str = "create_action"
+    _boto_load_method: str = "describe_action"
+    _boto_update_method: str = "update_action"
+    _boto_delete_method: str = "delete_action"
+
+    _boto_update_members = [
+        "action_name",
+        "description",
+        "status",
+        "properties",
+        "properties_to_remove",
+    ]
+
+    _boto_delete_members = ["action_name"]
+
+    _custom_boto_types = {"source": (_api_types.ActionSource, False)}
+
+    def save(self) -> "Action":
+        """Save the state of this Action to SageMaker.
+
+        Returns:
+            Action: A SageMaker ``Action``object.
+        """
+        return self._invoke_api(self._boto_update_method, self._boto_update_members)
+
+    def delete(self, disassociate: bool = False):
+        """Delete the action.
+
+        Args:
+            disassociate (bool): When set to true, disassociate incoming and outgoing association.
+
+        """
+        if disassociate:
+            _utils._disassociate(
+                source_arn=self.action_arn, sagemaker_session=self.sagemaker_session
+            )
+            _utils._disassociate(
+                destination_arn=self.action_arn,
+                sagemaker_session=self.sagemaker_session,
+            )
+
+        self._invoke_api(self._boto_delete_method, self._boto_delete_members)
+
+    @classmethod
+    def load(cls, action_name: str, sagemaker_session=None) -> "Action":
+        """Load an existing action and return an ``Action`` object representing it.
+
+        Args:
+            action_name (str): Name of the action
+            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:
+            Action: A SageMaker ``Action`` object
+        """
+        result = cls._construct(
+            cls._boto_load_method,
+            action_name=action_name,
+            sagemaker_session=sagemaker_session,
+        )
+        return result
+
+    def set_tag(self, tag=None):
+        """Add a tag to the object.
+
+        Args:
+
+        Returns:
+            list({str:str}): a list of key value pairs
+        """
+        return self._set_tags(resource_arn=self.action_arn, tags=[tag])
+
+    def set_tags(self, tags=None):
+        """Add tags to the object.
+
+        Args:
+            tags (Optional[Tags]): list of key value pairs.
+
+        Returns:
+            list({str:str}): a list of key value pairs
+        """
+        return self._set_tags(resource_arn=self.action_arn, tags=format_tags(tags))
+
+    @classmethod
+    def create(
+        cls,
+        action_name: str = None,
+        source_uri: str = None,
+        source_type: str = None,
+        action_type: str = None,
+        description: str = None,
+        status: str = None,
+        properties: dict = None,
+        tags: dict = None,
+        sagemaker_session: Session = None,
+    ) -> "Action":
+        """Create an action and return an ``Action`` object representing it.
+
+        Args:
+            action_name (str): Name of the action
+            source_uri (str): Source URI of the action
+            source_type (str): Source type of the action
+            action_type (str): The type of the action
+            description (str): Description of the action
+            status (str): Status of the action.
+            properties (dict): key/value properties
+            tags (dict): AWS tags for the action
+            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:
+            Action: A SageMaker ``Action`` object.
+        """
+        return super(Action, cls)._construct(
+            cls._boto_create_method,
+            action_name=action_name,
+            source=_api_types.ActionSource(source_uri=source_uri, source_type=source_type),
+            action_type=action_type,
+            description=description,
+            status=status,
+            properties=properties,
+            tags=tags,
+            sagemaker_session=sagemaker_session,
+        )
+
+    @classmethod
+    def list(
+        cls,
+        source_uri: Optional[str] = None,
+        action_type: Optional[str] = None,
+        created_after: Optional[datetime] = None,
+        created_before: Optional[datetime] = None,
+        sort_by: Optional[str] = None,
+        sort_order: Optional[str] = None,
+        sagemaker_session: Session = None,
+        max_results: Optional[int] = None,
+        next_token: Optional[str] = None,
+    ) -> Iterator[ActionSummary]:
+        """Return a list of action summaries.
+
+        Args:
+            source_uri (str, optional): A source URI.
+            action_type (str, optional): An action type.
+            created_before (datetime.datetime, optional): Return actions created before this
+                instant.
+            created_after (datetime.datetime, optional): Return actions 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 actions 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[ActionSummary]: An iterator
+                over ``ActionSummary`` objects.
+        """
+        return super(Action, cls)._list(
+            "list_actions",
+            _api_types.ActionSummary.from_boto,
+            "ActionSummaries",
+            source_uri=source_uri,
+            action_type=action_type,
+            created_before=created_before,
+            created_after=created_after,
+            sort_by=sort_by,
+            sort_order=sort_order,
+            sagemaker_session=sagemaker_session,
+            max_results=max_results,
+            next_token=next_token,
+        )
+
+    def artifacts(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.BOTH
+    ) -> List[Artifact]:
+        """Use a lineage query to retrieve all artifacts that use this action.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of Artifacts: Artifacts.
+        """
+        query_filter = LineageFilter(entities=[LineageEntityEnum.ARTIFACT])
+        query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.action_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]
+
+
+class ModelPackageApprovalAction(Action):
+    """An Amazon SageMaker model package approval action, which is part of a SageMaker lineage."""
+
+    def datasets(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
+    ) -> List[Artifact]:
+        """Use a lineage query to retrieve all upstream datasets that use this action.
+
+        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.action_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]
+
+    def model_package(self):
+        """Get model package from model package approval action.
+
+        Returns:
+            Model package.
+        """
+        source_uri = self.source.source_uri
+        if source_uri is None:
+            return None
+
+        model_package_name = source_uri.split("/")[1]
+        return self.sagemaker_session.sagemaker_client.describe_model_package(
+            ModelPackageName=model_package_name
+        )
+
+    def endpoints(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.DESCENDANTS
+    ) -> List:
+        """Use a lineage query to retrieve downstream endpoint contexts that use this action.
+
+        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.action_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        return [vertex.to_lineage_object() for vertex in query_result.vertices]