[serving] Read x-synchronus and x-starting-token from input payload

deepjavalibrary · Apr 18, 2023 · 3f8cef3 · 3f8cef3
1 parent 0776f3b
commit 3f8cef3
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/serving/docker/Dockerfile b/serving/docker/Dockerfile
@@ -61,7 +61,7 @@ RUN scripts/install_python.sh && \
     echo "${djl_version} cpufull" > /opt/djl/bin/telemetry && \
     djl-serving -i ai.djl.mxnet:mxnet-native-mkl:1.9.1:linux-x86_64 && \
     djl-serving -i ai.djl.pytorch:pytorch-native-cpu:$torch_version:linux-x86_64 && \
-    djl-serving -i ai.djl.tensorflow:tensorflow-native-cpu:2.7.4:linux-x86_64 && \
+    djl-serving -i ai.djl.tensorflow:tensorflow-native-cpu:2.10.1:linux-x86_64 && \
     scripts/patch_oss_dlc.sh python && \
     rm -rf /opt/djl/logs && \
     chown -R djl:djl /opt/djl && \

diff --git a/serving/src/main/java/ai/djl/serving/http/InferenceRequestHandler.java b/serving/src/main/java/ai/djl/serving/http/InferenceRequestHandler.java
@@ -71,6 +71,7 @@ public class InferenceRequestHandler extends HttpRequestHandler {
     private static final String X_STARTING_TOKEN = "x-starting-token";
     private static final String X_NEXT_TOKEN = "x-next-token";
     private static final String X_MAX_ITEMS = "x-max-items";
+    private static final String X_CUSTOM_ATTRIBUTES = "X-Amzn-SageMaker-Custom-Attributes";
 
     private RequestParser requestParser;
 
@@ -292,11 +293,13 @@ void runJob(
             pending.setMessage("The model result is not yet available");
             pending.setCode(202);
             pending.addProperty(X_NEXT_TOKEN, nextToken);
+            pending.addProperty(X_CUSTOM_ATTRIBUTES, X_NEXT_TOKEN + '=' + nextToken);
             cache.put(nextToken, pending);
 
             // Send back token to user
             Output out = new Output();
             out.addProperty(X_NEXT_TOKEN, nextToken);
+            out.addProperty(X_CUSTOM_ATTRIBUTES, X_NEXT_TOKEN + '=' + nextToken);
             sendOutput(out, ctx);
 
             // Run model

diff --git a/serving/src/main/java/ai/djl/serving/http/RequestParser.java b/serving/src/main/java/ai/djl/serving/http/RequestParser.java
@@ -60,7 +60,15 @@ public Input parseRequest(FullHttpRequest req, QueryStringDecoder decoder) {
 
         for (Map.Entry<String, String> entry : req.headers().entries()) {
             String key = entry.getKey();
-            if (!HttpHeaderNames.CONTENT_TYPE.contentEqualsIgnoreCase(key)) {
+            if ("X-Amzn-SageMaker-Custom-Attributes".equalsIgnoreCase(key)) {
+                String[] tokens = entry.getValue().split(";");
+                for (String token : tokens) {
+                    String[] pair = token.split("=", 2);
+                    if (pair.length == 2) {
+                        input.addProperty(pair[0], pair[1]);
+                    }
+                }
+            } else if (!HttpHeaderNames.CONTENT_TYPE.contentEqualsIgnoreCase(key)) {
                 input.addProperty(key, entry.getValue());
             }
         }