当前位置：首页 > news >正文

Python 协程池设计模式：在 RAG 场景下复用 Embedding 请求的最佳实践

news 2026/7/8 14:38:03

Python 协程池设计模式：在 RAG 场景下复用 Embedding 请求的最佳实践

一、深度引言与场景痛点

RAG 系统的 Embedding 调用有一个非常典型的特点：短请求、高频次、有状态。每次调用只需要几十毫秒，但调用的频次可以高到每秒几百次，而且调用方通常需要在多次请求之间复用同一个 HTTP 连接（否则每次新建连接的开销会把延迟吃掉一大半）。

一个很常见的反模式是：用aiohttp.ClientSession在每次 embedding 调用时创建和销毁。伪代码大概是这样的：

async def embed(text): async with aiohttp.ClientSession() as s: async with s.post("https://api.embedding.com/v1/embed", json={"text": text}) as resp: return await resp.json()

这在 QPS=10 时没问题，但当 QPS 升到 100+ 时，每个请求创建新 Session 意味着新建 TCP 连接、TLS 握手，每次额外开销 50-200ms。你 embedding 调用本身才 80ms，一半的时间都在建连接。

协程连接池就是用来解决这个问题的。它的核心思想是：把 HTTP 连接视为一种可复用的资源，用连接池管理它的生命周期。Python 的 asyncio 生态里已经有很多成熟的连接池实现（比如 aiohttp 的TCPConnector），但直接裸用它们还有一些工程细节需要注意——比如连接池的大小设置、空闲连接的回收、异常连接的自动剔除等。

本文就来深入剖析在 RAG 场景下，如何设计一个生产级的协程连接池，让 Embedding 请求的延迟和资源消耗降到最低。

二、底层机制与原理深度剖析

flowchart TD A[Embedding 请求到达] --> B{连接池} subgraph Pool [协程连接池架构] B -->|有空闲连接| C1[直接复用] B -->|无空闲但未达上限| C2[创建新连接] B -->|已达上限| C3[排队等待] C1 --> D[发送 Embedding 请求] C2 --> D C3 -->|等待| C4{超时?} C4 -->|否| B C4 -->|是| C5[抛出 TimeoutError] D --> E{请求结果} E -->|成功| F1[连接标记为健康] E -->|连接错误| F2[连接标记为异常] E -->|超时| F3[连接标记为慢连接] F1 --> G[连接回池] F2 --> H[连接销毁] F3 --> I{重试决策} I -->|重试| C2 I -->|放弃| H end subgraph HealthCheck [健康检查线程] J[定时扫描] --> K{空闲超时?} K -->|是| H K -->|否| L{连接异常?} L -->|是| H L -->|否| M[保持] end style Pool fill:#e8f5e9

连接池的关键机制有三层：

第一层：获取连接。请求到达时先从池中取空闲连接。有空闲的——直接拿；没有但还能建新的（没到上限）——建一个；没有且到上限了——排队等。排队不是无限等待，有超时机制。

第二层：返还连接。请求完成后，根据连接的健康状态决定命运。成功的连接标记为健康、回池复用；连接级别的错误（如 ConnectionReset）直接销毁连接；超时的请求会让连接背负"慢连接"标签，通常也不回池——因为它可能已经是半死不活的状态了。

第三层：后台维护。一个独立的定时任务扫描池中的空闲连接，销毁长时间不活跃的连接（节约资源）和状态异常但未被标记的连接。

三、生产级代码实现

import asyncio import time import logging from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple from collections import deque import aiohttp logger = logging.getLogger(__name__) @dataclass class PoolConfig: """连接池配置""" # 连接数限制 max_connections: int = 50 min_connections: int = 5 # 最小保活连接数 # 超时控制 acquire_timeout: float = 10.0 # 获取连接的最大等待时间 request_timeout: float = 30.0 # 单次请求的最大时间 idle_timeout: float = 300.0 # 空闲连接的最大存活时间（秒） # 健康检查 health_check_interval: float = 60.0 # 健康检查间隔（秒） max_retries: int = 2 # 单次请求的最大重试次数 # 连接参数 tcp_keepalive: bool = True tcp_nodelay: bool = True @dataclass class ConnectionState: """单个连接的状态""" session: aiohttp.ClientSession created_at: float = field(default_factory=time.monotonic) last_used_at: float = field(default_factory=time.monotonic) request_count: int = 0 error_count: int = 0 is_busy: bool = False @property def idle_seconds(self) -> float: return time.monotonic() - self.last_used_at @property def is_healthy(self) -> bool: # 连续失败 3 次的连接视为不健康 return self.error_count < 3 class EmbeddingConnectionPool: """Embedding 服务的协程连接池 核心特性： 1. 连接复用（减少 TCP/TLS 握手开销） 2. 连接数限制（保护下游服务） 3. 自动健康检查（剔除异常连接） 4. 空闲回收（释放不活跃连接） """ def __init__(self, config: Optional[PoolConfig] = None): self.config = config or PoolConfig() self._idle: deque[ConnectionState] = deque() self._busy: Dict[int, ConnectionState] = {} self._available = asyncio.Semaphore(self.config.max_connections) self._health_check_task: Optional[asyncio.Task] = None self._shutdown = False # 监控指标 self.metrics = { "acquired": 0, "released": 0, "created": 0, "destroyed": 0, "timeouts": 0, "errors": 0, } async def start(self): """启动连接池（预创建最小连接数 + 启动健康检查）""" # 预创建最小连接数 for _ in range(self.config.min_connections): conn = await self._create_connection() self._idle.append(conn) self.metrics["created"] += 1 # 启动后台健康检查 self._health_check_task = asyncio.create_task(self._health_check_loop()) logger.info(f"连接池启动: {len(self._idle)} 空闲连接") async def shutdown(self): """关闭连接池""" self._shutdown = True if self._health_check_task: self._health_check_task.cancel() try: await self._health_check_task except asyncio.CancelledError: pass # 关闭所有空闲连接 while self._idle: conn = self._idle.popleft() await conn.session.close() self.metrics["destroyed"] += 1 # 强制关闭忙碌连接（生产环境应等待一定时间） for conn in self._busy.values(): await conn.session.close() self.metrics["destroyed"] += 1 self._busy.clear() logger.info("连接池已关闭") async def _create_connection(self) -> ConnectionState: """创建新连接""" connector = aiohttp.TCPConnector( limit=0, # 不限制单个 connector 的连接数（connector 属于单连接的） ttl_dns_cache=300, force_close=True, # 确保连接可以被正确复用 ) timeout = aiohttp.ClientTimeout(total=self.config.request_timeout) session = aiohttp.ClientSession( connector=connector, timeout=timeout, ) return ConnectionState(session=session) async def acquire(self) -> ConnectionState: """获取一个可用连接（阻塞直到成功或超时）""" start = time.monotonic() try: # 等待信号量（控制总连接数） await asyncio.wait_for( self._available.acquire(), timeout=self.config.acquire_timeout, ) except asyncio.TimeoutError: self.metrics["timeouts"] += 1 raise asyncio.TimeoutError( f"获取连接超时 ({self.config.acquire_timeout}s)，" f"当前连接数: {len(self._idle) + len(self._busy)}" ) wait_time = time.monotonic() - start if wait_time > 1.0: logger.warning(f"连接获取耗时较长: {wait_time:.2f}s") # 尝试从空闲池取连接 while self._idle: conn = self._idle.popleft() if conn.is_healthy: conn.is_busy = True conn.last_used_at = time.monotonic() conn.request_count += 1 self._busy[id(conn)] = conn self.metrics["acquired"] += 1 return conn else: # 不健康的连接直接销毁 await conn.session.close() self.metrics["destroyed"] += 1 # 空闲池空了，创建新连接 conn = await self._create_connection() conn.is_busy = True self._busy[id(conn)] = conn self.metrics["created"] += 1 self.metrics["acquired"] += 1 return conn def release(self, conn: ConnectionState): """释放连接回池""" if id(conn) in self._busy: del self._busy[id(conn)] conn.is_busy = False conn.last_used_at = time.monotonic() conn.error_count = 0 # 成功释放时重置错误计数 if conn.is_healthy: self._idle.append(conn) else: # 不健康的连接不回池 asyncio.create_task(conn.session.close()) self.metrics["destroyed"] += 1 self._available.release() self.metrics["released"] += 1 def mark_error(self, conn: ConnectionState): """标记连接发生错误""" conn.error_count += 1 self.metrics["errors"] += 1 async def _health_check_loop(self): """后台健康检查循环""" while not self._shutdown: await asyncio.sleep(self.config.health_check_interval) # 清理长时间空闲的连接 to_remove = [] while self._idle: conn = self._idle.popleft() if conn.idle_seconds > self.config.idle_timeout: await conn.session.close() self.metrics["destroyed"] += 1 to_remove.append(conn) else: # 健康的连接放回去 self._idle.appendleft(conn) break if to_remove: logger.info(f"健康检查: 清理 {len(to_remove)} 个空闲连接") def get_stats(self) -> Dict[str, Any]: """获取连接池统计信息""" return { **self.metrics, "idle_count": len(self._idle), "busy_count": len(self._busy), "total_count": len(self._idle) + len(self._busy), "semaphore_waiting": ( self.config.max_connections - self._available._value ), } # ================== Embedding 服务封装 ================== class EmbeddingService: """基于连接池的 Embedding 服务""" def __init__(self, pool: EmbeddingConnectionPool, base_url: str): self.pool = pool self.base_url = base_url.rstrip("/") self.embed_endpoint = f"{self.base_url}/v1/embeddings" async def embed(self, text: str, retries: Optional[int] = None) -> List[float]: """对单条文本生成 Embedding 使用连接池复用 HTTP 连接，支持自动重试 """ max_retries = retries if retries is not None else self.pool.config.max_retries for attempt in range(max_retries + 1): conn = None try: conn = await self.pool.acquire() async with conn.session.post( self.embed_endpoint, json={"input": text, "model": "text-embedding-3-small"}, ) as resp: if resp.status == 200: data = await resp.json() embedding = data["data"][0]["embedding"] self.pool.release(conn) return embedding elif resp.status == 429: # 速率限制，等待后重试 retry_after = int(resp.headers.get("Retry-After", 2)) self.pool.release(conn) logger.warning(f"速率限制，{retry_after}s 后重试 (第 {attempt+1}/{max_retries+1} 次)") await asyncio.sleep(retry_after) continue else: body = await resp.text() raise aiohttp.ClientResponseError( resp.request_info, resp.history, status=resp.status, message=body ) except (aiohttp.ClientConnectionError, asyncio.TimeoutError) as e: if conn: self.pool.mark_error(conn) self.pool.release(conn) if attempt < max_retries: wait = 2 ** attempt # 指数退避 logger.warning(f"连接错误，{wait}s 后重试: {e}") await asyncio.sleep(wait) else: raise except Exception as e: if conn: self.pool.mark_error(conn) self.pool.release(conn) raise async def embed_batch(self, texts: List[str]) -> List[List[float]]: """批量生成 Embedding（使用连接池并发）""" tasks = [self.embed(text) for text in texts] results = await asyncio.gather(*tasks, return_exceptions=True) embeddings = [] for i, result in enumerate(results): if isinstance(result, Exception): logger.error(f"批量 Embedding 失败 [{i}/{len(texts)}]: {result}") embeddings.append([]) # 占位 else: embeddings.append(result) return embeddings # ================== 使用示例 ================== async def main(): pool = EmbeddingConnectionPool( PoolConfig( max_connections=50, min_connections=5, acquire_timeout=5.0, request_timeout=10.0, idle_timeout=120.0, ) ) await pool.start() service = EmbeddingService(pool, "https://api.openai.com") # 单条 embedding vec = await service.embed("Python asyncio 协程池最佳实践") print(f"向量维度: {len(vec)}") # 批量 embedding texts = [f"测试文本 {i}" for i in range(100)] start = time.monotonic() vectors = await service.embed_batch(texts) elapsed = time.monotonic() - start print(f"批量 {len(texts)} 条 embedding: {elapsed:.2f}s") print(f"连接池统计: {pool.get_stats()}") await pool.shutdown() if __name__ == "__main__": asyncio.run(main())

四、边界分析与架构权衡

1. 连接池大小怎么定？

大原则是：不超过下游服务的连接上限，也不超过系统能承受的并发连接数。一个估算公式：

pool_size ≈ 目标 QPS × avg_request_latency

举例：目标 QPS = 200，单次请求延迟约 0.1s，那么pool_size ≈ 200 × 0.1 = 20。但这个公式假设延迟是常数，实际中要留 30-50% 的余量。建议初始值设 30，然后通过压测逐步调优。

2. 连接池是否也需要背压？

需要。信号量_available本身就是一种背压——当连接数达到上限时，新请求必须排队等待。但信号量只限制了连接获取，没有限制请求的排队长度。如果短时间内涌入大量请求而又没有响应超时机制，排队的协程会越积越多，最终耗尽内存。对策是配合全局限制（类似上一篇文章的全局信号量）给请求本身设置上限。

3. 空闲连接回收的"惊群"问题

如果健康检查清掉了 10 个空闲连接，紧接着一批新请求涌入，它们会同时尝试创建新连接。如果创建连接很慢（涉及 DNS 解析和 TLS 握手），会导致瞬时延迟尖峰。解决方案是预热——健康检查在清掉连接前，先确保池中仍然有min_connections个连接可用。

4. 多服务场景下的连接池隔离

如果你的 RAG 系统调用了多个不同的 API（Embedding API、Rerank API、LLM API），它们各自的延迟和并发特性完全不同。应该为每个服务单独建一个连接池，而不是共用一个。因为一个 LLM 请求可能耗时 5 秒占据一个连接，如果它和 Embedding 请求共享连接池（60 个连接），那么长尾的 LLM 请求会挤占 Embedding 的"快车道"。