fix(#5956): ModelTrainer (V3) does not support output_kms_key for source code uploads

[sagemaker-bot]

Resolves aws#5956.

Draft PR authored by the SageMaker PySDK issue agent. The reproduction, fix, and validation below were produced and checked automatically; a human review is required before this is marked ready.

---

Agent Run Report — issue aws#5956

ModelTrainer (V3) does not support output_kms_key for source code uploads

Outcome: ready_for_pr

1. Triage

• SDK version: v3
• Sub-package: sagemaker-train
• Kind: bug
• Feature area: ModelTrainer source code S3 upload KMS encryption

2. Reproduction

• Tier: tier0_hermetic
• Repro test: sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py

Repro test

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Reproduction test for issue #5956.

ModelTrainer (V3) does not respect a user-provided KMS key when uploading
source code / driver files to S3. In secure environments where S3 bucket
policies enforce server-side encryption with KMS, these uploads (which use
``Session.upload_data`` with no encryption ExtraArgs) are rejected, so
ModelTrainer cannot be used with custom training scripts.

This test configures an ``output_data_config`` KMS key and asserts that the
source-code uploads are performed with that KMS key. On the unfixed code no
KMS key is ever passed to ``upload_data``, so the test fails.
"""
from __future__ import absolute_import

from unittest.mock import patch, MagicMock

from sagemaker.core.helper.session_helper import Session
from sagemaker.train.model_trainer import ModelTrainer
from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE
from sagemaker.train.configs import (
    Compute,
    StoppingCondition,
    OutputDataConfig,
    SourceCode,
)
from tests.unit import DATA_DIR

DEFAULT_IMAGE = "000000000000.dkr.ecr.us-west-2.amazonaws.com/dummy-image:latest"
DEFAULT_BUCKET = "sagemaker-us-west-2-000000000000"
DEFAULT_ROLE = "arn:aws:iam::000000000000:role/test-role"
DEFAULT_BUCKET_PREFIX = "sample-prefix"
DEFAULT_REGION = "us-west-2"
DEFAULT_SOURCE_DIR = f"{DATA_DIR}/script_mode"
KMS_KEY_ID = "arn:aws:kms:us-west-2:000000000000:key/my-kms-key"


def _make_session():
    session = MagicMock(spec=Session)
    session.default_bucket.return_value = DEFAULT_BUCKET
    session.get_caller_identity_arn.return_value = DEFAULT_ROLE
    session.default_bucket_prefix = DEFAULT_BUCKET_PREFIX
    session.boto_session = MagicMock(spec="boto3.session.Session")
    session.boto_region_name = DEFAULT_REGION
    session.upload_data.return_value = f"s3://{DEFAULT_BUCKET}/code"
    return session


@patch("sagemaker.train.defaults.resolve_and_validate_role", return_value=DEFAULT_ROLE)
@patch("sagemaker.train.model_trainer.TrainingJob")
def test_source_code_upload_uses_output_kms_key(mock_training_job, mock_resolve_role):
    """Source-code / driver uploads must use the configured output KMS key.

    Regression for issue #5956: ModelTrainer uploads source code and driver
    files with ``Session.upload_data`` without any KMS encryption, so it cannot
    be used in environments that enforce SSE-KMS on the staging bucket.
    """
    session = _make_session()

    trainer = ModelTrainer(
        training_image=DEFAULT_IMAGE,
        role=DEFAULT_ROLE,
        sagemaker_session=session,
        compute=Compute(instance_type=DEFAULT_INSTANCE_TYPE, instance_count=1),
        stopping_condition=StoppingCondition(max_runtime_in_seconds=3600),
        output_data_config=OutputDataConfig(
            s3_output_path=f"s3://{DEFAULT_BUCKET}/{DEFAULT_BUCKET_PREFIX}/output",
            kms_key_id=KMS_KEY_ID,
        ),
        source_code=SourceCode(
            source_dir=DEFAULT_SOURCE_DIR,
            entry_script="custom_script.py",
        ),
    )

    # Avoid the IAM SimulatePrincipalPolicy round-trip in staging bucket resolution.
    with patch.object(trainer, "_resolve_staging_bucket", return_value=(DEFAULT_BUCKET, None)):
        trainer._create_training_job_args()

    # The source code / sm_drivers channels trigger local -> S3 uploads.
    assert session.upload_data.call_count >= 1, "expected source code to be uploaded to S3"

    # Every upload of source artifacts must carry the KMS key so that buckets
    # enforcing SSE-KMS accept the object.
    for call in session.upload_data.call_args_list:
        kwargs = call.kwargs
        extra_args = kwargs.get("extra_args") or (
            call.args[4] if len(call.args) > 4 else None
        )
        assert extra_args is not None, (
            "ModelTrainer uploaded source code without encryption ExtraArgs; "
            "the configured output KMS key is ignored (issue #5956)."
        )
        assert extra_args.get("SSEKMSKeyId") == KMS_KEY_ID
        assert extra_args.get("ServerSideEncryption") == "aws:kms"

Failing output (before fix)

F                                                                        [100%]
=================================== FAILURES ===================================
_________________ test_source_code_upload_uses_output_kms_key __________________

mock_training_job = <MagicMock name='TrainingJob' id='140116692150912'>
mock_resolve_role = <MagicMock name='resolve_and_validate_role' id='140116684286112'>

    @patch("sagemaker.train.defaults.resolve_and_validate_role", return_value=DEFAULT_ROLE)
    @patch("sagemaker.train.model_trainer.TrainingJob")
    def test_source_code_upload_uses_output_kms_key(mock_training_job, mock_resolve_role):
        """Source-code / driver uploads must use the configured output KMS key.
    
        Regression for issue #5956: ModelTrainer uploads source code and driver
        files with ``Session.upload_data`` without any KMS encryption, so it cannot
        be used in environments that enforce SSE-KMS on the staging bucket.
        """
        session = _make_session()
    
        trainer = ModelTrainer(
            training_image=DEFAULT_IMAGE,
            role=DEFAULT_ROLE,
            sagemaker_session=session,
            compute=Compute(instance_type=DEFAULT_INSTANCE_TYPE, instance_count=1),
            stopping_condition=StoppingCondition(max_runtime_in_seconds=3600),
            output_data_config=OutputDataConfig(
                s3_output_path=f"s3://{DEFAULT_BUCKET}/{DEFAULT_BUCKET_PREFIX}/output",
                kms_key_id=KMS_KEY_ID,
            ),
            source_code=SourceCode(
                source_dir=DEFAULT_SOURCE_DIR,
                entry_script="custom_script.py",
            ),
        )
    
        # Avoid the IAM SimulatePrincipalPolicy round-trip in staging bucket resolution.
        with patch.object(trainer, "_resolve_staging_bucket", return_value=(DEFAULT_BUCKET, None)):
            trainer._create_training_job_args()
    
        # The source code / sm_drivers channels trigger local -> S3 uploads.
        assert session.upload_data.call_count >= 1, "expected source code to be uploaded to S3"
    
        # Every upload of source artifacts must carry the KMS key so that buckets
        # enforcing SSE-KMS accept the object.
        for call in session.upload_data.call_args_list:
            kwargs = call.kwargs
            extra_args = kwargs.get("extra_args") or (
                call.args[4] if len(call.args) > 4 else None
            )
>           assert extra_args is not None, (
                "ModelTrainer uploaded source code without encryption ExtraArgs; "
                "the configured output KMS key is ignored (issue #5956)."
            )
E           AssertionError: ModelTrainer uploaded source code without encryption ExtraArgs; the configured output KMS key is ignored (issue #5956).
E           assert None is not None

sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py:101: AssertionError
=============================== warnings summary ===============================
sagemaker-train/src/sagemaker/train/model_trainer.py:52
  /tmp/agent-sbi_8ug7/issue-5956/sagemaker-train/src/sagemaker/train/model_trainer.py:52: DeprecationWarning: sagemaker.train.configs has been moved to sagemaker.core.training.configs. Please update your imports. This shim will be removed in a future version.
    from sagemaker.train import configs

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=========================== short test summary info ============================
FAILED sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py::test_source_code_upload_uses_output_kms_key
1 failed, 1 warning in 2.93s

3. Fix

Fixed issue aws#5956: ModelTrainer (V3) now respects the configured output KMS key when uploading source code / driver files to S3. Previously create_input_data_channel called Session.upload_data with no encryption ExtraArgs, so in secure environments that enforce SSE-KMS on the staging bucket the source-code and driver uploads were rejected and ModelTrainer could not be used with custom training scripts. I added a _upload_extra_args() helper on ModelTrainer that derives {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": <key>} from output_data_config.kms_key_id (returning None when no key is configured, preserving prior behavior), and threaded the resulting extra_args through both upload_data calls in create_input_data_channel. Because this adds an extra_args keyword to the upload call, I updated the four directly-affected local mock_upload_data stub signatures in test_model_trainer.py to accept the new optional extra_args parameter. The reproduction test and the full model_trainer unit suite pass.

Changed files:

• sagemaker-train/src/sagemaker/train/model_trainer.py
• sagemaker-train/tests/unit/train/test_model_trainer.py
• sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py

Diff

diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py
index 241b6ac..34a2016 100644
--- a/sagemaker-train/src/sagemaker/train/model_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/model_trainer.py
@@ -863,6 +863,22 @@ class ModelTrainer(BaseModel):
         if self._temp_code_dir is not None:
             self._temp_code_dir.cleanup()
 
+    def _upload_extra_args(self) -> Optional[Dict[str, str]]:
+        """Build S3 ``ExtraArgs`` for source-code uploads.
+
+        When an output KMS key is configured (via ``output_data_config.kms_key_id``),
+        source-code and driver uploads must be encrypted with that key so that buckets
+        enforcing SSE-KMS accept the objects. Returns ``None`` when no KMS key is set,
+        preserving the default (unencrypted-arg) upload behavior.
+        """
+        kms_key_id = getattr(self.output_data_config, "kms_key_id", None)
+        if not kms_key_id:
+            return None
+        return {
+            "ServerSideEncryption": "aws:kms",
+            "SSEKMSKeyId": kms_key_id,
+        }
+
     def _resolve_staging_bucket(self) -> tuple[str,str]:
         """Resolve the S3 bucket and key prefix for staging training artifacts.
 
@@ -1042,6 +1058,9 @@ class ModelTrainer(BaseModel):
                     # Resolve staging bucket based on training role permissions
                     staging_bucket, staging_prefix = self._resolve_staging_bucket()
                     effective_prefix = f"{staging_prefix}/{key_prefix}" if staging_prefix else key_prefix
+                    # Apply the configured output KMS key to source-code uploads so that
+                    # staging buckets that enforce SSE-KMS accept the objects.
+                    extra_args = self._upload_extra_args()
                     if ignore_patterns and _is_valid_path(data_source, path_type="Directory"):
                         tmp_dir = TemporaryDirectory()
                         copied_path = os.path.join(
@@ -1057,12 +1076,14 @@ class ModelTrainer(BaseModel):
                             path=copied_path,
                             bucket=staging_bucket,
                             key_prefix=effective_prefix,
+                            extra_args=extra_args,
                         )
                     else:
                         s3_uri = self.sagemaker_session.upload_data(
                             path=data_source,
                             bucket=staging_bucket,
                             key_prefix=effective_prefix,
+                            extra_args=extra_args,
                         )
                     channel = Channel(
                         channel_name=channel_name,
diff --git a/sagemaker-train/tests/unit/train/test_model_trainer.py b/sagemaker-train/tests/unit/train/test_model_trainer.py
index ce5d208..d2f67c1 100644
--- a/sagemaker-train/tests/unit/train/test_model_trainer.py
+++ b/sagemaker-train/tests/unit/train/test_model_trainer.py
@@ -875,7 +875,7 @@ def test_remote_debug_config(mock_training_job, modules_session):
 @patch("sagemaker.train.model_trainer._get_unique_name")
 @patch("sagemaker.train.model_trainer.TrainingJob")
 def test_model_trainer_full_init(mock_training_job, mock_unique_name, modules_session):
-    def mock_upload_data(path, bucket, key_prefix):
+    def mock_upload_data(path, bucket, key_prefix, extra_args=None):
         return f"s3://{bucket}/{key_prefix}"
 
     modules_session.upload_data.side_effect = mock_upload_data
@@ -1407,7 +1407,7 @@ def test_hyperparameters_invalid(mock_exists, modules_session):
 @patch("sagemaker.train.model_trainer._get_unique_name")
 @patch("sagemaker.train.model_trainer.TrainingJob")
 def test_model_trainer_default_paths(mock_training_job, mock_unique_name, modules_session):
-    def mock_upload_data(path, bucket, key_prefix):
+    def mock_upload_data(path, bucket, key_prefix, extra_args=None):
         return f"s3://{bucket}/{key_prefix}"
 
     unique_name = "base-job-0123456789"
@@ -1510,7 +1510,7 @@ def test_metric_definitions(mock_training_job, modules_session):
 @patch("sagemaker.train.model_trainer._get_unique_name")
 @patch("sagemaker.core.resources.TrainingJob")
 def test_nova_recipe(mock_training_job, mock_unique_name, modules_session):
-    def mock_upload_data(path, bucket, key_prefix):
+    def mock_upload_data(path, bucket, key_prefix, extra_args=None):
         if os.path.isfile(path):
             file_name = os.path.basename(path)
             return f"s3://{bucket}/{key_prefix}/{file_name}"
@@ -1773,7 +1773,7 @@ def test_nova_recipe_model_package_config_only_mpg_from_recipe(modules_session):
 @patch("sagemaker.train.model_trainer._get_unique_name")
 @patch("sagemaker.train.model_trainer.TrainingJob")
 def test_llmft_recipe(mock_training_job, mock_unique_name, modules_session):
-    def mock_upload_data(path, bucket, key_prefix):
+    def mock_upload_data(path, bucket, key_prefix, extra_args=None):
         if os.path.isfile(path):
             file_name = os.path.basename(path)
             return f"s3://{bucket}/{key_prefix}/{file_name}"
diff --git a/sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py b/sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py
new file mode 100644
index 0000000..db306f4
--- /dev/null
+++ b/sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py
@@ -0,0 +1,106 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitatio
... [truncated]

4. Validation

• Repro test passes after fix: yes
• Pre-existing failures ignored (baseline): 4
• New failures introduced: none

Post-fix test output

.                                                                        [100%]
=============================== warnings summary ===============================
sagemaker-train/src/sagemaker/train/model_trainer.py:52
  /tmp/agent-sbi_8ug7/issue-5956/sagemaker-train/src/sagemaker/train/model_trainer.py:52: DeprecationWarning: sagemaker.train.configs has been moved to sagemaker.core.training.configs. Please update your imports. This shim will be removed in a future version.
    from sagemaker.train import configs

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
1 passed, 1 warning in 2.83s

5. Review

• Reviewer approved: yes
• Rounds: 1 · confidence: 0.85

5b. Adversarial verification

• Confirmed (symptom resolved, no regression): yes
• Notes: Verified the fix resolves the reported symptom (ModelTrainer V3 uploading source code/driver files to S3 without KMS encryption). The fix adds ModelTrainer._upload_extra_args() which builds {ServerSideEncryption: aws:kms, SSEKMSKeyId: } when output_data_config.kms_key_id is set, and passes it as extra_args to both Session.upload_data call sites in create_input_data_channel. Confirmed Session.upload_data forwards extra_args to s3.Object.upload_file(ExtraArgs=...) — so the KMS encryption is genuinely applied to the real S3 PUT (the user-observable outcome), not just an internal flag. The Unassigned() default of kms_key_id has bool=False so if not kms_key_id correctly treats unset as no-KMS; None output_data_config is also handled via getattr default.

I wrote independent probe tests (not the author's) that: (1) build the full training request end-to-end via _create_training_job_args with a KMS key and assert EVERY upload carries SSEKMSKeyId/aws:kms AND that the resulting request actually contains the 'code' and 'sm_drivers' channels; (2) assert the no-KMS path still uploads with extra_args=None (regression guard); (3) probe the None output_data_config path; (4) probe the S3-URI source_dir case (code channel points at the given S3 URI; sm_drivers still uploaded with KMS extra_args). I confirmed the KMS probes FAIL deterministically on the unfixed code (temporarily reverted extra_args to None -> 2 failures) and PASS with the fix, while the no-KMS probe passes on both. Ran the full sagemaker-train/tests/unit/train suite: 2397 passed, 3 pre-existing failures unrelated to this fix (real-AWS AccessDeniedException on DescribeHubContent and 403 on HeadBucket — network/credential-dependent, no relation to source-code upload or KMS). The 4 author-edited mock_upload_data signatures now accept extra_args=None consistent with the new call convention; other tests patch upload_data as a plain MagicMock which accepts any kwargs. Removed my scratch probe file contents (emptied) as no delete tool is available.

6. Routing

• Owning team: ?
• CTI: ?
• Discovery candidates: sagemaker-train/src/sagemaker/train/model_trainer.py

Timeline

• triage: v3/sagemaker-train/bug area='ModelTrainer source code S3 upload KMS encryption'
• workspace acquired: /tmp/agent-sbi_8ug7/issue-5956 (install_all=True)
• discovery: candidate files ['sagemaker-train/src/sagemaker/train/model_trainer.py']
• reproduced (tier0_hermetic) via sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py
• note: 4 pre-existing failure(s) in baseline (ignored)
• fix authored (round 1, turns<=40); shipping: sagemaker-train/src/sagemaker/train/model_trainer.py, sagemaker-train/tests/unit/train/test_model_trainer.py, sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py
• validated (round 1): repro passes + no new failures
• review (round 1): APPROVED
• verification: confirmed=True regressions=[] — Verified the fix resolves the reported symptom (ModelTrainer V3 uploading source code/driver files to S3 without KMS encryption). The fix adds ModelTrainer._upl
• routed to SageMaker Codex (AWS/SageMaker/SageMaker Python SDK)

Agent actions (68 tool calls)

Tool-by-tool transcript

1. search {'pattern': 'def _create_input_data_channel|S3Uploader|upload|kms_key|def create_data_channels', 'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py'} → (no matches)
2. search {'pattern': 'upload|kms|SourceCode|_prepare_train', 'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py'} → (no matches)
3. list_dir {'path': 'sagemaker-train/src/sagemaker/train'} → init.py agent_rft_job.py aws_batch base_trainer.py common.py common_utils configs.py constants.py container_drivers cpt_trainer.py custom_agent_lambda.py da
4. search {'pattern': 'upload|kms|SourceCode|def _prepare|S3DataSource|create_data_channel', 'path': 'sagemaker-train/src/sagemaker/train/base_trainer.py'} → (no matches)
5. search {'pattern': 'upload|kms|def create_data_channel', 'path': 'sagemaker-train/src/sagemaker/train'} → sagemaker-train/src/sagemaker/train/rlaif_trainer.py:113: kms_key_id (Optional[str]): sagemaker-train/src/sagemaker/train/rlaif_trainer.py:149:
6. read_file {'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
7. search {'pattern': 'def upload_data', 'path': '.'} → sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py:242: def upload_dataframe_to_s3( sagemaker-core/src/sagemaker/core/helper/session_helper.py:4
8. read_file {'path': 'sagemaker-core/src/sagemaker/core/helper/session_helper.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
9. list_dir {'path': 'sagemaker-train/tests/unit'} → init.py ai_registry train
10. search {'pattern': 'upload_data|create_input_data_channel|def test', 'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → (no matches)
11. list_dir {'path': 'sagemaker-train/tests/unit/train'} → init.py aws_batch common_utils container_drivers evaluate local remote_function sm_recipes test_agent_rft_job.py test_base_trainer_compute.py test_base_trai
12. search {'pattern': 'create_input_data_channel|upload_data|source_code|def test_create_input', 'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → (no matches)
13. read_file {'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
14. search {'pattern': 'DATA_DIR', 'path': 'sagemaker-train/tests/unit/init.py'} → (no matches)
15. read_file {'path': 'sagemaker-train/tests/unit/init.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
16. write_file {'path': 'sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py', 'bytes': 4485} → wrote 4485 bytes to sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py
17. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py'} → exit=1 F [100%] =================================== FAILURES ============================
18. edit_file {'path': 'sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py', 'old_bytes': 491, 'new_bytes': 598} → edited sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py (replaced 491→598 bytes)
19. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py'} → exit=1 F [100%] =================================== FAILURES ============================
20. read_file {'path': 'sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
21. read_file {'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
22. search {'pattern': 'def upload_data', 'path': 'sagemaker-core/src'} → sagemaker-core/src/sagemaker/core/helper/session_helper.py:400: def upload_data(self, path, bucket=None, key_prefix="data", callback=None, extra_args=None):
23. read_file {'path': 'sagemaker-core/src/sagemaker/core/helper/session_helper.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
24. edit_file {'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py', 'old_bytes': 1385, 'new_bytes': 1714} → edited sagemaker-train/src/sagemaker/train/model_trainer.py (replaced 1385→1714 bytes)
25. edit_file {'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py', 'old_bytes': 56, 'new_bytes': 755} → edited sagemaker-train/src/sagemaker/train/model_trainer.py (replaced 56→755 bytes)
26. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py'} → exit=0 . [100%] =============================== warnings summary ========================
27. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → exit=1 ......................................F.........F......F........ [100%] =================================== FAILURES ============================
28. search {'pattern': 'def mock_upload_data\(path, bucket, key_prefix\)', 'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → (no matches)
29. search {'pattern': 'def mock_upload_data', 'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → (no matches)
30. search {'pattern': 'mock_upload_data', 'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → (no matches)
31. search {'pattern': 'def mock_upload_data', 'path': 'sagemaker-train/tests/unit/train'} → sagemaker-train/tests/unit/train/test_model_trainer.py:878: def mock_upload_data(path, bucket, key_prefix): sagemaker-train/tests/unit/train/test_model_trai
32. read_file {'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
33. edit_file {'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py', 'old_bytes': 184, 'new_bytes': 201} → edited sagemaker-train/tests/unit/train/test_model_trainer.py (replaced 184→201 bytes)
34. edit_file {'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py', 'old_bytes': 188, 'new_bytes': 205} → edited sagemaker-train/tests/unit/train/test_model_trainer.py (replaced 188→205 bytes)
35. edit_file {'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py', 'old_bytes': 160, 'new_bytes': 177} → edited sagemaker-train/tests/unit/train/test_model_trainer.py (replaced 160→177 bytes)
36. edit_file {'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py', 'old_bytes': 161, 'new_bytes': 178} → edited sagemaker-train/tests/unit/train/test_model_trainer.py (replaced 161→178 bytes)
37. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_model_trainer.py sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py'} → exit=0 ................................................................. [100%] =============================== warnings summary ========================
38. search {'pattern': 'def create_input_data_channel', 'path': '.'} → sagemaker-train/src/sagemaker/train/model_trainer.py:970: def create_input_data_channel(
39. read_file {'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
40. search {'pattern': 'def upload_data', 'path': '.'} → sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py:242: def upload_dataframe_to_s3( sagemaker-core/src/sagemaker/core/helper/session_helper.py:4
41. search {'pattern': 'kms_key_id', 'path': 'sagemaker-train/src/sagemaker/train/configs.py'} → (no matches)
42. search {'pattern': 'class OutputDataConfig', 'path': '.'} → sagemaker-core/src/sagemaker/core/training/configs.py:128: class OutputDataConfig(shapes.OutputDataConfig): sagemaker-core/src/sagemaker/core/training/configs.p
43. read_file {'path': 'sagemaker-core/src/sagemaker/core/helper/session_helper.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
44. read_file {'path': 'sagemaker-core/src/sagemaker/core/shapes/shapes.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
45. search {'pattern': 'class Unassigned|def bool|eq', 'path': 'sagemaker-core/src/sagemaker/core/utils/utils.py'} → (no matches)
46. search {'pattern': 'class Unassigned', 'path': '.'} → sagemaker-core/src/sagemaker/core/utils/utils.py:291: class Unassigned:
47. read_file {'path': 'sagemaker-core/src/sagemaker/core/utils/utils.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
48. search {'pattern': 'DEFAULT_SOURCE_DIR|script_mode|custom_script', 'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → (no matches)
49. search {'pattern': 'DATA_DIR|source_dir|SourceCode\(', 'path': 'sagemaker-train/tests/unit/train/test_model_trainer.py'} → (no matches)
50. search {'pattern': 'DATA_DIR', 'path': 'sagemaker-train/tests/unit/init.py'} → (no matches)
51. list_dir {'path': 'sagemaker-train/tests/data/script_mode'} → init.py code.tar.gz custom_script.py data pytorch_model_def.py requirements.txt
52. read_file {'path': 'sagemaker-train/tests/unit/init.py'} → # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this
53. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py -v'} → exit=0 ============================= test session starts ============================== platform linux -- Python 3.12.14, pytest-8.4.2, pluggy-1.6.0 rootdir: /t
54. write_file {'path': 'sagemaker-train/tests/unit/train/test_probe_kms_verify.py', 'bytes': 6124} → wrote 6124 bytes to sagemaker-train/tests/unit/train/test_probe_kms_verify.py
55. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_probe_kms_verify.py -v'} → exit=0 ============================= test session starts ============================== platform linux -- Python 3.12.14, pytest-8.4.2, pluggy-1.6.0 rootdir: /t
56. edit_file {'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py', 'old_bytes': 321, 'new_bytes': 316} → edited sagemaker-train/src/sagemaker/train/model_trainer.py (replaced 321→316 bytes)
57. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_probe_kms_verify.py -v'} → exit=1 ============================= test session starts ============================== platform linux -- Python 3.12.14, pytest-8.4.2, pluggy-1.6.0 rootdir: /t
58. edit_file {'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py', 'old_bytes': 148, 'new_bytes': 153} → edited sagemaker-train/src/sagemaker/train/model_trainer.py (replaced 148→153 bytes)
59. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_probe_kms_verify.py sagemaker-train/tests/unit/train/test_model_trainer.py'} → exit=0 .................................................................... [100%] =============================== warnings summary ========================
60. search {'pattern': 'def mock_upload_data\(path, bucket, key_prefix\)', 'path': '.'} → sagemaker-train/tests/unit/train/test_model_trainer.py:1156: # def mock_upload_data(path, bucket, key_prefix):
61. search {'pattern': 'upload_data\.side_effect|def mock_upload_data|upload_data =|\.upload_data\b', 'path': 'sagemaker-train/tests'} → sagemaker-train/tests/integ/train/test_local_model_trainer.py:110: session.upload_data( sagemaker-train/tests/integ/train/test_local_model_trainer.p
62. run_pytest {'selection': 'sagemaker-train/tests/unit/train'} → exit=1 ........................................................................ [ 2%] ........................................................................
63. read_file {'path': 'sagemaker-train/tests/unit/train/test_probe_kms_verify.py'} → # Probe test for issue ModelTrainer (V3) does not support output_kms_key for source code uploads aws/sagemaker-python-sdk#5956 verification. from future import absolute_import from unittest.mock import patch, MagicMock import pytest from sagemaker.cor
64. write_file {'path': 'sagemaker-train/tests/unit/train/test_probe_kms_verify.py', 'bytes': 0} → wrote 0 bytes to sagemaker-train/tests/unit/train/test_probe_kms_verify.py
65. run_pytest {'selection': 'sagemaker-train/tests/unit/train/test_model_trainer_source_code_kms.py sagemaker-train/tests/unit/train/test_model_trainer.py'} → exit=0 ................................................................. [100%] =============================== warnings summary ========================
66. search {'pattern': 'extra_args = self._upload_extra_args\(\)|def _upload_extra_args', 'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py'} → (no matches)
67. search {'pattern': '_upload_extra_args', 'path': 'sagemaker-train/src/sagemaker/train/model_trainer.py'} → (no matches)
68. search {'pattern': 'upload_extra_args', 'path': '.'} → sagemaker-train/src/sagemaker/train/model_trainer.py:866: def _upload_extra_args(self) -> Optional[Dict[str, str]]: sagemaker-train/src/sagemaker/train/mode