diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index f76faba505..21724ec6fd 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/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) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 556525a45e..ec239445bb 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/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, diff --git a/sagemaker-serve/tests/unit/test_model_builder_deploy_tags.py b/sagemaker-serve/tests/unit/test_model_builder_deploy_tags.py new file mode 100644 index 0000000000..a52239bc5d --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_deploy_tags.py @@ -0,0 +1,256 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""``ModelBuilder.deploy(tags=...)`` is one tag contract for every deploy path. + +Before this change ``tags`` reached ``deploy()`` only through ``**kwargs``, so +``inspect.signature`` could not validate it, the fine-tuned (model customization) and +Nova paths ignored it, and ``add_tags`` raised ``KeyError: 'Key'`` on the +``{"key": ..., "value": ...}`` form that the ``sagemaker.core.resources`` classes use. +""" + +from __future__ import absolute_import + +import inspect +import unittest +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from sagemaker.core.common_utils import _validate_new_tags, format_tags, remove_tag_with_key +from sagemaker.core.enums import EndpointType +from sagemaker.core.inference_config import ResourceRequirements +from sagemaker.core.resources import Endpoint +from sagemaker.core.shapes import Tag +from sagemaker.serve.model_builder import ModelBuilder +from sagemaker.serve.utils.types import ModelServer + +ROLE_ARN = "arn:aws:iam::123456789012:role/TestRole" +REGION = "us-east-2" +IMAGE_URI = "123456789012.dkr.ecr.us-east-2.amazonaws.com/inference:latest" +BOTO_TAG = {"Key": "project", "Value": "example"} +CORE_TAG = {"key": "project", "value": "example"} + + +def _session(): + session = Mock() + session.boto_region_name = REGION + session.default_bucket.return_value = "test-bucket" + session.default_bucket_prefix = None + session.config = {} + session.sagemaker_config = {} + session.settings = Mock(include_jumpstart_tags=False) + session.settings._local_download_dir = None + session.sagemaker_client = MagicMock() + session.boto_session = Mock() + session.boto_session.region_name = REGION + return session + + +def _make_builder(): + builder = ModelBuilder( + model=Mock(), + model_server=ModelServer.TORCHSERVE, + image_uri=IMAGE_URI, + s3_model_data_url="s3://test-bucket/model.tar.gz", + role_arn=ROLE_ARN, + instance_type="ml.m5.large", + sagemaker_session=_session(), + ) + builder.region = REGION + builder.built_model = Mock(model_name="built-model") + builder.model_server = None + return builder + + +class TestTagNormalization(unittest.TestCase): + """``format_tags`` and ``add_tags`` accept every tag representation the SDK uses.""" + + def test_format_tags_accepts_every_representation(self): + expected = [BOTO_TAG] + + self.assertEqual(format_tags({"project": "example"}), expected) + self.assertEqual(format_tags(BOTO_TAG), expected) + self.assertEqual(format_tags([BOTO_TAG]), expected) + self.assertEqual(format_tags([CORE_TAG]), expected) + self.assertEqual(format_tags([Tag(key="project", value="example")]), expected) + self.assertEqual( + format_tags([BOTO_TAG, {"key": "team", "value": "ml"}]), + [BOTO_TAG, {"Key": "team", "Value": "ml"}], + ) + self.assertIsNone(format_tags(None)) + + def test_format_tags_rejects_a_tag_without_a_key(self): + with pytest.raises(TypeError, match="'Key'"): + format_tags([{"Value": "no-key"}]) + with pytest.raises(TypeError, match="Tag object"): + format_tags(["project=example"]) + + def test_validate_new_tags_merges_mixed_casing_by_key(self): + merged = _validate_new_tags([CORE_TAG, {"key": "team", "value": "ml"}], [BOTO_TAG]) + + self.assertEqual(merged, [BOTO_TAG, {"Key": "team", "Value": "ml"}]) + + def test_validate_new_tags_with_tag_objects(self): + merged = _validate_new_tags([Tag(key="team", value="ml")], [CORE_TAG]) + + self.assertEqual(merged, [BOTO_TAG, {"Key": "team", "Value": "ml"}]) + + def test_remove_tag_with_key_accepts_core_casing(self): + self.assertIsNone(remove_tag_with_key("project", [CORE_TAG])) + + def test_add_tags_normalizes_core_casing_to_boto(self): + builder = _make_builder() + + builder.add_tags([CORE_TAG]) + builder.add_tags({"team": "ml"}) + builder.add_tags([Tag(key="env", value="prod")]) + + self.assertEqual( + builder._tags, + [BOTO_TAG, {"Key": "team", "Value": "ml"}, {"Key": "env", "Value": "prod"}], + ) + + +class TestDeployTagsParameter(unittest.TestCase): + """``deploy(tags=...)`` is a named parameter and reaches every deploy path.""" + + def test_deploy_has_a_tags_parameter(self): + parameters = inspect.signature(ModelBuilder.deploy).parameters + + self.assertIn("tags", parameters) + self.assertIsNone(parameters["tags"].default) + + def test_deploy_merges_tags_into_the_builder_before_the_core_path(self): + builder = _make_builder() + builder.add_tags({"team": "ml"}) + seen = {} + + def _capture(**kwargs): + seen["tags"] = list(builder._tags) + seen["kwargs_tags"] = kwargs.get("tags") + return Mock(spec=Endpoint) + + with patch.object(ModelBuilder, "_deploy", side_effect=_capture), patch.object( + ModelBuilder, "_is_model_customization", return_value=False + ), patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None): + builder.deploy(endpoint_name="ep", tags=[CORE_TAG]) + + self.assertEqual(seen["tags"], [{"Key": "team", "Value": "ml"}, BOTO_TAG]) + # The core path reads self._tags; no duplicate copy travels through kwargs. + self.assertIsNone(seen["kwargs_tags"]) + + def test_deploy_core_endpoint_normalizes_tags_passed_through_kwargs(self): + builder = _make_builder() + builder.built_model.model_name = "built-model" + builder.model_name = "built-model" + builder.sagemaker_session.endpoint_in_service_or_not.return_value = False + seen = {} + + def _capture(**kwargs): + seen["tags"] = kwargs.get("tags") + return "ep" + + builder.sagemaker_session.endpoint_from_production_variants = Mock(side_effect=_capture) + with patch( + "sagemaker.serve.model_builder.session_helper.production_variant", + return_value={"VariantName": "AllTraffic"}, + ), patch( + "sagemaker.serve.model_builder.Endpoint.get", return_value=Mock(spec=Endpoint) + ), patch.object( + builder, "_wait_for_endpoint" + ): + builder._deploy_core_endpoint( + instance_type="ml.g5.xlarge", + initial_instance_count=1, + endpoint_name="ep", + endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, + resources=ResourceRequirements(requests={"memory": 1024, "copies": 1}), + inference_component_name="ic", + tags=[CORE_TAG, Tag(key="team", value="ml")], + ) + + self.assertEqual(seen["tags"], [BOTO_TAG, {"Key": "team", "Value": "ml"}]) + self.assertEqual(builder._tags, [BOTO_TAG, {"Key": "team", "Value": "ml"}]) + + def test_deploy_passes_the_builder_tags_to_a_recommendation_deploy(self): + builder = _make_builder() + builder._recommendation_job = Mock() + seen = {} + + def _capture(**kwargs): + seen["tags"] = kwargs["tags"] + return Mock(spec=Endpoint) + + with patch.object(ModelBuilder, "_deploy_recommendation", side_effect=_capture): + builder.deploy(tags={"project": "example"}) + + self.assertEqual(seen["tags"], [BOTO_TAG]) + + def test_core_endpoint_tags_use_the_core_resource_casing(self): + builder = _make_builder() + builder.add_tags([BOTO_TAG, Tag(key="team", value="ml")]) + + self.assertEqual( + builder._core_endpoint_tags(), + [CORE_TAG, {"key": "team", "value": "ml"}], + ) + + def test_model_customization_restricted_package_endpoint_is_tagged(self): + builder = _make_builder() + builder.instance_type = "ml.g5.xlarge" + created = {} + + def _create_endpoint(**kwargs): + created.update(kwargs) + endpoint = Mock(spec=Endpoint) + endpoint.wait_for_status = Mock() + return endpoint + + with patch.object(ModelBuilder, "_is_nova_model", return_value=False), patch.object( + ModelBuilder, "_fetch_model_package", return_value=Mock() + ), patch( + "sagemaker.serve.utils.model_package_utils.is_restricted_model_package", + return_value=True, + ), patch( + "sagemaker.serve.model_builder.EndpointConfig.create" + ), patch( + "sagemaker.serve.model_builder.Endpoint.create", side_effect=_create_endpoint + ): + builder.add_tags([CORE_TAG]) + builder._deploy_model_customization(endpoint_name="ep", wait=False) + + self.assertEqual(created["tags"], [CORE_TAG]) + + def test_nova_endpoint_merges_the_builder_tags(self): + builder = _make_builder() + builder.instance_type = "ml.g5.xlarge" + created = {} + + def _create_endpoint(**kwargs): + created.update(kwargs) + return Mock(spec=Endpoint) + + with patch.object( + ModelBuilder, "_base_model_name", return_value="nova-micro" + ), patch.object(ModelBuilder, "_fetch_model_package", return_value=None), patch( + "sagemaker.serve.model_builder.EndpointConfig.create" + ), patch( + "sagemaker.serve.model_builder.Endpoint.create", side_effect=_create_endpoint + ): + builder.add_tags([CORE_TAG, {"Key": "sagemaker-sdk:jumpstart-model-id", "Value": "x"}]) + builder._deploy_nova_model(endpoint_name="ep", wait=False) + + self.assertEqual( + created["tags"], + [{"key": "sagemaker-sdk:jumpstart-model-id", "value": "nova-micro"}, CORE_TAG], + )