Added support for Sagemaker ENV variables + custom formatter scripts

Author: smouaa

engines/python/setup/djl_python/input_parser.py

@@ -36,6 +36,28 @@ def input_formatter(function):
     return function
 
 
+def predict_formatter(function):
+    """
+    Decorator for predict_formatter. User just need to annotate @predict_formatter 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
+    return function
+
+
+def model_loading_formatter(function):
+    """
+    Decorator for model_loading_formatter. User just need to annotate @model_loading_formatter 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
+    return function
+
+
 @dataclass
 class ParsedInput:
     errors: dict = field(default_factory=lambda: {})

engines/python/setup/djl_python/sklearn_handler.py

@@ -18,7 +18,7 @@
 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
+from djl_python.utils import find_model_file, get_sagemaker_function
 from djl_python.service_loader import get_annotated_function
 from djl_python.import_utils import joblib, cloudpickle, skops_io as sio
 
@@ -31,6 +31,49 @@ def __init__(self):
         self.custom_input_formatter = None
         self.custom_output_formatter = None
         self.custom_predict_formatter = None
+        self.custom_model_loading_formatter = 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", "")
@@ -46,6 +89,8 @@ def _get_trusted_types(self, properties: dict):
         return trusted_types
 
     def initialize(self, properties: dict):
+        # Store initialization properties for use during inference
+        self.init_properties = properties.copy()
         model_dir = properties.get("model_dir")
         model_format = properties.get("model_format", "skops")
 
@@ -62,43 +107,51 @@ def initialize(self, properties: dict):
                 f"Unsupported model format: {model_format}. Supported formats: skops, joblib, pickle, cloudpickle"
             )
 
-        model_file = find_model_file(model_dir, extensions)
-        if not model_file:
-            raise FileNotFoundError(
-                f"No model file found with format '{model_format}' in {model_dir}"
-            )
+        # Load custom formatters
+        self._load_custom_formatters(model_dir)
 
-        if model_format == "skops":
-            trusted_types = self._get_trusted_types(properties)
-            self.model = sio.load(model_file, trusted=trusted_types)
+        # Load model
+        if self.custom_model_loading_formatter:
+            self.model = self.custom_model_loading_formatter(model_dir)
         else:
-            if properties.get("trust_insecure_model_files",
-                              "false").lower() != "true":
-                raise ValueError(
-                    f"option.trust_insecure_model_files must be set to 'true' to use {model_format} format (only skops is secure by default)"
+            model_file = find_model_file(model_dir, extensions)
+            if not model_file:
+                raise FileNotFoundError(
+                    f"No model file found with format '{model_format}' in {model_dir}"
                 )
 
-            if model_format == "joblib":
-                self.model = joblib.load(model_file)
-            elif model_format == "pickle":
-                with open(model_file, 'rb') as f:
-                    self.model = pickle.load(f)
-            elif model_format == "cloudpickle":
-                with open(model_file, 'rb') as f:
-                    self.model = cloudpickle.load(f)
-
-        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 model_format == "skops":
+                trusted_types = self._get_trusted_types(properties)
+                self.model = sio.load(model_file, trusted=trusted_types)
+            else:
+                if properties.get("trust_insecure_model_files",
+                                  "false").lower() != "true":
+                    raise ValueError(
+                        f"option.trust_insecure_model_files must be set to 'true' to use {model_format} format (only skops is secure by default)"
+                    )
+
+                if model_format == "joblib":
+                    self.model = joblib.load(model_file)
+                elif model_format == "pickle":
+                    with open(model_file, 'rb') as f:
+                        self.model = pickle.load(f)
+                elif model_format == "cloudpickle":
+                    with open(model_file, 'rb') as f:
+                        self.model = cloudpickle.load(f)
 
         self.initialized = True
 
     def inference(self, inputs: Input) -> Output:
         content_type = inputs.get_property("Content-Type")
-        accept = inputs.get_property("Accept") or "application/json"
+        properties = inputs.get_properties()
+        default_accept = self.init_properties.get("default_accept",
+                                                  "application/json")
+
+        accept = inputs.get_property("Accept")
+
+        # If no accept type is specified in the request, use default
+        if accept == "*/*":
+            accept = default_accept
 
         # Validate accept type (skip validation if custom output formatter is provided)
         if not self.custom_output_formatter:
@@ -112,7 +165,11 @@ def inference(self, inputs: Input) -> Output:
         # Input processing
         X = None
         if self.custom_input_formatter:
-            X = self.custom_input_formatter(inputs)
+            if self.is_sagemaker_script:
+                X = self.custom_input_formatter(inputs.get_as_bytes(),
+                                                content_type)
+            else:
+                X = self.custom_input_formatter(inputs)
         elif "text/csv" in content_type:
             X = decode(inputs, content_type, require_csv_headers=False)
         else:
@@ -129,17 +186,20 @@ def inference(self, inputs: Input) -> Output:
             X = X.reshape(1, -1)
 
         if self.custom_predict_formatter:
-            predictions = self.custom_predict_formatter(self.model, X)
+            predictions = self.custom_predict_formatter(X, self.model)
         else:
             predictions = self.model.predict(X)
 
         # Output processing
-        if self.custom_output_formatter:
-            return self.custom_output_formatter(predictions)
-
-        # Supports CSV/JSON outputs by default
         outputs = Output()
-        if "text/csv" in accept:
+        if self.custom_output_formatter:
+            if self.is_sagemaker_script:
+                data = self.custom_output_formatter(predictions, accept)
+                outputs.add_property("Content-Type", accept)
+            else:
+                data = self.custom_output_formatter(predictions)
+            outputs.add(data)
+        elif "text/csv" in accept:
             csv_buffer = StringIO()
             np.savetxt(csv_buffer, predictions, fmt='%s', delimiter=',')
             outputs.add(csv_buffer.getvalue().rstrip())

engines/python/setup/djl_python/tests/test_encode_decode.py

@@ -115,7 +115,8 @@ def test_decode_text_csv(self):
             mock_decode_csv.return_value = {"inputs": ["test input"]}
             result = decode(self.mock_input, "text/csv")
 
-        mock_decode_csv.assert_called_once_with(self.mock_input)
+        mock_decode_csv.assert_called_once_with(self.mock_input,
+                                                require_headers=True)
         self.assertEqual(result, {"inputs": ["test input"]})
 
     def test_decode_text_plain(self):

engines/python/setup/djl_python/utils.py

@@ -13,10 +13,21 @@
 import glob
 import logging
 import os
-from typing import Optional, List
+import inspect
+import importlib.util
+from typing import Optional, List, Callable
 
 from djl_python import Output
 from djl_python.inputs import Input
+from djl_python.service_loader import load_model_service, has_function_in_module
+
+# SageMaker function signatures for validation
+SAGEMAKER_SIGNATURES = {
+    'model_fn': ['model_dir'],
+    'input_fn': ['request_body', 'content_type'],
+    'predict_fn': ['input_data', 'model'],
+    'output_fn': ['prediction', 'accept']
+}
 
 
 class IdCounter:
@@ -188,3 +199,59 @@ def find_model_file(model_dir: str, extensions: List[str]) -> Optional[str]:
         )
 
     return all_matches[0] if all_matches else None
+
+
+def _validate_sagemaker_function(
+        module, func_name: str,
+        expected_params: List[str]) -> Optional[Callable]:
+    """
+    Validate that function exists and has correct signature
+    Returns the function if valid, None otherwise
+    """
+    if not hasattr(module, func_name):
+        return None
+
+    func = getattr(module, func_name)
+    if not callable(func):
+        return None
+
+    try:
+        sig = inspect.signature(func)
+        param_names = list(sig.parameters.keys())
+
+        # Check parameter count and names match exactly
+        if param_names == expected_params:
+            return func
+    except (ValueError, TypeError):
+        # Handle cases where signature inspection fails
+        pass
+
+    return None
+
+
+def get_sagemaker_function(model_dir: str,
+                           func_name: str) -> Optional[Callable]:
+    """
+    Load and validate SageMaker-style formatter function from model.py
+    
+    :param model_dir: model directory containing model.py
+    :param func_name: SageMaker function name (model_fn, input_fn, predict_fn, output_fn)
+    :return: Validated function or None if not found/invalid
+    """
+
+    if func_name not in SAGEMAKER_SIGNATURES:
+        return None
+
+    try:
+        service = load_model_service(model_dir, "model.py", -1)
+        if has_function_in_module(service.module, func_name):
+            func = getattr(service.module, func_name)
+            # Optional: validate signature
+            expected_params = SAGEMAKER_SIGNATURES[func_name]
+            if _validate_sagemaker_function(service.module, func_name,
+                                            expected_params):
+                return func
+
+    except Exception as e:
+        logging.debug(f"Failed to load {func_name} from model.py: {e}")
+        return None