【Pocket Flow】源码剖析（二）：批量与异步——BatchNode、AsyncNode 与并行执行

【Pocket Flow】源码剖析二批量与异步——BatchNode、AsyncNode 与并行执行写在前面第一篇我们拆解了 Pocket Flow 的三大核心抽象——Node、Flow 和 Shared Store理解了 100 行代码的骨架。今天我们进入这 100 行代码中最有料的部分——批量Batch与异步Async体系。Pocket Flow 只有 100 行代码却定义了10 个类其中 6 个与批量或异步相关。这些类不是凑数的——它们用极少的代码覆盖了从逐项串行到并行并发的所有执行模式。理解了它们你就理解了 Pocket Flow 为什么能用 100 行代码实现 MapReduce、RAG 和 Multi-Agent。 文章目录 一、为什么需要批量与异步 二、BatchNode逐项执行的数据并行 三、BatchFlow重跑流程的任务并行⚡ 四、AsyncNode异步 I/O 的优雅封装 五、AsyncFlow兼容同步与异步的编排器 六、并行执行asyncio.gather 的威力 七、系列预告 一、为什么需要批量与异步1.1 两个真实场景场景一RAG 文档处理。你有 10 个 PDF 文件每个需要加载文本 → 分块 → 生成 Embedding → 存入向量库。如果逐个处理每个文件耗时 5 秒总共 50 秒。但如果 10 个文件并行处理理论上只需 5 秒。场景二Multi-Agent 协作。两个 Agent 需要实时对话——一个出题一个猜词。出题 Agent 等待猜词 Agent 的输入猜词 Agent 等待出题 Agent 的提示。这是典型的异步 I/O 场景——如果用同步代码一个 Agent 的input()会阻塞整个线程另一个 Agent 永远得不到执行机会。1.2 Pocket Flow 的解法Pocket Flow 用6 个类覆盖了这两个场景类场景核心机制BatchNode同步逐项处理列表推导[exec(i) for i in items]BatchFlow同步重跑流程for bp in pr: _orch(shared, params)AsyncNode异步 I/Oprep_async / exec_async / post_asyncAsyncFlow异步编排await _orch_async() 兼容同步节点AsyncBatchNode异步逐项[await exec(i) for i in items]AsyncParallelBatchNode异步并行asyncio.gather(*[exec(i) for i in items])这 6 个类的源码加起来不到30 行但覆盖了从串行到并行的完整光谱。 二、BatchNode逐项执行的数据并行2.1 BatchNode 是什么BatchNode 继承自 Node但改变了prep和exec的语义prep(shared)返回一个可迭代对象如列表而不是单个值exec(item)对列表中的每个 item 调用一次而不是对整个 prep 结果调用一次post(shared, prep_res, exec_res_list)接收所有 exec 结果的列表2.2 源码解读classBatchNode(Node):def_exec(self,items):return[super(BatchNode,self)._exec(i)foriin(itemsor[])]就这一行。_exec接收prep返回的列表对每个 item 调用父类 Node 的_exec含重试逻辑收集结果到新列表。items or []处理了 prep 返回 None 的情况——优雅的防御性编程。2.3 实战示例文档分块摘要classMapSummaries(BatchNode):defprep(self,shared):contentshared[data]chunk_size10000return[content[i:ichunk_size]foriinrange(0,len(content),chunk_size)]defexec(self,chunk):promptfSummarize this chunk in 10 words:{chunk}returncall_llm(prompt)defpost(self,shared,prep_res,exec_res_list):shared[summary]\n.join(exec_res_list)returndefault这个模式就是经典的Map阶段——把大任务拆成小任务逐个处理。如果加上一个汇总节点就是完整的MapReduce。2.4 BatchNode 的局限BatchNode 是顺序执行的——[exec(i) for i in items]是列表推导一个接一个。如果每个 item 需要调用 LLM耗时 2 秒10 个 item 就需要 20 秒。要加速需要异步并行——我们第六节再讲。 三、BatchFlow重跑流程的任务并行3.1 BatchFlow 是什么BatchNode 是一个节点处理多个 item。BatchFlow 是一个流程跑多遍每遍参数不同。听起来相似但本质不同维度BatchNodeBatchFlow粒度单个 Node 的 exec整个 Flow 的 _orch输入prep 返回的 item 列表prep 返回的参数字典列表状态写入同一个 shared写入同一个 shared注意 key 冲突子节点不涉及通过 self.params 访问参数适合数据并行同逻辑不同数据任务并行同流程不同参数3.2 源码解读classBatchFlow(Flow):def_run(self,shared):prself.prep(shared)or[]forbpinpr:self._orch(shared,{**self.params,**bp})returnself.post(shared,pr,None)三行代码。prep返回参数字典列表for bp in pr逐组执行{**self.params, **bp}合并全局参数和当前批次参数。每次_orch都是从 start_node 开始跑一遍完整的 Flow。3.3 实战示例多文件摘要classSummarizeFiles(BatchFlow):defprep(self,shared):return[{fn:f}forfinshared[files]]summarizeSummarizeFiles(startload_nodesummarize_node)prep返回[{fn: a.txt}, {fn: b.txt}, {fn: c.txt}]BatchFlow 会跑 3 遍子流程每次self.params中有当前文件名。子节点通过self.params[fn]访问。3.4 BatchNode vs BatchFlow 的选择简单规则如果只需要一个节点处理多个 item用 BatchNode如果需要一串节点处理多组参数用 BatchFlow。BatchNode 是一把刀切多块肉BatchFlow 是整条流水线跑多轮。⚡ 四、AsyncNode异步 I/O 的优雅封装4.1 为什么需要异步LLM 调用是 I/O 密集型操作——你发一个请求等 2 秒收到响应。如果用同步代码这 2 秒线程完全空闲。如果用异步代码这 2 秒可以同时发 10 个请求等它们一起回来。4.2 AsyncNode 的生命周期AsyncNode 与 Node 的生命周期完全对称只是每个方法都加了_async后缀同步异步说明prep(shared)prep_async(shared)准备数据exec(prep_res)exec_async(prep_res)执行核心逻辑LLM 调用post(shared, p, e)post_async(shared, p, e)后处理 返回 action_run(shared)_run_async(shared)完整生命周期4.3 源码解读classAsyncNode(BaseNode):asyncdefprep_async(self,shared):passasyncdefexec_async(self,prep_res):passasyncdefpost_async(self,shared,prep_res,exec_res):passasyncdef_exec_async(self,prep_res):forself.cur_retryinrange(self.max_retries):try:returnawaitself.exec_async(prep_res)exceptExceptionase:ifself.cur_retryself.max_retries-1:returnawaitself.exec_fallback_async(prep_res,e)ifself.wait0:awaitasyncio.sleep(self.wait)asyncdef_run_async(self,shared):pawaitself.prep_async(shared)eawaitself._exec_async(p)returnawaitself.post_async(shared,p,e)注意_exec_async完整复刻了 Node 的重试逻辑只是加了await和asyncio.sleep。同步和异步的行为完全一致——这是 Pocket Flow 极简但不简陋的体现。4.4 实战示例异步 LLM 调用classAsyncLLMNode(AsyncNode):asyncdefprep_async(self,shared):returnshared[question]asyncdefexec_async(self,question):responseawaitasync_call_llm(question)returnresponseasyncdefpost_async(self,shared,prep_res,exec_res):shared[answer]exec_resreturndefault 五、AsyncFlow兼容同步与异步的编排器5.1 最精巧的一行代码AsyncFlow 的_orch_async方法中有一行代码是整个异步体系最精巧的设计last_actionawaitcurr._run_async(shared)ifisinstance(curr,AsyncNode)elsecurr._run(shared)这行代码让 AsyncFlow可以混用同步和异步节点——遇到 AsyncNode 就 await遇到普通 Node 就直接 run。不需要两套图定义一个 AsyncFlow 搞定一切。5.2 为什么这很重要在实际项目中你不可能把所有节点都改成异步——有些节点是纯计算不需要 await有些节点是异步 I/O必须 await。如果框架不支持混用你就得维护两套图或者把所有节点都改成异步即使不需要。Pocket Flow 用一行isinstance检查优雅地解决了这个问题。5.3 _orch_async 完整源码asyncdef_orch_async(self,shared,paramsNone):curr,p,last_actioncopy.copy(self.start_node),(paramsor{**self.params}),Nonewhilecurr:curr.set_params(p)last_actionawaitcurr._run_async(shared)ifisinstance(curr,AsyncNode)elsecurr._run(shared)currcopy.copy(self.get_next_node(curr,last_action))returnlast_action与同步版_orch唯一的区别就是那一行isinstance检查。其余逻辑——while 循环、copy.copy、get_next_node——完全相同。同步和异步共享同一套路由逻辑只是执行方式不同。 六、并行执行asyncio.gather 的威力6.1 三种执行模式Pocket Flow 提供了三种批量执行模式覆盖从串行到并行的完整光谱模式一BatchNode顺序执行classBatchNode(Node):def_exec(self,items):return[super()._exec(i)foriin(itemsor[])]列表推导逐个执行。简单可靠但无法并发。模式二AsyncBatchNode异步顺序执行classAsyncBatchNode(AsyncNode,BatchNode):asyncdef_exec(self,items):return[awaitsuper(AsyncBatchNode,self)._exec(i)foriinitems]逐个 await但每个 item 是异步的。兼顾顺序和异步 I/O。模式三AsyncParallelBatchNode异步并行执行classAsyncParallelBatchNode(AsyncNode,BatchNode):asyncdef_exec(self,items):returnawaitasyncio.gather(*(super(AsyncParallelBatchNode,self)._exec(i)foriinitems))asyncio.gather同时执行所有 item。速度最快但需注意并发限制。6.2 性能对比假设每个 item 耗时 2 秒LLM 调用10 个 item模式耗时适用场景BatchNode~20sAPI 限速必须排队AsyncBatchNode~20s异步 I/O 但需保序AsyncParallelBatchNode~2s无并发限制追求速度6.3 选型决策简单规则能并行就并行不能并行就顺序需要异步 I/O 就异步。如果你的 LLM API 有并发限制比如 OpenAI 的 RPM用 BatchNode 或 AsyncBatchNode如果没有限制用 AsyncParallelBatchNode。6.4 BatchFlow 的异步变体BatchFlow 也有异步和并行变体classAsyncBatchFlow(AsyncFlow,BatchFlow):asyncdef_run_async(self,shared):prawaitself.prep_async(shared)or[]forbpinpr:awaitself._orch_async(shared,{**self.params,**bp})returnawaitself.post_async(shared,pr,None)classAsyncParallelBatchFlow(AsyncFlow,BatchFlow):asyncdef_run_async(self,shared):prawaitself.prep_async(shared)or[]awaitasyncio.gather(*(self._orch_async(shared,{**self.params,**bp})forbpinpr))returnawaitself.post_async(shared,pr,None)AsyncBatchFlow逐组 awaitAsyncParallelBatchFlow用asyncio.gather并行跑所有组。同样的模式同样的极简实现。 七、系列预告第三篇也是最终篇我们将用 Pocket Flow实战构建两个完整的 LLM 应用应用用到的类设计模式ReAct AgentNode FlowAgent Loop 条件路由RAG 系统BatchNode FlowMapReduce 向量检索关注我不要错过最终篇 总结速查卡批量体系速查类执行方式核心源码适用场景BatchNode同步顺序[exec(i) for i in items]数据并行API 限速BatchFlow同步重跑for bp in pr: _orch(shared, params)任务并行同流程不同参数AsyncBatchNode异步顺序[await exec(i) for i in items]异步 I/O需保序AsyncParallelBatchNode异步并行asyncio.gather(*[exec(i) for i in items])无并发限制追求速度AsyncBatchFlow异步重跑for bp in pr: await _orch_async()异步任务并行AsyncParallelBatchFlow异步并行重跑asyncio.gather(*[_orch_async() for bp in pr])异步并行任务异步 vs 同步速查维度同步异步生命周期prep → exec → postprep_async → exec_async → post_async执行入口_run(shared)_run_async(shared)编排循环_orch(shared)_orch_async(shared)重试time.sleep(wait)asyncio.sleep(wait)混用不支持✅ isinstance 检查兼容同步节点一句话总结Pocket Flow 的批量与异步体系用不到 30 行代码覆盖了从逐项串行到并行并发的完整执行光谱。BatchNode 用列表推导实现数据并行BatchFlow 用 for 循环实现任务并行AsyncNode 用 async/await 封装异步 I/OAsyncParallelBatchNode 用 asyncio.gather 实现真正的并行执行。最精巧的设计是 AsyncFlow 的 isinstance 检查——一行代码让同步和异步节点可以在同一个图中混用。极简但不简陋。参考链接Pocket Flow GitHub 仓库Pocket Flow Batch 文档Pocket Flow Async 文档Pocket Flow Parallel 文档Pocket Flow Multi-Agent 文档