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

sagemaker/core/collection.py

@@ -0,0 +1,467 @@
+# 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 related to Amazon SageMaker Collections in the Model Registry.
+
+Use these methods to help you create and maintain your Collections.
+"""
+
+from __future__ import absolute_import
+import json
+import time
+from typing import List
+
+
+from botocore.exceptions import ClientError
+from sagemaker.core.helper.session_helper import Session
+
+
+class Collection(object):
+    """Sets up an Amazon SageMaker Collection."""
+
+    def __init__(self, sagemaker_session):
+        """Initializes a Collection instance.
+
+        A Collection is a logical grouping of Model Groups.
+
+        Args:
+            sagemaker_session (sagemaker.session.Session): A Session object which
+                manages interactions between Amazon SageMaker APIs and other
+                AWS services needed. If unspecified, a session is created using
+                the default AWS configuration chain.
+        """
+
+        self.sagemaker_session = sagemaker_session or Session()
+
+    def _check_access_error(self, err: ClientError):
+        """Checks if the error is related to the access error and provide the relevant message.
+
+        Args:
+            err: The client error to check.
+        """
+
+        error_code = err.response["Error"]["Code"]
+        if error_code == "AccessDeniedException":
+            raise Exception(
+                f"{error_code}: This account needs to attach a custom policy "
+                "to the user role to gain access to Collections. Refer - "
+                "https://docs.aws.amazon.com/sagemaker/latest/dg/modelcollections-permissions.html"
+            )
+
+    def _add_model_group(self, model_package_group, tag_rule_key, tag_rule_value):
+        """Adds a Model Group to a Collection.
+
+        Args:
+            model_package_group (str): The name of the Model Group.
+            tag_rule_key (str): The tag key of the destination collection.
+            tag_rule_value (str): The tag value of the destination collection.
+        """
+        model_group_details = self.sagemaker_session.sagemaker_client.describe_model_package_group(
+            ModelPackageGroupName=model_package_group
+        )
+        self.sagemaker_session.sagemaker_client.add_tags(
+            ResourceArn=model_group_details["ModelPackageGroupArn"],
+            Tags=[
+                {
+                    "Key": tag_rule_key,
+                    "Value": tag_rule_value,
+                }
+            ],
+        )
+
+    def _remove_model_group(self, model_package_group, tag_rule_key):
+        """Removes a Model Group from a Collection.
+
+        Args:
+            model_package_group (str): The name of the Model Group
+            tag_rule_key (str): The tag key of the Collection from which to remove the Model Group.
+        """
+        model_group_details = self.sagemaker_session.sagemaker_client.describe_model_package_group(
+            ModelPackageGroupName=model_package_group
+        )
+        self.sagemaker_session.sagemaker_client.delete_tags(
+            ResourceArn=model_group_details["ModelPackageGroupArn"], TagKeys=[tag_rule_key]
+        )
+
+    def create(self, collection_name: str, parent_collection_name: str = None):
+        """Creates a Collection.
+
+        Args:
+            collection_name (str): The name of the Collection to create.
+            parent_collection_name (str): The name of the parent Collection.
+                Is ``None`` if the Collection is created at the root level.
+        """
+
+        tag_rule_key = f"sagemaker:collection-path:{int(time.time() * 1000)}"
+        tags_on_collection = {
+            "sagemaker:collection": "true",
+            "sagemaker:collection-path:root": "true",
+        }
+        tag_rule_values = [collection_name]
+
+        if parent_collection_name is not None:
+            parent_tag_rules = self._get_collection_tag_rule(collection_name=parent_collection_name)
+            parent_tag_rule_key = parent_tag_rules["tag_rule_key"]
+            parent_tag_value = parent_tag_rules["tag_rule_value"]
+            tags_on_collection = {
+                parent_tag_rule_key: parent_tag_value,
+                "sagemaker:collection": "true",
+            }
+            tag_rule_values = [f"{parent_tag_value}/{collection_name}"]
+        try:
+            resource_filters = [
+                "AWS::SageMaker::ModelPackageGroup",
+                "AWS::ResourceGroups::Group",
+            ]
+
+            tag_filters = [
+                {
+                    "Key": tag_rule_key,
+                    "Values": tag_rule_values,
+                }
+            ]
+            resource_query = {
+                "Query": json.dumps(
+                    {"ResourceTypeFilters": resource_filters, "TagFilters": tag_filters}
+                ),
+                "Type": "TAG_FILTERS_1_0",
+            }
+            collection_create_response = self.sagemaker_session.create_group(
+                collection_name, resource_query, tags_on_collection
+            )
+            return {
+                "Name": collection_create_response["Group"]["Name"],
+                "Arn": collection_create_response["Group"]["GroupArn"],
+            }
+        except ClientError as e:
+            message = e.response["Error"]["Message"]
+            error_code = e.response["Error"]["Code"]
+
+            if error_code == "BadRequestException" and "group already exists" in message:
+                raise ValueError("Collection with the given name already exists")
+            self._check_access_error(err=e)
+            raise
+
+    def delete(self, collections: List[str]):
+        """Deletes a list of Collections.
+
+        Args:
+            collections (List[str]): A list of Collections to delete.
+                Only deletes a Collection if it is empty.
+        """
+
+        if len(collections) > 10:
+            raise ValueError("Can delete upto 10 collections at a time")
+
+        delete_collection_failures = []
+        deleted_collection = []
+        collection_filter = [
+            {
+                "Name": "resource-type",
+                "Values": ["AWS::ResourceGroups::Group", "AWS::SageMaker::ModelPackageGroup"],
+            },
+        ]
+
+        # loops over the list of collection and deletes one at a time.
+        for collection in collections:
+            try:
+                collection_details = self.sagemaker_session.list_group_resources(
+                    group=collection, filters=collection_filter
+                )
+            except ClientError as e:
+                self._check_access_error(err=e)
+                delete_collection_failures.append(
+                    {"collection": collection, "message": e.response["Error"]["Message"]}
+                )
+                continue
+            if collection_details.get("Resources") and len(collection_details["Resources"]) > 0:
+                delete_collection_failures.append(
+                    {"collection": collection, "message": "Validation error: Collection not empty"}
+                )
+            else:
+                try:
+                    self.sagemaker_session.delete_resource_group(group=collection)
+                    deleted_collection.append(collection)
+                except ClientError as e:
+                    self._check_access_error(err=e)
+                    delete_collection_failures.append(
+                        {"collection": collection, "message": e.response["Error"]["Message"]}
+                    )
+        return {
+            "deleted_collections": deleted_collection,
+            "delete_collection_failures": delete_collection_failures,
+        }
+
+    def _get_collection_tag_rule(self, collection_name: str):
+        """Returns the tag rule key and value for a Collection."""
+
+        if collection_name is not None:
+            try:
+                group_query = self.sagemaker_session.get_resource_group_query(group=collection_name)
+            except ClientError as e:
+                error_code = e.response["Error"]["Code"]
+
+                if error_code == "NotFoundException":
+                    raise ValueError(f"Cannot find collection: {collection_name}")
+                self._check_access_error(err=e)
+                raise
+            if group_query.get("GroupQuery"):
+                tag_rule_query = json.loads(
+                    group_query["GroupQuery"].get("ResourceQuery", {}).get("Query", "")
+                )
+                tag_rule = tag_rule_query.get("TagFilters", [])[0]
+                if not tag_rule:
+                    raise "Unsupported parent_collection_name"
+                tag_rule_value = tag_rule["Values"][0]
+                tag_rule_key = tag_rule["Key"]
+
+            return {
+                "tag_rule_key": tag_rule_key,
+                "tag_rule_value": tag_rule_value,
+            }
+        raise ValueError("Collection name is required")
+
+    def add_model_groups(self, collection_name: str, model_groups: List[str]):
+        """Adds a list of Model Groups to a Collection.
+
+        Args:
+            collection_name (str): The name of the Collection.
+            model_groups (List[str]): The names of the Model Groups to add to the Collection.
+        """
+        if len(model_groups) > 10:
+            raise Exception("Model groups can have a maximum length of 10")
+        tag_rules = self._get_collection_tag_rule(collection_name=collection_name)
+        tag_rule_key = tag_rules["tag_rule_key"]
+        tag_rule_value = tag_rules["tag_rule_value"]
+
+        add_groups_success = []
+        add_groups_failure = []
+        if tag_rule_key is not None and tag_rule_value is not None:
+            for model_group in model_groups:
+                try:
+                    self._add_model_group(
+                        model_package_group=model_group,
+                        tag_rule_key=tag_rule_key,
+                        tag_rule_value=tag_rule_value,
+                    )
+                    add_groups_success.append(model_group)
+                except ClientError as e:
+                    self._check_access_error(err=e)
+                    message = e.response["Error"]["Message"]
+                    add_groups_failure.append(
+                        {
+                            "model_group": model_group,
+                            "failure_reason": message,
+                        }
+                    )
+        return {
+            "added_groups": add_groups_success,
+            "failure": add_groups_failure,
+        }
+
+    def remove_model_groups(self, collection_name: str, model_groups: List[str]):
+        """Removes a list of Model Groups from a Collection.
+
+        Args:
+            collection_name (str): The name of the Collection.
+            model_groups (List[str]): The names of the Model Groups to remove.
+        """
+
+        if len(model_groups) > 10:
+            raise Exception("Model groups can have a maximum length of 10")
+        tag_rules = self._get_collection_tag_rule(collection_name=collection_name)
+
+        tag_rule_key = tag_rules["tag_rule_key"]
+        tag_rule_value = tag_rules["tag_rule_value"]
+
+        remove_groups_success = []
+        remove_groups_failure = []
+        if tag_rule_key is not None and tag_rule_value is not None:
+            for model_group in model_groups:
+                try:
+                    self._remove_model_group(
+                        model_package_group=model_group,
+                        tag_rule_key=tag_rule_key,
+                    )
+                    remove_groups_success.append(model_group)
+                except ClientError as e:
+                    self._check_access_error(err=e)
+                    message = e.response["Error"]["Message"]
+                    remove_groups_failure.append(
+                        {
+                            "model_group": model_group,
+                            "failure_reason": message,
+                        }
+                    )
+        return {
+            "removed_groups": remove_groups_success,
+            "failure": remove_groups_failure,
+        }
+
+    def move_model_group(
+        self, source_collection_name: str, model_group: str, destination_collection_name: str
+    ):
+        """Moves a Model Group from one Collection to another.
+
+        Args:
+            source_collection_name (str): The name of the source Collection.
+            model_group (str): The name of the Model Group to move.
+            destination_collection_name (str): The name of the destination Collection.
+        """
+        remove_details = self.remove_model_groups(
+            collection_name=source_collection_name, model_groups=[model_group]
+        )
+        if len(remove_details["failure"]) == 1:
+            raise Exception(remove_details["failure"][0]["failure"])
+
+        added_details = self.add_model_groups(
+            collection_name=destination_collection_name, model_groups=[model_group]
+        )
+
+        if len(added_details["failure"]) == 1:
+            # adding the Model Group back to the source collection in case of an add failure
+            self.add_model_groups(
+                collection_name=source_collection_name, model_groups=[model_group]
+            )
+            raise Exception(added_details["failure"][0]["failure"])
+
+        return {
+            "moved_success": model_group,
+        }
+
+    def _convert_tag_collection_response(self, tag_collections: List[str]):
+        """Converts a Collection response from the tag api to a Collection list response.
+
+        Args:
+            tag_collections List[dict]: The Collection list response from the tag api.
+        """
+        collection_details = []
+        for collection in tag_collections:
+            collection_arn = collection["ResourceARN"]
+            collection_name = collection_arn.split("group/")[1]
+            collection_details.append(
+                {
+                    "Name": collection_name,
+                    "Arn": collection_arn,
+                    "Type": "Collection",
+                }
+            )
+        return collection_details
+
+    def _convert_group_resource_response(
+        self, group_resource_details: List[dict], is_model_group: bool = False
+    ):
+        """Converts a Collection response from the resource group api to a Collection list response.
+
+        Args:
+            group_resource_details (List[dict]): The Collection list response from the
+                resource group api.
+            is_model_group (bool): Indicates if the response is of Collection or Model Group type.
+        """
+        collection_details = []
+        if group_resource_details["Resources"]:
+            for resource_group in group_resource_details["Resources"]:
+                collection_arn = resource_group["Identifier"]["ResourceArn"]
+                collection_name = collection_arn.split("group/")[1]
+                collection_details.append(
+                    {
+                        "Name": collection_name,
+                        "Arn": collection_arn,
+                        "Type": (
+                            resource_group["Identifier"]["ResourceType"]
+                            if is_model_group
+                            else "Collection"
+                        ),
+                    }
+                )
+        return collection_details
+
+    def _get_full_list_resource(self, collection_name, collection_filter):
+        """Iterates the full resource group list and returns the appended paginated response.
+
+        Args:
+            collection_name (str): The name of the Collection from which to get details.
+            collection_filter (dict): Filter details to pass to get the resource list.
+        """
+        list_group_response = self.sagemaker_session.list_group_resources(
+            group=collection_name, filters=collection_filter
+        )
+        next_token = list_group_response.get("NextToken")
+        while next_token is not None:
+
+            paginated_group_response = self.sagemaker_session.list_group_resources(
+                group=collection_name,
+                filters=collection_filter,
+                next_token=next_token,
+            )
+            list_group_response["Resources"] = (
+                list_group_response["Resources"] + paginated_group_response["Resources"]
+            )
+            list_group_response["ResourceIdentifiers"] = (
+                list_group_response["ResourceIdentifiers"]
+                + paginated_group_response["ResourceIdentifiers"]
+            )
+            next_token = paginated_group_response.get("NextToken")
+
+        return list_group_response
+
+    def list_collection(self, collection_name: str = None):
+        """Lists the contents of the specified Collection.
+
+        If there is no Collection with the name ``collection_name``, lists all the
+        Collections at the root level.
+
+        Args:
+            collection_name (str): The name of the Collection whose contents are listed.
+        """
+        collection_content = []
+        if collection_name is None:
+            tag_filters = [
+                {
+                    "Key": "sagemaker:collection-path:root",
+                    "Values": ["true"],
+                },
+            ]
+            resource_type_filters = ["resource-groups:group"]
+            tag_collections = self.sagemaker_session.get_tagging_resources(
+                tag_filters=tag_filters, resource_type_filters=resource_type_filters
+            )
+
+            return self._convert_tag_collection_response(tag_collections)
+
+        collection_filter = [
+            {
+                "Name": "resource-type",
+                "Values": ["AWS::ResourceGroups::Group"],
+            },
+        ]
+        list_group_response = self._get_full_list_resource(
+            collection_name=collection_name, collection_filter=collection_filter
+        )
+        collection_content = self._convert_group_resource_response(list_group_response)
+
+        collection_filter = [
+            {
+                "Name": "resource-type",
+                "Values": ["AWS::SageMaker::ModelPackageGroup"],
+            },
+        ]
+        list_group_response = self._get_full_list_resource(
+            collection_name=collection_name, collection_filter=collection_filter
+        )
+
+        collection_content = collection_content + self._convert_group_resource_response(
+            list_group_response, True
+        )
+
+        return collection_content