Fix ValueError in concatenation of result_by_req_id with zero-dimensional arrays, when there are 1-length arrays in the inference result.

model_repository/sample/1/model.py

@@ -17,7 +17,7 @@ def execute(self, requests):
         responses = [None for _ in requests]
         for idx, request in enumerate(requests):
             current_add_value = int(json.loads(request.parameters()).get("add", 0))
-            in_tensor = [item.as_numpy() + current_add_value for item in request.inputs() if item.name() == "model_in"]
+            in_tensor = [item.as_numpy() + current_add_value for item in request.inputs() if "model_in" in item.name()]
             out_tensor = [
                 pb_utils.Tensor(output_name, x.astype(output_dtype))
                 for x, output_name, output_dtype in zip(in_tensor, self.output_name_list, self.output_dtype_list)

model_repository/sample_multiple_hybrid_dims/1/model.py

@@ -0,0 +1,38 @@
+import json
+
+import numpy as np
+import triton_python_backend_utils as pb_utils
+
+
+class TritonPythonModel:
+    def initialize(self, args):
+        self.model_config = model_config = json.loads(args["model_config"])
+        output_configs = model_config["output"]
+
+        self.output_name_list = [output_config["name"] for output_config in output_configs]
+        self.output_dtype_list = [
+            pb_utils.triton_string_to_numpy(output_config["data_type"]) for output_config in output_configs
+        ]
+
+    def execute(self, requests):
+        responses = [None for _ in requests]
+        for idx, request in enumerate(requests):
+            current_add_value = int(json.loads(request.parameters()).get("add", 0))
+            in_tensor = [item.as_numpy() + current_add_value for item in request.inputs() if "model_in" in item.name()]
+
+            out_tensor = [
+                pb_utils.Tensor(output_name, x.astype(output_dtype))
+                for x, output_name, output_dtype in zip(in_tensor, self.output_name_list, self.output_dtype_list)
+            ]
+
+            out_tensor.append(
+                pb_utils.Tensor(
+                    "model_out2",
+                    np.array([current_add_value], dtype=self.output_dtype_list[2]),
+                )
+            )
+
+            inference_response = pb_utils.InferenceResponse(output_tensors=out_tensor)
+
+            responses[idx] = inference_response
+        return responses

model_repository/sample_multiple_hybrid_dims/config.pbtxt

@@ -0,0 +1,63 @@
+name: "sample_multiple_hybrid_dims"
+backend: "python"
+max_batch_size: 2
+
+parameters [
+  {
+    key: "add",
+    value: { string_value: "0" }
+  }
+]
+
+input [
+{
+    name: "model_in0"
+    data_type: TYPE_FP32
+    dims: [ -1 ]
+},
+{
+    name: "model_in1"
+    data_type: TYPE_FP32
+    dims: [ -1 ]
+}
+]
+
+output [
+{
+    name: "model_out0"
+    data_type: TYPE_FP32
+    dims: [ -1 ]
+},
+{
+    name: "model_out1"
+    data_type: TYPE_FP32
+    dims: [ -1 ]
+},
+{
+    name: "model_out2"
+    data_type: TYPE_FP32
+    dims: [ 1 ]
+}
+]
+
+instance_group [{ kind: KIND_CPU, count: 1 }]
+
+model_warmup {
+  name: "RandomSampleInput"
+  batch_size: 1
+  inputs [{
+      key: "model_in0"
+      value: {
+        data_type: TYPE_FP32
+        dims: [ 10 ]
+        random_data: true
+      }
+   }, {
+      key: "model_in1"
+      value: {
+        data_type: TYPE_FP32
+        dims: [ 10 ]
+        zero_data: true
+      }
+   }]
+}

tests/test_model_call.py

@@ -69,3 +69,22 @@ def test_reload_model_spec(config):
     sample = np.random.rand(8, 100).astype(np.float32)
     result = client(sample)
     assert np.isclose(result, sample).all()
+
+
+def test_with_multiple_hybrid_dims(config):
+    client = get_client(*config, model_name="sample_multiple_hybrid_dims")
+
+    samples = [np.random.rand(1, 100).astype(np.float32) for _ in range(2)]
+    ADD_VALUE = 1
+    result = client(
+        {
+            client.default_model_spec.model_input[0].name: samples[0],
+            client.default_model_spec.model_input[1].name: samples[1],
+        },
+        parameters={"add": ADD_VALUE},
+    )
+
+    assert np.isclose(result[0], samples[0][0] + ADD_VALUE).all()
+    assert np.isclose(result[1], samples[1][0] + ADD_VALUE).all()
+
+    assert result[2] == ADD_VALUE

tritony/tools.py

@@ -66,6 +66,18 @@ async def data_generator(data: list[np.ndarray], batch_size: int, queue: asyncio
     stop.set()
 
 
+def safe_concatenate(arrays, axis=0):
+    arrays = [
+        (
+            np.expand_dims(arr, axis=0)
+            if arr is not None and arr.ndim == 0
+            else (arr if arr is not None else np.array([]))
+        )
+        for arr in arrays
+    ]
+    return np.concatenate(arrays, axis=axis)
+
+
 async def send_request_async(
     inference_client: InferenceClient,
     data_queue,
@@ -410,10 +422,11 @@ async def _call_async_item(
             ret = sorted(itertools.chain(*ret[:ASYNC_TASKS]))
             result_by_req_id = [output_result_list for req_id, output_result_list in ret]
 
+            zipped_result = list(zip(*result_by_req_id))
             if model_spec.max_batch_size == 0:
-                result_by_output_name = list(zip(*result_by_req_id))
+                result_by_output_name = zipped_result
             else:
-                result_by_output_name = list(map(lambda ll: np.concatenate(ll, axis=0), zip(*result_by_req_id)))
+                result_by_output_name = list(map(lambda ll: safe_concatenate(ll, axis=0), zipped_result))
 
             if len(result_by_output_name) == 1:
                 result_by_output_name = result_by_output_name[0]
@@ -476,10 +489,11 @@ def _call_request(
                         )
                     )
 
+                zipped_result = list(zip(*result_by_req_id))
                 if model_spec.max_batch_size == 0:
-                    result_by_output_name = list(zip(*result_by_req_id))
+                    result_by_output_name = zipped_result
                 else:
-                    result_by_output_name = list(map(lambda ll: np.concatenate(ll, axis=0), zip(*result_by_req_id)))
+                    result_by_output_name = list(map(lambda ll: safe_concatenate(ll, axis=0), zipped_result))
 
                 if len(result_by_output_name) == 1:
                     result_by_output_name = result_by_output_name[0]