Enable the Gradio server to call inference services through the RESTful API (#287)

* app use async engine

* add stop logic

* app update cancel

* app support restful-api

* update doc and use the right model name

* set doc url root

* add comments

* add an example

* renew_session

* update readme.md

* resolve comments

* Update restful_api.md

* Update restful_api.md

* Update restful_api.md

---------
Co-authored-by: tpoisonooo <khj.application@aliyun.com>

README.md

@@ -133,6 +133,32 @@ python3 -m lmdeploy.serve.gradio.app ./workspace
 
 ![](https://github.com/InternLM/lmdeploy/assets/67539920/08d1e6f2-3767-44d5-8654-c85767cec2ab)
 
+#### Serving with Restful API
+
+Launch inference server by:
+
+```shell
+python3 -m lmdeploy.serve.openai.api_server ./workspace server_ip server_port --instance_num 32 --tp 1
+```
+
+Then, you can communicate with it by command line,
+
+```shell
+# restful_api_url is what printed in api_server.py, e.g. http://localhost:23333
+python -m lmdeploy.serve.openai.api_client restful_api_url
+```
+
+or webui,
+
+```shell
+# restful_api_url is what printed in api_server.py, e.g. http://localhost:23333
+# server_ip and server_port here are for gradio ui
+# example: python -m lmdeploy.serve.gradio.app http://localhost:23333 localhost 6006 --restful_api True
+python -m lmdeploy.serve.gradio.app restful_api_url server_ip --restful_api True
+```
+
+Refer to [restful_api.md](docs/en/restful_api.md) for more details.
+
 #### Serving with Triton Inference Server
 
 Launch inference server by:

README_zh-CN.md

@@ -133,6 +133,32 @@ python3 -m lmdeploy.serve.gradio.app ./workspace
 
 ![](https://github.com/InternLM/lmdeploy/assets/67539920/08d1e6f2-3767-44d5-8654-c85767cec2ab)
 
+#### 通过 Restful API 部署服务
+
+使用下面的命令启动推理服务：
+
+```shell
+python3 -m lmdeploy.serve.openai.api_server ./workspace server_ip server_port --instance_num 32 --tp 1
+```
+
+你可以通过命令行方式与推理服务进行对话：
+
+```shell
+# restful_api_url is what printed in api_server.py, e.g. http://localhost:23333
+python -m lmdeploy.serve.openai.api_client restful_api_url
+```
+
+也可以通过 WebUI 方式来对话：
+
+```shell
+# restful_api_url is what printed in api_server.py, e.g. http://localhost:23333
+# server_ip and server_port here are for gradio ui
+# example: python -m lmdeploy.serve.gradio.app http://localhost:23333 localhost 6006 --restful_api True
+python -m lmdeploy.serve.gradio.app restful_api_url server_ip --restful_api True
+```
+
+更多详情可以查阅 [restful_api.md](docs/zh_cn/restful_api.md)。
+
 #### 通过容器部署推理服务
 
 使用下面的命令启动推理服务：

docs/en/restful_api.md

@@ -3,10 +3,10 @@
 ### Launch Service
 
 ```shell
-python3 -m lmdeploy.serve.openai.api_server ./workspace server_name server_port --instance_num 32 --tp 1
+python3 -m lmdeploy.serve.openai.api_server ./workspace 0.0.0.0 server_port --instance_num 32 --tp 1
 ```
 
-Then, the user can open the swagger UI: http://{server_name}:{server_port}/docs for the detailed api usage.
+Then, the user can open the swagger UI: `http://{server_ip}:{server_port}` for the detailed api usage.
 We provide four restful api in total. Three of them are in OpenAI format. However, we recommend users try
 our own api which provides more arguments for users to modify. The performance is comparatively better.
 
@@ -50,16 +50,29 @@ def get_streaming_response(prompt: str,
 
 
 for output, tokens in get_streaming_response(
-        "Hi, how are you?", "http://{server_name}:{server_port}/generate", 0,
+        "Hi, how are you?", "http://{server_ip}:{server_port}/generate", 0,
         512):
     print(output, end='')
 ```
 
-### Golang/Rust
+### Java/Golang/Rust
 
-Golang can also build a http request to use the service. You may refer
-to [the blog](https://pkg.go.dev/net/http) for details to build own client.
-Besides, Rust supports building a client in [many ways](https://blog.logrocket.com/best-rust-http-client/).
+May use [openapi-generator-cli](https://github.com/OpenAPITools/openapi-generator-cli) to convert `http://{server_ip}:{server_port}/openapi.json` to java/rust/golang client.
+Here is an example:
+
+```shell
+$ docker run -it --rm -v ${PWD}:/local openapitools/openapi-generator-cli generate -i /local/openapi.json -g rust -o /local/rust
+
+$ ls rust/*
+rust/Cargo.toml  rust/git_push.sh  rust/README.md
+
+rust/docs:
+ChatCompletionRequest.md  EmbeddingsRequest.md  HttpValidationError.md  LocationInner.md  Prompt.md
+DefaultApi.md             GenerateRequest.md    Input.md                Messages.md       ValidationError.md
+
+rust/src:
+apis  lib.rs  models
+```
 
 ### cURL
 
@@ -68,13 +81,13 @@ cURL is a tool for observing the output of the api.
 List Models:
 
 ```bash
-curl http://{server_name}:{server_port}/v1/models
+curl http://{server_ip}:{server_port}/v1/models
 ```
 
 Generate:
 
 ```bash
-curl http://{server_name}:{server_port}/generate \
+curl http://{server_ip}:{server_port}/generate \
   -H "Content-Type: application/json" \
   -d '{
     "model": "internlm-chat-7b",
@@ -87,7 +100,7 @@ curl http://{server_name}:{server_port}/generate \
 Chat Completions:
 
 ```bash
-curl http://{server_name}:{server_port}/v1/chat/completions \
+curl http://{server_ip}:{server_port}/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
     "model": "internlm-chat-7b",
@@ -98,7 +111,7 @@ curl http://{server_name}:{server_port}/v1/chat/completions \
 Embeddings:
 
 ```bash
-curl http://{server_name}:{server_port}/v1/embeddings \
+curl http://{server_ip}:{server_port}/v1/embeddings \
   -H "Content-Type: application/json" \
   -d '{
     "model": "internlm-chat-7b",
@@ -106,6 +119,26 @@ curl http://{server_name}:{server_port}/v1/embeddings \
   }'
 ```
 
+### CLI client
+
+There is a client script for restful api server.
+
+```shell
+# restful_api_url is what printed in api_server.py, e.g. http://localhost:23333
+python -m lmdeploy.serve.openai.api_client restful_api_url
+```
+
+### webui
+
+You can also test restful-api through webui.
+
+```shell
+# restful_api_url is what printed in api_server.py, e.g. http://localhost:23333
+# server_ip and server_port here are for gradio ui
+# example: python -m lmdeploy.serve.gradio.app http://localhost:23333 localhost 6006 --restful_api True
+python -m lmdeploy.serve.gradio.app restful_api_url server_ip --restful_api True
+```
+
 ### FAQ
 
 1. When user got `"finish_reason":"length"` which means the session is too long to be continued.

docs/zh_cn/restful_api.md

@@ -5,10 +5,10 @@
 运行脚本
 
 ```shell
-python3 -m lmdeploy.serve.openai.api_server ./workspace server_name server_port --instance_num 32 --tp 1
+python3 -m lmdeploy.serve.openai.api_server ./workspace 0.0.0.0 server_port --instance_num 32 --tp 1
 ```
 
-然后用户可以打开 swagger UI: http://{server_name}:{server_port}/docs 详细查看所有的 API 及其使用方法。
+然后用户可以打开 swagger UI: `http://{server_ip}:{server_port}` 详细查看所有的 API 及其使用方法。
 我们一共提供四个 restful api，其中三个仿照 OpenAI 的形式。不过，我们建议用户用我们提供的另一个 API: `generate`。
 它有更好的性能，提供更多的参数让用户自定义修改。
 
@@ -52,15 +52,29 @@ def get_streaming_response(prompt: str,
 
 
 for output, tokens in get_streaming_response(
-        "Hi, how are you?", "http://{server_name}:{server_port}/generate", 0,
+        "Hi, how are you?", "http://{server_ip}:{server_port}/generate", 0,
         512):
     print(output, end='')
 ```
 
-### Golang/Rust
+### Java/Golang/Rust
 
-Golang 也可以建立 http 请求使用启动的服务，用户可以参考[这篇博客](https://pkg.go.dev/net/http)构建自己的客户端。
-Rust 也有许多[方法](https://blog.logrocket.com/best-rust-http-client/)构建客户端，使用服务。
+可以使用代码生成工具 [openapi-generator-cli](https://github.com/OpenAPITools/openapi-generator-cli) 将 `http://{server_ip}:{server_port}/openapi.json` 转成 java/rust/golang 客户端。
+下面是一个使用示例：
+
+```shell
+$ docker run -it --rm -v ${PWD}:/local openapitools/openapi-generator-cli generate -i /local/openapi.json -g rust -o /local/rust
+
+$ ls rust/*
+rust/Cargo.toml  rust/git_push.sh  rust/README.md
+
+rust/docs:
+ChatCompletionRequest.md  EmbeddingsRequest.md  HttpValidationError.md  LocationInner.md  Prompt.md
+DefaultApi.md             GenerateRequest.md    Input.md                Messages.md       ValidationError.md
+
+rust/src:
+apis  lib.rs  models
+```
 
 ### cURL
 
@@ -69,13 +83,13 @@ cURL 也可以用于查看 API 的输出结果
 查看模型列表：
 
 ```bash
-curl http://{server_name}:{server_port}/v1/models
+curl http://{server_ip}:{server_port}/v1/models
 ```
 
 使用 generate:
 
 ```bash
-curl http://{server_name}:{server_port}/generate \
+curl http://{server_ip}:{server_port}/generate \
   -H "Content-Type: application/json" \
   -d '{
     "model": "internlm-chat-7b",
@@ -88,7 +102,7 @@ curl http://{server_name}:{server_port}/generate \
 Chat Completions:
 
 ```bash
-curl http://{server_name}:{server_port}/v1/chat/completions \
+curl http://{server_ip}:{server_port}/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
     "model": "internlm-chat-7b",
@@ -99,7 +113,7 @@ curl http://{server_name}:{server_port}/v1/chat/completions \
 Embeddings:
 
 ```bash
-curl http://{server_name}:{server_port}/v1/embeddings \
+curl http://{server_ip}:{server_port}/v1/embeddings \
   -H "Content-Type: application/json" \
   -d '{
     "model": "internlm-chat-7b",
@@ -107,6 +121,26 @@ curl http://{server_name}:{server_port}/v1/embeddings \
   }'
 ```
 
+### CLI client
+
+restful api 服务可以通过客户端测试，例如
+
+```shell
+# restful_api_url 就是 api_server 产生的，比如 http://localhost:23333
+python -m lmdeploy.serve.openai.api_client restful_api_url
+```
+
+### webui
+
+也可以直接用 webui 测试使用 restful-api。
+
+```shell
+# restful_api_url 就是 api_server 产生的，比如 http://localhost:23333
+# server_ip 和 server_port 是用来提供 gradio ui 访问服务的
+# 例子: python -m lmdeploy.serve.gradio.app http://localhost:23333 localhost 6006 --restful_api True
+python -m lmdeploy.serve.gradio.app restful_api_url server_ip --restful_api True
+```
+
 ### FAQ
 
 1. 当返回结果结束原因为 `"finish_reason":"length"`，这表示回话长度超过最大值。

lmdeploy/serve/openai/async_engine.py → lmdeploy/serve/async_engine.py

@@ -3,11 +3,11 @@ import asyncio
 import dataclasses
 import os.path as osp
 import random
+from contextlib import contextmanager
 from typing import Literal, Optional
 
 from lmdeploy import turbomind as tm
 from lmdeploy.model import MODELS, BaseModel
-from lmdeploy.turbomind.tokenizer import Tokenizer
 
 
 @dataclasses.dataclass
@@ -30,6 +30,7 @@ class AsyncEngine:
     """
 
     def __init__(self, model_path, instance_num=32, tp=1) -> None:
+        from lmdeploy.turbomind.tokenizer import Tokenizer
         tokenizer_model_path = osp.join(model_path, 'triton_models',
                                         'tokenizer')
         tokenizer = Tokenizer(tokenizer_model_path)
@@ -46,15 +47,22 @@ class AsyncEngine:
         self.starts = [None] * instance_num
         self.steps = {}
 
+    @contextmanager
+    def safe_run(self, instance_id: int, stop: bool = False):
+        self.available[instance_id] = False
+        yield
+        self.available[instance_id] = True
+
     async def get_embeddings(self, prompt):
         prompt = self.model.get_prompt(prompt)
         input_ids = self.tokenizer.encode(prompt)
         return input_ids
 
-    async def get_generator(self, instance_id):
+    async def get_generator(self, instance_id: int, stop: bool = False):
         """Only return the model instance if it is available."""
-        while self.available[instance_id] is False:
-            await asyncio.sleep(0.1)
+        if not stop:
+            while self.available[instance_id] is False:
+                await asyncio.sleep(0.1)
         return self.generators[instance_id]
 
     async def generate(
@@ -104,43 +112,43 @@ class AsyncEngine:
         prompt = self.model.messages2prompt(messages, sequence_start)
         input_ids = self.tokenizer.encode(prompt)
         finish_reason = 'stop' if stop else None
-        if self.steps[str(session_id)] + len(
+        if not sequence_end and self.steps[str(session_id)] + len(
                 input_ids) >= self.tm_model.session_len:
             finish_reason = 'length'
             yield GenOut('', self.steps[str(session_id)], len(input_ids), 0,
                          finish_reason)
         else:
-            generator = await self.get_generator(instance_id)
-            self.available[instance_id] = False
-            response_size = 0
-            async for outputs in generator.async_stream_infer(
-                    session_id=session_id,
-                    input_ids=[input_ids],
-                    stream_output=stream_response,
-                    request_output_len=request_output_len,
-                    sequence_start=(sequence_start),
-                    sequence_end=sequence_end,
-                    step=self.steps[str(session_id)],
-                    stop=stop,
-                    top_k=top_k,
-                    top_p=top_p,
-                    temperature=temperature,
-                    repetition_penalty=repetition_penalty,
-                    ignore_eos=ignore_eos,
-                    random_seed=seed if sequence_start else None):
-                res, tokens = outputs[0]
-                # decode res
-                response = self.tokenizer.decode(res[response_size:])
-                # response, history token len, input token len, gen token len
-                yield GenOut(response, self.steps[str(session_id)],
-                             len(input_ids), tokens, finish_reason)
-                response_size = tokens
+            generator = await self.get_generator(instance_id, stop)
+            with self.safe_run(instance_id):
+                response_size = 0
+                async for outputs in generator.async_stream_infer(
+                        session_id=session_id,
+                        input_ids=[input_ids],
+                        stream_output=stream_response,
+                        request_output_len=request_output_len,
+                        sequence_start=(sequence_start),
+                        sequence_end=sequence_end,
+                        step=self.steps[str(session_id)],
+                        stop=stop,
+                        top_k=top_k,
+                        top_p=top_p,
+                        temperature=temperature,
+                        repetition_penalty=repetition_penalty,
+                        ignore_eos=ignore_eos,
+                        random_seed=seed if sequence_start else None):
+                    res, tokens = outputs[0]
+                    # decode res
+                    response = self.tokenizer.decode(res)[response_size:]
+                    # response, history token len,
+                    # input token len, gen token len
+                    yield GenOut(response, self.steps[str(session_id)],
+                                 len(input_ids), tokens, finish_reason)
+                    response_size += len(response)
 
-            # update step
-            self.steps[str(session_id)] += len(input_ids) + tokens
-            if sequence_end:
-                self.steps[str(session_id)] = 0
-            self.available[instance_id] = True
+                # update step
+                self.steps[str(session_id)] += len(input_ids) + tokens
+                if sequence_end or stop:
+                    self.steps[str(session_id)] = 0
 
     async def generate_openai(
         self,
@@ -180,13 +188,11 @@ class AsyncEngine:
         sequence_start = False
         generator = await self.get_generator(instance_id)
         self.available[instance_id] = False
-        if renew_session and str(session_id) in self.steps and self.steps[str(
-                session_id)] > 0:  # renew a session
-            empty_prompt = self.model.messages2prompt('', False)
-            empty_input_ids = self.tokenizer.encode(empty_prompt)
+        if renew_session:  # renew a session
+            empty_input_ids = self.tokenizer.encode('')
             for outputs in generator.stream_infer(session_id=session_id,
                                                   input_ids=[empty_input_ids],
-                                                  request_output_len=1,
+                                                  request_output_len=0,
                                                   sequence_start=False,
                                                   sequence_end=True):
                 pass

lmdeploy/serve/openai/api_client.py

@@ -6,6 +6,15 @@ import fire
 import requests
 
 
+def get_model_list(api_url: str):
+    response = requests.get(api_url)
+    if hasattr(response, 'text'):
+        model_list = json.loads(response.text)
+        model_list = model_list.pop('data', [])
+        return [item['id'] for item in model_list]
+    return None
+
+
 def get_streaming_response(prompt: str,
                            api_url: str,
                            instance_id: int,
@@ -13,7 +22,8 @@ def get_streaming_response(prompt: str,
                            stream: bool = True,
                            sequence_start: bool = True,
                            sequence_end: bool = True,
-                           ignore_eos: bool = False) -> Iterable[List[str]]:
+                           ignore_eos: bool = False,
+                           stop: bool = False) -> Iterable[List[str]]:
     headers = {'User-Agent': 'Test Client'}
     pload = {
         'prompt': prompt,
@@ -22,7 +32,8 @@ def get_streaming_response(prompt: str,
         'request_output_len': request_output_len,
         'sequence_start': sequence_start,
         'sequence_end': sequence_end,
-        'ignore_eos': ignore_eos
+        'ignore_eos': ignore_eos,
+        'stop': stop
     }
     response = requests.post(api_url,
                              headers=headers,
@@ -33,9 +44,9 @@ def get_streaming_response(prompt: str,
                                      delimiter=b'\0'):
         if chunk:
             data = json.loads(chunk.decode('utf-8'))
-            output = data['text']
-            tokens = data['tokens']
-            finish_reason = data['finish_reason']
+            output = data.pop('text', '')
+            tokens = data.pop('tokens', 0)
+            finish_reason = data.pop('finish_reason', None)
             yield output, tokens, finish_reason
 
 
@@ -46,7 +57,7 @@ def input_prompt():
     return '\n'.join(iter(input, sentinel))
 
 
-def main(server_name: str, server_port: int, session_id: int = 0):
+def main(restful_api_url: str, session_id: int = 0):
     nth_round = 1
     while True:
         prompt = input_prompt()
@@ -55,7 +66,7 @@ def main(server_name: str, server_port: int, session_id: int = 0):
         else:
             for output, tokens, finish_reason in get_streaming_response(
                     prompt,
-                    f'http://{server_name}:{server_port}/generate',
+                    f'{restful_api_url}/generate',
                     instance_id=session_id,
                     request_output_len=512,
                     sequence_start=(nth_round == 1),

lmdeploy/serve/openai/api_server.py

@@ -10,7 +10,7 @@ import uvicorn
 from fastapi import BackgroundTasks, FastAPI, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 
-from lmdeploy.serve.openai.async_engine import AsyncEngine
+from lmdeploy.serve.async_engine import AsyncEngine
 from lmdeploy.serve.openai.protocol import (  # noqa: E501
     ChatCompletionRequest, ChatCompletionResponse,
     ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
@@ -27,7 +27,7 @@ class VariableInterface:
     request_hosts = []
 
 
-app = FastAPI()
+app = FastAPI(docs_url='/')
 
 
 def get_model_list():
@@ -253,11 +253,12 @@ async def generate(request: GenerateRequest, raw_request: Request = None):
 
     The request should be a JSON object with the following fields:
     - prompt: the prompt to use for the generation.
-    - stream: whether to stream the results or not.
-    - sequence_start (bool): indicator for starting a sequence.
-    - sequence_end (bool): indicator for ending a sequence
     - instance_id: determine which instance will be called. If not specified
         with a value other than -1, using host ip directly.
+    - sequence_start (bool): indicator for starting a sequence.
+    - sequence_end (bool): indicator for ending a sequence
+    - stream: whether to stream the results or not.
+    - stop: whether to stop the session response or not.
     - request_output_len (int): output token nums
     - step (int): the offset of the k/v cache
     - top_p (float): If set to float < 1, only the smallest set of most
@@ -283,6 +284,7 @@ async def generate(request: GenerateRequest, raw_request: Request = None):
         request_output_len=request.request_output_len,
         top_p=request.top_p,
         top_k=request.top_k,
+        stop=request.stop,
         temperature=request.temperature,
         repetition_penalty=request.repetition_penalty,
         ignore_eos=request.ignore_eos)

lmdeploy/serve/openai/protocol.py

@@ -189,11 +189,12 @@ class EmbeddingsResponse(BaseModel):
 
 class GenerateRequest(BaseModel):
     """Generate request."""
-    prompt: str
+    prompt: Union[str, List[Dict[str, str]]]
     instance_id: int = -1
     sequence_start: bool = True
     sequence_end: bool = False
     stream: bool = False
+    stop: bool = False
     request_output_len: int = 512
     top_p: float = 0.8
     top_k: int = 40

requirements.txt

@@ -9,6 +9,7 @@ pybind11
 safetensors
 sentencepiece
 setuptools
+shortuuid
 tiktoken
 torch
 transformers