From 89860467fa0cd741599e5bf6707c3c88f8642fc0 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Tue, 29 Sep 2026 16:57:18 -0700 Subject: [PATCH 1/2] fix(processing): encrypt sourcedir.tar.gz with resolved kms_key FrameworkProcessor._pack_and_upload_code computed the output_kms_key fallback (evaluated_kms_key) only after calling _package_code, and passed the raw kms_key to _package_code. FrameworkProcessor.run does not apply the fallback before delegating here, so with output_kms_key set and no explicit kms_key, sourcedir.tar.gz -- the largest, most sensitive artifact -- was uploaded unencrypted, while install_requirements.py and runproc.sh were encrypted. This left #4874 unfixed for FrameworkProcessor. Resolve the key once, before packaging, and pass it into _package_code so all three artifacts share the same key. An explicitly passed kms_key still wins; behaviour when both are unset is unchanged. Add unit tests asserting _package_code and the helper-script uploads all receive the resolved key on the fallback path (regression for #4874) and that an explicit key overrides output_kms_key. Relates-to #4874 --- .../src/sagemaker/core/processing.py | 11 ++- sagemaker-core/tests/unit/test_processing.py | 85 +++++++++++++++++++ 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index e1725a3c81..0abdbf89ea 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -1428,6 +1428,14 @@ def _pack_and_upload_code( if job_name is None: job_name = self._generate_current_job_name(job_name) + # Resolve the KMS key once, before any upload, so every artifact this method + # produces -- the source bundle, install_requirements.py and runproc.sh -- is + # encrypted with the same key. An explicitly passed kms_key always wins; when it + # is not set we fall back to the configured output_kms_key. This must happen + # before _package_code so sourcedir.tar.gz (the largest, most sensitive artifact) + # is not left unencrypted (issue #4874). + evaluated_kms_key = kms_key if kms_key else self.output_kms_key + # Package and upload code s3_payload = self._package_code( entry_point=code, @@ -1435,7 +1443,7 @@ def _pack_and_upload_code( dependencies=dependencies, requirements=requirements, job_name=job_name, - kms_key=kms_key, + kms_key=evaluated_kms_key, ) inputs = self._patch_inputs_with_payload(inputs, s3_payload) @@ -1471,7 +1479,6 @@ def _pack_and_upload_code( # Upload install_requirements helper import sagemaker.core.utils.install_requirements as _ir_mod - evaluated_kms_key = kms_key if kms_key else self.output_kms_key with open(_ir_mod.__file__, "r") as _ir_file: _ir_body = _ir_file.read() s3.S3Uploader.upload_string_as_file_body( diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index d54d2c4e81..525fca0a12 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -1399,6 +1399,91 @@ def test_pack_and_upload_code_with_local_file(self, mock_session): assert any("runproc.sh" in uri for uri in upload_uris) assert mock_upload.call_count == 2 + def test_pack_and_upload_code_falls_back_to_output_kms_key(self, mock_session): + """kms_key falls back to output_kms_key for ALL artifacts, incl. sourcedir.tar.gz. + + Regression test for #4874: the source bundle is packaged and uploaded by + ``_package_code`` before the helper scripts. Previously the fallback was + computed after that call, so ``sourcedir.tar.gz`` (the largest, most sensitive + artifact) was uploaded unencrypted while install_requirements.py / runproc.sh + were encrypted. All three must use the resolved key. + """ + kms_key_arn = "arn:aws:kms:us-west-2:123456789012:key/output-key" + processor = FrameworkProcessor( + role="arn:aws:iam::123456789012:role/SageMakerRole", + image_uri="test-image:latest", + instance_count=1, + instance_type="ml.m5.xlarge", + sagemaker_session=mock_session, + output_kms_key=kms_key_arn, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + entry_point = os.path.join(tmpdir, "train.py") + with open(entry_point, "w") as f: + f.write("print('training')") + + with patch.object( + processor, "_package_code", return_value="s3://bucket/code/sourcedir.tar.gz" + ) as mock_package_code: + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://bucket/runproc.sh", + ) as mock_upload: + processor._pack_and_upload_code( + code=entry_point, + source_dir=None, + requirements=None, + job_name=None, + inputs=None, + kms_key=None, + ) + + # The source bundle upload (via _package_code) must receive the resolved key. + assert mock_package_code.call_args.kwargs["kms_key"] == kms_key_arn + # install_requirements.py and runproc.sh must use the same resolved key. + upload_kms = [call.kwargs.get("kms_key") for call in mock_upload.call_args_list] + assert upload_kms, "expected helper-script uploads" + assert all(k == kms_key_arn for k in upload_kms) + + def test_pack_and_upload_code_explicit_kms_key_wins(self, mock_session): + """An explicitly passed kms_key overrides output_kms_key for every artifact.""" + output_kms = "arn:aws:kms:us-west-2:123456789012:key/output-key" + explicit_kms = "arn:aws:kms:us-west-2:123456789012:key/explicit-key" + processor = FrameworkProcessor( + role="arn:aws:iam::123456789012:role/SageMakerRole", + image_uri="test-image:latest", + instance_count=1, + instance_type="ml.m5.xlarge", + sagemaker_session=mock_session, + output_kms_key=output_kms, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + entry_point = os.path.join(tmpdir, "train.py") + with open(entry_point, "w") as f: + f.write("print('training')") + + with patch.object( + processor, "_package_code", return_value="s3://bucket/code/sourcedir.tar.gz" + ) as mock_package_code: + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://bucket/runproc.sh", + ) as mock_upload: + processor._pack_and_upload_code( + code=entry_point, + source_dir=None, + requirements=None, + job_name=None, + inputs=None, + kms_key=explicit_kms, + ) + + assert mock_package_code.call_args.kwargs["kms_key"] == explicit_kms + upload_kms = [call.kwargs.get("kms_key") for call in mock_upload.call_args_list] + assert upload_kms and all(k == explicit_kms for k in upload_kms) + class TestProcessingInputOutputHelpers: def test_processing_input_with_app_managed(self): From 6ec5c6a61796f6c4dd03f2b70a858abee343310c Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Tue, 29 Sep 2026 17:03:29 -0700 Subject: [PATCH 2/2] fix(processing): reject pipeline variable for requirements argument FrameworkProcessor.run bakes the requirements path into the generated runproc.sh script, so it cannot be a pipeline variable. Passing a ParameterString previously surfaced as an opaque "TypeError: expected str, bytes or os.PathLike object, not ParameterString" from os.path.isabs deep in script generation. Add an explicit is_pipeline_variable guard at the top of FrameworkProcessor.run, raising a clear ValueError, matching the existing code and submit_app guards. Add a unit test. Relates-to #5805 --- .../src/sagemaker/core/processing.py | 7 +++++ sagemaker-core/tests/unit/test_processing.py | 26 +++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index 0abdbf89ea..a6eb1248b1 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -1386,6 +1386,13 @@ def run( None or pipeline step arguments in case the Processor instance is built with :class:`~sagemaker.workflow.pipeline_context.PipelineSession` """ + if requirements and is_pipeline_variable(requirements): + raise ValueError( + "requirements argument has to be a path relative to source_dir " + + "rather than a pipeline variable, because it is baked into the " + + "generated runproc.sh script" + ) + s3_runproc_sh, inputs, job_name = self._pack_and_upload_code( code, source_dir, diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index 525fca0a12..f294d0d634 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -1335,6 +1335,32 @@ def test_run_with_local_code(self, mock_session): processor.run(code=entry_point, wait=False) assert processor.latest_job == mock_job + def test_run_rejects_pipeline_variable_requirements(self, mock_session): + """requirements must not be a pipeline variable; it is baked into runproc.sh. + + Previously a pipeline variable raised an opaque TypeError out of + os.path.isabs; the guard now raises a clear ValueError like the code + and submit_app guards. + """ + processor = FrameworkProcessor( + role="arn:aws:iam::123456789012:role/SageMakerRole", + image_uri="test-image:latest", + instance_count=1, + instance_type="ml.m5.xlarge", + sagemaker_session=mock_session, + ) + + with patch( + "sagemaker.core.processing.is_pipeline_variable", + side_effect=lambda v: v == "reqs.txt", + ): + with pytest.raises(ValueError, match="requirements argument"): + processor.run( + code="s3://bucket/train.py", + requirements="reqs.txt", + wait=False, + ) + class TestFrameworkProcessorPackAndUpload: def test_pack_and_upload_code_with_s3_uri(self, mock_session):