feature: add support for spark ml serving container version 2.4

Author: mahithsuresh

src/sagemaker/sparkml/model.py

@@ -59,7 +59,7 @@ class SparkMLModel(Model):
     model .
     """
 
-    def __init__(self, model_data, role=None, spark_version=2.2, sagemaker_session=None, **kwargs):
+    def __init__(self, model_data, role=None, spark_version=2.4, sagemaker_session=None, **kwargs):
         """Initialize a SparkMLModel.
 
         Args:
@@ -73,7 +73,7 @@ def __init__(self, model_data, role=None, spark_version=2.2, sagemaker_session=N
                 artifacts. After the endpoint is created, the inference code
                 might use the IAM role, if it needs to access an AWS resource.
             spark_version (str): Spark version you want to use for executing the
-                inference (default: '2.2').
+                inference (default: '2.4').
             sagemaker_session (sagemaker.session.Session): Session object which
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, the estimator creates one

tests/integ/test_sparkml_serving.py

@@ -17,19 +17,18 @@
 
 import pytest
 
+from botocore.errorfactory import ClientError
+
 from sagemaker.sparkml.model import SparkMLModel
 from sagemaker.utils import sagemaker_timestamp
 from tests.integ import DATA_DIR
 from tests.integ.timeout import timeout_and_delete_endpoint_by_name
 
 
 @pytest.mark.canary_quick
-@pytest.mark.skip(
-    reason="This test has always failed, but the failure was masked by a bug. "
-    "This test should be fixed. Details in https://github.com/aws/sagemaker-python-sdk/pull/968"
-)
 def test_sparkml_model_deploy(sagemaker_session, cpu_instance_type):
-    # Uploads an MLeap serialized MLeap model to S3 and use that to deploy a SparkML model to perform inference
+    # Uploads an MLeap serialized MLeap model to S3 and use that to deploy
+    # a SparkML model to perform inference
     data_path = os.path.join(DATA_DIR, "sparkml_model")
     endpoint_name = "test-sparkml-deploy-{}".format(sagemaker_timestamp())
     model_data = sagemaker_session.upload_data(
@@ -59,7 +58,8 @@ def test_sparkml_model_deploy(sagemaker_session, cpu_instance_type):
         predictor = model.deploy(1, cpu_instance_type, endpoint_name=endpoint_name)
 
         valid_data = "1.0,C,38.0,71.5,1.0,female"
-        assert predictor.predict(valid_data) == "1.0,0.0,38.0,1.0,71.5,0.0,1.0"
+        assert predictor.predict(valid_data) == b"1.0,0.0,38.0,1.0,71.5,0.0,1.0"
 
         invalid_data = "1.0,28.0,C,38.0,71.5,1.0"
-        assert predictor.predict(invalid_data) is None
+        with pytest.raises(ClientError):
+            predictor.predict(invalid_data)

tests/unit/test_sparkml_serving.py

@@ -49,7 +49,7 @@ def sagemaker_session():
 
 def test_sparkml_model(sagemaker_session):
     sparkml = SparkMLModel(sagemaker_session=sagemaker_session, model_data=MODEL_DATA, role=ROLE)
-    assert sparkml.image_uri == image_uris.retrieve("sparkml-serving", REGION, version="2.2")
+    assert sparkml.image_uri == image_uris.retrieve("sparkml-serving", REGION, version="2.4")
 
 
 def test_predictor_type(sagemaker_session):