fix: make ModelBuilder.deploy(tags=...) a named parameter that reaches every deploy path

sagemaker-core/src/sagemaker/core/common_utils.py

@@ -1688,12 +1688,53 @@ def create_paginator_config(max_items: int = None, page_size: int = None) -> Dic
     }
 
 
+def _normalize_tag(tag: Any) -> TagsDict:
+    """Return one tag as a boto-style ``{"Key": ..., "Value": ...}`` dict.
+
+    Accepts the boto form, the ``{"key": ..., "value": ...}`` form that the
+    ``sagemaker.core.resources`` classes take, and a ``Tag`` shape object (anything
+    with ``key`` and ``value`` attributes). Any other input raises ``TypeError`` so a
+    malformed tag fails at the call site instead of as a ``KeyError`` deep in a
+    deploy path.
+    """
+    if isinstance(tag, dict):
+        if "Key" in tag:
+            return {"Key": tag["Key"], "Value": tag.get("Value")}
+        if "key" in tag:
+            return {"Key": tag["key"], "Value": tag.get("value")}
+        raise TypeError(
+            "A tag dict must have a 'Key' (or 'key') entry, "
+            f"got {sorted(tag)}. Pass a {{key: value}} mapping to set many tags at once."
+        )
+    if hasattr(tag, "key") and hasattr(tag, "value"):
+        return {"Key": tag.key, "Value": tag.value}
+    raise TypeError(
+        "A tag must be a {'Key': ..., 'Value': ...} dict, a {'key': ..., 'value': ...} dict, "
+        f"or a Tag object, got {type(tag).__name__}."
+    )
+
+
 def format_tags(tags: Tags) -> List[TagsDict]:
-    """Process tags to turn them into the expected format for Sagemaker."""
+    """Process tags to turn them into the expected format for Sagemaker.
+
+    Every accepted representation comes out as a list of boto-style
+    ``{"Key": ..., "Value": ...}`` dicts:
+
+    * a ``{key: value}`` mapping, one tag per entry;
+    * a single tag dict or ``Tag`` object;
+    * a list of tag dicts in either casing, ``Tag`` objects, or a mix.
+
+    ``None`` is returned unchanged.
+    """
+    if tags is None:
+        return tags
     if isinstance(tags, dict):
+        if "Key" in tags or "key" in tags:
+            return [_normalize_tag(tags)]
         return [{"Key": str(k), "Value": str(v)} for k, v in tags.items()]
-
-    return tags
+    if isinstance(tags, (list, tuple)):
+        return [_normalize_tag(tag) for tag in tags]
+    return [_normalize_tag(tags)]
 
 
 def _get_resolved_path(path):
@@ -2124,8 +2165,9 @@ def tag_exists(tag: TagsDict, curr_tags: Optional[Tags]) -> bool:
     if curr_tags is None:
         return False
 
-    for curr_tag in curr_tags:
-        if tag["Key"] == curr_tag["Key"]:
+    key = _normalize_tag(tag)["Key"]
+    for curr_tag in format_tags(curr_tags):
+        if key == curr_tag["Key"]:
             return True
 
     return False
@@ -2141,21 +2183,18 @@ def _validate_new_tags(new_tags: Optional[Tags], curr_tags: Optional[Tags]) -> O
     Returns:
         Optional[Tags]: The updated tags.
     """
-    if curr_tags is None:
-        return new_tags
+    if new_tags is None:
+        return format_tags(curr_tags) if curr_tags is not None else None
 
-    if curr_tags and isinstance(curr_tags, dict):
-        curr_tags = [curr_tags]
+    if curr_tags is None:
+        return format_tags(new_tags)
 
-    if isinstance(new_tags, dict):
-        if not tag_exists(new_tags, curr_tags):
-            curr_tags.append(new_tags)
-    elif isinstance(new_tags, list):
-        for new_tag in new_tags:
-            if not tag_exists(new_tag, curr_tags):
-                curr_tags.append(new_tag)
+    merged = format_tags(curr_tags)
+    for new_tag in format_tags(new_tags):
+        if not tag_exists(new_tag, merged):
+            merged.append(new_tag)
 
-    return curr_tags
+    return merged
 
 
 def remove_tag_with_key(key: str, tags: Optional[Tags]) -> Optional[Tags]:
@@ -2170,11 +2209,9 @@ def remove_tag_with_key(key: str, tags: Optional[Tags]) -> Optional[Tags]:
     """
     if tags is None:
         return tags
-    if isinstance(tags, dict):
-        tags = [tags]
 
     updated_tags = []
-    for tag in tags:
+    for tag in format_tags(tags):
         if tag["Key"] != key:
             updated_tags.append(tag)
 

sagemaker-serve/src/sagemaker/serve/model_builder.py

@@ -3935,8 +3935,9 @@ def _deploy_core_endpoint(self, **kwargs):
             self.endpoint_name = endpoint_name
 
         tags = kwargs.get("tags", getattr(self, "_tags", None))
-        if "tags" in kwargs:
-            self._tags = tags
+        if "tags" in kwargs and tags is not None:
+            self.add_tags(tags)
+            tags = self._tags
 
         kms_key = kwargs.get("kms_key", getattr(self, "kms_key", None))
         if "kms_key" in kwargs:
@@ -4564,14 +4565,11 @@ def build(
         # Tag the Model resource that build() creates. Without this the only way to
         # set _tags was deploy(), so every model built without an endpoint -- a
         # pipeline ModelStep, batch transform, evaluation, an inference pipeline --
-        # was created untagged with no public way to fix it (#6501).
-        # format_tags() expands the {key: value} mapping form; serialize() turns
-        # sagemaker.core.shapes.Tag objects into the boto-style {"Key", "Value"}
-        # dicts the CreateModel request needs, so both representations are accepted.
+        # was created untagged with no public way to fix it (#6501). add_tags()
+        # normalizes every accepted representation to the boto-style {"Key", "Value"}
+        # dicts the CreateModel request needs.
         if tags is not None:
-            from sagemaker.core.utils.utils import serialize
-
-            self.add_tags(serialize(format_tags(tags)))
+            self.add_tags(tags)
 
         if hasattr(self, "built_model") and self.built_model is not None:
             logger.warning(
@@ -6077,6 +6075,7 @@ def deploy(
         role: Optional[str] = None,
         model_name: Optional[str] = None,
         endpoint_config_name: Optional[str] = None,
+        tags: Optional[Tags] = None,
         **kwargs,
     ) -> Union[Endpoint, LocalEndpoint, Transformer]:
         """Deploy the built model to an ``Endpoint``.
@@ -6152,6 +6151,13 @@ def deploy(
             endpoint_config_name (str, optional): Recommendation deploy only. Name for the created
                 EndpointConfig; auto-generated if omitted. Ignored when deploying a normally-built
                 model.
+            tags (Optional[Tags]): Tags to attach to the endpoint and endpoint config this call
+                creates, merged with any added through ``add_tags()`` or ``build(tags=...)``.
+                Accepts a ``{key: value}`` mapping, boto-style ``{"Key": ..., "Value": ...}``
+                dicts, ``{"key": ..., "value": ...}`` dicts, or ``sagemaker.core.shapes.Tag``
+                objects. The same tags apply on every deploy path: a built model, a fine-tuned
+                model package, a Nova model, an inference component, and a recommendation.
+                (Default: None).
         Returns:
             Union[Endpoint, LocalEndpoint, Transformer]: A ``sagemaker.core.resources.Endpoint``
                 resource representing the deployed endpoint, a ``LocalEndpoint`` for local mode,
@@ -6177,6 +6183,12 @@ def deploy(
             )
         self._deployed = True
 
+        # Normalize the caller's tags once, in boto form, and merge them into the
+        # builder's accumulated tags. Every deploy path below reads self._tags, so a
+        # tag passed here reaches the endpoint whichever path the model takes.
+        if tags is not None:
+            self.add_tags(tags)
+
         # Recommendation-mode deploy: when a recommendation job is attached
         # (from a prior generate_deployment_recommendations call, or because the
         # builder was hydrated via ModelBuilder.from_recommendation_job), deploy
@@ -6223,7 +6235,7 @@ def deploy(
                 role=role,
                 instance_type=instance_type,
                 initial_instance_count=initial_instance_count,
-                tags=kwargs.get("tags"),
+                tags=format_tags(getattr(self, "_tags", None)),
                 wait=wait,
                 auto_approve=auto_approve,
             )
@@ -6436,6 +6448,17 @@ def deploy(
 
         raise ValueError("Deployment Options not supported")
 
+    def _core_endpoint_tags(self) -> List[Dict[str, str]]:
+        """Return the builder's accumulated tags in the form ``Endpoint.create`` takes.
+
+        ``add_tags`` stores tags in boto form (``{"Key": ..., "Value": ...}``); the
+        ``sagemaker.core.resources`` classes take ``{"key": ..., "value": ...}``.
+        """
+        return [
+            {"key": tag["Key"], "value": tag["Value"]}
+            for tag in format_tags(getattr(self, "_tags", None) or [])
+        ]
+
     def _deploy_model_customization(
         self,
         endpoint_name: str,
@@ -6506,7 +6529,9 @@ def _deploy_model_customization(
                 ],
             )
             endpoint = Endpoint.create(
-                endpoint_name=endpoint_name, endpoint_config_name=endpoint_name
+                endpoint_name=endpoint_name,
+                endpoint_config_name=endpoint_name,
+                tags=self._core_endpoint_tags() or None,
             )
             if kwargs.get("wait", True):
                 endpoint.wait_for_status("InService")
@@ -6593,18 +6618,13 @@ def _deploy_model_customization(
                 enable_network_isolation=enable_network_isolation,
             )
             logger.info("Endpoint core call starting")
-            # Apply tags accumulated via add_tags (e.g. the model-source reuse
-            # tag) to the endpoint so it is discoverable. Stored tags are in
-            # {"Key":..,"Value":..} form; normalize to the key/value form the
-            # core resource expects.
-            endpoint_tags = [
-                {"key": tag["Key"], "value": tag["Value"]}
-                for tag in format_tags(getattr(self, "_tags", None) or [])
-            ]
+            # Apply tags accumulated via add_tags (the caller's deploy(tags=...),
+            # build(tags=...), and the model-source reuse tag) to the endpoint so it
+            # is discoverable.
             endpoint = Endpoint.create(
                 endpoint_name=endpoint_name,
                 endpoint_config_name=endpoint_name,
-                tags=endpoint_tags or None,
+                tags=self._core_endpoint_tags() or None,
             )
             endpoint.wait_for_status("InService")
         else:
@@ -6828,15 +6848,13 @@ def _deploy_nova_model(
             if base_model is not None and base_model.recipe_name:
                 tags.append({"key": "sagemaker-sdk:recipe-name", "value": base_model.recipe_name})
 
-        # Merge tags accumulated via add_tags (e.g. the model-source reuse tag).
-        # Those are stored in {"Key": ..., "Value": ...} form, so normalize to the
-        # {"key": ..., "value": ...} form Endpoint.create expects and de-duplicate.
+        # Merge tags accumulated via add_tags (the caller's deploy(tags=...),
+        # build(tags=...), and the model-source reuse tag) and de-duplicate by key.
         existing_keys = {tag["key"] for tag in tags}
-        for tag in format_tags(getattr(self, "_tags", None) or []):
-            key = tag["Key"]
-            if key not in existing_keys:
-                tags.append({"key": key, "value": tag["Value"]})
-                existing_keys.add(key)
+        for tag in self._core_endpoint_tags():
+            if tag["key"] not in existing_keys:
+                tags.append(tag)
+                existing_keys.add(tag["key"])
 
         endpoint = Endpoint.create(
             endpoint_name=endpoint_name,