add hub and hubcontent support in retrieval function for jumpstart model cache

Author: bencrabtree

src/sagemaker/jumpstart/cache.py

@@ -465,6 +465,26 @@ def get_specs(self, model_id: str, semantic_version_str: str) -> JumpStartModelS
                 )
             )
         return specs.formatted_content
+    
+    def get_hub_model(self, hub_model_arn: str) -> JumpStartModelSpecs:
+        """Return JumpStart-compatible specs for a given Hub model
+        
+        Args:
+            hub_model_arn (str): Arn for the Hub model to get specs for
+        """
+
+        specs, _ = self._content_cache.get(JumpStartCachedContentKey(HubDataType.MODEL, hub_model_arn))
+        return specs.formatted_content
+    
+    def get_hub(self, hub_arn: str) -> Dict[str, Any]:
+        """Return descriptive info for a given Hub
+        
+        Args:
+            hub_arn (str): Arn for the Hub to get info for
+        """
+        
+        manifest, _ = self._content_cache.get(JumpStartCachedContentKey(HubDataType.HUB, hub_arn))
+        return manifest.formatted_content
 
     def get_hub_model(self, hub_model_arn: str) -> JumpStartModelSpecs:
         """Return JumpStart-compatible specs for a given Hub model

src/sagemaker/jumpstart/curated_hub/constants.py

@@ -0,0 +1,15 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from botocore.config import Config
+
+DEFAULT_CLIENT_CONFIG = Config(retries={"max_attempts": 10, "mode": "standard"})

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -18,6 +18,8 @@
 
 from sagemaker.session import Session
 
+from sagemaker.jumpstart.curated_hub.constants import DEFAULT_CLIENT_CONFIG
+
 
 class CuratedHub:
     """Class for creating and managing a curated JumpStart hub"""

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -860,12 +860,16 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
 def test_jumpstart_cache_get_hub_model():
     cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
 
-    model_arn = "arn:aws:sagemaker:us-west-2:000000000000:hub-content/HubName/Model/huggingface-mock-model-123/1.2.3"
+    model_arn = (
+        "arn:aws:sagemaker:us-west-2:000000000000:hub-content/HubName/Model/huggingface-mock-model-123/1.2.3"
+    )
     assert get_spec_from_base_spec(
         model_id="huggingface-mock-model-123", version="1.2.3"
     ) == cache.get_hub_model(hub_model_arn=model_arn)
 
-    model_arn = "arn:aws:sagemaker:us-west-2:000000000000:hub-content/HubName/Model/pytorch-mock-model-123/*"
-    assert get_spec_from_base_spec(
-        model_id="pytorch-mock-model-123", version="*"
-    ) == cache.get_hub_model(hub_model_arn=model_arn)
+    model_arn = (
+        "arn:aws:sagemaker:us-west-2:000000000000:hub-content/HubName/Model/pytorch-mock-model-123/*"
+    )
+    assert get_spec_from_base_spec(model_id="pytorch-mock-model-123", version="*") == cache.get_hub_model(
+        hub_model_arn=model_arn
+    )