delete hubutils, refactor

Author: jinyoung-lim

src/sagemaker/jumpstart/cache.py

@@ -44,7 +44,7 @@
     JumpStartModelSpecs,
     JumpStartS3FileType,
     JumpStartVersionedModelId,
-    HubDataType,
+    HubContentType,
 )
 from sagemaker.jumpstart import utils
 from sagemaker.utilities.cache import LRUCache
@@ -338,7 +338,7 @@ 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
             )
@@ -353,14 +353,14 @@ def _retrieval_function(
             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)
         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 +474,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 +485,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

@@ -20,8 +20,8 @@
     JUMPSTART_DEFAULT_REGION_NAME,
 )
 
-from sagemaker.jumpstart.types import HubDataType
-import sagemaker.jumpstart.curated_hub.utils as hubutils
+from sagemaker.jumpstart.types import HubContentType
+import sagemaker.jumpstart.session_utils as session_utils
 
 
 class CuratedHub:
@@ -50,32 +50,60 @@ def create(
     ) -> Dict[str, str]:
         """Creates a hub with the given description"""
 
-        return hubutils.create_hub(
+        bucket_name = session_utils.create_hub_bucket_if_it_does_not_exist(
+            bucket_name, self._session
+        )
+
+        return self._session.create_hub(
             hub_name=self.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(self) -> Dict[str, Any]:
+        """Returns descriptive information about the Hub"""
+
+        hub_info = self._session.describe_hub(hub_name=self.name)
+
+        return hub_info
+
+    def list_models(self, **kwargs) -> Dict[str, Any]:
+        """Lists the models in this Curated Hub
+
+        **kwargs: Passed to invocation of ``Session:list_hub_contents``.
+        """
+        # TODO: Validate kwargs and fast-fail?
+
+        hub_content_summaries = self._session.list_hub_contents(
+            hub_name=self.hub_name, hub_content_type=HubContentType.MODEL, **kwargs
+        )
+        # TODO: Handle pagination
+        return hub_content_summaries
+
     def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
         """Returns descriptive information about the Hub Model"""
 
-        hub_content = hubutils.describe_hub_content(
+        hub_content = self._session.describe_hub_content(
             hub_name=self.name,
-            content_name=model_name,
-            content_type=HubDataType.MODEL,
-            content_version=model_version,
-            sagemaker_session=self._session,
+            hub_content_name=model_name,
+            hub_content_version=model_version,
+            hub_content_type=HubContentType.MODEL,
         )
 
         return hub_content
 
-    def describe(self) -> Dict[str, Any]:
-        """Returns descriptive information about the Hub"""
-
-        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._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._session.delete_hub(self.hub_name)

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -208,3 +208,45 @@ def get_model_id_version_from_training_job(
         )
 
     return model_id, model_version
+
+
+def generate_default_hub_bucket_name(
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+) -> str:
+    """Return the name of the default bucket to use in relevant Amazon SageMaker Hub interactions.
+
+    Returns:
+        str: The name of the default bucket. If the name was not explicitly specified through
+            the Session or sagemaker_config, the bucket will take the form:
+            ``sagemaker-hubs-{region}-{AWS account ID}``.
+    """
+
+    region: str = sagemaker_session.boto_region_name
+    account_id: str = sagemaker_session.account_id()
+
+    # TODO: Validate and fast fail
+
+    return f"sagemaker-hubs-{region}-{account_id}"
+
+
+def create_hub_bucket_if_it_does_not_exist(
+    bucket_name: str = None, sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+):
+    """Return the name of the default bucket to use in relevant Amazon SageMaker Hub interactions.
+    This function will create the s3 bucket if it does not exist.
+
+    Returns:
+        str: The name of the default bucket. Takes the form:
+            ``sagemaker-hubs-{region}-{AWS account ID}``.
+    """
+
+    region: str = sagemaker_session.boto_region_name
+    if bucket_name is None:
+        bucket_name: str = generate_default_hub_bucket_name(sagemaker_session)
+
+    sagemaker_session._create_s3_bucket_if_it_does_not_exist(
+        bucket_name=bucket_name,
+        region=region,
+    )
+
+    return bucket_name