diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index e1725a3c81..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, @@ -1428,6 +1435,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 +1450,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 +1486,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..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): @@ -1399,6 +1425,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):