Issue 1120

docs/configuration.md

@@ -284,6 +284,7 @@ the backend workers convert "Bytearray to utf-8 string" when the Content-Type of
 * `model_server_home` : Torchserve home directory. 
 * `max_request_size` : The maximum allowable request size that the Torchserve accepts, in bytes. Default: 6553500
 * `max_response_size` : The maximum allowable response size that the Torchserve sends, in bytes. Default: 6553500
+* `limit_max_image_pixels` : Default value is true (Use default [PIL.Image.MAX_IMAGE_PIXELS](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.MAX_IMAGE_PIXELS)). If this is set to "false", set PIL.Image.MAX_IMAGE_PIXELS = None in backend default vision handler for large image payload. 
 * `allowed_urls` : Comma separated regex of allowed source URL(s) from where models can be registered. Default: "file://.*|http(s)?://.*" (all URLs and local file system)
 eg : To allow base URLs `https://s3.amazonaws.com/` and `https://torchserve.pytorch.org/` use the following regex string `allowed_urls=https://s3.amazonaws.com/.*,https://torchserve.pytorch.org/.*`
 * `workflow_store` : Path of workflow store directory. Defaults to model store directory.

frontend/server/src/main/java/org/pytorch/serve/http/api/rest/InferenceRequestHandler.java

@@ -25,6 +25,7 @@
 import org.pytorch.serve.openapi.OpenApiUtils;
 import org.pytorch.serve.servingsdk.ModelServerEndpoint;
 import org.pytorch.serve.util.ApiUtils;
+import org.pytorch.serve.util.ConfigManager;
 import org.pytorch.serve.util.NettyUtils;
 import org.pytorch.serve.util.messages.InputParameter;
 import org.pytorch.serve.util.messages.RequestInput;
@@ -252,7 +253,8 @@ private static RequestInput parseRequest(
         if (HttpPostRequestDecoder.isMultipart(req)
                 || HttpHeaderValues.APPLICATION_X_WWW_FORM_URLENCODED.contentEqualsIgnoreCase(
                         contentType)) {
-            HttpDataFactory factory = new DefaultHttpDataFactory(6553500);
+            HttpDataFactory factory =
+                    new DefaultHttpDataFactory(ConfigManager.getInstance().getMaxRequestSize());
             HttpPostRequestDecoder form = new HttpPostRequestDecoder(factory, req);
             try {
                 while (form.hasNext()) {

frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java

@@ -78,6 +78,7 @@ public final class ConfigManager {
     private static final String TS_PRIVATE_KEY_FILE = "private_key_file";
     private static final String TS_MAX_REQUEST_SIZE = "max_request_size";
     private static final String TS_MAX_RESPONSE_SIZE = "max_response_size";
+    private static final String TS_LIMIT_MAX_IMAGE_PIXELS = "limit_max_image_pixels";
     private static final String TS_DEFAULT_SERVICE_HANDLER = "default_service_handler";
     private static final String TS_SERVICE_ENVELOPE = "service_envelope";
     private static final String TS_MODEL_SERVER_HOME = "model_server_home";
@@ -604,6 +605,8 @@ public String dumpConfigurations() {
                 + prop.getProperty(TS_MAX_RESPONSE_SIZE, "6553500")
                 + "\nMaximum Request Size: "
                 + prop.getProperty(TS_MAX_REQUEST_SIZE, "6553500")
+                + "\nLimit Maximum Image Pixels: "
+                + prop.getProperty(TS_LIMIT_MAX_IMAGE_PIXELS, "true")
                 + "\nPrefer direct buffer: "
                 + prop.getProperty(TS_PREFER_DIRECT_BUFFER, "false")
                 + "\nAllowed Urls: "
@@ -636,6 +639,10 @@ public int getMaxRequestSize() {
         return getIntProperty(TS_MAX_REQUEST_SIZE, 6553500);
     }
 
+    public boolean isLimitMaxImagePixels() {
+        return Boolean.parseBoolean(prop.getProperty(TS_LIMIT_MAX_IMAGE_PIXELS, "true"));
+    }
+
     public void setProperty(String key, String value) {
         prop.setProperty(key, value);
     }

frontend/server/src/main/java/org/pytorch/serve/util/codec/ModelRequestEncoder.java

@@ -60,6 +60,7 @@ protected void encode(ChannelHandlerContext ctx, BaseModelRequest msg, ByteBuf o
             out.writeInt(buf.length);
             out.writeBytes(buf);
 
+            out.writeBoolean(request.isLimitMaxImagePixels());
         } else if (msg instanceof ModelInferenceRequest) {
             out.writeByte('I');
             ModelInferenceRequest request = (ModelInferenceRequest) msg;

frontend/server/src/main/java/org/pytorch/serve/util/messages/ModelLoadModelRequest.java

@@ -1,5 +1,6 @@
 package org.pytorch.serve.util.messages;
 
+import org.pytorch.serve.util.ConfigManager;
 import org.pytorch.serve.wlm.Model;
 
 public class ModelLoadModelRequest extends BaseModelRequest {
@@ -14,6 +15,7 @@ public class ModelLoadModelRequest extends BaseModelRequest {
     private String envelope;
     private int batchSize;
     private int gpuId;
+    private boolean limitMaxImagePixels;
 
     public ModelLoadModelRequest(Model model, int gpuId) {
         super(WorkerCommands.LOAD, model.getModelName());
@@ -22,6 +24,7 @@ public ModelLoadModelRequest(Model model, int gpuId) {
         handler = model.getModelArchive().getManifest().getModel().getHandler();
         envelope = model.getModelArchive().getManifest().getModel().getEnvelope();
         batchSize = model.getBatchSize();
+        limitMaxImagePixels = ConfigManager.getInstance().isLimitMaxImagePixels();
     }
 
     public String getModelPath() {
@@ -43,4 +46,8 @@ public int getBatchSize() {
     public int getGpuId() {
         return gpuId;
     }
+
+    public boolean isLimitMaxImagePixels() {
+        return limitMaxImagePixels;
+    }
 }

frontend/server/src/main/java/org/pytorch/serve/wlm/WorkerThread.java

@@ -182,6 +182,7 @@ public void run() {
                 req = aggregator.getRequest(workerId, state);
 
                 long wtStartTime = System.currentTimeMillis();
+                logger.info("Flushing req. to backend at: " + wtStartTime);
                 backendChannel.writeAndFlush(req).sync();
 
                 long begin = System.currentTimeMillis();

frontend/server/src/main/java/org/pytorch/serve/workflow/api/http/WorkflowInferenceRequestHandler.java

@@ -18,6 +18,7 @@
 import org.pytorch.serve.http.HttpRequestHandlerChain;
 import org.pytorch.serve.http.ResourceNotFoundException;
 import org.pytorch.serve.http.StatusResponse;
+import org.pytorch.serve.util.ConfigManager;
 import org.pytorch.serve.util.NettyUtils;
 import org.pytorch.serve.util.messages.InputParameter;
 import org.pytorch.serve.util.messages.RequestInput;
@@ -96,7 +97,8 @@ private static RequestInput parseRequest(ChannelHandlerContext ctx, FullHttpRequ
         if (HttpPostRequestDecoder.isMultipart(req)
                 || HttpHeaderValues.APPLICATION_X_WWW_FORM_URLENCODED.contentEqualsIgnoreCase(
                         contentType)) {
-            HttpDataFactory factory = new DefaultHttpDataFactory(6553500);
+            HttpDataFactory factory =
+                    new DefaultHttpDataFactory(ConfigManager.getInstance().getMaxRequestSize());
             HttpPostRequestDecoder form = new HttpPostRequestDecoder(factory, req);
             try {
                 while (form.hasNext()) {

ts/context.py

@@ -11,19 +11,21 @@ class Context(object):
     Some fixed during load times and some
     """
 
-    def __init__(self, model_name, model_dir, manifest, batch_size, gpu, mms_version):
+    def __init__(self, model_name, model_dir, manifest, batch_size, gpu, mms_version, limit_max_image_pixels=True):
         self.model_name = model_name
         self.manifest = manifest
         self._system_properties = {
             "model_dir": model_dir,
             "gpu_id": gpu,
             "batch_size": batch_size,
             "server_name": "MMS",
-            "server_version": mms_version
+            "server_version": mms_version,
+            "limit_max_image_pixels": limit_max_image_pixels,
         }
         self.request_ids = None
         self.request_processor = None
         self._metrics = None
+        self._limit_max_image_pixels = True
 
     @property
     def system_properties(self):

ts/model_loader.py

@@ -34,7 +34,7 @@ class ModelLoader(object):
     __metaclass__ = ABCMeta
 
     @abstractmethod
-    def load(self, model_name, model_dir, handler, gpu_id, batch_size, envelope=None):
+    def load(self, model_name, model_dir, handler, gpu_id, batch_size, envelope=None, limit_max_image_pixels=True):
         """
         Load model from file.
 
@@ -44,6 +44,7 @@ def load(self, model_name, model_dir, handler, gpu_id, batch_size, envelope=None
         :param gpu_id:
         :param batch_size:
         :param envelope:
+        :param limit_max_image_pixels:
         :return: Model
         """
         # pylint: disable=unnecessary-pass
@@ -55,7 +56,7 @@ class TsModelLoader(ModelLoader):
     TorchServe 1.0 Model Loader
     """
 
-    def load(self, model_name, model_dir, handler, gpu_id, batch_size, envelope=None):
+    def load(self, model_name, model_dir, handler, gpu_id, batch_size, envelope=None, limit_max_image_pixels=True):
         """
         Load TorchServe 1.0 model from file.
 
@@ -65,6 +66,7 @@ def load(self, model_name, model_dir, handler, gpu_id, batch_size, envelope=None
         :param gpu_id:
         :param batch_size:
         :param envelope:
+        :param limit_max_image_pixels:
         :return:
         """
         logging.debug("Loading model - working dir: %s", os.getcwd())
@@ -89,7 +91,7 @@ def load(self, model_name, model_dir, handler, gpu_id, batch_size, envelope=None
 
         if hasattr(module, function_name):
             entry_point = getattr(module, function_name)
-            service = Service(model_name, model_dir, manifest, entry_point, gpu_id, batch_size)
+            service = Service(model_name, model_dir, manifest, entry_point, gpu_id, batch_size, limit_max_image_pixels)
 
         envelope_class = None
         if envelope is not None:
@@ -105,7 +107,7 @@ def load(self, model_name, model_dir, handler, gpu_id, batch_size, envelope=None
             envelope_instance = envelope_class(entry_point)
             entry_point = envelope_instance.handle
 
-        service = Service(model_name, model_dir, manifest, entry_point, gpu_id, batch_size)
+        service = Service(model_name, model_dir, manifest, entry_point, gpu_id, batch_size, limit_max_image_pixels)
         service.context.metrics = metrics
         initialize_fn(service.context)
 

ts/model_service_worker.py

@@ -65,6 +65,7 @@ def load_model(load_model_request):
             "handler" : service handler entry point if provided, string
             "envelope" : name of wrapper/unwrapper of request data if provided, string
             "batchSize" : batch size, int
+            "limitMaxImagePixels": limit pillow image max_image_pixels, bool
         }
 
         :param load_model_request:
@@ -86,8 +87,13 @@ def load_model(load_model_request):
             if "gpu" in load_model_request:
                 gpu = int(load_model_request["gpu"])
 
+            limit_max_image_pixels = True
+            if "limitMaxImagePixels" in load_model_request:
+                limit_max_image_pixels = bool(load_model_request["limitMaxImagePixels"])
+
             model_loader = ModelLoaderFactory.get_model_loader()
-            service = model_loader.load(model_name, model_dir, handler, gpu, batch_size, envelope)
+            service = model_loader.load(model_name, model_dir, handler, gpu,
+                                        batch_size, envelope, limit_max_image_pixels)
 
             logging.debug("Model %s loaded.", model_name)
 

ts/protocol/otf_message_handler.py

@@ -11,8 +11,10 @@
 import io
 from builtins import bytearray
 from builtins import bytes
+import time
 import torch
 
+bool_size = 1
 int_size = 4
 END_OF_LIST = -1
 LOAD_MSG = b'L'
@@ -32,6 +34,7 @@ def retrieve_msg(conn):
         msg = _retrieve_load_msg(conn)
     elif cmd == PREDICT_MSG:
         msg = _retrieve_inference_msg(conn)
+        logging.info("Backend received inference at: %d", time.time())
     else:
         raise ValueError("Invalid command: {}".format(cmd))
 
@@ -173,6 +176,9 @@ def _retrieve_int(conn):
     data = _retrieve_buffer(conn, int_size)
     return struct.unpack("!i", data)[0]
 
+def _retrieve_bool(conn):
+    data = _retrieve_buffer(conn, bool_size)
+    return struct.unpack("!?", data)[0]
 
 def _retrieve_load_msg(conn):
     """
@@ -184,6 +190,7 @@ def _retrieve_load_msg(conn):
     | int batch-size length |
     | int handler length | handler value |
     | int gpu id |
+    | bool limitMaxImagePixels |
 
     :param conn:
     :return:
@@ -202,6 +209,7 @@ def _retrieve_load_msg(conn):
 
     length = _retrieve_int(conn)
     msg["envelope"] = _retrieve_buffer(conn, length)
+    msg["limitMaxImagePixels"] = _retrieve_bool(conn)
 
     return msg
 

ts/service.py

@@ -22,8 +22,9 @@ class Service(object):
     Wrapper for custom entry_point
     """
 
-    def __init__(self, model_name, model_dir, manifest, entry_point, gpu, batch_size):
-        self._context = Context(model_name, model_dir, manifest, batch_size, gpu, ts.__version__)
+    def __init__(self, model_name, model_dir, manifest, entry_point, gpu, batch_size, limit_max_image_pixels=True):
+        self._context = Context(model_name, model_dir, manifest,
+                                batch_size, gpu, ts.__version__, limit_max_image_pixels)
         self._entry_point = entry_point
 
     @property

ts/tests/unit_tests/test_otf_codec_protocol.py

@@ -34,7 +34,8 @@ def test_retrieve_msg_unknown(self, socket_patches):
     def test_retrieve_msg_load_gpu(self, socket_patches):
         expected = {"modelName": b"model_name", "modelPath": b"model_path",
                     "batchSize": 1, "handler": b"handler", "gpu": 1,
-                    "envelope": b"envelope"}
+                    "envelope": b"envelope",
+                    "limitMaxImagePixels": True}
 
         socket_patches.socket.recv.side_effect = [
             b"L",
@@ -43,7 +44,8 @@ def test_retrieve_msg_load_gpu(self, socket_patches):
             b"\x00\x00\x00\x01",
             b"\x00\x00\x00\x07", b"handler",
             b"\x00\x00\x00\x01",
-            b"\x00\x00\x00\x08", b"envelope"
+            b"\x00\x00\x00\x08", b"envelope",
+            b"\x01"
         ]
         cmd, ret = codec.retrieve_msg(socket_patches.socket)
 
@@ -53,7 +55,8 @@ def test_retrieve_msg_load_gpu(self, socket_patches):
     def test_retrieve_msg_load_no_gpu(self, socket_patches):
         expected = {"modelName": b"model_name", "modelPath": b"model_path",
                     "batchSize": 1, "handler": b"handler",
-                    "envelope": b"envelope"}
+                    "envelope": b"envelope",
+                    "limitMaxImagePixels": True}
 
         socket_patches.socket.recv.side_effect = [
             b"L",
@@ -62,7 +65,8 @@ def test_retrieve_msg_load_no_gpu(self, socket_patches):
             b"\x00\x00\x00\x01",
             b"\x00\x00\x00\x07", b"handler",
             b"\xFF\xFF\xFF\xFF",
-            b"\x00\x00\x00\x08", b"envelope"
+            b"\x00\x00\x00\x08", b"envelope",
+            b"\x01"
         ]
         cmd, ret = codec.retrieve_msg(socket_patches.socket)
 

ts/torch_handler/vision_handler.py

@@ -20,6 +20,9 @@ def initialize(self, context):
         super().initialize(context)
         self.ig = IntegratedGradients(self.model)
         self.initialized = True
+        properties = context.system_properties
+        if not properties.get("limit_max_image_pixels") :
+            Image.MAX_IMAGE_PIXELS = None
 
     def preprocess(self, data):
         """The preprocess function of MNIST program converts the input data to a float tensor