feat: Adding Bedrock Store model support for HubService (#1539)

* feat: Adding BRS support

Author: chrstfu

src/sagemaker/jumpstart/enums.py

@@ -82,6 +82,12 @@ class VariableTypes(str, Enum):
     BOOL = "bool"
 
 
+class HubContentCapability(str, Enum):
+    """Enum class for HubContent capabilities."""
+
+    BEDROCK_CONSOLE = "BEDROCK_CONSOLE"
+
+
 class JumpStartTag(str, Enum):
     """Enum class for tag keys to apply to JumpStart models."""
 
@@ -99,6 +105,8 @@ class JumpStartTag(str, Enum):
 
     HUB_CONTENT_ARN = "sagemaker-sdk:hub-content-arn"
 
+    BEDROCK = "sagemaker-sdk:bedrock"
+
 
 class SerializerType(str, Enum):
     """Enum class for serializers associated with JumpStart models."""

src/sagemaker/jumpstart/factory/model.py

@@ -42,7 +42,7 @@
 from sagemaker.model_metrics import ModelMetrics
 from sagemaker.metadata_properties import MetadataProperties
 from sagemaker.drift_check_baselines import DriftCheckBaselines
-from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartModelType
+from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartModelType, HubContentCapability
 from sagemaker.jumpstart.types import (
     HubContentType,
     JumpStartModelDeployKwargs,
@@ -53,6 +53,7 @@
 from sagemaker.jumpstart.utils import (
     add_hub_content_arn_tags,
     add_jumpstart_model_info_tags,
+    add_bedrock_store_tags,
     get_default_jumpstart_session_with_user_agent_suffix,
     get_top_ranked_config_name,
     update_dict_if_key_not_present,
@@ -495,6 +496,10 @@ def _add_tags_to_kwargs(kwargs: JumpStartModelDeployKwargs) -> Dict[str, Any]:
             )
         kwargs.tags = add_hub_content_arn_tags(kwargs.tags, hub_content_arn=hub_content_arn)
 
+    if hasattr(kwargs.specs, "capabilities") and kwargs.specs.capabilities is not None:
+        if HubContentCapability.BEDROCK_CONSOLE in kwargs.specs.capabilities:
+            kwargs.tags = add_bedrock_store_tags(kwargs.tags, compatibility="compatible")
+
     return kwargs
 
 

src/sagemaker/jumpstart/hub/parser_utils.py

@@ -18,10 +18,17 @@
 from typing import Any, Dict, List, Optional
 
 
+<<<<<<< HEAD
 def camel_to_snake(camel_case_string: str) -> str:
     """Converts camelCase to snake_case_string using a regex.
 
     This regex cannot handle whitespace ("camelString TwoWords")
+=======
+def pascal_to_snake(camel_case_string: str) -> str:
+    """Converts PascalCase to snake_case_string using a regex.
+
+    This regex cannot handle whitespace ("PascalString TwoWords")
+>>>>>>> ff3eae05 (feat: Adding Bedrock Store model support for HubService (#1539))
     """
     return re.sub(r"(?<!^)(?=[A-Z])", "_", camel_case_string).lower()
 

src/sagemaker/jumpstart/hub/parsers.py

@@ -27,7 +27,7 @@
     HubModelDocument,
 )
 from sagemaker.jumpstart.hub.parser_utils import (
-    camel_to_snake,
+    pascal_to_snake,
     snake_to_upper_camel,
     walk_and_apply_json,
 )
@@ -86,7 +86,7 @@ def get_model_spec_arg_keys(
         arg_keys = []
 
     if naming_convention == NamingConventionType.SNAKE_CASE:
-        arg_keys = [camel_to_snake(key) for key in arg_keys]
+        arg_keys = [pascal_to_snake(key) for key in arg_keys]
     elif naming_convention == NamingConventionType.UPPER_CAMEL_CASE:
         return arg_keys
     else:
@@ -207,7 +207,7 @@ def make_model_specs_from_describe_hub_content_response(
     default_payloads: Dict[str, Any] = {}
     if hub_model_document.default_payloads is not None:
         for alias, payload in hub_model_document.default_payloads.items():
-            default_payloads[alias] = walk_and_apply_json(payload.to_json(), camel_to_snake)
+            default_payloads[alias] = walk_and_apply_json(payload.to_json(), pascal_to_snake)
         specs["default_payloads"] = default_payloads
     specs["gated_bucket"] = hub_model_document.gated_bucket
     specs["inference_volume_size"] = hub_model_document.inference_volume_size

src/sagemaker/jumpstart/hub/utils.py

@@ -219,17 +219,31 @@ def get_hub_model_version(
     except Exception as ex:
         raise Exception(f"Failed calling list_hub_content_versions: {str(ex)}")
 
+<<<<<<< HEAD
+=======
+    marketplace_hub_content_version = _get_hub_model_version_for_marketplace_version(
+        hub_content_summaries, hub_model_version
+    )
+
+>>>>>>> ff3eae05 (feat: Adding Bedrock Store model support for HubService (#1539))
     try:
         return _get_hub_model_version_for_open_weight_version(
             hub_content_summaries, hub_model_version
         )
+<<<<<<< HEAD
     except KeyError:
         marketplace_hub_content_version = _get_hub_model_version_for_marketplace_version(
             hub_content_summaries, hub_model_version
         )
         if marketplace_hub_content_version:
             return marketplace_hub_content_version
         raise
+=======
+    except KeyError as e:
+        if marketplace_hub_content_version:
+            return marketplace_hub_content_version
+        raise e
+>>>>>>> ff3eae05 (feat: Adding Bedrock Store model support for HubService (#1539))
 
 
 def _get_hub_model_version_for_open_weight_version(

src/sagemaker/jumpstart/types.py

@@ -40,7 +40,7 @@
 from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
 from sagemaker.enums import EndpointType
 from sagemaker.jumpstart.hub.parser_utils import (
-    camel_to_snake,
+    pascal_to_snake,
     walk_and_apply_json,
 )
 from sagemaker.model_life_cycle import ModelLifeCycle
@@ -241,7 +241,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             return
 
         if self._is_hub_content:
-            json_obj = walk_and_apply_json(json_obj, camel_to_snake)
+            json_obj = walk_and_apply_json(json_obj, pascal_to_snake)
 
         self.framework = json_obj.get("framework")
         self.framework_version = json_obj.get("framework_version")
@@ -295,7 +295,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         """
 
         if self._is_hub_content:
-            json_obj = walk_and_apply_json(json_obj, camel_to_snake)
+            json_obj = walk_and_apply_json(json_obj, pascal_to_snake)
         self.name = json_obj["name"]
         self.type = json_obj["type"]
         self.default = json_obj["default"]
@@ -363,7 +363,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         Args:
             json_obj (Dict[str, Any]): Dictionary representation of environment variable.
         """
-        json_obj = walk_and_apply_json(json_obj, camel_to_snake)
+        json_obj = walk_and_apply_json(json_obj, pascal_to_snake)
         self.name = json_obj["name"]
         self.type = json_obj["type"]
         self.default = json_obj["default"]
@@ -413,7 +413,7 @@ def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None:
             return
 
         if self._is_hub_content:
-            json_obj = walk_and_apply_json(json_obj, camel_to_snake)
+            json_obj = walk_and_apply_json(json_obj, pascal_to_snake)
         self.default_content_type = json_obj["default_content_type"]
         self.supported_content_types = json_obj["supported_content_types"]
         self.default_accept_type = json_obj["default_accept_type"]
@@ -467,7 +467,7 @@ def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None:
             return
 
         if self._is_hub_content:
-            json_obj = walk_and_apply_json(json_obj, camel_to_snake)
+            json_obj = walk_and_apply_json(json_obj, pascal_to_snake)
         self.raw_payload = json_obj
         self.content_type = json_obj["content_type"]
         self.body = json_obj.get("body")
@@ -540,7 +540,7 @@ def from_describe_hub_content_response(self, response: Optional[Dict[str, Any]])
         if response is None:
             return
 
-        response = walk_and_apply_json(response, camel_to_snake)
+        response = walk_and_apply_json(response, pascal_to_snake)
         self.aliases: Optional[dict] = response.get("aliases")
         self.regional_aliases = None
         self.variants: Optional[dict] = response.get("variants")
@@ -1180,7 +1180,7 @@ def __init__(self, spec: Optional[Dict[str, Any]], is_hub_content=False):
             spec (Dict[str, Any]): Dictionary representation of training config ranking.
         """
         if is_hub_content:
-            spec = walk_and_apply_json(spec, camel_to_snake)
+            spec = walk_and_apply_json(spec, pascal_to_snake)
         self.from_json(spec)
 
     def from_json(self, json_obj: Dict[str, Any]) -> None:
@@ -1286,7 +1286,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             json_obj (Dict[str, Any]): Dictionary representation of spec.
         """
         if self._is_hub_content:
-            json_obj = walk_and_apply_json(json_obj, camel_to_snake)
+            json_obj = walk_and_apply_json(json_obj, pascal_to_snake)
         self.model_id: str = json_obj.get("model_id")
         self.url: str = json_obj.get("url")
         self.version: str = json_obj.get("version")
@@ -1515,7 +1515,7 @@ def __init__(
             ValueError: If the component field is invalid.
         """
         if is_hub_content:
-            component = walk_and_apply_json(component, camel_to_snake)
+            component = walk_and_apply_json(component, pascal_to_snake)
         self.component_name = component_name
         super().__init__(component, is_hub_content)
         self.from_json(component)
@@ -1568,8 +1568,8 @@ def __init__(
                 The list of components that are used to construct the resolved config.
         """
         if is_hub_content:
-            config = walk_and_apply_json(config, camel_to_snake)
-            base_fields = walk_and_apply_json(base_fields, camel_to_snake)
+            config = walk_and_apply_json(config, pascal_to_snake)
+            base_fields = walk_and_apply_json(base_fields, pascal_to_snake)
         self.base_fields = base_fields
         self.config_components: Dict[str, JumpStartConfigComponent] = config_components
         self.benchmark_metrics: Dict[str, List[JumpStartBenchmarkStat]] = (
@@ -1735,7 +1735,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         """
         super().from_json(json_obj)
         if self._is_hub_content:
-            json_obj = walk_and_apply_json(json_obj, camel_to_snake)
+            json_obj = walk_and_apply_json(json_obj, pascal_to_snake)
         self.inference_config_components: Optional[Dict[str, JumpStartConfigComponent]] = (
             {
                 component_name: JumpStartConfigComponent(component_name, component)

src/sagemaker/jumpstart/utils.py

@@ -35,7 +35,7 @@
 
 from sagemaker.jumpstart import constants, enums
 from sagemaker.jumpstart import accessors
-from sagemaker.jumpstart.hub.parser_utils import camel_to_snake, snake_to_upper_camel
+from sagemaker.jumpstart.hub.parser_utils import pascal_to_snake, snake_to_upper_camel
 from sagemaker.s3 import parse_s3_url
 from sagemaker.jumpstart.exceptions import (
     DeprecatedJumpStartModelError,
@@ -455,6 +455,21 @@ def add_hub_content_arn_tags(
     return tags
 
 
+def add_bedrock_store_tags(
+    tags: Optional[List[TagsDict]],
+    compatibility: str,
+) -> Optional[List[TagsDict]]:
+    """Adds custom Hub arn tag to JumpStart related resources."""
+
+    tags = add_single_jumpstart_tag(
+        compatibility,
+        enums.JumpStartTag.BEDROCK,
+        tags,
+        is_uri=False,
+    )
+    return tags
+
+
 def add_jumpstart_uri_tags(
     tags: Optional[List[TagsDict]] = None,
     inference_model_uri: Optional[Union[str, dict]] = None,
@@ -1163,7 +1178,7 @@ def get_jumpstart_configs(
         return (
             {
                 config_name: metadata_configs.configs[
-                    camel_to_snake(snake_to_upper_camel(config_name))
+                    pascal_to_snake(snake_to_upper_camel(config_name))
                 ]
                 for config_name in config_names
             }