fix: allow kms_key to be passed for processing step

Author: jayatalr

src/sagemaker/workflow/steps.py

@@ -461,6 +461,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.
 
@@ -486,6 +487,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
@@ -496,6 +499,7 @@ def __init__(
         self.job_arguments = job_arguments
         self.code = code
         self.property_files = property_files
+        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.
@@ -518,8 +522,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

@@ -516,6 +516,7 @@ def test_processing_step_normalizes_args(mock_normalize_args, sagemaker_session)
         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()
@@ -524,6 +525,7 @@ def test_processing_step_normalizes_args(mock_normalize_args, sagemaker_session)
         inputs=step.inputs,
         outputs=step.outputs,
         code=step.code,
+        kms_key=step.kms_key,
     )