Interim commit: support all processor types in ProcessingStep

Author: Payton Staub

src/sagemaker/processing.py

@@ -123,6 +123,25 @@ def __init__(
 
         self.sagemaker_session = sagemaker_session or Session()
 
+    def get_run_args(
+        self,
+        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,
+            kms_key=kms_key,
+            outputs=outputs,
+        )
+
+        return RunArgs(inputs=normalized_inputs, outputs=normalized_outputs, code=None)
+
     def run(
         self,
         inputs=None,
@@ -442,6 +461,27 @@ def __init__(
             network_config=network_config,
         )
 
+    def get_run_args(
+        self,
+        code,
+        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,
+        )
+
+        return RunArgs(inputs=normalized_inputs, outputs=normalized_outputs, code=code)
+
     def run(
         self,
         code,
@@ -1144,6 +1184,40 @@ def _to_request_dict(self):
         return s3_output_request
 
 
+class RunArgs(object):
+    """Accepts parameters that specify an Amazon S3 output for a processing job.
+
+    It also provides a method to turn those parameters into a dictionary.
+    """
+
+    def __init__(
+        self,
+        inputs=None,
+        outputs=None,
+        code=None,
+    ):
+        """Initializes a ``ProcessingOutput`` instance.
+
+        ``ProcessingOutput`` accepts parameters that specify an Amazon S3 output for a
+        processing job and provides a method to turn those parameters into a dictionary.
+
+        Args:
+            source (str): The source for the output.
+            destination (str): The destination of the output. If a destination
+                is not provided, one will be generated:
+                "s3://<default-bucket-name>/<job-name>/output/<output-name>".
+            output_name (str): The name of the output. If a name
+                is not provided, one will be generated (eg. "output-1").
+            s3_upload_mode (str): Valid options are "EndOfJob" or "Continuous".
+            app_managed (bool): Whether the input are managed by SageMaker or application
+            feature_store_output (:class:`~sagemaker.processing.FeatureStoreOutput`)
+                Configuration for processing job outputs of FeatureStore.
+        """
+        self.inputs = inputs
+        self.outputs = outputs
+        self.code = code
+
+
 class FeatureStoreOutput(ApiObject):
     """Configuration for processing job outputs in Amazon SageMaker Feature Store."""
 

src/sagemaker/spark/processing.py

@@ -33,7 +33,7 @@
 
 from sagemaker import image_uris
 from sagemaker.local.image import _ecr_login_if_needed, _pull_image
-from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor
+from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor, RunArgs
 from sagemaker.s3 import S3Uploader
 from sagemaker.session import Session
 from sagemaker.spark import defaults
@@ -171,6 +171,25 @@ def __init__(
             network_config=network_config,
         )
 
+    def get_run_args(
+        self,
+        submit_app,
+        inputs=None,
+        outputs=None,
+        arguments=None,
+        job_name=None,
+        kms_key=None,
+    ):
+        # TODO: description
+        return super().get_run_args(
+            code=submit_app,
+            inputs=inputs,
+            outputs=outputs,
+            arguments=arguments,
+            job_name=job_name,
+            kms_key=kms_key,
+        )
+
     def run(
         self,
         submit_app,
@@ -685,6 +704,46 @@ def __init__(
             network_config=network_config,
         )
 
+    def get_run_args(
+        self,
+        submit_app,
+        submit_py_files=None,
+        submit_jars=None,
+        submit_files=None,
+        inputs=None,
+        outputs=None,
+        arguments=None,
+        job_name=None,
+        configuration=None,
+        spark_event_logs_s3_uri=None,
+        kms_key=None,
+    ):
+        self._current_job_name = self._generate_current_job_name(job_name=job_name)
+        self.command = [_SparkProcessorBase._default_command]
+
+        if not submit_app:
+            raise ValueError("submit_app is required")
+
+        extended_inputs, extended_outputs = self._extend_processing_args(
+            inputs,
+            outputs,
+            submit_py_files=submit_py_files,
+            submit_jars=submit_jars,
+            submit_files=submit_files,
+            configuration=configuration,
+            spark_event_logs_s3_uri=spark_event_logs_s3_uri,
+        )
+
+        # TODO: description
+        return super().get_run_args(
+            submit_app=submit_app,
+            inputs=extended_inputs,
+            outputs=extended_outputs,
+            arguments=arguments,
+            job_name=self._current_job_name,
+            kms_key=kms_key,
+        )
+
     def run(
         self,
         submit_app,
@@ -866,6 +925,46 @@ def __init__(
             network_config=network_config,
         )
 
+    def get_run_args(
+        self,
+        submit_app,
+        submit_class=None,
+        submit_jars=None,
+        submit_files=None,
+        inputs=None,
+        outputs=None,
+        arguments=None,
+        job_name=None,
+        configuration=None,
+        spark_event_logs_s3_uri=None,
+        kms_key=None,
+    ):
+        self._current_job_name = self._generate_current_job_name(job_name=job_name)
+        self.command = [_SparkProcessorBase._default_command]
+
+        if not submit_app:
+            raise ValueError("submit_app is required")
+
+        extended_inputs, extended_outputs = self._extend_processing_args(
+            inputs,
+            outputs,
+            submit_class=submit_class,
+            submit_jars=submit_jars,
+            submit_files=submit_files,
+            configuration=configuration,
+            spark_event_logs_s3_uri=spark_event_logs_s3_uri,
+        )
+
+        # TODO: description
+        return super().get_run_args(
+            submit_app=submit_app,
+            inputs=extended_inputs,
+            outputs=extended_outputs,
+            arguments=arguments,
+            job_name=self._current_job_name,
+            kms_key=kms_key,
+        )
+
     def run(
         self,
         submit_app,