fix: jumpstart predictor class support following integ tests

Author: evakravi

src/sagemaker/jumpstart/constants.py

@@ -26,8 +26,8 @@
     BaseSerializer,
     CSVSerializer,
     DataSerializer,
+    IdentitySerializer,
     JSONSerializer,
-    SimpleBaseSerializer,
 )
 
 
@@ -162,12 +162,12 @@
 }
 
 SERIALIZER_TYPE_TO_CLASS_MAP: Dict[SerializerType, BaseSerializer] = {
-    SerializerType.RAW_BYTES: DataSerializer,
-    SerializerType.JSON: JSONSerializer,
-    SerializerType.TEXT: SimpleBaseSerializer,
-    SerializerType.CSV: CSVSerializer,
+    SerializerType.RAW_BYTES: DataSerializer(),
+    SerializerType.JSON: JSONSerializer(),
+    SerializerType.TEXT: IdentitySerializer(),
+    SerializerType.CSV: CSVSerializer(),
 }
 
 DESERIALIZER_TYPE_TO_CLASS_MAP: Dict[DeserializerType, BaseDeserializer] = {
-    DeserializerType.JSON: JSONDeserializer,
+    DeserializerType.JSON: JSONDeserializer(),
 }

src/sagemaker/predictor.py

@@ -182,10 +182,14 @@ def _create_request_args(
             args["EndpointName"] = self.endpoint_name
 
         if "ContentType" not in args:
-            args["ContentType"] = self.content_type
+            args["ContentType"] = (
+                self.content_type
+                if isinstance(self.content_type, str)
+                else ", ".join(self.content_type)
+            )
 
         if "Accept" not in args:
-            args["Accept"] = ", ".join(self.accept)
+            args["Accept"] = self.accept if isinstance(self.accept, str) else ", ".join(self.accept)
 
         if target_model:
             args["TargetModel"] = target_model

src/sagemaker/serializers.py

@@ -20,6 +20,7 @@
 import io
 import json
 import numpy as np
+from pandas import DataFrame
 from six import with_metaclass
 
 from sagemaker.utils import DeferredError
@@ -100,14 +101,17 @@ def serialize(self, data):
 
         Args:
             data (object): Data to be serialized. Can be a NumPy array, list,
-                file, or buffer.
+                file, Pandas DataFrame, or buffer.
 
         Returns:
             str: The data serialized as a CSV-formatted string.
         """
         if hasattr(data, "read"):
             return data.read()
 
+        if isinstance(data, DataFrame):
+            return data.to_csv(header=False, index=False)
+
         is_mutable_sequence_like = self._is_sequence_like(data) and hasattr(data, "__setitem__")
         has_multiple_rows = len(data) > 0 and self._is_sequence_like(data[0])
 

tests/integ/sagemaker/jumpstart/script_mode_class/test_inference.py

@@ -17,14 +17,14 @@
 from sagemaker.jumpstart.constants import INFERENCE_ENTRY_POINT_SCRIPT_NAME
 from sagemaker.jumpstart.artifacts import _retrieve_kwargs
 from sagemaker.jumpstart.enums import EnvVariableUseCase, KwargUseCase
+from sagemaker.jumpstart.predictor import JumpStartPredictor
 from sagemaker.model import Model
 from tests.integ.sagemaker.jumpstart.constants import (
     ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
     JUMPSTART_TAG,
     InferenceTabularDataname,
 )
 from tests.integ.sagemaker.jumpstart.utils import (
-    EndpointInvoker,
     download_inference_assets,
     get_sm_session,
     get_tabular_data,
@@ -91,13 +91,13 @@ def test_jumpstart_inference_model_class(setup):
         **deploy_kwargs,
     )
 
-    endpoint_invoker = EndpointInvoker(
-        endpoint_name=model.endpoint_name,
+    predictor = JumpStartPredictor(
+        endpoint_name=model.endpoint_name, model_id=model_id, model_version=model_version
     )
 
     download_inference_assets()
     ground_truth_label, features = get_tabular_data(InferenceTabularDataname.MULTICLASS)
 
-    response = endpoint_invoker.invoke_tabular_endpoint(features)
+    response = predictor.predict(features)
 
     assert response is not None

tests/unit/sagemaker/jumpstart/test_predictor.py

@@ -7,7 +7,7 @@
 
 
 from sagemaker.jumpstart.utils import verify_model_region_and_return_specs
-from sagemaker.serializers import SimpleBaseSerializer
+from sagemaker.serializers import IdentitySerializer
 from tests.unit.sagemaker.jumpstart.utils import (
     get_special_model_spec,
 )
@@ -28,7 +28,7 @@ def test_list_jumpstart_scripts(
     predictor = JumpStartPredictor(endpoint_name="blah", model_id=model_id)
 
     assert predictor.content_type == MIMEType.X_TEXT
-    assert predictor.serializer == SimpleBaseSerializer
+    assert isinstance(predictor.serializer, IdentitySerializer)
 
-    assert predictor.deserializer == JSONDeserializer
+    assert isinstance(predictor.deserializer, JSONDeserializer)
     assert predictor.accept == MIMEType.JSON