feature(runner): add multiple output support (#2912)

## What does this PR address?
<!--
Thanks for sending a pull request!

Congrats for making it this far! Here's a 🍱 for you. There are still a
few steps ahead.

Please make sure to read the contribution guidelines, then fill out the
blanks below before requesting a code review.

Name your Pull Request with one of the following prefixes, e.g. "feat:
add support for PyTorch", to indicate the type of changes proposed. This
is based on the [Conventional Commits
specification](https://www.conventionalcommits.org/en/v1.0.0/#summary).
  - feat: (new feature for the user, not a new feature for build script)
  - fix: (bug fix for the user, not a fix to a build script)
  - docs: (changes to the documentation)
- style: (formatting, missing semicolons, etc; no production code
change)
  - refactor: (refactoring production code, eg. renaming a variable)
  - perf: (code changes that improve performance)
- test: (adding missing tests, refactoring tests; no production code
change)
  - chore: (updating grunt tasks etc; no production code change)
- build: (changes that affect the build system or external dependencies)
  - ci: (changes to configuration files and scripts)
  - revert: (reverts a previous commit)

Describe your changes in detail. Attach screenshots here if appropriate.

Once you're done with this, someone from BentoML team or community
member will help review your PR (see "Who can help review?" section for
potential reviewers.). If no one has reviewed your PR after a week have
passed, don't hesitate to post a new comment and ping @-the same person.
Notifications sometimes get lost 🥲.
-->

<!-- Remove if not applicable -->
Fixes #(issue)

## Before submitting:
<!--- Go over all the following points, and put an `x` in all the boxes
that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're
here to help! -->
<!--- If you plan to update documentation or tests in follow-up, please
note -->
- [ ] Does the Pull Request follow [Conventional Commits
specification](https://www.conventionalcommits.org/en/v1.0.0/#summary)
naming? Here are [GitHub's

guide](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request)
on how to create a pull request.
- [ ] Does the code follow BentoML's code style, both `make format` and
`make lint` script have passed
([instructions](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#style-check-auto-formatting-type-checking))?
- [ ] Did you read through [contribution
guidelines](https://github.com/bentoml/BentoML/blob/main/CONTRIBUTING.md#ways-to-contribute)
and follow [development
guidelines](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#start-developing)?
- [ ] Did your changes require updates to the documentation? Have you
updated
those accordingly? Here are [documentation
guidelines](https://github.com/bentoml/BentoML/tree/main/docs) and [tips
on writting
docs](https://github.com/bentoml/BentoML/tree/main/docs#writing-documentation).
- [ ] Did you write tests to cover your changes?

## Who can help review?

Feel free to tag members/contributors who can help review your PR.
<!--
Feel free to ping any of the BentoML members for help on your issue, but
don't ping more than three people 😊.
If you know how to use git blame, that is probably the easiest way.

Team members that you can ping:
- @parano
- @yubozhao
- @bojiang
- @ssheng
- @aarnphm
- @sauyon
- @larme
- @yetone
- @jjmachan
-->

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",