[serving] make input header key case-insensitive (#1134)

1. make it easy to get content-type header 2. improve error handling Change-Id: I12a1ea9f12e3565649d0f218870fff6d111bc2fe
deepjavalibrary · Aug 2, 2021 · 14bf240 · 14bf240
1 parent 91cce8e
commit 14bf240
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 7 deletions.
diff --git a/api/src/main/java/ai/djl/modality/Input.java b/api/src/main/java/ai/djl/modality/Input.java
@@ -13,6 +13,7 @@
 package ai.djl.modality;
 
 import ai.djl.util.PairList;
+import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -82,7 +83,7 @@ public void addProperty(String key, String value) {
      * @return the value to which the specified key is mapped
      */
     public String getProperty(String key, String defaultValue) {
-        return properties.getOrDefault(key, defaultValue);
+        return properties.getOrDefault(key.toLowerCase(Locale.ROOT), defaultValue);
     }
 
     /**

diff --git a/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java b/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
@@ -247,6 +247,7 @@ public Output processOutput(TranslatorContext ctx, NDList list) throws Exception
             } else {
                 output.setContent(JsonUtils.GSON_PRETTY.toJson(obj) + '\n');
             }
+            output.addProperty("Content-Type", "application/json");
             return output;
         }
 
@@ -262,8 +263,12 @@ public NDList processInput(TranslatorContext ctx, Input input) throws Exception
             if (data == null) {
                 data = input.getContent().valueAt(0);
             }
-            Image image = factory.fromInputStream(new ByteArrayInputStream(data));
-            return translator.processInput(ctx, image);
+            try {
+                Image image = factory.fromInputStream(new ByteArrayInputStream(data));
+                return translator.processInput(ctx, image);
+            } catch (IOException e) {
+                throw new TranslateException("Input is not an Image data type", e);
+            }
         }
 
         /** {@inheritDoc} */
@@ -307,6 +312,7 @@ public Output processOutput(TranslatorContext ctx, NDList list) {
             Input input = (Input) ctx.getAttachment("input");
             Output output = new Output(input.getRequestId(), 200, "OK");
             output.setContent(list.encode());
+            output.addProperty("Content-Type", "tensor/ndlist");
             return output;
         }
     }

diff --git a/serving/serving/src/main/java/ai/djl/serving/http/RequestParser.java b/serving/serving/src/main/java/ai/djl/serving/http/RequestParser.java
@@ -24,6 +24,7 @@
 import io.netty.handler.codec.http.multipart.HttpPostRequestDecoder;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -58,9 +59,13 @@ public Input parseRequest(
             }
         }
 
-        CharSequence contentType = HttpUtil.getMimeType(req);
         for (Map.Entry<String, String> entry : req.headers().entries()) {
-            input.addProperty(entry.getKey(), entry.getValue());
+            input.addProperty(entry.getKey().toLowerCase(Locale.ROOT), entry.getValue());
+        }
+        CharSequence contentType = HttpUtil.getMimeType(req);
+        if (contentType != null) {
+            // use normalized content type
+            input.addProperty("content-type", contentType.toString());
         }
 
         if (HttpPostRequestDecoder.isMultipart(req)

diff --git a/serving/serving/src/main/java/ai/djl/serving/wlm/WorkerThread.java b/serving/serving/src/main/java/ai/djl/serving/wlm/WorkerThread.java
@@ -12,6 +12,7 @@
  */
 package ai.djl.serving.wlm;
 
+import ai.djl.engine.EngineException;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
@@ -67,13 +68,17 @@ public void run() {
         currentThread.set(thread);
         this.state = WorkerState.WORKER_STARTED;
         List<Input> req = null;
+        String errorMessage = "Worker shutting down";
         try {
             while (isRunning() && !aggregator.isFinished()) {
                 req = aggregator.getRequest();
                 if (req != null && !req.isEmpty()) {
                     try {
                         List<Output> reply = predictor.batchPredict(req);
                         aggregator.sendResponse(reply);
+                    } catch (EngineException e) {
+                        logger.warn("Failed to predict", e);
+                        aggregator.sendError(HttpResponseStatus.INTERNAL_SERVER_ERROR, e);
                     } catch (TranslateException e) {
                         logger.warn("Failed to predict", e);
                         aggregator.sendError(HttpResponseStatus.BAD_REQUEST, e);
@@ -85,12 +90,13 @@ public void run() {
             logger.debug("Shutting down the thread .. Scaling down.");
         } catch (Throwable t) {
             logger.error("Server error", t);
+            errorMessage = t.getMessage();
         } finally {
             logger.debug("Shutting down worker thread .. {}", currentThread.get().getName());
             currentThread.set(null);
             shutdown(WorkerState.WORKER_STOPPED);
             if (req != null) {
-                Exception e = new InternalServerException("Server shutting down");
+                Exception e = new InternalServerException(errorMessage);
                 aggregator.sendError(HttpResponseStatus.INTERNAL_SERVER_ERROR, e);
             }
         }
@@ -122,7 +128,7 @@ public void shutdown(WorkerState state) {
         Thread thread = currentThread.getAndSet(null);
         if (thread != null) {
             thread.interrupt();
-            Exception e = new InternalServerException("Server shutting down");
+            Exception e = new InternalServerException("Worker shutting down");
             aggregator.sendError(HttpResponseStatus.INTERNAL_SERVER_ERROR, e);
         }
         predictor.close();