fix: allow kms_key to be passed for processing step (#2779)

Author: jayatalr

src/sagemaker/workflow/steps.py

@@ -486,6 +486,7 @@ def __init__(
         cache_config: CacheConfig = None,
         depends_on: Union[List[str], List[Step]] = None,
         retry_policies: List[RetryPolicy] = None,
+        kms_key=None,
     ):
         """Construct a ProcessingStep, given a `Processor` instance.
 
@@ -511,6 +512,8 @@ def __init__(
             depends_on (List[str] or List[Step]): A list of step names or step instance
                 this `sagemaker.workflow.steps.ProcessingStep` depends on
             retry_policies (List[RetryPolicy]):  A list of retry policy
+            kms_key (str): The ARN of the KMS key that is used to encrypt the
+                user code file. Defaults to `None`.
         """
         super(ProcessingStep, self).__init__(
             name, StepTypeEnum.PROCESSING, display_name, description, depends_on, retry_policies
@@ -522,6 +525,7 @@ def __init__(
         self.code = code
         self.property_files = property_files
         self.job_name = None
+        self.kms_key = kms_key
 
         # Examine why run method in sagemaker.processing.Processor mutates the processor instance
         # by setting the instance's arguments attribute. Refactor Processor.run, if possible.
@@ -556,8 +560,8 @@ def arguments(self) -> RequestType:
             inputs=self.inputs,
             outputs=self.outputs,
             code=self.code,
+            kms_key=self.kms_key,
         )
-
         process_args = ProcessingJob._get_process_args(
             self.processor, normalized_inputs, normalized_outputs, experiment_config=dict()
         )

tests/unit/sagemaker/workflow/test_steps.py

@@ -617,6 +617,7 @@ def test_processing_step_normalizes_args_with_local_code(mock_normalize_args, sc
         inputs=step.inputs,
         outputs=step.outputs,
         code=step.code,
+        kms_key=None,
     )
 
 
@@ -643,6 +644,7 @@ def test_processing_step_normalizes_args_with_s3_code(mock_normalize_args, scrip
         outputs=outputs,
         job_arguments=["arg1", "arg2"],
         cache_config=cache_config,
+        kms_key="arn:aws:kms:us-west-2:012345678901:key/s3-kms-key",
     )
     mock_normalize_args.return_value = [step.inputs, step.outputs]
     step.to_request()
@@ -652,6 +654,7 @@ def test_processing_step_normalizes_args_with_s3_code(mock_normalize_args, scrip
         inputs=step.inputs,
         outputs=step.outputs,
         code=step.code,
+        kms_key=step.kms_key,
     )
 
 
@@ -686,6 +689,7 @@ def test_processing_step_normalizes_args_with_no_code(mock_normalize_args, scrip
         inputs=step.inputs,
         outputs=step.outputs,
         code=None,
+        kms_key=None,
     )