Repository: bentoml/BentoML
Pull request #3764 (2023-04-17 10:01:26)

fix: passing serve_cmd and passthrough kwargs
Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>

## 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 -->

This PR addresses two issues:
1. Currently on main, `bentoml.Server` defaults to use `serve-http` as serve_cmd, which means GrpcServer won't work. This PR fixes that by passing `serve_cmd` to the server args.
2. After creating the server, there is no way to pass down subprocess args to the actual process. This PR adds `**kwargs` as passthrough to server.process 

## 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?



Patch:
diff --git a/src/bentoml/__init__.py b/src/bentoml/__init__.py
index 684c63f6..03d63968 100644
--- a/src/bentoml/__init__.py
+++ b/src/bentoml/__init__.py
@@ -91,6 +91,7 @@ if TYPE_CHECKING:
     from . import client  # Client API
     from . import batch  # Batch API
     from . import exceptions  # BentoML exceptions
+    from . import server  # Server API
 
     # isort: on
 else:
@@ -142,6 +143,7 @@ else:
     metrics = _LazyLoader("bentoml.metrics", globals(), "bentoml.metrics")
     container = _LazyLoader("bentoml.container", globals(), "bentoml.container")
     client = _LazyLoader("bentoml.client", globals(), "bentoml.client")
+    server = _LazyLoader("bentoml.server", globals(), "bentoml.server")
     exceptions = _LazyLoader("bentoml.exceptions", globals(), "bentoml.exceptions")
 
     del _LazyLoader
@@ -156,6 +158,7 @@ __all__ = [
     "metrics",
     "container",
     "client",
+    "server",
     "io",
     "Tag",
     "Model",
diff --git a/src/bentoml/_internal/client/http.py b/src/bentoml/_internal/client/http.py
index 5c7d1896..f4e6cda7 100644
--- a/src/bentoml/_internal/client/http.py
+++ b/src/bentoml/_internal/client/http.py
@@ -163,7 +163,7 @@ class HTTPClient(Client):
 
         async with aiohttp.ClientSession(self.server_url) as sess:
             async with sess.post(
-                "/" + api.route,
+                "/" + api.route if not api.route.startswith("/") else api.route,
                 data=req_body,
                 headers={"content-type": fake_resp.headers["content-type"]},
             ) as resp:
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..7a58981a 100644
--- a/src/bentoml/server.py
+++ b/src/bentoml/server.py
@@ -64,7 +64,7 @@ class Server(ABC):
             sys.executable,
             "-m",
             "bentoml",
-            "serve-http",
+            serve_cmd,
             bento_str,
             "--host",
             host,
@@ -93,7 +93,21 @@ class Server(ABC):
 
     def _create_startmanager(server):  # type: ignore # not calling self self
         class _StartManager:
-            def __init__(self, blocking: bool = False):
+            def __init__(
+                self, blocking: bool = False, env: dict[str, str] | None = None
+            ):
+                """Start the server programmatically.
+
+                To get the client, use the context manager.
+
+                .. note::
+
+                   ``blocking=True`` and using ``start()`` as a context manager is mutually exclusive.
+
+                Args:
+                    blocking: If True, the server will block until it is stopped.
+                    env: A dictionary of environment variables to pass to the server. Default to ``None``.
+                """
                 logger.info(f"starting server with arguments: {server.args}")
 
                 server.process = subprocess.Popen(
@@ -101,6 +115,7 @@ class Server(ABC):
                     stdout=subprocess.PIPE,
                     stderr=subprocess.PIPE,
                     stdin=subprocess.PIPE,
+                    env=env,
                 )
 
                 if blocking:
@@ -268,8 +283,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/requirements/tests-requirements.txt b/requirements/tests-requirements.txt
index ab0cf604..9fa5927c 100644
--- a/requirements/tests-requirements.txt
+++ b/requirements/tests-requirements.txt
@@ -9,7 +9,7 @@ setuptools>=63
 isort==5.11.4
 ruff==0.0.261
 pytest-cov==4.0.0
-pytest==7.3.0
+pytest==7.3.1
 pytest-xdist[psutil]==3.2.1
 pytest-asyncio==0.21.0
 imageio==2.27.0
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()