move kms_key to end of arg list

Author: Chuyang Deng

src/sagemaker/processing.py

@@ -113,22 +113,20 @@ def __init__(
     def run(
         self,
         inputs=None,
-        kms_key=None,
         outputs=None,
         arguments=None,
         wait=True,
         logs=True,
         job_name=None,
         experiment_config=None,
+        kms_key=None,
     ):
         """Runs a processing job.
 
         Args:
             inputs (list[:class:`~sagemaker.processing.ProcessingInput`]): Input files for
                 the processing job. These must be provided as
                 :class:`~sagemaker.processing.ProcessingInput` objects (default: None).
-            kms_key (str): The ARN of the KMS key that is used to encrypt the
-                user code file (default: None).
             outputs (list[:class:`~sagemaker.processing.ProcessingOutput`]): Outputs for
                 the processing job. These can be specified as either path strings or
                 :class:`~sagemaker.processing.ProcessingOutput` objects (default: None).
@@ -142,6 +140,8 @@ def run(
             experiment_config (dict[str, str]): Experiment management configuration.
                 Dictionary contains three optional keys:
                 'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
+            kms_key (str): The ARN of the KMS key that is used to encrypt the
+                user code file (default: None).
 
         Raises:
             ValueError: if ``logs`` is True but ``wait`` is False.
@@ -175,7 +175,13 @@ def _extend_processing_args(self, inputs, outputs, **kwargs):  # pylint: disable
         return inputs, outputs
 
     def _normalize_args(
-        self, job_name=None, arguments=None, inputs=None, kms_key=None, outputs=None, code=None
+        self,
+        job_name=None,
+        arguments=None,
+        inputs=None,
+        outputs=None,
+        code=None,
+        kms_key=None,
     ):
         """Normalizes the arguments so that they can be passed to the job run
 
@@ -417,13 +423,13 @@ def run(
         self,
         code,
         inputs=None,
-        kms_key=None,
         outputs=None,
         arguments=None,
         wait=True,
         logs=True,
         job_name=None,
         experiment_config=None,
+        kms_key=None,
     ):
         """Runs a processing job.
 
@@ -433,8 +439,6 @@ def run(
             inputs (list[:class:`~sagemaker.processing.ProcessingInput`]): Input files for
                 the processing job. These must be provided as
                 :class:`~sagemaker.processing.ProcessingInput` objects (default: None).
-            kms_key (str): The ARN of the KMS key that is used to encrypt the
-                user code file (default: None).
             outputs (list[:class:`~sagemaker.processing.ProcessingOutput`]): Outputs for
                 the processing job. These can be specified as either path strings or
                 :class:`~sagemaker.processing.ProcessingOutput` objects (default: None).
@@ -448,14 +452,16 @@ def run(
             experiment_config (dict[str, str]): Experiment management configuration.
                 Dictionary contains three optional keys:
                 'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
+            kms_key (str): The ARN of the KMS key that is used to encrypt the
+                user code file (default: None).
         """
         normalized_inputs, normalized_outputs = self._normalize_args(
             job_name=job_name,
             arguments=arguments,
             inputs=inputs,
-            kms_key=kms_key,
             outputs=outputs,
             code=code,
+            kms_key=kms_key,
         )
 
         self.latest_job = ProcessingJob.start_new(

src/sagemaker/spark/processing.py

@@ -174,13 +174,13 @@ def run(
         self,
         submit_app,
         inputs=None,
-        kms_key=None,
         outputs=None,
         arguments=None,
         wait=True,
         logs=True,
         job_name=None,
         experiment_config=None,
+        kms_key=None,
     ):
         """Runs a processing job.
 
@@ -189,8 +189,6 @@ def run(
             inputs (list[:class:`~sagemaker.processing.ProcessingInput`]): Input files for
                 the processing job. These must be provided as
                 :class:`~sagemaker.processing.ProcessingInput` objects (default: None).
-            kms_key (str): The ARN of the KMS key that is used to encrypt the
-                user code file (default: None).
             outputs (list[:class:`~sagemaker.processing.ProcessingOutput`]): Outputs for
                 the processing job. These can be specified as either path strings or
                 :class:`~sagemaker.processing.ProcessingOutput` objects (default: None).
@@ -204,19 +202,21 @@ def run(
             experiment_config (dict[str, str]): Experiment management configuration.
                 Dictionary contains three optional keys:
                 'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
+            kms_key (str): The ARN of the KMS key that is used to encrypt the
+                user code file (default: None).
         """
         self._current_job_name = self._generate_current_job_name(job_name=job_name)
 
         super().run(
             submit_app,
             inputs,
-            kms_key,
             outputs,
             arguments,
             wait,
             logs,
             job_name,
             experiment_config,
+            kms_key,
         )
 
     def _extend_processing_args(self, inputs, outputs, **kwargs):
@@ -699,6 +699,7 @@ def run(
         experiment_config=None,
         configuration=None,
         spark_event_logs_s3_uri=None,
+        kms_key=None,
     ):
         """Runs a processing job.
 
@@ -732,6 +733,8 @@ def run(
                 https://docs.aws.amazon.com/emr/latest/ReleaseGuide/emr-configure-apps.html
             spark_event_logs_s3_uri (str): S3 path where spark application events will
                 be published to.
+            kms_key (str): The ARN of the KMS key that is used to encrypt the
+                user code file (default: None).
         """
         self._current_job_name = self._generate_current_job_name(job_name=job_name)
 
@@ -876,6 +879,7 @@ def run(
         experiment_config=None,
         configuration=None,
         spark_event_logs_s3_uri=None,
+        kms_key=None,
     ):
         """Runs a processing job.
 
@@ -909,6 +913,8 @@ def run(
                 https://docs.aws.amazon.com/emr/latest/ReleaseGuide/emr-configure-apps.html
             spark_event_logs_s3_uri (str): S3 path where spark application events will
                 be published to.
+            kms_key (str): The ARN of the KMS key that is used to encrypt the
+                user code file (default: None).
         """
         self._current_job_name = self._generate_current_job_name(job_name=job_name)
 
@@ -934,6 +940,7 @@ def run(
             logs=logs,
             job_name=self._current_job_name,
             experiment_config=experiment_config,
+            kms_key=kms_key,
         )
 
     def _extend_processing_args(self, inputs, outputs, **kwargs):

tests/unit/sagemaker/spark/test_processing.py

@@ -200,7 +200,7 @@ def test_configuration_validation(config, expected, sagemaker_session) -> None:
 def test_spark_processor_base_run(mock_super_run, spark_processor_base):
     spark_processor_base.run(submit_app="app")
 
-    mock_super_run.assert_called_with("app", None, None, None, None, True, True, None, None)
+    mock_super_run.assert_called_with("app", None, None, None, True, True, None, None, None)
 
 
 @pytest.mark.parametrize(
@@ -854,6 +854,7 @@ def test_spark_jar_processor_run(
             logs=True,
             job_name="jobName",
             experiment_config=None,
+            kms_key=None,
         )