Update black-check version, add support for Spark 3.1 Processing (#3034)

Author: shreyapandit

src/sagemaker/image_uri_config/spark.json

@@ -59,6 +59,35 @@
                         "us-gov-west-1": "271483468897"
                     },
                     "repository": "sagemaker-spark-processing"
+            },
+            "3.1": {
+                "py_versions": ["py37"],
+                "registries": {
+                    "me-south-1": "750251592176",
+                    "ap-south-1": "105495057255",
+                    "eu-north-1": "330188676905",
+                    "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"
             }
         }
     }

tests/integ/test_spark_processing.py

@@ -87,6 +87,19 @@ def spark_py_processor(sagemaker_session, cpu_instance_type):
     return spark_py_processor
 
 
+@pytest.fixture(scope="module")
+def spark_v3_py_processor(sagemaker_session, cpu_instance_type):
+    spark_py_processor = PySparkProcessor(
+        role="SageMakerRole",
+        instance_count=2,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+        framework_version="3.1",
+    )
+
+    return spark_py_processor
+
+
 @pytest.fixture(scope="module")
 def spark_jar_processor(sagemaker_session, cpu_instance_type):
     spark_jar_processor = SparkJarProcessor(
@@ -100,6 +113,19 @@ def spark_jar_processor(sagemaker_session, cpu_instance_type):
     return spark_jar_processor
 
 
+@pytest.fixture(scope="module")
+def spark_v3_jar_processor(sagemaker_session, cpu_instance_type):
+    spark_jar_processor = SparkJarProcessor(
+        role="SageMakerRole",
+        instance_count=2,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+        framework_version="3.1",
+    )
+
+    return spark_jar_processor
+
+
 @pytest.fixture
 def configuration() -> list:
     configuration = [
@@ -180,6 +206,15 @@ def configuration() -> list:
     return configuration
 
 
+def test_sagemaker_pyspark_v3(
+    spark_v3_py_processor, spark_v3_jar_processor, sagemaker_session, configuration, build_jar
+):
+    test_sagemaker_pyspark_multinode(spark_v3_py_processor, sagemaker_session, configuration)
+    test_sagemaker_java_jar_multinode(
+        spark_v3_jar_processor, sagemaker_session, configuration, build_jar
+    )
+
+
 def test_sagemaker_pyspark_multinode(spark_py_processor, sagemaker_session, configuration):
     """Test that basic multinode case works on 32KB of data"""
     bucket = spark_py_processor.sagemaker_session.default_bucket()

tox.ini

@@ -134,7 +134,7 @@ commands = doc8
 setenv =
     LC_ALL=C.UTF-8
     LANG=C.UTF-8
-deps = black==22.1.0
+deps = black==22.3.0
 commands =
     black -l 100 ./
 
@@ -143,7 +143,7 @@ commands =
 setenv =
     LC_ALL=C.UTF-8
     LANG=C.UTF-8
-deps = black==22.1.0
+deps = black==22.3.0
 commands =
     black -l 100 --check ./