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 + )