diff --git a/sagemaker-core/src/sagemaker/core/resources.py b/sagemaker-core/src/sagemaker/core/resources.py index c0780dc553..9935657215 100644 --- a/sagemaker-core/src/sagemaker/core/resources.py +++ b/sagemaker-core/src/sagemaker/core/resources.py @@ -13830,6 +13830,7 @@ class HubContent(Base): hub_content_dependencies: The location of any dependencies that the hub content has, such as scripts, model artifacts, datasets, or notebooks. failure_reason: The failure reason if importing hub content failed. last_modified_time: The last modified time of the hub content. + original_creation_time: The date and time when the hub content was originally created, before any updates or revisions. """ @@ -13854,6 +13855,7 @@ class HubContent(Base): creation_time: Optional[datetime.datetime] = Unassigned() last_modified_time: Optional[datetime.datetime] = Unassigned() hub_name: Optional[str] = Unassigned() + original_creation_time: Optional[datetime.datetime] = Unassigned() def get_name(self) -> str: attributes = vars(self) @@ -14210,6 +14212,85 @@ def load( region=region, ) + @classmethod + @Base.add_validate_call + def get_all( + cls, + hub_name: StrPipeVar, + hub_content_type: StrPipeVar, + name_contains: Optional[StrPipeVar] = Unassigned(), + max_schema_version: Optional[StrPipeVar] = Unassigned(), + creation_time_before: Optional[datetime.datetime] = Unassigned(), + creation_time_after: Optional[datetime.datetime] = Unassigned(), + sort_by: Optional[StrPipeVar] = Unassigned(), + sort_order: Optional[StrPipeVar] = Unassigned(), + session: Optional[Session] = None, + region: Optional[StrPipeVar] = None, + ) -> ResourceIterator["HubContent"]: + """ + Get all HubContent resources + + Parameters: + hub_name: The name of the hub to list the contents of. + hub_content_type: The type of hub content to list. + name_contains: Only list hub content if the name contains the specified string. + max_schema_version: The upper bound of the hub content schema verion. + creation_time_before: Only list hub content that was created before the time specified. + creation_time_after: Only list hub content that was created after the time specified. + sort_by: Sort hub content versions by either name or creation time. + sort_order: Sort hubs by ascending or descending order. + max_results: The maximum amount of hub content to list. + next_token: If the response to a previous ListHubContents request was truncated, the response includes a NextToken. To retrieve the next set of hub content, use the token in the next request. + session: Boto3 session. + region: Region name. + + Returns: + Iterator for listed HubContent resources. + + Raises: + botocore.exceptions.ClientError: This exception is raised for AWS service related errors. + The error message and error code can be parsed from the exception as follows: + ``` + try: + # AWS service call here + except botocore.exceptions.ClientError as e: + error_message = e.response['Error']['Message'] + error_code = e.response['Error']['Code'] + ``` + ResourceNotFound: Resource being access is not found. + """ + + client = Base.get_sagemaker_client( + session=session, region_name=region, service_name="sagemaker" + ) + + operation_input_args = { + "HubName": hub_name, + "HubContentType": hub_content_type, + "NameContains": name_contains, + "MaxSchemaVersion": max_schema_version, + "CreationTimeBefore": creation_time_before, + "CreationTimeAfter": creation_time_after, + "SortBy": sort_by, + "SortOrder": sort_order, + } + + # serialize the input request + operation_input_args = serialize(operation_input_args) + logger.debug(f"Serialized input request: {operation_input_args}") + + return ResourceIterator( + client=client, + session=session, + region=region, + list_method="list_hub_contents", + summaries_key="HubContentSummaries", + summary_name="HubContentInfo", + resource_cls=HubContent, + extra_init_data={"hub_name": hub_name}, + list_method_kwargs=operation_input_args, + ) + @Base.add_validate_call def get_all_versions( self, @@ -14279,6 +14360,7 @@ def get_all_versions( summary_name="HubContentInfo", resource_cls=HubContent, list_method_kwargs=operation_input_args, + extra_init_data={"hub_name": self.hub_name}, ) diff --git a/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py b/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py index 22428f3006..166f915ee0 100644 --- a/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py +++ b/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py @@ -106,6 +106,25 @@ "ModelPackage": {"ModelPackageName": "model_package_arn"}, } +# Some List APIs return a summary member that the Describe API does not. A resource class +# is generated from the Describe output, so ResourceIterator drops such members when it +# builds the resource from the summary (`extra="forbid"` leaves no other option). This +# maps a resource to the summary members it keeps as optional attributes. ListHubContents +# returns OriginalCreationTime, the model-level catalog date, while DescribeHubContent +# returns only the version-level CreationTime. +LIST_SUMMARY_EXTRA_ATTRIBUTES = { + "HubContent": {"OriginalCreationTime": "Optional[datetime.datetime] = Unassigned()"}, +} + +# Some List APIs take an input member that the resource's Describe input requires but the +# list summary does not return (ListHubContents takes HubName, HubContentInfo omits it). +# get_all() cannot build a resource from such a summary unless the list argument is +# copied onto every resource. This maps a resource to the list-input members to copy, as +# {summary-less member: resource attribute}. +GET_ALL_INPUT_ATTRIBUTES = { + "HubContent": {"HubName": "hub_name"}, +} + class ResourcesCodeGen: """A class for generating resources based on a service JSON file. @@ -535,8 +554,15 @@ def _get_class_attributes(self, resource_name: str, class_methods: list) -> tupl class_attributes["hub_name"] = "Optional[str] = Unassigned()" class_attributes_string = class_attributes_string.replace("hub_name: str", "") class_attributes_string = ( - class_attributes_string + "hub_name: Optional[str] = Unassigned()" + class_attributes_string + "hub_name: Optional[str] = Unassigned()\n" ) + for member, annotation in LIST_SUMMARY_EXTRA_ATTRIBUTES.get(resource_name, {}).items(): + attribute = convert_to_snake_case(member) + class_attributes[attribute] = annotation + class_attributes_string += f"{attribute}: {annotation}\n" + attributes_and_documentation[member] = self.shapes_extractor.combined_shapes[ + self._list_summary_shape(resource_name) + ]["members"][member].get("documentation") if resource_name == "Endpoint": class_attributes["serializer"] = "Optional[BaseSerializer] = None" class_attributes_string = class_attributes_string.replace( @@ -1714,6 +1740,21 @@ def generate_additional_get_all_method(self, method: Method, resource_attributes "list_method_kwargs=operation_input_args", ] + # An object method lists resources of its own type under its own identifiers + # (ListHubContentVersions lists the versions of one HubContent). The summaries + # omit the identifiers the list call was made with, so copy them onto every + # resource or the first refresh() fails with a missing required parameter. + if ( + method.method_type == MethodType.OBJECT.value + and method.return_type == method.resource_name + ): + input_attributes = GET_ALL_INPUT_ATTRIBUTES.get(method.resource_name, {}) + if input_attributes: + extra_init = ", ".join( + f"'{attribute}': self.{attribute}" for attribute in input_attributes.values() + ) + resource_iterator_args_list.append(f"extra_init_data={{{extra_init}}}") + resource_iterator_args = ",\n".join(resource_iterator_args_list) resource_iterator_args = add_indent(resource_iterator_args, 8) serialize_operation_input = SERIALIZE_INPUT_TEMPLATE.format( @@ -1934,6 +1975,14 @@ def generate_wait_for_delete_method(self, resource_name: str) -> str: ) return formatted_method + def _list_summary_shape(self, resource_name: str) -> str: + """Return the summary shape name of the resource's List API output.""" + operation_name = "List" + resource_name + "s" + output_shape = self.operations[operation_name]["output"]["shape"] + output_members = self.shapes[output_shape]["members"] + summaries_key = next(key for key in output_members if key != "NextToken") + return self.shapes[output_members[summaries_key]["shape"]]["member"]["shape"] + def generate_get_all_method(self, resource_name: str) -> str: """Auto-Generate 'get_all' class Method [list API] for a resource. @@ -1974,7 +2023,13 @@ def generate_get_all_method(self, resource_name: str) -> str: get_operation_required_input = [] custom_key_mapping_str = "" - if any(member not in summary_members for member in get_operation_required_input): + input_attributes = GET_ALL_INPUT_ATTRIBUTES.get(resource_name, {}) + missing_members = [ + member + for member in get_operation_required_input + if member not in summary_members and member not in input_attributes + ] + if missing_members: if "MonitoringJobDefinitionSummary" == summary_name: custom_key_mapping = { "monitoring_job_definition_name": "job_definition_name", @@ -2002,6 +2057,13 @@ def generate_get_all_method(self, resource_name: str) -> str: if custom_key_mapping_str: resource_iterator_args_list.append("custom_key_mapping=custom_key_mapping") + if input_attributes: + extra_init = ", ".join( + f"'{attribute}': {convert_to_snake_case(member)}" + for member, attribute in input_attributes.items() + ) + resource_iterator_args_list.append(f"extra_init_data={{{extra_init}}}") + exclude_list = ["next_token", "max_results"] get_all_args = self._generate_method_args(operation_input_shape_name, exclude_list) diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index dc6b96378d..7a686c86cc 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -546,6 +546,7 @@ def __init__( custom_key_mapping: dict = None, session=None, region: Optional[str] = None, + extra_init_data: dict = None, ): """Initialize a ResourceIterator object @@ -563,6 +564,9 @@ def __init__( Each returned resource is bound to it, so its refresh and other object methods use the same account. Defaults to None (the process default client). region (str, optional): The region the list call was made in. Defaults to None. + extra_init_data (dict, optional): Attributes set on every resource in addition to + the summary, for a list argument the summary does not echo back (for example + the ``hub_name`` a ``ListHubContents`` call was made with). Defaults to None. """ self.summaries_key = summaries_key self.summary_name = summary_name @@ -572,6 +576,7 @@ def __init__( self.custom_key_mapping = custom_key_mapping self.session = session self.region = region + self.extra_init_data = extra_init_data or {} self.resource_cls = resource_cls self.index = 0 @@ -605,6 +610,8 @@ def __next__(self) -> T: # Filter out the fields that are not in the resource class fields = self.resource_cls.__annotations__ init_data = {k: v for k, v in init_data.items() if k in fields} + for key, value in self.extra_init_data.items(): + init_data.setdefault(key, value) resource_object = self.resource_cls(**init_data) diff --git a/sagemaker-core/tests/unit/test_hub_content_get_all.py b/sagemaker-core/tests/unit/test_hub_content_get_all.py new file mode 100644 index 0000000000..19a11a4508 --- /dev/null +++ b/sagemaker-core/tests/unit/test_hub_content_get_all.py @@ -0,0 +1,185 @@ +# 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. +"""``HubContent.get_all`` lists a hub's contents as bound resources. + +Before this method the SDK had no resource-level ``ListHubContents``: callers paged the +boto3 client by hand (six places in the SDK did the same), ``ResourceIterator`` dropped +``OriginalCreationTime`` because the ``HubContent`` class had no such field, and +``get_all_versions()`` built resources without the ``hub_name`` the list call used, so +the first ``refresh()`` failed with ``Missing required parameter in input: "HubName"``. +""" + +from __future__ import absolute_import + +import datetime +from unittest.mock import MagicMock, patch + +import boto3 +import pytest + +from sagemaker.core.resources import Base, HubContent +from sagemaker.core.utils.utils import ResourceIterator, SageMakerClient + +HUB = "SageMakerPublicHub" +CREATED = datetime.datetime(2024, 3, 1, tzinfo=datetime.timezone.utc) +ORIGINAL = datetime.datetime(2023, 6, 1, tzinfo=datetime.timezone.utc) + + +def _summary(name, version="1.0.0"): + return { + "HubContentName": name, + "HubContentArn": f"arn:aws:sagemaker:us-west-2:aws:hub-content/{HUB}/Model/{name}/{version}", + "HubContentVersion": version, + "HubContentType": "Model", + "DocumentSchemaVersion": "2.0.0", + "HubContentStatus": "Available", + "HubContentSearchKeywords": ["@task:text-generation", "@license:apache-2.0"], + "CreationTime": CREATED, + "OriginalCreationTime": ORIGINAL, + } + + +def _describe(**kwargs): + if "HubName" not in kwargs: + raise ValueError('Missing required parameter in input: "HubName"') + summary = _summary(kwargs["HubContentName"], kwargs.get("HubContentVersion", "1.0.0")) + summary.pop("OriginalCreationTime") + return { + **summary, + "HubName": kwargs["HubName"], + "HubArn": f"arn:aws:sagemaker:us-west-2:aws:hub/{HUB}", + "HubContentDocument": "{}", + } + + +@pytest.fixture(autouse=True) +def _reset_client_cache(): + SageMakerClient.reset() + yield + SageMakerClient.reset() + + +@pytest.fixture +def client(): + fake = MagicMock() + fake.list_hub_contents.side_effect = [ + {"HubContentSummaries": [_summary("model-a")], "NextToken": "page-2"}, + {"HubContentSummaries": [_summary("model-b")]}, + ] + fake.list_hub_content_versions.return_value = { + "HubContentSummaries": [_summary("model-a", "1.0.0"), _summary("model-a", "2.0.0")] + } + fake.describe_hub_content.side_effect = _describe + with patch.object(Base, "get_sagemaker_client", return_value=fake): + yield fake + + +def test_get_all_pages_through_the_hub_and_binds_hub_name(client): + contents = list(HubContent.get_all(hub_name=HUB, hub_content_type="Model")) + + assert [content.hub_content_name for content in contents] == ["model-a", "model-b"] + assert all(content.hub_name == HUB for content in contents) + assert client.list_hub_contents.call_args_list[0].kwargs == { + "HubName": HUB, + "HubContentType": "Model", + } + assert client.list_hub_contents.call_args_list[1].kwargs["NextToken"] == "page-2" + + +def test_get_all_forwards_the_list_filters(client): + client.list_hub_contents.side_effect = [{"HubContentSummaries": []}] + + list( + HubContent.get_all( + hub_name=HUB, + hub_content_type="Model", + name_contains="llama", + sort_by="HubContentName", + sort_order="Ascending", + ) + ) + + assert client.list_hub_contents.call_args.kwargs == { + "HubName": HUB, + "HubContentType": "Model", + "NameContains": "llama", + "SortBy": "HubContentName", + "SortOrder": "Ascending", + } + + +def test_get_all_keeps_original_creation_time_from_the_summary(client): + contents = list(HubContent.get_all(hub_name=HUB, hub_content_type="Model")) + + assert contents[0].original_creation_time == ORIGINAL + assert contents[0].creation_time == CREATED + assert contents[0].hub_content_search_keywords == [ + "@task:text-generation", + "@license:apache-2.0", + ] + + +def test_get_all_refreshes_each_resource_with_the_hub_name(client): + list(HubContent.get_all(hub_name=HUB, hub_content_type="Model")) + + for call in client.describe_hub_content.call_args_list: + assert call.kwargs["HubName"] == HUB + + +def test_get_all_binds_the_session_and_region_to_each_resource(client): + session = boto3.Session( + aws_access_key_id="AKID", aws_secret_access_key="secret", region_name="eu-west-1" + ) + + content = next( + iter( + HubContent.get_all( + hub_name=HUB, hub_content_type="Model", session=session, region="eu-west-1" + ) + ) + ) + + assert content._session is session + assert content._region == "eu-west-1" + + +def test_get_all_versions_binds_the_hub_name_of_the_parent(client): + parent = HubContent(hub_name=HUB, hub_content_type="Model", hub_content_name="model-a") + + versions = list(parent.get_all_versions()) + + assert [version.hub_content_version for version in versions] == ["1.0.0", "2.0.0"] + assert all(version.hub_name == HUB for version in versions) + for call in client.describe_hub_content.call_args_list: + assert call.kwargs["HubName"] == HUB + + +def test_resource_iterator_extra_init_data_does_not_override_the_summary(): + fake = MagicMock() + fake.list_hub_contents.return_value = {"HubContentSummaries": [_summary("model-a")]} + fake.describe_hub_content.side_effect = _describe + + iterator = ResourceIterator( + client=fake, + summaries_key="HubContentSummaries", + summary_name="HubContentInfo", + resource_cls=HubContent, + list_method="list_hub_contents", + list_method_kwargs={"HubName": HUB, "HubContentType": "Model"}, + extra_init_data={"hub_name": HUB, "hub_content_name": "ignored-when-present"}, + ) + with patch.object(Base, "get_sagemaker_client", return_value=fake): + content = next(iterator) + + assert content.hub_content_name == "model-a" + assert content.hub_name == HUB