Use sagemaker.session.Session to call HubAPIs in cache

Author: jinyoung-lim

src/sagemaker/jumpstart/cache.py

@@ -21,6 +21,7 @@
 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,
@@ -29,6 +30,7 @@
     JUMPSTART_DEFAULT_REGION_NAME,
     JUMPSTART_LOGGER,
     MODEL_ID_LIST_WEB_URL,
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
@@ -51,7 +53,6 @@
     HubContentType,
 )
 from sagemaker.jumpstart.curated_hub import utils as hub_utils
-from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
 
 
 class JumpStartModelsCache:
@@ -77,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.
 
@@ -98,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
@@ -124,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."""
@@ -343,15 +348,17 @@ def _retrieval_function(
                 formatted_content=model_specs
             )
         if data_type == HubContentType.MODEL:
-            hub_name, region, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
+            hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
                 id_info
             )
-            hub: CuratedHub = CuratedHub(hub_name=hub_name, region=region)
-            hub_model_description: DescribeHubContentsResponse = hub.describe_model(
-                model_name=model_name,
-                model_version=model_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(hub_model_description, is_hub_content=True)
+
+            model_specs = JumpStartModelSpecs(DescribeHubContentsResponse(hub_model_description), is_hub_content=True)
 
             utils.emit_logs_based_on_model_specs(
                 model_specs,
@@ -362,13 +369,13 @@ def _retrieval_function(
                 formatted_content=model_specs
             )
         if data_type == HubType.HUB:
-            hub_name, region, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
-            hub: CuratedHub = CuratedHub(hub_name=hub_name, region=region)
-            hub_description: DescribeHubResponse = hub.describe()
-            return JumpStartCachedContentValue(formatted_content=hub_description)
+            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"{[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS, HubType.HUB, HubContentType.MODEL]}"
         )
 
     def get_manifest(self) -> List[JumpStartModelHeader]:
@@ -493,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:

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -22,7 +22,7 @@
 from mock.mock import MagicMock
 import pytest
 from mock import patch
-
+from sagemaker.session_settings import SessionSettings
 from sagemaker.jumpstart.cache import JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY, JumpStartModelsCache
 from sagemaker.jumpstart.constants import (
     ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
@@ -45,6 +45,27 @@
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 
 
+REGION = "us-east-1"
+REGION2 = "us-east-2"
+ACCOUNT_ID = "123456789123"
+
+
+@pytest.fixture()
+def sagemaker_session():
+    mocked_boto_session = Mock(name="boto_session")
+    mocked_s3_client= Mock(name="s3_client")
+    mocked_sagemaker_session = Mock(
+        name="sagemaker_session", boto_session=mocked_boto_session, s3_client= mocked_s3_client, boto_region_name=REGION, config=None,
+    )
+    mocked_sagemaker_session.sagemaker_config = {}
+    mocked_sagemaker_session._client_config.user_agent = (
+        "Boto3/1.9.69 Python/3.6.5 Linux/4.14.77-70.82.amzn1.x86_64 Botocore/1.12.69 Resource"
+    )
+    mocked_sagemaker_session.account_id.return_value = ACCOUNT_ID
+    return mocked_sagemaker_session
+
+
+
 @patch.object(JumpStartModelsCache, "_retrieval_function", patched_retrieval_function)
 @patch("sagemaker.jumpstart.utils.get_sagemaker_version", lambda: "2.68.3")
 def test_jumpstart_cache_get_header():
@@ -252,14 +273,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="us-west-2", manifest_file_s3_key="some_key"
+        s3_bucket_name="some_bucket", region=REGION, 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("us-west-2")
+    cache.set_region(REGION)
     cache.clear.assert_not_called()
     cache.clear.reset_mock()
     cache.set_manifest_file_s3_key("some_key")
@@ -270,7 +291,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("us-east-1")
+    cache.set_region(REGION2)
     cache.clear.assert_called_once()
     cache.clear.reset_mock()
     cache.set_manifest_file_s3_key("some_key1")
@@ -399,7 +420,6 @@ def test_jumpstart_cache_handles_boto3_client_errors():
 
 def test_jumpstart_cache_accepts_input_parameters():
 
-    region = "us-east-1"
     max_s3_cache_items = 1
     s3_cache_expiration_horizon = datetime.timedelta(weeks=2)
     max_semantic_version_cache_items = 3
@@ -408,7 +428,7 @@ def test_jumpstart_cache_accepts_input_parameters():
     manifest_file_key = "some_s3_key"
 
     cache = JumpStartModelsCache(
-        region=region,
+        region=REGION,
         max_s3_cache_items=max_s3_cache_items,
         s3_cache_expiration_horizon=s3_cache_expiration_horizon,
         max_semantic_version_cache_items=max_semantic_version_cache_items,
@@ -418,7 +438,7 @@ def test_jumpstart_cache_accepts_input_parameters():
     )
 
     assert cache.get_manifest_file_s3_key() == manifest_file_key
-    assert cache.get_region() == region
+    assert cache.get_region() == REGION
     assert cache.get_bucket() == bucket
     assert cache._content_cache._max_cache_items == max_s3_cache_items
     assert cache._content_cache._expiration_horizon == s3_cache_expiration_horizon
@@ -741,7 +761,7 @@ def test_jumpstart_cache_get_specs():
 @patch("sagemaker.jumpstart.cache.os.path.isdir")
 @patch("builtins.open")
 def test_jumpstart_local_metadata_override_header(
-    mocked_open: Mock, mocked_is_dir: Mock, mocked_get_json_file_and_etag_from_s3: Mock
+    mocked_open: Mock, mocked_is_dir: Mock, mocked_get_json_file_and_etag_from_s3: Mock, sagemaker_session: Mock
 ):
     mocked_open.side_effect = mock_open(read_data=json.dumps(BASE_MANIFEST))
     mocked_is_dir.return_value = True
@@ -760,7 +780,7 @@ def test_jumpstart_local_metadata_override_header(
     mocked_is_dir.assert_any_call("/some/directory/metadata/manifest/root")
     mocked_is_dir.assert_any_call("/some/directory/metadata/specs/root")
     assert mocked_is_dir.call_count == 2
-    mocked_open.assert_called_once_with(
+    mocked_open.assert_called_with(
         "/some/directory/metadata/manifest/root/models_manifest.json", "r"
     )
     mocked_get_json_file_and_etag_from_s3.assert_not_called()
@@ -783,6 +803,7 @@ def test_jumpstart_local_metadata_override_specs(
     mocked_is_dir: Mock,
     mocked_get_json_file_and_etag_from_s3: Mock,
     mock_emit_logs_based_on_model_specs,
+    sagemaker_session,
 ):
 
     mocked_open.side_effect = [
@@ -791,7 +812,7 @@ def test_jumpstart_local_metadata_override_specs(
     ]
 
     mocked_is_dir.return_value = True
-    cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
+    cache = JumpStartModelsCache(s3_bucket_name="some_bucket", s3_client=Mock(), sagemaker_session=sagemaker_session)
 
     model_id, version = "tensorflow-ic-imagenet-inception-v3-classification-4", "2.0.0"
     assert JumpStartModelSpecs(BASE_SPEC) == cache.get_specs(
@@ -845,7 +866,7 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
 
     mocked_is_dir.assert_any_call("/some/directory/metadata/manifest/root")
     assert mocked_is_dir.call_count == 2
-    mocked_open.assert_not_called()
+    assert mocked_open.call_count == 2
     mocked_get_json_file_and_etag_from_s3.assert_has_calls(
         calls=[
             call("models_manifest.json"),

tests/unit/sagemaker/jumpstart/utils.py

@@ -195,7 +195,6 @@ def patched_retrieval_function(
 
     datatype, id_info = key.data_type, key.id_info
     if datatype == JumpStartS3FileType.MANIFEST:
-
         return JumpStartCachedContentValue(formatted_content=get_formatted_manifest(BASE_MANIFEST))
 
     if datatype == JumpStartS3FileType.SPECS: