From 37d2b09f460ec913a129781376d77031c3ea4227 Mon Sep 17 00:00:00 2001 From: Rohan Gujarathi Date: Wed, 12 Aug 2026 23:48:12 +0000 Subject: [PATCH] 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 applies the deploy-time tags kwarg. Those paths built their endpoint tags from self._tags alone, so they saw only the tags ModelBuilder had added itself. Resolve the endpoint tags once on entry from both sources and use them for whichever path creates the endpoint. 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, Tag objects, 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(). The merge deliberately does not go through add_tags(): tag_exists() compares tags as tag["Key"] == curr_tag["Key"], so writing the lowercase form these create() calls require back onto the accumulator raises KeyError as soon as self._tags is non-empty, which is the normal state after build(). Resolving into a local list instead also keeps deploy() from mutating builder state as a side effect. SageMaker rejects duplicate tag keys, so on a collision the tag ModelBuilder manages is kept and the caller's duplicate is dropped rather than displacing a tag the SDK relies on. --- .../src/sagemaker/serve/model_builder.py | 99 +++++++++-- .../tests/unit/test_model_builder_deploy.py | 167 +++++++++++++++++- 2 files changed, 247 insertions(+), 19 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 8472e6a584..9730f48cb3 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/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: diff --git a/sagemaker-serve/tests/unit/test_model_builder_deploy.py b/sagemaker-serve/tests/unit/test_model_builder_deploy.py index 3a0fca3d8e..cae156a141 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_deploy.py +++ b/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 + )