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 @@
     JumpStartVersionedModelId,
     HubType,
     HubContentType,
+    HubDataType,
 )
 from sagemaker.jumpstart.curated_hub import utils as hub_utils
 from sagemaker.jumpstart.curated_hub.interfaces import (

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/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

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,