Repository navigation
fix: let ModelBuilder.build() tag the SageMaker Model it creates - #6510
mohamedzeidan2021 merged 1 commit into
Conversation
There was a problem hiding this comment.
🤖 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.
Issue
Fixes #6501
ModelBuilder.build()creates a SageMaker Model resource and sends its tags:but
_tagsis not reachable from any public API:init=Falsekeeps it out of the constructor, and the only thing that assigned it wasdeploy(). So every caller that builds without deploying created an untagged Model with no supported way to fix it:build()called beforedeploy()build()with no endpoint at allbuild()captured in a pipelineModelStepbuild(tags=...)raisedTypeError; dropping the argument succeeded and produced aCreateModelrequest with noTags.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
tagsparameter tobuild():Three details:
add_tags()rather than assignment, so tags passed tobuild()merge with any already added through the existingadd_tags()API instead of silently replacing them.format_tags()expands the{key: value}mapping form into{"Key", "Value"}dicts.serialize()normalisessagemaker.core.shapes.Tagobjects to boto-style dicts.format_tags()passes lists straight through, so without this a v3Tagwould travel as a pydantic object all the way into the botocreate_modelcall. This is theserialize()-is-the-right-tool point fromModelTraineremits invalid lowercase tag keys in pipeline definitions #6474 applied at the new entry point — a user reaching for atagsparameter on a v3 API should be able to hand it the v3Tagtype.Set before the rest of
build()runs, so the tags are in place for both the interactiveCreateModeland the request aPipelineSessioncaptures for aModelStep.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-fixbuild()has notagsparameter; post-fix the tags reachself._tagsand appear asTagsin the capturedCreateModelrequest.sagemaker-serve/tests/unit/test_model_builder_build_tags.py(7 tests):build()exposestags; boto-style dicts, coreTagobjects and the{key: value}mapping form all arrive as[{"Key": ..., "Value": ...}];build(tags=...)merges with earlieradd_tags()rather than overwriting; omittingtagsleaves_tagsasNone; and the tags survive into thecreate_modelcall on thePipelineSessionpath, which is theModelStepcapture the issue was reported against.sagemaker-serve/tests/integ/test_model_builder_build_tags_integration.py: builds two Models (one with boto-style tags, one with coreTagobjects) and reads the tags back withListTagson 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.sagemaker-serve/tests/unit: 2179 passed, 25 skipped. The 5 failures inmarshalling/test_triton_translator.py::TestTensorflowTensorTranslatorare pre-existing and reproduce identically on unmodifiedmasterin the same environment (no TensorFlow installed).black(100) clean,flake8(120) clean,pylint9.89/10 — identical to the unmodified baseline, with no message in the changed file.Backwards compatibility
Purely additive:
tagsis a new keyword parameter defaulting toNone, placed after every existing parameter, so positional callers are unaffected. Withtagsomitted, behaviour is byte-identical to before.deploy(tags=...)is unchanged and still overwrites_tagsfor the endpoint.