首页 > Python资料博客日记

LLM部署，并发控制，流式响应（Python，Qwen2+FastAPI）

2024-09-19 00:00:07Python资料围观173次

这篇文章介绍了LLM部署，并发控制，流式响应（Python，Qwen2+FastAPI），分享给大家做个参考，收藏Python资料网收获更多编程知识

前言

随着生成式人工智能的快速发展，部分场景希望能过自主部署大型语言模型（LLM）服务器用于推理服务，而相关教程博文尽管很多，但存在孤立零散现象，各功能没有打通实现，使得在工程实践中令人困惑。本文从工程实践的角度，着重从“并发控制”与“流式响应”两方面展开。

FastAPI 是一个简单高效的 Web 框架，可以快速实现大型语言模型（LLM）服务器独立部署，目前很多博客也写了相关教程。但在工程实践中，我们希望LLM服务器可以同时处理多个请求，并实现“当请求达到一定数量后，直接拒绝后续的推理请求”功能，防止服务器过载以及排队时间过长影响用户体验（不如直接告知服务器繁忙），本文在“并发控制”上做了进一步的相关工作。

此外，对于较为复杂的问题（推理复杂、文本过长），希望能够“边推理便输出内容”，即“流式响应”，可以极大提高用户体验（与之相对的我称之为“一次响应”，即生成完后一次性全部返回），并在前述的并发控制框架之下实现。

因此，本文实现了工程实践下并发控制下LLM服务器部署并提供流式响应。

为便于理解，首先给出所需要的库（ LLM 服务器环境）。

# Web
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
from sse_starlette.sse import EventSourceResponse
import uvicorn
# 异步库
import asyncio
# 线程池执行器
from concurrent.futures import ThreadPoolExecutor
from threading import Thread
# LLM
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
from transformers import TextStreamer, TextIteratorStreamer # 后者是可迭代的
import torch

import json
import datetime

客户端以 Python 支撑的 Web 为例，仅使用 requests 库即可。

（本文创作于2024年7月）

一、并发控制

模型的加载

使用 FastAPI 对 LLM 进行封装提供 API 服务。

（为便于演示，使用 Qwen2-0.5B，实际应用可以替换为其他模型）

# 创建FastAPI应用
app = FastAPI()

# 主函数入口
if __name__ == '__main__':
    # 加载预训练的分词器和模型
    now_model_place = "E:\LLM\Qwen2_0.5B\Qwen2-0.5B-Instruct"
    model_name_or_path = now_model_place

    # 加载分词器
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
    # 加载预训练模型，数据类型和自动选择设备
    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map="auto", torch_dtype=torch.bfloat16).eval()

    # 启动FastAPI应用
    # 用6006端口可以将autodl的端口映射到本地，从而在本地使用api
    uvicorn.run(app, host='0.0.0.0', port=6006, workers=1)  # 在指定端口和主机上启动应用

其中需要注意的是，.eval() 是将模型设置为评估模式，在预测或评估时使用，实现禁用特定于训练的操作如 Dropout 防止误修改模型权重。

排队法

对于基于 Python 的 FastAPI，有多种并发与限流方法，比如 asyncio.Semaphore（信号量）、令牌桶/漏桶算法等。在这里，本文采用排队法（基于 asyncio.Queue）来跟踪正在处理与等待的请求，最终实现同时处理多个请求，并实现“当请求达到一定数量后，直接拒绝后续的推理请求”的功能。

采用 async-await 的异步形式。

首先创建队列，队列大小为10，即允许最大并发量为10。然后监听 post 请求，解析给定的 prompt 输入内容。如果当前队里已满，直接返回服务器繁忙提示；队列未满则加入队列，执行推理，结束后释放相应的队列占用，伪代码如下：

# 创建队列来跟踪请求数量（并发队列，数量为10）
asyncio_queue = asyncio.Queue(maxsize = 10)

# 处理POST请求的端点
@app.post("/")
async def create_item(request: Request):

    # 获取 prompt （LLM的输入内容）
    json_post_raw = await request.json()  # 获取POST请求的JSON数据
    json_post = json.dumps(json_post_raw)  # 将JSON数据转换为字符串
    json_post_list = json.loads(json_post)  # 将字符串转换为Python对象
    prompt = json_post_list.get('prompt')  # 获取请求中的提示

    # 获取了令牌许可：加入了请求队列（并发队列）
    # 队列已满，则返回 503HOT 的“繁忙”提示
    if asyncio_queue.full():
        now = datetime.datetime.now()  # 获取当前时间
        time = now.strftime("%Y-%m-%d %H:%M:%S")  # 格式化时间为字符串
        # 连接成功，但是LLM服务器繁忙
        return {"response": "503HOT！LLM服务器推理繁忙，请稍后再试",
                "status": 200,
                "time": time}
        
    # 将请求加入队列
    await asyncio_queue.put(None)
    try:
        # 进行LLM推理
        return answer  # 返回响应
    finally:
        # 任务完成，从队列中移除请求
        await asyncio_queue.get()
        asyncio_queue.task_done()

需要注意的是，本文所列方法均为单个工作进程，即在 Uvicorn 中 workers=1。若依托多核采用多个工作进程，需要使用外部的队列等缓存机制来实现分布式下的状态共享。

异步与同步

（重点）

若在上述“ LLM 推理”直接调用传统的推理形式，会发现在实际使用中，请求推理会逐个执行，同时后续请求会一直等待，并没有返回“LLM服务器繁忙”的提示。

这是由于，“ LLM 推理”属于同步任务，其基于深度学习，涉及到了不支持异步的库（如 PyTorch），造成阻塞。Python 的异步特性是基于事件循环的，而在单个线程中，当执行阻塞操作时，事件循环会被阻塞，无法执行其他任务。因此，即使使用了异步编程技术，如果底层的模型推理库不支持异步操作，造成在异步函数中调用了同步函数，那么在执行模型推理时仍然可能会出现阻塞。

但实际上，LLM 推理涉及过多，无法重新写成完全异步的形式。查阅 AI，给出了如下解决方案：

使用支持异步操作的库进行 LLM 推理
使用批处理，把多个推理请求合并成一个，一次性执行
使用分布式系统，多个服务器分散负载
使用 asyncio.run_in_executor 将同步任务放入线程池或进程池

显然，在当前情况下，且资源受限，采用方案4为最佳选择，同时可以提高单机的并发能力，即解决方案：将同步的模型推理代码放入线程池中执行，从而避免阻塞事件循环，集成到异步框架下。

    try:
        # 创建一个线程执行器
        loop = asyncio.get_event_loop()
        with ThreadPoolExecutor() as executor:
            # 使用执行器执行同步的函数
            # LLM_carry为LLM推理函数，prompt为输入内容
            answer = await loop.run_in_executor(executor, LLM_carry, prompt)
        #answer = LLM_carry(prompt)
        return answer  # 返回响应
    finally:
        # 任务完成，从队列中移除请求
        await asyncio_queue.get()
        asyncio_queue.task_done()

在这里，使用了 concurrent.futures.ThreadPoolExecutor 来提供线程池执行器，线程池执行器允许你将可调用的对象（通常是函数）提交到线程池中执行，可以重用线程，这减少了线程创建和销毁的开销，并且在处理大量短生命周期的任务时能够提高性能。

在实际使用中，发现当使用 LLM（Qwen2-7B）运行时，GPU 占用首先直接达到 16GB 左右保持稳定（无推理任务）。尔后随并发下推理任务 GPU 占用逐渐小幅度增加，推理结束后存在回落（即：GPU 内存加载时大幅度占用，并发推理时小幅度增加与回落）。并且可以明显发现，LLM 推理确实是并行执行（现象：A\B\C 顺序请求，一段时间后 B 先返回响应，紧接着 A\B 返回响应）。推测原因：

（1）大模型在初始加载时占用大量显存。因为模型参数和必要的数据结构需要被加载到 GPU 内存中。这些数据包括模型的权重、梯度、优化器状态等。一旦模型被加载，这些显存就会一直被占用，直到模型被卸载或者 Python 进程结束。

（2）后续的并发推理中，显存占用的小幅度增加。因为在模型推理过程中，可能会生成一些临时变量，这些变量用于存储中间计算结果等。

上述技术方法与实验现象，为单机并发推理提供了理论与实践的支撑。

二、流式响应

当前，最容易实现的，是 LLM 服务器进行推理后，将生成的结果一次性返回到客户端（“一次响应”）。但当上下文过长、问题复杂时，推理时间会很长，导致用户使用体验感不佳（较长时间盯着加载图标转圈）。因此，第二部分着重解决“边推理便输出内容”（“流式响应”）。

LLM的流式输出

首先给出“一次响应”下的代码（改编自 Qwen2 示例代码）

# 大模型推理（这应该是一个同步的阻塞任务，导致该线程只进行此任务，无法处理其他的async异步）
def LLM_carry(prompt):
    global model, tokenizer  # 声明全局变量以便在函数内部使用模型和分词器
    messages = [
        {"role": "system", "content": "你是一个热情、积极向上而客观严谨的接待员，为客人提供问答服务。"},
        #{"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": prompt}
    ]

    torch_gc()  # 执行GPU内存清理

    # 使用分词器的apply_chat_template方法来格式化消息
    input_ids = tokenizer.apply_chat_template(
        messages,  # 要格式化的消息
        tokenize=False, # 不进行分词
        add_generation_prompt=True # 添加生成提示
        )
    
    # 将格式化后的文本转换为模型输入，并转换为PyTorch张量，然后移动到指定的设备（cuda）
    model_inputs = tokenizer([input_ids], return_tensors="pt").to('cuda')

    # 使用model.generate()方法直接生成文本（一次响应）
    # generated_ids 是由数字构成的 tensor
    generated_ids = model.generate(
        model_inputs.input_ids,  # 模型输入的input_ids
        max_new_tokens=512 # 最大新生成的token数量
    )
    # 从生成的ID中提取新生成的ID部分
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    # 使用分词器的batch_decode方法将生成的ID解码回文本，并跳过特殊token
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]   

    now = datetime.datetime.now()  # 获取当前时间
    time = now.strftime("%Y-%m-%d %H:%M:%S")  # 格式化时间为字符串
    # 构建响应JSON
    answer = {
        "response": response,
        "status": 200,
        "time": time
    }
    # 构建日志信息
    log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
    print(log)  # 打印日志
    torch_gc()  # 执行GPU内存清理
    return answer  # 返回响应

可以看出，演示代码主要是使用了 model.generate() 方法来直接进行文本生成。

对于 Qwen2 的流式输出，存在 swift.llm 的解决方案【1】（同为阿里在 LLM 领域布局），但截止2024年7月初，本人参考 Qwen-7B-Chat 教程对 Qwen2 进行流式输出尚未成功（确实本人水平有限）。因此只能采用其他方法。

目前（2024年7月），LLM 领域虽然百花齐放，但大多数投入开发的仍属于 Transformer 结构之下，因此，也离不开 transformers 这一开源库。经资料搜集，transformers 确实提供了模型推理的流式输出【2】，并且确实可以用于Qwen2【3】，因此采用 transformers 来实现Qwen2的“流式响应”。

（重点）

transformers 库有两个流式输出函数，TextStreamer 与 TextIterateStreamer，前者属于终端进行输出，后者是返回一个可迭代对象，更加具有自定义性，因此我们使用 TextIterateStreamer。

# 大模型推理（这应该是一个同步的阻塞任务，导致该线程只进行此任务，无法处理其他的async异步）
def LLM_carry(prompt):
    global model, tokenizer  # 声明全局变量以便在函数内部使用模型和分词器
    messages = [
        {"role": "system", "content": "你是一个热情、积极向上而客观严谨的接待员，为客人提供问答服务。"},
        #{"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": prompt}
    ]

    torch_gc()  # 执行GPU内存清理

    # 使用分词器的apply_chat_template方法来格式化消息
    input_ids = tokenizer.apply_chat_template(
        messages,  # 要格式化的消息
        tokenize=False, # 不进行分词
        add_generation_prompt=True # 添加生成提示
        )
    
    # 将格式化后的文本转换为模型输入，并转换为PyTorch张量，然后移动到指定的设备（cuda）
    model_inputs = tokenizer([input_ids], return_tensors="pt").to('cuda')

    # 使用流式传输模式（更加流畅和动态的交互体验）
    # 自定义、可迭代的流式输出
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    #另一种执行流式输出的写法
    #generation_kwargs = dict(model_inputs, streamer=streamer, max_new_tokens=512)
    #thread = Thread(target=model.generate, kwargs=generation_kwargs)
    #thread.start()
    
    # 执行流式输出
    model.generate(model_inputs.input_ids, max_new_tokens=512, streamer=streamer)

    # FastAPI流式响应
    async def stream_response():
        for text in streamer:
            if text:
                print(text)

此时，在服务器的终端上就可以流式（逐 token）输出生成的文本了。

FastAPI传输

生成的文本不仅需要流式输出，还需要通过 FastAPI 来进行流式传输到客户端，即服务器实时响应，经收集整理可分为以下几种方案【4】【5】【6】【7】：

HTTP：传统的“请求—响应”模型，客户端主动发请求，服务端被动地返回响应；可采用长轮询形式（需要维护大量长连接）。
WebSocket：全双工通信，实现实时双向通信；太“重”了。
流式传输（基于 HTTP /1.1的分块传输）：FastAPI 中的 StreamingResponse ，通常用于流式传输原始数据，如文件流或视频流，一般要求要传输的数据是确定的，适用于需要流式传输大量数据或长时间数据的场景，更有通用性。
SSE（Server-sent Events）：允许服务器主动向客户端发送数据的技术，使用标准的 HTTP 协议，仅提供服务端到客户端的单向通信。FastAPI中的 EventSourceResponse（需要安装 sse_starlette 进行拓展），提供了更适合 LLM 推理场景的特性，如事件驱动的设计和更好的实时性，同时更方便传输复杂的数据结构。（重点）

显然，选择3、4方案可以有效解决流式传输问题，并且在实际使用中都可以实现需要的效果。考虑当前主流方向，采用 SSE 的主动推送方案（代码注释部分中也包括实现了的方案3）

    # FastAPI流式响应（StreamingResponse 与 EventSourceResponse）
    async def stream_response():
        for text in streamer:
            if text:
                #print(repr(text))
                # 用yield迭代推送对话内容
                # StreamingResponse下的返回信息
                #yield text
                # EventSourceResponse下的返回信息
                yield {
                    "event": "message",
                    "retry": 15000,
                    "data": repr(text)  # 防止\n换行符等传输过程中丢失
                }

    #return StreamingResponse(stream_response(), media_type="application/octet-stream")
    return EventSourceResponse(stream_response())

需要说明的是，在采用 EventSourceResponse 直接返回 text 时，会出现换行符丢失的情况，产生的根本原因不明，这在以md语法输出的LLM推理中是致命问题。

SSE响应内容格式定义：event: {event}\r\ndata: {data}\r\nretry: {retry}\r\n\r\n

使用 repr() 显示字符形式来进行检查，发现LLM推理流式输出的token，在传输过程中被拆分了（原因不明），导致客户端解析错误。

因此，通过返回字符形式的 repr(text) 使得能够在客户端正确解析。

客户端获取

来自服务器的流式响应需要使用合适的方式进行接收，本文以基于 Python 的 Web 为例，使用 requests 库的 post 请求，分别实现接收来自“一次响应”、StreamingResponse 的流式传输、EventSourceResponse 的SSE的数据。

# 直接返回
def get_completion(prompt):
    headers = {'Content-Type': 'application/json'}
    data = {"prompt": prompt}
    response = requests.post(url='http://127.0.0.1:6006', headers=headers, data=json.dumps(data))
    return response.json()['response']


# 流式响应
def get_completion_stream(prompt):
    headers = {'Content-Type': 'application/octet-stream'}
    data = {"prompt": prompt}
    
    # LLM服务器采用 StreamingResponse
    try:
        response = requests.post(url='http://127.0.0.1:6006', headers=headers, data=json.dumps(data), stream=True)
        if response.status_code == 200:
            # 列表，用于拼接流式返回的生成文本
            all_chunk_response = []
            # chunk_size: 默认为1，正常情况下要设置一个比较大的值，否则获取到一个字节数据就会走到下面的处理逻辑
            # #decode_unicode: iter_content() 函数遍历的数据是bytes类型的，这个参数可以控制是否将bytes转为str
            for chunk in response.iter_content(chunk_size=512, decode_unicode=True):
                #print(chunk)
                # StreamingResponse下的返回信息处理
                #all_chunk_response.append(chunk)
                # EventSourceResponse下的返回信息处理
                # 解析SSE响应内容格式，分割出所需数据
                chunk_data = chunk.split("\r\ndata: ")[1].split("\r\nretry: ")[0]
                
                yield chunk_data

                all_chunk_response.append(chunk_data)
                all_chunk_response_text = ''.join(all_chunk_response)
                # 打印日志
                now = datetime.datetime.now()  # 获取当前时间
                time = now.strftime("%Y-%m-%d %H:%M:%S.%f")  # 格式化时间为字符串（微秒）
                log = "[" + time + "] " + all_chunk_response_text
                #print(log)  
        else:
            print(response)
    except requests.RequestException as e:
        print(f"Request failed: {e}")
    

if __name__ == '__main__':
    #print(get_completion('你好'))
    for data in get_completion_stream('什么是大模型技术？有什么影响？'):
        print(data, end="")  # 流式显示数据

在实际使用过程中，将客户端示例代码的主函数部分使用 yield 包装成迭代器调用即可。

三、代码实现

该部分给出本文实现的完整代码。

LLM服务器（服务端）

待进一步整理后完整开源。

Web服务器（客户端）

同二、客户端获取部分代码。

四、小结

本文实现了工程实践下并发控制下LLM服务器部署并提供流式响应。使得LLM服务器可以同时处理多个请求，并实现“当请求达到一定数量后，直接拒绝后续的推理请求”功能，防止服务器过载以及排队时间过长影响用户体验（直接告知服务器繁忙）。此外，对于较为复杂的问题（推理复杂、文本过长）生成回答，在并发架构下实现了“边推理便输出内容”的“流式响应”，可有效提高用户体验，优化人机交互实现。