Add deployment tuning and chat load test
This commit is contained in:
@@ -5,6 +5,10 @@ API_PREFIX=/api
|
||||
|
||||
DATABASE_URL=mysql+pymysql://ai_kb:ai_kb@127.0.0.1:3306/ai_knowledge_base_v2?charset=utf8mb4
|
||||
REDIS_URL=
|
||||
DB_POOL_SIZE=10
|
||||
DB_MAX_OVERFLOW=20
|
||||
DB_POOL_TIMEOUT=30
|
||||
DB_POOL_RECYCLE=1800
|
||||
AUTO_CREATE_TABLES=false
|
||||
|
||||
JWT_SECRET_KEY=change-me-in-production
|
||||
|
||||
@@ -12,4 +12,4 @@ COPY . .
|
||||
|
||||
EXPOSE 8100
|
||||
|
||||
CMD ["sh", "-c", "alembic upgrade head && uvicorn app.main:app --host 0.0.0.0 --port 8100"]
|
||||
CMD ["sh", "-c", "alembic upgrade head && uvicorn app.main:app --host 0.0.0.0 --port 8100 --workers ${BACKEND_WORKERS:-1}"]
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
@@ -31,6 +32,7 @@ from app.services.chat_queue_service import load_chat_queue_config
|
||||
from app.services.chat_stream_service import ChatStreamService
|
||||
|
||||
router = APIRouter()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@router.post("/session")
|
||||
@@ -155,6 +157,7 @@ async def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("Chat stream failed")
|
||||
yield _sse_event("error", message="AI 回复生成失败,请稍后再试。")
|
||||
finally:
|
||||
if acquired:
|
||||
|
||||
@@ -26,6 +26,10 @@ class Settings(BaseSettings):
|
||||
|
||||
database_url: str = "mysql+pymysql://ai_kb:ai_kb@127.0.0.1:3306/ai_knowledge_base_v2?charset=utf8mb4"
|
||||
redis_url: str = ""
|
||||
db_pool_size: int = 10
|
||||
db_max_overflow: int = 20
|
||||
db_pool_timeout: int = 30
|
||||
db_pool_recycle: int = 1800
|
||||
auto_create_tables: bool = False
|
||||
|
||||
jwt_secret_key: str = "change-me-in-production"
|
||||
|
||||
@@ -10,7 +10,15 @@ from app.models import Base
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
engine = create_engine(settings.database_url, pool_pre_ping=True, future=True)
|
||||
engine = create_engine(
|
||||
settings.database_url,
|
||||
pool_pre_ping=True,
|
||||
pool_size=settings.db_pool_size,
|
||||
max_overflow=settings.db_max_overflow,
|
||||
pool_timeout=settings.db_pool_timeout,
|
||||
pool_recycle=settings.db_pool_recycle,
|
||||
future=True,
|
||||
)
|
||||
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True)
|
||||
|
||||
|
||||
|
||||
@@ -474,3 +474,41 @@
|
||||
1. dev 环境默认仍使用 1 个 worker,避免本地调试时多 worker 日志和热更新混乱。
|
||||
2. 生产可以通过环境变量调高 worker,不需要改代码。
|
||||
3. 压测脚本只作为研发/验收工具,不进入应用启动链路。
|
||||
|
||||
## 阶段 5 实施记录
|
||||
|
||||
2026-07-08 完成 worker、连接池和压测工具第一版。
|
||||
|
||||
已完成:
|
||||
|
||||
1. 后端启动支持 `BACKEND_WORKERS` 环境变量。
|
||||
2. Dockerfile 启动命令支持 `uvicorn --workers ${BACKEND_WORKERS:-1}`。
|
||||
3. SQLAlchemy 连接池支持环境变量配置:
|
||||
- `DB_POOL_SIZE`
|
||||
- `DB_MAX_OVERFLOW`
|
||||
- `DB_POOL_TIMEOUT`
|
||||
- `DB_POOL_RECYCLE`
|
||||
4. `.env.example` 和 dev compose 已补齐对应配置。
|
||||
5. 新增 `scripts/chat_load_test.py`,支持并发请求 `/chat/completions` 并输出:
|
||||
- 单请求结果。
|
||||
- `sessionId`。
|
||||
- 首事件耗时。
|
||||
- 总耗时。
|
||||
- 成功数和失败数汇总。
|
||||
6. `/chat/completions` 增加后端异常日志,用户端仍保持通用错误提示,便于压测时查堆栈。
|
||||
|
||||
验证记录:
|
||||
|
||||
1. 本机 `compileall` 检查通过。
|
||||
2. 压测脚本 `--help` 检查通过。
|
||||
3. `git diff --check` 通过。
|
||||
4. 后端镜像重建通过。
|
||||
5. 后端健康检查通过。
|
||||
6. 单请求压测通过:1 个请求成功,首事件约 20ms,总耗时约 11.66s。
|
||||
7. 2 并发、2 请求压测暴露当前模型链路风险:两次复测均出现 1 成功、1 失败,失败请求在约 60s 返回“AI 回复生成失败”。该现象更像当前模型配置或上游模型并发/超时限制,需要后续结合后端异常日志和模型管理配置继续调参。
|
||||
|
||||
阶段限制:
|
||||
|
||||
1. 当前没有做大并发压测;阶段 5 只补齐调参能力和轻量验证脚本。
|
||||
2. 生产 worker 数不能只看 CPU,需要结合模型接口限流、飞书频控、DB 连接数和 Redis 队列等待时间一起定。
|
||||
3. 当前建议先把 `chat_max_active_requests` 调低到当前模型真实可承受范围,再逐步放大压测。
|
||||
|
||||
@@ -44,6 +44,11 @@ services:
|
||||
DEBUG: "true"
|
||||
DATABASE_URL: mysql+pymysql://ai_kb:ai_kb@mysql:3306/ai_knowledge_base_v2?charset=utf8mb4
|
||||
REDIS_URL: redis://redis:6379/0
|
||||
BACKEND_WORKERS: "1"
|
||||
DB_POOL_SIZE: "10"
|
||||
DB_MAX_OVERFLOW: "20"
|
||||
DB_POOL_TIMEOUT: "30"
|
||||
DB_POOL_RECYCLE: "1800"
|
||||
JWT_SECRET_KEY: local-dev-secret-change-before-production
|
||||
MOCK_SMS_ENABLED: "true"
|
||||
MOCK_SMS_CODE: "123456"
|
||||
|
||||
145
ai_knowledge_base_v2/scripts/chat_load_test.py
Normal file
145
ai_knowledge_base_v2/scripts/chat_load_test.py
Normal file
@@ -0,0 +1,145 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import statistics
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Run a lightweight chat SSE load test.")
|
||||
parser.add_argument("--base-url", default="http://127.0.0.1:8100/api")
|
||||
parser.add_argument("--phone", required=True)
|
||||
parser.add_argument("--code", default="123456")
|
||||
parser.add_argument("--message", default="请用一句话回复:压测验证")
|
||||
parser.add_argument("--concurrency", type=int, default=5)
|
||||
parser.add_argument("--requests", type=int, default=20)
|
||||
parser.add_argument("--timeout", type=int, default=90)
|
||||
args = parser.parse_args()
|
||||
|
||||
token = login(args.base_url, args.phone, args.code, args.timeout)
|
||||
lock = threading.Lock()
|
||||
results: list[dict] = []
|
||||
|
||||
started_at = time.perf_counter()
|
||||
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
|
||||
futures = [
|
||||
executor.submit(run_one, args.base_url, token, args.message, args.timeout)
|
||||
for _ in range(args.requests)
|
||||
]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
with lock:
|
||||
results.append(result)
|
||||
print(json.dumps(result, ensure_ascii=False))
|
||||
|
||||
elapsed = time.perf_counter() - started_at
|
||||
print_summary(results, elapsed)
|
||||
|
||||
|
||||
def login(base_url: str, phone: str, code: str, timeout: int) -> str:
|
||||
body = post_json(f"{base_url}/auth/login", {"phone": phone, "code": code}, timeout=timeout)
|
||||
return body["data"]["token"]
|
||||
|
||||
|
||||
def run_one(base_url: str, token: str, message: str, timeout: int) -> dict:
|
||||
started_at = time.perf_counter()
|
||||
try:
|
||||
session_body = post_json(
|
||||
f"{base_url}/chat/session",
|
||||
{"title": "压测会话"},
|
||||
token=token,
|
||||
timeout=timeout,
|
||||
)
|
||||
session_id = session_body["data"]["sessionId"]
|
||||
first_event_ms = None
|
||||
content_events = 0
|
||||
error_message = ""
|
||||
|
||||
request = urllib.request.Request(
|
||||
f"{base_url}/chat/completions",
|
||||
data=json.dumps({"sessionId": session_id, "message": message}).encode("utf-8"),
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {token}",
|
||||
},
|
||||
)
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
for raw_line in response:
|
||||
line = raw_line.decode("utf-8").strip()
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
if first_event_ms is None:
|
||||
first_event_ms = elapsed_ms(started_at)
|
||||
if line == "data: [DONE]":
|
||||
break
|
||||
payload = json.loads(line[5:].strip())
|
||||
if payload.get("type") == "content":
|
||||
content_events += 1
|
||||
if payload.get("type") == "error":
|
||||
error_message = payload.get("message", "unknown error")
|
||||
|
||||
return {
|
||||
"ok": not error_message and content_events > 0,
|
||||
"sessionId": session_id,
|
||||
"firstEventMs": first_event_ms,
|
||||
"totalMs": elapsed_ms(started_at),
|
||||
"contentEvents": content_events,
|
||||
"error": error_message,
|
||||
}
|
||||
except (urllib.error.URLError, TimeoutError, KeyError, ValueError) as exc:
|
||||
return {
|
||||
"ok": False,
|
||||
"sessionId": None,
|
||||
"firstEventMs": None,
|
||||
"totalMs": elapsed_ms(started_at),
|
||||
"contentEvents": 0,
|
||||
"error": str(exc),
|
||||
}
|
||||
|
||||
|
||||
def post_json(url: str, payload: dict, *, token: str | None = None, timeout: int) -> dict:
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if token:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
request = urllib.request.Request(url, data=json.dumps(payload).encode("utf-8"), headers=headers)
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
return json.loads(response.read().decode("utf-8"))
|
||||
|
||||
|
||||
def print_summary(results: list[dict], elapsed: float) -> None:
|
||||
ok_results = [result for result in results if result["ok"]]
|
||||
failed_results = [result for result in results if not result["ok"]]
|
||||
first_event_values = [result["firstEventMs"] for result in ok_results if result["firstEventMs"] is not None]
|
||||
total_values = [result["totalMs"] for result in ok_results]
|
||||
summary = {
|
||||
"totalRequests": len(results),
|
||||
"success": len(ok_results),
|
||||
"failed": len(failed_results),
|
||||
"elapsedSeconds": round(elapsed, 2),
|
||||
"firstEventMsAvg": round(statistics.mean(first_event_values), 2) if first_event_values else None,
|
||||
"firstEventMsP95": percentile(first_event_values, 95),
|
||||
"totalMsAvg": round(statistics.mean(total_values), 2) if total_values else None,
|
||||
"totalMsP95": percentile(total_values, 95),
|
||||
}
|
||||
print("SUMMARY " + json.dumps(summary, ensure_ascii=False))
|
||||
|
||||
|
||||
def percentile(values: list[int], percentile_value: int) -> int | None:
|
||||
if not values:
|
||||
return None
|
||||
sorted_values = sorted(values)
|
||||
index = min(len(sorted_values) - 1, round((percentile_value / 100) * (len(sorted_values) - 1)))
|
||||
return sorted_values[index]
|
||||
|
||||
|
||||
def elapsed_ms(started_at: float) -> int:
|
||||
return int((time.perf_counter() - started_at) * 1000)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user