fix(processing): encrypt sourcedir.tar.gz + reject pipeline-variable requirements

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,14 +1435,22 @@ 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,
             source_dir=source_dir,
             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(

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):