easy-rag/tests/performance_monitor.py

223 lines
7.3 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""
简单的性能监控脚本
监控并发测试期间的系统资源使用情况
"""
import asyncio
import aiohttp
import time
import psutil
import json
from typing import Dict, List
from datetime import datetime
class SimplePerformanceMonitor:
"""简单性能监控器"""
def __init__(self):
self.metrics = []
self.start_time = None
async def start_monitoring(self, duration: int = 60, interval: float = 1.0):
"""开始监控系统资源"""
self.start_time = time.time()
print(f"🔍 开始性能监控 (持续 {duration} 秒)")
print("-" * 50)
end_time = self.start_time + duration
while time.time() < end_time:
# 获取系统指标
cpu_percent = psutil.cpu_percent(interval=0.1)
memory = psutil.virtual_memory()
metric = {
"timestamp": time.time(),
"relative_time": time.time() - self.start_time,
"cpu_percent": cpu_percent,
"memory_percent": memory.percent,
"memory_used_mb": memory.used / 1024 / 1024,
"memory_available_mb": memory.available / 1024 / 1024
}
self.metrics.append(metric)
# 实时显示
print(f"⏱️ {metric['relative_time']:6.1f}s | "
f"CPU: {cpu_percent:5.1f}% | "
f"内存: {memory.percent:5.1f}% | "
f"已用: {memory.used/1024/1024:6.0f}MB")
await asyncio.sleep(interval)
def generate_summary(self):
"""生成性能摘要"""
if not self.metrics:
print("❌ 没有性能数据")
return
cpu_values = [m["cpu_percent"] for m in self.metrics]
memory_values = [m["memory_percent"] for m in self.metrics]
print("\n" + "=" * 50)
print("📊 性能监控摘要")
print("=" * 50)
print(f"监控时长: {self.metrics[-1]['relative_time']:.1f}")
print(f"采样点数: {len(self.metrics)}")
print(f"\nCPU 使用率:")
print(f" 平均: {sum(cpu_values) / len(cpu_values):5.1f}%")
print(f" 最大: {max(cpu_values):5.1f}%")
print(f" 最小: {min(cpu_values):5.1f}%")
print(f"\n内存使用率:")
print(f" 平均: {sum(memory_values) / len(memory_values):5.1f}%")
print(f" 最大: {max(memory_values):5.1f}%")
print(f" 最小: {min(memory_values):5.1f}%")
# 检查性能警告
avg_cpu = sum(cpu_values) / len(cpu_values)
max_cpu = max(cpu_values)
avg_memory = sum(memory_values) / len(memory_values)
print(f"\n🔍 性能评估:")
if avg_cpu > 80:
print(f"⚠️ 平均 CPU 使用率较高: {avg_cpu:.1f}%")
elif avg_cpu < 20:
print(f"✅ CPU 使用率正常: {avg_cpu:.1f}%")
else:
print(f"✅ CPU 使用率适中: {avg_cpu:.1f}%")
if max_cpu > 95:
print(f"⚠️ CPU 峰值过高: {max_cpu:.1f}%")
else:
print(f"✅ CPU 峰值正常: {max_cpu:.1f}%")
if avg_memory > 85:
print(f"⚠️ 内存使用率较高: {avg_memory:.1f}%")
else:
print(f"✅ 内存使用率正常: {avg_memory:.1f}%")
def save_metrics(self, filename: str = None):
"""保存性能指标"""
if not filename:
filename = f"performance_metrics_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
with open(filename, 'w', encoding='utf-8') as f:
json.dump({
"monitoring_info": {
"start_time": self.start_time,
"duration": self.metrics[-1]['relative_time'] if self.metrics else 0,
"sample_count": len(self.metrics)
},
"metrics": self.metrics
}, f, indent=2, ensure_ascii=False)
print(f"💾 性能数据已保存: {filename}")
async def run_load_test_with_monitoring():
"""运行负载测试并监控性能"""
print("🚀 负载测试 + 性能监控")
print("=" * 50)
# 创建监控器
monitor = SimplePerformanceMonitor()
# 启动监控任务
monitor_task = asyncio.create_task(
monitor.start_monitoring(duration=30, interval=0.5)
)
# 等待一下让监控开始
await asyncio.sleep(1)
# 运行负载测试
async with aiohttp.ClientSession() as session:
print("🔥 开始并发负载...")
tasks = []
# 并发上传任务
for i in range(5):
content = f"性能测试文档 {i+1}" + "这是测试内容。" * 50
tasks.append(upload_document(session, content, f"perf_test_{i+1}.txt"))
# 并发查询任务
for i in range(15):
tasks.append(chat_query(session, f"测试问题 {i+1}"))
# 执行所有任务
print(f"📤 启动 {len(tasks)} 个并发任务...")
results = await asyncio.gather(*tasks, return_exceptions=True)
# 统计结果
successful = sum(1 for r in results if isinstance(r, dict) and r.get("success", False))
print(f"✅ 负载测试完成: {successful}/{len(tasks)} 成功")
# 等待监控完成
await monitor_task
# 生成报告
monitor.generate_summary()
monitor.save_metrics()
async def upload_document(session: aiohttp.ClientSession, content: str, filename: str):
"""上传文档"""
import tempfile
import os
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False, encoding='utf-8') as f:
f.write(content)
temp_path = f.name
try:
with open(temp_path, 'rb') as f:
data = aiohttp.FormData()
data.add_field('file', f, filename=filename, content_type='text/plain')
async with session.post("http://localhost:8000/api/upload", data=data) as response:
return {
"success": response.status == 200,
"type": "upload",
"filename": filename
}
except Exception as e:
return {"success": False, "type": "upload", "error": str(e)}
finally:
if os.path.exists(temp_path):
os.unlink(temp_path)
async def chat_query(session: aiohttp.ClientSession, question: str):
"""聊天查询"""
try:
payload = {"question": question, "top_k": 3, "temperature": 0.7}
async with session.post(
"http://localhost:8000/api/chat",
json=payload,
headers={"Content-Type": "application/json"}
) as response:
return {
"success": response.status == 200,
"type": "chat",
"question": question
}
except Exception as e:
return {"success": False, "type": "chat", "error": str(e)}
if __name__ == "__main__":
try:
asyncio.run(run_load_test_with_monitoring())
except KeyboardInterrupt:
print("\n⏹️ 监控被中断")
except Exception as e:
print(f"❌ 监控失败: {e}")
import traceback
traceback.print_exc()