Encapsulated custom handler/formatter logic and added SageMaker endpoint integration tests

Author: smouaa

.gitignore

@@ -31,6 +31,7 @@ node_modules/
 
 # dir
 tests/integration/models/
+tests/integration/sagemaker_test_results.txt
 engines/python/setup/djl_python/tests/resources*
 
 tests/integration/awscurl
@@ -39,3 +40,4 @@ __pycache__
 dist/
 *.egg-info/
 *.pt
+

engines/python/setup/djl_python/custom_formatter_handling.py

@@ -11,8 +11,12 @@
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
 import logging
+import os
+from dataclasses import dataclass
+from typing import Optional, Callable
 
 from djl_python.service_loader import get_annotated_function
+from djl_python.utils import get_sagemaker_function
 
 logger = logging.getLogger(__name__)
 
@@ -26,31 +30,103 @@ def __init__(self, message: str, original_exception: Exception):
         self.__cause__ = original_exception
 
 
+@dataclass
+class CustomFormatters:
+    """Container for input/output formatting functions"""
+    input_formatter: Optional[Callable] = None
+    output_formatter: Optional[Callable] = None
+
+
+@dataclass
+class CustomHandlers:
+    """Container for prediction/initialization handler functions"""
+    prediction_handler: Optional[Callable] = None
+    init_handler: Optional[Callable] = None
+
+
+@dataclass
+class CustomCode:
+    """Container for all custom formatters and handlers"""
+    formatters: CustomFormatters
+    handlers: CustomHandlers
+    is_sagemaker_script: bool = False
+
+    def __init__(self):
+        self.formatters = CustomFormatters()
+        self.handlers = CustomHandlers()
+        self.is_sagemaker_script = False
+
+
 class CustomFormatterHandler:
 
     def __init__(self):
-        self.output_formatter = None
-        self.input_formatter = None
+        self.custom_code = CustomCode()
 
-    def load_formatters(self, model_dir: str):
-        """Load custom formatters from model.py"""
+    def load_formatters(self, model_dir: str) -> CustomCode:
+        """Load custom formatters/handlers from model.py with SageMaker detection"""
         try:
-            self.input_formatter = get_annotated_function(
+            self.custom_code.formatters.input_formatter = get_annotated_function(
                 model_dir, "is_input_formatter")
-            self.output_formatter = get_annotated_function(
+            self.custom_code.formatters.output_formatter = get_annotated_function(
                 model_dir, "is_output_formatter")
+            self.custom_code.handlers.prediction_handler = get_annotated_function(
+                model_dir, "is_prediction_handler")
+            self.custom_code.handlers.init_handler = get_annotated_function(
+                model_dir, "is_init_handler")
+
+            # Detect SageMaker script pattern for backward compatibility
+            self._detect_sagemaker_functions(model_dir)
+
             logger.info(
-                f"Loaded formatters - input: {self.input_formatter}, output: {self.output_formatter}"
+                f"Loaded formatters - input: {bool(self.custom_code.formatters.input_formatter)}, "
+                f"output: {bool(self.custom_code.formatters.output_formatter)}"
             )
+            logger.info(
+                f"Loaded handlers - prediction: {bool(self.custom_code.handlers.prediction_handler)}, "
+                f"init: {bool(self.custom_code.handlers.init_handler)}, "
+                f"sagemaker: {self.custom_code.is_sagemaker_script}")
+            return self.custom_code
         except Exception as e:
             raise CustomFormatterError(
-                f"Failed to load custom formatters from {model_dir}", e)
+                f"Failed to load custom code from {model_dir}", e)
+
+    def _detect_sagemaker_functions(self, model_dir: str):
+        """Detect and load SageMaker-style functions for backward compatibility"""
+        # If no decorator-based code found, check for SageMaker functions
+        if not any([
+                self.custom_code.formatters.input_formatter,
+                self.custom_code.formatters.output_formatter,
+                self.custom_code.handlers.prediction_handler,
+                self.custom_code.handlers.init_handler
+        ]):
+            sagemaker_model_fn = get_sagemaker_function(model_dir, 'model_fn')
+            sagemaker_input_fn = get_sagemaker_function(model_dir, 'input_fn')
+            sagemaker_predict_fn = get_sagemaker_function(
+                model_dir, 'predict_fn')
+            sagemaker_output_fn = get_sagemaker_function(
+                model_dir, 'output_fn')
+
+            if any([
+                    sagemaker_model_fn, sagemaker_input_fn,
+                    sagemaker_predict_fn, sagemaker_output_fn
+            ]):
+                self.custom_code.is_sagemaker_script = True
+                if sagemaker_model_fn:
+                    self.custom_code.handlers.init_handler = sagemaker_model_fn
+                if sagemaker_input_fn:
+                    self.custom_code.formatters.input_formatter = sagemaker_input_fn
+                if sagemaker_predict_fn:
+                    self.custom_code.handlers.prediction_handler = sagemaker_predict_fn
+                if sagemaker_output_fn:
+                    self.custom_code.formatters.output_formatter = sagemaker_output_fn
+                logger.info("Loaded SageMaker-style functions")
 
     def apply_input_formatter(self, decoded_payload, **kwargs):
         """Apply input formatter if available"""
-        if self.input_formatter:
+        if self.custom_code.formatters.input_formatter:
             try:
-                return self.input_formatter(decoded_payload, **kwargs)
+                return self.custom_code.formatters.input_formatter(
+                    decoded_payload, **kwargs)
             except Exception as e:
                 logger.exception("Custom input formatter failed")
                 raise CustomFormatterError(
@@ -59,9 +135,9 @@ def apply_input_formatter(self, decoded_payload, **kwargs):
 
     def apply_output_formatter(self, output):
         """Apply output formatter if available"""
-        if self.output_formatter:
+        if self.custom_code.formatters.output_formatter:
             try:
-                return self.output_formatter(output)
+                return self.custom_code.formatters.output_formatter(output)
             except Exception as e:
                 logger.exception("Custom output formatter failed")
                 raise CustomFormatterError(
@@ -79,3 +155,9 @@ async def apply_output_formatter_streaming_raw(self, stream_generator):
             logger.exception("Streaming formatter failed")
             raise CustomFormatterError(
                 "Custom streaming formatter execution failed", e)
+
+
+def load_custom_code(model_dir: str) -> CustomCode:
+    """Load custom code, checking DJL decorators first, then SageMaker functions"""
+    handler = CustomFormatterHandler()
+    return handler.load_formatters(model_dir)

engines/python/setup/djl_python/input_parser.py

@@ -36,25 +36,25 @@ def input_formatter(function):
     return function
 
 
-def predict_formatter(function):
+def prediction_handler(function):
     """
-    Decorator for predict_formatter. User just need to annotate @predict_formatter for their custom defined function.
+    Decorator for prediction_handler. User just need to annotate @prediction_handler for their custom defined function.
     :param function:  Decorator takes in the function and adds an attribute.
     :return:
     """
     # adding an attribute to the function, which is used to find the decorated function.
-    function.is_predict_formatter = True
+    function.is_prediction_handler = True
     return function
 
 
-def model_loading_formatter(function):
+def init_handler(function):
     """
-    Decorator for model_loading_formatter. User just need to annotate @model_loading_formatter for their custom defined function.
+    Decorator for init_handler. User just need to annotate @init_handler for their custom defined function.
     :param function:  Decorator takes in the function and adds an attribute.
     :return:
     """
     # adding an attribute to the function, which is used to find the decorated function.
-    function.is_model_loading_formatter = True
+    function.is_init_handler = True
     return function
 
 

engines/python/setup/djl_python/sklearn_handler.py

@@ -18,8 +18,8 @@
 from typing import Optional
 from djl_python import Input, Output
 from djl_python.encode_decode import decode
-from djl_python.utils import find_model_file, get_sagemaker_function
-from djl_python.service_loader import get_annotated_function
+from djl_python.utils import find_model_file
+from djl_python.custom_formatter_handling import load_custom_code
 from djl_python.import_utils import joblib, cloudpickle, skops_io as sio
 
 
@@ -28,52 +28,8 @@ class SklearnHandler:
     def __init__(self):
         self.model = None
         self.initialized = False
-        self.custom_input_formatter = None
-        self.custom_output_formatter = None
-        self.custom_predict_formatter = None
-        self.custom_model_loading_formatter = None
+        self.custom_code = None
         self.init_properties = None
-        self.is_sagemaker_script = False
-
-    def _load_custom_formatters(self, model_dir: str):
-        """Load custom formatters, checking DJL decorators first, then SageMaker functions."""
-        # Check for DJL decorator-based custom formatters first
-        self.custom_model_loading_formatter = get_annotated_function(
-            model_dir, "is_model_loading_formatter")
-        self.custom_input_formatter = get_annotated_function(
-            model_dir, "is_input_formatter")
-        self.custom_output_formatter = get_annotated_function(
-            model_dir, "is_output_formatter")
-        self.custom_predict_formatter = get_annotated_function(
-            model_dir, "is_predict_formatter")
-
-        # If no decorator-based formatters found, check for SageMaker-style formatters
-        if not any([
-                self.custom_input_formatter, self.custom_output_formatter,
-                self.custom_predict_formatter,
-                self.custom_model_loading_formatter
-        ]):
-
-            sagemaker_model_fn = get_sagemaker_function(model_dir, 'model_fn')
-            sagemaker_input_fn = get_sagemaker_function(model_dir, 'input_fn')
-            sagemaker_predict_fn = get_sagemaker_function(
-                model_dir, 'predict_fn')
-            sagemaker_output_fn = get_sagemaker_function(
-                model_dir, 'output_fn')
-
-            if any([
-                    sagemaker_model_fn, sagemaker_input_fn,
-                    sagemaker_predict_fn, sagemaker_output_fn
-            ]):
-                self.is_sagemaker_script = True
-                if sagemaker_model_fn:
-                    self.custom_model_loading_formatter = sagemaker_model_fn
-                if sagemaker_input_fn:
-                    self.custom_input_formatter = sagemaker_input_fn
-                if sagemaker_predict_fn:
-                    self.custom_predict_formatter = sagemaker_predict_fn
-                if sagemaker_output_fn:
-                    self.custom_output_formatter = sagemaker_output_fn
 
     def _get_trusted_types(self, properties: dict):
         trusted_types_str = properties.get("skops_trusted_types", "")
@@ -107,12 +63,12 @@ def initialize(self, properties: dict):
                 f"Unsupported model format: {model_format}. Supported formats: skops, joblib, pickle, cloudpickle"
             )
 
-        # Load custom formatters
-        self._load_custom_formatters(model_dir)
+        # Load custom code
+        self.custom_code = load_custom_code(model_dir)
 
         # Load model
-        if self.custom_model_loading_formatter:
-            self.model = self.custom_model_loading_formatter(model_dir)
+        if self.custom_code.handlers.init_handler:
+            self.model = self.custom_code.handlers.init_handler(model_dir)
         else:
             model_file = find_model_file(model_dir, extensions)
             if not model_file:
@@ -154,7 +110,7 @@ def inference(self, inputs: Input) -> Output:
             accept = default_accept
 
         # Validate accept type (skip validation if custom output formatter is provided)
-        if not self.custom_output_formatter:
+        if not self.custom_code.formatters.output_formatter:
             supported_accept_types = ["application/json", "text/csv"]
             if not any(supported_type in accept
                        for supported_type in supported_accept_types):
@@ -164,12 +120,12 @@ def inference(self, inputs: Input) -> Output:
 
         # Input processing
         X = None
-        if self.custom_input_formatter:
-            if self.is_sagemaker_script:
-                X = self.custom_input_formatter(inputs.get_as_bytes(),
-                                                content_type)
+        if self.custom_code.formatters.input_formatter:
+            if self.custom_code.is_sagemaker_script:
+                X = self.custom_code.formatters.input_formatter(
+                    inputs.get_as_bytes(), content_type)
             else:
-                X = self.custom_input_formatter(inputs)
+                X = self.custom_code.formatters.input_formatter(inputs)
         elif "text/csv" in content_type:
             X = decode(inputs, content_type, require_csv_headers=False)
         else:
@@ -185,19 +141,22 @@ def inference(self, inputs: Input) -> Output:
         if X.ndim == 1:
             X = X.reshape(1, -1)
 
-        if self.custom_predict_formatter:
-            predictions = self.custom_predict_formatter(X, self.model)
+        if self.custom_code.handlers.prediction_handler:
+            predictions = self.custom_code.handlers.prediction_handler(
+                X, self.model)
         else:
             predictions = self.model.predict(X)
 
         # Output processing
         outputs = Output()
-        if self.custom_output_formatter:
-            if self.is_sagemaker_script:
-                data = self.custom_output_formatter(predictions, accept)
+        if self.custom_code.formatters.output_formatter:
+            if self.custom_code.is_sagemaker_script:
+                data = self.custom_code.formatters.output_formatter(
+                    predictions, accept)
                 outputs.add_property("Content-Type", accept)
             else:
-                data = self.custom_output_formatter(predictions)
+                data = self.custom_code.formatters.output_formatter(
+                    predictions)
             outputs.add(data)
         elif "text/csv" in accept:
             csv_buffer = StringIO()