feat: add hub and hubcontent support in retrieval function for jumpstart model cache (aws#4438)

Author: bencrabtree

src/sagemaker/jumpstart/cache.py

@@ -57,6 +57,7 @@
     DescribeHubContentsResponse,
     HubType,
     HubContentType,
+    HubDataType,
 )
 from sagemaker.jumpstart.curated_hub import utils as hub_utils
 from sagemaker.jumpstart.enums import JumpStartModelType
@@ -420,6 +421,7 @@ def _retrieval_function(
         """
 
         data_type, id_info = key.data_type, key.id_info
+
         if data_type in {
             JumpStartS3FileType.OPEN_WEIGHT_MANIFEST,
             JumpStartS3FileType.PROPRIETARY_MANIFEST,

src/sagemaker/jumpstart/constants.py

@@ -172,7 +172,8 @@
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
 JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY = "proprietary-sdk-manifest.json"
 
-HUB_CONTENT_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub-content/(.*?)/(.*?)/(.*?)/(.*?)$"
+# works cross-partition
+HUB_MODEL_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub_content/(.*?)/Model/(.*?)/(.*?)$"
 HUB_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub/(.*?)$"
 
 INFERENCE_ENTRY_POINT_SCRIPT_NAME = "inference.py"

src/sagemaker/jumpstart/types.py

@@ -972,6 +972,25 @@ def from_hub_content_doc(self, hub_content_doc: Dict[str, Any]) -> None:
         """
         # TODO: Implement
 
+    def to_json(self) -> Dict[str, Any]:
+        """Returns json representation of JumpStartModelSpecs object."""
+        json_obj = {}
+        for att in self.__slots__:
+            if hasattr(self, att):
+                cur_val = getattr(self, att)
+                if issubclass(type(cur_val), JumpStartDataHolderType):
+                    json_obj[att] = cur_val.to_json()
+                elif isinstance(cur_val, list):
+                    json_obj[att] = []
+                    for obj in cur_val:
+                        if issubclass(type(obj), JumpStartDataHolderType):
+                            json_obj[att].append(obj.to_json())
+                        else:
+                            json_obj[att].append(obj)
+                else:
+                    json_obj[att] = cur_val
+        return json_obj
+
     def supports_prepacked_inference(self) -> bool:
         """Returns True if the model has a prepacked inference artifact."""
         return getattr(self, "hosting_prepacked_artifact_key", None) is not None

src/sagemaker/jumpstart/utils.py

@@ -16,6 +16,7 @@
 import os
 import re
 from typing import Any, Dict, List, Set, Optional, Tuple, Union
+import re
 from urllib.parse import urlparse
 import boto3
 from packaging.version import Version
@@ -866,7 +867,6 @@ def generate_studio_spec_file_prefix(model_id: str, model_version: str) -> str:
     """Returns the Studio Spec file prefix given a model ID and version."""
     return f"studio_models/{model_id}/studio_specs_v{model_version}.json"
 
-
 def extract_info_from_hub_content_arn(
     arn: str,
 ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -1206,6 +1206,35 @@ def test_mime_type_enum_from_str():
             assert MIMEType.from_suffixed_type(mime_type_with_suffix) == mime_type
 
 
+def test_extract_info_from_hub_content_arn():
+    model_arn = (
+        "arn:aws:sagemaker:us-west-2:000000000000:hub_content/MockHub/Model/my-mock-model/1.0.2"
+    )
+    assert utils.extract_info_from_hub_content_arn(model_arn) == (
+        "MockHub",
+        "us-west-2",
+        "my-mock-model",
+        "1.0.2",
+    )
+
+    hub_arn = "arn:aws:sagemaker:us-west-2:000000000000:hub/MockHub"
+    assert utils.extract_info_from_hub_content_arn(hub_arn) == ("MockHub", "us-west-2", None, None)
+
+    invalid_arn = "arn:aws:sagemaker:us-west-2:000000000000:endpoint/my-endpoint-123"
+    assert utils.extract_info_from_hub_content_arn(invalid_arn) == (None, None, None, None)
+
+    invalid_arn = "nonsense-string"
+    assert utils.extract_info_from_hub_content_arn(invalid_arn) == (None, None, None, None)
+
+    invalid_arn = ""
+    assert utils.extract_info_from_hub_content_arn(invalid_arn) == (None, None, None, None)
+
+    invalid_arn = (
+        "arn:aws:sagemaker:us-west-2:000000000000:hub-content/MyHub/Notebook/my-notebook/1.0.0"
+    )
+    assert utils.extract_info_from_hub_content_arn(invalid_arn) == (None, None, None, None)
+
+
 class TestIsValidModelId(TestCase):
     @patch("sagemaker.jumpstart.utils.accessors.JumpStartModelsAccessor._get_manifest")
     @patch("sagemaker.jumpstart.utils.accessors.JumpStartModelsAccessor.get_model_specs")

tests/unit/sagemaker/jumpstart/utils.py

@@ -22,6 +22,7 @@
     JUMPSTART_REGION_NAME_SET,
 )
 from sagemaker.jumpstart.types import (
+    HubDataType,
     JumpStartCachedContentKey,
     JumpStartCachedContentValue,
     JumpStartModelSpecs,