remove region dependency for curatedhub

Author: jinyoung-lim

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -31,24 +31,17 @@ class CuratedHub:
     def __init__(
         self,
         hub_name: str,
-        region: str,
         sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     ):
         """Instantiates a SageMaker ``CuratedHub``.
 
         Args:
             hub_name (str): The name of the Hub to create.
-            region (str): The region in which the CuratedHub is in.
             sagemaker_session (sagemaker.session.Session): A SageMaker Session
                 object, used for SageMaker interactions.
         """
         self.hub_name = hub_name
-        if sagemaker_session.boto_region_name != region:
-            raise ValueError(
-                f"Cannot have conflicting regions for region=[{region}] and ",
-                f"sagemaker_session region=[{str(sagemaker_session.boto_region_name)}].",
-            )
-        self.region = region
+        self.region = sagemaker_session.boto_region_name
         self._sagemaker_session = sagemaker_session
 
     def create(
@@ -75,9 +68,11 @@ def create(
     def describe(self) -> DescribeHubResponse:
         """Returns descriptive information about the Hub"""
 
-        hub_description = self._sagemaker_session.describe_hub(hub_name=self.hub_name)
+        hub_description: DescribeHubResponse = self._sagemaker_session.describe_hub(
+            hub_name=self.hub_name
+        )
 
-        return DescribeHubResponse(hub_description)
+        return hub_description
 
     def list_models(self, **kwargs) -> Dict[str, Any]:
         """Lists the models in this Curated Hub

tests/unit/sagemaker/jumpstart/curated_hub/test_curated_hub.py

@@ -34,21 +34,12 @@ def sagemaker_session():
 
 
 def test_instantiates(sagemaker_session):
-    hub = CuratedHub(hub_name=HUB_NAME, region=REGION, sagemaker_session=sagemaker_session)
+    hub = CuratedHub(hub_name=HUB_NAME, sagemaker_session=sagemaker_session)
     assert hub.hub_name == HUB_NAME
     assert hub.region == "us-east-1"
     assert hub._sagemaker_session == sagemaker_session
 
 
-def test_instantiates_handles_conflicting_regions(sagemaker_session):
-    conflicting_region = "us-east-2"
-
-    with pytest.raises(ValueError):
-        CuratedHub(
-            hub_name=HUB_NAME, region=conflicting_region, sagemaker_session=sagemaker_session
-        )
-
-
 @pytest.mark.parametrize(
     ("hub_name,hub_description,hub_bucket_name,hub_display_name,hub_search_keywords,tags"),
     [
@@ -74,7 +65,7 @@ def test_create_with_no_bucket_name(
 ):
     create_hub = {"HubArn": f"arn:aws:sagemaker:us-east-1:123456789123:hub/{hub_name}"}
     sagemaker_session.create_hub = Mock(return_value=create_hub)
-    hub = CuratedHub(hub_name=hub_name, region=REGION, sagemaker_session=sagemaker_session)
+    hub = CuratedHub(hub_name=hub_name, sagemaker_session=sagemaker_session)
     request = {
         "hub_name": hub_name,
         "hub_description": hub_description,
@@ -119,7 +110,7 @@ def test_create_with_bucket_name(
 ):
     create_hub = {"HubArn": f"arn:aws:sagemaker:us-east-1:123456789123:hub/{hub_name}"}
     sagemaker_session.create_hub = Mock(return_value=create_hub)
-    hub = CuratedHub(hub_name=hub_name, region=REGION, sagemaker_session=sagemaker_session)
+    hub = CuratedHub(hub_name=hub_name, sagemaker_session=sagemaker_session)
     request = {
         "hub_name": hub_name,
         "hub_description": hub_description,

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -252,14 +252,14 @@ def test_jumpstart_cache_handles_boto3_issues(mock_boto3_client):
 @patch("boto3.client")
 def test_jumpstart_cache_gets_cleared_when_params_are_set(mock_boto3_client):
     cache = JumpStartModelsCache(
-        s3_bucket_name="some_bucket", region="some_region", manifest_file_s3_key="some_key"
+        s3_bucket_name="some_bucket", region="us-west-2", manifest_file_s3_key="some_key"
     )
 
     cache.clear = MagicMock()
     cache.set_s3_bucket_name("some_bucket")
     cache.clear.assert_not_called()
     cache.clear.reset_mock()
-    cache.set_region("some_region")
+    cache.set_region("us-west-2")
     cache.clear.assert_not_called()
     cache.clear.reset_mock()
     cache.set_manifest_file_s3_key("some_key")
@@ -270,7 +270,7 @@ def test_jumpstart_cache_gets_cleared_when_params_are_set(mock_boto3_client):
     cache.set_s3_bucket_name("some_bucket1")
     cache.clear.assert_called_once()
     cache.clear.reset_mock()
-    cache.set_region("some_region1")
+    cache.set_region("us-east-1")
     cache.clear.assert_called_once()
     cache.clear.reset_mock()
     cache.set_manifest_file_s3_key("some_key1")

tests/unit/sagemaker/jumpstart/utils.py

@@ -27,6 +27,7 @@
     JumpStartModelSpecs,
     JumpStartS3FileType,
     JumpStartModelHeader,
+    HubType,
     HubContentType,
 )
 
@@ -211,7 +212,7 @@ def patched_retrieval_function(
         )
 
     # TODO: Implement
-    if datatype == HubContentType.HUB:
+    if datatype == HubType.HUB:
         return None
 
     raise ValueError(f"Bad value for filetype: {datatype}")