使用flask/fastapi部署模型（工业级）

前言

LLM部署，首先介绍fastapi和flask，两者都是用于构建Web应用程序的框架。再介绍工业级部署策略。简单调试可以用第二小节内容，真正部署可以用第三节内容。

使用fastapi或者flask

笔者更喜欢flask框架，觉得他更接地气，且易于使用。

使用fastapi实现

使用了客户端和服务端的分离，使得你可以在需要的时候发送推理请求。

# start_server.py
from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModel
import uvicorn, json, datetime
import torch

from tqdm import tqdm
from transformers import LlamaTokenizer, AutoTokenizer, AutoModelForCausalLM, AutoConfig
from peft import  PeftModel

DEVICE = "cuda"
DEVICE_ID = "0"
CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE
print(CUDA_DEVICE)

def torch_gc():
    if torch.cuda.is_available():
        with torch.cuda.device(CUDA_DEVICE):
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

app = FastAPI()

@app.post("/")
async def create_item(request: Request):
    global model, tokenizer
    json_post_raw = await request.json()
    json_post = json.dumps(json_post_raw)
    json_post_list = json.loads(json_post)
    prompt = json_post_list.get('prompt')  # input_text

	##### 以下为模型推理代码，替换成自己的代码  #####
    inputs = tokenizer(
        prompt,
        add_special_tokens=False,
        return_tensors="pt"
    )
    generation_output = model.generate(
        input_ids = inputs["input_ids"].to(CUDA_DEVICE), 
        **generation_config
    )  # [0]

    for i, line in enumerate(generation_output):
        generate_text = tokenizer.decode(line,skip_special_tokens=True)

    infer_result = generate_text.split('Assistant:\n')[-1].strip() + '\n'
	##### 以上为模型推理代码，替换成自己的代码 #####

    answer = {
        "response": infer_result,  # 返回的结果
    }
    

    torch_gc()
    return jsonify(answer)
    # return json.dumps({'result': result})
    # 注： 返回结果answer中如果有np.array格式，需要转成list格式，才可以返回，否则会报错TypeError: Object of type ndarray is not JSON serializable


if __name__ == '__main__':
	# 以下加载你的模型
    tokenizer = AutoTokenizer.from_pretrained('/workspace/BELLE/train/dataset/llama2_7b_chat_hf')
    tokenizer.pad_token_id = 0
    tokenizer.bos_token_id = 1
    tokenizer.eos_token_id = 2
    tokenizer.padding_side = "left"
    load_type = torch.float16 #Sometimes may need torch.float32
    model_config = AutoConfig.from_pretrained('/workspace/BELLE/train/dataset/llama2_7b_chat_hf')
    model = AutoModelForCausalLM.from_pretrained('/workspace/BELLE/train/output/saved_models/llama2_7b_chat_hf/checkpoint-67008', torch_dtype=load_type, config=model_config, device_map='auto')
    model.eval()
    # 以上加载你的模型
    uvicorn.run(app, host='0.0.0.0', port=7999, workers=1)

客户端代码如下：

# client.py
import requests

# 定义请求URL
url = "http://0.0.0.0:7999"

# 定义请求头
headers = {
    "Content-Type": "application/json"
}

# 定义请求体数据
data = {
    "prompt": "Human: \nYou are now a proficient people, \n\nAssistant:\n",
    # "history": []
}

# 发送POST请求
response = requests.post(url, headers=headers, json=data)

# 打印响应
print(response.text)

# 在shell输入下面指令，获取ip
hostname -I 

# 在终端输入下面指令调用，上述获取的ip替换下面的localhost，输入-d中数据中的引号需要进行转义，否则会报错。<p>The browser (or proxy) sent a request that this server could not understand.
curl -X POST "http://localhost:7999" -H 'Content-Type: application/json' -d '{\"prompt\": \"你好\", \"history\": []}'

使用flask实现

使用flask，记得安装相应的库pip install Flask requests
将服务启动和推理分开为两个脚本。以下是一个简单的示例，其中一个脚本 start_server.py 负责启动服务，另一个脚本 client.py 负责进行推理：

# start_server.py
from flask import Flask, request, jsonify
from transformers import GPT2LMHeadModel, GPT2Tokenizer

app = Flask(__name__)


# 以下为模型初始化,加载你的模型
model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
# 以上为模型初始化,加载你的模型


@app.route('/', methods=['POST'])
def infer():
    data = request.json
    input_text = data.get('input_text')

    # 处理推理请求
    result = generate_response(input_text)

    return jsonify({'result': result})
    # return json.dumps({'result': result})
    # 注： 返回结果answer中如果有np.array格式，需要转成list格式，才可以返回，否则会报错TypeError: Object of type ndarray is not JSON serializable

def generate_response(input_text):
    # 以下实现模型推理逻辑
    # 例如，使用预训练的 GPT-2 模型进行文本生成
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    output = model.generate(input_ids, max_length=100, num_return_sequences=1)
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text
	# 以上实现模型推理逻辑
	
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

客户端代码如下：

# client.py
import requests

def run_client(input_text):
    url = 'http://localhost:5000'
    data = {'input_text': input_text}
    response = requests.post(url, json=data)
    
    if response.status_code == 200:  # 200 表示请求成功
        result = response.json()
        print(result['result'])
    else:
        print('Error:', response.status_code)

if __name__ == '__main__':
    input_text = "Human: \nYou are now a proficient people, \n\nAssistant:\n"
    run_client(input_text)

# 也可以直接在终端输入下面指令调用
curl -X POST "http://localhost:5000" -H 'Content-Type: application/json' -d '{\"prompt\": \"你好\", \"history\": []}'

工业级改造建议

方案A：用 Gunicorn + Uvicorn + FastAPI 替代 Flask

原理：Gunicorn 做多进程管理，每个进程绑定一块 GPU，互不干扰，能安全高并发。
实现步骤：

1. 用 FastAPI 重写 infer 逻辑（大部分和 Flask 类似）

   # server.py
   from fastapi import FastAPI, Request
   from transformers import T5ForConditionalGeneration, T5Tokenizer
   import torch
   
   app = FastAPI()
   
   # 这里假设每个进程只用一个GPU（外部启动进程时设置CUDA_VISIBLE_DEVICES）
   model_id = "/home/jerry/workspace/code/t5/saved/checkpoint_199500"
   model = T5ForConditionalGeneration.from_pretrained(model_id).cuda()
   tokenizer = T5Tokenizer.from_pretrained(model_id)
   
   @app.post("/")
   async def infer(request: Request):
       data = await request.json()
       input_text = data.get('input_smile')
       beam_size = data.get('beam_size', 1)
   
       # infer 逻辑同之前
       input_ids = tokenizer(input_text, return_tensors="pt", max_length=256, truncation=True).input_ids.cuda()
       outputs = model.generate(
           input_ids=input_ids, return_dict_in_generate=True, output_scores=True,
           do_sample=False, max_length=256, num_beams=20, num_return_sequences=20,
           diversity_penalty=1.5, num_beam_groups=10
       )
       decoded_outputs = tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True)
       decoded_scores = [float(s) for s in outputs.sequences_scores.cpu().numpy()]
       # 你的 normalize、process_final_result逻辑...
       # ...
       return {"topk": decoded_outputs, "score": decoded_scores}

2. 使用tmux启动多个进程，分别绑定不同 GPU
   用 shell 脚本 run_server.py，假设你有4张卡：

   #run_server.py
   
   CUDA_VISIBLE_DEVICES=0 gunicorn server:app -k uvicorn.workers.UvicornWorker -b 0.0.0.0:9007 --workers 1 &
   CUDA_VISIBLE_DEVICES=1 gunicorn server:app -k uvicorn.workers.UvicornWorker -b 0.0.0.0:9008 --workers 1 &
   CUDA_VISIBLE_DEVICES=2 gunicorn server:app -k uvicorn.workers.UvicornWorker -b 0.0.0.0:9009 --workers 1 &
   CUDA_VISIBLE_DEVICES=3 gunicorn server:app -k uvicorn.workers.UvicornWorker -b 0.0.0.0:9010 --workers 1 &

   其中server是第一步中的server.py文件去掉 .py，这里要保持一致。需要在 server.py 所在目录执行 gunicorn，或者确保 server:app 可以被 import 到！也就是直接把run_server.py放在server.py所在目录执行gunicorn就行了。

3. 配置 nginx 轮询端口分流（负载均衡到4个端口）

   upstream t5_backend {
       server 127.0.0.1:9007;
       server 127.0.0.1:9008;
       server 127.0.0.1:9009;
       server 127.0.0.1:9010;
   }
   
   server {
       listen 9006;
       location / {
           proxy_pass http://t5_backend;
           proxy_set_header Host $host;
           proxy_set_header X-Real-IP $remote_addr;
       }
   }

   新建 t5_backend.conf 放在你的nginx 配置目录下，然后 sudo nginx -t, sudo nginx -s reload 重启 nginx 即可。过程中如果有端口被占用，假如是80端口，可以使用下面指令杀掉：sudo kill -9 sudo lsof -i:80 -t。

这样，client 只需要访问 http://服务器IP:9006，nginx 自动转发到不同 GPU 实例，天然支持高并发和多线程！ curl http://127.0.0.1:9006即可，返回 {“detail”:“Method Not Allowed”} 其实是 FastAPI/服务端的正常提示！

方案B：用 Triton Inference Server / vLLM / TensorRT-LLM（推荐大模型/生产环境）

如果你要“卷得更专业”，用工业级推理服务器：

• vLLM: 专为 LLM 高吞吐、并发推理设计，兼容 transformers API，自动管理多卡。
• Triton: NVIDIA 官方推理引擎，可动态批量、热升级、Prometheus 监控。
• 优势：自动 batch 请求、显存隔离、API 并发量超 Flask/FastAPI 百倍，性能极强。
• 迁移：你只需写一个 handler/model wrapper，然后按官方文档启动多 GPU，RESTful API 无缝调用。

后续继续优化中…

client 端并发优化建议

• 线程池/进程池调用，提高本地 CPU 并发
• 不用手动指定 gpu_id，由负载均衡自动分发
• 如果响应慢，client 端可以用 retry 机制自动重试
  eg:

from concurrent.futures import ThreadPoolExecutor

def worker(data):
    return run_client(data)

with ThreadPoolExecutor(max_workers=8) as pool:
    results = list(pool.map(worker, all_data_list))

总结与推荐

• 如果要和 OpenAI、GLM、Llama 一样的高并发体验，必须用专业推理服务+多进程/多卡绑定+负载均衡
• Flask/单进程/线程只是调试用，不能用于生产高并发
• 最优推荐：FastAPI + Gunicorn + 多进程（每进程单卡） + nginx 负载均衡
• 进阶推荐：直接用 vLLM、Triton，完全无缝支持 Huggingface/transformers，简单几行代码即可