refactor, add types for hub/hubcontent descriptions, add helpers.

Author: jinyoung-lim

src/sagemaker/jumpstart/cache.py

@@ -29,7 +29,6 @@
     JUMPSTART_LOGGER,
     MODEL_ID_LIST_WEB_URL,
 )
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
     JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
@@ -44,9 +43,12 @@
     JumpStartModelSpecs,
     JumpStartS3FileType,
     JumpStartVersionedModelId,
-    HubDataType,
+    HubContentType,
+    HubDescription,
+    HubContentDescription,
 )
 from sagemaker.jumpstart import utils
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 from sagemaker.utilities.cache import LRUCache
 
 
@@ -338,29 +340,33 @@ def _retrieval_function(
             return JumpStartCachedContentValue(
                 formatted_content=model_specs
             )
-        if data_type == HubDataType.MODEL:
+        if data_type == HubContentType.MODEL:
             hub_name, region, model_name, model_version = utils.extract_info_from_hub_content_arn(
                 id_info
             )
-            hub = CuratedHub(hub_name=hub_name, region=region)
-            hub_content = hub.describe_model(model_name=model_name, model_version=model_version)
+            hub_model_description: HubContentDescription = hub_utils.describe_model(
+                hub_name=hub_name,
+                region=region,
+                model_name=model_name,
+                model_version=model_version
+            )
+            model_specs = JumpStartModelSpecs(hub_model_description, is_hub_content=True)
             utils.emit_logs_based_on_model_specs(
-                hub_content.content_document,
+                model_specs,
                 self.get_region(),
                 self._s3_client
             )
-            model_specs = JumpStartModelSpecs(hub_content.content_document, is_hub_content=True)
+            # TODO: Parse HubContentDescription
             return JumpStartCachedContentValue(
                 formatted_content=model_specs
             )
-        if data_type == HubDataType.HUB:
+        if data_type == HubContentType.HUB:
             hub_name, region, _, _ = utils.extract_info_from_hub_content_arn(id_info)
-            hub = CuratedHub(hub_name=hub_name, region=region)
-            hub_info = hub.describe()
-            return JumpStartCachedContentValue(formatted_content=hub_info)
+            hub_description: HubDescription = hub_utils.describe(hub_name=hub_name, region=region)
+            return JumpStartCachedContentValue(formatted_content=hub_description)
         raise ValueError(
             f"Bad value for key '{key}': must be in",
-            f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubDataType.HUB, HubDataType.MODEL]}"
+            f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubContentType.HUB, HubContentType.MODEL]}"
         )
 
     def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -474,7 +480,7 @@ def get_hub_model(self, hub_model_arn: str) -> JumpStartModelSpecs:
         """
 
         details, _ = self._content_cache.get(
-            JumpStartCachedContentKey(HubDataType.MODEL, hub_model_arn)
+            JumpStartCachedContentKey(HubContentType.MODEL, hub_model_arn)
         )
         return details.formatted_content
 
@@ -485,7 +491,7 @@ def get_hub(self, hub_arn: str) -> Dict[str, Any]:
             hub_arn (str): Arn for the Hub to get info for
         """
 
-        details, _ = self._content_cache.get(JumpStartCachedContentKey(HubDataType.HUB, hub_arn))
+        details, _ = self._content_cache.get(JumpStartCachedContentKey(HubContentType.HUB, hub_arn))
         return details.formatted_content
 
     def clear(self) -> None:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,32 +13,32 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
+from typing import Any, Dict, Optional
 import boto3
 from sagemaker.session import Session
 from sagemaker.jumpstart.constants import (
     JUMPSTART_DEFAULT_REGION_NAME,
 )
 
-from sagemaker.jumpstart.types import HubDataType
-import sagemaker.jumpstart.curated_hub.utils as hubutils
+from sagemaker.jumpstart.types import HubDescription, HubContentType, HubContentDescription
+import sagemaker.jumpstart.session_utils as session_utils
 
 
 class CuratedHub:
     """Class for creating and managing a curated JumpStart hub"""
 
     def __init__(
         self,
-        name: str,
+        hub_name: str,
         region: str = JUMPSTART_DEFAULT_REGION_NAME,
-        session: Optional[Session] = None,
+        sagemaker_session: Optional[Session] = None,
     ):
-        self.name = name
-        if session.boto_region_name != region:
+        self.hub_name = hub_name
+        if sagemaker_session.boto_region_name != region:
             # TODO: Handle error
             pass
         self.region = region
-        self._session = session or Session(boto3.Session(region_name=region))
+        self._sagemaker_session = sagemaker_session or Session(boto3.Session(region_name=region))
 
     def create(
         self,
@@ -50,32 +50,60 @@ def create(
     ) -> Dict[str, str]:
         """Creates a hub with the given description"""
 
-        return hubutils.create_hub(
-            hub_name=self.name,
+        bucket_name = session_utils.create_hub_bucket_if_it_does_not_exist(
+            bucket_name, self._sagemaker_session
+        )
+
+        return self._sagemaker_session.create_hub(
+            hub_name=self.hub_name,
             hub_description=description,
             hub_display_name=display_name,
             hub_search_keywords=search_keywords,
             hub_bucket_name=bucket_name,
             tags=tags,
-            sagemaker_session=self._session,
         )
 
-    def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
-        """Returns descriptive information about the Hub Model"""
+    def describe(self) -> HubDescription:
+        """Returns descriptive information about the Hub"""
+
+        hub_description = self._sagemaker_session.describe_hub(hub_name=self.hub_name)
+
+        return HubDescription(hub_description)
+
+    def list_models(self, **kwargs) -> Dict[str, Any]:
+        """Lists the models in this Curated Hub
 
-        hub_content = hubutils.describe_hub_content(
-            hub_name=self.name,
-            content_name=model_name,
-            content_type=HubDataType.MODEL,
-            content_version=model_version,
-            sagemaker_session=self._session,
+        **kwargs: Passed to invocation of ``Session:list_hub_contents``.
+        """
+        # TODO: Validate kwargs and fast-fail?
+
+        hub_content_summaries = self._sagemaker_session.list_hub_contents(
+            hub_name=self.hub_name, hub_content_type=HubContentType.MODEL, **kwargs
         )
+        # TODO: Handle pagination
+        return hub_content_summaries
 
-        return hub_content
+    def describe_model(self, model_name: str, model_version: str = "*") -> HubContentDescription:
+        """Returns descriptive information about the Hub Model"""
 
-    def describe(self) -> Dict[str, Any]:
-        """Returns descriptive information about the Hub"""
+        hub_content_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+            hub_name=self.hub_name,
+            hub_content_name=model_name,
+            hub_content_version=model_version,
+            hub_content_type=HubContentType.MODEL,
+        )
+
+        return HubContentDescription(hub_content_description)
 
-        hub_info = hubutils.describe_hub(hub_name=self.name, sagemaker_session=self._session)
+    def delete_model(self, model_name: str, model_version: str = "*") -> None:
+        """Deletes a model from this CuratedHub."""
+        return self._sagemaker_session.delete_hub_content(
+            hub_content_name=model_name,
+            hub_content_version=model_version,
+            hub_content_type=HubContentType.MODEL,
+            hub_name=self.hub_name,
+        )
 
-        return hub_info
+    def delete(self) -> None:
+        """Deletes this Curated Hub"""
+        return self._sagemaker_session.delete_hub(self.hub_name)