feature: add spark image_uri_retriever support for framework_version 3.0 (aws#450)

* add spark image_uri_retriever support for framework_version 3.0

* changing framework_version to 2.4 in spark unit tests

* fix spark integration

Author: guoqiao1992

src/sagemaker/image_uri_config/spark.json

@@ -30,6 +30,35 @@
                     "us-gov-west-1": "271483468897"
                 },
                 "repository": "sagemaker-spark-processing"
+            },
+            "3.0": {
+                    "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

@@ -75,26 +75,26 @@ def build_jar():
 
 
 @pytest.fixture(scope="module")
-def spark_py_processor(sagemaker_session, cpu_instance_type, spark_processing_latest_version):
+def spark_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=spark_processing_latest_version,
+        framework_version="2.4",
     )
 
     return spark_py_processor
 
 
 @pytest.fixture(scope="module")
-def spark_jar_processor(sagemaker_session, cpu_instance_type, spark_processing_latest_version):
+def spark_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=spark_processing_latest_version,
+        framework_version="2.4",
     )
 
     return spark_jar_processor

tests/unit/sagemaker/spark/test_processing.py

@@ -68,7 +68,7 @@ def spark_processor_base(sagemaker_session) -> _SparkProcessorBase:
     spark_processor_base = _SparkProcessorBase(
         base_job_name="sm-spark",
         role="AmazonSageMaker-ExecutionRole",
-        framework_version="0.1.0",
+        framework_version="2.4",
         instance_count=1,
         instance_type="ml.c5.xlarge",
         image_uri="790336243319.dkr.ecr.us-west-2.amazonaws.com/sagemaker-spark:0.1",
@@ -83,7 +83,7 @@ def py_spark_processor(sagemaker_session) -> PySparkProcessor:
     spark = PySparkProcessor(
         base_job_name="sm-spark",
         role="AmazonSageMaker-ExecutionRole",
-        framework_version="0.1.0",
+        framework_version="2.4",
         instance_count=1,
         instance_type="ml.c5.xlarge",
         image_uri="790336243319.dkr.ecr.us-west-2.amazonaws.com/sagemaker-spark:0.1",
@@ -125,7 +125,7 @@ def test_pyspark_processor_instantiation(sagemaker_session):
     PySparkProcessor(
         base_job_name="sm-spark",
         role="AmazonSageMaker-ExecutionRole",
-        framework_version="0.1.0",
+        framework_version="2.4",
         instance_count=1,
         instance_type="ml.c5.xlarge",
         sagemaker_session=sagemaker_session,
@@ -183,7 +183,7 @@ def test_configuration_validation(config, expected, sagemaker_session) -> None:
     spark = PySparkProcessor(
         base_job_name="sm-spark",
         role="AmazonSageMaker-ExecutionRole",
-        framework_version="0.1.0",
+        framework_version="2.4",
         instance_count=1,
         instance_type="ml.c5.xlarge",
         sagemaker_session=sagemaker_session,
@@ -362,6 +362,16 @@ def test_stage_submit_deps(mock_s3_uploader, py_spark_processor, jar_file, confi
             },
             "153931337802.dkr.ecr.us-west-2.amazonaws.com/sagemaker-spark-processing:2.4-cpu",
         ),
+        (
+            {
+                "image_uri": None,
+                "framework_version": "3.0",
+                "py_version": None,
+                "container_version": None,
+                "instance_type": "ml.c5.xlarge",
+            },
+            "153931337802.dkr.ecr.us-west-2.amazonaws.com/sagemaker-spark-processing:3.0-cpu",
+        ),
         (
             {
                 "image_uri": "image_uri",
@@ -372,6 +382,16 @@ def test_stage_submit_deps(mock_s3_uploader, py_spark_processor, jar_file, confi
             },
             "image_uri",
         ),
+        (
+            {
+                "image_uri": None,
+                "framework_version": "invalidFrameworkVersion",
+                "py_version": None,
+                "container_version": None,
+                "instance_type": "ml.c5.xlarge",
+            },
+            ValueError,
+        ),
         (
             {
                 "image_uri": None,
@@ -848,7 +868,7 @@ def test_spark_jar_processor_run(
     spark_jar_processor = SparkJarProcessor(
         base_job_name="sm-spark",
         role="AmazonSageMaker-ExecutionRole",
-        framework_version="0.1.0",
+        framework_version="2.4",
         instance_count=1,
         instance_type="ml.c5.xlarge",
         image_uri="790336243319.dkr.ecr.us-west-2.amazonaws.com/sagemaker-spark:0.1",