说真的,这次翻车来得猝不及防,但回过头看又觉得特别“活该”——因为我从一开始就没认真做过容量测试,只是在几条本地文档上跑通就觉得万事大吉了。我想把整个过程还原出来,包括我的思考、误判,以及那种“内部小工具突然被推上风口浪尖”的狼狈。希望能给同样在中小厂、要自己兼任测试和运维的兄弟们提个醒。
前言:
我在做一个自己使用的 RAG 知识库应用,最初的打算就是最多几个人用,帮忙检索技术文档和日常的学习笔记。为了弥补单纯向量检索在“专有名词匹配”上的不足,我决定引入 BM25 算法来做混合检索。
考虑到 Python 生态的便利性,我选择了非常轻量级的rank_bm25库(具体用的是BM25Okapi)。在本地用几个测试文档跑的时候,效果出奇的好,关键词召回率大幅提升。既然是内部小工具,我心里就没绷着“生产环境”那根弦,觉得功能跑通、结果好看就完事了,什么压力测试、资源预估基本是空白。
结果没想到,,一位带课的老师觉得这个工具挺适合学生期末复习,在课堂上随口推荐了一下。期末考试前一晚,几十号学生突然涌进来,疯狂上传课件、笔记、历年真题。更要命的是,很多文件其实是一样的——同一本教材的电子版、同一份习题集,你传一份我传一份。而我们的用户隔离做得极其简单,仅仅是上传时在数据库里加个user_id字段,完全没有做任何全局去重。
于是服务从“岁月静好”直接进入“垂死挣扎”模式,接口超时、内存告警,最后进程被 OOM killer 一刀带走。复盘下来,元凶并非网络 I/O 或者大文件上传,而是rank_bm25不支持增量更新的致命缺陷,以及我在内存数据结构上对资源消耗的预估不足。
架构还原:
BM25Okapi本身只是一个纯粹的算法实现对象,它基于全量语料库预计算了 IDF(逆文档频率)、文档长度等统计信息,且不支持动态插入新文档。为了让系统运转起来,我不得不在内存中为每个用户维护三个核心数据结构:
# 1. BM25 索引对象,检索时直接用 _bm25_indices: Dict[str, BM25Okapi] 2. 原始文档分块列表,用来根据得分找回文本 _bm25_corpus: Dict[str, List[DocumentChunk]] 3. 分词后的语料,重建索引时避免重复分词 _bm25_tokenized: Dict[str, List[List[str]]]为什么要这样冗余?因为BM25Okapi.get_scores()只返回一个分数数组,我必须用分数的索引去原始分块列表里捞对应的文本,才能喂给 LLM。而保存分词结果,是为了重建索引时不用再把所有原始文本重新切一遍词——这里其实已经暴露了我对性能的焦虑,但当时的想法只是“省点儿 CPU 算力”。
全量重建与重复文件带来的问题
因为不支持增量更新,我的更新逻辑是这样的:每次用户上传新文件或删除文件,都要将该用户的索引推倒重来。
async def _rebuild_bm25_index(self, user_id: str) -> None: # 每次调用都从 ChromaDB 获取该用户的所有文档 results = await asyncio.to_thread( self._collection.get, where={"user_id": user_id} ) # 然后重新分词、重新构建索引在没有大量用户和文档的情况下,这个架构毫无问题。但当老师推荐后,几十个学生同时操作,画风就彻底变了。
1. O(n) 的时间复杂度
很多学生一股脑上传了数百兆的课件包,我们的文档分块策略又切得比较细,一个用户的文档分块数轻轻松松破千。这时候,哪怕他只是新增了一小段笔记,系统也要从数据库捞出之前的上千条分块,重新跑一遍分词和 BM25 预计算。
这个操作的耗时从毫秒级直接膨胀到秒级,直接堵住了上传接口。更别提几十个学生在同时操作,大量这样的全量重建任务并行堆叠,请求链路一个接一个超时,整个服务几乎不可用。
2. 每用户 14MB 堆叠
上述的三重数据结构,实质上是在内存中保存了文档的三份拷贝(原始文本、分词结果、统计模型)。让我们算一笔账:
如果只是几个用户,这点内存完全不算什么。但几十个学生同时在线,并且大量上传了相同的文件——因为用户隔离只是简单的user_id,同一本教材的电子版在不同用户下被分别存储、分别构建索引,各自独立扛着一份完整的内存副本。假设同时有 70 个活跃学生在查询或上传,单纯维护 BM25 的字典就会吃掉将近 1 GB 内存。而我们的服务器只是个普普通通的 4 GB 虚拟机,这还不包括 FastAPI 框架、ChromaDB 客户端以及并发请求本身的消耗。物理内存很快被打满,OOM 随之而来。
反思:个人开发和实际场景的偏离
这次事故给我最大的教训,不是代码写得烂,而是真实的使用场景和个人的开发测试的视角不一致。
在中小团队,开发、测试、运维常常是同一个人,我们太容容易忽略真实上线时可能出现的流量高峰和用户计数,我们时测试时,很可能为了方便,只模拟一两个用户、每人几份文档,觉得接口返回正常就算过了。根本没有去想:万一突然涌进来五十个用户呢?万一他们传的都是几百页的 PDF 呢?万一大量文件是重复的呢?
这里老师的一句推荐,就是一个从天而降的压力测试场景,而我的系统毫无招架之力。如果我能提前用脚本模拟这个量级的数据和并发——比如创建几十个用户,每个用户塞进去 800–1000 条分块,再混入大量重复文件,然后并发上传和检索——这个 OOM 的问题会立刻暴露,根本轮不到真实场景来给我们上课。
我们很容易高估自己代码的健壮性,因为平时开发看到的数据量太小了。几个文档跑得好好的,不代表一千个文档、几十个用户同时操作时还能站得住。没有数据压力和流量峰值的测试,本质上就是没测试。
问题修复
这次事故让我彻底放弃了“用 Python 字典手撸搜索引擎”的念头。不是说rank_bm25不好,而是它本来就不是为这种场景设计的。后续的优化方向很明确:
- 把 BM25 交给真正的检索引擎:比如 Elasticsearch 或 Meilisearch,它们原生支持 BM25,并且具备增量倒排索引更新和磁盘持久化。把这块逻辑从应用进程里剥离出去,内存压力和重建开销都能解决。
这次踩坑让我明白一件事:在 AI 应用开发里,我们太容易被 Prompt、模型效果这些“耀眼”的东西吸引,而忽略了基础的工程原则。再好的算法,如果扛不住真实流量,就是纸上谈兵。而对我们这种小团队来说,自己写的程序,自己就得是它最严厉的测试者——测试不光是看功能对不对,更要看它在数据和并发面前会不会倒下。