feat: add pre-processing and post-processing logic to inference_spec (#1560)

* add pre-processing and post-processing logic to inference_spec

* fix format

* make  accept_type and content_type optional

* remove accept_type and content_type from pre/post processing

* correct typo

Author: pravali96

src/sagemaker/serve/model_server/multi_model_server/inference.py

@@ -44,13 +44,19 @@ def input_fn(input_data, content_type):
     """Deserializes the bytes that were received from the model server"""
     try:
         if hasattr(schema_builder, "custom_input_translator"):
-            return schema_builder.custom_input_translator.deserialize(
+            deserialized_data = schema_builder.custom_input_translator.deserialize(
                 io.BytesIO(input_data), content_type
             )
         else:
-            return schema_builder.input_deserializer.deserialize(
+            deserialized_data = schema_builder.input_deserializer.deserialize(
                 io.BytesIO(input_data), content_type[0]
             )
+
+        # Check if preprocess method is defined and call it
+        if hasattr(inference_spec, "preprocess"):
+            return inference_spec.preprocess(deserialized_data)
+
+        return deserialized_data
     except Exception as e:
         logger.error("Encountered error: %s in deserialize_response." % e)
         raise Exception("Encountered error in deserialize_request.") from e
@@ -64,6 +70,8 @@ def predict_fn(input_data, predict_callable):
 def output_fn(predictions, accept_type):
     """Prediction is serialized to bytes and sent back to the customer"""
     try:
+        if hasattr(inference_spec, "postprocess"):
+            predictions = inference_spec.postprocess(predictions)
         if hasattr(schema_builder, "custom_output_translator"):
             return schema_builder.custom_output_translator.serialize(predictions, accept_type)
         else:

src/sagemaker/serve/model_server/torchserve/inference.py

@@ -66,13 +66,19 @@ def input_fn(input_data, content_type):
     """Placeholder docstring"""
     try:
         if hasattr(schema_builder, "custom_input_translator"):
-            return schema_builder.custom_input_translator.deserialize(
+            deserialized_data = schema_builder.custom_input_translator.deserialize(
                 io.BytesIO(input_data), content_type
             )
         else:
-            return schema_builder.input_deserializer.deserialize(
+            deserialized_data = schema_builder.input_deserializer.deserialize(
                 io.BytesIO(input_data), content_type[0]
             )
+
+        # Check if preprocess method is defined and call it
+        if hasattr(inference_spec, "preprocess"):
+            return inference_spec.preprocess(deserialized_data)
+
+        return deserialized_data
     except Exception as e:
         raise Exception("Encountered error in deserialize_request.") from e
 
@@ -85,6 +91,8 @@ def predict_fn(input_data, predict_callable):
 def output_fn(predictions, accept_type):
     """Placeholder docstring"""
     try:
+        if hasattr(inference_spec, "postprocess"):
+            predictions = inference_spec.postprocess(predictions)
         if hasattr(schema_builder, "custom_output_translator"):
             return schema_builder.custom_output_translator.serialize(predictions, accept_type)
         else:

src/sagemaker/serve/spec/inference_spec.py

@@ -28,6 +28,12 @@ def invoke(self, input_object: object, model: object):
             model (object): The model object
         """
 
+    def preprocess(self, input_data: object):
+        """Custom pre-processing function"""
+
+    def postprocess(self, predictions: object):
+        """Custom post-processing function"""
+
     def prepare(self, *args, **kwargs):
         """Custom prepare function"""