fix: emit boto-style tag keys in ModelTrainer pipeline requests

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

@@ -882,6 +882,11 @@ def _create_training_job_args(
         if boto3 or isinstance(self.sagemaker_session, PipelineSession):
             if isinstance(self.sagemaker_session, PipelineSession):
                 training_request.pop("training_job_name", None)
+            # ``serialize`` renames the fields of Tag models to the boto-style
+            # {"Key": ..., "Value": ...} form the service requires, but passes plain dicts
+            # through untouched, so hand it the validated Tag objects. ``TrainingJob.create``
+            # takes the dicts instead, which it revalidates into its own Tag model.
+            training_request["tags"] = self.tags
             # Convert snake_case to PascalCase for AWS API
             pipeline_request = {to_pascal_case(k): v for k, v in training_request.items()}
             serialized_request = serialize(pipeline_request)

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

@@ -25,7 +25,7 @@
 from unittest.mock import patch, MagicMock, ANY, mock_open
 from tempfile import NamedTemporaryFile
 
-from sagemaker.core.resources import TrainingJob
+from sagemaker.core.resources import Base, TrainingJob
 from sagemaker.core.shapes import (
     ResourceConfig,
     VpcConfig,
@@ -45,6 +45,7 @@
 )
 from sagemaker.core.config.config_manager import SageMakerConfig
 from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.workflow.pipeline_context import PipelineSession
 from sagemaker.train.model_trainer import ModelTrainer, Mode
 from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE
 from sagemaker.train.constants import (
@@ -2616,3 +2617,106 @@ def test_output_data_strips_trailing_slash_on_fallback(model_trainer):
     assert model_trainer.output_data == (
         f"s3://{DEFAULT_BUCKET}/{DEFAULT_BUCKET_PREFIX}" "/my-training-job/output/output.tar.gz"
     )
+
+
+# ---------------------------------------------------------------------------
+# Tag serialization (issue #6474)
+# ---------------------------------------------------------------------------
+
+TAG = Tag(key="project", value="ml-project")
+BOTO_TAG = {"Key": "project", "Value": "ml-project"}
+
+
+def _pipeline_session():
+    """A real PipelineSession with only its outbound edges stubbed."""
+    session = PipelineSession.__new__(PipelineSession)
+    session._context = None
+    session.boto_session = MagicMock()
+    session.boto_session.region_name = DEFAULT_REGION
+    session.sagemaker_config = {}
+    session.default_bucket = lambda: DEFAULT_BUCKET
+    session.default_bucket_prefix = DEFAULT_BUCKET_PREFIX
+    return session
+
+
+def _pipeline_trainer(tags):
+    """A trainer whose session makes ``_create_training_job_args`` serialize the request."""
+    return ModelTrainer(
+        training_image=DEFAULT_IMAGE,
+        role=DEFAULT_ROLE,
+        sagemaker_session=_pipeline_session(),
+        compute=DEFAULT_COMPUTE_CONFIG,
+        stopping_condition=DEFAULT_STOPPING_CONDITION,
+        output_data_config=DEFAULT_OUTPUT_DATA_CONFIG,
+        tags=tags,
+    )
+
+
+def _train_wire_tags(trainer):
+    """The Tags CreateTrainingJob really receives, measured not derived."""
+    mock_client = MagicMock()
+    with patch.object(Base, "get_sagemaker_client", return_value=mock_client):
+        try:
+            trainer.train(wait=False)
+        except ValidationError:
+            # Building the TrainingJob resource out of a MagicMock response fails after
+            # the call. The call itself is what is under test.
+            pass
+    assert (
+        mock_client.create_training_job.call_args is not None
+    ), "create_training_job was not reached"
+    return mock_client.create_training_job.call_args.kwargs.get("Tags")
+
+
+def test_pipeline_request_uses_boto_style_tag_keys():
+    """Regression for #6474.
+
+    The service requires Tags of the form {"Key": ..., "Value": ...}. ``serialize``
+    renames the fields of Tag models but passes plain dicts through untouched, so
+    dumping the tags to dicts first left lowercase keys in the pipeline definition
+    and CreatePipeline rejected it with "Unknown Argument member 'key'".
+    """
+    args = _pipeline_trainer([TAG])._create_training_job_args()
+
+    assert args["Tags"] == [BOTO_TAG]
+
+
+def test_pipeline_request_omits_tags_when_unset():
+    """An untagged trainer must not add an empty Tags member to the request."""
+    args = _pipeline_trainer(None)._create_training_job_args()
+
+    assert "Tags" not in args
+
+
+def test_boto3_request_uses_boto_style_tag_keys(model_trainer):
+    """The boto3 payload (used by TrainingQueue) needs the same boto-style tags."""
+    model_trainer.tags = [TAG]
+
+    args = model_trainer._create_training_job_args(boto3=True)
+
+    assert args["Tags"] == [BOTO_TAG]
+
+
+def test_interactive_request_keeps_tags_training_job_create_accepts(model_trainer):
+    """The interactive request carries dicts, not Tag models.
+
+    ``TrainingJob.create`` annotates ``tags`` with ``sagemaker.core.resources.Tag``,
+    which shadows the shape of the same name, so a ``sagemaker.core.shapes.Tag``
+    instance fails its validation. The dicts are revalidated into that model instead,
+    and come out boto-style on the wire.
+    """
+    model_trainer.tags = [TAG]
+
+    args = model_trainer._create_training_job_args()
+
+    assert args["tags"] == [{"key": "project", "value": "ml-project"}]
+    assert _train_wire_tags(model_trainer) == [BOTO_TAG]
+
+
+def test_pipeline_and_interactive_emit_identical_tags(model_trainer):
+    """Both paths must put the same tags on the request."""
+    model_trainer.tags = [TAG]
+
+    pipeline_tags = _pipeline_trainer([TAG])._create_training_job_args()["Tags"]
+
+    assert pipeline_tags == _train_wire_tags(model_trainer)