curated_hub/types.py to types.py

Author: jinyoung-lim

src/sagemaker/jumpstart/cache.py

@@ -45,10 +45,9 @@
     JumpStartModelSpecs,
     JumpStartS3FileType,
     JumpStartVersionedModelId,
-)
-from sagemaker.jumpstart.curated_hub.types import (
     DescribeHubResponse,
     DescribeHubContentsResponse,
+    HubType,
     HubContentType,
 )
 from sagemaker.jumpstart.curated_hub import utils as hub_utils
@@ -362,7 +361,7 @@ def _retrieval_function(
             return JumpStartCachedContentValue(
                 formatted_content=model_specs
             )
-        if data_type == HubContentType.HUB:
+        if data_type == HubType.HUB:
             hub_name, region, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
             hub: CuratedHub = CuratedHub(hub_name=hub_name, region=region)
             hub_description: DescribeHubResponse = hub.describe()

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -17,12 +17,12 @@
 from typing import Any, Dict, Optional
 from sagemaker.session import Session
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
-from sagemaker.jumpstart.curated_hub.utils import create_hub_bucket_if_it_does_not_exist
-from sagemaker.jumpstart.curated_hub.types import (
+from sagemaker.jumpstart.types import (
     DescribeHubResponse,
     DescribeHubContentsResponse,
     HubContentType,
 )
+from sagemaker.jumpstart.curated_hub.utils import create_hub_bucket_if_it_does_not_exist
 
 
 class CuratedHub:

src/sagemaker/jumpstart/curated_hub/types.py

@@ -16,11 +16,11 @@
 from typing import Optional
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
-from sagemaker.jumpstart import constants
-from sagemaker.jumpstart.curated_hub.types import (
+from sagemaker.jumpstart.types import (
     HubContentType,
     HubArnExtractedInfo,
 )
+from sagemaker.jumpstart import constants
 
 
 def get_info_from_hub_resource_arn(