Address PR comments

Author: Payton Staub

src/sagemaker/processing.py

@@ -20,6 +20,7 @@
 
 import os
 import pathlib
+import attr
 
 from six.moves.urllib.parse import urlparse
 from six.moves.urllib.request import url2pathname
@@ -207,13 +208,13 @@ def _normalize_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).
             code (str): This can be an S3 URI or a local path to a file with the framework
                 script to run (default: None). A no op in the base class.
+            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)
 
@@ -444,7 +445,7 @@ def __init__(
 
     def get_run_args(
         self,
-        code=None,
+        code,
         inputs=None,
         outputs=None,
         arguments=None,
@@ -1172,39 +1173,31 @@ def _to_request_dict(self):
         return s3_output_request
 
 
+@attr.s
 class RunArgs(object):
     """Accepts parameters that correspond to ScriptProcessors.
 
     An instance of this class is returned from the ``get_run_args()`` method on processors,
     and is used for normalizing the arguments so that they can be passed to
     :class:`~sagemaker.workflow.steps.ProcessingStep`
-    """
 
-    def __init__(
-        self,
-        code=None,
-        inputs=None,
-        outputs=None,
-        arguments=None,
-    ):
-        """Initializes a ``RunArgs`` instance.
+    Args:
+        code (str): This can be an S3 URI or a local path to a file with the framework
+            script to 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).
+        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).
+        arguments (list[str]): A list of string arguments to be passed to a
+            processing job (default: None).
+    """
 
-        Args:
-            code (str): This can be an S3 URI or a local path to a file with the framework
-                script to 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).
-            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).
-            arguments (list[str]): A list of string arguments to be passed to a
-                processing job (default: None).
-        """
-        self.inputs = inputs
-        self.outputs = outputs
-        self.code = code
-        self.arguments = arguments
+    code = attr.ib()
+    inputs = attr.ib(default=None)
+    outputs = attr.ib(default=None)
+    arguments = attr.ib(default=None)
 
 
 class FeatureStoreOutput(ApiObject):

src/sagemaker/spark/processing.py

@@ -173,7 +173,7 @@ def __init__(
 
     def get_run_args(
         self,
-        code=None,
+        code,
         inputs=None,
         outputs=None,
         arguments=None,
@@ -768,8 +768,8 @@ def get_run_args(
             raise ValueError("submit_app is required")
 
         extended_inputs, extended_outputs = self._extend_processing_args(
-            inputs,
-            outputs,
+            inputs=inputs,
+            outputs=outputs,
             submit_py_files=submit_py_files,
             submit_jars=submit_jars,
             submit_files=submit_files,
@@ -842,8 +842,8 @@ def run(
             raise ValueError("submit_app is required")
 
         extended_inputs, extended_outputs = self._extend_processing_args(
-            inputs,
-            outputs,
+            inputs=inputs,
+            outputs=outputs,
             submit_py_files=submit_py_files,
             submit_jars=submit_jars,
             submit_files=submit_files,
@@ -1016,8 +1016,8 @@ def get_run_args(
             raise ValueError("submit_app is required")
 
         extended_inputs, extended_outputs = self._extend_processing_args(
-            inputs,
-            outputs,
+            inputs=inputs,
+            outputs=outputs,
             submit_class=submit_class,
             submit_jars=submit_jars,
             submit_files=submit_files,
@@ -1090,8 +1090,8 @@ def run(
             raise ValueError("submit_app is required")
 
         extended_inputs, extended_outputs = self._extend_processing_args(
-            inputs,
-            outputs,
+            inputs=inputs,
+            outputs=outputs,
             submit_class=submit_class,
             submit_jars=submit_jars,
             submit_files=submit_files,

src/sagemaker/workflow/steps.py

@@ -320,7 +320,6 @@ def __init__(
         code: str = None,
         property_files: List[PropertyFile] = None,
         cache_config: CacheConfig = None,
-        kms_key=None,
     ):
         """Construct a ProcessingStep, given a `Processor` instance.
 
@@ -341,8 +340,6 @@ def __init__(
             property_files (List[PropertyFile]): A list of property files that workflow looks
                 for and resolves from the configured processing output list.
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
-            kms_key (str): The ARN of the KMS key that is used to encrypt the
-                user code file (default: None)
         """
         super(ProcessingStep, self).__init__(name, StepTypeEnum.PROCESSING)
         self.processor = processor
@@ -351,7 +348,6 @@ 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.
@@ -374,7 +370,6 @@ def arguments(self) -> RequestType:
             inputs=self.inputs,
             outputs=self.outputs,
             code=self.code,
-            kms_key=self.kms_key,
         )
 
         process_args = ProcessingJob._get_process_args(

tests/unit/sagemaker/workflow/test_steps.py

@@ -262,7 +262,6 @@ def test_processing_step_normalizes_args(mock_normalize_args, sagemaker_session)
         outputs=outputs,
         job_arguments=["arg1", "arg2"],
         cache_config=cache_config,
-        kms_key="key",
     )
     mock_normalize_args.return_value = [step.inputs, step.outputs]
     step.to_request()
@@ -271,7 +270,6 @@ 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,
     )