fix: propagate deploy tags on the model-customization endpoint paths

Tags passed to ModelBuilder.deploy() were accepted but never reached the
endpoints created on the model-customization paths, so they could not be used
for tag-based resource association.

deploy() branches into _deploy_model_customization() for fine-tuned and Nova
models and returns before reaching _deploy_core_endpoint(), which is the only
place that merges the deploy-time tags kwarg into self._tags. Those paths build
their endpoint tags from self._tags, so they saw only the tags ModelBuilder had
added itself. Merge the kwarg into the accumulator on entry so all three
endpoint-creation paths below pick it up.

The restricted-model-package path created its endpoint with no tags at all;
it now receives them like the other two.

Adds _tags_as_key_value_list() to normalize tags into the lowercase key/value
list the sagemaker.core.resources create() calls are typed for. self._tags can
hold either casing or a bare {key: value} mapping, and the previous inline
reads indexed tag["Key"] directly after format_tags(), which returns a list
unchanged whatever its casing -- so a lowercase list raised KeyError. That was
already reachable through build(tags=...) and add_tags().

Author: Rohan Gujarathi

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

@@ -118,6 +118,7 @@
 from sagemaker.core.enums import EndpointType
 from sagemaker.core.common_utils import (
     Tags,
+    TagsDict,
     ModelApprovalStatusEnum,
     _resolve_routing_config,
     format_tags,
@@ -180,6 +181,44 @@
 SAGEMAKER_OUTPUT_LOCATION = "sagemaker_s3_output"
 
 
+def _tags_as_key_value_list(tags: Optional[Tags]) -> List[TagsDict]:
+    """Normalize any accepted tag form into the ``{"key": ..., "value": ...}`` list form.
+
+    ``deploy()`` and ``add_tags()`` accept tags either as a list of dicts using
+    ``key``/``value`` or ``Key``/``Value``, or as a single ``{key: value}`` mapping, so
+    ``self._tags`` can hold any of those. The ``sagemaker.core.resources`` ``create()``
+    calls are typed ``List[Tag]`` whose fields are lowercase and which forbid extra
+    fields, so they reject the capitalized form and reject a bare mapping outright.
+
+    Note this differs from :func:`format_tags`, which emits the capitalized form the
+    legacy ``sagemaker_session.create_*`` APIs expect and passes a list through
+    unchanged whatever its casing.
+
+    Args:
+        tags: Tags in any accepted form, or None.
+
+    Returns:
+        The tags as a list of lowercase key/value dicts; empty when none were supplied.
+        Entries missing a key or value are skipped, since SageMaker rejects them.
+    """
+    if not tags:
+        return []
+    if isinstance(tags, dict):
+        return [{"key": k, "value": v} for k, v in tags.items()]
+
+    normalized = []
+    for tag in tags:
+        if isinstance(tag, dict):
+            key = tag.get("key", tag.get("Key"))
+            value = tag.get("value", tag.get("Value"))
+        else:
+            key = getattr(tag, "key", None)
+            value = getattr(tag, "value", None)
+        if key is not None and value is not None:
+            normalized.append({"key": key, "value": value})
+    return normalized
+
+
 @dataclass
 class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuilderUtils):
     """Unified interface for building and deploying machine learning models.
@@ -6185,6 +6224,12 @@ def _deploy_model_customization(
         from sagemaker.core.resources import InferenceComponent
         from sagemaker.core.resources import Tag as CoreTag
 
+        # Tags passed to deploy() arrive here in kwargs. Merge them into the tag
+        # accumulator so every endpoint-creation path below picks them up; the
+        # standard deploy path does the equivalent in _deploy_core_endpoint().
+        if kwargs.get("tags"):
+            self.add_tags(_tags_as_key_value_list(kwargs["tags"]))
+
         # An inference_config of ResourceRequirements requests an inference
         # component deployment; otherwise the model is placed directly on the
         # production variant.
@@ -6226,7 +6271,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=_tags_as_key_value_list(getattr(self, "_tags", None)) or None,
             )
             if kwargs.get("wait", True):
                 endpoint.wait_for_status("InService")
@@ -6266,10 +6313,7 @@ def _deploy_model_customization(
             # 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 [])
-            ]
+            endpoint_tags = _tags_as_key_value_list(getattr(self, "_tags", None))
             endpoint = Endpoint.create(
                 endpoint_name=endpoint_name,
                 endpoint_config_name=endpoint_name,
@@ -6541,11 +6585,10 @@ def _deploy_nova_model(
         # Those are stored in {"Key": ..., "Value": ...} form, so normalize to the
         # {"key": ..., "value": ...} form Endpoint.create expects and de-duplicate.
         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 _tags_as_key_value_list(getattr(self, "_tags", None)):
+            if tag["key"] not in existing_keys:
+                tags.append(tag)
+                existing_keys.add(tag["key"])
 
         endpoint = Endpoint.create(
             endpoint_name=endpoint_name,

sagemaker-serve/tests/unit/test_model_builder_deploy.py

@@ -7,7 +7,7 @@
 from unittest.mock import Mock, patch, MagicMock, call
 import tempfile
 
-from sagemaker.serve.model_builder import ModelBuilder
+from sagemaker.serve.model_builder import ModelBuilder, _tags_as_key_value_list
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.constants import Framework
@@ -602,3 +602,115 @@ def test_reset_build_state_clears_upload_state(self):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+class TestTagsAsKeyValueList(unittest.TestCase):
+    """Tests for the tag normalization used by the model-customization deploy paths."""
+
+    def test_lowercase_list_passes_through(self):
+        self.assertEqual(
+            _tags_as_key_value_list([{"key": "a", "value": "b"}]), [{"key": "a", "value": "b"}]
+        )
+
+    def test_capitalized_list_is_lowercased(self):
+        self.assertEqual(
+            _tags_as_key_value_list([{"Key": "a", "Value": "b"}]), [{"key": "a", "value": "b"}]
+        )
+
+    def test_bare_mapping_is_expanded(self):
+        self.assertEqual(_tags_as_key_value_list({"a": "b"}), [{"key": "a", "value": "b"}])
+
+    def test_none_and_empty_yield_empty_list(self):
+        self.assertEqual(_tags_as_key_value_list(None), [])
+        self.assertEqual(_tags_as_key_value_list([]), [])
+
+    def test_empty_value_is_kept(self):
+        """Tag.Value has a minimum length of 0, so an empty value is a legal tag."""
+        self.assertEqual(
+            _tags_as_key_value_list([{"key": "a", "value": ""}]), [{"key": "a", "value": ""}]
+        )
+
+    def test_entry_missing_key_or_value_is_skipped(self):
+        self.assertEqual(
+            _tags_as_key_value_list([{"key": "a"}, {"value": "b"}, {"key": "c", "value": "d"}]),
+            [{"key": "c", "value": "d"}],
+        )
+
+    def test_output_validates_as_core_tag(self):
+        """The output must satisfy the List[Tag] shape the create() calls are typed with."""
+        from pydantic import TypeAdapter
+        from typing import List as TypingList
+        from sagemaker.core.shapes import Tag
+
+        coerced = TypeAdapter(TypingList[Tag]).validate_python(
+            _tags_as_key_value_list([{"Key": "a", "Value": "b"}])
+        )
+        self.assertEqual([(t.key, t.value) for t in coerced], [("a", "b")])
+
+
+class TestDeployNovaModelTags(unittest.TestCase):
+    """Tag propagation through the Nova model-customization deploy path.
+
+    ``deploy(tags=...)`` previously never reached ``self._tags`` on the
+    model-customization paths, so user tags were dropped before Endpoint.create.
+    """
+
+    PROJECT_TAG = {"key": "sagemaker:project-id", "value": "p-12345"}
+
+    def _make_builder(self, tags=None):
+        builder = ModelBuilder.__new__(ModelBuilder)
+        builder.instance_type = "ml.g5.12xlarge"
+        builder.built_model = MagicMock(model_name="model-1")
+        builder._tags = None
+        base_model = MagicMock(hub_content_name="amazon-nova-lite-v1", recipe_name="nova-sft")
+        package = MagicMock()
+        package.inference_specification.containers = [MagicMock(base_model=base_model)]
+        builder._fetch_model_package = MagicMock(return_value=package)
+        builder._base_model_name = MagicMock(return_value="amazon-nova-lite-v1")
+        builder._is_raw_s3_model = MagicMock(return_value=False)
+        if tags:
+            builder.add_tags(_tags_as_key_value_list(tags))
+        return builder
+
+    def _created_tags(self, builder):
+        with patch("sagemaker.serve.model_builder.EndpointConfig"), patch(
+            "sagemaker.serve.model_builder.Endpoint"
+        ) as mock_endpoint:
+            mock_endpoint.create.return_value = MagicMock()
+            builder._deploy_nova_model(endpoint_name="endpoint-1", wait=False)
+            return mock_endpoint.create.call_args.kwargs["tags"]
+
+    def test_user_tags_merged_with_jumpstart_tags(self):
+        tags = self._created_tags(self._make_builder(tags=[self.PROJECT_TAG]))
+        self.assertIn(self.PROJECT_TAG, tags)
+        self.assertTrue(any("jumpstart-model-id" in tag["key"] for tag in tags))
+
+    def test_capitalized_user_tags_accepted(self):
+        builder = self._make_builder(tags=[{"Key": "sagemaker:project-id", "Value": "p-12345"}])
+        self.assertIn(self.PROJECT_TAG, self._created_tags(builder))
+
+    def test_no_user_tags_leaves_jumpstart_tags_only(self):
+        tags = self._created_tags(self._make_builder())
+        self.assertEqual(len(tags), 2)
+        self.assertTrue(all(tag["key"].startswith("sagemaker-sdk:") for tag in tags))
+
+    def test_deploy_tags_reach_the_endpoint_through_model_customization(self):
+        """Covers the full chain: deploy(tags=...) -> kwargs -> self._tags -> Endpoint.create.
+
+        _deploy_model_customization() is the entry point the customization deploy path
+        takes, and it is where the deploy-time tags are merged into the accumulator.
+        """
+        builder = self._make_builder()
+        builder._is_nova_model = MagicMock(return_value=True)
+
+        with patch("sagemaker.serve.model_builder.EndpointConfig"), patch(
+            "sagemaker.serve.model_builder.Endpoint"
+        ) as mock_endpoint:
+            mock_endpoint.create.return_value = MagicMock()
+            builder._deploy_model_customization(
+                endpoint_name="endpoint-1", wait=False, tags=[self.PROJECT_TAG]
+            )
+            tags = mock_endpoint.create.call_args.kwargs["tags"]
+
+        self.assertIn(self.PROJECT_TAG, tags)
+        self.assertTrue(any("jumpstart-model-id" in tag["key"] for tag in tags))