fix: restore main runtime sources and cache consistency
This commit is contained in:
User
273
scripts/analyze_rag_performance.py
Normal file
273
scripts/analyze_rag_performance.py
Normal file
@@ -0,0 +1,273 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
RAG 性能分析脚本
|
||||
|
||||
分析 RAG 流程各阶段耗时，帮助定位性能瓶颈
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import argparse
|
||||
from datetime import datetime
|
||||
from typing import Dict
|
||||
|
||||
# 添加项目根目录到路径
|
||||
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, PROJECT_ROOT)
|
||||
|
||||
|
||||
def call_rag_stream_api(
|
||||
question: str,
|
||||
kb_name: str = "public_kb",
|
||||
base_url: str = "http://localhost:5001",
|
||||
) -> Dict:
|
||||
"""
|
||||
调用 RAG 流式 API 并收集各阶段事件
|
||||
|
||||
Returns:
|
||||
包含各阶段耗时信息的字典
|
||||
"""
|
||||
import requests
|
||||
|
||||
url = f"{base_url.rstrip('/')}/rag"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": "Bearer mock-token-admin"
|
||||
}
|
||||
data = {
|
||||
"message": question,
|
||||
"collections": [kb_name] if kb_name else []
|
||||
}
|
||||
|
||||
events = []
|
||||
timing = {
|
||||
'total_duration_ms': 0,
|
||||
'search_time_ms': 0,
|
||||
'rerank_time_ms': 0,
|
||||
'llm_time_ms': 0,
|
||||
'stages': []
|
||||
}
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
response = requests.post(url, json=data, headers=headers, stream=True, timeout=300)
|
||||
response.raise_for_status()
|
||||
|
||||
for line in response.iter_lines():
|
||||
if line:
|
||||
line_str = line.decode('utf-8')
|
||||
if line_str.startswith('data: '):
|
||||
try:
|
||||
event = json.loads(line_str[6:])
|
||||
event['_received_at'] = time.time()
|
||||
events.append(event)
|
||||
|
||||
# 分析 finish 事件中的耗时信息
|
||||
if event.get('type') == 'finish':
|
||||
timing['total_duration_ms'] = event.get('duration_ms', 0)
|
||||
if 'timing' in event:
|
||||
timing['search_time_ms'] = event['timing'].get('total_search_ms', 0)
|
||||
timing['rerank_time_ms'] = event['timing'].get('rerank_ms', 0)
|
||||
timing['rerank_cached'] = event['timing'].get('rerank_cached', False)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
end_time = time.time()
|
||||
timing['actual_elapsed_ms'] = int((end_time - start_time) * 1000)
|
||||
|
||||
# 计算 LLM 生成时间
|
||||
if timing['total_duration_ms'] > 0 and timing['search_time_ms'] > 0:
|
||||
timing['llm_time_ms'] = timing['total_duration_ms'] - timing['search_time_ms']
|
||||
|
||||
finish_event = next((event for event in events if event.get('type') == 'finish'), None)
|
||||
error_event = next((event for event in events if event.get('type') == 'error'), None)
|
||||
return {
|
||||
'success': finish_event is not None and error_event is None,
|
||||
'error': error_event.get('message', '') if error_event else (
|
||||
'' if finish_event else 'SSE 流结束但未收到 finish 事件'
|
||||
),
|
||||
'events': events,
|
||||
'timing': timing,
|
||||
'answer': finish_event.get('answer', '') if finish_event else '',
|
||||
}
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': str(e),
|
||||
'events': events,
|
||||
'timing': timing
|
||||
}
|
||||
|
||||
|
||||
def analyze_performance(result: Dict, question: str) -> None:
|
||||
"""分析并打印性能数据"""
|
||||
print("\n" + "=" * 80)
|
||||
print(f"问题: {question}")
|
||||
print("=" * 80)
|
||||
|
||||
if not result['success']:
|
||||
print(f"[X] 请求失败: {result.get('error', '未知错误')}")
|
||||
return
|
||||
|
||||
timing = result['timing']
|
||||
events = result['events']
|
||||
|
||||
# 打印各阶段事件时间线
|
||||
print("\n[事件时间线]")
|
||||
print("-" * 80)
|
||||
first_event_time = None
|
||||
for event in events:
|
||||
event_type = event.get('type', 'unknown')
|
||||
received_at = event.get('_received_at', 0)
|
||||
|
||||
if first_event_time is None:
|
||||
first_event_time = received_at
|
||||
relative_time = 0
|
||||
else:
|
||||
relative_time = (received_at - first_event_time) * 1000
|
||||
|
||||
if event_type == 'finish':
|
||||
print(f" {relative_time:>8.0f}ms | {event_type:20s} | 总耗时: {event.get('duration_ms', 0)}ms")
|
||||
elif event_type == 'sources':
|
||||
sources = event.get('sources', [])
|
||||
print(f" {relative_time:>8.0f}ms | {event_type:20s} | 找到 {len(sources)} 个来源")
|
||||
elif event_type == 'chunks_retrieved':
|
||||
chunks = event.get('data', {}).get('chunks', [])
|
||||
print(f" {relative_time:>8.0f}ms | {event_type:20s} | 召回 {len(chunks)} 个切片")
|
||||
elif event_type == 'chunk':
|
||||
# 流式输出，只显示第一个
|
||||
if not hasattr(analyze_performance, '_chunk_printed'):
|
||||
print(f" {relative_time:>8.0f}ms | {event_type:20s} | 开始流式输出...")
|
||||
analyze_performance._chunk_printed = True
|
||||
else:
|
||||
print(f" {relative_time:>8.0f}ms | {event_type:20s}")
|
||||
|
||||
if hasattr(analyze_performance, '_chunk_printed'):
|
||||
delattr(analyze_performance, '_chunk_printed')
|
||||
|
||||
# 打印耗时统计
|
||||
print("\n[耗时统计]")
|
||||
print("-" * 80)
|
||||
total = timing['total_duration_ms']
|
||||
search = timing['search_time_ms']
|
||||
rerank = timing['rerank_time_ms']
|
||||
llm = timing['llm_time_ms']
|
||||
actual = timing.get('actual_elapsed_ms', total)
|
||||
|
||||
if total > 0:
|
||||
print(f" 总耗时 (API报告): {total:>8}ms ({total/1000:.2f}s)")
|
||||
print(f" 实际耗时 (本地测量): {actual:>8}ms ({actual/1000:.2f}s)")
|
||||
print()
|
||||
print(f" 检索阶段: {search:>8}ms ({search/total*100:.1f}%)")
|
||||
if rerank > 0:
|
||||
cached_flag = " [缓存]" if timing.get('rerank_cached') else ""
|
||||
print(f" 重排序阶段: {rerank:>8}ms ({rerank/total*100:.1f}%){cached_flag}")
|
||||
print(f" LLM生成阶段: {llm:>8}ms ({llm/total*100:.1f}%)")
|
||||
|
||||
# 性能诊断
|
||||
print("\n[性能诊断]")
|
||||
print("-" * 80)
|
||||
if total > 10000:
|
||||
print(" [!] 总耗时超过 10 秒，需要优化")
|
||||
|
||||
if search > 3000:
|
||||
print(f" [!] 检索耗时过长 ({search}ms)，可能原因:")
|
||||
print(" - 向量库数据量过大")
|
||||
print(" - 未命中查询缓存")
|
||||
print(" - BM25 索引加载慢")
|
||||
|
||||
if rerank > 2000 and not timing.get('rerank_cached'):
|
||||
print(f" [!] 重排序耗时过长 ({rerank}ms)，可能原因:")
|
||||
print(" - Rerank 模型计算量大")
|
||||
print(" - 候选切片数量过多")
|
||||
|
||||
if llm > 5000:
|
||||
print(f" [!] LLM生成耗时过长 ({llm}ms)，可能原因:")
|
||||
print(" - 模型生成速度慢")
|
||||
print(" - 输出 token 数量多")
|
||||
print(" - 网络延迟")
|
||||
print(" 建议: 检查 MiMo thinking 是否关闭、上下文长度及 max_tokens")
|
||||
|
||||
if total < 5000:
|
||||
print(" [OK] 性能良好")
|
||||
else:
|
||||
print(" [!] 未能获取耗时数据")
|
||||
|
||||
|
||||
def run_performance_tests(
|
||||
base_url: str = "http://localhost:5001",
|
||||
kb_name: str = "public_kb",
|
||||
question: str = None,
|
||||
):
|
||||
"""执行性能测试"""
|
||||
# 测试问题（不同类型）
|
||||
test_questions = [
|
||||
{"question": "智启科技成立于哪一年？", "type": "精确匹配"},
|
||||
{"question": "公司的愿景是什么？", "type": "语义理解"},
|
||||
{"question": "公司有哪些分公司，分别在哪些城市？", "type": "跨文档关联"},
|
||||
{"question": "一个入职3年的员工，累计病假1个月，能拿到多少病假工资？", "type": "复杂推理"},
|
||||
{"question": "P4级工程师的年薪范围是多少？", "type": "表格数据"},
|
||||
]
|
||||
if question:
|
||||
test_questions = [{"question": question, "type": "自定义问题"}]
|
||||
|
||||
print("=" * 80)
|
||||
print("RAG Performance Analysis")
|
||||
print(f"Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
print("=" * 80)
|
||||
|
||||
results = []
|
||||
for i, test in enumerate(test_questions):
|
||||
print(f"\n\n[{i+1}/{len(test_questions)}] 类型: {test['type']}")
|
||||
result = call_rag_stream_api(test['question'], kb_name=kb_name, base_url=base_url)
|
||||
analyze_performance(result, test['question'])
|
||||
results.append({
|
||||
'question': test['question'],
|
||||
'type': test['type'],
|
||||
'success': result['success'],
|
||||
'timing': result['timing']
|
||||
})
|
||||
|
||||
# 汇总统计
|
||||
print("\n\n" + "=" * 80)
|
||||
print("[Performance Summary]")
|
||||
print("=" * 80)
|
||||
|
||||
successful = [r for r in results if r['success']]
|
||||
if successful:
|
||||
total_times = [r['timing']['total_duration_ms'] for r in successful]
|
||||
search_times = [r['timing']['search_time_ms'] for r in successful]
|
||||
llm_times = [r['timing']['llm_time_ms'] for r in successful]
|
||||
|
||||
print(f"\n成功请求数: {len(successful)}/{len(results)}")
|
||||
print(f"\n平均总耗时: {sum(total_times)/len(total_times):.0f}ms")
|
||||
print(f"平均检索耗时: {sum(search_times)/len(search_times):.0f}ms")
|
||||
print(f"平均LLM耗时: {sum(llm_times)/len(llm_times):.0f}ms")
|
||||
print(f"\n最快总耗时: {min(total_times)}ms")
|
||||
print(f"最慢总耗时: {max(total_times)}ms")
|
||||
|
||||
# 保存详细报告
|
||||
report_dir = os.path.join(PROJECT_ROOT, ".data", "performance")
|
||||
os.makedirs(report_dir, exist_ok=True)
|
||||
report_file = os.path.join(
|
||||
report_dir,
|
||||
f"performance_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
|
||||
)
|
||||
with open(report_file, 'w', encoding='utf-8') as f:
|
||||
json.dump({
|
||||
'test_time': datetime.now().isoformat(),
|
||||
'results': results
|
||||
}, f, ensure_ascii=False, indent=2)
|
||||
print(f"\n详细报告已保存: {report_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="RAG SSE 链路耗时分析")
|
||||
parser.add_argument("--base-url", default="http://localhost:5001")
|
||||
parser.add_argument("--kb", default="public_kb", help="目标知识库名称")
|
||||
parser.add_argument("--question", help="只测试一个自定义问题")
|
||||
args = parser.parse_args()
|
||||
run_performance_tests(args.base_url, args.kb, args.question)
|
||||
Reference in New Issue
Block a user