fix: let ModelBuilder.build() tag the SageMaker Model it creates

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. "

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

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)