fix(processing): Fix V3 shape regressions + KMS/requirements gaps in processors

sagemaker-core/src/sagemaker/core/processing.py

@@ -333,7 +333,8 @@ def run(
                 is built with :class:`~sagemaker.workflow.pipeline_context.PipelineSession`.
                 However, the value of `TrialComponentDisplayName` is honored for display in Studio.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
-                user code file (default: None).
+                user code file. If not provided, the processor's configured
+                ``output_kms_key`` is used (default: None).
         Returns:
             None or pipeline step arguments in case the Processor instance is built with
             :class:`~sagemaker.workflow.pipeline_context.PipelineSession`
@@ -344,6 +345,12 @@ def run(
             raise ValueError("""Logs can only be shown if wait is set to True.
                 Please either set wait to True or set logs to False.""")
 
+        # When no explicit code-encryption key is given, fall back to the
+        # configured output KMS key so the uploaded code and job outputs are
+        # encrypted with the same key.
+        if kms_key is None:
+            kms_key = self.output_kms_key
+
         normalized_inputs, normalized_outputs = self._normalize_args(
             job_name=job_name,
             arguments=arguments,
@@ -915,7 +922,8 @@ def run(
                 is built with :class:`~sagemaker.workflow.pipeline_context.PipelineSession`.
                 However, the value of `TrialComponentDisplayName` is honored for display in Studio.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
-                user code file (default: None).
+                user code file. If not provided, the processor's configured
+                ``output_kms_key`` is used (default: None).
         Returns:
             None or pipeline step arguments in case the Processor instance is built with
             :class:`~sagemaker.workflow.pipeline_context.PipelineSession`
@@ -926,7 +934,7 @@ def run(
             inputs=inputs,
             outputs=outputs,
             code=code,
-            kms_key=kms_key,
+            kms_key=kms_key if kms_key is not None else self.output_kms_key,
         )
 
         experiment_config = check_and_get_run_experiment_config(experiment_config)
@@ -1406,6 +1414,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,
@@ -1448,14 +1463,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)
@@ -1491,7 +1514,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(
@@ -1509,6 +1531,7 @@ def _pack_and_upload_code(
             entry_point,
             source_dir,
             install_requirements_dir,
+            requirements=requirements,
         )
 
         return s3_runproc_sh, inputs, job_name
@@ -1554,13 +1577,14 @@ def _create_and_upload_runproc(
         entry_point=None,
         source_dir=None,
         install_requirements_dir=None,
+        requirements=None,
     ):
         """Create runproc shell script and upload to S3 bucket."""
         from sagemaker.core.workflow.utilities import _pipeline_config, hash_object
 
         if _pipeline_config and _pipeline_config.pipeline_name:
             runproc_file_str = self._generate_framework_script(
-                user_script, entry_point, source_dir, install_requirements_dir
+                user_script, entry_point, source_dir, install_requirements_dir, requirements
             )
             runproc_file_hash = hash_object(runproc_file_str)
             s3_uri = s3.s3_path_join(
@@ -1579,7 +1603,7 @@ def _create_and_upload_runproc(
         else:
             s3_runproc_sh = s3.S3Uploader.upload_string_as_file_body(
                 self._generate_framework_script(
-                    user_script, entry_point, source_dir, install_requirements_dir
+                    user_script, entry_point, source_dir, install_requirements_dir, requirements
                 ),
                 desired_s3_uri=entrypoint_s3_uri,
                 kms_key=kms_key,
@@ -1588,20 +1612,36 @@ def _create_and_upload_runproc(
 
         return s3_runproc_sh
 
+    @staticmethod
+    def _requirements_file_in_container(requirements: Optional[str]) -> str:
+        """Return the requirements path as it appears inside the extracted source bundle.
+
+        ``requirements`` is documented as relative to ``source_dir`` and ``_package_code``
+        preserves the directory layout, so a relative path (``reqs/cpu.txt``) is kept. An
+        absolute path cannot be located inside the bundle, so only its basename is used.
+        """
+        if not requirements:
+            return "requirements.txt"
+        if os.path.isabs(requirements):
+            return os.path.basename(requirements)
+        return os.path.normpath(requirements).replace(os.sep, "/")
+
     def _generate_framework_script(
         self,
         user_script: str,
         entry_point: str = None,
         source_dir: str = None,
         install_requirements_dir: str = None,
+        requirements: str = None,
     ) -> str:
         """Generate the framework entrypoint file (as text) for a processing job."""
         if entry_point:
             return self._generate_custom_framework_script(
-                user_script, entry_point, source_dir, install_requirements_dir
+                user_script, entry_point, source_dir, install_requirements_dir, requirements
             )
 
         install_requirements_dir = install_requirements_dir or self._SOURCE_CODE_CONTAINER_DIR
+        requirements_file = self._requirements_file_in_container(requirements)
 
         return dedent("""\
             #!/bin/bash
@@ -1625,16 +1665,17 @@ def _generate_framework_script(
                 exit 1
             fi
 
-            if [[ -f 'requirements.txt' ]]; then
+            if [[ -f '{requirements_file}' ]]; then
                 # Some py3 containers has typing, which may breaks pip install
                 pip uninstall --yes typing
 
-                python3 {install_requirements_dir}/install_requirements.py requirements.txt
+                python3 {install_requirements_dir}/install_requirements.py {requirements_file}
             fi
 
             {entry_point_command} {entry_point} "$@"
         """).format(
             install_requirements_dir=install_requirements_dir,
+            requirements_file=requirements_file,
             entry_point_command=" ".join(self.command),
             entry_point=user_script,
         )
@@ -1645,6 +1686,7 @@ def _generate_custom_framework_script(
         entry_point: str,
         source_dir: str = None,
         install_requirements_dir: str = None,
+        requirements: str = None,
     ) -> str:
         """Generate a custom framework script with a user-provided entrypoint embedded.
 
@@ -1658,6 +1700,8 @@ def _generate_custom_framework_script(
                 is relative, it will be combined with source_dir.
             install_requirements_dir (str): Container directory that holds
                 ``install_requirements.py`` (default: the extracted source code dir).
+            requirements (str): Path to the requirements file relative to source_dir
+                (default: ``requirements.txt``).
 
         Returns:
             str: The generated script content
@@ -1667,6 +1711,7 @@ def _generate_custom_framework_script(
         # source bundle on the container.
         if self._is_s3_uri(source_dir):
             install_requirements_dir = install_requirements_dir or self._SOURCE_CODE_CONTAINER_DIR
+            requirements_file = self._requirements_file_in_container(requirements)
             return dedent("""\
                 #!/bin/bash
 
@@ -1683,9 +1728,9 @@ def _generate_custom_framework_script(
                     exit 1
                 fi
 
-                if [[ -f 'requirements.txt' ]]; then
+                if [[ -f '{requirements_file}' ]]; then
                     pip uninstall --yes typing
-                    python3 {install_requirements_dir}/install_requirements.py requirements.txt
+                    python3 {install_requirements_dir}/install_requirements.py {requirements_file}
                 fi
 
                 # Execute custom entrypoint
@@ -1695,6 +1740,7 @@ def _generate_custom_framework_script(
                 {entry_point_command} {user_script} "$@"
             """).format(
                 install_requirements_dir=install_requirements_dir,
+                requirements_file=requirements_file,
                 entry_point=entry_point,
                 entry_point_command=" ".join(self.command),
                 user_script=user_script,

sagemaker-core/src/sagemaker/core/spark/processing.py

@@ -39,6 +39,7 @@
 from sagemaker.core import s3
 from sagemaker.core.local.image import _ecr_login_if_needed, _pull_image
 from sagemaker.core.processing import ProcessingInput, ProcessingOutput, ScriptProcessor
+from sagemaker.core.shapes import ProcessingS3Input, ProcessingS3Output
 from sagemaker.core.s3 import S3Uploader
 from sagemaker.core.helper.session_helper import Session
 from sagemaker.core.network import NetworkConfig
@@ -330,9 +331,12 @@ def _extend_processing_args(self, inputs, outputs, **kwargs):
             )
 
             output = ProcessingOutput(
-                source=_SparkProcessorBase._spark_event_log_default_local_path,
-                destination=spark_event_logs_s3_uri,
-                s3_upload_mode="Continuous",
+                output_name="spark-event-logs",
+                s3_output=ProcessingS3Output(
+                    s3_uri=spark_event_logs_s3_uri,
+                    local_path=_SparkProcessorBase._spark_event_log_default_local_path,
+                    s3_upload_mode="Continuous",
+                ),
             )
 
             extended_outputs.append(output)
@@ -444,9 +448,13 @@ def _stage_configuration(self, configuration):
         )
 
         conf_input = ProcessingInput(
-            source=s3_uri,
-            destination=f"{self._conf_container_base_path}{self._conf_container_input_name}",
             input_name=_SparkProcessorBase._conf_container_input_name,
+            s3_input=ProcessingS3Input(
+                s3_uri=s3_uri,
+                local_path=f"{self._conf_container_base_path}{self._conf_container_input_name}",
+                s3_data_type="S3Prefix",
+                s3_input_mode="File",
+            ),
         )
         return conf_input
 
@@ -473,6 +481,11 @@ def _stage_submit_deps(self, submit_deps, input_channel_name):
             )
         if not input_channel_name:
             raise ValueError("input_channel_name value may not be empty.")
+        if not isinstance(submit_deps, (list, tuple)):
+            raise ValueError(
+                f"submit_deps must be a list of one or more paths, but got "
+                f"{type(submit_deps).__name__}. {self._submit_deps_error_message}"
+            )
 
         use_input_channel = False
         spark_opt_s3_uris = []
@@ -544,15 +557,20 @@ def _stage_submit_deps(self, submit_deps, input_channel_name):
         # them to the Spark container  and form the spark-submit option from a
         # combination of S3 URIs and container's local input path
         if use_input_channel:
+            input_channel_local_path = f"{self._conf_container_base_path}{input_channel_name}"
             input_channel = ProcessingInput(
-                source=input_channel_s3_uri,
-                destination=f"{self._conf_container_base_path}{input_channel_name}",
                 input_name=input_channel_name,
+                s3_input=ProcessingS3Input(
+                    s3_uri=input_channel_s3_uri,
+                    local_path=input_channel_local_path,
+                    s3_data_type="S3Prefix",
+                    s3_input_mode="File",
+                ),
             )
             spark_opt = (
-                Join(on=",", values=spark_opt_s3_uris + [input_channel.destination])
+                Join(on=",", values=spark_opt_s3_uris + [input_channel_local_path])
                 if spark_opt_s3_uris_has_pipeline_var
-                else ",".join(spark_opt_s3_uris + [input_channel.destination])
+                else ",".join(spark_opt_s3_uris + [input_channel_local_path])
             )
         # If no local files were uploaded, form the spark-submit option from a list of S3 URIs
         else:

sagemaker-core/tests/data/spark/code/python/hello_py_spark/hello_py_spark_app.py

@@ -0,0 +1,55 @@
+# 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.
+import argparse
+import time
+
+from pyspark.sql import SparkSession, SQLContext
+from pyspark.sql.functions import udf
+from pyspark.sql.types import IntegerType
+
+# Import local module to test spark-submit--py-files dependencies
+import hello_py_spark_udfs as udfs
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="inputs and outputs")
+    parser.add_argument("--input", type=str, help="path to input data")
+    parser.add_argument("--output", required=False, type=str, help="path to output data")
+    args = parser.parse_args()
+    spark = SparkSession.builder.appName("SparkContainerTestApp").getOrCreate()
+    sqlContext = SQLContext(spark.sparkContext)
+
+    # Load test data set
+    inputPath = args.input
+    salesDF = spark.read.json(inputPath)
+    salesDF.printSchema()
+
+    salesDF.createOrReplaceTempView("sales")
+    topDF = spark.sql("SELECT date, sale FROM sales WHERE sale > 750")
+    # Show the first 20 rows of the dataframe
+    topDF.show()
+    time.sleep(60)
+
+    # Calculate average sales by date
+    averageSalesPerDay = salesDF.groupBy("date").avg().collect()
+    print(averageSalesPerDay)
+
+    outputPath = args.output
+
+    # Define a UDF that doubles an integer column
+    # The UDF function is imported from local module to test spark-submit--py-files dependencies
+    double_udf_int = udf(udfs.double_x, IntegerType())
+
+    # Save transformed data set to disk
+    salesDF.select("date", "sale", double_udf_int("sale").alias("sale_double")).write.json(
+        outputPath
+    )

sagemaker-core/tests/data/spark/code/python/hello_py_spark/hello_py_spark_udfs.py

@@ -0,0 +1,2 @@
+def double_x(x):
+    return x + x