Use CuratedHub class instead of using a middleware utilities

Author: jinyoung-lim

src/sagemaker/jumpstart/cache.py

@@ -51,6 +51,7 @@
 )
 from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.curated_hub import utils as hub_utils
+from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
 from sagemaker.utilities.cache import LRUCache
 
 
@@ -346,9 +347,8 @@ def _retrieval_function(
             hub_name, region, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
                 id_info
             )
-            hub_model_description: DescribeHubContentsResponse = hub_utils.describe_model(
-                hub_name=hub_name,
-                region=region,
+            hub: CuratedHub = CuratedHub(hub_name=hub_name, region=region)
+            hub_model_description: DescribeHubContentsResponse = hub.describe_model(
                 model_name=model_name,
                 model_version=model_version
             )
@@ -364,10 +364,8 @@ def _retrieval_function(
             )
         if data_type == HubContentType.HUB:
             hub_name, region, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
-            hub_description: DescribeHubResponse = hub_utils.describe(
-                hub_name=hub_name,
-                region=region
-            )
+            hub: CuratedHub = CuratedHub(hub_name=hub_name, region=region)
+            hub_description: DescribeHubResponse = hub.describe()
             return JumpStartCachedContentValue(formatted_content=hub_description)
         raise ValueError(
             f"Bad value for key '{key}': must be in",

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -13,43 +13,16 @@
 """This module contains utilities related to SageMaker JumpStart CuratedHub."""
 from __future__ import absolute_import
 import re
-from typing import Any, Dict, Optional
-import boto3
+from typing import Optional
 from sagemaker.session import Session
 from sagemaker.jumpstart import constants
 from sagemaker.utils import aws_partition
 from sagemaker.jumpstart.curated_hub.types import (
-    DescribeHubResponse,
     HubContentType,
-    DescribeHubContentsResponse,
     HubArnExtractedInfo,
 )
 
 
-def describe(hub_name: str, region: str) -> DescribeHubResponse:
-    """Returns descriptive information about the Hub."""
-
-    sagemaker_session = Session(boto3.Session(region_name=region))
-    hub_description = sagemaker_session.describe_hub(hub_name=hub_name)
-    return DescribeHubResponse(hub_description)
-
-
-def describe_model(
-    hub_name: str, region: str, model_name: str, model_version: str = "*"
-) -> DescribeHubContentsResponse:
-    """Returns descriptive information about the Hub model."""
-
-    sagemaker_session = Session(boto3.Session(region_name=region))
-    hub_content_description: Dict[str, Any] = sagemaker_session.describe_hub_content(
-        hub_name=hub_name,
-        hub_content_name=model_name,
-        hub_content_version=model_version,
-        hub_content_type=HubContentType.MODEL,
-    )
-
-    return DescribeHubContentsResponse(hub_content_description)
-
-
 def get_info_from_hub_resource_arn(
     arn: str,
 ) -> HubArnExtractedInfo: