Use np.object_ when serializing BYTES type (triton-inference-server#38)

* Fix docstring

* Fix trailing zero issue for Python backend

* Revert

Author: Tabrizian

src/resources/startup.py

@@ -53,29 +53,33 @@
 
 def serialize_byte_tensor(input_tensor):
     """
-        Serializes a bytes tensor into a flat numpy array of length prepend bytes.
-        Can pass bytes tensor as numpy array of bytes with dtype of np.bytes_,
-        numpy strings with dtype of np.str_ or python strings with dtype of np.object.
-        Parameters
-        ----------
-        input_tensor : np.array
-            The bytes tensor to serialize.
-        Returns
-        -------
-        serialized_bytes_tensor : np.array
-            The 1-D numpy array of type uint8 containing the serialized bytes in 'C' order.
-        Raises
-        ------
-        InferenceServerException
-            If unable to serialize the given tensor.
-        """
+    Serializes a bytes tensor into a flat numpy array of length prepended
+    bytes. The numpy array should use dtype of np.object_. For np.bytes_,
+    numpy will remove trailing zeros at the end of byte sequence and because
+    of this it should be avoided.
+
+    Parameters
+    ----------
+    input_tensor : np.array
+        The bytes tensor to serialize.
+
+    Returns
+    -------
+    serialized_bytes_tensor : np.array
+        The 1-D numpy array of type uint8 containing the serialized bytes in 'C' order.
+
+    Raises
+    ------
+    InferenceServerException
+        If unable to serialize the given tensor.
+    """
 
     if input_tensor.size == 0:
         return np.empty([0])
 
-    # If the input is a tensor of string/bytes objects, then must flatten those
-    # into a 1-dimensional array containing the 4-byte byte size followed by
-    # the actual element bytes. All elements are concatenated together in "C"
+    # If the input is a tensor of string/bytes objects, then must flatten those into
+    # a 1-dimensional array containing the 4-byte byte size followed by the
+    # actual element bytes. All elements are concatenated together in "C"
     # order.
     if (input_tensor.dtype == np.object) or (input_tensor.dtype.type
                                              == np.bytes_):
@@ -84,18 +88,15 @@ def serialize_byte_tensor(input_tensor):
             # If directly passing bytes to BYTES type,
             # don't convert it to str as Python will encode the
             # bytes which may distort the meaning
-            if obj.dtype.type == np.bytes_:
-                if type(obj.item()) == bytes:
-                    s = obj.item()
-                else:
-                    s = bytes(obj)
+            if type(obj.item()) == bytes:
+                s = obj.item()
             else:
-                s = str(obj).encode('utf-8')
+                s = bytes(obj)
             flattened += struct.pack("<I", len(s))
             flattened += s
-        flattened_array = np.asarray(flattened)
+        flattened_array = np.asarray(flattened, dtype=np.object_)
         if not flattened_array.flags['C_CONTIGUOUS']:
-            flattened_array = np.ascontiguousarray(flattened_array)
+            flattened_array = np.ascontiguousarray(flattened_array, dtype=np.object_)
         return flattened_array
     else:
         raise TritonModelException(
@@ -128,7 +129,7 @@ def deserialize_bytes_tensor(encoded_tensor):
         sb = struct.unpack_from("<{}s".format(l), val_buf, offset)[0]
         offset += l
         strs.append(sb)
-    return (np.array(strs, dtype=bytes))
+    return (np.array(strs, dtype=np.object_))
 
 
 def parse_startup_arguments():
@@ -149,7 +150,6 @@ def parse_startup_arguments():
 class PythonHost(PythonInterpreterServicer):
     """This class handles inference request for python script.
     """
-
     def __init__(self, module_path, *args, **kwargs):
         super(PythonInterpreterServicer, self).__init__(*args, **kwargs)
 
@@ -305,13 +305,18 @@ def Execute(self, request, context):
                 # We need to serialize TYPE_STRING
                 if output_np_array.dtype == np.object or output_np_array.dtype.type is np.bytes_:
                     output_np_array = serialize_byte_tensor(output_np_array)
-
-                tensor = Tensor(name=output_tensor.name(),
-                                dtype=tpb_utils.numpy_to_triton_type(
-                                    output_np_array.dtype.type),
-                                dims=output_shape,
-                                raw_data=output_np_array.tobytes())
-
+                    raw_data = output_np_array.item() if output_np_array.size > 0 else b''
+                    tensor = Tensor(name=output_tensor.name(),
+                                    dtype=tpb_utils.numpy_to_triton_type(
+                                        output_np_array.dtype.type),
+                                    dims=output_shape,
+                                    raw_data=raw_data)
+                else:
+                    tensor = Tensor(name=output_tensor.name(),
+                                    dtype=tpb_utils.numpy_to_triton_type(
+                                        output_np_array.dtype.type),
+                                    dims=output_shape,
+                                    raw_data=output_np_array.tobytes())
                 response_tensors.append(tensor)
             exec_responses.append(InferenceResponse(outputs=response_tensors))
         execute_response = ExecuteResponse(responses=exec_responses)

src/resources/triton_python_backend_utils.py

@@ -52,7 +52,7 @@
     # TRITONSERVER_TYPE_FP64
     12: np.float64,
     # TRITONSERVER_TYPE_STRING
-    13: np.bytes_
+    13: np.object_
 }
 
 TRITON_STRING_TO_NUMPY = {
@@ -68,7 +68,7 @@
     'TYPE_FP16': np.float16,
     'TYPE_FP32': np.float32,
     'TYPE_FP64': np.float64,
-    'TYPE_STRING': np.bytes_
+    'TYPE_STRING': np.object_
 }
 
 NUMPY_TO_TRITON_TYPE = {v: k for k, v in TRITON_TO_NUMPY_TYPE.items()}
@@ -346,9 +346,6 @@ def triton_to_numpy_type(data_type):
 
 
 def numpy_to_triton_type(data_type):
-    if data_type == np.object_:
-        return 'TYPE_STRING'
-
     return NUMPY_TO_TRITON_TYPE[data_type]