feature: add spark processor

Author: guoqiao1992

src/sagemaker/image_uri_config/spark.json

@@ -0,0 +1,35 @@
+{
+    "processing": {
+        "processors": ["cpu"],
+        "versions": {
+            "2.4": {
+                "py_versions": ["py37"],
+                "registries": {
+                    "me-south-1": "750251592176",
+                    "ap-south-1": "105495057255",
+                    "eu-west-3": "136845547031",
+                    "us-east-2": "314815235551",
+                    "eu-west-1": "571004829621",
+                    "eu-central-1": "906073651304",
+                    "sa-east-1": "737130764395",
+                    "ap-east-1": "732049463269",
+                    "us-east-1": "173754725891",
+                    "ap-northeast-2": "860869212795",
+                    "eu-west-2": "836651553127",
+                    "ap-northeast-1": "411782140378",
+                    "us-west-2": "153931337802",
+                    "us-west-1": "667973535471",
+                    "ap-southeast-1": "759080221371",
+                    "ap-southeast-2": "440695851116",
+                    "ca-central-1": "446299261295",
+                    "cn-north-1": "671472414489",
+                    "cn-northwest-1": "844356804704",
+                    "eu-south-1": "753923664805",
+                    "af-south-1": "309385258863",
+                    "us-gov-west-1": "271483468897"
+                },
+                "repository": "sagemaker-spark-processing"
+            }
+        }
+    }
+}

src/sagemaker/image_uris.py

@@ -19,6 +19,7 @@
 import re
 
 from sagemaker import utils
+from sagemaker.spark import defaults
 
 logger = logging.getLogger(__name__)
 
@@ -33,6 +34,7 @@ def retrieve(
     instance_type=None,
     accelerator_type=None,
     image_scope=None,
+    container_version=None,
 ):
     """Retrieves the ECR URI for the Docker image matching the given arguments.
 
@@ -51,6 +53,7 @@ def retrieve(
         image_scope (str): The image type, i.e. what it is used for.
             Valid values: "training", "inference", "eia". If ``accelerator_type`` is set,
             ``image_scope`` is ignored.
+        container_version (str): the version of docker image
 
     Returns:
         str: the ECR URI for the corresponding SageMaker Docker image.
@@ -63,7 +66,7 @@ def retrieve(
     version = _validate_version_and_set_if_needed(version, config, framework)
     version_config = config["versions"][_version_for_config(version, config)]
 
-    py_version = _validate_py_version_and_set_if_needed(py_version, version_config)
+    py_version = _validate_py_version_and_set_if_needed(py_version, version_config, framework)
     version_config = version_config.get(py_version) or version_config
 
     registry = _registry_from_region(region, version_config["registries"])
@@ -74,7 +77,9 @@ def retrieve(
     processor = _processor(
         instance_type, config.get("processors") or version_config.get("processors")
     )
-    tag = _format_tag(version_config.get("tag_prefix", version), processor, py_version)
+    tag = _format_tag(
+        version_config.get("tag_prefix", version), processor, py_version, container_version
+    )
 
     if tag:
         repo += ":{}".format(tag)
@@ -103,7 +108,7 @@ def _config_for_framework_and_scope(framework, image_scope, accelerator_type=Non
                 available_scopes[0],
                 image_scope,
             )
-        image_scope = available_scopes[0]
+        image_scope = list(available_scopes)[0]
 
     if not image_scope and "scope" in config and set(available_scopes) == {"training", "inference"}:
         logger.info(
@@ -212,7 +217,7 @@ def _processor(instance_type, available_processors):
     return processor
 
 
-def _validate_py_version_and_set_if_needed(py_version, version_config):
+def _validate_py_version_and_set_if_needed(py_version, version_config, framework):
     """Checks if the Python version is one of the supported versions."""
     if "repository" in version_config:
         available_versions = version_config.get("py_versions")
@@ -224,6 +229,9 @@ def _validate_py_version_and_set_if_needed(py_version, version_config):
             logger.info("Ignoring unnecessary Python version: %s.", py_version)
         return None
 
+    if py_version is None and defaults.SPARK_NAME == framework:
+        return None
+
     if py_version is None and len(available_versions) == 1:
         logger.info("Defaulting to only available Python version: %s", available_versions[0])
         return available_versions[0]
@@ -242,6 +250,6 @@ def _validate_arg(arg, available_options, arg_name):
         )
 
 
-def _format_tag(tag_prefix, processor, py_version):
+def _format_tag(tag_prefix, processor, py_version, container_version):
     """Creates a tag for the image URI."""
-    return "-".join([x for x in (tag_prefix, processor, py_version) if x])
+    return "-".join([x for x in (tag_prefix, processor, py_version, container_version) if x])

src/sagemaker/processing.py

@@ -166,6 +166,10 @@ def run(
         if wait:
             self.latest_job.wait(logs=logs)
 
+    def _extend_processing_args(self, inputs, outputs, **kwargs):  # pylint: disable=W0613
+        """Extend inputs and outputs based on extra parameters"""
+        return inputs, outputs
+
     def _normalize_args(self, job_name=None, arguments=None, inputs=None, outputs=None, code=None):
         """Normalizes the arguments so that they can be passed to the job run
 

src/sagemaker/spark/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Placeholder docstring"""
+from __future__ import absolute_import
+
+from sagemaker.spark.processing import PySparkProcessor, SparkJarProcessor  # noqa: F401

src/sagemaker/spark/defaults.py

@@ -0,0 +1,16 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Default constants used by Spark processing."""
+from __future__ import absolute_import
+
+SPARK_NAME = "spark"