From 09e49b22b2ce3ec12985bc8959bca6dbd542f807 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Tue, 20 Oct 2026 16:15:58 -0700 Subject: [PATCH] fix: let ModelBuilder.build() tag the SageMaker Model it creates --- .../src/sagemaker/serve/model_builder.py | 21 +++ ...st_model_builder_build_tags_integration.py | 130 +++++++++++++++ .../unit/test_model_builder_build_tags.py | 150 ++++++++++++++++++ 3 files changed, 301 insertions(+) create mode 100644 sagemaker-serve/tests/integ/test_model_builder_build_tags_integration.py create mode 100644 sagemaker-serve/tests/unit/test_model_builder_build_tags.py diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 8675245d0f..556525a45e 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -4510,6 +4510,7 @@ def build( sagemaker_session: Optional[Session] = None, region: Optional[str] = None, reuse_resources: bool = False, + tags: Optional[Tags] = None, ) -> Union[Model, "ModelBuilder", None]: """Build a deployable ``Model`` instance with ``ModelBuilder``. @@ -4538,6 +4539,14 @@ def build( no new Model and sets ``built_model`` to the existing Model; on a miss, it creates a new Model. Endpoint reuse is handled independently by passing ``reuse_resources=True`` to deploy(). (Default: False). + tags (Optional[Tags]): Tags to attach to the SageMaker Model resource that + build() creates, as boto-style ``{"Key": ..., "Value": ...}`` dicts or + ``sagemaker.core.shapes.Tag`` objects. Tags passed here are merged with + any added through ``add_tags()`` and are preserved when the build is + captured by a pipeline ``ModelStep``. Pass tags here rather than to + ``deploy()`` when the model is built without an endpoint — for a pipeline + step, batch transform, evaluation, or an inference pipeline. + (Default: None). Returns: Union[Model, ModelBuilder, None]: A ``sagemaker.core.resources.Model`` resource @@ -4552,6 +4561,18 @@ def build( """ self._built_model_was_reused = False + # 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. + if tags is not None: + from sagemaker.core.utils.utils import serialize + + self.add_tags(serialize(format_tags(tags))) + if hasattr(self, "built_model") and self.built_model is not None: logger.warning( "ModelBuilder.build() has already been called. " diff --git a/sagemaker-serve/tests/integ/test_model_builder_build_tags_integration.py b/sagemaker-serve/tests/integ/test_model_builder_build_tags_integration.py new file mode 100644 index 0000000000..456e5eadb0 --- /dev/null +++ b/sagemaker-serve/tests/integ/test_model_builder_build_tags_integration.py @@ -0,0 +1,130 @@ +# 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. +"""Integration test for tagging the Model that ``ModelBuilder.build()`` creates. + +Regression for issue #6501. ``build()`` took no ``tags`` argument, so every model +built without deploying -- a pipeline ``ModelStep``, batch transform, evaluation, +an inference pipeline -- was created untagged. Unit tests can only assert the +request the SDK sends; this reads the tags back off the created Model with +``ListTags`` to prove they reached the service. build() only (no deploy), so it +runs in seconds. +""" + +from __future__ import absolute_import + +import os +import tarfile +import tempfile +import uuid + +import boto3 + +from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.shapes import Tag +from sagemaker.serve.mode.function_pointers import Mode +from sagemaker.serve.model_builder import ModelBuilder + +MODEL_NAME_PREFIX = "mb-build-tags" + + +def _upload_artifact(s3_client, bucket, prefix): + """Upload a minimal model.tar.gz and return its URI and key.""" + tmp = tempfile.mkdtemp() + model_file = os.path.join(tmp, "model.json") + with open(model_file, "w") as handle: + handle.write('{"weights": [1, 1, 1, 1]}') + tar_path = os.path.join(tmp, "model.tar.gz") + with tarfile.open(tar_path, "w:gz") as archive: + archive.add(model_file, arcname="model.json") + key = "{}/model.tar.gz".format(prefix) + s3_client.upload_file(tar_path, bucket, key) + return "s3://{}/{}".format(bucket, key), key + + +def _model_builder(session, image_uri, artifact_uri, role, unique_id): + builder = ModelBuilder( + image_uri=image_uri, + s3_model_data_url=artifact_uri, + role_arn=role, + instance_type="ml.m5.large", + sagemaker_session=session, + ) + builder.model_path = "/tmp/sagemaker/model-builder/{}".format(unique_id) + os.makedirs(builder.model_path, exist_ok=True) + builder.dependencies = [] + return builder + + +def _tags_on(sagemaker_client, model_arn): + return { + tag["Key"]: tag["Value"] + for tag in sagemaker_client.list_tags(ResourceArn=model_arn)["Tags"] + } + + +def test_build_tags_reach_the_created_model(): + """build(tags=...) must tag the Model, for both tag representations.""" + session = Session() + region = session.boto_region_name + bucket = session.default_bucket() + role = get_execution_role(sagemaker_session=session) + s3_client = boto3.client("s3", region_name=region) + sagemaker_client = boto3.client("sagemaker", region_name=region) + + unique_id = uuid.uuid4().hex[:8] + prefix = "{}/{}".format(MODEL_NAME_PREFIX, unique_id) + s3_keys = [] + created_models = [] + + try: + from sagemaker.core import image_uris + + image_uri = image_uris.retrieve( + framework="sklearn", + region=region, + version="1.2-1", + instance_type="ml.m5.large", + image_scope="inference", + ) + artifact_uri, key = _upload_artifact(s3_client, bucket, prefix) + s3_keys.append(key) + + # Boto-style tags. + boto_model = _model_builder(session, image_uri, artifact_uri, role, unique_id).build( + model_name="{}-boto-{}".format(MODEL_NAME_PREFIX, unique_id), + mode=Mode.SAGEMAKER_ENDPOINT, + tags=[{"Key": "project", "Value": "build-tags"}], + ) + created_models.append(boto_model) + assert _tags_on(sagemaker_client, boto_model.model_arn)["project"] == "build-tags" + + # v3 core Tag objects must work too, not just the boto-style dicts. + core_model = _model_builder(session, image_uri, artifact_uri, role, unique_id).build( + model_name="{}-core-{}".format(MODEL_NAME_PREFIX, unique_id), + mode=Mode.SAGEMAKER_ENDPOINT, + tags=[Tag(key="project", value="build-tags-core")], + ) + created_models.append(core_model) + assert _tags_on(sagemaker_client, core_model.model_arn)["project"] == "build-tags-core" + + finally: + for model in created_models: + try: + model.delete() + except Exception: # noqa + pass + for s3_key in s3_keys: + try: + s3_client.delete_object(Bucket=bucket, Key=s3_key) + except Exception: # noqa + pass diff --git a/sagemaker-serve/tests/unit/test_model_builder_build_tags.py b/sagemaker-serve/tests/unit/test_model_builder_build_tags.py new file mode 100644 index 0000000000..de93653670 --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_build_tags.py @@ -0,0 +1,150 @@ +"""Unit tests for tagging the Model that ``ModelBuilder.build()`` creates (issue #6501). + +``build()`` creates a SageMaker Model resource and sends ``tags=format_tags(self._tags)``, +but ``_tags`` was only ever set by ``deploy()``. Any caller that builds without deploying +-- a pipeline ``ModelStep``, batch transform, evaluation, an inference pipeline -- created +an untagged Model with no public way to fix it. +""" + +import inspect +import unittest +from unittest.mock import MagicMock, Mock, patch + +from sagemaker.core.shapes import Tag +from sagemaker.core.workflow.pipeline_context import PipelineSession +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" +MODEL_NAME = "modelbuilder-build-tags" +BOTO_TAGS = [{"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() + 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.model_name = MODEL_NAME + builder.region = REGION + return builder + + +class TestBuildAcceptsTags(unittest.TestCase): + """``build()`` must expose a public way to tag the Model it creates.""" + + def test_build_has_a_tags_parameter(self): + self.assertIn("tags", inspect.signature(ModelBuilder.build).parameters) + + def _build_with(self, **kwargs): + """Run build() with model creation stubbed, returning the tags it would send.""" + builder = _make_builder() + seen = {} + + def _capture(*args, **call_kwargs): + seen["tags"] = builder._tags + builder.built_model = Mock(model_arn=f"arn:aws:sagemaker:{REGION}:1:model/x") + return builder.built_model + + with patch.object( + ModelBuilder, "_build_single_modelbuilder", side_effect=_capture + ), patch.object(ModelBuilder, "_get_serve_setting"): + builder.build(**kwargs) + + return builder, seen.get("tags") + + def test_boto_style_tags_reach_the_build_path(self): + _, tags = self._build_with(tags=BOTO_TAGS) + + self.assertEqual(tags, BOTO_TAGS) + + def test_core_tag_objects_are_normalised_to_boto_keys(self): + """A v3 ``Tag`` must reach CreateModel as {"Key", "Value"}, not {"key", "value"}.""" + _, tags = self._build_with(tags=[Tag(key="project", value="example")]) + + self.assertEqual(tags, BOTO_TAGS) + + def test_mapping_form_is_expanded(self): + _, tags = self._build_with(tags={"project": "example"}) + + self.assertEqual(tags, BOTO_TAGS) + + def test_build_tags_merge_with_add_tags(self): + builder = _make_builder() + builder.add_tags({"Key": "team", "Value": "ml-platform"}) + seen = {} + + def _capture(*args, **kwargs): + seen["tags"] = list(builder._tags) + builder.built_model = Mock(model_arn=f"arn:aws:sagemaker:{REGION}:1:model/x") + return builder.built_model + + with patch.object( + ModelBuilder, "_build_single_modelbuilder", side_effect=_capture + ), patch.object(ModelBuilder, "_get_serve_setting"): + builder.build(tags=BOTO_TAGS) + + self.assertEqual( + seen["tags"], + [{"Key": "team", "Value": "ml-platform"}, {"Key": "project", "Value": "example"}], + ) + + def test_omitting_tags_leaves_them_unset(self): + _, tags = self._build_with() + + self.assertIsNone(tags) + + +@patch("sagemaker.serve.model_builder.resolve_nested_dict_value_from_config") +@patch("sagemaker.serve.model_builder.resolve_value_from_config") +@patch.object(ModelBuilder, "_init_sagemaker_session_if_does_not_exist") +@patch.object(ModelBuilder, "_ensure_base_name_if_needed") +@patch.object(ModelBuilder, "_prepare_container_def") +class TestCreateModelRequestCarriesTags(unittest.TestCase): + """The tags must survive into the CreateModel request, pipeline capture included.""" + + def _create_model_kwargs(self, mocks, tags): + mock_prepare, mock_resolve, mock_resolve_nested = mocks + mock_prepare.return_value = {"Image": IMAGE_URI, "Environment": {}} + mock_resolve.side_effect = lambda value, *args, **kwargs: value + mock_resolve_nested.side_effect = lambda value, *args, **kwargs: value + + builder = _make_builder() + # A PipelineSession captures the request instead of calling CreateModel, + # which is the ModelStep path the defect was reported against. + builder.sagemaker_session = Mock(spec=PipelineSession) + builder._tags = tags + + builder._create_sagemaker_model() + + builder.sagemaker_session.create_model.assert_called_once() + return builder.sagemaker_session.create_model.call_args.kwargs + + def test_pipeline_create_model_request_has_tags( + self, mock_prepare, mock_base_name, mock_init, mock_resolve, mock_resolve_nested + ): + kwargs = self._create_model_kwargs( + (mock_prepare, mock_resolve, mock_resolve_nested), BOTO_TAGS + ) + + self.assertEqual(kwargs["tags"], BOTO_TAGS)