From efd03ff9df4429afd4fcd18540dae11a2ca59161 Mon Sep 17 00:00:00 2001 From: Nelson <1475262689@qq.com> Date: Wed, 8 Jul 2026 17:45:18 +0800 Subject: [PATCH] Add deployment tuning and chat load test --- .../apps/backend/.env.example | 4 + ai_knowledge_base_v2/apps/backend/Dockerfile | 2 +- .../apps/backend/app/api/chat.py | 3 + .../apps/backend/app/core/config.py | 4 + .../apps/backend/app/core/database.py | 10 +- ...026-07-08-concurrency-optimization-plan.md | 38 +++++ ai_knowledge_base_v2/docker-compose.dev.yml | 5 + .../scripts/chat_load_test.py | 145 ++++++++++++++++++ 8 files changed, 209 insertions(+), 2 deletions(-) create mode 100644 ai_knowledge_base_v2/scripts/chat_load_test.py diff --git a/ai_knowledge_base_v2/apps/backend/.env.example b/ai_knowledge_base_v2/apps/backend/.env.example index a0fa612..70d2439 100644 --- a/ai_knowledge_base_v2/apps/backend/.env.example +++ b/ai_knowledge_base_v2/apps/backend/.env.example @@ -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 diff --git a/ai_knowledge_base_v2/apps/backend/Dockerfile b/ai_knowledge_base_v2/apps/backend/Dockerfile index 1eb680f..dbe907d 100644 --- a/ai_knowledge_base_v2/apps/backend/Dockerfile +++ b/ai_knowledge_base_v2/apps/backend/Dockerfile @@ -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}"] diff --git a/ai_knowledge_base_v2/apps/backend/app/api/chat.py b/ai_knowledge_base_v2/apps/backend/app/api/chat.py index 8da25f5..0e626ad 100644 --- a/ai_knowledge_base_v2/apps/backend/app/api/chat.py +++ b/ai_knowledge_base_v2/apps/backend/app/api/chat.py @@ -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: diff --git a/ai_knowledge_base_v2/apps/backend/app/core/config.py b/ai_knowledge_base_v2/apps/backend/app/core/config.py index 1d47930..e37c0ae 100644 --- a/ai_knowledge_base_v2/apps/backend/app/core/config.py +++ b/ai_knowledge_base_v2/apps/backend/app/core/config.py @@ -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" diff --git a/ai_knowledge_base_v2/apps/backend/app/core/database.py b/ai_knowledge_base_v2/apps/backend/app/core/database.py index 87384b1..5dc1385 100644 --- a/ai_knowledge_base_v2/apps/backend/app/core/database.py +++ b/ai_knowledge_base_v2/apps/backend/app/core/database.py @@ -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) diff --git a/ai_knowledge_base_v2/development_records/2026-07-08-concurrency-optimization-plan.md b/ai_knowledge_base_v2/development_records/2026-07-08-concurrency-optimization-plan.md index 58ed4ff..a542543 100644 --- a/ai_knowledge_base_v2/development_records/2026-07-08-concurrency-optimization-plan.md +++ b/ai_knowledge_base_v2/development_records/2026-07-08-concurrency-optimization-plan.md @@ -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` 调低到当前模型真实可承受范围，再逐步放大压测。 diff --git a/ai_knowledge_base_v2/docker-compose.dev.yml b/ai_knowledge_base_v2/docker-compose.dev.yml index 10313cb..f31476c 100644 --- a/ai_knowledge_base_v2/docker-compose.dev.yml +++ b/ai_knowledge_base_v2/docker-compose.dev.yml @@ -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" diff --git a/ai_knowledge_base_v2/scripts/chat_load_test.py b/ai_knowledge_base_v2/scripts/chat_load_test.py new file mode 100644 index 0000000..eb09bf2 --- /dev/null +++ b/ai_knowledge_base_v2/scripts/chat_load_test.py @@ -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()