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

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,61 @@
 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
+
+
+def _merge_tags(*tag_sets: Optional[Tags]) -> List[TagsDict]:
+    """Merge tag sets into one key/value list, keeping the first entry for a repeated key.
+
+    SageMaker rejects duplicate tag keys, so a later set cannot override an earlier one.
+    Callers pass the tags SageMaker manages first so a caller-supplied tag cannot displace
+    one the SDK relies on (for example the model-source tag used for resource reuse).
+    """
+    merged: List[TagsDict] = []
+    seen = set()
+    for tags in tag_sets:
+        for tag in _tags_as_key_value_list(tags):
+            if tag["key"] not in seen:
+                merged.append(tag)
+                seen.add(tag["key"])
+    return merged
+
+
 @dataclass
 class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuilderUtils):
     """Unified interface for building and deploying machine learning models.
@@ -6185,6 +6241,14 @@ 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, and the tags ModelBuilder
+        # manages itself are on self._tags. Merge both for the endpoint created
+        # below, whichever path creates it. This deliberately avoids add_tags():
+        # tag_exists() indexes tags as tag["Key"], so writing the lowercase form
+        # the resource create() calls require back onto the accumulator raises
+        # KeyError once self._tags is non-empty (which build() leaves it).
+        endpoint_tags = _merge_tags(getattr(self, "_tags", None), kwargs.get("tags"))
+
         # An inference_config of ResourceRequirements requests an inference
         # component deployment; otherwise the model is placed directly on the
         # production variant.
@@ -6201,6 +6265,7 @@ def _deploy_model_customization(
                 endpoint_name=endpoint_name,
                 initial_instance_count=initial_instance_count,
                 wait=kwargs.get("wait", True),
+                tags=endpoint_tags,
             )
 
         # The model package may be absent (e.g. a Nova CPTTrainer or raw-S3
@@ -6226,7 +6291,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=endpoint_tags or None,
             )
             if kwargs.get("wait", True):
                 endpoint.wait_for_status("InService")
@@ -6266,10 +6333,6 @@ 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 = Endpoint.create(
                 endpoint_name=endpoint_name,
                 endpoint_config_name=endpoint_name,
@@ -6500,6 +6563,7 @@ def _deploy_nova_model(
         endpoint_name: str,
         initial_instance_count: int = 1,
         wait: bool = True,
+        tags: Optional[Tags] = None,
     ) -> Endpoint:
         """Deploy a Nova model directly to an endpoint without inference components.
 
@@ -6528,29 +6592,28 @@ def _deploy_nova_model(
         # The jumpstart-model-id tag always applies (resolved from the model
         # package or the trainer's base_model_name). The recipe-name tag is only
         # available when a model package is present.
-        tags = [
+        nova_tags = [
             {"key": "sagemaker-sdk:jumpstart-model-id", "value": self._base_model_name()},
         ]
         model_package = self._fetch_model_package()
         if model_package is not None:
             base_model = model_package.inference_specification.containers[0].base_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.
-        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)
+                nova_tags.append(
+                    {"key": "sagemaker-sdk:recipe-name", "value": base_model.recipe_name}
+                )
+
+        # ``tags`` carries the tags resolved by the caller: those ModelBuilder
+        # manages plus any passed to deploy(). When this method is called directly,
+        # fall back to the accumulator so its tags are still applied.
+        endpoint_tags = _merge_tags(
+            nova_tags, tags if tags is not None else getattr(self, "_tags", None)
+        )
 
         endpoint = Endpoint.create(
             endpoint_name=endpoint_name,
             endpoint_config_name=endpoint_name,
-            tags=tags,
+            tags=endpoint_tags,
         )
 
         if wait:

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,168 @@ 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 -> Endpoint.create.
+
+        _deploy_model_customization() is the entry point the customization deploy path
+        takes, and it is where the deploy-time tags are resolved for every endpoint
+        branch below it.
+        """
+        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))
+
+    def test_deploy_tags_merge_with_tags_already_on_the_builder(self):
+        """The realistic build()-then-deploy() flow: self._tags is already populated.
+
+        build() leaves capitalized entries on self._tags (JumpStart tags) and can leave
+        Tag objects there for trainer-backed models, so the deploy-time merge has to cope
+        with a non-empty accumulator in a different form than the caller's tags.
+        """
+        from sagemaker.core.shapes import Tag as CoreTag
+
+        builder = self._make_builder()
+        builder._is_nova_model = MagicMock(return_value=True)
+        builder._tags = [
+            {"Key": "sagemaker-sdk:jumpstart-model-id", "Value": "llama"},
+            CoreTag(key="from-trainer", value="yes"),
+        ]
+
+        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.assertIn({"key": "from-trainer", "value": "yes"}, tags)
+        self.assertTrue(any(tag["key"] == "sagemaker-sdk:jumpstart-model-id" for tag in tags))
+
+    def test_a_managed_tag_is_not_displaced_by_a_colliding_deploy_tag(self):
+        """SageMaker rejects duplicate keys, so the tag the SDK relies on must win."""
+        builder = self._make_builder()
+        builder._is_nova_model = MagicMock(return_value=True)
+        builder._tags = [{"Key": "sagemaker-sdk:jumpstart-model-id", "Value": "managed"}]
+
+        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=[{"key": "sagemaker-sdk:jumpstart-model-id", "value": "user-override"}],
+            )
+            tags = mock_endpoint.create.call_args.kwargs["tags"]
+
+        keys = [tag["key"] for tag in tags]
+        self.assertEqual(keys.count("sagemaker-sdk:jumpstart-model-id"), 1)
+        self.assertNotIn(
+            {"key": "sagemaker-sdk:jumpstart-model-id", "value": "user-override"}, tags
+        )