fix(core): restore the Session hub methods used by the JumpStart Hub class

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"""

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.

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"