Refactor hub related stuff from session utils

Author: jinyoung-lim

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -15,16 +15,14 @@
 
 
 from typing import Any, Dict, Optional
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
-
 from sagemaker.session import Session
-
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 from sagemaker.jumpstart.curated_hub.types import (
     DescribeHubResponse,
     HubContentType,
     DescribeHubContentsResponse,
 )
-import sagemaker.jumpstart.session_utils as session_utils
 
 
 class CuratedHub:
@@ -53,7 +51,7 @@ def create(
     ) -> Dict[str, str]:
         """Creates a hub with the given description"""
 
-        bucket_name = session_utils.create_hub_bucket_if_it_does_not_exist(
+        bucket_name = hub_utils.create_hub_bucket_if_it_does_not_exist(
             bucket_name, self._sagemaker_session
         )
 

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -110,3 +110,45 @@ def generate_hub_arn_for_estimator_init_kwargs(
         else:
             hub_arn = construct_hub_arn_from_name(hub_name=hub_name, region=region, session=session)
     return hub_arn
+
+
+def generate_default_hub_bucket_name(
+    sagemaker_session: Session = constants.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: Optional[str] = None,
+    sagemaker_session: Session = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+) -> str:
+    """Creates the default SageMaker Hub 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

src/sagemaker/jumpstart/session_utils.py

@@ -208,45 +208,3 @@ 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: Optional[str] = None,
-    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
-) -> str:
-    """Creates the default SageMaker Hub 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