Repository: xorbitsai/inference
Pull request: ENH: Safe iterate stream of ggml model
- [x] Use FastAPI stream.
- [x] Fix parallel ggml stream undefined behaviour.

Fix: https://github.com/xorbitsai/inference/issues/448

Problem statement:
BUG: Incorrect traceback if raises an exception when stream iteration 
### Describe the bug
A clear and concise description of what the bug is.

### To Reproduce
To help us to reproduce this bug, please provide information below:

```python
xinference/tests/test_client.py:188 (test_RESTful_client)
self = <urllib3.response.HTTPResponse object at 0x2b3f437c0>

    def _update_chunk_length(self):
        # First, we'll figure out length of a chunk and then
        # we'll try to read it from socket.
        if self.chunk_left is not None:
            return
        line = self._fp.fp.readline()
        line = line.split(b";", 1)[0]
        try:
>           self.chunk_left = int(line, 16)
E           ValueError: invalid literal for int() with base 16: b''

../../.pyenv/versions/3.11.4/lib/python3.11/site-packages/urllib3/response.py:761: ValueError

During handling of the above exception, another exception occurred:

self = <urllib3.response.HTTPResponse object at 0x2b3f437c0>

    @contextmanager
    def _error_catcher(self):
        """
        Catch low-level python exceptions, instead re-raising urllib3
        variants, so that low-level exceptions are not leaked in the
        high-level api.
    
        On exit, release the connection back to the pool.
        """
        clean_exit = False
    
        try:
            try:
>               yield

../../.pyenv/versions/3.11.4/lib/python3.11/site-packages/urllib3/response.py:444: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../.pyenv/versions/3.11.4/lib/python3.11/site-packages/urllib3/response.py:828: in read_chunked
    self._update_chunk_length()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <urllib3.response.HTTPResponse object at 0x2b3f437c0>

    def _update_chunk_length(self):
        # First, we'll figure out length of a chunk and then
        # we'll try to read it from socket.
        if self.chunk_left is not None:
            return
        line = self._fp.fp.readline()
        line = line.split(b";", 1)[0]
        try:
            self.chunk_left = int(line, 16)
        except ValueError:
            # Invalid chunked protocol response, abort.
            self.close()
>           raise InvalidChunkLength(self, line)
E           urllib3.exceptions.InvalidChunkLength: InvalidChunkLength(got length b'', 0 bytes read)

../../.pyenv/versions/3.11.4/lib/python3.11/site-packages/urllib3/response.py:765: InvalidChunkLength

During handling of the above exception, another exception occurred:

    def generate():
        # Special case for urllib3.
        if hasattr(self.raw, "stream"):
            try:
>               yield from self.raw.stream(chunk_size, decode_content=True)

../../.pyenv/versions/3.11.4/lib/python3.11/site-packages/requests/models.py:816: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../.pyenv/versions/3.11.4/lib/python3.11/site-packages/urllib3/response.py:624: in stream
    for line in self.read_chunked(amt, decode_content=decode_content):
../../.pyenv/versions/3.11.4/lib/python3.11/site-packages/urllib3/response.py:816: in read_chunked
    with self._error_catcher():
../../.pyenv/versions/3.11.4/lib/python3.11/contextlib.py:155: in __exit__
    self.gen.throw(typ, value, traceback)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <urllib3.response.HTTPResponse object at 0x2b3f437c0>

    @contextmanager
    def _error_catcher(self):
        """
        Catch low-level python exceptions, instead re-raising urllib3
        variants, so that low-level exceptions are not leaked in the
        high-level api.
    
        On exit, release the connection back to the pool.
        """
        clean_exit = False
    
        try:
            try:
                yield
    
            except SocketTimeout:
                # FIXME: Ideally we'd like to include the url in the ReadTimeoutError but
                # there is yet no clean way to get at it from this context.
                raise ReadTimeoutError(self._pool, None, "Read timed out.")
    
            except BaseSSLError as e:
                # FIXME: Is there a better way to differentiate between SSLErrors?
                if "read operation timed out" not in str(e):
                    # SSL errors related to framing/MAC get wrapped and reraised here
                    raise SSLError(e)
    
                raise ReadTimeoutError(self._pool, None, "Read timed out.")
    
            except (HTTPException, SocketError) as e:
                # This includes IncompleteRead.
>               raise ProtocolError("Connection broken: %r" % e, e)
E               urllib3.exceptions.ProtocolError: ("Connection broken: InvalidChunkLength(got length b'', 0 bytes read)", InvalidChunkLength(got length b'', 0 bytes read))

../../.pyenv/versions/3.11.4/lib/python3.11/site-packages/urllib3/response.py:461: ProtocolError
```

1. Your Python version.
2. The version of xinference you use.
3. Versions of crucial packages.
4. Full stack of the error.
5. Minimized code to reproduce the error.

### Expected behavior
A clear and concise description of what you expected to happen.

### Additional context
Add any other context about the problem here.

Hints from issue discussion:

Commit message: Avoid parallel stream iteration on ggml model

Patch to review:
diff --git a/setup.cfg b/setup.cfg
index 7b92e1f6dd..2acfd1a976 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -34,7 +34,6 @@ install_requires =
     pydantic<2
     fastapi
     uvicorn
-    sse_starlette
     huggingface-hub>=0.14.1,<1.0
     typing_extensions
     fsspec
diff --git a/xinference/client.py b/xinference/client.py
index a69579f52d..c7af0b8747 100644
--- a/xinference/client.py
+++ b/xinference/client.py
@@ -208,15 +208,15 @@ def streaming_response_iterator(
 
 # Duplicate code due to type hint issues
 def chat_streaming_response_iterator(
-    response_lines: Iterator[bytes],
+    response_chunk: Iterator[bytes],
 ) -> Iterator["ChatCompletionChunk"]:
     """
     Create an Iterator to handle the streaming type of generation.
 
     Parameters
     ----------
-    response_lines: Iterator[bytes]
-        Generated lines by the Model Generator.
+    response_chunk: Iterator[bytes]
+        Generated chunk by the Model Generator.
 
     Returns
     -------
@@ -225,11 +225,12 @@ def chat_streaming_response_iterator(
 
     """
 
-    for line in response_lines:
-        line = line.strip()
-        if line.startswith(b"data:"):
-            data = json.loads(line.decode("utf-8").replace("data: ", "", 1))
-            yield data
+    for line in response_chunk:
+        content = json.loads(line.decode("utf-8"))
+        error = content.get("error", None)
+        if error is not None:
+            raise Exception(str(error))
+        yield content
 
 
 class RESTfulModelHandle:
@@ -405,7 +406,9 @@ def chat(
             )
 
         if stream:
-            return chat_streaming_response_iterator(response.iter_lines())
+            return chat_streaming_response_iterator(
+                response.iter_content(chunk_size=None)
+            )
 
         response_data = response.json()
         return response_data
@@ -469,7 +472,9 @@ def chat(
             )
 
         if stream:
-            return chat_streaming_response_iterator(response.iter_lines())
+            return chat_streaming_response_iterator(
+                response.iter_content(chunk_size=None)
+            )
 
         response_data = response.json()
         return response_data
diff --git a/xinference/core/model.py b/xinference/core/model.py
index 5bde398808..b228dabcea 100644
--- a/xinference/core/model.py
+++ b/xinference/core/model.py
@@ -110,6 +110,9 @@ def load(self):
 
     async def _wrap_generator(self, ret: Any):
         if inspect.isgenerator(ret):
+            if self._lock is not None:
+                # Make sure only one iterator is valid.
+                self._generators.clear()
             generator_uid = str(uuid.uuid1())
             self._generators[generator_uid] = ret
             return IteratorWrapper(
@@ -168,6 +171,11 @@ async def next(
 
         def _wrapper():
             try:
+                if self._lock is not None and generator_uid not in self._generators:
+                    raise Exception(
+                        f"The generator {generator_uid} is invalid, "
+                        f"parallel iteration is not supported by ggml."
+                    )
                 return next(self._generators[generator_uid])
             except StopIteration:
                 return stop
diff --git a/xinference/core/restful_api.py b/xinference/core/restful_api.py
index 59426d0522..b09fe70e37 100644
--- a/xinference/core/restful_api.py
+++ b/xinference/core/restful_api.py
@@ -19,19 +19,15 @@
 import sys
 import threading
 import warnings
-from functools import partial
 from typing import Any, Dict, List, Literal, Optional, Union
 
-import anyio
 import gradio as gr
 import xoscar as xo
-from anyio.streams.memory import MemoryObjectSendStream
 from fastapi import APIRouter, FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, Field
-from sse_starlette.sse import EventSourceResponse
 from starlette.responses import RedirectResponse
 from typing_extensions import NotRequired, TypedDict
 from uvicorn import Config, Server
@@ -517,32 +513,16 @@ async def create_completion(self, request: Request, body: CreateCompletionReques
             raise HTTPException(status_code=500, detail=str(e))
 
         if body.stream:
-            # create a pair of memory object streams
-            send_chan, recv_chan = anyio.create_memory_object_stream(10)
-
-            async def event_publisher(inner_send_chan: MemoryObjectSendStream):
-                async with inner_send_chan:
-                    try:
-                        iterator = await model.generate(body.prompt, kwargs)
-                        async for chunk in iterator:
-                            await inner_send_chan.send(dict(data=json.dumps(chunk)))
-                            if await request.is_disconnected():
-                                raise anyio.get_cancelled_exc_class()()
-                    except anyio.get_cancelled_exc_class() as e:
-                        logger.warning("disconnected")
-                        with anyio.move_on_after(1, shield=True):
-                            logger.warning(
-                                f"Disconnected from client (via refresh/close) {request.client}"
-                            )
-                            await inner_send_chan.send(dict(closing=True))
-                            raise e
-                    except Exception as e:
-                        raise HTTPException(status_code=500, detail=str(e))
-
-            return EventSourceResponse(
-                recv_chan, data_sender_callable=partial(event_publisher, send_chan)
-            )
 
+            async def encode_iterator():
+                try:
+                    iterator = await model.generate(body.prompt, kwargs)
+                    async for item in iterator:
+                        yield json.dumps(item)
+                except Exception as ex:
+                    yield json.dumps({"error": str(ex)})
+
+            return StreamingResponse(encode_iterator())
         else:
             try:
                 return await model.generate(body.prompt, kwargs)
@@ -640,37 +620,21 @@ async def create_chat_completion(
             )
 
         if body.stream:
-            # create a pair of memory object streams
-            send_chan, recv_chan = anyio.create_memory_object_stream(10)
-
-            async def event_publisher(inner_send_chan: MemoryObjectSendStream):
-                async with inner_send_chan:
-                    try:
-                        if is_chatglm_ggml:
-                            iterator = await model.chat(prompt, chat_history, kwargs)
-                        else:
-                            iterator = await model.chat(
-                                prompt, system_prompt, chat_history, kwargs
-                            )
-                        async for chunk in iterator:
-                            await inner_send_chan.send(dict(data=json.dumps(chunk)))
-                            if await request.is_disconnected():
-                                raise anyio.get_cancelled_exc_class()()
-                    except anyio.get_cancelled_exc_class() as e:
-                        logger.warning("disconnected")
-                        with anyio.move_on_after(1, shield=True):
-                            logger.warning(
-                                f"Disconnected from client (via refresh/close) {request.client}"
-                            )
-                            await inner_send_chan.send(dict(closing=True))
-                            raise e
-                    except Exception as e:
-                        raise HTTPException(status_code=500, detail=str(e))
-
-            return EventSourceResponse(
-                recv_chan, data_sender_callable=partial(event_publisher, send_chan)
-            )
 
+            async def encode_iterator():
+                try:
+                    if is_chatglm_ggml:
+                        iterator = await model.chat(prompt, chat_history, kwargs)
+                    else:
+                        iterator = await model.chat(
+                            prompt, system_prompt, chat_history, kwargs
+                        )
+                    async for item in iterator:
+                        yield json.dumps(item)
+                except Exception as ex:
+                    yield json.dumps({"error": str(ex)})
+
+            return StreamingResponse(encode_iterator())
         else:
             try:
                 if is_chatglm_ggml:
diff --git a/xinference/tests/test_client.py b/xinference/tests/test_client.py
index 5fe48e16b0..71f61f480b 100644
--- a/xinference/tests/test_client.py
+++ b/xinference/tests/test_client.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+from concurrent.futures import ThreadPoolExecutor
 
 import pytest
 
@@ -227,12 +228,30 @@ def test_RESTful_client(setup):
     completion = model.chat("What is the capital of France?")
     assert "content" in completion["choices"][0]["message"]
 
-    streaming_response = model.chat(
-        prompt="What is the capital of France?", generate_config={"stream": True}
-    )
-
-    for chunk in streaming_response:
-        assert "content" or "role" in chunk["choices"][0]["delta"]
+    def _check_stream():
+        streaming_response = model.chat(
+            prompt="What is the capital of France?", generate_config={"stream": True}
+        )
+        for chunk in streaming_response:
+            print(chunk)
+            assert "content" or "role" in chunk["choices"][0]["delta"]
+
+    _check_stream()
+
+    results = []
+    with ThreadPoolExecutor() as executor:
+        for _ in range(3):
+            r = executor.submit(_check_stream)
+            results.append(r)
+    # Parallel iterates on a ggml model, only one can be success.
+    error_count = 0
+    for r in results:
+        try:
+            r.result()
+        except Exception as ex:
+            assert "parallel iteration" in str(ex)
+            error_count += 1
+    assert error_count == 2
 
     client.terminate_model(model_uid=model_uid)
     assert len(client.list_models()) == 0


Code review comments:
diff_hunk: @@ -517,32 +513,16 @@ async def create_completion(self, request: Request, body: CreateCompletionReques
             raise HTTPException(status_code=500, detail=str(e))
 
         if body.stream:
-            # create a pair of memory object streams
-            send_chan, recv_chan = anyio.create_memory_object_stream(10)
-
-            async def event_publisher(inner_send_chan: MemoryObjectSendStream):
-                async with inner_send_chan:
-                    try:
-                        iterator = await model.generate(body.prompt, kwargs)
-                        async for chunk in iterator:
-                            await inner_send_chan.send(dict(data=json.dumps(chunk)))
-                            if await request.is_disconnected():
-                                raise anyio.get_cancelled_exc_class()()
-                    except anyio.get_cancelled_exc_class() as e:
-                        logger.warning("disconnected")
-                        with anyio.move_on_after(1, shield=True):
-                            logger.warning(
-                                f"Disconnected from client (via refresh/close) {request.client}"
-                            )
-                            await inner_send_chan.send(dict(closing=True))
-                            raise e
-                    except Exception as e:
-                        raise HTTPException(status_code=500, detail=str(e))
-
-            return EventSourceResponse(
-                recv_chan, data_sender_callable=partial(event_publisher, send_chan)
-            )
 
+            async def encode_iterator():
+                try:
+                    iterator = await model.generate(body.prompt, kwargs)
+                    async for item in iterator:
+                        yield json.dumps(item)
+                except Exception as ex:
+                    yield json.dumps({"error": str(ex)})
path: xinference/core/restful_api.py
text: @user1:
Better log the error for debugging since `str(ex)` may not contain the full invocation stack.

@author:
Good suggestion.

diff_hunk: @@ -110,6 +110,9 @@ def load(self):
 
     async def _wrap_generator(self, ret: Any):
         if inspect.isgenerator(ret):
+            if self._lock is not None:
+                # Make sure only one iterator is valid.
+                self._generators.clear()
path: xinference/core/model.py
text: @user1:
How about simply raise an exception when self._lock is not None and self._generators is not empty?

@author:
Current behavior similars to C++ iterator, append data invalidate old iterators.

@author:
Fixed. Thanks.

diff_hunk: @@ -517,32 +513,16 @@ async def create_completion(self, request: Request, body: CreateCompletionReques
             raise HTTPException(status_code=500, detail=str(e))
 
         if body.stream:
-            # create a pair of memory object streams
-            send_chan, recv_chan = anyio.create_memory_object_stream(10)
-
-            async def event_publisher(inner_send_chan: MemoryObjectSendStream):
-                async with inner_send_chan:
-                    try:
-                        iterator = await model.generate(body.prompt, kwargs)
-                        async for chunk in iterator:
-                            await inner_send_chan.send(dict(data=json.dumps(chunk)))
-                            if await request.is_disconnected():
-                                raise anyio.get_cancelled_exc_class()()
-                    except anyio.get_cancelled_exc_class() as e:
-                        logger.warning("disconnected")
-                        with anyio.move_on_after(1, shield=True):
-                            logger.warning(
-                                f"Disconnected from client (via refresh/close) {request.client}"
-                            )
-                            await inner_send_chan.send(dict(closing=True))
-                            raise e
-                    except Exception as e:
-                        raise HTTPException(status_code=500, detail=str(e))
-
-            return EventSourceResponse(
-                recv_chan, data_sender_callable=partial(event_publisher, send_chan)
-            )
 
+            async def encode_iterator():
path: xinference/core/restful_api.py
text: @user1:
What does `encode_iterator` mean? Since GPTs are decoder-only models, maybe `decode_iterator` sounds more reasonable?

Or simply name it `stream_results`?

@author:
OK

diff_hunk: @@ -168,6 +171,11 @@ async def next(
 
         def _wrapper():
             try:
+                if self._lock is not None and generator_uid not in self._generators:
+                    raise Exception(
path: xinference/core/model.py
text: @user1:
How about raising an RuntimeException since this is totally unexpected？

@author:
OK