feature(runner): add multiple output support

bentoml/_internal/frameworks/keras.py

@@ -327,7 +327,7 @@ def _mapping(item: "KerasArgType") -> "tf_ext.TensorLike":
 
         def _run_method(
             runnable_self: KerasRunnable, *args: "KerasArgType"
-        ) -> "ext.NpNDArray":
+        ) -> "ext.NpNDArray" | t.Tuple["ext.NpNDArray", ...]:
 
             params = Params["KerasArgType"](*args)
 
@@ -345,6 +345,8 @@ def _run_method(
                 ).isinstance(res):
                     return t.cast("ext.NpNDArray", res.numpy())
 
+                if isinstance(res, list):
+                    return tuple(res)
                 return res
 
         return _run_method

bentoml/_internal/runner/runnable.py

@@ -150,5 +150,5 @@ def __set_name__(self, owner: t.Any, name: str):
 class RunnableMethodConfig:
     batchable: bool
     batch_dim: tuple[int, int]
-    input_spec: AnyType | t.Tuple[AnyType, ...] | None = None
+    input_spec: AnyType | tuple[AnyType, ...] | None = None
     output_spec: AnyType | None = None

bentoml/_internal/runner/runner_handle/remote.py

@@ -189,6 +189,10 @@ async def async_run_method(
                 f"Bento payload decode error: invalid Content-Type '{content_type}'."
             )
 
+        if content_type == "application/vnd.bentoml.multiple_outputs":
+            payloads = pickle.loads(body)
+            return tuple(AutoContainer.from_payload(payload) for payload in payloads)
+
         container = content_type.strip("application/vnd.bentoml.")
 
         try:

bentoml/_internal/server/runner_app.py

@@ -181,6 +181,31 @@ async def _run(requests: t.Iterable[Request]) -> list[Response]:
                 *batched_params.args, **batched_params.kwargs
             )
 
+            server_str = f"BentoML-Runner/{self.runner.name}/{runner_method.name}/{self.worker_index}"
+
+            # multiple output branch
+            if isinstance(batch_ret, tuple):
+                output_num = len(batch_ret)
+                payloadss = [
+                    AutoContainer.batch_to_payloads(
+                        batch_ret[idx], indices, batch_dim=output_batch_dim
+                    )
+                    for idx in range(output_num)
+                ]
+
+                return [
+                    Response(
+                        pickle.dumps(payloads),
+                        headers={
+                            PAYLOAD_META_HEADER: json.dumps({}),
+                            "Content-Type": "application/vnd.bentoml.multiple_outputs",
+                            "Server": server_str,
+                        },
+                    )
+                    for payloads in zip(*payloadss)
+                ]
+
+            # single output branch
             payloads = AutoContainer.batch_to_payloads(
                 batch_ret,
                 indices,
@@ -193,7 +218,7 @@ async def _run(requests: t.Iterable[Request]) -> list[Response]:
                     headers={
                         PAYLOAD_META_HEADER: json.dumps(payload.meta),
                         "Content-Type": f"application/vnd.bentoml.{payload.container}",
-                        "Server": f"BentoML-Runner/{self.runner.name}/{runner_method.name}/{self.worker_index}",
+                        "Server": server_str,
                     },
                 )
                 for payload in payloads

tests/e2e/bento_server_general_features/service.py

@@ -79,6 +79,17 @@ async def predict_ndarray_enforce_dtype(
     return await py_model.predict_ndarray.async_run(inp)
 
 
+@svc.api(
+    input=NumpyNdarray(),
+    output=NumpyNdarray(),
+)
+async def predict_ndarray_multi_output(
+    inp: "np.ndarray[t.Any, np.dtype[t.Any]]",
+) -> "np.ndarray[t.Any, np.dtype[t.Any]]":
+    out1, out2 = await py_model.echo_multi_ndarray.async_run(inp, inp)
+    return out1 + out2
+
+
 @svc.api(
     input=PandasDataFrame(dtype={"col1": "int64"}, orient="records"),
     output=PandasDataFrame(),

tests/e2e/bento_server_general_features/tests/test_io.py

@@ -22,6 +22,14 @@ async def test_numpy(host):
         assert_status=200,
         assert_data=b"[[2, 4], [6, 8]]",
     )
+    await async_request(
+        "POST",
+        f"http://{host}/predict_ndarray_multi_output",
+        headers={"Content-Type": "application/json"},
+        data="[[1,2],[3,4]]",
+        assert_status=200,
+        assert_data=b"[[2, 4], [6, 8]]",
+    )
     await async_request(
         "POST",
         f"http://{host}/predict_ndarray_enforce_shape",