feature: JumpStart CuratedHub class creation and function definitions

src/sagemaker/jumpstart/cache.py

@@ -21,34 +21,38 @@
 import botocore
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from sagemaker.session import Session
+from sagemaker.utilities.cache import LRUCache
 from sagemaker.jumpstart.constants import (
  ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
  ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
  JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
  JUMPSTART_DEFAULT_REGION_NAME,
  JUMPSTART_LOGGER,
  MODEL_ID_LIST_WEB_URL,
+ DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.utils import get_info_from_hub_resource_arn
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
  JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
  JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
  JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
  JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
 )
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.types import (
  JumpStartCachedContentKey,
  JumpStartCachedContentValue,
  JumpStartModelHeader,
  JumpStartModelSpecs,
  JumpStartS3FileType,
  JumpStartVersionedModelId,
+ DescribeHubResponse,
+ DescribeHubContentsResponse,
+ HubType,
  HubContentType,
 )
-from sagemaker.jumpstart import utils
-from sagemaker.utilities.cache import LRUCache
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 
 
 class JumpStartModelsCache:
@@ -74,6 +78,7 @@ def __init__(
  s3_bucket_name: Optional[str] = None,
  s3_client_config: Optional[botocore.config.Config] = None,
  s3_client: Optional[boto3.client] = None,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ) -> None: # fmt: on
  """Initialize a ``JumpStartModelsCache`` instance.
 
@@ -95,6 +100,8 @@ def __init__(
  s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
  Default: None (no config).
  s3_client (Optional[boto3.client]): s3 client to use. Default: None.
+ sagemaker_session (Optional[sagemaker.session.Session]): A SageMaker Session object,
+ used for SageMaker interactions. Default: Session in region associated with boto3 session.
  """
 
  self._region = region
@@ -121,6 +128,7 @@ def __init__(
  if s3_client_config
  else boto3.client("s3", region_name=self._region)
  )
+ self._sagemaker_session = sagemaker_session
 
  def set_region(self, region: str) -> None:
  """Set region for cache. Clears cache after new region is set."""
@@ -340,32 +348,34 @@ def _retrieval_function(
  formatted_content=model_specs
  )
  if data_type == HubContentType.MODEL:
-info = get_info_from_hub_resource_arn(
+hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
  id_info
  )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_content = hub.describe_model(
- model_name=info.hub_content_name, model_version=info.hub_content_version
+ hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+ hub_name=hub_name,
+ hub_content_name=model_name,
+ hub_content_version=model_version,
+ hub_content_type=data_type
  )
+
+ model_specs = JumpStartModelSpecs(DescribeHubContentsResponse(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)
  return JumpStartCachedContentValue(
  formatted_content=model_specs
  )
- if data_type == HubContentType.HUB:
- info = get_info_from_hub_resource_arn(
- id_info
- )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_info = hub.describe()
- return JumpStartCachedContentValue(formatted_content=hub_info)
+ if data_type == HubType.HUB:
+ hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
+ response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)
+ hub_description = DescribeHubResponse(response)
+ return JumpStartCachedContentValue(formatted_content=DescribeHubResponse(hub_description))
  raise ValueError(
- f"Bad value for key '{key}': must be in",
- f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubContentType.HUB, HubContentType.MODEL]}"
+ f"Bad value for key '{key}': must be in",
+ f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubType.HUB, HubContentType.MODEL]}"
  )
 
  def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -490,7 +500,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(HubContentType.HUB, hub_arn))
+ details, _ = self._content_cache.get(JumpStartCachedContentKey(HubType.HUB, hub_arn))
  return details.formatted_content
 
  def clear(self) -> None:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,10 +13,16 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
+from typing import Any, Dict, Optional
 from sagemaker.session import Session
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+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:
@@ -25,30 +31,85 @@ class CuratedHub:
  def __init__(
  self,
  hub_name: str,
- region: str,
- session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ):
+ """Instantiates a SageMaker ``CuratedHub``.
+
+ Args:
+ hub_name (str): The name of the Hub to create.
+ sagemaker_session (sagemaker.session.Session): A SageMaker Session
+ object, used for SageMaker interactions.
+ """
  self.hub_name = hub_name
- self.region = region
- self._sm_session = session
+ self.region = sagemaker_session.boto_region_name
+ self._sagemaker_session = sagemaker_session
 
- def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
- """Returns descriptive information about the Hub Model"""
+ def create(
+ self,
+ description: str,
+ display_name: Optional[str] = None,
+ search_keywords: Optional[str] = None,
+ bucket_name: Optional[str] = None,
+ tags: Optional[str] = None,
+ ) -> Dict[str, str]:
+ """Creates a hub with the given description"""
 
- hub_content = self._sm_session.describe_hub_content(
- model_name, "Model", self.hub_name, model_version
+ bucket_name = 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,
  )
 
- # TODO: Parse HubContent
-# TODO: Parse HubContentDocument
+def describe(self) -> DescribeHubResponse:
+"""Returns descriptive information about the Hub"""
 
- return hub_content
+ hub_description: DescribeHubResponse = self._sagemaker_session.describe_hub(
+ hub_name=self.hub_name
+ )
 
- def describe(self) -> Dict[str, Any]:
- """Returns descriptive information about the Hub"""
+ return hub_description
 
- hub_info = self._sm_session.describe_hub(hub_name=self.hub_name)
+ def list_models(self, **kwargs) -> Dict[str, Any]:
+ """Lists the models in this Curated Hub
 
- # TODO: Validations?
+ **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
+
+ def describe_model(
+ self, model_name: str, model_version: str = "*"
+ ) -> DescribeHubContentsResponse:
+ """Returns descriptive information about the Hub Model"""
+
+ 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 DescribeHubContentsResponse(hub_content_description)
+
+ 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)

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -14,12 +14,13 @@
 from __future__ import absolute_import
 import re
 from typing import Optional
-from sagemaker.jumpstart import constants
-
-from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
-from sagemaker.jumpstart.types import HubContentType
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.types import (
+ HubContentType,
+ HubArnExtractedInfo,
+)
+from sagemaker.jumpstart import constants
 
 
 def get_info_from_hub_resource_arn(
@@ -109,3 +110,45 @@ def generate_hub_arn_for_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

@@ -15,11 +15,11 @@
 from __future__ import absolute_import
 
 from typing import Optional, Tuple
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
-from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 
 
 def get_model_id_version_from_endpoint(

src/sagemaker/jumpstart/cache.py

@@ -21,34 +21,38 @@
 import botocore
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from sagemaker.session import Session
+from sagemaker.utilities.cache import LRUCache
 from sagemaker.jumpstart.constants import (
  ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
  ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
  JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
  JUMPSTART_DEFAULT_REGION_NAME,
  JUMPSTART_LOGGER,
  MODEL_ID_LIST_WEB_URL,
+ DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.utils import get_info_from_hub_resource_arn
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
  JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
  JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
  JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
  JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
 )
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.types import (
  JumpStartCachedContentKey,
  JumpStartCachedContentValue,
  JumpStartModelHeader,
  JumpStartModelSpecs,
  JumpStartS3FileType,
  JumpStartVersionedModelId,
+ DescribeHubResponse,
+ DescribeHubContentsResponse,
+ HubType,
  HubContentType,
 )
-from sagemaker.jumpstart import utils
-from sagemaker.utilities.cache import LRUCache
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 
 
 class JumpStartModelsCache:
@@ -74,6 +78,7 @@ def __init__(
  s3_bucket_name: Optional[str] = None,
  s3_client_config: Optional[botocore.config.Config] = None,
  s3_client: Optional[boto3.client] = None,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ) -> None: # fmt: on
  """Initialize a ``JumpStartModelsCache`` instance.
 
@@ -95,6 +100,8 @@ def __init__(
  s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
  Default: None (no config).
  s3_client (Optional[boto3.client]): s3 client to use. Default: None.
+ sagemaker_session (Optional[sagemaker.session.Session]): A SageMaker Session object,
+ used for SageMaker interactions. Default: Session in region associated with boto3 session.
  """
 
  self._region = region
@@ -121,6 +128,7 @@ def __init__(
  if s3_client_config
  else boto3.client("s3", region_name=self._region)
  )
+ self._sagemaker_session = sagemaker_session
 
  def set_region(self, region: str) -> None:
  """Set region for cache. Clears cache after new region is set."""
@@ -340,32 +348,34 @@ def _retrieval_function(
  formatted_content=model_specs
  )
  if data_type == HubContentType.MODEL:
-info = get_info_from_hub_resource_arn(
+hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
  id_info
  )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_content = hub.describe_model(
- model_name=info.hub_content_name, model_version=info.hub_content_version
+ hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+ hub_name=hub_name,
+ hub_content_name=model_name,
+ hub_content_version=model_version,
+ hub_content_type=data_type
  )
+
+ model_specs = JumpStartModelSpecs(DescribeHubContentsResponse(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)
  return JumpStartCachedContentValue(
  formatted_content=model_specs
  )
- if data_type == HubContentType.HUB:
- info = get_info_from_hub_resource_arn(
- id_info
- )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_info = hub.describe()
- return JumpStartCachedContentValue(formatted_content=hub_info)
+ if data_type == HubType.HUB:
+ hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
+ response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)
+ hub_description = DescribeHubResponse(response)
+ return JumpStartCachedContentValue(formatted_content=DescribeHubResponse(hub_description))
  raise ValueError(
- f"Bad value for key '{key}': must be in",
- f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubContentType.HUB, HubContentType.MODEL]}"
+ f"Bad value for key '{key}': must be in",
+ f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubType.HUB, HubContentType.MODEL]}"
  )
 
  def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -490,7 +500,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(HubContentType.HUB, hub_arn))
+ details, _ = self._content_cache.get(JumpStartCachedContentKey(HubType.HUB, hub_arn))
  return details.formatted_content
 
  def clear(self) -> None:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,10 +13,16 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
+from typing import Any, Dict, Optional
 from sagemaker.session import Session
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+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:
@@ -25,30 +31,85 @@ class CuratedHub:
  def __init__(
  self,
  hub_name: str,
- region: str,
- session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ):
+ """Instantiates a SageMaker ``CuratedHub``.
+
+ Args:
+ hub_name (str): The name of the Hub to create.
+ sagemaker_session (sagemaker.session.Session): A SageMaker Session
+ object, used for SageMaker interactions.
+ """
  self.hub_name = hub_name
- self.region = region
- self._sm_session = session
+ self.region = sagemaker_session.boto_region_name
+ self._sagemaker_session = sagemaker_session
 
- def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
- """Returns descriptive information about the Hub Model"""
+ def create(
+ self,
+ description: str,
+ display_name: Optional[str] = None,
+ search_keywords: Optional[str] = None,
+ bucket_name: Optional[str] = None,
+ tags: Optional[str] = None,
+ ) -> Dict[str, str]:
+ """Creates a hub with the given description"""
 
- hub_content = self._sm_session.describe_hub_content(
- model_name, "Model", self.hub_name, model_version
+ bucket_name = 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,
  )
 
- # TODO: Parse HubContent
-# TODO: Parse HubContentDocument
+def describe(self) -> DescribeHubResponse:
+"""Returns descriptive information about the Hub"""
 
- return hub_content
+ hub_description: DescribeHubResponse = self._sagemaker_session.describe_hub(
+ hub_name=self.hub_name
+ )
 
- def describe(self) -> Dict[str, Any]:
- """Returns descriptive information about the Hub"""
+ return hub_description
 
- hub_info = self._sm_session.describe_hub(hub_name=self.hub_name)
+ def list_models(self, **kwargs) -> Dict[str, Any]:
+ """Lists the models in this Curated Hub
 
- # TODO: Validations?
+ **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
+
+ def describe_model(
+ self, model_name: str, model_version: str = "*"
+ ) -> DescribeHubContentsResponse:
+ """Returns descriptive information about the Hub Model"""
+
+ 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 DescribeHubContentsResponse(hub_content_description)
+
+ 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)

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -14,12 +14,13 @@
 from __future__ import absolute_import
 import re
 from typing import Optional
-from sagemaker.jumpstart import constants
-
-from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
-from sagemaker.jumpstart.types import HubContentType
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.types import (
+ HubContentType,
+ HubArnExtractedInfo,
+)
+from sagemaker.jumpstart import constants
 
 
 def get_info_from_hub_resource_arn(
@@ -109,3 +110,45 @@ def generate_hub_arn_for_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

@@ -15,11 +15,11 @@
 from __future__ import absolute_import
 
 from typing import Optional, Tuple
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
-from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 
 
 def get_model_id_version_from_endpoint(

src/sagemaker/jumpstart/cache.py

@@ -21,34 +21,38 @@
 import botocore
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from sagemaker.session import Session
+from sagemaker.utilities.cache import LRUCache
 from sagemaker.jumpstart.constants import (
  ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
  ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
  JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
  JUMPSTART_DEFAULT_REGION_NAME,
  JUMPSTART_LOGGER,
  MODEL_ID_LIST_WEB_URL,
+ DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.utils import get_info_from_hub_resource_arn
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
  JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
  JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
  JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
  JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
 )
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.types import (
  JumpStartCachedContentKey,
  JumpStartCachedContentValue,
  JumpStartModelHeader,
  JumpStartModelSpecs,
  JumpStartS3FileType,
  JumpStartVersionedModelId,
+ DescribeHubResponse,
+ DescribeHubContentsResponse,
+ HubType,
  HubContentType,
 )
-from sagemaker.jumpstart import utils
-from sagemaker.utilities.cache import LRUCache
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 
 
 class JumpStartModelsCache:
@@ -74,6 +78,7 @@ def __init__(
  s3_bucket_name: Optional[str] = None,
  s3_client_config: Optional[botocore.config.Config] = None,
  s3_client: Optional[boto3.client] = None,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ) -> None: # fmt: on
  """Initialize a ``JumpStartModelsCache`` instance.
 
@@ -95,6 +100,8 @@ def __init__(
  s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
  Default: None (no config).
  s3_client (Optional[boto3.client]): s3 client to use. Default: None.
+ sagemaker_session (Optional[sagemaker.session.Session]): A SageMaker Session object,
+ used for SageMaker interactions. Default: Session in region associated with boto3 session.
  """
 
  self._region = region
@@ -121,6 +128,7 @@ def __init__(
  if s3_client_config
  else boto3.client("s3", region_name=self._region)
  )
+ self._sagemaker_session = sagemaker_session
 
  def set_region(self, region: str) -> None:
  """Set region for cache. Clears cache after new region is set."""
@@ -340,32 +348,34 @@ def _retrieval_function(
  formatted_content=model_specs
  )
  if data_type == HubContentType.MODEL:
-info = get_info_from_hub_resource_arn(
+hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
  id_info
  )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_content = hub.describe_model(
- model_name=info.hub_content_name, model_version=info.hub_content_version
+ hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+ hub_name=hub_name,
+ hub_content_name=model_name,
+ hub_content_version=model_version,
+ hub_content_type=data_type
  )
+
+ model_specs = JumpStartModelSpecs(DescribeHubContentsResponse(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)
  return JumpStartCachedContentValue(
  formatted_content=model_specs
  )
- if data_type == HubContentType.HUB:
- info = get_info_from_hub_resource_arn(
- id_info
- )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_info = hub.describe()
- return JumpStartCachedContentValue(formatted_content=hub_info)
+ if data_type == HubType.HUB:
+ hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
+ response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)
+ hub_description = DescribeHubResponse(response)
+ return JumpStartCachedContentValue(formatted_content=DescribeHubResponse(hub_description))
  raise ValueError(
- f"Bad value for key '{key}': must be in",
- f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubContentType.HUB, HubContentType.MODEL]}"
+ f"Bad value for key '{key}': must be in",
+ f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubType.HUB, HubContentType.MODEL]}"
  )
 
  def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -490,7 +500,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(HubContentType.HUB, hub_arn))
+ details, _ = self._content_cache.get(JumpStartCachedContentKey(HubType.HUB, hub_arn))
  return details.formatted_content
 
  def clear(self) -> None:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,10 +13,16 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
+from typing import Any, Dict, Optional
 from sagemaker.session import Session
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+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:
@@ -25,30 +31,85 @@ class CuratedHub:
  def __init__(
  self,
  hub_name: str,
- region: str,
- session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ):
+ """Instantiates a SageMaker ``CuratedHub``.
+
+ Args:
+ hub_name (str): The name of the Hub to create.
+ sagemaker_session (sagemaker.session.Session): A SageMaker Session
+ object, used for SageMaker interactions.
+ """
  self.hub_name = hub_name
- self.region = region
- self._sm_session = session
+ self.region = sagemaker_session.boto_region_name
+ self._sagemaker_session = sagemaker_session
 
- def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
- """Returns descriptive information about the Hub Model"""
+ def create(
+ self,
+ description: str,
+ display_name: Optional[str] = None,
+ search_keywords: Optional[str] = None,
+ bucket_name: Optional[str] = None,
+ tags: Optional[str] = None,
+ ) -> Dict[str, str]:
+ """Creates a hub with the given description"""
 
- hub_content = self._sm_session.describe_hub_content(
- model_name, "Model", self.hub_name, model_version
+ bucket_name = 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,
  )
 
- # TODO: Parse HubContent
-# TODO: Parse HubContentDocument
+def describe(self) -> DescribeHubResponse:
+"""Returns descriptive information about the Hub"""
 
- return hub_content
+ hub_description: DescribeHubResponse = self._sagemaker_session.describe_hub(
+ hub_name=self.hub_name
+ )
 
- def describe(self) -> Dict[str, Any]:
- """Returns descriptive information about the Hub"""
+ return hub_description
 
- hub_info = self._sm_session.describe_hub(hub_name=self.hub_name)
+ def list_models(self, **kwargs) -> Dict[str, Any]:
+ """Lists the models in this Curated Hub
 
- # TODO: Validations?
+ **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
+
+ def describe_model(
+ self, model_name: str, model_version: str = "*"
+ ) -> DescribeHubContentsResponse:
+ """Returns descriptive information about the Hub Model"""
+
+ 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 DescribeHubContentsResponse(hub_content_description)
+
+ 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)

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -14,12 +14,13 @@
 from __future__ import absolute_import
 import re
 from typing import Optional
-from sagemaker.jumpstart import constants
-
-from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
-from sagemaker.jumpstart.types import HubContentType
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.types import (
+ HubContentType,
+ HubArnExtractedInfo,
+)
+from sagemaker.jumpstart import constants
 
 
 def get_info_from_hub_resource_arn(
@@ -109,3 +110,45 @@ def generate_hub_arn_for_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

@@ -15,11 +15,11 @@
 from __future__ import absolute_import
 
 from typing import Optional, Tuple
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
-from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 
 
 def get_model_id_version_from_endpoint(

src/sagemaker/jumpstart/cache.py

@@ -21,34 +21,38 @@
 import botocore
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from sagemaker.session import Session
+from sagemaker.utilities.cache import LRUCache
 from sagemaker.jumpstart.constants import (
  ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
  ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
  JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
  JUMPSTART_DEFAULT_REGION_NAME,
  JUMPSTART_LOGGER,
  MODEL_ID_LIST_WEB_URL,
+ DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.utils import get_info_from_hub_resource_arn
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
  JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
  JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
  JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
  JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
 )
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.types import (
  JumpStartCachedContentKey,
  JumpStartCachedContentValue,
  JumpStartModelHeader,
  JumpStartModelSpecs,
  JumpStartS3FileType,
  JumpStartVersionedModelId,
+ DescribeHubResponse,
+ DescribeHubContentsResponse,
+ HubType,
  HubContentType,
 )
-from sagemaker.jumpstart import utils
-from sagemaker.utilities.cache import LRUCache
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 
 
 class JumpStartModelsCache:
@@ -74,6 +78,7 @@ def __init__(
  s3_bucket_name: Optional[str] = None,
  s3_client_config: Optional[botocore.config.Config] = None,
  s3_client: Optional[boto3.client] = None,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ) -> None: # fmt: on
  """Initialize a ``JumpStartModelsCache`` instance.
 
@@ -95,6 +100,8 @@ def __init__(
  s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
  Default: None (no config).
  s3_client (Optional[boto3.client]): s3 client to use. Default: None.
+ sagemaker_session (Optional[sagemaker.session.Session]): A SageMaker Session object,
+ used for SageMaker interactions. Default: Session in region associated with boto3 session.
  """
 
  self._region = region
@@ -121,6 +128,7 @@ def __init__(
  if s3_client_config
  else boto3.client("s3", region_name=self._region)
  )
+ self._sagemaker_session = sagemaker_session
 
  def set_region(self, region: str) -> None:
  """Set region for cache. Clears cache after new region is set."""
@@ -340,32 +348,34 @@ def _retrieval_function(
  formatted_content=model_specs
  )
  if data_type == HubContentType.MODEL:
-info = get_info_from_hub_resource_arn(
+hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
  id_info
  )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_content = hub.describe_model(
- model_name=info.hub_content_name, model_version=info.hub_content_version
+ hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+ hub_name=hub_name,
+ hub_content_name=model_name,
+ hub_content_version=model_version,
+ hub_content_type=data_type
  )
+
+ model_specs = JumpStartModelSpecs(DescribeHubContentsResponse(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)
  return JumpStartCachedContentValue(
  formatted_content=model_specs
  )
- if data_type == HubContentType.HUB:
- info = get_info_from_hub_resource_arn(
- id_info
- )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_info = hub.describe()
- return JumpStartCachedContentValue(formatted_content=hub_info)
+ if data_type == HubType.HUB:
+ hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
+ response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)
+ hub_description = DescribeHubResponse(response)
+ return JumpStartCachedContentValue(formatted_content=DescribeHubResponse(hub_description))
  raise ValueError(
- f"Bad value for key '{key}': must be in",
- f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubContentType.HUB, HubContentType.MODEL]}"
+ f"Bad value for key '{key}': must be in",
+ f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubType.HUB, HubContentType.MODEL]}"
  )
 
  def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -490,7 +500,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(HubContentType.HUB, hub_arn))
+ details, _ = self._content_cache.get(JumpStartCachedContentKey(HubType.HUB, hub_arn))
  return details.formatted_content
 
  def clear(self) -> None:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,10 +13,16 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
+from typing import Any, Dict, Optional
 from sagemaker.session import Session
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+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:
@@ -25,30 +31,85 @@ class CuratedHub:
  def __init__(
  self,
  hub_name: str,
- region: str,
- session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ):
+ """Instantiates a SageMaker ``CuratedHub``.
+
+ Args:
+ hub_name (str): The name of the Hub to create.
+ sagemaker_session (sagemaker.session.Session): A SageMaker Session
+ object, used for SageMaker interactions.
+ """
  self.hub_name = hub_name
- self.region = region
- self._sm_session = session
+ self.region = sagemaker_session.boto_region_name
+ self._sagemaker_session = sagemaker_session
 
- def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
- """Returns descriptive information about the Hub Model"""
+ def create(
+ self,
+ description: str,
+ display_name: Optional[str] = None,
+ search_keywords: Optional[str] = None,
+ bucket_name: Optional[str] = None,
+ tags: Optional[str] = None,
+ ) -> Dict[str, str]:
+ """Creates a hub with the given description"""
 
- hub_content = self._sm_session.describe_hub_content(
- model_name, "Model", self.hub_name, model_version
+ bucket_name = 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,
  )
 
- # TODO: Parse HubContent
-# TODO: Parse HubContentDocument
+def describe(self) -> DescribeHubResponse:
+"""Returns descriptive information about the Hub"""
 
- return hub_content
+ hub_description: DescribeHubResponse = self._sagemaker_session.describe_hub(
+ hub_name=self.hub_name
+ )
 
- def describe(self) -> Dict[str, Any]:
- """Returns descriptive information about the Hub"""
+ return hub_description
 
- hub_info = self._sm_session.describe_hub(hub_name=self.hub_name)
+ def list_models(self, **kwargs) -> Dict[str, Any]:
+ """Lists the models in this Curated Hub
 
- # TODO: Validations?
+ **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
+
+ def describe_model(
+ self, model_name: str, model_version: str = "*"
+ ) -> DescribeHubContentsResponse:
+ """Returns descriptive information about the Hub Model"""
+
+ 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 DescribeHubContentsResponse(hub_content_description)
+
+ 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)

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -14,12 +14,13 @@
 from __future__ import absolute_import
 import re
 from typing import Optional
-from sagemaker.jumpstart import constants
-
-from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
-from sagemaker.jumpstart.types import HubContentType
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.types import (
+ HubContentType,
+ HubArnExtractedInfo,
+)
+from sagemaker.jumpstart import constants
 
 
 def get_info_from_hub_resource_arn(
@@ -109,3 +110,45 @@ def generate_hub_arn_for_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

@@ -15,11 +15,11 @@
 from __future__ import absolute_import
 
 from typing import Optional, Tuple
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
-from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 
 
 def get_model_id_version_from_endpoint(

src/sagemaker/jumpstart/cache.py

@@ -21,34 +21,38 @@
 import botocore
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from sagemaker.session import Session
+from sagemaker.utilities.cache import LRUCache
 from sagemaker.jumpstart.constants import (
  ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
  ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
  JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
  JUMPSTART_DEFAULT_REGION_NAME,
  JUMPSTART_LOGGER,
  MODEL_ID_LIST_WEB_URL,
+ DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.utils import get_info_from_hub_resource_arn
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
  JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
  JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
  JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
  JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
 )
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.types import (
  JumpStartCachedContentKey,
  JumpStartCachedContentValue,
  JumpStartModelHeader,
  JumpStartModelSpecs,
  JumpStartS3FileType,
  JumpStartVersionedModelId,
+ DescribeHubResponse,
+ DescribeHubContentsResponse,
+ HubType,
  HubContentType,
 )
-from sagemaker.jumpstart import utils
-from sagemaker.utilities.cache import LRUCache
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 
 
 class JumpStartModelsCache:
@@ -74,6 +78,7 @@ def __init__(
  s3_bucket_name: Optional[str] = None,
  s3_client_config: Optional[botocore.config.Config] = None,
  s3_client: Optional[boto3.client] = None,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ) -> None: # fmt: on
  """Initialize a ``JumpStartModelsCache`` instance.
 
@@ -95,6 +100,8 @@ def __init__(
  s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
  Default: None (no config).
  s3_client (Optional[boto3.client]): s3 client to use. Default: None.
+ sagemaker_session (Optional[sagemaker.session.Session]): A SageMaker Session object,
+ used for SageMaker interactions. Default: Session in region associated with boto3 session.
  """
 
  self._region = region
@@ -121,6 +128,7 @@ def __init__(
  if s3_client_config
  else boto3.client("s3", region_name=self._region)
  )
+ self._sagemaker_session = sagemaker_session
 
  def set_region(self, region: str) -> None:
  """Set region for cache. Clears cache after new region is set."""
@@ -340,32 +348,34 @@ def _retrieval_function(
  formatted_content=model_specs
  )
  if data_type == HubContentType.MODEL:
-info = get_info_from_hub_resource_arn(
+hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
  id_info
  )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_content = hub.describe_model(
- model_name=info.hub_content_name, model_version=info.hub_content_version
+ hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+ hub_name=hub_name,
+ hub_content_name=model_name,
+ hub_content_version=model_version,
+ hub_content_type=data_type
  )
+
+ model_specs = JumpStartModelSpecs(DescribeHubContentsResponse(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)
  return JumpStartCachedContentValue(
  formatted_content=model_specs
  )
- if data_type == HubContentType.HUB:
- info = get_info_from_hub_resource_arn(
- id_info
- )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_info = hub.describe()
- return JumpStartCachedContentValue(formatted_content=hub_info)
+ if data_type == HubType.HUB:
+ hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
+ response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)
+ hub_description = DescribeHubResponse(response)
+ return JumpStartCachedContentValue(formatted_content=DescribeHubResponse(hub_description))
  raise ValueError(
- f"Bad value for key '{key}': must be in",
- f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubContentType.HUB, HubContentType.MODEL]}"
+ f"Bad value for key '{key}': must be in",
+ f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubType.HUB, HubContentType.MODEL]}"
  )
 
  def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -490,7 +500,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(HubContentType.HUB, hub_arn))
+ details, _ = self._content_cache.get(JumpStartCachedContentKey(HubType.HUB, hub_arn))
  return details.formatted_content
 
  def clear(self) -> None:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,10 +13,16 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
+from typing import Any, Dict, Optional
 from sagemaker.session import Session
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+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:
@@ -25,30 +31,85 @@ class CuratedHub:
  def __init__(
  self,
  hub_name: str,
- region: str,
- session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ):
+ """Instantiates a SageMaker ``CuratedHub``.
+
+ Args:
+ hub_name (str): The name of the Hub to create.
+ sagemaker_session (sagemaker.session.Session): A SageMaker Session
+ object, used for SageMaker interactions.
+ """
  self.hub_name = hub_name
- self.region = region
- self._sm_session = session
+ self.region = sagemaker_session.boto_region_name
+ self._sagemaker_session = sagemaker_session
 
- def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
- """Returns descriptive information about the Hub Model"""
+ def create(
+ self,
+ description: str,
+ display_name: Optional[str] = None,
+ search_keywords: Optional[str] = None,
+ bucket_name: Optional[str] = None,
+ tags: Optional[str] = None,
+ ) -> Dict[str, str]:
+ """Creates a hub with the given description"""
 
- hub_content = self._sm_session.describe_hub_content(
- model_name, "Model", self.hub_name, model_version
+ bucket_name = 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,
  )
 
- # TODO: Parse HubContent
-# TODO: Parse HubContentDocument
+def describe(self) -> DescribeHubResponse:
+"""Returns descriptive information about the Hub"""
 
- return hub_content
+ hub_description: DescribeHubResponse = self._sagemaker_session.describe_hub(
+ hub_name=self.hub_name
+ )
 
- def describe(self) -> Dict[str, Any]:
- """Returns descriptive information about the Hub"""
+ return hub_description
 
- hub_info = self._sm_session.describe_hub(hub_name=self.hub_name)
+ def list_models(self, **kwargs) -> Dict[str, Any]:
+ """Lists the models in this Curated Hub
 
- # TODO: Validations?
+ **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
+
+ def describe_model(
+ self, model_name: str, model_version: str = "*"
+ ) -> DescribeHubContentsResponse:
+ """Returns descriptive information about the Hub Model"""
+
+ 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 DescribeHubContentsResponse(hub_content_description)
+
+ 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)

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -14,12 +14,13 @@
 from __future__ import absolute_import
 import re
 from typing import Optional
-from sagemaker.jumpstart import constants
-
-from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
-from sagemaker.jumpstart.types import HubContentType
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.types import (
+ HubContentType,
+ HubArnExtractedInfo,
+)
+from sagemaker.jumpstart import constants
 
 
 def get_info_from_hub_resource_arn(
@@ -109,3 +110,45 @@ def generate_hub_arn_for_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

@@ -15,11 +15,11 @@
 from __future__ import absolute_import
 
 from typing import Optional, Tuple
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
-from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 
 
 def get_model_id_version_from_endpoint(

src/sagemaker/jumpstart/cache.py

@@ -21,34 +21,38 @@
 import botocore
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from sagemaker.session import Session
+from sagemaker.utilities.cache import LRUCache
 from sagemaker.jumpstart.constants import (
  ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
  ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
  JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
  JUMPSTART_DEFAULT_REGION_NAME,
  JUMPSTART_LOGGER,
  MODEL_ID_LIST_WEB_URL,
+ DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.utils import get_info_from_hub_resource_arn
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
  JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
  JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
  JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
  JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
 )
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.types import (
  JumpStartCachedContentKey,
  JumpStartCachedContentValue,
  JumpStartModelHeader,
  JumpStartModelSpecs,
  JumpStartS3FileType,
  JumpStartVersionedModelId,
+ DescribeHubResponse,
+ DescribeHubContentsResponse,
+ HubType,
  HubContentType,
 )
-from sagemaker.jumpstart import utils
-from sagemaker.utilities.cache import LRUCache
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 
 
 class JumpStartModelsCache:
@@ -74,6 +78,7 @@ def __init__(
  s3_bucket_name: Optional[str] = None,
  s3_client_config: Optional[botocore.config.Config] = None,
  s3_client: Optional[boto3.client] = None,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ) -> None: # fmt: on
  """Initialize a ``JumpStartModelsCache`` instance.
 
@@ -95,6 +100,8 @@ def __init__(
  s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
  Default: None (no config).
  s3_client (Optional[boto3.client]): s3 client to use. Default: None.
+ sagemaker_session (Optional[sagemaker.session.Session]): A SageMaker Session object,
+ used for SageMaker interactions. Default: Session in region associated with boto3 session.
  """
 
  self._region = region
@@ -121,6 +128,7 @@ def __init__(
  if s3_client_config
  else boto3.client("s3", region_name=self._region)
  )
+ self._sagemaker_session = sagemaker_session
 
  def set_region(self, region: str) -> None:
  """Set region for cache. Clears cache after new region is set."""
@@ -340,32 +348,34 @@ def _retrieval_function(
  formatted_content=model_specs
  )
  if data_type == HubContentType.MODEL:
-info = get_info_from_hub_resource_arn(
+hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
  id_info
  )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_content = hub.describe_model(
- model_name=info.hub_content_name, model_version=info.hub_content_version
+ hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+ hub_name=hub_name,
+ hub_content_name=model_name,
+ hub_content_version=model_version,
+ hub_content_type=data_type
  )
+
+ model_specs = JumpStartModelSpecs(DescribeHubContentsResponse(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)
  return JumpStartCachedContentValue(
  formatted_content=model_specs
  )
- if data_type == HubContentType.HUB:
- info = get_info_from_hub_resource_arn(
- id_info
- )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_info = hub.describe()
- return JumpStartCachedContentValue(formatted_content=hub_info)
+ if data_type == HubType.HUB:
+ hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
+ response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)
+ hub_description = DescribeHubResponse(response)
+ return JumpStartCachedContentValue(formatted_content=DescribeHubResponse(hub_description))
  raise ValueError(
- f"Bad value for key '{key}': must be in",
- f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubContentType.HUB, HubContentType.MODEL]}"
+ f"Bad value for key '{key}': must be in",
+ f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubType.HUB, HubContentType.MODEL]}"
  )
 
  def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -490,7 +500,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(HubContentType.HUB, hub_arn))
+ details, _ = self._content_cache.get(JumpStartCachedContentKey(HubType.HUB, hub_arn))
  return details.formatted_content
 
  def clear(self) -> None:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,10 +13,16 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
+from typing import Any, Dict, Optional
 from sagemaker.session import Session
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+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:
@@ -25,30 +31,85 @@ class CuratedHub:
  def __init__(
  self,
  hub_name: str,
- region: str,
- session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ):
+ """Instantiates a SageMaker ``CuratedHub``.
+
+ Args:
+ hub_name (str): The name of the Hub to create.
+ sagemaker_session (sagemaker.session.Session): A SageMaker Session
+ object, used for SageMaker interactions.
+ """
  self.hub_name = hub_name
- self.region = region
- self._sm_session = session
+ self.region = sagemaker_session.boto_region_name
+ self._sagemaker_session = sagemaker_session
 
- def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
- """Returns descriptive information about the Hub Model"""
+ def create(
+ self,
+ description: str,
+ display_name: Optional[str] = None,
+ search_keywords: Optional[str] = None,
+ bucket_name: Optional[str] = None,
+ tags: Optional[str] = None,
+ ) -> Dict[str, str]:
+ """Creates a hub with the given description"""
 
- hub_content = self._sm_session.describe_hub_content(
- model_name, "Model", self.hub_name, model_version
+ bucket_name = 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,
  )
 
- # TODO: Parse HubContent
-# TODO: Parse HubContentDocument
+def describe(self) -> DescribeHubResponse:
+"""Returns descriptive information about the Hub"""
 
- return hub_content
+ hub_description: DescribeHubResponse = self._sagemaker_session.describe_hub(
+ hub_name=self.hub_name
+ )
 
- def describe(self) -> Dict[str, Any]:
- """Returns descriptive information about the Hub"""
+ return hub_description
 
- hub_info = self._sm_session.describe_hub(hub_name=self.hub_name)
+ def list_models(self, **kwargs) -> Dict[str, Any]:
+ """Lists the models in this Curated Hub
 
- # TODO: Validations?
+ **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
+
+ def describe_model(
+ self, model_name: str, model_version: str = "*"
+ ) -> DescribeHubContentsResponse:
+ """Returns descriptive information about the Hub Model"""
+
+ 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 DescribeHubContentsResponse(hub_content_description)
+
+ 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)

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -14,12 +14,13 @@
 from __future__ import absolute_import
 import re
 from typing import Optional
-from sagemaker.jumpstart import constants
-
-from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
-from sagemaker.jumpstart.types import HubContentType
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.types import (
+ HubContentType,
+ HubArnExtractedInfo,
+)
+from sagemaker.jumpstart import constants
 
 
 def get_info_from_hub_resource_arn(
@@ -109,3 +110,45 @@ def generate_hub_arn_for_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

@@ -15,11 +15,11 @@
 from __future__ import absolute_import
 
 from typing import Optional, Tuple
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
-from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 
 
 def get_model_id_version_from_endpoint(

src/sagemaker/jumpstart/cache.py

@@ -21,34 +21,38 @@
 import botocore
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from sagemaker.session import Session
+from sagemaker.utilities.cache import LRUCache
 from sagemaker.jumpstart.constants import (
  ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
  ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
  JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
  JUMPSTART_DEFAULT_REGION_NAME,
  JUMPSTART_LOGGER,
  MODEL_ID_LIST_WEB_URL,
+ DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.utils import get_info_from_hub_resource_arn
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
  JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
  JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
  JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
  JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
 )
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.types import (
  JumpStartCachedContentKey,
  JumpStartCachedContentValue,
  JumpStartModelHeader,
  JumpStartModelSpecs,
  JumpStartS3FileType,
  JumpStartVersionedModelId,
+ DescribeHubResponse,
+ DescribeHubContentsResponse,
+ HubType,
  HubContentType,
 )
-from sagemaker.jumpstart import utils
-from sagemaker.utilities.cache import LRUCache
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 
 
 class JumpStartModelsCache:
@@ -74,6 +78,7 @@ def __init__(
  s3_bucket_name: Optional[str] = None,
  s3_client_config: Optional[botocore.config.Config] = None,
  s3_client: Optional[boto3.client] = None,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ) -> None: # fmt: on
  """Initialize a ``JumpStartModelsCache`` instance.
 
@@ -95,6 +100,8 @@ def __init__(
  s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
  Default: None (no config).
  s3_client (Optional[boto3.client]): s3 client to use. Default: None.
+ sagemaker_session (Optional[sagemaker.session.Session]): A SageMaker Session object,
+ used for SageMaker interactions. Default: Session in region associated with boto3 session.
  """
 
  self._region = region
@@ -121,6 +128,7 @@ def __init__(
  if s3_client_config
  else boto3.client("s3", region_name=self._region)
  )
+ self._sagemaker_session = sagemaker_session
 
  def set_region(self, region: str) -> None:
  """Set region for cache. Clears cache after new region is set."""
@@ -340,32 +348,34 @@ def _retrieval_function(
  formatted_content=model_specs
  )
  if data_type == HubContentType.MODEL:
-info = get_info_from_hub_resource_arn(
+hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
  id_info
  )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_content = hub.describe_model(
- model_name=info.hub_content_name, model_version=info.hub_content_version
+ hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+ hub_name=hub_name,
+ hub_content_name=model_name,
+ hub_content_version=model_version,
+ hub_content_type=data_type
  )
+
+ model_specs = JumpStartModelSpecs(DescribeHubContentsResponse(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)
  return JumpStartCachedContentValue(
  formatted_content=model_specs
  )
- if data_type == HubContentType.HUB:
- info = get_info_from_hub_resource_arn(
- id_info
- )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_info = hub.describe()
- return JumpStartCachedContentValue(formatted_content=hub_info)
+ if data_type == HubType.HUB:
+ hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
+ response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)
+ hub_description = DescribeHubResponse(response)
+ return JumpStartCachedContentValue(formatted_content=DescribeHubResponse(hub_description))
  raise ValueError(
- f"Bad value for key '{key}': must be in",
- f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubContentType.HUB, HubContentType.MODEL]}"
+ f"Bad value for key '{key}': must be in",
+ f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubType.HUB, HubContentType.MODEL]}"
  )
 
  def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -490,7 +500,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(HubContentType.HUB, hub_arn))
+ details, _ = self._content_cache.get(JumpStartCachedContentKey(HubType.HUB, hub_arn))
  return details.formatted_content
 
  def clear(self) -> None:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,10 +13,16 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
+from typing import Any, Dict, Optional
 from sagemaker.session import Session
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+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:
@@ -25,30 +31,85 @@ class CuratedHub:
  def __init__(
  self,
  hub_name: str,
- region: str,
- session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ):
+ """Instantiates a SageMaker ``CuratedHub``.
+
+ Args:
+ hub_name (str): The name of the Hub to create.
+ sagemaker_session (sagemaker.session.Session): A SageMaker Session
+ object, used for SageMaker interactions.
+ """
  self.hub_name = hub_name
- self.region = region
- self._sm_session = session
+ self.region = sagemaker_session.boto_region_name
+ self._sagemaker_session = sagemaker_session
 
- def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
- """Returns descriptive information about the Hub Model"""
+ def create(
+ self,
+ description: str,
+ display_name: Optional[str] = None,
+ search_keywords: Optional[str] = None,
+ bucket_name: Optional[str] = None,
+ tags: Optional[str] = None,
+ ) -> Dict[str, str]:
+ """Creates a hub with the given description"""
 
- hub_content = self._sm_session.describe_hub_content(
- model_name, "Model", self.hub_name, model_version
+ bucket_name = 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,
  )
 
- # TODO: Parse HubContent
-# TODO: Parse HubContentDocument
+def describe(self) -> DescribeHubResponse:
+"""Returns descriptive information about the Hub"""
 
- return hub_content
+ hub_description: DescribeHubResponse = self._sagemaker_session.describe_hub(
+ hub_name=self.hub_name
+ )
 
- def describe(self) -> Dict[str, Any]:
- """Returns descriptive information about the Hub"""
+ return hub_description
 
- hub_info = self._sm_session.describe_hub(hub_name=self.hub_name)
+ def list_models(self, **kwargs) -> Dict[str, Any]:
+ """Lists the models in this Curated Hub
 
- # TODO: Validations?
+ **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
+
+ def describe_model(
+ self, model_name: str, model_version: str = "*"
+ ) -> DescribeHubContentsResponse:
+ """Returns descriptive information about the Hub Model"""
+
+ 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 DescribeHubContentsResponse(hub_content_description)
+
+ 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)

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -14,12 +14,13 @@
 from __future__ import absolute_import
 import re
 from typing import Optional
-from sagemaker.jumpstart import constants
-
-from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
-from sagemaker.jumpstart.types import HubContentType
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.types import (
+ HubContentType,
+ HubArnExtractedInfo,
+)
+from sagemaker.jumpstart import constants
 
 
 def get_info_from_hub_resource_arn(
@@ -109,3 +110,45 @@ def generate_hub_arn_for_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

@@ -15,11 +15,11 @@
 from __future__ import absolute_import
 
 from typing import Optional, Tuple
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
-from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 
 
 def get_model_id_version_from_endpoint(

src/sagemaker/jumpstart/cache.py

@@ -21,34 +21,38 @@
 import botocore
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from sagemaker.session import Session
+from sagemaker.utilities.cache import LRUCache
 from sagemaker.jumpstart.constants import (
  ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
  ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
  JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
  JUMPSTART_DEFAULT_REGION_NAME,
  JUMPSTART_LOGGER,
  MODEL_ID_LIST_WEB_URL,
+ DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.utils import get_info_from_hub_resource_arn
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
  JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
  JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
  JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
  JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
 )
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.types import (
  JumpStartCachedContentKey,
  JumpStartCachedContentValue,
  JumpStartModelHeader,
  JumpStartModelSpecs,
  JumpStartS3FileType,
  JumpStartVersionedModelId,
+ DescribeHubResponse,
+ DescribeHubContentsResponse,
+ HubType,
  HubContentType,
 )
-from sagemaker.jumpstart import utils
-from sagemaker.utilities.cache import LRUCache
+from sagemaker.jumpstart.curated_hub import utils as hub_utils
 
 
 class JumpStartModelsCache:
@@ -74,6 +78,7 @@ def __init__(
  s3_bucket_name: Optional[str] = None,
  s3_client_config: Optional[botocore.config.Config] = None,
  s3_client: Optional[boto3.client] = None,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ) -> None: # fmt: on
  """Initialize a ``JumpStartModelsCache`` instance.
 
@@ -95,6 +100,8 @@ def __init__(
  s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
  Default: None (no config).
  s3_client (Optional[boto3.client]): s3 client to use. Default: None.
+ sagemaker_session (Optional[sagemaker.session.Session]): A SageMaker Session object,
+ used for SageMaker interactions. Default: Session in region associated with boto3 session.
  """
 
  self._region = region
@@ -121,6 +128,7 @@ def __init__(
  if s3_client_config
  else boto3.client("s3", region_name=self._region)
  )
+ self._sagemaker_session = sagemaker_session
 
  def set_region(self, region: str) -> None:
  """Set region for cache. Clears cache after new region is set."""
@@ -340,32 +348,34 @@ def _retrieval_function(
  formatted_content=model_specs
  )
  if data_type == HubContentType.MODEL:
-info = get_info_from_hub_resource_arn(
+hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
  id_info
  )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_content = hub.describe_model(
- model_name=info.hub_content_name, model_version=info.hub_content_version
+ hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+ hub_name=hub_name,
+ hub_content_name=model_name,
+ hub_content_version=model_version,
+ hub_content_type=data_type
  )
+
+ model_specs = JumpStartModelSpecs(DescribeHubContentsResponse(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)
  return JumpStartCachedContentValue(
  formatted_content=model_specs
  )
- if data_type == HubContentType.HUB:
- info = get_info_from_hub_resource_arn(
- id_info
- )
- hub = CuratedHub(hub_name=info.hub_name, region=info.region)
- hub_info = hub.describe()
- return JumpStartCachedContentValue(formatted_content=hub_info)
+ if data_type == HubType.HUB:
+ hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
+ response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)
+ hub_description = DescribeHubResponse(response)
+ return JumpStartCachedContentValue(formatted_content=DescribeHubResponse(hub_description))
  raise ValueError(
- f"Bad value for key '{key}': must be in",
- f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubContentType.HUB, HubContentType.MODEL]}"
+ f"Bad value for key '{key}': must be in",
+ f"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubType.HUB, HubContentType.MODEL]}"
  )
 
  def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -490,7 +500,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(HubContentType.HUB, hub_arn))
+ details, _ = self._content_cache.get(JumpStartCachedContentKey(HubType.HUB, hub_arn))
  return details.formatted_content
 
  def clear(self) -> None:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,10 +13,16 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
+from typing import Any, Dict, Optional
 from sagemaker.session import Session
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+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:
@@ -25,30 +31,85 @@ class CuratedHub:
  def __init__(
  self,
  hub_name: str,
- region: str,
- session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+ sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  ):
+ """Instantiates a SageMaker ``CuratedHub``.
+
+ Args:
+ hub_name (str): The name of the Hub to create.
+ sagemaker_session (sagemaker.session.Session): A SageMaker Session
+ object, used for SageMaker interactions.
+ """
  self.hub_name = hub_name
- self.region = region
- self._sm_session = session
+ self.region = sagemaker_session.boto_region_name
+ self._sagemaker_session = sagemaker_session
 
- def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
- """Returns descriptive information about the Hub Model"""
+ def create(
+ self,
+ description: str,
+ display_name: Optional[str] = None,
+ search_keywords: Optional[str] = None,
+ bucket_name: Optional[str] = None,
+ tags: Optional[str] = None,
+ ) -> Dict[str, str]:
+ """Creates a hub with the given description"""
 
- hub_content = self._sm_session.describe_hub_content(
- model_name, "Model", self.hub_name, model_version
+ bucket_name = 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,
  )
 
- # TODO: Parse HubContent
-# TODO: Parse HubContentDocument
+def describe(self) -> DescribeHubResponse:
+"""Returns descriptive information about the Hub"""
 
- return hub_content
+ hub_description: DescribeHubResponse = self._sagemaker_session.describe_hub(
+ hub_name=self.hub_name
+ )
 
- def describe(self) -> Dict[str, Any]:
- """Returns descriptive information about the Hub"""
+ return hub_description
 
- hub_info = self._sm_session.describe_hub(hub_name=self.hub_name)
+ def list_models(self, **kwargs) -> Dict[str, Any]:
+ """Lists the models in this Curated Hub
 
- # TODO: Validations?
+ **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
+
+ def describe_model(
+ self, model_name: str, model_version: str = "*"
+ ) -> DescribeHubContentsResponse:
+ """Returns descriptive information about the Hub Model"""
+
+ 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 DescribeHubContentsResponse(hub_content_description)
+
+ 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)

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -14,12 +14,13 @@
 from __future__ import absolute_import
 import re
 from typing import Optional
-from sagemaker.jumpstart import constants
-
-from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
-from sagemaker.jumpstart.types import HubContentType
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.types import (
+ HubContentType,
+ HubArnExtractedInfo,
+)
+from sagemaker.jumpstart import constants
 
 
 def get_info_from_hub_resource_arn(
@@ -109,3 +110,45 @@ def generate_hub_arn_for_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

@@ -15,11 +15,11 @@
 from __future__ import absolute_import
 
 from typing import Optional, Tuple
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
-from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 
 
 def get_model_id_version_from_endpoint(