Repository: InternLM/lmdeploy

Issue:
bug-fix: when using stream is False, continuous batching doesn't work
## Motivation

bug fix : when using stream=False, the continuous batching is not working.
See details in this issue: https://github.com/InternLM/lmdeploy/issues/308

## Modification

In non-streaming mode, the main thread which is responsible for receiving request is stuck in the following code since self.que.get() will wait until the queue is not empty.

But when using non-stream mode, the queue is always empty unless a request is finished.

So the newly coming request won't be processed before the last request is done. So the batching doesn't work.

```python
 while True:
            while self.que.qsize() > 1:
                self.que.get()
            finish, tm_outputs = self.que.get()
            outputs = _tm_dict_to_torch_dict(tm_outputs)
```

I modify here to avoid the long wait. just see the commits.

I'm not sure if there is a better way to fix this. If any, please comment.

I have a question, what is this line for?  Won't it cause output tokens being lost???
```python
while self.que.qsize() > 1: self.que.get()
```

## BC-breaking (Optional)

Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.

## Use cases (Optional)

If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.

## Checklist

1. Pre-commit or other linting tools are used to fix the potential lint issues.
2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
4. The documentation has been modified accordingly, like docstring or example tutorials.



Discussion:
@sleepwalker2017 May fix linting problems. You can install pre-commit in your host, which can help fix them 
```shell
pip install pre-commit
# Move the the root path of lmdeploy repository and execute the following commands
pre-commit install
```
You can run `pre-commit run --all-files`. It will report and try to fix linting errors

Another method for this is we do `stream=True` inside the restful API even though the client passed `stream=False`. And we return the final result all at once.
> Another method for this is we do `stream=True` inside the restful API even though the client passed `stream=False`. And we return the final result all at once.

Hi, I think it's a better solution. As my solution keeps the main thread running in the while loop doing little work.

I tried this method, it works well. 

Does the openai generate function also needs to be modified? 

I'm not familiar with the openai generate function and don't know how to test it.

If it need more changes, you can refer to my changes and send a new pr to replace this one.

Thank you.

Not a final solution. Maybe there are other methods. @grimoire Any comments?
> Not a final solution. Maybe there are other methods. @grimoire Any comments?

The core problem is that: the main thread waits for the output queue until there is some output data generated.

So when non-stream request come, the thread can't process new requests.

I have a question, what is this line for?

```
while self.que.qsize() > 1: self.que.get()
```
why do you lose all results until it only has 1 left?

@AllentDan 




The turbomind instance is designed to process the session in it's **own threads**. I think it is the **coroutine** that block the processing of continuous batching. Await the queue getting might be helpful.

> why do you lose all results until it only has 1 left?

assume we have queue:

```
que = [
  [1]
  [1, 2]
  [1, 2, 3]
  [1, 2, 3, 4]
]
```

Output the last one `[1,2,3,4]` and throw other tokens away is enough for our task. We just need to perform memory copy once.
> The turbomind instance is designed to process the session in it's **own threads**. I think it is the **coroutine** that block the processing of continuous batching. Await the queue getting might be helpful.
> 
> > why do you lose all results until it only has 1 left?
> 
> assume we have queue:
> 
> ```
> que = [
>   [1]
>   [1, 2]
>   [1, 2, 3]
>   [1, 2, 3, 4]
> ]
> ```
> 
> Output the last one `[1,2,3,4]` and throw other tokens away is enough for our task. We just need to perform memory copy once.

I try this, but it doesn't work. Always 1 thread appending tasks. Is there anything wrong with it?
 I'm not familiar with the coroutine api in python. I use python 3.8, the api may be different from  3.10.

<img width="638" alt="image" src="https://github.com/InternLM/lmdeploy/assets/26128514/d6ed304e-b06a-4bfe-a68a-ad6af4fe222b">

Can we warp the queue get while loop inside the wrapper?
Or we can initialize a counter with current queue size n, pop the top n-1 from the queue and output next one?
> 

the problem comes from the situation when the queue is empty.

I non-stream mode, the queue is empty for a long time until a sequence is finished. So the line waits for a long time and can't process new requests. That's the problem.

I think this line is the cause for non-stream modes.
```
            finish, tm_outputs = self.que.get()
```

I see. I guess that Queue is not awaitable since there is another Queue in module `asyncio` (which we can not use since it is not thread safe).
Can we use threads instead of coroutine in restful API? Or can we make the queue awaitable?
> I see. I guess that Queue is not awaitable since there is another Queue in module `asyncio` (which we can not use since it is not thread safe). Can we use threads instead of coroutine in restful API? Or can we make the queue awaitable?

The code I posted above is trying to wrap the que.get() in a awaitable mode, but it doesn't work, I wonder why is that. 

Using multiple threads is ok, the api_server file may be modified a lot or it needs to be refactored.
An object is awaitable only when it has `__await__` attribute.
I think we can merge this PR as a quick fix and find a better solution in the future(threads or something better).



Resolving patch:
diff --git a/lmdeploy/serve/openai/api_server.py b/lmdeploy/serve/openai/api_server.py
--- a/lmdeploy/serve/openai/api_server.py
+++ b/lmdeploy/serve/openai/api_server.py
@@ -57,9 +57,10 @@ def create_error_response(status: HTTPStatus, message: str):
         status (HTTPStatus): HTTP status codes and reason phrases
         message (str): error message
     """
-    return JSONResponse(ErrorResponse(message=message,
-                                      type='invalid_request_error').dict(),
-                        status_code=status.value)
+    return JSONResponse(
+        ErrorResponse(message=message,
+                      type='invalid_request_error',
+                      code=status.value).dict())
 
 
 async def check_request(request) -> Optional[JSONResponse]:
@@ -117,7 +118,7 @@ async def chat_completions_v1(request: ChatCompletionRequest,
     result_generator = VariableInterface.async_engine.generate_openai(
         request.messages,
         instance_id,
-        request.stream,
+        True,  # always use stream to enable batching
         request.renew_session,
         request_output_len=request.max_tokens if request.max_tokens else 512,
         stop=request.stop,
@@ -130,7 +131,7 @@ async def abort_request() -> None:
         async for _ in VariableInterface.async_engine.generate_openai(
                 request.messages,
                 instance_id,
-                request.stream,
+                True,
                 request.renew_session,
                 stop=True):
             pass
@@ -188,6 +189,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]:
 
     # Non-streaming response
     final_res = None
+    text = ''
     async for res in result_generator:
         if await raw_request.is_disconnected():
             # Abort the request if the client disconnects.
@@ -195,11 +197,12 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]:
             return create_error_response(HTTPStatus.BAD_REQUEST,
                                          'Client disconnected')
         final_res = res
+        text += res.response
     assert final_res is not None
     choices = []
     choice_data = ChatCompletionResponseChoice(
         index=0,
-        message=ChatMessage(role='assistant', content=final_res.response),
+        message=ChatMessage(role='assistant', content=text),
         finish_reason=final_res.finish_reason,
     )
     choices.append(choice_data)
@@ -308,7 +311,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
         finish_reason = None
         async for out in generation:
             text += out.response
-            tokens += out.generate_token_len
+            tokens = out.generate_token_len
             finish_reason = out.finish_reason
         ret = {'text': text, 'tokens': tokens, 'finish_reason': finish_reason}
         return JSONResponse(ret)


Test patch: