From 4bc18d06e0991c3a20c7e6163369eb213e30b100 Mon Sep 17 00:00:00 2001
From: Akshay Shah <akshayshah0803@gmail.com>
Date: Fri, 9 Oct 2026 13:06:42 -0500
Subject: [PATCH] fix(train): send Key/Value tag casing in ModelTrainer's boto3
 payload

_create_training_job_args converts Tag objects to snake_case dicts
({'key', 'value'}) for TrainingJob.create. For the boto3/PipelineSession
payload the request is passed through serialize(), which leaves plain dict
keys unchanged, so the wire payload carried 'Tags': [{'key': ..., 'value': ...}].
TrainingQueue.submit then failed with KeyError: 'Key' when merging those tags
into the Batch job tags, and the pipeline TrainingStep arguments carried the
wrong casing as well.

Rebuild Tag objects in that payload so serialize() emits Key/Value.
---
 .../src/sagemaker/train/model_trainer.py      |  7 ++++++
 .../tests/unit/train/test_model_trainer.py    | 23 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py
index 8816ea68f9..dd11aa9915 100644
--- a/sagemaker-train/src/sagemaker/train/model_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/model_trainer.py
@@ -884,6 +884,13 @@ def _create_training_job_args(
                 training_request.pop("training_job_name", None)
             # Convert snake_case to PascalCase for AWS API
             pipeline_request = {to_pascal_case(k): v for k, v in training_request.items()}
+            # serialize() leaves plain dict keys as-is, so turn the snake_case tag dicts
+            # back into Tag objects to get the API's Key/Value casing.
+            if pipeline_request.get("Tags"):
+                pipeline_request["Tags"] = [
+                    Tag(key=tag["key"], value=tag["value"]) if "key" in tag else tag
+                    for tag in pipeline_request["Tags"]
+                ]
             serialized_request = serialize(pipeline_request)
             return serialized_request
 
diff --git a/sagemaker-train/tests/unit/train/test_model_trainer.py b/sagemaker-train/tests/unit/train/test_model_trainer.py
index f990b7ab88..2a0d4afeb0 100644
--- a/sagemaker-train/tests/unit/train/test_model_trainer.py
+++ b/sagemaker-train/tests/unit/train/test_model_trainer.py
@@ -2559,3 +2559,26 @@ 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"
     )
+
+
+def test_create_training_job_args_boto3_tags_use_api_casing():
+    """The boto3/pipeline payload must use the API's Key/Value tag casing."""
+    from sagemaker.train.aws_batch.batch_api_helper import _submit_service_job
+
+    trainer = ModelTrainer(
+        training_image=DEFAULT_IMAGE,
+        role=DEFAULT_ROLE,
+        compute=DEFAULT_COMPUTE_CONFIG,
+        stopping_condition=DEFAULT_STOPPING_CONDITION,
+        output_data_config=DEFAULT_OUTPUT_DATA_CONFIG,
+        tags=[Tag(key="team", value="ml")],
+    )
+
+    args = trainer._create_training_job_args(boto3=True)
+
+    assert args["Tags"] == [{"Key": "team", "Value": "ml"}]
+
+    # TrainingQueue.submit merges these tags into the Batch job tags.
+    with patch("sagemaker.train.aws_batch.batch_api_helper.get_batch_boto_client") as mock_client:
+        _submit_service_job(args, "job-name", "queue-name")
+    assert mock_client.return_value.submit_service_job.call_args.kwargs["tags"] == {"team": "ml"}