Repository: bentoml/BentoML
Pull request #4375 (2023-12-28 21:49:04)

chore: cherry-pick SSE utils into 1.2 branch
## 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, `pre-commit run -a` script has 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/_internal/io_descriptors/utils.py b/src/bentoml/_internal/io_descriptors/utils.py
new file mode 100644
index 00000000..04258edd
--- /dev/null
+++ b/src/bentoml/_internal/io_descriptors/utils.py
@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import io
+import typing as t
+
+import attr
+
+from bentoml.exceptions import BentoMLException
+
+if t.TYPE_CHECKING:
+
+    class SSEArgs(t.TypedDict, total=False):
+        data: str
+        event: str
+        id: str
+        retry: int
+
+
+@attr.s(auto_attribs=True)
+class SSE:
+    data: str
+    event: str | None = None
+    id: str | None = None
+    retry: int | None = None
+
+    def marshal(self) -> str:
+        with io.StringIO() as buffer:
+            if self.event:
+                if "\n" in self.event:
+                    raise BentoMLException("SSE event name should not contain new line")
+                buffer.write(f"event: {self.event}\n")
+            if self.id is not None:
+                buffer.write(f"id: {self.id}\n")
+            if self.retry is not None:
+                buffer.write(f"retry: {self.retry}\n")
+            for line in self.data.split("\n"):
+                buffer.write(f"data: {line}\n")
+            buffer.write("\n")
+            return buffer.getvalue()
+
+    @classmethod
+    def _read_sse(cls, buffer: io.StringIO) -> SSE:
+        event: SSEArgs = {}
+        data_buffer: io.StringIO | None = None
+        first_data_line = True
+
+        while True:
+            line = buffer.readline()
+            if not line:
+                break
+            if line == "\n":
+                break
+            if line.startswith("data: "):
+                if first_data_line:
+                    first_data_line = False
+                    event["data"] = line[6:-1]
+                else:
+                    if data_buffer is None:
+                        # only init data_buffer when there is more than one data line
+                        data_buffer = io.StringIO()
+                        data_buffer.write(event.get("data", ""))
+                    data_buffer.write("\n")
+                    data_buffer.write(line[6:-1])
+            if line.startswith("event: "):
+                event["event"] = line[7:].strip()
+            if line.startswith("id: "):
+                event["id"] = line[4:].strip()
+            if line.startswith("retry: "):
+                event["retry"] = int(line[7:].strip())
+
+        if data_buffer is not None:
+            event["data"] = data_buffer.getvalue()
+            data_buffer.close()
+
+        return SSE(**event)
+
+    @classmethod
+    async def from_iterator(
+        cls,
+        async_iterator: t.AsyncIterator[str],
+    ) -> t.AsyncGenerator[SSE, None]:
+        with io.StringIO() as buffer:
+            read_cursor = 0
+            last_chunk = ""
+            async for chunk in async_iterator:
+                buffer.write(chunk)
+                if (
+                    "\n\n" in chunk
+                    or chunk
+                    and chunk[0] == "\n"
+                    and last_chunk
+                    and last_chunk[-1] == "\n"
+                ):
+                    buffer.seek(read_cursor)
+                    yield cls._read_sse(buffer)
+                    read_cursor = buffer.tell()
+                last_chunk = chunk
diff --git a/src/bentoml/_internal/runner/runner.py b/src/bentoml/_internal/runner/runner.py
index a33d1ec5..94a3b7a7 100644
--- a/src/bentoml/_internal/runner/runner.py
+++ b/src/bentoml/_internal/runner/runner.py
@@ -57,7 +57,7 @@ class RunnerMethod(t.Generic[T, P, R]):
 
     def async_stream(
         self, *args: P.args, **kwargs: P.kwargs
-    ) -> t.AsyncGenerator[R, None]:
+    ) -> t.AsyncGenerator[str, None]:
         return self.runner._runner_handle.async_stream_method(self, *args, **kwargs)
 
 
diff --git a/src/bentoml/_internal/runner/runner_handle/remote.py b/src/bentoml/_internal/runner/runner_handle/remote.py
index 8f5dd926..5a3245a5 100644
--- a/src/bentoml/_internal/runner/runner_handle/remote.py
+++ b/src/bentoml/_internal/runner/runner_handle/remote.py
@@ -276,10 +276,10 @@ class RemoteRunnerClient(RunnerHandle):
 
     async def async_stream_method(
         self,
-        __bentoml_method: RunnerMethod[t.Any, P, R],
+        __bentoml_method: RunnerMethod[t.Any, P, str],
         *args: P.args,
         **kwargs: P.kwargs,
-    ) -> t.AsyncGenerator[R, None]:
+    ) -> t.AsyncGenerator[str, None]:
         import aiohttp
 
         from ...runner.container import AutoContainer
@@ -326,18 +326,9 @@ class RemoteRunnerClient(RunnerHandle):
                 data=data,
                 headers=headers,
             ) as resp:
-                buffer = bytearray()
-                async for b, end_of_http_chunk in resp.content.iter_chunks():
-                    buffer.extend(b)
-
-                    # This is to handle large payload that is split into multiple chunks
-                    if end_of_http_chunk and len(buffer) > 0:
-                        # TODO: To remove pickling so that we can stream data as it is
-                        payload = pickle.loads(buffer)
-                        yield AutoContainer.from_payload(payload)
-
-                        # Clearing the buffer for the next data
-                        buffer = bytearray()
+                async for b in resp.content.iter_any():
+                    # convert bytes to string
+                    yield b.decode()
 
         except aiohttp.ClientOSError as e:
             raise RemoteException("Failed to connect to runner server.") from e
diff --git a/src/bentoml/_internal/server/runner_app.py b/src/bentoml/_internal/server/runner_app.py
index ae3856ff..a833c9ef 100644
--- a/src/bentoml/_internal/server/runner_app.py
+++ b/src/bentoml/_internal/server/runner_app.py
@@ -201,7 +201,7 @@ class RunnerAppFactory(BaseAppFactory):
             # Streaming does not have batching implemented yet
             async def infer_stream(
                 paramss: t.Sequence[Params[t.Any]],
-            ) -> t.Sequence[t.AsyncGenerator[Payload, None]]:
+            ) -> t.Sequence[t.AsyncGenerator[str, None]]:
                 async def inner():
                     # This is a workaround to allow infer stream to return a iterable of
                     # async generator, to align with how non stream inference works
@@ -212,8 +212,7 @@ class RunnerAppFactory(BaseAppFactory):
                         traceback.print_exc()
                         raise
                     async for data in ret:
-                        payload = AutoContainer.to_payload(data, 0)
-                        yield payload
+                        yield data
 
                 return (inner(),)
 
@@ -350,16 +349,8 @@ class RunnerAppFactory(BaseAppFactory):
             from starlette.responses import StreamingResponse
 
             if runner_method.config.is_stream:
-
-                async def streamer() -> t.AsyncGenerator[bytes, None]:
-                    """
-                    Extract Data from a AsyncGenerator[Payload, None]
-                    """
-                    async for p in payload:
-                        yield pickle.dumps(p)
-
                 return StreamingResponse(
-                    streamer(),
+                    stream_encoder(payload),
                     media_type="text/plain",
                     headers={
                         PAYLOAD_META_HEADER: json.dumps({}),
@@ -375,6 +366,16 @@ class RunnerAppFactory(BaseAppFactory):
         return _request_handler
 
 
+async def stream_encoder(
+    payload: t.AsyncGenerator[str, None]
+) -> t.AsyncGenerator[bytes, None]:
+    """
+    Extract Data from a AsyncGenerator[str, None]
+    """
+    async for p in payload:
+        yield p.encode("utf-8")
+
+
 def _deserialize_single_param(request: Request, bs: bytes) -> Params[t.Any]:
     container = request.headers["Payload-Container"]
     meta = json.loads(request.headers["Payload-Meta"])
diff --git a/src/bentoml/io.py b/src/bentoml/io.py
index ca862c21..1157c9f2 100644
--- a/src/bentoml/io.py
+++ b/src/bentoml/io.py
@@ -10,6 +10,7 @@ from ._internal.io_descriptors.numpy import NumpyNdarray
 from ._internal.io_descriptors.pandas import PandasDataFrame
 from ._internal.io_descriptors.pandas import PandasSeries
 from ._internal.io_descriptors.text import Text
+from ._internal.io_descriptors.utils import SSE
 
 __all__ = [
     "File",
@@ -22,4 +23,5 @@ __all__ = [
     "PandasSeries",
     "Text",
     "from_spec",
+    "SSE",
 ]


Test patch:
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
new file mode 100644
index 00000000..32e0138c
--- /dev/null
+++ b/tests/unit/test_utils.py
@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+from bentoml._internal.server.runner_app import stream_encoder
+from bentoml.io import SSE
+
+if TYPE_CHECKING:
+    from typing import AsyncGenerator
+    from typing import AsyncIterator
+
+
+async def reverse_proxy(
+    iterator: AsyncIterator[bytes],
+    chunk_size: int = 5,
+) -> AsyncGenerator[bytes, None]:
+    """
+    A function that will split the data into chunks. This is a mimic of a
+    real HTTP reverse proxy that will split the data into chunks.
+    """
+    async for data in iterator:
+        if len(data) > chunk_size:
+            while data:
+                yield data[:chunk_size]
+                data = data[chunk_size:]
+        else:
+            yield data
+
+
+async def stream_decoder(
+    iterator: AsyncIterator[bytes],
+) -> AsyncGenerator[str, None]:
+    """
+    Converting bytes to string
+    """
+    async for data in iterator:
+        yield data.decode("utf-8")
+
+
+@pytest.mark.asyncio  # type: ignore
+async def test_sse():
+    # \n\n is a forbidden sequence in SSE
+    test_datas = [
+        SSE(data=""),
+        SSE(data=" "),
+        SSE(data=":"),
+        SSE(data="long data i am longer than buffer"),
+        SSE(data="my data", event="my event"),
+        SSE(data="my data", id="a123cbx"),
+        SSE(data="my data", retry=123),
+        SSE(data="my data", event="my event", id="a123cbx", retry=123),
+        SSE(data="my data\nanother line", event="my event", id="a123cbx", retry=123),
+        SSE(
+            data="my data\nanother line\nand another line",
+            event="my event",
+            id="a123cbx",
+            retry=123,
+        ),
+    ]
+
+    async def runner_func():
+        for data in test_datas:
+            yield data.marshal()
+
+    runner_iterator = stream_decoder(
+        reverse_proxy(stream_encoder(runner_func()), chunk_size=5)
+    )
+
+    async def api_server_func() -> AsyncGenerator[SSE, None]:
+        async for data in SSE.from_iterator(runner_iterator):
+            yield data
+
+    results = [data async for data in api_server_func()]
+    expected_results = [data for data in test_datas]
+    assert results == expected_results