diff --git a/sagemaker-core/src/sagemaker/core/helper/session_helper.py b/sagemaker-core/src/sagemaker/core/helper/session_helper.py index e8eab15519..7ebcdca6e7 100644 --- a/sagemaker-core/src/sagemaker/core/helper/session_helper.py +++ b/sagemaker-core/src/sagemaker/core/helper/session_helper.py @@ -2167,6 +2167,101 @@ def list_hub_contents(self, hub_name, hub_content_type, **kwargs): request["NextToken"] = next_token return self.sagemaker_client.list_hub_contents(**request) + def create_hub( + self, + hub_name: str, + hub_description: str, + hub_display_name: str = None, + hub_search_keywords: List[str] = None, + s3_storage_config: Dict[str, Any] = None, + tags: List[Dict[str, Any]] = None, + ) -> Dict[str, str]: + """Creates a SageMaker Hub. + + Args: + hub_name (str): The name of the Hub to create. + hub_description (str): A description of the Hub. + hub_display_name (str): The display name of the Hub. Defaults to ``hub_name``. + hub_search_keywords (list): The searchable keywords for the Hub. + s3_storage_config (dict): The Amazon S3 storage configuration for the Hub. + tags (list): Any tags to associate with the Hub. + + Returns: + dict: Response from the CreateHub API. + """ + request = { + "HubName": hub_name, + "HubDescription": hub_description, + "HubDisplayName": hub_display_name or hub_name, + } + if hub_search_keywords: + request["HubSearchKeywords"] = hub_search_keywords + if s3_storage_config: + request["S3StorageConfig"] = s3_storage_config + if tags: + request["Tags"] = tags + return self.sagemaker_client.create_hub(**request) + + def describe_hub(self, hub_name: str) -> Dict[str, Any]: + """Describes a SageMaker Hub. + + Args: + hub_name (str): The name of the Hub to describe. + + Returns: + dict: Response from the DescribeHub API. + """ + return self.sagemaker_client.describe_hub(HubName=hub_name) + + def delete_hub(self, hub_name: str) -> None: + """Deletes a SageMaker Hub. + + Args: + hub_name (str): The name of the Hub to delete. + """ + return self.sagemaker_client.delete_hub(HubName=hub_name) + + def create_hub_content_reference( + self, + hub_name: str, + source_hub_content_arn: str, + hub_content_name: str = None, + min_version: str = None, + ) -> Dict[str, str]: + """Creates a reference to public hub content in a SageMaker Hub. + + Args: + hub_name (str): The name of the Hub to add the reference to. + source_hub_content_arn (str): The ARN of the hub content in the public Hub. + hub_content_name (str): The name of the reference to add to the Hub. + min_version (str): The minimum version of the hub content to reference. + + Returns: + dict: Response from the CreateHubContentReference API. + """ + request = {"HubName": hub_name, "SageMakerPublicHubContentArn": source_hub_content_arn} + if hub_content_name: + request["HubContentName"] = hub_content_name + if min_version: + request["MinVersion"] = min_version + return self.sagemaker_client.create_hub_content_reference(**request) + + def delete_hub_content_reference( + self, hub_name: str, hub_content_type: str, hub_content_name: str + ) -> None: + """Deletes a hub content reference from a SageMaker Hub. + + Args: + hub_name (str): The name of the Hub to delete the reference from. + hub_content_type (str): The type of the hub content reference. + hub_content_name (str): The name of the hub content reference. + """ + return self.sagemaker_client.delete_hub_content_reference( + HubName=hub_name, + HubContentType=hub_content_type, + HubContentName=hub_content_name, + ) + def _expand_container_def(c_def): """Placeholder docstring""" diff --git a/sagemaker-core/src/sagemaker/core/jumpstart/hub/hub.py b/sagemaker-core/src/sagemaker/core/jumpstart/hub/hub.py index 3d871ff22b..c33dba170f 100644 --- a/sagemaker-core/src/sagemaker/core/jumpstart/hub/hub.py +++ b/sagemaker-core/src/sagemaker/core/jumpstart/hub/hub.py @@ -68,12 +68,12 @@ def __init__( object, used for SageMaker interactions. """ self.hub_name = hub_name - self.region = sagemaker_session.boto_region_name self.bucket_name = bucket_name self._sagemaker_session = ( sagemaker_session or utils.get_default_jumpstart_session_with_user_agent_suffix(is_hub_content=True) ) + self.region = self._sagemaker_session.boto_region_name def _get_latest_model_version(self, model_id: str) -> str: """Populates the lastest version of a model from specs no matter what is passed. diff --git a/sagemaker-core/tests/unit/jumpstart/hub/test_hub.py b/sagemaker-core/tests/unit/jumpstart/hub/test_hub.py index 8a6c404f5c..2d1fac54a9 100644 --- a/sagemaker-core/tests/unit/jumpstart/hub/test_hub.py +++ b/sagemaker-core/tests/unit/jumpstart/hub/test_hub.py @@ -57,3 +57,81 @@ def list_hub_contents(**kwargs): "model-c", ] assert pages == {} + + +def _hub_with_session_client(): + """A Hub backed by a real Session whose SageMaker client is mocked.""" + from sagemaker.core.helper.session_helper import Session + + session = Session.__new__(Session) + session._region_name = "us-west-2" + session.sagemaker_client = Mock() + return Hub(hub_name=HUB_NAME, sagemaker_session=session), session.sagemaker_client + + +def test_create_calls_create_hub(): + hub, client = _hub_with_session_client() + + hub.create(description="my hub", search_keywords=["a"], tags=[{"Key": "k", "Value": "v"}]) + + client.create_hub.assert_called_once_with( + HubName=HUB_NAME, + HubDescription="my hub", + HubDisplayName=HUB_NAME, + HubSearchKeywords=["a"], + Tags=[{"Key": "k", "Value": "v"}], + ) + + +def test_describe_calls_describe_hub(): + hub, client = _hub_with_session_client() + + hub.describe() + + client.describe_hub.assert_called_once_with(HubName=HUB_NAME) + + +def test_delete_calls_delete_hub(): + hub, client = _hub_with_session_client() + + hub.delete() + + client.delete_hub.assert_called_once_with(HubName=HUB_NAME) + + +def test_create_model_reference_calls_create_hub_content_reference(): + hub, client = _hub_with_session_client() + model_arn = "arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/m/1.0.0" + + hub.create_model_reference(model_arn=model_arn, model_name="m", min_version="1.0.0") + + client.create_hub_content_reference.assert_called_once_with( + HubName=HUB_NAME, + SageMakerPublicHubContentArn=model_arn, + HubContentName="m", + MinVersion="1.0.0", + ) + + +def test_delete_model_reference_calls_delete_hub_content_reference(): + hub, client = _hub_with_session_client() + + hub.delete_model_reference(model_name="m") + + client.delete_hub_content_reference.assert_called_once_with( + HubName=HUB_NAME, HubContentType="ModelReference", HubContentName="m" + ) + + +def test_init_without_session_uses_default_jumpstart_session(): + from unittest.mock import patch + + default_session = Mock(boto_region_name="us-east-2") + with patch( + "sagemaker.core.jumpstart.hub.hub.utils.get_default_jumpstart_session_with_user_agent_suffix", + return_value=default_session, + ): + hub = Hub(hub_name=HUB_NAME, sagemaker_session=None) + + assert hub._sagemaker_session is default_session + assert hub.region == "us-east-2"