Repository: deepjavalibrary/djl
Pull request #1258 (2021-09-29 16:26:23)

[api] refactor input/output interface
Change-Id: I695054aa63a8366b1b1160953faa37c608bc2fee

## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/api/src/main/java/ai/djl/modality/Input.java b/api/src/main/java/ai/djl/modality/Input.java
index 729b8ec4e..5b6525b95 100644
--- a/api/src/main/java/ai/djl/modality/Input.java
+++ b/api/src/main/java/ai/djl/modality/Input.java
@@ -12,7 +12,12 @@
  */
 package ai.djl.modality;
 
+import ai.djl.ndarray.BytesSupplier;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
 import ai.djl.util.PairList;
+import java.nio.charset.StandardCharsets;
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -20,30 +25,15 @@ import java.util.concurrent.ConcurrentHashMap;
 /** A class stores the generic input data for inference. */
 public class Input {
 
-    private String requestId;
     private Map<String, String> properties;
-    private PairList<String, byte[]> content;
+    private PairList<String, BytesSupplier> content;
 
-    /**
-     * Constructs a {@code Input} with specified {@code requestId}.
-     *
-     * @param requestId the requestId of the input
-     */
-    public Input(String requestId) {
-        this.requestId = requestId;
+    /** Constructs a new {@code Input} instance. */
+    public Input() {
         properties = new ConcurrentHashMap<>();
         content = new PairList<>();
     }
 
-    /**
-     * Returns the requestId of the input.
-     *
-     * @return the requestId of the input
-     */
-    public String getRequestId() {
-        return requestId;
-    }
-
     /**
      * Returns the properties of the input.
      *
@@ -69,9 +59,6 @@ public class Input {
      * @param value value to be added with the specified key
      */
     public void addProperty(String key, String value) {
-        if (properties == null) {
-            properties = new ConcurrentHashMap<>();
-        }
         properties.put(key, value);
     }
 
@@ -93,7 +80,7 @@ public class Input {
      *
      * @return the content of the input
      */
-    public PairList<String, byte[]> getContent() {
+    public PairList<String, BytesSupplier> getContent() {
         return content;
     }
 
@@ -102,7 +89,7 @@ public class Input {
      *
      * @param content the content of the input
      */
-    public void setContent(PairList<String, byte[]> content) {
+    public void setContent(PairList<String, BytesSupplier> content) {
         this.content = content;
     }
 
@@ -111,8 +98,26 @@ public class Input {
      *
      * @param data data to be added
      */
-    public void addData(byte[] data) {
-        addData(null, data);
+    public void add(byte[] data) {
+        add(BytesSupplier.wrap(data));
+    }
+
+    /**
+     * Appends an item at the end of the input.
+     *
+     * @param data data to be added
+     */
+    public void add(String data) {
+        add(BytesSupplier.wrap(data.getBytes(StandardCharsets.UTF_8)));
+    }
+
+    /**
+     * Appends an item at the end of the input.
+     *
+     * @param data data to be added
+     */
+    public void add(BytesSupplier data) {
+        add(null, data);
     }
 
     /**
@@ -121,10 +126,27 @@ public class Input {
      * @param key key with which the specified data is to be added
      * @param data data to be added with the specified key
      */
-    public void addData(String key, byte[] data) {
-        if (content == null) {
-            content = new PairList<>();
-        }
+    public void add(String key, byte[] data) {
+        add(key, BytesSupplier.wrap(data));
+    }
+
+    /**
+     * Adds a key/value pair to the input content.
+     *
+     * @param key key with which the specified data is to be added
+     * @param data data to be added with the specified key
+     */
+    public void add(String key, String data) {
+        add(key, BytesSupplier.wrap(data));
+    }
+
+    /**
+     * Adds a key/value pair to the input content.
+     *
+     * @param key key with which the specified data is to be added
+     * @param data data to be added with the specified key
+     */
+    public void add(String key, BytesSupplier data) {
         content.add(key, data);
     }
 
@@ -132,12 +154,167 @@ public class Input {
      * Inserts the specified element at the specified position in the input.
      *
      * @param index the index at which the specified element is to be inserted
+     * @param key key with which the specified data is to be added
      * @param data data to be added with the specified key
      */
-    public void addData(int index, byte[] data) {
-        if (content == null) {
-            content = new PairList<>();
+    public void add(int index, String key, BytesSupplier data) {
+        content.add(index, key, data);
+    }
+
+    /**
+     * Returns the default data item.
+     *
+     * @return the default data item
+     */
+    public BytesSupplier getData() {
+        BytesSupplier data = get("data");
+        if (data == null) {
+            return get(0);
+        }
+        return data;
+    }
+
+    /**
+     * Returns the default data as {@code NDList}.
+     *
+     * @param manager {@link NDManager} used to create this {@code NDArray}
+     * @return the default data as {@code NDList}
+     */
+    public NDList getDataAsNDList(NDManager manager) {
+        int index = content.indexOf("data");
+        if (index < 0) {
+            index = 0;
+        }
+        return getAsNDList(manager, index);
+    }
+
+    /**
+     * Returns the element for the first key found in the {@code Input}.
+     *
+     * @param key the key of the element to get
+     * @return the element for the first key found in the {@code Input}
+     */
+    public BytesSupplier get(String key) {
+        return content.get(key);
+    }
+
+    /**
+     * Returns the element at the specified position in the {@code Input}.
+     *
+     * @param index the index of the element to return
+     * @return the element at the specified position in the {@code Input}
+     */
+    public BytesSupplier get(int index) {
+        return content.valueAt(index);
+    }
+
+    /**
+     * Returns the value as {@code byte[]} for the first key found in the {@code Input}.
+     *
+     * @param key the key of the element to get
+     * @return the value as {@code byte[]} for the first key found in the {@code Input}
+     */
+    public byte[] getAsBytes(String key) {
+        BytesSupplier data = content.get(key);
+        if (data == null) {
+            return null;
+        }
+        return data.getAsBytes();
+    }
+
+    /**
+     * Returns the value as {@code byte[]} at the specified position in the {@code Input}.
+     *
+     * @param index the index of the element to return
+     * @return the value as {@code byte[]} at the specified position in the {@code Input}
+     */
+    public byte[] getAsBytes(int index) {
+        return content.valueAt(index).getAsBytes();
+    }
+
+    /**
+     * Returns the value as {@code byte[]} for the first key found in the {@code Input}.
+     *
+     * @param key the key of the element to get
+     * @return the value as {@code byte[]} for the first key found in the {@code Input}
+     */
+    public String getAsString(String key) {
+        BytesSupplier data = content.get(key);
+        if (data == null) {
+            return null;
+        }
+        return data.getAsString();
+    }
+
+    /**
+     * Returns the value as {@code byte[]} at the specified position in the {@code Input}.
+     *
+     * @param index the index of the element to return
+     * @return the value as {@code byte[]} at the specified position in the {@code Input}
+     */
+    public String getAsString(int index) {
+        return content.valueAt(index).getAsString();
+    }
+
+    /**
+     * Returns the value as {@code NDArray} for the first key found in the {@code Input}.
+     *
+     * @param manager {@link NDManager} used to create this {@code NDArray}
+     * @param key the key of the element to get
+     * @return the value as {@code NDArray} for the first key found in the {@code Input}
+     */
+    public NDArray getAsNDArray(NDManager manager, String key) {
+        int index = content.indexOf(key);
+        if (index < 0) {
+            return null;
+        }
+        return getAsNDArray(manager, index);
+    }
+
+    /**
+     * Returns the value as {@code NDArray} at the specified position in the {@code Input}.
+     *
+     * @param manager {@link NDManager} used to create this {@code NDArray}
+     * @param index the index of the element to return
+     * @return the value as {@code NDArray} at the specified position in the {@code Input}
+     */
+    public NDArray getAsNDArray(NDManager manager, int index) {
+        BytesSupplier data = content.valueAt(index);
+        if (data instanceof NDArray) {
+            return (NDArray) data;
+        }
+        return NDArray.decode(manager, data.getAsBytes());
+    }
+
+    /**
+     * Returns the value as {@code NDList} for the first key found in the {@code Input}.
+     *
+     * @param manager {@link NDManager} used to create this {@code NDArray}
+     * @param key the key of the element to get
+     * @return the value as {@code NDList} for the first key found in the {@code Input}
+     */
+    public NDList getAsNDList(NDManager manager, String key) {
+        int index = content.indexOf(key);
+        if (index < 0) {
+            return null;
+        }
+        return getAsNDList(manager, index);
+    }
+
+    /**
+     * Returns the value as {@code NDList} at the specified position in the {@code Input}.
+     *
+     * @param manager {@link NDManager} used to create this {@code NDArray}
+     * @param index the index of the element to return
+     * @return the value as {@code NDList} at the specified position in the {@code Input}
+     */
+    public NDList getAsNDList(NDManager manager, int index) {
+        BytesSupplier data = content.valueAt(index);
+        if (data instanceof NDList) {
+            return (NDList) data;
+        } else if (data instanceof NDArray) {
+            return new NDList((NDArray) data);
         }
-        content.add(index, null, data);
+        return NDList.decode(manager, data.getAsBytes());
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/Output.java b/api/src/main/java/ai/djl/modality/Output.java
index be8f942c7..a22c18e00 100644
--- a/api/src/main/java/ai/djl/modality/Output.java
+++ b/api/src/main/java/ai/djl/modality/Output.java
@@ -12,19 +12,11 @@
  */
 package ai.djl.modality;
 
-import java.nio.charset.StandardCharsets;
-import java.util.Collections;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-
 /** A class stores the generic inference results. */
-public class Output {
+public class Output extends Input {
 
-    private String requestId;
     private int code;
     private String message;
-    private Map<String, String> properties;
-    private byte[] content;
 
     /**
      * Constructs a {@code Output} with specified {@code requestId}, {@code code} and {@code
@@ -38,24 +30,6 @@ public class Output {
         this.message = message;
     }
 
-    /**
-     * Sets the requestId of the output.
-     *
-     * @param requestId the requestId of the output
-     */
-    public void setRequestId(String requestId) {
-        this.requestId = requestId;
-    }
-
-    /**
-     * Returns the requestId of the output.
-     *
-     * @return the requestId of the output
-     */
-    public String getRequestId() {
-        return requestId;
-    }
-
     /**
      * Returns the status code of the output.
      *
@@ -91,65 +65,4 @@ public class Output {
     public void setMessage(String message) {
         this.message = message;
     }
-
-    /**
-     * Returns the properties of the output.
-     *
-     * @return the properties of the output
-     */
-    public Map<String, String> getProperties() {
-        if (properties == null) {
-            return Collections.emptyMap();
-        }
-        return properties;
-    }
-
-    /**
-     * Sets the properties of the output.
-     *
-     * @param properties the properties of the output
-     */
-    public void setProperties(Map<String, String> properties) {
-        this.properties = properties;
-    }
-
-    /**
-     * Adds a property to the output.
-     *
-     * @param key key with which the specified value is to be added
-     * @param value value to be added with the specified key
-     */
-    public void addProperty(String key, String value) {
-        if (properties == null) {
-            properties = new ConcurrentHashMap<>();
-        }
-        properties.put(key, value);
-    }
-
-    /**
-     * Returns the content of the input.
-     *
-     * @return the content of the input
-     */
-    public byte[] getContent() {
-        return content;
-    }
-
-    /**
-     * Sets the content of the input.
-     *
-     * @param content the content of the input
-     */
-    public void setContent(byte[] content) {
-        this.content = content;
-    }
-
-    /**
-     * Sets the content of the input with string value.
-     *
-     * @param content the content of the input in string
-     */
-    public void setContent(String content) {
-        this.content = content.getBytes(StandardCharsets.UTF_8);
-    }
 }
diff --git a/api/src/main/java/ai/djl/ndarray/BytesSupplier.java b/api/src/main/java/ai/djl/ndarray/BytesSupplier.java
new file mode 100644
index 000000000..a9753b975
--- /dev/null
+++ b/api/src/main/java/ai/djl/ndarray/BytesSupplier.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.ndarray;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+
+/** Represents a supplier of {@code byte[]}. */
+public interface BytesSupplier {
+
+    /**
+     * Returns the {@code byte[]} presentation of the object.
+     *
+     * @return the {@code byte[]} presentation of the object
+     */
+    default byte[] getAsBytes() {
+        ByteBuffer bb = toByteBuffer();
+        if (bb.hasArray() && bb.remaining() == bb.capacity()) {
+            return bb.array();
+        }
+        byte[] buf = new byte[bb.remaining()];
+        bb.get(buf);
+        return buf;
+    }
+
+    /**
+     * Returns the {@code String} presentation of the object.
+     *
+     * @return the {@code String} presentation of the object
+     */
+    default String getAsString() {
+        return new String(getAsBytes(), StandardCharsets.UTF_8);
+    }
+
+    /**
+     * Returns the {@code ByteBuffer} presentation of the object.
+     *
+     * @return the {@code ByteBuffer} presentation of the object
+     */
+    ByteBuffer toByteBuffer();
+
+    /**
+     * Wraps a byte array into a {code BytesSupplier}.
+     *
+     * @param buf the byte array that will back this {code BytesSupplier}
+     * @return a {@code ByteBuffer}
+     */
+    static BytesSupplier wrap(byte[] buf) {
+        return new BytesSupplierImpl(buf);
+    }
+
+    /**
+     * Wraps a string into a {code BytesSupplier}.
+     *
+     * @param value the string that will back this {code BytesSupplier}
+     * @return a {@code ByteBuffer}
+     */
+    static BytesSupplier wrap(String value) {
+        return new BytesSupplierImpl(value);
+    }
+}
diff --git a/api/src/main/java/ai/djl/ndarray/BytesSupplierImpl.java b/api/src/main/java/ai/djl/ndarray/BytesSupplierImpl.java
new file mode 100644
index 000000000..21823bda0
--- /dev/null
+++ b/api/src/main/java/ai/djl/ndarray/BytesSupplierImpl.java
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.ndarray;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+
+class BytesSupplierImpl implements BytesSupplier {
+
+    private byte[] buf;
+    private String value;
+
+    BytesSupplierImpl(byte[] buf) {
+        this.buf = buf;
+    }
+
+    BytesSupplierImpl(String value) {
+        this.value = value;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public byte[] getAsBytes() {
+        if (buf == null) {
+            buf = value.getBytes(StandardCharsets.UTF_8);
+        }
+        return buf;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public String getAsString() {
+        if (value == null) {
+            value = new String(buf, StandardCharsets.UTF_8);
+        }
+        return value;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public ByteBuffer toByteBuffer() {
+        return ByteBuffer.wrap(getAsBytes());
+    }
+}
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index 00b90ff12..838d2e6fe 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -43,7 +43,7 @@ import java.util.stream.LongStream;
  * href="https://github.com/deepjavalibrary/djl/blob/master/docs/development/memory_management.md">NDArray
  * Memory Management Guide</a>
  */
-public interface NDArray extends NDResource {
+public interface NDArray extends NDResource, BytesSupplier {
 
     /**
      * Decodes {@code NDArray} from bytes.
@@ -396,13 +396,6 @@ public interface NDArray extends NDResource {
         }
     }
 
-    /**
-     * Converts this {@code NDArray} to a ByteBuffer.
-     *
-     * @return a ByteBuffer
-     */
-    ByteBuffer toByteBuffer();
-
     /**
      * Sets this {@code NDArray} value from {@link Buffer}.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDList.java b/api/src/main/java/ai/djl/ndarray/NDList.java
index 3ec54aaf4..8852501c1 100644
--- a/api/src/main/java/ai/djl/ndarray/NDList.java
+++ b/api/src/main/java/ai/djl/ndarray/NDList.java
@@ -22,6 +22,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.PushbackInputStream;
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -37,7 +38,7 @@ import java.util.zip.ZipOutputStream;
  *
  * @see NDArray
  */
-public class NDList extends ArrayList<NDArray> implements NDResource {
+public class NDList extends ArrayList<NDArray> implements NDResource, BytesSupplier {
 
     private static final long serialVersionUID = 1L;
 
@@ -343,6 +344,18 @@ public class NDList extends ArrayList<NDArray> implements NDResource {
         dos.flush();
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public byte[] getAsBytes() {
+        return encode();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public ByteBuffer toByteBuffer() {
+        return ByteBuffer.wrap(encode());
+    }
+
     /**
      * Gets all of shapes in the {@code NDList}.
      *
diff --git a/api/src/main/java/ai/djl/translate/NDArraySupplier.java b/api/src/main/java/ai/djl/translate/NDArraySupplier.java
index 846da4918..a3a5fdb3b 100644
--- a/api/src/main/java/ai/djl/translate/NDArraySupplier.java
+++ b/api/src/main/java/ai/djl/translate/NDArraySupplier.java
@@ -22,11 +22,10 @@ import ai.djl.ndarray.NDManager;
  * invoked.
  *
  * <p>This is a functional interface whose functional method is {@link #get(NDManager)}.
- *
- * @since 1.8
  */
 @FunctionalInterface
 public interface NDArraySupplier {
+
     /**
      * Gets an {@link NDArray} from the given {@link NDManager}.
      *
diff --git a/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java b/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
index 714ae2859..0bb42b64e 100644
--- a/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
@@ -20,6 +20,7 @@ import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
 import ai.djl.modality.cv.translator.ImageClassificationTranslator;
 import ai.djl.modality.cv.translator.SingleShotDetectionTranslator;
+import ai.djl.ndarray.BytesSupplier;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.util.JsonSerializable;
@@ -246,9 +247,9 @@ public class ServingTranslatorFactory implements TranslatorFactory {
             Output output = new Output(200, "OK");
             Object obj = translator.processOutput(ctx, list);
             if (obj instanceof JsonSerializable) {
-                output.setContent(((JsonSerializable) obj).toJson() + '\n');
+                output.add(((JsonSerializable) obj).toJson() + '\n');
             } else {
-                output.setContent(JsonUtils.GSON_PRETTY.toJson(obj) + '\n');
+                output.add(JsonUtils.GSON_PRETTY.toJson(obj) + '\n');
             }
             output.addProperty("Content-Type", "application/json");
             return output;
@@ -258,8 +259,8 @@ public class ServingTranslatorFactory implements TranslatorFactory {
         @Override
         public NDList processInput(TranslatorContext ctx, Input input) throws Exception {
             ctx.setAttachment("input", input);
-            PairList<String, byte[]> inputs = input.getContent();
-            byte[] data = inputs.get("data");
+            PairList<String, BytesSupplier> inputs = input.getContent();
+            BytesSupplier data = inputs.get("data");
             if (data == null) {
                 data = inputs.get("body");
             }
@@ -267,7 +268,7 @@ public class ServingTranslatorFactory implements TranslatorFactory {
                 data = input.getContent().valueAt(0);
             }
             try {
-                Image image = factory.fromInputStream(new ByteArrayInputStream(data));
+                Image image = factory.fromInputStream(new ByteArrayInputStream(data.getAsBytes()));
                 return translator.processInput(ctx, image);
             } catch (IOException e) {
                 throw new TranslateException("Input is not an Image data type", e);
@@ -298,17 +299,9 @@ public class ServingTranslatorFactory implements TranslatorFactory {
         /** {@inheritDoc} */
         @Override
         public NDList processInput(TranslatorContext ctx, Input input) throws TranslateException {
-            PairList<String, byte[]> inputs = input.getContent();
-            byte[] data = inputs.get("data");
-            if (data == null) {
-                data = inputs.get("body");
-            }
-            if (data == null) {
-                data = input.getContent().valueAt(0);
-            }
             NDManager manager = ctx.getNDManager();
             try {
-                return NDList.decode(manager, data);
+                return input.getDataAsNDList(manager);
             } catch (IllegalArgumentException e) {
                 throw new TranslateException("Input is not a NDList data type", e);
             }
@@ -318,7 +311,7 @@ public class ServingTranslatorFactory implements TranslatorFactory {
         @Override
         public Output processOutput(TranslatorContext ctx, NDList list) {
             Output output = new Output(200, "OK");
-            output.setContent(list.encode());
+            output.add(list);
             output.addProperty("Content-Type", "tensor/ndlist");
             return output;
         }
diff --git a/api/src/main/java/ai/djl/util/PairList.java b/api/src/main/java/ai/djl/util/PairList.java
index 120e0b31a..e101d3f1b 100644
--- a/api/src/main/java/ai/djl/util/PairList.java
+++ b/api/src/main/java/ai/djl/util/PairList.java
@@ -180,6 +180,18 @@ public class PairList<K, V> implements Iterable<Pair<K, V>> {
         return values.get(index);
     }
 
+    /**
+     * Returns the index of the first occurrence of the specified element in this list, or -1 if
+     * this list does not contain the element.
+     *
+     * @param key – element to search for
+     * @return the index of the first occurrence of the specified element in this list, or -1 if
+     *     this list does not contain the element
+     */
+    public int indexOf(K key) {
+        return keys.indexOf(key);
+    }
+
     /**
      * Returns the key at the specified position in this list.
      *
diff --git a/tools/gradle/jacoco.gradle b/tools/gradle/jacoco.gradle
index 53b43fa60..5eb909d5c 100644
--- a/tools/gradle/jacoco.gradle
+++ b/tools/gradle/jacoco.gradle
@@ -14,7 +14,7 @@ task jacocoMergeTestData(type: JacocoMerge) {
     }
 }
 
-def exclusions = [":examples", ":integration"]
+def exclusions = [":examples", ":integration", ":engines:tensorrt", ":djl-zero"]
 
 task jacocoRootReport(type: JacocoReport) {
     dependsOn jacocoMergeTestData


Test patch:
diff --git a/api/src/test/java/ai/djl/DeviceTest.java b/api/src/test/java/ai/djl/DeviceTest.java
index 45c06b0cd..8cd4ed3c7 100644
--- a/api/src/test/java/ai/djl/DeviceTest.java
+++ b/api/src/test/java/ai/djl/DeviceTest.java
@@ -33,5 +33,7 @@ public class DeviceTest {
         Assert.assertNotEquals(Device.cpu(), Device.gpu());
         Device dev = Device.of("myDevice", 1);
         Assert.assertEquals(dev.getDeviceType(), "myDevice");
+
+        Engine.debugEnvironment();
     }
 }
diff --git a/api/src/test/java/ai/djl/modality/InputOutputTest.java b/api/src/test/java/ai/djl/modality/InputOutputTest.java
new file mode 100644
index 000000000..0f61e80b4
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/InputOutputTest.java
@@ -0,0 +1,74 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality;
+
+import ai.djl.ndarray.BytesSupplier;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import java.nio.charset.StandardCharsets;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class InputOutputTest {
+
+    @Test
+    public void testOutput() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            Output output = new Output(200, "OK");
+            output.addProperty("content-type", "tensor/ndlist");
+            output.add("bytes".getBytes(StandardCharsets.UTF_8));
+            output.add("str");
+            output.add(manager.ones(new Shape(1)));
+            output.add(new NDList(manager.ones(new Shape(2))));
+            output.add("input_1", "str1");
+            output.add("input_2", "bytes1".getBytes(StandardCharsets.UTF_8));
+            output.add("input_3", manager.ones(new Shape(3)));
+            output.add("input_4", new NDList(manager.zeros(new Shape(2))));
+            output.add(1, "data", new NDList(manager.zeros(new Shape(1))));
+
+            Assert.assertEquals(output.getProperty("Content-Type", null), "tensor/ndlist");
+
+            // non-exists key
+            Assert.assertNull(output.get("null"));
+            Assert.assertNull(output.getAsString("null"));
+            Assert.assertNull(output.getAsBytes("null"));
+            Assert.assertNull(output.getAsNDArray(manager, "null"));
+            Assert.assertNull(output.getAsNDList(manager, "null"));
+
+            Assert.assertEquals(output.getAsBytes(2), "str".getBytes(StandardCharsets.UTF_8));
+            Assert.assertEquals(output.getAsString(0), "bytes");
+            Assert.assertEquals(
+                    output.getAsBytes("input_1"), "str1".getBytes(StandardCharsets.UTF_8));
+            Assert.assertEquals(output.getAsString("input_1"), "str1");
+            Shape shape = output.getAsNDArray(manager, "input_3").getShape();
+            Assert.assertEquals(shape.get(0), 3);
+            shape = output.getAsNDList(manager, "input_3").head().getShape();
+            Assert.assertEquals(shape.get(0), 3);
+            Assert.assertEquals(output.getAsBytes("data").length, 58);
+            Assert.assertEquals(output.getAsBytes("input_3").length, 12);
+
+            BytesSupplier data = output.getData();
+            Assert.assertEquals(data, output.getDataAsNDList(manager));
+
+            data = output.get(3);
+            Assert.assertTrue(data instanceof NDArray);
+
+            output = new Output(200, "OK");
+            output.add(manager.zeros(new Shape(1)));
+            NDList ndlist = output.getDataAsNDList(manager);
+            Assert.assertEquals(ndlist.size(), 1);
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/package-info.java b/api/src/test/java/ai/djl/modality/package-info.java
new file mode 100644
index 000000000..842d19be8
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/package-info.java
@@ -0,0 +1,14 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+/** Contains tests for {@link ai.djl.modality}. */
+package ai.djl.modality;
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/CustomTranslatorTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/CustomTranslatorTest.java
index 464a0f809..a4f21a16e 100644
--- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/CustomTranslatorTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/CustomTranslatorTest.java
@@ -32,6 +32,7 @@ import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ModelNotFoundException;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
 import ai.djl.util.JsonUtils;
 import ai.djl.util.Utils;
 import ai.djl.util.ZipUtils;
@@ -41,7 +42,6 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.Writer;
 import java.lang.reflect.Type;
-import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -118,7 +118,8 @@ public class CustomTranslatorTest {
         arguments.put("height", "28");
         arguments.put("flag", Image.Flag.GRAYSCALE.name());
         arguments.put("applySoftmax", "true");
-        runImageClassification(Application.CV.IMAGE_CLASSIFICATION, arguments);
+        runImageClassification(
+                Application.CV.IMAGE_CLASSIFICATION, arguments, "ImageServingTranslator");
 
         Path libDir = modelDir.resolve("lib");
         Path classesDir = libDir.resolve("classes");
@@ -134,7 +135,7 @@ public class CustomTranslatorTest {
         try (Writer writer = Files.newBufferedWriter(confFile)) {
             prop.store(writer, "");
         }
-        runImageClassification(Application.UNDEFINED, null);
+        runImageClassification(Application.UNDEFINED, null, "ImageServingTranslator");
 
         Path libsDir = modelDir.resolve("libs");
         Files.move(libDir, libsDir, StandardCopyOption.REPLACE_EXISTING);
@@ -144,14 +145,14 @@ public class CustomTranslatorTest {
         Files.copy(srcFile, destFile, StandardCopyOption.REPLACE_EXISTING);
 
         // load translator from classes folder
-        runImageClassification(Application.UNDEFINED, null);
+        runImageClassification(Application.UNDEFINED, null, "MyTranslator");
 
         Path jarFile = libsDir.resolve("example.jar");
         ZipUtils.zip(classesDir, jarFile, false);
         Utils.deleteQuietly(classesDir);
 
         // load translator from jar file
-        runImageClassification(Application.UNDEFINED, null);
+        runImageClassification(Application.UNDEFINED, null, "MyTranslator");
     }
 
     @Test
@@ -191,18 +192,19 @@ public class CustomTranslatorTest {
 
         try (ZooModel<Input, Output> model = criteria.loadModel();
                 Predictor<Input, Output> predictor = model.newPredictor()) {
-            Input input = new Input("1");
-            input.addData(buf);
+            Input input = new Input();
+            input.add(buf);
             Output output = predictor.predict(input);
             Assert.assertEquals(output.getCode(), 200);
-            String content = new String(output.getContent(), StandardCharsets.UTF_8);
+            String content = output.getAsString(0);
             Type type = new TypeToken<List<Classification>>() {}.getType();
             List<Classification> result = JsonUtils.GSON.fromJson(content, type);
             Assert.assertEquals(result.get(0).getClassName(), "car");
         }
     }
 
-    private void runImageClassification(Application application, Map<String, Object> arguments)
+    private void runImageClassification(
+            Application application, Map<String, Object> arguments, String translatorName)
             throws IOException, ModelException, TranslateException {
         Criteria<Input, Output> criteria =
                 Criteria.builder()
@@ -214,11 +216,14 @@ public class CustomTranslatorTest {
 
         try (ZooModel<Input, Output> model = criteria.loadModel();
                 Predictor<Input, Output> predictor = model.newPredictor()) {
-            Input input = new Input("1");
-            input.addData("body", data);
+            Translator<Input, Output> translator = model.getTranslator();
+            Assert.assertEquals(translator.getClass().getSimpleName(), translatorName);
+
+            Input input = new Input();
+            input.add("data", data);
             Output output = predictor.predict(input);
             Assert.assertEquals(output.getCode(), 200);
-            String content = new String(output.getContent(), StandardCharsets.UTF_8);
+            String content = output.getAsString(0);
             Type type = new TypeToken<List<Classification>>() {}.getType();
             List<Classification> result = JsonUtils.GSON.fromJson(content, type);
             Assert.assertEquals(result.get(0).getClassName(), "0");
@@ -243,13 +248,13 @@ public class CustomTranslatorTest {
             array = NDImageUtils.toTensor(array).expandDims(0);
             NDList list = new NDList(array);
 
-            Input input = new Input("1");
-            input.addData(0, list.encode());
+            Input input = new Input();
+            input.add(list);
             Output output = predictor.predict(input);
             Assert.assertEquals(output.getCode(), 200);
 
             // manually post process
-            list = NDList.decode(manager, output.getContent());
+            list = output.getDataAsNDList(manager);
             NDArray probabilities = list.singletonOrThrow().get(0).softmax(0);
 
             List<String> classes = model.getArtifact("synset.txt", Utils::readLines);
diff --git a/integration/src/test/translator/MyTranslator.java b/integration/src/test/translator/MyTranslator.java
index 69d2e901a..011e526f2 100644
--- a/integration/src/test/translator/MyTranslator.java
+++ b/integration/src/test/translator/MyTranslator.java
@@ -28,7 +28,7 @@ public class MyTranslator implements ServingTranslator {
 
     @Override
     public NDList processInput(TranslatorContext ctx, Input input) throws Exception {
-        byte[] data = input.getContent().valueAt(0);
+        byte[] data = input.getAsBytes(0);
         ImageFactory factory = ImageFactory.getInstance();
         Image image = factory.fromInputStream(new ByteArrayInputStream(data));
 
@@ -47,7 +47,7 @@ public class MyTranslator implements ServingTranslator {
         Classifications classifications = new Classifications(classes, probabilitiesNd);
 
         Output output = new Output(200, "OK");
-        output.setContent(classifications.toJson());
+        output.add(classifications.toJson());
         return output;
     }