如何设计一个请求限流模块

本地大模型 → RESTful API（问答接口）全流程

目标：把已部署在本地（CPU/GPU）的大模型（如 LLaMA、ChatGLM、GPT‑4‑Lite 等）暴露为可以通过 curl、Postman 或前端直接调用的 POST /query 接口。

技术栈：

• 框架：FastAPI（轻量、async 原生）
• 推理后端：vLLM（支持大模型并行、低延迟）或 transformers + torch（适合小模型）
• 容器化：Docker（可选）
• Web 服务器：uvicorn（或 gunicorn + uvicorn workers）
• 安全/性能：速率限制、JWT、日志、水平扩展

下面给出两套实现方案：
1️⃣ 适合 GPU + 超大模型：vLLM + FastAPI
2️⃣ 适合 CPU / 小模型：transformers + FastAPI

---

1️⃣ 方案 1：GPU + 超大模型（vLLM）

1.1 环境准备

# 基础镜像
python -m venv venv
source venv/bin/activate
pip install --upgrade pip

# 依赖
pip install fastapi uvicorn vllm

vLLM：支持 GPU 多核、batch、beam search，极大降低推理延迟。
若你使用的是 Nvidia GPU，请先确认 torch 与 cuda 已正确安装。

1.2 快速示例代码

# app.py
import json
from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
from vllm import LLM, SamplingParams

app = FastAPI(title="Local LLM REST API")

# 1️⃣ 加载模型（一次性，程序启动时）
model_name_or_path = "path/to/your/llama-7b"  # 本地路径或 huggingface hub
llm = LLM(
    model=model_name_or_path,
    gpu_memory_utilization=0.9,    # 留一点显存给其他进程
    tensor_parallel_size=1,       # 如果多卡可改为 >1
)

# 2️⃣ 定义请求/响应格式
class QueryRequest(BaseModel):
    prompt: str
    max_tokens: int | None = 200
    temperature: float | None = 0.7
    top_p: float | None = 0.95
    stop: list[str] | None = None

class QueryResponse(BaseModel):
    prompt: str
    answer: str

# 3️⃣ 接口实现
@app.post("/query", response_model=QueryResponse)
async def query(req: QueryRequest):
    # 构造 vLLM SamplingParams
    sampling = SamplingParams(
        max_tokens=req.max_tokens,
        temperature=req.temperature,
        top_p=req.top_p,
        stop=req.stop or [],
    )
    # 异步调用（vLLM 负责多请求并发）
    outputs = await llm.generate_async(req.prompt, sampling_params=sampling)
    # vLLM 的返回结构为 generator，取第一个完成的 output
    answer = next(outputs).choices[0].text

    return QueryResponse(prompt=req.prompt, answer=answer)

# 4️⃣ 启动指令：uvicorn app:app --host 0.0.0.0 --port 8000

说明

• generate_async 支持异步并发，可通过 asyncio.gather 进一步聚合多请求。
• 通过 SamplingParams 可以灵活控制生成长度、温度、top‑p 等超参。
• 你也可以把 stop 参数默认设为 ["<|end_of_text|>"] 或模型自带 stop token。

1.3 运行

uvicorn app:app --host 0.0.0.0 --port 8000

1.4 调用示例

curl -X POST http://localhost:8000/query \
-H "Content-Type: application/json" \
-d '{
  "prompt": "写一段关于机器学习的简短介绍",
  "max_tokens": 150,
  "temperature": 0.6
}'

返回：

{
  "prompt": "写一段关于机器学习的简短介绍",
  "answer": "机器学习是一种让计算机从数据中学习并做出预测或决策的技术..."
}

1.5 可选：Docker 镜像

# Dockerfile
FROM nvcr.io/nvidia/pytorch:23.06-py3

WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY app.py .

EXPOSE 8000
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

requirements.txt：

fastapi
uvicorn
vllm

构建 & 运行：

docker build -t local-llm-api .
docker run -d -p 8000:8000 --gpus all local-llm-api

---

2️⃣ 方案 2：CPU 或小模型（transformers + torch）

适用于：Llama-2-7B-Chat 之类的模型，但 CPU 推理时速度慢；适合实验、demo。

2.1 环境

python -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip install fastapi uvicorn transformers torch

若你已安装 accelerate 或 bitsandbytes，可进一步提升性能。

2.2 示例代码

# app.py
import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

app = FastAPI(title="Local LLM REST API (CPU / small model)")

# 加载 tokenizer / model
model_name_or_path = "path/to/your/model"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",  # 自动分配到 GPU 或 CPU
)

# 如果你使用的是 Llama 系列，记得 set pad_token_id
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# 生成器
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=0 if torch.cuda.is_available() else -1,
)

class QueryRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 150
    temperature: float = 0.7
    top_p: float = 0.95
    stop: list[str] | None = None

class QueryResponse(BaseModel):
    prompt: str
    answer: str

@app.post("/query", response_model=QueryResponse)
def query(req: QueryRequest):
    # 调用 pipeline
    outputs = pipe(
        req.prompt,
        max_new_tokens=req.max_new_tokens,
        temperature=req.temperature,
        top_p=req.top_p,
        eos_token_id=tokenizer.eos_token_id if req.stop else None,
    )
    answer = outputs[0]["generated_text"][len(req.prompt):]  # 去掉 prompt
    return QueryResponse(prompt=req.prompt, answer=answer)

2.3 运行

uvicorn app:app --host 0.0.0.0 --port 8000

2.4 调用示例

与方案 1 相同的 curl 调用即可。

---

3️⃣ 性能与可扩展性建议

方向	方案	实现细节
并发	FastAPI + uvicorn	uvicorn --workers 4 或 gunicorn -k uvicorn.workers.UvicornWorker
GPU 负载	vLLM 自动管理 batch	batch_size=8 等参数可写在 LLM 初始化或 SamplingParams
多节点	Docker Swarm / Kubernetes	使用 deployment.yaml + horizontal pod autoscaler
速率限制	slowapi / starlette-limiter	限制 requests/second
安全	JWT + HTTPS	fastapi.security、uvicorn --ssl-keyfile
日志	loguru / structlog	统一日志格式，方便监控
监控	Prometheus + Grafana	导出 /metrics endpoint

---

4️⃣ 典型错误与排查

问题	可能原因	解决方案
显存不足	vLLM GPU 利用率 100%	tensor_parallel_size 小一点；gpu_memory_utilization=0.7；把部分层放到 CPU（device_map）
请求卡顿	CPU 模型慢	升级到 GPU；或换用 vLLM + transformers 的 inference 方式
API 超时	uvicorn worker 未足够	增加 workers，或改用 gunicorn
错误返回	输入 JSON 格式不符	检查 Content-Type: application/json 并使用 pydantic 进行校验
安全漏洞	未做 IP 限制	加入 IP白名单、API Key 认证

---

5️⃣ 小结

1. 选择后端

   • 超大模型 + GPU → vLLM + FastAPI
   • 小模型或 CPU → transformers + FastAPI

2. 一次性加载

   • model = LLM(...) 或 model = AutoModelForCausalLM(...)
   • 放在全局，避免每次请求都重新加载。

3. 接口设计

   • POST /query，JSON 包含 prompt、max_tokens、temperature 等。
   • 返回 prompt 与 answer，便于前端展示。

4. 部署

   • uvicorn app:app --workers 4 --host 0.0.0.0 --port 8000
   • Docker + GPU：--gpus all，若不想手动拉起请使用 docker-compose。

5. 监控与安全

   • 用 Prometheus + Grafana 监控 latency、throughput、GPU 负载。
   • 用 JWT、IP 限制或反向代理做防护。

只要按上述步骤走，你就能把本地大模型变成一个稳定、可扩展、易调用的 RESTful 服务。祝你顺利部署，问答顺畅！ 🚀

---

转载请注明来源，欢迎对文章中的引用来源进行考证，欢迎指出任何有错误或不够清晰的表达。可以在下面评论区评论，也可以邮件至 1056615746@qq.com