feat: Adding tagging to service changes

Author: chrstfu

src/sagemaker/jumpstart/enums.py

@@ -82,6 +82,11 @@ 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 +104,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

@@ -41,7 +41,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,
@@ -51,6 +51,7 @@
 )
 from sagemaker.jumpstart.utils import (
     add_hub_content_arn_tags,
+    add_bedrock_store_tags,
     add_jumpstart_model_info_tags,
     get_default_jumpstart_session_with_user_agent_suffix,
     get_neo_content_bucket,
@@ -488,6 +489,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"):
+      if HubContentCapability.BEDROCK_CONSOLE in kwargs.specs.capabilities:
+        kwargs.tags = add_bedrock_store_tags(kwargs.tags, compatibility="compatible")
+
     return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -1200,6 +1200,7 @@ class JumpStartMetadataBaseFields(JumpStartDataHolderType):
         "url",
         "version",
         "min_sdk_version",
+        "capabilities",
         "incremental_training_supported",
         "hosting_ecr_specs",
         "hosting_ecr_uri",
@@ -1287,6 +1288,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             json_obj.get("incremental_training_supported", False)
         )
         if self._is_hub_content:
+            self.capabilities: Optional[str] = json_obj.get("capabilities")
             self.hosting_ecr_uri: Optional[str] = json_obj.get("hosting_ecr_uri")
             self._non_serializable_slots.append("hosting_ecr_specs")
         else:

src/sagemaker/jumpstart/utils.py

@@ -451,6 +451,20 @@ 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,