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文件路径队列\nmaxsize=1000"] end subgraph "处理层 — 消费者 + 生产者" Q1 --> PARSER["文档解析器\nThreadPoolExecutor"] PARSER --> CLEAN["文本清洗器\nThreadPoolExecutor"] CLEAN --> SPLITTER["文本切分器"] SPLITTER --> Q2["asyncio.Queue\n文本块队列\nmaxsize=5000"] 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( maxsize=self.cfg.file_queue_size, ) chunk_queue: asyncio.Queue[TextChunk] = asyncio.Queue( maxsize=self.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_path=str(file_path), file_type=file_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(encoding="utf-8") except UnicodeDecodeError: text = Path(task.file_path).read_text(encoding="gbk", errors="ignore") # 简单切分（生产环境用 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( text=chunk_text, chunk_index=i // chunk_size, source_file=task.file_path, metadata=task.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(), timeout=5.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( # vectors=embeddings, # 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("文本块队列深度过高: %d，Embedding 可能是瓶颈", chunk_depth) # ===== 使用示例 ===== async def demo(): import tempfile # 创建测试文件 test_files = [] for i in range(10): with tempfile.NamedTemporaryFile( mode="w", suffix=".txt", delete=False, ) as f: f.write(f"这是测试文档 {i} 的内容。" * 100) test_files.append(f.name) config = PipelineConfig( file_queue_size=100, chunk_queue_size=500, embed_concurrency=5, parser_workers=2, embed_batch_size=10, ) 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 就跑得很稳了。
