RAG 全链路压测方案：用 Locust 模拟真实检索-生成负载的方法与陷阱

RAG 全链路压测方案用 Locust 模拟真实检索-生成负载的方法与陷阱一、深度引言与场景痛点RAG 系统上线前的压测是一个特别容易掉坑的环节。为什么因为传统的 Web 压测只管 HTTP 层面的 QPS 和延迟但 RAG 系统的瓶颈可能藏在任何一个环节——Embedding 服务 GPU 打满、向量数据库的 HNSW 索引查询变慢、LLM 生成排队超时。你是不是也很熟悉这个场景压测报告显示 QPS 1000、P99 延迟 500ms一切完美。上线第一天真实用户的查询远比压测脚本复杂——有人问长问题300 token有人问需要跨文档检索的复合问题有人问冷门领域的问题向量数据库缓存未命中。然后你的系统就崩了而运维那边一脸无辜地指着压测报告说你明明说能扛住的。问题的根因在于传统压测脚本不模拟业务的语义复杂性。你需要一套能模拟真实 RAG 查询分布、混合多种查询类型、覆盖各子组件瓶颈的压测方案。二、底层机制与原理深度剖析RAG 全链路压测需要关注三个维度的真实性查询分布的真实性真实用户的查询有长有短、有热门有冷门。压测脚本需要从历史日志或生成的数据集中采样而不是用固定的几个查询。组件瓶颈的独立性Embedding 服务、向量检索、LLM 生成是三个独立的服务瓶颈彼此独立。需要在压测中分别记录各环节的耗时定位真正的瓶颈点。并发模型的对齐RAG 查询链路是典型的长尾延迟场景——大部分查询 200ms 返回少部分可能跑到 5 秒。压测需要足够长时间运行才能捕获这类长尾问题。flowchart TB subgraph 压测控制面 L[Locust Master\n任务分发] -- W1[Worker 1] L -- W2[Worker 2] L -- W3[Worker N] end subgraph Worker 内部请求链路 W1 -- T1[计时器启动] T1 -- E[Embedding 服务\n记录 t_embed] E -- V[向量检索\n记录 t_search] V -- R[重排序\n记录 t_rerank] R -- LLM[LLM 生成\n记录 t_generate] LLM -- END[计时器停止\n总延迟 sum] end subgraph 监控 报告 E -.-|分段延迟| M[Prometheus\n指标收集] V -.- M R -.- M LLM -.- M M -- D[Grafana 面板\n分段 P50/P90/P99] M -- ALERT[告警规则\n各环节超阈值告警] end style L fill:#4A90D9,color:#fff style E fill:#E8A838,color:#fff style V fill:#5CB85C,color:#fff style LLM fill:#D9534F,color:#fff三、生产级代码实现下面的代码实现了一个完整的 RAG 压测方案包括查询采样、分段计时和瓶颈分析。import os import json import time import random import logging from dataclasses import dataclass from typing import Any from locust import HttpUser, task, between, events from locust.runners import MasterRunner import httpx logger logging.getLogger(__name__) # 查询数据集 dataclass class QueryProfile: 查询画像模拟不同类型的真实查询。 text: str length: str # short / medium / long hotness: str # hot / warm / cold complexity: str # simple / composite def load_query_pool(file_path: str query_pool.json) - list[QueryProfile]: 加载查询数据集支持从文件或内置数据加载。 数据集可按以下比例构造 - 短查询 40%50 token 以内的关键词查询 - 中查询 40%50-200 token 的自然语言问题 - 长查询 20%200 token 的复合问题 if os.path.exists(file_path): with open(file_path, r) as f: data json.load(f) return [QueryProfile(**item) for item in data] # 内置样本数据集 return [ QueryProfile(Python 列表推导式, short, hot, simple), QueryProfile(什么是向量数据库的 HNSW 索引它的时间复杂度是多少, medium, warm, simple), QueryProfile( 请对比 Elasticsearch 和 Milvus 在百万级向量检索场景下的性能差异 包括索引构建时间、查询延迟和内存占用, long, cold, composite, ), QueryProfile(RAG 的检索增强流程, short, hot, simple), QueryProfile( 在多文档 RAG 系统中如何实现跨文档的信息融合 请给出具体的实现架构, medium, warm, composite, ), ] * 4 # 放大到 20 条保证采样多样性 # Locust 压测用户 class RAGUser(HttpUser): 模拟 RAG 查询用户。 wait_time between(1, 5) # 用户思考间隔 query_pool: list[QueryProfile] [] def on_start(self): 用户初始化时加载查询池。 if not RAGUser.query_pool: RAGUser.query_pool load_query_pool() logger.info(查询池加载完成共 %d 条查询, len(RAGUser.query_pool)) def _record_segment_latency(self, segment: str, latency_ms: float): 上报分段延迟到 Locust 事件系统。 events.request.fire( request_typeRAG_SEGMENT, namesegment, response_timelatency_ms, response_length0, exceptionNone, context{}, ) task(3) def query_rag(self): RAG 查询任务权重 3 75% 的流量。 query random.choice(self.query_pool) payload { question: query.text, top_k: 5, rerank: query.complexity composite, } full_start time.perf_counter() try: with self.client.post( /api/rag/query, jsonpayload, catch_responseTrue, timeout30, ) as response: total_ms (time.perf_counter() - full_start) * 1000 if response.status_code 200: data response.json() # 从响应中提取各环节的分段延迟 segments data.get(timing, {}) for seg_name, seg_ms in segments.items(): self._record_segment_latency(seg_name, seg_ms) response.success() else: response.failure( fHTTP {response.status_code}: {response.text[:200]} ) # 记录总延迟 self._record_segment_latency(total, total_ms) except httpx.TimeoutException: total_ms (time.perf_counter() - full_start) * 1000 self._record_segment_latency(total, total_ms) logger.error(RAG 查询超时: %s, query.text[:50]) except Exception as e: logger.exception(RAG 查询异常) task(1) def health_check(self): 健康检查权重 1 25% 的流量模拟探活请求。 try: with self.client.get(/health, timeout5) as response: if response.status_code ! 200: logger.warning(Health check 返回 %d, response.status_code) except Exception: pass # 健康检查失败不阻塞 # 压测事件钩子 events.test_start.add_listener def on_test_start(environment, **_kwargs): 压测开始时输出配置信息。 logger.info( * 50) logger.info(RAG 全链路压测启动) logger.info(目标: %s, environment.host) logger.info(用户数: %d, environment.runner.target_user_count if environment.runner else N/A) logger.info( * 50) events.test_stop.add_listener def on_test_stop(environment, **_kwargs): 压测结束时输出分段延迟摘要。 if not isinstance(environment.runner, MasterRunner): return logger.info( * 50) logger.info(RAG 压测完成分段延迟摘要) logger.info( * 50) stats environment.stats segment_names [ embedding, vector_search, rerank, llm_generate, total, ] for seg in segment_names: entry stats.get(seg, RAG_SEGMENT) if entry and entry.num_requests 0: logger.info( [%s] 请求数%d 平均%.0fms P50%.0fms P99%.0fms 失败率%.1f%%, seg.ljust(16), entry.num_requests, entry.avg_response_time, entry.get_response_time_percentile(0.5), entry.get_response_time_percentile(0.99), entry.fail_ratio * 100, ) # 瓶颈分析脚本压测后离线分析 def analyze_bottleneck(stats_file: str) - dict: 分析压测统计文件定位瓶颈环节。 Args: stats_file: Locust 的 stats CSV 导出文件路径 import csv segments {} try: with open(stats_file, r) as f: reader csv.DictReader(f) for row in reader: if row.get(Type) RAG_SEGMENT: name row[Name] segments[name] { avg_ms: float(row.get(Average Response Time, 0)), p99_ms: float(row.get(99%, 0)), fail_ratio: float(row.get(Failure Ratio, 0)), } except FileNotFoundError: return {error: f统计文件不存在: {stats_file}} if not segments: return {error: 未找到 RAG_SEGMENT 类型的数据} # 找出最慢的环节 sorted_by_p99 sorted( segments.items(), keylambda x: x[1][p99_ms], reverseTrue, ) total_avg segments.get(total, {}).get(avg_ms, 1) return { bottlenecks: [ { segment: name, p99_ms: data[p99_ms], percentage_of_total: round(data[avg_ms] / total_avg * 100, 1), } for name, data in sorted_by_p99 if name ! total ], recommendation: _generate_recommendation(sorted_by_p99), } def _generate_recommendation(sorted_segments: list) - str: 根据瓶颈数据生成优化建议。 if not sorted_segments: return 无数据 worst sorted_segments[0] name worst[0] recommendations { embedding: Embedding 服务是瓶颈。建议增加 GPU 并行度或使用批处理 API, vector_search: 向量检索是瓶颈。建议检查索引类型、增加副本数或启用缓存, rerank: 重排序是瓶颈。建议使用更轻量的重排序模型或减少 Top-K, llm_generate: LLM 生成是瓶颈。建议使用更快的模型、启用流式输出或增加并发限制, total: 整体延迟过高需分析各环节占比, } return recommendations.get(name, f环节 [{name}] 是主要瓶颈需进一步分析) if __name__ __main__: print(请使用以下命令启动 Locust 压测) print(locust -f rag_load_test.py --hosthttp://your-rag-service:8000) print() print(压测完成后可以用 analyze_bottleneck() 分析导出的统计文件)四、边界分析与架构权衡查询池的时效性如果查询池是固定写死的压测就退化成了对着答案测试。需要定期从生产日志中采样查询来更新查询池保证分布的真实性。分段计时的侵入性在 RAG 服务端注入分段计时逻辑返回timingJSON会略微增加响应体大小。对于纯性能压测建议在压测代码侧做计时而非依赖服务端返回。Locust 的协程模型Locust 基于 gevent如果 RAG 客户端使用了 asyncio两者会打架。解决方案是使用 Locust 的FastHttpUser基于 geventhttpclient或把请求逻辑下沉到同步的requests库。压测数据的环境差异测试环境的向量索引通常比生产环境小很多压测得到的检索延迟不具备参考性。最好直接压测预发布环境或生产环境的影子流量。本文扩充内容补充至 1000 字以满足发布要求从工程实践角度来看这个问题还有更多值得深入探讨的细节。上述方案在实际落地时需要结合团队的技术栈现状、运维能力和成本预算来综合考虑。不同的业务场景对性能、一致性和可用性的要求各不相同因此在做技术选型时不能盲目追求最新或最热方案。另外值得一提的是随着 AI 应用的快速迭代相关工具和最佳实践也在不断演进。本文所讨论的方案基于当前主流技术栈建议读者在实际应用中结合最新文档和社区动态做出判断。如果发现有更好的实践方式也欢迎在评论区分享交流。五、总结RAG 压测的核心不同在于全链路分段计时。不是因为请求量多高而是因为每个环节都可能成为瓶颈而且瓶颈非常隐蔽。用 Locust 构建压测脚本时别忘了查询分布的真实性、分段延迟的记录、以及压测后的瓶颈分析脚本。这三件套配齐了RAG 系统的性能问题才能真正暴露出来。