Repository: bentoml/BentoML
Pull request #3736 (2023-03-31 01:33:00)

feat: apply pep574 out-of-band pickling to DefaultContainer
## What does this PR address?

This PR is an improvement of #3606 . Instead of using pickle5 to avoid memory copying of a single numpy ndarray, codes improved by this PR can avoid copying of any composite data structure (e.g. a list of ndarray) containing numpy ndarray or pandas dataframe.

<!--
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?


codecov[bot]: ## [Codecov](https://codecov.io/gh/bentoml/BentoML/pull/3736?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bentoml) Report
> Merging [#3736](https://codecov.io/gh/bentoml/BentoML/pull/3736?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bentoml) (f4fd2cc) into [main](https://codecov.io/gh/bentoml/BentoML/commit/05a45a5a36b9b2a7d902c47f334f57da4f6ade12?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bentoml) (05a45a5) will **decrease** coverage by `31.60%`.
> The diff coverage is `0.00%`.

[![Impacted file tree graph](https://codecov.io/gh/bentoml/BentoML/pull/3736/graphs/tree.svg?width=650&height=150&src=pr&token=GMzgXdpE5b&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bentoml)](https://codecov.io/gh/bentoml/BentoML/pull/3736?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bentoml)

```diff
@@            Coverage Diff             @@
##             main   #3736       +/-   ##
==========================================
- Coverage   31.59%   0.00%   -31.60%     
==========================================
  Files         149     149               
  Lines       12187   12179        -8     
  Branches     2008       0     -2008     
==========================================
- Hits         3850       0     -3850     
- Misses       8054   12179     +4125     
+ Partials      283       0      -283     
```

| [Impacted Files](https://codecov.io/gh/bentoml/BentoML/pull/3736?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bentoml) | Coverage Δ | |
|---|---|---|
| [src/bentoml/\_internal/runner/container.py](https://codecov.io/gh/bentoml/BentoML/pull/3736?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bentoml#diff-c3JjL2JlbnRvbWwvX2ludGVybmFsL3J1bm5lci9jb250YWluZXIucHk=) | `0.00% <0.00%> (ø)` | |
| [src/bentoml/\_internal/utils/pickle.py](https://codecov.io/gh/bentoml/BentoML/pull/3736?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bentoml#diff-c3JjL2JlbnRvbWwvX2ludGVybmFsL3V0aWxzL3BpY2tsZS5weQ==) | `0.00% <0.00%> (ø)` | |

... and [83 files with indirect coverage changes](https://codecov.io/gh/bentoml/BentoML/pull/3736/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bentoml)

bojiang: Potential:
- [ ] pandas DataFrame doing
- [ ] binary file (need investigation, not now)
- [ ] PIL.Image (not supported naturally)
bojiang: Will continue to review after dataframe supported
aarnphm: @larme should we also mention here the performance cost for this?

Patch:
diff --git a/src/bentoml/_internal/runner/container.py b/src/bentoml/_internal/runner/container.py
index 1e8d95ed..43d5a4e1 100644
--- a/src/bentoml/_internal/runner/container.py
+++ b/src/bentoml/_internal/runner/container.py
@@ -1,23 +1,15 @@
 from __future__ import annotations
 
 import abc
-import sys
 import base64
+import pickle
 import typing as t
 import itertools
 
-from simple_di import inject
-from simple_di import Provide
-
 from ..types import LazyType
 from ..utils import LazyLoader
-from ..configuration.containers import BentoMLContainer
-
-if sys.version_info < (3, 8):
-    import pickle5 as pickle
-else:
-    import pickle
-
+from ..utils.pickle import pep574_dumps
+from ..utils.pickle import pep574_loads
 
 SingleType = t.TypeVar("SingleType")
 BatchType = t.TypeVar("BatchType")
@@ -41,7 +33,7 @@ else:
 
 class Payload(t.NamedTuple):
     data: bytes
-    meta: dict[str, bool | int | float | str]
+    meta: dict[str, bool | int | float | str | list[int]]
     container: str
     batch_size: int = -1
 
@@ -52,7 +44,7 @@ class DataContainer(t.Generic[SingleType, BatchType]):
         cls,
         data: bytes,
         batch_size: int,
-        meta: dict[str, bool | int | float | str] | None = None,
+        meta: dict[str, bool | int | float | str | list[int]] | None = None,
     ) -> Payload:
         return Payload(data, meta or {}, container=cls.__name__, batch_size=batch_size)
 
@@ -289,21 +281,23 @@ class NdarrayContainer(DataContainer["ext.NpNDArray", "ext.NpNDArray"]):
         # skip 0-dimensional array
         if batch.shape:
 
-            buffers: list[pickle.PickleBuffer] = []
             if not (batch.flags["C_CONTIGUOUS"] or batch.flags["F_CONTIGUOUS"]):
                 # TODO: use fortan contiguous if it's faster
                 batch = np.ascontiguousarray(batch)
-            bs = pickle.dumps(batch, protocol=5, buffer_callback=buffers.append)
+
+            bs: bytes
+            concat_buffer_bs: bytes
+            indices: list[int]
+            bs, concat_buffer_bs, indices = pep574_dumps(batch)
             bs_str = base64.b64encode(bs).decode("ascii")
-            buffer_bs = buffers[0].raw().tobytes()
-            # release memory
-            buffers[0].release()
+
             return cls.create_payload(
-                buffer_bs,
+                concat_buffer_bs,
                 batch.shape[batch_dim],
                 {
                     "format": "pickle5",
-                    "pickle_bytes": bs_str,
+                    "pickle_bytes_str": bs_str,
+                    "indices": indices,
                 },
             )
 
@@ -320,10 +314,10 @@ class NdarrayContainer(DataContainer["ext.NpNDArray", "ext.NpNDArray"]):
     ) -> ext.NpNDArray:
         format = payload.meta.get("format", "default")
         if format == "pickle5":
-            bs_str = payload.meta["pickle_bytes"]
+            bs_str = t.cast(str, payload.meta["pickle_bytes_str"])
             bs = base64.b64decode(bs_str)
-            recovered_buffers = [pickle.PickleBuffer(payload.data)]
-            return pickle.loads(bs, buffers=recovered_buffers)
+            indices = t.cast(t.List[int], payload.meta["indices"])
+            return t.cast("ext.NpNDArray", pep574_loads(bs, payload.data, indices))
 
         return pickle.loads(payload.data)
 
@@ -340,7 +334,6 @@ class NdarrayContainer(DataContainer["ext.NpNDArray", "ext.NpNDArray"]):
         return payloads
 
     @classmethod
-    @inject
     def from_batch_payloads(
         cls,
         payloads: t.Sequence[Payload],
@@ -387,12 +380,10 @@ class PandasDataFrameContainer(
         ]
 
     @classmethod
-    @inject
     def to_payload(
         cls,
         batch: ext.PdDataFrame | ext.PdSeries,
         batch_dim: int,
-        plasma_db: ext.PlasmaClient | None = Provide[BentoMLContainer.plasma_db],
     ) -> Payload:
         import pandas as pd
 
@@ -403,59 +394,60 @@ class PandasDataFrameContainer(
         if isinstance(batch, pd.Series):
             batch = pd.DataFrame([batch])
 
-        if plasma_db:
-            return cls.create_payload(
-                plasma_db.put(batch).binary(),
-                batch.size,
-                {"plasma": True},
-            )
+        meta: dict[str, bool | int | float | str | list[int]] = {"format": "pickle5"}
+
+        bs: bytes
+        concat_buffer_bs: bytes
+        indices: list[int]
+        bs, concat_buffer_bs, indices = pep574_dumps(batch)
+
+        if indices:
+            meta["with_buffer"] = True
+            data = concat_buffer_bs
+            meta["pickle_bytes_str"] = base64.b64encode(bs).decode("ascii")
+            meta["indices"] = indices
+        else:
+            meta["with_buffer"] = False
+            data = bs
 
         return cls.create_payload(
-            pickle.dumps(batch),
+            data,
             batch.size,
-            {"plasma": False},
+            meta=meta,
         )
 
     @classmethod
-    @inject
     def from_payload(
         cls,
         payload: Payload,
-        plasma_db: ext.PlasmaClient | None = Provide[BentoMLContainer.plasma_db],
     ) -> ext.PdDataFrame:
-        if payload.meta.get("plasma"):
-            import pyarrow.plasma as plasma
-
-            assert plasma_db
-            return plasma_db.get(plasma.ObjectID(payload.data))
-
-        return pickle.loads(payload.data)
+        if payload.meta["with_buffer"]:
+            bs_str = t.cast(str, payload.meta["pickle_bytes_str"])
+            bs = base64.b64decode(bs_str)
+            indices = t.cast(t.List[int], payload.meta["indices"])
+            return pep574_loads(bs, payload.data, indices)
+        else:
+            return pep574_loads(payload.data, b"", [])
 
     @classmethod
-    @inject
     def batch_to_payloads(
         cls,
         batch: ext.PdDataFrame,
         indices: t.Sequence[int],
         batch_dim: int = 0,
-        plasma_db: ext.PlasmaClient | None = Provide[BentoMLContainer.plasma_db],
     ) -> list[Payload]:
         batches = cls.batch_to_batches(batch, indices, batch_dim)
 
-        payloads = [
-            cls.to_payload(subbatch, batch_dim, plasma_db) for subbatch in batches
-        ]
+        payloads = [cls.to_payload(subbatch, batch_dim) for subbatch in batches]
         return payloads
 
     @classmethod
-    @inject
     def from_batch_payloads(  # pylint: disable=arguments-differ
         cls,
         payloads: t.Sequence[Payload],
         batch_dim: int = 0,
-        plasma_db: ext.PlasmaClient | None = Provide[BentoMLContainer.plasma_db],
     ) -> tuple[ext.PdDataFrame, list[int]]:
-        batches = [cls.from_payload(payload, plasma_db) for payload in payloads]
+        batches = [cls.from_payload(payload) for payload in payloads]
         return cls.batches_to_batch(batches, batch_dim)
 
 
@@ -487,20 +479,45 @@ class DefaultContainer(DataContainer[t.Any, t.List[t.Any]]):
     def to_payload(cls, batch: t.Any, batch_dim: int) -> Payload:
         if isinstance(batch, t.Generator):  # Generators can't be pickled
             batch = list(t.cast(t.Generator[t.Any, t.Any, t.Any], batch))
+
+        meta: dict[str, bool | int | float | str | list[int]] = {"format": "pickle5"}
+
+        bs: bytes
+        concat_buffer_bs: bytes
+        indices: list[int]
+        bs, concat_buffer_bs, indices = pep574_dumps(batch)
+
+        if indices:
+            meta["with_buffer"] = True
+            data = concat_buffer_bs
+            meta["pickle_bytes_str"] = base64.b64encode(bs).decode("ascii")
+            meta["indices"] = indices
+        else:
+            meta["with_buffer"] = False
+            data = bs
+
         if isinstance(batch, list):
-            return cls.create_payload(
-                pickle.dumps(batch), len(t.cast(t.List[t.Any], batch))
-            )
+            batch_size = len(t.cast(t.List[t.Any], batch))
         else:
-            return cls.create_payload(pickle.dumps(batch), 1)
+            batch_size = 1
+
+        return cls.create_payload(
+            data=data,
+            batch_size=batch_size,
+            meta=meta,
+        )
 
     @classmethod
-    @inject
     def from_payload(cls, payload: Payload) -> t.Any:
-        return pickle.loads(payload.data)
+        if payload.meta["with_buffer"]:
+            bs_str = t.cast(str, payload.meta["pickle_bytes_str"])
+            bs = base64.b64decode(bs_str)
+            indices = t.cast(t.List[int], payload.meta["indices"])
+            return pep574_loads(bs, payload.data, indices)
+        else:
+            return pep574_loads(payload.data, b"", [])
 
     @classmethod
-    @inject
     def batch_to_payloads(
         cls,
         batch: list[t.Any],
@@ -513,7 +530,6 @@ class DefaultContainer(DataContainer[t.Any, t.List[t.Any]]):
         return payloads
 
     @classmethod
-    @inject
     def from_batch_payloads(
         cls,
         payloads: t.Sequence[Payload],
diff --git a/src/bentoml/_internal/utils/pickle.py b/src/bentoml/_internal/utils/pickle.py
new file mode 100644
index 00000000..a0e54194
--- /dev/null
+++ b/src/bentoml/_internal/utils/pickle.py
@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import sys
+import typing as t
+
+if sys.version_info < (3, 8):
+    import pickle5 as pickle
+else:
+    import pickle
+
+
+# Pickle protocol 5 with out-of-band data
+# https://peps.python.org/pep-0574/
+def pep574_dumps(obj: t.Any) -> tuple[bytes, bytes, list[int]]:
+    buffers: list[pickle.PickleBuffer] = []
+    main_bytes: bytes = pickle.dumps(obj, protocol=5, buffer_callback=buffers.append)
+
+    if not buffers:
+        return main_bytes, b"", []
+
+    buffer_bytess: list[bytes] = [buff.raw().tobytes() for buff in buffers]
+
+    for buff in buffers:
+        buff.release()
+
+    indices: list[int] = [0]
+    for buff_bytes in buffer_bytess:
+        start = indices[-1]
+        end = start + len(buff_bytes)
+        indices.append(end)
+
+    concat_buffer_bytes: bytes = b"".join(buffer_bytess)
+    return main_bytes, concat_buffer_bytes, indices
+
+
+def pep574_loads(
+    main_bytes: bytes, concat_buffer_bytes: bytes, indices: list[int]
+) -> t.Any:
+
+    if not indices:
+        return pickle.loads(main_bytes)
+
+    mem = memoryview(concat_buffer_bytes)
+    partitions = zip(indices, indices[1:])
+    recover_buffers: list[pickle.PickleBuffer] = []
+    for partition in partitions:
+        buff = pickle.PickleBuffer(mem[slice(*partition)])
+        recover_buffers.append(buff)
+
+    return pickle.loads(main_bytes, buffers=recover_buffers)
diff --git a/src/bentoml/server.py b/src/bentoml/server.py
index f063e50d..8427accb 100644
--- a/src/bentoml/server.py
+++ b/src/bentoml/server.py
@@ -268,8 +268,8 @@ class GrpcServer(Server):
         reload: bool = False,
         production: bool = False,
         env: t.Literal["conda"] | None = None,
-        host: str = Provide[BentoMLContainer.http.host],
-        port: int = Provide[BentoMLContainer.http.port],
+        host: str = Provide[BentoMLContainer.grpc.host],
+        port: int = Provide[BentoMLContainer.grpc.port],
         working_dir: str | None = None,
         api_workers: int | None = Provide[BentoMLContainer.api_server_workers],
         backlog: int = Provide[BentoMLContainer.api_server_config.backlog],


Test patch:
diff --git a/tests/e2e/bento_server_http/tests/test_meta.py b/tests/e2e/bento_server_http/tests/test_meta.py
index e95b6e67..a244a971 100644
--- a/tests/e2e/bento_server_http/tests/test_meta.py
+++ b/tests/e2e/bento_server_http/tests/test_meta.py
@@ -12,6 +12,17 @@ import bentoml
 from bentoml.testing.utils import async_request
 
 
+@pytest.mark.asyncio
+async def test_api_server_load(host: str):
+    for _ in range(20):
+        status, _, _ = await async_request(
+            "POST",
+            f"http://{host}/echo_json",
+            headers={"Content-Type": "application/json"},
+            data='"hi"',
+        )
+
+
 @pytest.mark.asyncio
 async def test_api_server_meta(host: str) -> None:
     status, _, _ = await async_request("GET", f"http://{host}/")
diff --git a/tests/unit/_internal/utils/test_pickle.py b/tests/unit/_internal/utils/test_pickle.py
new file mode 100644
index 00000000..e3ebadc6
--- /dev/null
+++ b/tests/unit/_internal/utils/test_pickle.py
@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import typing as t
+
+if t.TYPE_CHECKING:
+    import bentoml._internal.external_typing as ext
+
+
+def test_pep574_restore() -> None:
+    import numpy as np
+    import pandas as pd
+
+    from bentoml._internal.utils.pickle import pep574_dumps
+    from bentoml._internal.utils.pickle import pep574_loads
+
+    arr1: ext.NpNDArray = np.random.uniform(size=(20, 20))
+    arr2: ext.NpNDArray = np.random.uniform(size=(64, 64))
+    arr3: ext.NpNDArray = np.random.uniform(size=(72, 72))
+
+    lst = [arr1, arr2, arr3]
+
+    bs: bytes
+    concat_buffer_bs: bytes
+    indices: list[int]
+    bs, concat_buffer_bs, indices = pep574_dumps(lst)
+    restored = t.cast(
+        t.List["ext.NpNDArray"], pep574_loads(bs, concat_buffer_bs, indices)
+    )
+    for idx, arr in enumerate(lst):
+        assert np.isclose(arr, restored[idx]).all()
+
+    dic: dict[str, ext.NpNDArray] = dict(a=arr1, b=arr2, c=arr3)
+    bs, concat_buffer_bs, indices = pep574_dumps(dic)
+    restored = t.cast(
+        t.Dict[str, "ext.NpNDArray"], pep574_loads(bs, concat_buffer_bs, indices)
+    )
+    for key, arr in dic.items():
+        assert np.isclose(arr, restored[key]).all()
+
+    df1: ext.PdDataFrame = pd.DataFrame(arr1)
+    df2: ext.PdDataFrame = pd.DataFrame(arr2)
+    df3: ext.PdDataFrame = pd.DataFrame(arr3)
+
+    df_lst = [df1, df2, df3]
+
+    bs, concat_buffer_bs, indices = pep574_dumps(df_lst)
+    restored = t.cast(
+        t.List["ext.PdDataFrame"], pep574_loads(bs, concat_buffer_bs, indices)
+    )
+    for idx, df in enumerate(df_lst):
+        assert np.isclose(df.to_numpy(), restored[idx].to_numpy()).all()
+
+    df_dic: dict[str, ext.PdDataFrame] = dict(a=df1, b=df2, c=df3)
+    bs, concat_buffer_bs, indices = pep574_dumps(df_dic)
+    restored = t.cast(
+        t.Dict[str, "ext.PdDataFrame"], pep574_loads(bs, concat_buffer_bs, indices)
+    )
+    for key, df in df_dic.items():
+        assert np.isclose(df.to_numpy(), restored[key].to_numpy()).all()