Skip to content

Commit b04efc3

Browse files
mohamedzeidan2021Mohamed Zeidan
andauthored
fix: emit boto-style tag keys in ModelTrainer pipeline requests (#6485)
_create_training_job_args() called model_dump() on each Tag, which yields the model's own lowercase field names. serialize() renames the fields of Tag models to the boto-style {"Key": ..., "Value": ...} form the service requires, but passes plain dicts through untouched, so the dump happened too early and the pipeline definition kept lowercase keys: "Tags": [{"key": "project", "value": "ml-project"}] CreatePipeline rejected the definition with "Unable to parse pipeline definition. Unknown Argument member 'key'". The boto3 payload that TrainingQueue submits to AWS Batch carried the same invalid tags. Hand the validated Tag objects to the branch that serializes the request. The interactive request keeps the dicts, because TrainingJob.create annotates its tags parameter with sagemaker.core.resources.Tag, a resource class that shadows the shape of the same name, so a sagemaker.core.shapes.Tag instance fails its validation while the dicts are revalidated into it. Fixes #6474 Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>
1 parent c65ad99 commit b04efc3

2 files changed

Lines changed: 110 additions & 1 deletion

File tree

‎sagemaker-train/src/sagemaker/train/model_trainer.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -882,6 +882,11 @@ def _create_training_job_args(
882882
if boto3 or isinstance(self.sagemaker_session, PipelineSession):
883883
if isinstance(self.sagemaker_session, PipelineSession):
884884
training_request.pop("training_job_name", None)
885+
# ``serialize`` renames the fields of Tag models to the boto-style
886+
# {"Key": ..., "Value": ...} form the service requires, but passes plain dicts
887+
# through untouched, so hand it the validated Tag objects. ``TrainingJob.create``
888+
# takes the dicts instead, which it revalidates into its own Tag model.
889+
training_request["tags"] = self.tags
885890
# Convert snake_case to PascalCase for AWS API
886891
pipeline_request = {to_pascal_case(k): v for k, v in training_request.items()}
887892
serialized_request = serialize(pipeline_request)

‎sagemaker-train/tests/unit/train/test_model_trainer.py‎

Lines changed: 105 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
from unittest.mock import patch, MagicMock, ANY, mock_open
2626
from tempfile import NamedTemporaryFile
2727

28-
from sagemaker.core.resources import TrainingJob
28+
from sagemaker.core.resources import Base, TrainingJob
2929
from sagemaker.core.shapes import (
3030
ResourceConfig,
3131
VpcConfig,
@@ -45,6 +45,7 @@
4545
)
4646
from sagemaker.core.config.config_manager import SageMakerConfig
4747
from sagemaker.core.helper.session_helper import Session
48+
from sagemaker.core.workflow.pipeline_context import PipelineSession
4849
from sagemaker.train.model_trainer import ModelTrainer, Mode
4950
from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE
5051
from sagemaker.train.constants import (
@@ -2616,3 +2617,106 @@ def test_output_data_strips_trailing_slash_on_fallback(model_trainer):
26162617
assert model_trainer.output_data == (
26172618
f"s3://{DEFAULT_BUCKET}/{DEFAULT_BUCKET_PREFIX}" "/my-training-job/output/output.tar.gz"
26182619
)
2620+
2621+
2622+
# ---------------------------------------------------------------------------
2623+
# Tag serialization (issue #6474)
2624+
# ---------------------------------------------------------------------------
2625+
2626+
TAG = Tag(key="project", value="ml-project")
2627+
BOTO_TAG = {"Key": "project", "Value": "ml-project"}
2628+
2629+
2630+
def _pipeline_session():
2631+
"""A real PipelineSession with only its outbound edges stubbed."""
2632+
session = PipelineSession.__new__(PipelineSession)
2633+
session._context = None
2634+
session.boto_session = MagicMock()
2635+
session.boto_session.region_name = DEFAULT_REGION
2636+
session.sagemaker_config = {}
2637+
session.default_bucket = lambda: DEFAULT_BUCKET
2638+
session.default_bucket_prefix = DEFAULT_BUCKET_PREFIX
2639+
return session
2640+
2641+
2642+
def _pipeline_trainer(tags):
2643+
"""A trainer whose session makes ``_create_training_job_args`` serialize the request."""
2644+
return ModelTrainer(
2645+
training_image=DEFAULT_IMAGE,
2646+
role=DEFAULT_ROLE,
2647+
sagemaker_session=_pipeline_session(),
2648+
compute=DEFAULT_COMPUTE_CONFIG,
2649+
stopping_condition=DEFAULT_STOPPING_CONDITION,
2650+
output_data_config=DEFAULT_OUTPUT_DATA_CONFIG,
2651+
tags=tags,
2652+
)
2653+
2654+
2655+
def _train_wire_tags(trainer):
2656+
"""The Tags CreateTrainingJob really receives, measured not derived."""
2657+
mock_client = MagicMock()
2658+
with patch.object(Base, "get_sagemaker_client", return_value=mock_client):
2659+
try:
2660+
trainer.train(wait=False)
2661+
except ValidationError:
2662+
# Building the TrainingJob resource out of a MagicMock response fails after
2663+
# the call. The call itself is what is under test.
2664+
pass
2665+
assert (
2666+
mock_client.create_training_job.call_args is not None
2667+
), "create_training_job was not reached"
2668+
return mock_client.create_training_job.call_args.kwargs.get("Tags")
2669+
2670+
2671+
def test_pipeline_request_uses_boto_style_tag_keys():
2672+
"""Regression for #6474.
2673+
2674+
The service requires Tags of the form {"Key": ..., "Value": ...}. ``serialize``
2675+
renames the fields of Tag models but passes plain dicts through untouched, so
2676+
dumping the tags to dicts first left lowercase keys in the pipeline definition
2677+
and CreatePipeline rejected it with "Unknown Argument member 'key'".
2678+
"""
2679+
args = _pipeline_trainer([TAG])._create_training_job_args()
2680+
2681+
assert args["Tags"] == [BOTO_TAG]
2682+
2683+
2684+
def test_pipeline_request_omits_tags_when_unset():
2685+
"""An untagged trainer must not add an empty Tags member to the request."""
2686+
args = _pipeline_trainer(None)._create_training_job_args()
2687+
2688+
assert "Tags" not in args
2689+
2690+
2691+
def test_boto3_request_uses_boto_style_tag_keys(model_trainer):
2692+
"""The boto3 payload (used by TrainingQueue) needs the same boto-style tags."""
2693+
model_trainer.tags = [TAG]
2694+
2695+
args = model_trainer._create_training_job_args(boto3=True)
2696+
2697+
assert args["Tags"] == [BOTO_TAG]
2698+
2699+
2700+
def test_interactive_request_keeps_tags_training_job_create_accepts(model_trainer):
2701+
"""The interactive request carries dicts, not Tag models.
2702+
2703+
``TrainingJob.create`` annotates ``tags`` with ``sagemaker.core.resources.Tag``,
2704+
which shadows the shape of the same name, so a ``sagemaker.core.shapes.Tag``
2705+
instance fails its validation. The dicts are revalidated into that model instead,
2706+
and come out boto-style on the wire.
2707+
"""
2708+
model_trainer.tags = [TAG]
2709+
2710+
args = model_trainer._create_training_job_args()
2711+
2712+
assert args["tags"] == [{"key": "project", "value": "ml-project"}]
2713+
assert _train_wire_tags(model_trainer) == [BOTO_TAG]
2714+
2715+
2716+
def test_pipeline_and_interactive_emit_identical_tags(model_trainer):
2717+
"""Both paths must put the same tags on the request."""
2718+
model_trainer.tags = [TAG]
2719+
2720+
pipeline_tags = _pipeline_trainer([TAG])._create_training_job_args()["Tags"]
2721+
2722+
assert pipeline_tags == _train_wire_tags(model_trainer)

0 commit comments

Comments
 (0)