From 1ddec29ffb22b3f6b7ae64f584ec7ad90b923ad3 Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Tue, 16 Dec 2025 17:38:15 +0800 Subject: [PATCH] feat: optimize concurrency limit --- Dockerfile | 42 +++--- app/core/config.py | 6 +- app/services/grok/client.py | 16 ++- app/services/grok/token.py | 117 ++++++++++++--- docker-entrypoint.sh | 53 +++++++ main.py | 45 +++++- requirements.txt | 2 + test_concurrency.py | 276 ++++++++++++++++++++++++++++++++++++ test_concurrency.sh | 177 +++++++++++++++++++++++ 9 files changed, 688 insertions(+), 46 deletions(-) create mode 100644 docker-entrypoint.sh create mode 100755 test_concurrency.py create mode 100755 test_concurrency.sh diff --git a/Dockerfile b/Dockerfile index 2a3d598..f625296 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,13 +6,13 @@ WORKDIR /build # 安装依赖到独立目录 COPY requirements.txt . RUN pip install --no-cache-dir --only-binary=:all: --prefix=/install -r requirements.txt && \ - find /install -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true && \ - find /install -type d -name "tests" -exec rm -rf {} + 2>/dev/null || true && \ - find /install -type d -name "test" -exec rm -rf {} + 2>/dev/null || true && \ - find /install -type d -name "*.dist-info" -exec sh -c 'rm -f "$1"/RECORD "$1"/INSTALLER' _ {} \; && \ - find /install -type f -name "*.pyc" -delete && \ - find /install -type f -name "*.pyo" -delete && \ - find /install -name "*.so" -exec strip --strip-unneeded {} \; 2>/dev/null || true + find /install -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true && \ + find /install -type d -name "tests" -exec rm -rf {} + 2>/dev/null || true && \ + find /install -type d -name "test" -exec rm -rf {} + 2>/dev/null || true && \ + find /install -type d -name "*.dist-info" -exec sh -c 'rm -f "$1"/RECORD "$1"/INSTALLER' _ {} \; && \ + find /install -type f -name "*.pyc" -delete && \ + find /install -type f -name "*.pyo" -delete && \ + find /install -name "*.so" -exec strip --strip-unneeded {} \; 2>/dev/null || true # 运行阶段 - 使用最小镜像 FROM python:3.11-slim @@ -21,31 +21,35 @@ WORKDIR /app # 清理基础镜像中的冗余文件 RUN rm -rf /usr/share/doc/* \ - /usr/share/man/* \ - /usr/share/locale/* \ - /var/cache/apt/* \ - /var/lib/apt/lists/* \ - /tmp/* \ - /var/tmp/* + /usr/share/man/* \ + /usr/share/locale/* \ + /var/cache/apt/* \ + /var/lib/apt/lists/* \ + /tmp/* \ + /var/tmp/* # 从构建阶段复制已安装的包 COPY --from=builder /install /usr/local -# 创建必要的目录 +# 创建必要的目录（包括用于挂载的data目录） RUN mkdir -p /app/logs /app/data/temp/image /app/data/temp/video -# 复制应用代码和配置文件 +# 复制应用代码 COPY app/ ./app/ -COPY data/setting.toml ./data/setting.toml COPY main.py . -# 创建默认的 token.json 文件 -RUN echo '{"ssoNormal": {}, "ssoSuper": {}}' > /app/data/token.json +# 复制并设置 entrypoint 脚本 +COPY docker-entrypoint.sh /usr/local/bin/ +RUN chmod +x /usr/local/bin/docker-entrypoint.sh # 删除 Python 字节码和缓存 ENV PYTHONDONTWRITEBYTECODE=1 \ - PYTHONUNBUFFERED=1 + PYTHONUNBUFFERED=1 EXPOSE 8000 +# 使用 entrypoint 脚本初始化配置 +ENTRYPOINT ["docker-entrypoint.sh"] + +# 默认启动命令 CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/app/core/config.py b/app/core/config.py index 8036266..4b63ec6 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -28,7 +28,11 @@ DEFAULT_GLOBAL = { "admin_password": "admin", "admin_username": "admin", "image_cache_max_size_mb": 512, - "video_cache_max_size_mb": 1024 + "video_cache_max_size_mb": 1024, + "max_upload_concurrency": 20, # 最大并发上传数 + "max_request_concurrency": 50, # 最大并发请求数 + "batch_save_interval": 1.0, # 批量保存间隔（秒） + "batch_save_threshold": 10 # 触发批量保存的变更数阈值 } diff --git a/app/services/grok/client.py b/app/services/grok/client.py index 277d640..435a722 100644 --- a/app/services/grok/client.py +++ b/app/services/grok/client.py @@ -21,13 +21,23 @@ API_ENDPOINT = "https://grok.com/rest/app-chat/conversations/new" TIMEOUT = 120 BROWSER = "chrome133a" MAX_RETRY = 3 -MAX_UPLOADS = 5 +MAX_UPLOADS = 20 # 提高并发上传限制以支持更高并发 class GrokClient: """Grok API 客户端""" - _upload_sem = asyncio.Semaphore(MAX_UPLOADS) + _upload_sem = None # 延迟初始化 + + @staticmethod + def _get_upload_semaphore(): + """获取上传信号量（动态配置）""" + if GrokClient._upload_sem is None: + # 从配置读取，如果不可用则使用默认值 + max_concurrency = setting.global_config.get("max_upload_concurrency", MAX_UPLOADS) + GrokClient._upload_sem = asyncio.Semaphore(max_concurrency) + logger.debug(f"[Client] 初始化上传并发限制: {max_concurrency}") + return GrokClient._upload_sem @staticmethod async def openai_to_grok(request: dict): @@ -109,7 +119,7 @@ class GrokClient: return [], [] async def upload_limited(url): - async with GrokClient._upload_sem: + async with GrokClient._get_upload_semaphore(): return await ImageUploadManager.upload(url, token) results = await asyncio.gather(*[upload_limited(u) for u in urls], return_exceptions=True) diff --git a/app/services/grok/token.py b/app/services/grok/token.py index 95aee40..08ab037 100644 --- a/app/services/grok/token.py +++ b/app/services/grok/token.py @@ -4,6 +4,7 @@ import orjson import time import asyncio import aiofiles +import portalocker from pathlib import Path from curl_cffi.requests import AsyncSession from typing import Dict, Any, Optional, Tuple @@ -43,7 +44,13 @@ class GrokTokenManager: self._file_lock = asyncio.Lock() self.token_file.parent.mkdir(parents=True, exist_ok=True) self._storage = None - self._load_data() + self.token_data = None # 延迟加载 + + # 批量保存队列 + self._save_pending = False # 标记是否有待保存的数据 + self._save_task = None # 后台保存任务 + self._shutdown = False # 关闭标志 + self._initialized = True logger.debug(f"[Token] 初始化完成: {self.token_file}") @@ -51,14 +58,21 @@ class GrokTokenManager: """设置存储实例""" self._storage = storage - def _load_data(self) -> None: - """加载Token数据""" + async def _load_data(self) -> None: + """异步加载Token数据（支持多进程）""" default = {TokenType.NORMAL.value: {}, TokenType.SUPER.value: {}} try: if self.token_file.exists(): - with open(self.token_file, "r", encoding="utf-8") as f: - self.token_data = orjson.loads(f.read()) + # 使用进程锁读取文件 + async with self._file_lock: + with open(self.token_file, "r", encoding="utf-8") as f: + portalocker.lock(f, portalocker.LOCK_SH) # 共享锁（读） + try: + content = f.read() + self.token_data = orjson.loads(content) + finally: + portalocker.unlock(f) else: self.token_data = default logger.debug("[Token] 创建新数据文件") @@ -67,18 +81,69 @@ class GrokTokenManager: self.token_data = default async def _save_data(self) -> None: - """保存Token数据""" + """保存Token数据（支持多进程）""" try: if not self._storage: async with self._file_lock: - async with aiofiles.open(self.token_file, "w", encoding="utf-8") as f: - await f.write(orjson.dumps(self.token_data, option=orjson.OPT_INDENT_2).decode()) + # 使用进程锁写入文件 + with open(self.token_file, "w", encoding="utf-8") as f: + portalocker.lock(f, portalocker.LOCK_EX) # 独占锁（写） + try: + content = orjson.dumps(self.token_data, option=orjson.OPT_INDENT_2).decode() + f.write(content) + f.flush() # 确保写入磁盘 + finally: + portalocker.unlock(f) else: await self._storage.save_tokens(self.token_data) except IOError as e: logger.error(f"[Token] 保存失败: {e}") raise GrokApiException(f"保存失败: {e}", "TOKEN_SAVE_ERROR", {"file": str(self.token_file)}) + def _mark_dirty(self) -> None: + """标记有待保存的数据""" + self._save_pending = True + + async def _batch_save_worker(self) -> None: + """批量保存后台任务""" + from app.core.config import setting + + interval = setting.global_config.get("batch_save_interval", 1.0) + logger.info(f"[Token] 批量保存任务已启动，间隔: {interval}s") + + while not self._shutdown: + await asyncio.sleep(interval) + + if self._save_pending and not self._shutdown: + try: + await self._save_data() + self._save_pending = False + logger.debug("[Token] 批量保存完成") + except Exception as e: + logger.error(f"[Token] 批量保存失败: {e}") + + async def start_batch_save(self) -> None: + """启动批量保存任务""" + if self._save_task is None: + self._save_task = asyncio.create_task(self._batch_save_worker()) + logger.info("[Token] 批量保存任务已创建") + + async def shutdown(self) -> None: + """关闭并刷新所有待保存数据""" + self._shutdown = True + + if self._save_task: + self._save_task.cancel() + try: + await self._save_task + except asyncio.CancelledError: + pass + + # 最终刷新 + if self._save_pending: + await self._save_data() + logger.info("[Token] 关闭时刷新完成") + @staticmethod def _extract_sso(auth_token: str) -> Optional[str]: """提取SSO值""" @@ -117,7 +182,7 @@ class GrokTokenManager: } count += 1 - await self._save_data() + self._mark_dirty() # 批量保存 logger.info(f"[Token] 添加 {count} 个 {token_type.value} Token") async def delete_token(self, tokens: list[str], token_type: TokenType) -> None: @@ -131,7 +196,7 @@ class GrokTokenManager: del self.token_data[token_type.value][token] count += 1 - await self._save_data() + self._mark_dirty() # 批量保存 logger.info(f"[Token] 删除 {count} 个 {token_type.value} Token") async def update_token_tags(self, token: str, token_type: TokenType, tags: list[str]) -> None: @@ -141,7 +206,7 @@ class GrokTokenManager: cleaned = [t.strip() for t in tags if t and t.strip()] self.token_data[token_type.value][token]["tags"] = cleaned - await self._save_data() + self._mark_dirty() # 批量保存 logger.info(f"[Token] 更新标签: {token[:10]}... -> {cleaned}") async def update_token_note(self, token: str, token_type: TokenType, note: str) -> None: @@ -150,20 +215,40 @@ class GrokTokenManager: raise GrokApiException("Token不存在", "TOKEN_NOT_FOUND", {"token": token[:10]}) self.token_data[token_type.value][token]["note"] = note.strip() - await self._save_data() + self._mark_dirty() # 批量保存 logger.info(f"[Token] 更新备注: {token[:10]}...") def get_tokens(self) -> Dict[str, Any]: """获取所有Token""" return self.token_data.copy() + def _reload_if_needed(self) -> None: + """在多进程模式下重新加载数据（同步版本，用于select_token）""" + # 只在文件模式且多进程环境下才重新加载 + if self._storage: + return # 数据库模式不需要 + + try: + if self.token_file.exists(): + with open(self.token_file, "r", encoding="utf-8") as f: + portalocker.lock(f, portalocker.LOCK_SH) + try: + content = f.read() + self.token_data = orjson.loads(content) + finally: + portalocker.unlock(f) + except Exception as e: + logger.warning(f"[Token] 重新加载失败: {e}") + def get_token(self, model: str) -> str: """获取Token""" jwt = self.select_token(model) return f"sso-rw={jwt};sso={jwt}" def select_token(self, model: str) -> str: - """选择最优Token""" + """选择最优Token（多进程安全）""" + # 重新加载最新数据（多进程模式） + self._reload_if_needed() def select_best(tokens: Dict[str, Any], field: str) -> Tuple[Optional[str], Optional[int]]: """选择最佳Token""" unused, used = [], [] @@ -285,7 +370,7 @@ class GrokTokenManager: self.token_data[token_type][sso]["remainingQueries"] = normal if heavy is not None: self.token_data[token_type][sso]["heavyremainingQueries"] = heavy - await self._save_data() + self._mark_dirty() # 批量保存 logger.info(f"[Token] 更新限制: {sso[:10]}...") return logger.warning(f"[Token] 未找到: {sso[:10]}...") @@ -321,7 +406,7 @@ class GrokTokenManager: data["status"] = "expired" logger.error(f"[Token] 标记失效: {sso[:10]}... (连续{status}错误{data['failedCount']}次)") - await self._save_data() + self._mark_dirty() # 批量保存 except Exception as e: logger.error(f"[Token] 记录失败错误: {e}") @@ -341,7 +426,7 @@ class GrokTokenManager: data["failedCount"] = 0 data["lastFailureTime"] = None data["lastFailureReason"] = None - await self._save_data() + self._mark_dirty() # 批量保存 logger.info(f"[Token] 重置失败计数: {sso[:10]}...") except Exception as e: diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh new file mode 100644 index 0000000..b98e9f5 --- /dev/null +++ b/docker-entrypoint.sh @@ -0,0 +1,53 @@ +#!/bin/sh +set -e + +# 初始化配置文件（如果不存在） +echo "[Grok2API] 检查配置文件..." + +# 确保数据目录存在 +mkdir -p /app/data/temp/image /app/data/temp/video /app/logs + +# 如果 setting.toml 不存在，创建默认配置 +if [ ! -f /app/data/setting.toml ]; then + echo "[Grok2API] 初始化 setting.toml..." + cat > /app/data/setting.toml << 'EOF' +[global] +base_url = "http://localhost:8000" +log_level = "INFO" +image_mode = "url" +admin_password = "admin" +admin_username = "admin" +image_cache_max_size_mb = 512 +video_cache_max_size_mb = 1024 +max_upload_concurrency = 20 +max_request_concurrency = 50 +batch_save_interval = 1.0 +batch_save_threshold = 10 + +[grok] +api_key = "" +proxy_url = "" +cache_proxy_url = "" +cf_clearance = "" +x_statsig_id = "" +dynamic_statsig = true +filtered_tags = "xaiartifact,xai:tool_usage_card,grok:render" +stream_chunk_timeout = 120 +stream_total_timeout = 600 +stream_first_response_timeout = 30 +temporary = true +show_thinking = true +EOF +fi + +# 如果 token.json 不存在，创建空token文件 +if [ ! -f /app/data/token.json ]; then + echo "[Grok2API] 初始化 token.json..." + echo '{"ssoNormal": {}, "ssoSuper": {}}' > /app/data/token.json +fi + +echo "[Grok2API] 配置文件检查完成" +echo "[Grok2API] 启动应用..." + +# 执行传入的命令 +exec "$@" diff --git a/main.py b/main.py index 2b9c463..a042695 100644 --- a/main.py +++ b/main.py @@ -26,11 +26,14 @@ async def lifespan(app: FastAPI): """ 启动顺序: 1. 初始化核心服务 (storage, settings, token_manager) - 2. 启动MCP服务生命周期 + 2. 异步加载 token 数据 + 3. 启动批量保存任务 + 4. 启动MCP服务生命周期 关闭顺序 (LIFO): 1. 关闭MCP服务生命周期 - 2. 关闭核心服务 + 2. 关闭批量保存任务并刷新数据 + 3. 关闭核心服务 """ # --- 启动过程 --- # 1. 初始化核心服务 @@ -41,12 +44,18 @@ async def lifespan(app: FastAPI): setting.set_storage(storage) token_manager.set_storage(storage) - # 重新加载配置和token数据 + # 2. 重新加载配置 await setting.reload() - token_manager._load_data() logger.info("[Grok2API] 核心服务初始化完成") + + # 3. 异步加载 token 数据 + await token_manager._load_data() + logger.info("[Grok2API] Token数据加载完成") + + # 4. 启动批量保存任务 + await token_manager.start_batch_save() - # 2. 管理MCP服务的生命周期 + # 5. 管理MCP服务的生命周期 mcp_lifespan_context = mcp_app.lifespan(app) await mcp_lifespan_context.__aenter__() logger.info("[MCP] MCP服务初始化完成") @@ -61,7 +70,11 @@ async def lifespan(app: FastAPI): await mcp_lifespan_context.__aexit__(None, None, None) logger.info("[MCP] MCP服务已关闭") - # 2. 关闭核心服务 + # 2. 关闭批量保存任务并刷新数据 + await token_manager.shutdown() + logger.info("[Token] Token管理器已关闭") + + # 3. 关闭核心服务 await storage_manager.close() logger.info("[Grok2API] 应用关闭成功") @@ -111,4 +124,22 @@ app.mount("", mcp_app) if __name__ == "__main__": import uvicorn - uvicorn.run(app, host="0.0.0.0", port=8001) \ No newline at end of file + import os + + # 从环境变量读取 worker 数量，默认为 1 + workers = int(os.getenv("WORKERS", "1")) + + # 提示多进程模式 + if workers > 1: + logger.info( + f"[Grok2API] 多进程模式已启用 (workers={workers})。" + f"建议使用 Redis/MySQL 存储以获得最佳性能。" + ) + + uvicorn.run( + "main:app", # 使用字符串以支持多worker + host="0.0.0.0", + port=8001, + workers=workers, + loop="uvloop" if workers == 1 else "auto" # 单进程用uvloop，多进程用auto + ) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index cab633d..85145fa 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,12 +1,14 @@ toml==0.10.2 fastapi==0.119.0 uvicorn==0.37.0 +uvloop==0.21.0 python-dotenv==1.1.1 curl_cffi==0.13.0 requests==2.32.5 starlette==0.48.0 pydantic==2.12.2 aiofiles==25.1.0 +portalocker==3.0.0 aiomysql==0.2.0 redis==6.4.0 fastmcp==2.12.4 diff --git a/test_concurrency.py b/test_concurrency.py new file mode 100755 index 0000000..c0089b1 --- /dev/null +++ b/test_concurrency.py @@ -0,0 +1,276 @@ +#!/usr/bin/env python3 +""" +Grok2API 并发性能测试脚本 + +测试不同并发级别下的API性能表现 +""" + +import asyncio +import aiohttp +import time +import statistics +import argparse +from datetime import datetime +from typing import List, Dict, Any +import json + + +class ConcurrencyTester: + """并发测试器""" + + def __init__(self, base_url: str, api_key: str = None): + self.base_url = base_url.rstrip('/') + self.api_key = api_key + self.results: List[Dict[str, Any]] = [] + + async def test_request(self, session: aiohttp.ClientSession, request_id: int) -> Dict[str, Any]: + """发送单个测试请求""" + url = f"{self.base_url}/v1/chat/completions" + + headers = { + "Content-Type": "application/json" + } + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + + payload = { + "model": "grok-3-fast", + "messages": [ + {"role": "user", "content": f"测试请求 #{request_id}，请简短回复OK"} + ], + "stream": False, + "max_tokens": 10 + } + + start_time = time.time() + + try: + async with session.post(url, json=payload, headers=headers, timeout=30) as response: + status = response.status + + if status == 200: + data = await response.json() + elapsed = time.time() - start_time + + return { + "id": request_id, + "status": "success", + "http_status": status, + "elapsed": elapsed, + "response_length": len(json.dumps(data)) + } + else: + elapsed = time.time() - start_time + error_text = await response.text() + + return { + "id": request_id, + "status": "error", + "http_status": status, + "elapsed": elapsed, + "error": error_text[:200] + } + + except asyncio.TimeoutError: + elapsed = time.time() - start_time + return { + "id": request_id, + "status": "timeout", + "elapsed": elapsed, + "error": "Request timeout" + } + + except Exception as e: + elapsed = time.time() - start_time + return { + "id": request_id, + "status": "exception", + "elapsed": elapsed, + "error": str(e) + } + + async def run_concurrent_test(self, concurrency: int, total_requests: int): + """运行并发测试""" + print(f"\n{'='*60}") + print(f"📊 测试配置：并发数 {concurrency}, 总请求数 {total_requests}") + print(f"{'='*60}") + + connector = aiohttp.TCPConnector(limit=concurrency, limit_per_host=concurrency) + timeout = aiohttp.ClientTimeout(total=60) + + async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session: + # 预热 + print("🔥 预热中...") + await self.test_request(session, 0) + + # 开始测试 + print(f"🚀 开始并发测试...") + start_time = time.time() + + # 创建任务 + tasks = [] + for i in range(1, total_requests + 1): + task = asyncio.create_task(self.test_request(session, i)) + tasks.append(task) + + # 控制并发数 + if len(tasks) >= concurrency: + results = await asyncio.gather(*tasks) + self.results.extend(results) + tasks = [] + + # 显示进度 + print(f" 进度: {i}/{total_requests} ({i/total_requests*100:.1f}%)", end='\r') + + # 处理剩余任务 + if tasks: + results = await asyncio.gather(*tasks) + self.results.extend(results) + + total_time = time.time() - start_time + + # 统计和输出 + self.print_statistics(concurrency, total_requests, total_time) + + def print_statistics(self, concurrency: int, total_requests: int, total_time: float): + """打印统计信息""" + success_results = [r for r in self.results if r["status"] == "success"] + error_results = [r for r in self.results if r["status"] != "success"] + + success_count = len(success_results) + error_count = len(error_results) + + if success_results: + latencies = [r["elapsed"] for r in success_results] + avg_latency = statistics.mean(latencies) + min_latency = min(latencies) + max_latency = max(latencies) + p50_latency = statistics.median(latencies) + p95_latency = sorted(latencies)[int(len(latencies) * 0.95)] if len(latencies) > 1 else latencies[0] + p99_latency = sorted(latencies)[int(len(latencies) * 0.99)] if len(latencies) > 1 else latencies[0] + else: + avg_latency = min_latency = max_latency = p50_latency = p95_latency = p99_latency = 0 + + throughput = total_requests / total_time if total_time > 0 else 0 + + print(f"\n\n{'='*60}") + print(f"📈 测试结果统计") + print(f"{'='*60}") + print(f" 测试时间: {total_time:.2f}s") + print(f" 总请求数: {total_requests}") + print(f" 并发数: {concurrency}") + print(f"") + print(f" 成功请求: {success_count} ({success_count/total_requests*100:.1f}%)") + print(f" 失败请求: {error_count} ({error_count/total_requests*100:.1f}%)") + print(f"") + print(f" 吞吐量: {throughput:.2f} req/s") + print(f"") + print(f" 延迟统计:") + print(f" 最小: {min_latency*1000:.0f}ms") + print(f" 平均: {avg_latency*1000:.0f}ms") + print(f" 最大: {max_latency*1000:.0f}ms") + print(f" P50: {p50_latency*1000:.0f}ms") + print(f" P95: {p95_latency*1000:.0f}ms") + print(f" P99: {p99_latency*1000:.0f}ms") + + # 错误详情 + if error_results: + print(f"\n ⚠️ 错误详情:") + error_types = {} + for r in error_results: + error_type = r.get("status", "unknown") + error_types[error_type] = error_types.get(error_type, 0) + 1 + + for error_type, count in error_types.items(): + print(f" {error_type}: {count}") + + print(f"{'='*60}\n") + + # 性能评级 + self.print_performance_rating(throughput, avg_latency) + + def print_performance_rating(self, throughput: float, avg_latency: float): + """打印性能评级""" + print(f"🎯 性能评级:") + + # 吞吐量评级 + if throughput >= 100: + rating = "⭐⭐⭐⭐⭐ 优秀" + elif throughput >= 60: + rating = "⭐⭐⭐⭐ 良好" + elif throughput >= 30: + rating = "⭐⭐⭐ 中等" + elif throughput >= 10: + rating = "⭐⭐ 较低" + else: + rating = "⭐ 需优化" + + print(f" 吞吐量 ({throughput:.1f} req/s): {rating}") + + # 延迟评级 + if avg_latency < 0.5: + rating = "⭐⭐⭐⭐⭐ 优秀" + elif avg_latency < 1.0: + rating = "⭐⭐⭐⭐ 良好" + elif avg_latency < 2.0: + rating = "⭐⭐⭐ 中等" + elif avg_latency < 5.0: + rating = "⭐⭐ 较高" + else: + rating = "⭐ 需优化" + + print(f" 平均延迟 ({avg_latency*1000:.0f}ms): {rating}") + print() + + +async def main(): + """主函数""" + parser = argparse.ArgumentParser(description='Grok2API 并发性能测试') + parser.add_argument('--url', default='http://localhost:8001', help='API 基础URL') + parser.add_argument('--key', default='', help='API Key（可选）') + parser.add_argument('-c', '--concurrency', type=int, default=10, help='并发数') + parser.add_argument('-n', '--requests', type=int, default=50, help='总请求数') + parser.add_argument('--multi-test', action='store_true', help='运行多级并发测试') + + args = parser.parse_args() + + print(f""" +╔══════════════════════════════════════════════════════════╗ +║ Grok2API 并发性能测试工具 ║ +╚══════════════════════════════════════════════════════════╝ + +🔗 测试目标: {args.url} +🔑 API Key: {'已设置' if args.key else '未设置'} +⏰ 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} +""") + + tester = ConcurrencyTester(args.url, args.key) + + if args.multi_test: + # 多级并发测试 + test_configs = [ + (5, 20), # 5并发，20请求 + (10, 50), # 10并发，50请求 + (20, 100), # 20并发，100请求 + (50, 200), # 50并发，200请求 + ] + + for concurrency, requests in test_configs: + tester.results = [] # 清空结果 + await tester.run_concurrent_test(concurrency, requests) + await asyncio.sleep(2) # 间隔2秒 + else: + # 单次测试 + await tester.run_concurrent_test(args.concurrency, args.requests) + + print(f"\n✅ 测试完成！") + print(f"⏰ 结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + print("\n\n⚠️ 测试被用户中断") + except Exception as e: + print(f"\n\n❌ 测试失败: {e}") diff --git a/test_concurrency.sh b/test_concurrency.sh new file mode 100755 index 0000000..b100e2c --- /dev/null +++ b/test_concurrency.sh @@ -0,0 +1,177 @@ +#!/bin/bash + +# Grok2API 并发测试脚本（Shell版本） +# 使用 curl 和 GNU parallel 进行并发测试 + +set -e + +# 配置 +BASE_URL="${BASE_URL:-http://localhost:8001}" +API_KEY="${API_KEY:-}" +CONCURRENCY="${CONCURRENCY:-10}" +TOTAL_REQUESTS="${TOTAL_REQUESTS:-50}" + +# 颜色 +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +NC='\033[0m' # No Color + +echo -e "${BLUE}╔══════════════════════════════════════════════════════════╗${NC}" +echo -e "${BLUE}║ Grok2API 并发性能测试工具 (Shell版) ║${NC}" +echo -e "${BLUE}╚══════════════════════════════════════════════════════════╝${NC}" +echo "" +echo -e "${GREEN}🔗 测试目标:${NC} $BASE_URL" +echo -e "${GREEN}🔑 API Key:${NC} ${API_KEY:-(未设置)}" +echo -e "${GREEN}📊 并发数:${NC} $CONCURRENCY" +echo -e "${GREEN}📈 总请求数:${NC} $TOTAL_REQUESTS" +echo "" + +# 检查依赖 +if ! command -v curl &> /dev/null; then + echo -e "${RED}❌ 错误: 需要安装 curl${NC}" + exit 1 +fi + +# 创建临时目录 +TMP_DIR=$(mktemp -d) +trap "rm -rf $TMP_DIR" EXIT + +# 单个请求函数 +test_request() { + local request_id=$1 + local start_time=$(date +%s.%N) + + # 构建请求 + local headers="Content-Type: application/json" + if [ -n "$API_KEY" ]; then + headers="${headers}\nAuthorization: Bearer ${API_KEY}" + fi + + local response=$(curl -s -w "\n%{http_code}\n%{time_total}" \ + -X POST "${BASE_URL}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + ${API_KEY:+-H "Authorization: Bearer $API_KEY"} \ + -d "{ + \"model\": \"grok-3-fast\", + \"messages\": [{\"role\": \"user\", \"content\": \"测试请求 #${request_id}，请简短回复OK\"}], + \"stream\": false, + \"max_tokens\": 10 + }" 2>&1) + + local http_code=$(echo "$response" | tail -n 2 | head -n 1) + local time_total=$(echo "$response" | tail -n 1) + + # 记录结果 + echo "${request_id},${http_code},${time_total}" >> "$TMP_DIR/results.csv" + + # 显示进度 + echo -ne "\r 进度: ${request_id}/${TOTAL_REQUESTS}" +} + +# 导出函数供 parallel 使用 +export -f test_request +export BASE_URL API_KEY TMP_DIR + +# 清空结果文件 +echo "id,status,time" > "$TMP_DIR/results.csv" + +echo -e "${YELLOW}🚀 开始并发测试...${NC}" +START_TIME=$(date +%s.%N) + +# 使用 GNU parallel（如果可用），否则使用简单循环 +if command -v parallel &> /dev/null; then + seq 1 $TOTAL_REQUESTS | parallel -j $CONCURRENCY test_request {} +else + # 简单的后台任务并发 + for i in $(seq 1 $TOTAL_REQUESTS); do + test_request $i & + + # 控制并发数 + if (( i % CONCURRENCY == 0 )); then + wait + fi + done + wait +fi + +END_TIME=$(date +%s.%N) +TOTAL_TIME=$(echo "$END_TIME - $START_TIME" | bc) + +echo -e "\n" + +# 统计结果 +echo -e "${BLUE}═══════════════════════════════════════════════════════════${NC}" +echo -e "${BLUE}📈 测试结果统计${NC}" +echo -e "${BLUE}═══════════════════════════════════════════════════════════${NC}" + +# 统计成功/失败 +SUCCESS_COUNT=$(awk -F',' '$2 == 200 {count++} END {print count+0}' "$TMP_DIR/results.csv") +ERROR_COUNT=$((TOTAL_REQUESTS - SUCCESS_COUNT)) + +echo -e " 测试时间: ${TOTAL_TIME}s" +echo -e " 总请求数: ${TOTAL_REQUESTS}" +echo -e " 并发数: ${CONCURRENCY}" +echo "" +echo -e " 成功请求: ${GREEN}${SUCCESS_COUNT}${NC} ($(echo "scale=1; $SUCCESS_COUNT * 100 / $TOTAL_REQUESTS" | bc)%)" +echo -e " 失败请求: ${RED}${ERROR_COUNT}${NC} ($(echo "scale=1; $ERROR_COUNT * 100 / $TOTAL_REQUESTS" | bc)%)" +echo "" + +# 计算吞吐量 +THROUGHPUT=$(echo "scale=2; $TOTAL_REQUESTS / $TOTAL_TIME" | bc) +echo -e " 吞吐量: ${GREEN}${THROUGHPUT}${NC} req/s" +echo "" + +# 延迟统计（只统计成功的请求） +if [ $SUCCESS_COUNT -gt 0 ]; then + echo -e " 延迟统计:" + + # 提取成功请求的延迟时间 + awk -F',' '$2 == 200 {print $3}' "$TMP_DIR/results.csv" | sort -n > "$TMP_DIR/latencies.txt" + + MIN=$(head -n 1 "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + MAX=$(tail -n 1 "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + AVG=$(awk '{sum+=$1; count++} END {printf "%.0f", sum/count*1000}' "$TMP_DIR/latencies.txt") + + # P50 + P50_LINE=$((SUCCESS_COUNT / 2)) + P50=$(sed -n "${P50_LINE}p" "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + + # P95 + P95_LINE=$(echo "scale=0; $SUCCESS_COUNT * 0.95 / 1" | bc) + P95=$(sed -n "${P95_LINE}p" "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + + # P99 + P99_LINE=$(echo "scale=0; $SUCCESS_COUNT * 0.99 / 1" | bc) + P99=$(sed -n "${P99_LINE}p" "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + + echo -e " 最小: ${MIN}ms" + echo -e " 平均: ${AVG}ms" + echo -e " 最大: ${MAX}ms" + echo -e " P50: ${P50}ms" + echo -e " P95: ${P95}ms" + echo -e " P99: ${P99}ms" +fi + +echo -e "${BLUE}═══════════════════════════════════════════════════════════${NC}" + +# 性能评级 +echo -e "${YELLOW}🎯 性能评级:${NC}" + +if (( $(echo "$THROUGHPUT >= 100" | bc -l) )); then + RATING="⭐⭐⭐⭐⭐ 优秀" +elif (( $(echo "$THROUGHPUT >= 60" | bc -l) )); then + RATING="⭐⭐⭐⭐ 良好" +elif (( $(echo "$THROUGHPUT >= 30" | bc -l) )); then + RATING="⭐⭐⭐ 中等" +elif (( $(echo "$THROUGHPUT >= 10" | bc -l) )); then + RATING="⭐⭐ 较低" +else + RATING="⭐ 需优化" +fi + +echo -e " 吞吐量 (${THROUGHPUT} req/s): ${RATING}" + +echo "" +echo -e "${GREEN}✅ 测试完成！${NC}"