Skip to content

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

Merged
mohamedzeidan2021 merged 1 commit into
aws:masterfrom
mohamedzeidan2021:fix/issue-6501-modelbuilder-build-tags
Oct 21, 2026
Merged

mohamedzeidan2021 merged 1 commit into
aws:masterfrom
mohamedzeidan2021:fix/issue-6501-modelbuilder-build-tags

Conversation

@mohamedzeidan2021

Copy link
Copy Markdown
Collaborator

Issue

Fixes #6501

ModelBuilder.build() creates a SageMaker Model resource and sends its tags:

create_model_args = dict(
    ...
    tags=format_tags(self._tags),
)

but _tags is not reachable from any public API:

_tags: Optional[Tags] = field(default=None, init=False)

init=False keeps it out of the constructor, and the only thing that assigned it was deploy(). So every caller that builds without deploying created an untagged Model with no supported way to fix it:

  • build() called before deploy()
  • build() with no endpoint at all
  • build() captured in a pipeline ModelStep
  • models built for batch transform, evaluation, or inference pipelines

build(tags=...) raised TypeError; dropping the argument succeeded and produced a CreateModel request with no Tags.

Tagging is how the reporting customer attributes AWS spend across teams, so untagged resources are a governance gap rather than a cosmetic one.

Fix

Add a tags parameter to build():

if tags is not None:
    from sagemaker.core.utils.utils import serialize

    self.add_tags(serialize(format_tags(tags)))

Three details:

  • add_tags() rather than assignment, so tags passed to build() merge with any already added through the existing add_tags() API instead of silently replacing them.
  • format_tags() expands the {key: value} mapping form into {"Key", "Value"} dicts.
  • serialize() normalises sagemaker.core.shapes.Tag objects to boto-style dicts. format_tags() passes lists straight through, so without this a v3 Tag would travel as a pydantic object all the way into the boto create_model call. This is the serialize()-is-the-right-tool point from ModelTrainer emits invalid lowercase tag keys in pipeline definitions #6474 applied at the new entry point — a user reaching for a tags parameter on a v3 API should be able to hand it the v3 Tag type.

Set before the rest of build() runs, so the tags are in place for both the interactive CreateModel and the request a PipelineSession captures for a ModelStep.

This is the fifth tag-handling defect in this engagement (#6090, #6472, #6474, #6495, #6501) but a different mechanism from the other four: a missing public parameter rather than mis-normalisation. The central format_tags() normalisation work tracked from PR #6482 would not have prevented it, since there was no way to get tags onto this path at all.

Testing

Reproduced first with an offline script built from the repro in #6501 (mocked PipelineSession, no AWS): pre-fix build() has no tags parameter; post-fix the tags reach self._tags and appear as Tags in the captured CreateModel request.

  • New sagemaker-serve/tests/unit/test_model_builder_build_tags.py (7 tests): build() exposes tags; boto-style dicts, core Tag objects and the {key: value} mapping form all arrive as [{"Key": ..., "Value": ...}]; build(tags=...) merges with earlier add_tags() rather than overwriting; omitting tags leaves _tags as None; and the tags survive into the create_model call on the PipelineSession path, which is the ModelStep capture the issue was reported against.
  • New sagemaker-serve/tests/integ/test_model_builder_build_tags_integration.py: builds two Models (one with boto-style tags, one with core Tag objects) and reads the tags back with ListTags on the Model ARN, so the assertion is what the service stored rather than what the SDK sent. build() only, no deploy — runs in seconds, and cleans up both Models and the uploaded artifact.
  • Full sagemaker-serve/tests/unit: 2179 passed, 25 skipped. The 5 failures in marshalling/test_triton_translator.py::TestTensorflowTensorTranslator are pre-existing and reproduce identically on unmodified master in the same environment (no TensorFlow installed).
  • black (100) clean, flake8 (120) clean, pylint 9.89/10 — identical to the unmodified baseline, with no message in the changed file.

Backwards compatibility

Purely additive: tags is a new keyword parameter defaulting to None, placed after every existing parameter, so positional callers are unaffected. With tags omitted, behaviour is byte-identical to before. deploy(tags=...) is unchanged and still overwrites _tags for the endpoint.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 AI code review

This PR adds a tags parameter to ModelBuilder.build() so the SageMaker Model resource created during build() (e.g. for pipeline ModelStep, batch transform, or evaluation paths that never call deploy()) can be tagged. It normalizes the three tag representations via format_tags() + serialize() and merges them into _tags through the existing add_tags() API, with new unit and integration tests. The implementation is correct and consistent with the codebase: Tags is imported, serialize(format_tags(...)) produces boto-style {"Key","Value"} dicts for dict/mapping/core Tag inputs, add_tags merges rather than overwrites, and the @runnable_by_pipeline decorator forwards the new kwarg. No blocking issues found.
[]

Reviewed commit 09e49b2. Automated review; verify before acting.

@mohamedzeidan2021
mohamedzeidan2021 merged commit 6e99f9b into aws:master Oct 21, 2026
21 of 24 checks passed

This branch was previously deployed

1 inactive deployment
auto-approve — 09e49b22 Deployed Oct 20, 2026 by mohamedzeidan2021 via wait-for-approval #2070
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ModelBuilder.build() cannot tag the SageMaker model that it creates

2 participants