Repository: bentoml/BentoML
Pull request #4330 (2023-12-11 19:36:29)

fix(stream): streaming enable to work with proxy
## What does this PR address?
Relates to #4310 
1) streaming now just stream raw text instead of pickled data container
2) provides `to_sse` and `from_sse` to be used at boundaries (api server and runner) if user wants to stream text in SSE format

<!--
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?


aarnphm: @jianshen92 I think might be worth merging this to 1.2
jianshen92: > @jianshen92 I think might be worth merging this to 1.2

@bojiang should I be merging to 1.2 instead?
bojiang: ‌‌‌@aarnphm This is a bug fix. It's affecting the deployment of openllm under certain circumstances. There's still some time for 1.2 to release. Let's merge to main first?
bojiang: ‌‌‌‌By the way, I think what Aaron meant might not be that we don't offer these utils, but rather that we provide utils which then return SSE objects. @jianshen92 
I wrote some unit test for the utils. Let me push them here.
aarnphm: @bojiang @jianshen92 what is the status of this PR?
bojiang: @aarnphm Jianshen is still working on the utils.
jianshen92: Usage:

runner
```
class Runner(bentoml.Runnable):
    ...
    
    async def raw_text_streaming():
        yield "this is valid"
    
    async def sse_streaming():
        yield SSE(data="my data", event="my event", id="a123cbx", retry=123)
```

api-server
```
# if we expect raw text, nothing changes. we dont care if they get chunked 
# if we expect a SSE text stream, we can use the `from_chunks` method to reconstruct the chunked data
async for event in SSE.from_chunks(runner_iterator):
    yield event.to_str()
    yield event.data
```

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..ccb6ba32
--- /dev/null
+++ b/src/bentoml/_internal/io_descriptors/utils.py
@@ -0,0 +1,68 @@
+import io
+from typing import AsyncGenerator
+from typing import AsyncIterator
+from typing import Optional
+
+import attr
+
+from bentoml.exceptions import BentoMLException
+
+
+@attr.s(auto_attribs=True)
+class SSE:
+    data: str
+    event: Optional[str] = None
+    id: Optional[str] = None
+    retry: Optional[int] = None
+
+    def marshal(self) -> str:
+        sse_data = ""
+        if self.id is not None:
+            sse_data += f"id: {self.id}\n"
+        if self.event is not None:
+            sse_data += f"event: {self.event}\n"
+        if self.retry is not None:
+            sse_data += f"retry: {self.retry}\n"
+
+        # Handle multi-line data
+        data_lines = self.data.split("\n")
+        for line in data_lines:
+            sse_data += f"data: {line}\n"
+
+        # Add an extra newline to mark the end of the message
+        sse_data += "\n"
+        return sse_data
+
+    @staticmethod
+    async def from_iterator(
+        async_iterator: AsyncIterator[str],
+    ) -> AsyncGenerator["SSE", None]:
+        buffer = io.StringIO()
+        async for chunk in async_iterator:
+            buffer.write(chunk)
+            # Move to the beginning to read content
+            buffer.seek(0, io.SEEK_SET)
+            content = buffer.read()
+
+            while "\n\n" in content:
+                event_text, _, remaining = content.partition("\n\n")
+                fields = {"data": "", "event": None, "id": None, "retry": None}
+                data = []
+                for line in event_text.split("\n"):
+                    key, _, value = line.partition(": ")
+                    if key == "data":
+                        data.append(value)
+                    elif key == "retry":
+                        fields[key] = int(value)
+                    elif key in fields:
+                        fields[key] = value
+                    else:
+                        raise BentoMLException("Invalid SSE message")
+                fields["data"] = "\n".join(data)
+                yield SSE(**fields)
+
+                # Reset buffer
+                buffer.seek(0)
+                buffer.truncate()
+                buffer.write(remaining)
+                content = remaining
diff --git a/src/bentoml/_internal/runner/runner.py b/src/bentoml/_internal/runner/runner.py
index 59a6cd99..80bb5547 100644
--- a/src/bentoml/_internal/runner/runner.py
+++ b/src/bentoml/_internal/runner/runner.py
@@ -56,7 +56,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 6166ce6c..7861a2dd 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(),)
 
@@ -348,16 +347,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({}),
@@ -373,6 +364,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