Interim commit #3 - support all processors for ProcessingStep

Author: Payton Staub

src/sagemaker/processing.py

@@ -463,24 +463,29 @@ def __init__(
 
     def get_run_args(
         self,
-        code,
+        code=None,
         inputs=None,
         outputs=None,
         arguments=None,
-        job_name=None,
-        kms_key=None,
     ):
-        # TODO: description
-        normalized_inputs, normalized_outputs = self._normalize_args(
-            job_name=job_name,
-            arguments=arguments,
-            inputs=inputs,
-            outputs=outputs,
-            code=code,
-            kms_key=kms_key,
-        )
+        """Returns a RunArgs object. For processors (:class:`~sagemaker.spark.processing.PySparkProcessor`,
+            :class:`~sagemaker.spark.processing.SparkJar`) that have special
+            run() arguments, this object contains the normalized arguments for passing to
+            :class:`~sagemaker.workflow.steps.ProcessingStep`.
 
-        return RunArgs(inputs=normalized_inputs, outputs=normalized_outputs, code=code)
+        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).
+        """
+        return RunArgs(code=code, inputs=inputs, outputs=outputs, arguments=arguments)
 
     def run(
         self,
@@ -1195,6 +1200,7 @@ def __init__(
         inputs=None,
         outputs=None,
         code=None,
+        arguments=None,
     ):
         """Initializes a ``ProcessingOutput`` instance.
 
@@ -1216,6 +1222,7 @@ def __init__(
         self.inputs = inputs
         self.outputs = outputs
         self.code = code
+        self.arguments = arguments
 
 
 class FeatureStoreOutput(ApiObject):

src/sagemaker/spark/processing.py

@@ -173,21 +173,33 @@ def __init__(
 
     def get_run_args(
         self,
-        submit_app,
+        code=None,
         inputs=None,
         outputs=None,
         arguments=None,
-        job_name=None,
-        kms_key=None,
     ):
-        # TODO: description
+        """Returns a RunArgs object. For processors (:class:`~sagemaker.spark.processing.PySparkProcessor`,
+            :class:`~sagemaker.spark.processing.SparkJar`) that have special
+            run() arguments, this object contains the normalized arguments for passing to
+            :class:`~sagemaker.workflow.steps.ProcessingStep`.
+
+        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).
+        """
         return super().get_run_args(
-            code=submit_app,
+            code=code,
             inputs=inputs,
             outputs=outputs,
             arguments=arguments,
-            job_name=job_name,
-            kms_key=kms_key,
         )
 
     def run(
@@ -716,8 +728,35 @@ def get_run_args(
         job_name=None,
         configuration=None,
         spark_event_logs_s3_uri=None,
-        kms_key=None,
     ):
+        """Returns a RunArgs object. This object contains the normalized inputs, outputs
+            and arguments needed when creating using a ``PySparkProcessor`` in a :class:`~sagemaker.workflow.steps.ProcessingStep`.
+
+        Args:
+            submit_app (str): Path (local or S3) to Python file to submit to Spark
+                as the primary application
+            submit_py_files (list[str]): List of paths (local or S3) to provide for
+                `spark-submit --py-files` option
+            submit_jars (list[str]): List of paths (local or S3) to provide for
+                `spark-submit --jars` option
+            submit_files (list[str]): List of paths (local or S3) to provide for
+                `spark-submit --files` option
+            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).
+            job_name (str): Processing job name. If not specified, the processor generates
+                a default job name, based on the base job name and current timestamp.
+            configuration (list[dict] or dict): Configuration for Hadoop, Spark, or Hive.
+                List or dictionary of EMR-style classifications.
+                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.
+        """
         self._current_job_name = self._generate_current_job_name(job_name=job_name)
         self.command = [_SparkProcessorBase._default_command]
 
@@ -734,14 +773,11 @@ def get_run_args(
             spark_event_logs_s3_uri=spark_event_logs_s3_uri,
         )
 
-        # TODO: description
         return super().get_run_args(
-            submit_app=submit_app,
+            code=submit_app,
             inputs=extended_inputs,
             outputs=extended_outputs,
             arguments=arguments,
-            job_name=self._current_job_name,
-            kms_key=kms_key,
         )
 
     def run(
@@ -821,6 +857,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):
@@ -937,8 +974,35 @@ def get_run_args(
         job_name=None,
         configuration=None,
         spark_event_logs_s3_uri=None,
-        kms_key=None,
     ):
+        """Returns a RunArgs object. This object contains the normalized inputs, outputs
+            and arguments needed when creating using a ``SparkJarProcessor`` in a :class:`~sagemaker.workflow.steps.ProcessingStep`.
+
+        Args:
+            submit_app (str): Path (local or S3) to Python file to submit to Spark
+                as the primary application
+            submit_class (str): Java class reference to submit to Spark as the primary
+                application
+            submit_jars (list[str]): List of paths (local or S3) to provide for
+                `spark-submit --jars` option
+            submit_files (list[str]): List of paths (local or S3) to provide for
+                `spark-submit --files` option
+            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).
+            job_name (str): Processing job name. If not specified, the processor generates
+                a default job name, based on the base job name and current timestamp.
+            configuration (list[dict] or dict): Configuration for Hadoop, Spark, or Hive.
+                List or dictionary of EMR-style classifications.
+                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.
+        """
         self._current_job_name = self._generate_current_job_name(job_name=job_name)
         self.command = [_SparkProcessorBase._default_command]
 
@@ -955,14 +1019,11 @@ def get_run_args(
             spark_event_logs_s3_uri=spark_event_logs_s3_uri,
         )
 
-        # TODO: description
         return super().get_run_args(
-            submit_app=submit_app,
+            code=submit_app,
             inputs=extended_inputs,
             outputs=extended_outputs,
             arguments=arguments,
-            job_name=self._current_job_name,
-            kms_key=kms_key,
         )
 
     def run(

src/sagemaker/workflow/steps.py

@@ -320,6 +320,7 @@ def __init__(
         code: str = None,
         property_files: List[PropertyFile] = None,
         cache_config: CacheConfig = None,
+        kms_key=None,
     ):
         """Construct a ProcessingStep, given a `Processor` instance.
 
@@ -340,6 +341,8 @@ 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
@@ -348,6 +351,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.
@@ -370,7 +374,9 @@ 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()
         )