fix: include sparkml 2.4 in image uri config properly (#2054)

Author: metrizable

src/sagemaker/image_uri_config/sparkml-serving.json

@@ -29,8 +29,7 @@
                 "us-west-2": "246618743249"
             },
             "repository": "sagemaker-sparkml-serving"
-        }
-    },
+        },
         "2.4": {
             "registries": {
                 "af-south-1": "510948584623",
@@ -60,4 +59,5 @@
             },
             "repository": "sagemaker-sparkml-serving"
         }
+    }
 }

src/sagemaker/sparkml/model.py

@@ -70,7 +70,9 @@ class SparkMLModel(Model):
     model .
     """
 
-    def __init__(self, model_data, role=None, spark_version=2.4, sagemaker_session=None, **kwargs):
+    def __init__(
+        self, model_data, role=None, spark_version="2.4", sagemaker_session=None, **kwargs
+    ):
         """Initialize a SparkMLModel.
 
         Args:

tests/unit/sagemaker/image_uris/test_sparkml.py

@@ -12,6 +12,8 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import pytest
+
 from sagemaker import image_uris
 from tests.unit.sagemaker.image_uris import expected_uris, regions
 
@@ -41,14 +43,15 @@
     "us-west-1": "746614075791",
     "us-west-2": "246618743249",
 }
-VERSION = "2.2"
+VERSIONS = ["2.2", "2.4"]
 
 
-def test_sparkml():
+@pytest.mark.parametrize("version", VERSIONS)
+def test_sparkml(version):
     for region in regions.regions():
-        uri = image_uris.retrieve("sparkml-serving", region=region, version=VERSION)
+        uri = image_uris.retrieve("sparkml-serving", region=region, version=version)
 
         expected = expected_uris.algo_uri(
-            "sagemaker-sparkml-serving", ACCOUNTS[region], region, version=VERSION
+            "sagemaker-sparkml-serving", ACCOUNTS[region], region, version=version
         )
         assert expected == uri

tests/unit/test_pipeline_model.py

@@ -97,7 +97,7 @@ def test_prepare_container_def(tfo, time, sagemaker_session):
         {
             "Environment": {"SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT": "text/csv"},
             "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com"
-            + "/sagemaker-sparkml-serving:2.2",
+            + "/sagemaker-sparkml-serving:2.4",
             "ModelDataUrl": "s3://bucket/model_2.tar.gz",
         },
     ]
@@ -335,7 +335,7 @@ def test_network_isolation(tfo, time, sagemaker_session):
                 "ModelDataUrl": "s3://bucket/model_1.tar.gz",
             },
             {
-                "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:2.2",
+                "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:2.4",
                 "Environment": {},
                 "ModelDataUrl": "s3://bucket/model_2.tar.gz",
             },