Python 生产者-消费者模式：在 RAG 文档处理 pipeline 中的并发设计

Python 生产者-消费者模式在 RAG 文档处理 pipeline 中的并发设计一、深度引言与场景痛点RAG 系统的第一步永远是文档处理——把 PDF、Word、网页之类的东西变成向量化的文本块。这听起来是个简单的读文件、切文本、生成向量三步走但当你面对百万级文档时它就变成了一个系统工程问题。单线程处理百万文档的逻辑是读一个文件I/O 等待→ 切文本CPU 密集→ 生成向量网络 I/O 等待→ 存数据库网络 I/O 等待。每一步都在等CPU 大部分时间在摸鱼。更糟糕的是 I/O 密集型操作读文件、调 Embedding API和 CPU 密集型操作文本切分混合在一起你没法简单地说开多线程或开多进程就能解决。正确的策略是生产者-消费者模式把整个 pipeline 拆成独立的生产者和消费者中间用队列连接。生产者负责读文件I/O 密集用协程消费者负责切文本和生成向量混合型用线程池/进程池。队列充当天然的背压机制——当消费者处理不过来时生产者自动减速。二、底层机制与原理深度剖析生产者-消费者模式在 RAG 文档处理 pipeline 中的架构分为三层数据源层生产者负责从文件系统、S3、数据库等源头读取待处理文档。这层是纯 I/O 密集用asyncio的协程处理最合适。可以按批次读取每次往队列里推送一个批次的文件路径。处理层消费者接收文件路径依次执行解析文档 → 文本清洗 → 文本切分。这层混合了 I/O读文件和 CPU切分处理适合用线程池。如果切分逻辑特别复杂如大模型的 Tokenizer可以考虑进程池。向量化层消费者接收文本块调用 Embedding API 生成向量并存入向量数据库。纯网络 I/O 密集用异步协程处理配合并发数限制防止打爆 Embedding 服务。队列层Python 的asyncio.Queue提供异步队列天然支持背压。队列满时put自动等待队列空时get自动等待。可以通过maxsize限制队列长度防止内存溢出。flowchart TB subgraph 数据源层 — 生产者 FS[文件系统/S3] -- READER[Async 文档读取器\n(纯 I/O协程)] READER -- Q1[asyncio.Queue\n文件路径队列\nmaxsize1000] end subgraph 处理层 — 消费者 生产者 Q1 -- PARSER[文档解析器\nThreadPoolExecutor] PARSER -- CLEAN[文本清洗器\nThreadPoolExecutor] CLEAN -- SPLITTER[文本切分器] SPLITTER -- Q2[asyncio.Queue\n文本块队列\nmaxsize5000] end subgraph 向量化层 — 消费者 Q2 -- EMBED[Embedding API\nAsync 协程 信号量] EMBED -- Q3[asyncio.Queue\n结果队列] Q3 -- DB[向量数据库写入\nBatch 批量写入] end subgraph 监控 Q1 -.-|监控队列深度| MON[背压监控\n队列满→降低生产速率] Q2 -.-|监控队列深度| MON end style READER fill:#4A90D9,color:#fff style PARSER fill:#5CB85C,color:#fff style EMBED fill:#E8A838,color:#fff style DB fill:#D9534F,color:#fff三、生产级代码实现下面的代码实现了一个完整的 RAG 文档处理 pipeline采用多层生产者-消费者模式。关键是每个环节独立、有背压控制、有异常处理。import asyncio import logging import time from dataclasses import dataclass from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor from pathlib import Path from typing import Any logger logging.getLogger(__name__) # 数据模型 dataclass class DocumentTask: 待处理的文档任务。 file_path: str file_type: str pdf metadata: dict[str, Any] | None None dataclass class TextChunk: 处理后的文本块。 text: str chunk_index: int source_file: str metadata: dict[str, Any] | None None dataclass class PipelineConfig: Pipeline 配置。 file_queue_size: int 1000 # 文件队列最大长度 chunk_queue_size: int 5000 # 文本块队列最大长度 embed_concurrency: int 10 # Embedding 并发数 parser_workers: int 4 # 解析线程数 batch_size: int 32 # 批量写入大小 embed_batch_size: int 20 # Embedding 批量请求大小 # Pipeline 实现 class RAGDocumentPipeline: RAG 文档处理 Pipeline —— 生产者-消费者模式。 def __init__(self, config: PipelineConfig | None None): self.cfg config or PipelineConfig() self._stats { files_processed: 0, chunks_generated: 0, embeddings_stored: 0, errors: 0, start_time: 0.0, } async def run(self, file_paths: list[str]): 启动整个 Pipeline。 self._stats[start_time] time.monotonic() logger.info(Pipeline 启动待处理文件数%d, len(file_paths)) # 创建队列 file_queue: asyncio.Queue[DocumentTask] asyncio.Queue( maxsizeself.cfg.file_queue_size, ) chunk_queue: asyncio.Queue[TextChunk] asyncio.Queue( maxsizeself.cfg.chunk_queue_size, ) # 启动所有阶段的协程 async with asyncio.TaskGroup() as tg: # 阶段 1文件读取生产者 tg.create_task(self._read_files(file_paths, file_queue)) # 阶段 2文档解析和切分消费者 生产者 for _ in range(self.cfg.parser_workers): tg.create_task(self._parse_and_split(file_queue, chunk_queue)) # 阶段 3向量化和存储消费者 tg.create_task(self._embed_and_store(chunk_queue)) # 打印统计 elapsed time.monotonic() - self._stats[start_time] logger.info( Pipeline 完成: 文件%d, 文本块%d, 向量%d, 错误%d, 耗时%.1fs, self._stats[files_processed], self._stats[chunks_generated], self._stats[embeddings_stored], self._stats[errors], elapsed, ) async def _read_files( self, file_paths: list[str], file_queue: asyncio.Queue[DocumentTask], ): 阶段 1异步读取文件列表推入队列。 for path_str in file_paths: file_path Path(path_str) if not file_path.exists(): logger.warning(文件不存在跳过: %s, path_str) continue task DocumentTask( file_pathstr(file_path), file_typefile_path.suffix.lstrip(.).lower(), ) try: await file_queue.put(task) except asyncio.QueueFull: logger.warning(文件队列已满等待消费...) await file_queue.put(task) # 自动等待 # 发送终止信号与消费者约定None 表示输入结束 for _ in range(self.cfg.parser_workers): await file_queue.put(None) # type: ignore logger.info(文件读取完成共 %d 个文件入队, len(file_paths)) async def _parse_and_split( self, file_queue: asyncio.Queue[DocumentTask], chunk_queue: asyncio.Queue[TextChunk], ): 阶段 2解析文档并切分文本消费者 生产者。 使用线程池处理 CPU 密集型操作避免阻塞事件循环。 loop asyncio.get_running_loop() while True: task await file_queue.get() # 检查终止信号 if task is None: file_queue.task_done() break try: # 在线程池中执行解析和切分 chunks await loop.run_in_executor( None, self._parse_document_sync, task, ) for chunk in chunks: await chunk_queue.put(chunk) self._stats[files_processed] 1 self._stats[chunks_generated] len(chunks) except Exception as e: self._stats[errors] 1 logger.error(文件处理失败 [%s]: %s, task.file_path, e) finally: file_queue.task_done() logger.info(解析 worker 退出) staticmethod def _parse_document_sync(task: DocumentTask) - list[TextChunk]: 同步的文档解析和切分在线程池中执行。 # 实际实现 # 这里根据 file_type 选择不同的解析器 # - PDF → PyMuPDF / pdfplumber # - DOCX → python-docx # - HTML → BeautifulSoup # - TXT → 直接读取 # 简化演示读取文本文件 try: text Path(task.file_path).read_text(encodingutf-8) except UnicodeDecodeError: text Path(task.file_path).read_text(encodinggbk, errorsignore) # 简单切分生产环境用 RecursiveCharacterTextSplitter chunk_size 500 chunks [] for i in range(0, len(text), chunk_size): chunk_text text[i:i chunk_size] if chunk_text.strip(): chunks.append(TextChunk( textchunk_text, chunk_indexi // chunk_size, source_filetask.file_path, metadatatask.metadata, )) return chunks async def _embed_and_store( self, chunk_queue: asyncio.Queue[TextChunk], ): 阶段 3向量化和存储消费者。 使用信号量控制 Embedding API 的并发数避免打爆服务。 sem asyncio.Semaphore(self.cfg.embed_concurrency) batch: list[TextChunk] [] while True: try: # 批量收集文本块 chunk await asyncio.wait_for( chunk_queue.get(), timeout5.0, ) # 检查终止信号 if chunk is None: chunk_queue.task_done() break batch.append(chunk) if len(batch) self.cfg.embed_batch_size: async with sem: await self._process_batch(batch) batch [] self._stats[embeddings_stored] len(batch) chunk_queue.task_done() except asyncio.TimeoutError: # 队列空了处理剩余批次 if batch: async with sem: await self._process_batch(batch) self._stats[embeddings_stored] len(batch) batch [] async def _process_batch(self, chunks: list[TextChunk]): 批量处理向量化 存入数据库。 texts [c.text for c in chunks] try: # 调用 Embedding API实际调用替换为你的服务 # embeddings await self.embedding_service.embed(texts) # 模拟 Embedding 调用 await asyncio.sleep(0.1) embeddings [[0.0] * 768] * len(texts) # 写入向量数据库 # await self.vector_db.insert( # vectorsembeddings, # metadatas[c.metadata for c in chunks], # ids[f{c.source_file}_{c.chunk_index} for c in chunks], # ) logger.debug(批量写入 %d 条向量, len(chunks)) except Exception as e: self._stats[errors] len(chunks) logger.error(向量批量处理失败: %s, e) # 失败不丢数据——写入死信队列供后续重试 # await self._dead_letter_queue.put(chunks) # 监控辅助 class PipelineMonitor: Pipeline 监控器——背压和进度追踪。 def __init__(self, pipeline: RAGDocumentPipeline, interval: float 5.0): self.pipeline pipeline self.interval interval async def monitor(self, file_queue: asyncio.Queue, chunk_queue: asyncio.Queue): 定期输出 Pipeline 状态。 while True: await asyncio.sleep(self.interval) file_depth file_queue.qsize() chunk_depth chunk_queue.qsize() stats self.pipeline._stats # 计算吞吐 elapsed max(time.monotonic() - stats[start_time], 0.001) files_per_sec stats[files_processed] / elapsed logger.info( Pipeline 状态: 文件队列%d, 块队列%d, 已处理%d文件, %.1f文件/秒, 错误%d, file_depth, chunk_depth, stats[files_processed], files_per_sec, stats[errors], ) # 背压告警 if file_depth 800: logger.warning(文件队列深度过高: %d检测到背压, file_depth) if chunk_depth 4000: logger.warning(文本块队列深度过高: %dEmbedding 可能是瓶颈, chunk_depth) # 使用示例 async def demo(): import tempfile # 创建测试文件 test_files [] for i in range(10): with tempfile.NamedTemporaryFile( modew, suffix.txt, deleteFalse, ) as f: f.write(f这是测试文档 {i} 的内容。 * 100) test_files.append(f.name) config PipelineConfig( file_queue_size100, chunk_queue_size500, embed_concurrency5, parser_workers2, embed_batch_size10, ) pipeline RAGDocumentPipeline(config) # 启动监控 # monitor PipelineMonitor(pipeline) # asyncio.create_task(monitor.monitor(file_queue, chunk_queue)) await pipeline.run(test_files) # 清理测试文件 import os for f in test_files: os.unlink(f) if __name__ __main__: asyncio.run(demo())四、边界分析与架构权衡队列的终止信号设计在多消费者场景下如何优雅地通知所有消费者没有更多数据了上面的方案是每个消费者发一个None作为毒丸。当消费者数量可变时更好的方案是用一个共享的asyncio.Event标记输入已结束消费者检查 Event 和空队列的组合条件。进程池 vs 线程池文本切分如果是重度 CPU 操作如大型 Transformer 的 Tokenizer需要使用进程池而非线程池来绕过 GIL。但进程间通信的开销比线程间大得多需要传输的数据应该尽量轻量——传文件路径而不是文件内容。Embedding 批量大小调节批量过大 单次调用延迟高、失败影响面大批量过小 API 调用次数多、网络开销占比高。建议对 Embedding 服务做 benchmark找到吞吐量最大的批量大小而不是拍脑袋。死信队列的必要性生产环境中Embedding 调用或数据库写入可能偶发失败。需要为失败的数据设计死信队列定期重试而不是直接丢弃。死信队列也应该有 TTL防止永久失败的数据堆积。五、总结生产者-消费者模式在 RAG 文档处理 pipeline 中简直是量身定做的并发模型。文件的 I/O 读取、文本的 CPU 切分、向量的网络写入——三种不同特征的操作用三个独立的环节处理队列自动提供背压和缓冲。关键设计决策是I/O 密集用协程、CPU 密集用进程/线程池、网络 I/O 用协程 信号量控制并发。再加一个监控器观察队列深度Pipeline 就跑得很稳了。