Feat: list model collection

Author: Keshav Chandak

src/sagemaker/collection.py

@@ -147,7 +147,6 @@ def create(self, collection_name: str, parent_collection_name: str = None):
 
             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
 
@@ -337,3 +336,127 @@ def move_model_group(
         return {
             "moved_success": model_group,
         }
+
+    def _convert_tag_collection_response(self, tag_collections: List[str]):
+        """Converts collection response from tag api to collection list response
+
+        Args:
+            tag_collections List[dict]: Collections list response from 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 collection response from resource group api to collection list response
+
+        Args:
+            group_resource_details (List[dict]): Collections list response from resource group api
+            is_model_group (bool): If the reponse 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):
+        """Iterating to the full resource group list and returns appended paginated response
+
+        Args:
+            collection_name (str): Name of the collection to get the details
+            collection_filter (dict): Filter details to be passed 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):
+        """To all list the collections and content of the collections
+
+        In case there is no collection_name, it lists all the collections on the root level
+
+        Args:
+            collection_name (str): The name of the collection to list the contents of
+        """
+        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

src/sagemaker/session.py

@@ -205,6 +205,7 @@ def __init__(
         self.s3_resource = None
         self.s3_client = None
         self.resource_groups_client = None
+        self.resource_group_tagging_client = None
         self.config = None
         self.lambda_client = None
         self.settings = settings
@@ -3962,7 +3963,7 @@ def delete_model(self, model_name):
         LOGGER.info("Deleting model with name: %s", model_name)
         self.sagemaker_client.delete_model(ModelName=model_name)
 
-    def list_group_resources(self, group, filters):
+    def list_group_resources(self, group, filters, next_token: str = ""):
         """To list group resources with given filters
 
         Args:
@@ -3972,7 +3973,9 @@ def list_group_resources(self, group, filters):
         self.resource_groups_client = self.resource_groups_client or self.boto_session.client(
             "resource-groups"
         )
-        return self.resource_groups_client.list_group_resources(Group=group, Filters=filters)
+        return self.resource_groups_client.list_group_resources(
+            Group=group, Filters=filters, NextToken=next_token
+        )
 
     def delete_resource_group(self, group):
         """To delete a resource group
@@ -3996,6 +3999,39 @@ def get_resource_group_query(self, group):
         )
         return self.resource_groups_client.get_group_query(Group=group)
 
+    def get_tagging_resources(self, tag_filters, resource_type_filters):
+        """To list the complete resources for a particular resource group tag
+
+        tag_filters: filters for the tag
+        resource_type_filters: resource filter for the tag
+        """
+        self.resource_group_tagging_client = (
+            self.resource_group_tagging_client
+            or self.boto_session.client("resourcegroupstaggingapi")
+        )
+        resource_list = []
+
+        try:
+            resource_tag_response = self.resource_group_tagging_client.get_resources(
+                TagFilters=tag_filters, ResourceTypeFilters=resource_type_filters
+            )
+
+            resource_list = resource_list + resource_tag_response["ResourceTagMappingList"]
+
+            next_token = resource_tag_response.get("PaginationToken")
+            while next_token is not None and next_token != "":
+                resource_tag_response = self.resource_group_tagging_client.get_resources(
+                    TagFilters=tag_filters,
+                    ResourceTypeFilters=resource_type_filters,
+                    NextToken=next_token,
+                )
+                resource_list = resource_list + resource_tag_response["ResourceTagMappingList"]
+                next_token = resource_tag_response.get("PaginationToken")
+
+            return resource_list
+        except ClientError as error:
+            raise error
+
     def create_group(self, name, resource_query, tags):
         """To create a AWS Resource Group
 

tests/integ/test_collection.py

@@ -166,3 +166,33 @@ def test_move_model_groups_in_collection_success(sagemaker_session):
     sagemaker_session.sagemaker_client.delete_model_package_group(
         ModelPackageGroupName=model_group_name
     )
+
+
+def test_list_collection_success(sagemaker_session):
+    model_group_name = unique_name_from_base("test-model-group")
+    sagemaker_session.sagemaker_client.create_model_package_group(
+        ModelPackageGroupName=model_group_name
+    )
+    collection = Collection(sagemaker_session)
+    collection_name = unique_name_from_base("test-collection")
+    collection.create(collection_name)
+    model_groups = []
+    model_groups.append(model_group_name)
+    collection.add_model_groups(collection_name=collection_name, model_groups=model_groups)
+    child_collection_name = unique_name_from_base("test-collection")
+    collection.create(parent_collection_name=collection_name, collection_name=child_collection_name)
+    root_collections = collection.list_collection()
+    is_collection_found = False
+    for root_collection in root_collections:
+        if root_collection["Name"] == collection_name:
+            is_collection_found = True
+    assert is_collection_found
+
+    collection_content = collection.list_collection(collection_name)
+    assert len(collection_content) == 2
+
+    collection.remove_model_groups(collection_name=collection_name, model_groups=model_groups)
+    collection.delete([child_collection_name, collection_name])
+    sagemaker_session.sagemaker_client.delete_model_package_group(
+        ModelPackageGroupName=model_group_name
+    )