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
@@ -428,6 +429,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

@@ -15,6 +15,7 @@
 import logging
 import os
 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,3 +867,25 @@ def get_jumpstart_model_id_version_from_resource_arn(
 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]]:
+    """Extracts hub_name, content_name, and content_version from a HubContentArn"""
+
+    match = re.match(constants.HUB_MODEL_ARN_REGEX, arn)
+    if match:
+        hub_name = match.group(4)
+        hub_region = match.group(2)
+        content_name = match.group(5)
+        content_version = match.group(6)
+
+        return hub_name, hub_region, content_name, content_version
+
+    match = re.match(constants.HUB_ARN_REGEX, arn)
+    if match:
+        hub_name = match.group(4)
+        hub_region = match.group(2)
+        return hub_name, hub_region, None, None
+
+    return None, None, None, None

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -1214,6 +1214,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,
@@ -218,12 +219,10 @@ def patched_retrieval_function(
     datatype, id_info = key.data_type, key.id_info
     if datatype == JumpStartS3FileType.OPEN_WEIGHT_MANIFEST:
 
-        return JumpStartCachedContentValue(
-            formatted_content=get_formatted_manifest(BASE_MANIFEST)
-        )
+        return JumpStartCachedContentValue(formatted_content=get_formatted_manifest(BASE_MANIFEST))
 
-    if datatype == JumpStartCachedContentValue.OPEN_WEIGHT_SPECS:
-        _, model_id, specs_version = s3_key.split("/")
+    if datatype == JumpStartS3FileType.OPEN_WEIGHT_SPECS:
+        _, model_id, specs_version = id_info.split("/")
         version = specs_version.replace("specs_v", "").replace(".json", "")
         return JumpStartCachedContentValue(
             formatted_content=get_spec_from_base_spec(model_id=model_id, version=version)
@@ -245,7 +244,7 @@ def patched_retrieval_function(
         )
 
     if datatype == JumpStartS3FileType.PROPRIETARY_SPECS:
-        _, model_id, specs_version = s3_key.split("/")
+        _, model_id, specs_version = id_info.split("/")
         version = specs_version.replace("proprietary_specs_", "").replace(".json", "")
         return JumpStartCachedContentValue(
             formatted_content=get_spec_from_base_spec(