ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

离线推理框架核心链路拆解:Scheduler与ExecutionEngine的协作艺术

离线推理框架核心链路拆解:Scheduler与ExecutionEngine的协作艺术 前阵子读 cannn-recipes-infer 源码我习惯先把一条链路的入口文件铺在桌面上再顺着日志顺序把类调用过一遍。写这系列笔记第三篇的时候正好走到离线推理这条最核心的执行链路OfflineInference - Scheduler - ExecutionEngine - ModelWorker。刚接触这套框架的读者多半会有个困惑一个离线推理脚本看起来就是“读数据、进模型、出结果”三步为什么内部要拆出这么多角色搞懂这条链路之后你会发现每一个参与者都有不可替代的位置Scheduler 管“什么时候做、做哪批”ExecutionEngine 管“怎么拆、交给谁”ModelWorker 管“真刀真枪碰硬件”而 OfflineInference 则是这一切的总开关。这篇文章适合两类人一是被业务需求推着要用昇腾卡做推理想搞清楚 cannn-recipes-infer 到底怎么把请求送入模型的开发者二是对服务化推理框架的调度与执行架构感兴趣想拿一套实际工程实现做参照的同行。我会按执行顺序拆解每个角色的核心职责把关键数据结构、接口设计、调度策略和调试方法都摊开来讲尽量做到看完就能自己在源码里找到对应位置。1. 整体设计一条推理请求在框架里如何“走完全程”1.1 什么是 OfflineInference它和在线的区别在哪在聊链路之前得先明确 OfflineInference 的定位。它通常指“一次性喂一批输入等待全部或部分结果返回”的推理模式典型场景是数据清洗、批量打分、离线评测、离线增强而不是对话服务那种“来一个问句快速回一个答案”的在线推理。在线推理对首 token 延迟敏感对单请求响应时间有硬指标离线推理则更关注吞吐允许排队、允许动态组 batch甚至允许把不同长度的序列拼在一起执行。这个本质差异决定了代码结构。在线场景要专门处理持续连接、请求超时、会话状态离线场景则可以把精力集中在“批量构造、调度、执行、收集结果”这条干净的主线上。cannn-recipes-infer 里 OfflineInference 的出现本质上就是把这条主线做成一个可复用的编排器把调度和执行细节下沉到下面的组件里。1.2 四个核心类各自的职责边界先看一张职责划分的速览表后面再逐个展开组件核心职责关键产出OfflineInference入口编排、配置解析、请求初始化初始化后的推理会话、执行结果Scheduler请求排队、批决策、显存/优先级管理可执行的调度批次ScheduleBatchExecutionEngine调度批次到硬件任务的转换、资源协商任务执行上下文、任务句柄ModelWorker加载模型、执行算子、维护运行时状态推理结果、张量输出、日志回调它们不是简单的“调用链”而是围绕任务状态机互相协作。OfflineInference 更像“项目经理”Scheduler 是“排产计划员”ExecutionEngine 是“车间主任”ModelWorker 是“操作机器的工人”。每个角色都有明确的产出物下游只依赖上游的产出不关心上游内部怎么实现。1.3 为什么要拆成 Scheduler / ExecutionEngine / ModelWorker 三层拆层不是增加复杂度是为了让每个问题域都能独立演进。我自己的理解是Scheduler 和服务态、批处理策略绑定它会随负载特征变化不断调整必须和硬件执行细节解耦。ExecutionEngine 要管理“调度出来的批次怎么落到具体的卡上”它需要知道设备数量、通信拓扑、内存预算但不需要关心算子怎么实现。ModelWorker 最贴近硬件负责加载权重、运行模型、把张量结果搬回宿主内存它只服务 ExecutionEngine 下达的命令。这有点像后端开发里 controller - service - dao 的分层。不在 controller 里写 SQL也不在 dao 里做权限校验是因为边界清晰之后每一层都能单独测试和替换。cannn-recipes-infer 真正工程的点在于三层之间通过明确的数据对象通信并没有出现层层透传大字典的坏味道。2. OfflineInference 启动入口初始化与任务分发2.1 从执行脚本到 OfflineInference 构造打开 cannn-recipes-infer 的仓库常见的入口脚本大致是这样的形态def main(): config load_config(configs/example.yaml) infer OfflineInference(config) outputs infer.infer(prompts, generation_config) save_outputs(outputs)看起来简单但OfflineInference(config)这一步内部工作量很大。构造过程会做四件事解析混合数据源配置、初始化日志和运行时环境、创建 ExecutionEngine、向底层注册默认的调度策略。如果配置里包含了模型并行参数比如多卡场景下的 TPTensor Parallel或 PPPipeline Parallel配置这一步还会做资源拓扑协商确认当前机器有几张可用卡、显存多大、是否满足并发需求。有一个容易忽略的细节OfflineInference构造并不代表模型已经加载到显存。模型真正的加载发生在 ModelWorker 初始化阶段而 ModelWorker 的启动又在 ExecutionEngine 的 lazy init 流程里。这样设计的考虑是如果用户只做配置校验或者只构建任务图就没必要白白占住显存。2.2 配置解析与关键参数离线推理和在线服务相比配置文件里多了一批和批处理相关的参数。我整理过一组常用配置项按作用分组参数分组典型参数作用说明模型类model_path, tokenizer_path, dtype权重和分词器位置、精度资源类device_ids, tensor_parallel_size, block_size选择卡、并行度、KV Cache 块大小调度类max_num_batched_tokens, max_num_seqs, max_batch_size控制单轮调度能吃的请求量和 token 量生成类max_new_tokens, temperature, top_p生成策略参数运行时类log_level, cache_dir, stream_output日志、缓存、是否流式返回max_num_batched_tokens这个值很关键它决定一次调度最多往模型里塞多少个 token。过小会浪费多卡并行能力过大会导致单次 step 耗时变长甚至显存溢出。配参数的时候不能只盯着最大并发数要想清楚你手上的卡是 32G 还是 64G模型是 7B 还是 13B。实测下来7B 模型在 32G 单卡上如果 max_num_batched_tokens 给到 8192再用上 paged KV cache是能稳定运行并保住吞吐的。2.3 请求预处理与任务批次构建拿到一列 prompts 之后OfflineInference 不会立刻把所有数据原生塞给下游。它要做 tokenize、长度统计、按预分配策略分组。源码里通常能看到类似build_request_batch的过程它会把每个 prompt 转换成“请求对象”里面至少包含class InferRequest: request_id: str input_ids: List[int] sampling_params: dict arrival_time: float state: str # pending / running / finished带上arrival_time是很实用的设计调度器可以据此实现先到先服务或者结合优先级做老化处理防止低优先级请求永远饿死。这个对象就是 Scheduler 眼里最小的任务单元。预处理阶段还有一个值得注意的点长文本截断和 batch 内长度对齐。离线推理场景里prompt 长度方差往往很大有的可能是几十 token有的可能是上千 token。如果不做处理短序列和长序列混在一个 batch 里执行短序列会白白等待长序列跑完。处理方式有两种一种是把超长样本放进单独的 batch另一种是允许动态 batch让短序列先执行长序列慢慢跑完后再加入新请求。cannn-recipes-infer 的实现路径更接近后者因为连续 batching 本身就是这类框架的主打能力。2.4 触发执行infer 方法的内部编排用户调用infer()后OfflineInference 首先把请求批次交给 Scheduler 注册然后进入一个事件循环。这个循环的核心逻辑可以概括为def infer(self, prompts, generation_config): requests [self._build_request(p, generation_config) for p in prompts] scheduler self.engine.scheduler scheduler.add_requests(requests) while not scheduler.is_finished(): scheduler.schedule() engine.execute(scheduler.current_batch) engine.collect_results() return self._assemble_outputs()while not scheduler.is_finished()是离线推理的精髓。在线服务里请求是源源不断进来的循环永远不退出离线推理则有一个明确的终止条件所有请求都进入 finished 状态。每一轮循环里Scheduler 决定当前该执行哪些请求Engine 负责真正跑一版模型 forward然后收集结果如此反复直到全部完成。我读源码时的一个体会是这个循环看起来简单写好不容易。关键是要处理好“跑完一轮后哪些请求还需要继续生成下一个 token”这件事。不是所有请求一轮就结束多数请求要迭代很多轮每一轮都需要把新生成的 token 拼接回 KV cache 并继续调度。这套“迭代式生成 动态调度”的机制正是离线推理框架吞吐高的基石。3. Scheduler 调度核心排队、合并与批次决策3.1 调度器的内部状态与请求队列Scheduler 内部通常维护两个主要集合等待队列waiting queue和运行集合running set。等待队列里的请求还没有真正进入模型运行集合里的请求正在被 batch 执行或者已经执行了一部分。class Scheduler: def __init__(self, policy_typefcfs): self.waiting: List[InferRequest] [] self.running: List[InferRequest] [] self.finished: List[InferRequest] [] self.policy self._init_policy(policy_type)在进一步看 schedule 逻辑之前要先理解“一个请求从 waiting 到 running 意味着什么”。进入 running 意味着这个请求已经拿到了第一轮推理所需的显存空间包括 KV cache 里给它预留的块。显存不是无限的所以 Scheduler 最大的约束不是 CPU 算力而是“当前还能释放出多少显存块给新请求”。这也是为什么 Scheduler 要主动管理 KV cache 块而不是简单地把请求丢给执行器。3.2 连续 Batching离线场景下的默认调度策略早期离线推理框架常用静态 batching也就是把所有请求预先分成固定大小的 batch一个 batch 跑完再跑下一个。静态 batching 实现简单但弊端很明显一个 batch 里一旦有早完成生成的短请求它占的显存和计算位置依然被锁住直到整个 batch 都结束。cannn-recipes-infer 主推的是连续 batchingcontinuous batching思路每一轮调度都重新扫描 waiting 和 running 集合已经结束的请求立刻退出把显存块释放给新请求还在生成中的请求和新加入的请求按资源预算合并进本轮 batch。这样的效果就是模型每执行一次 forwardbatch 里的内容都可能发生变化GPU 永远不会被“等队友”拖慢。在源码里这个逻辑通常表现为一个循环def schedule(self) - Optional[ScheduleBatch]: self._preempt_overloaded_requests() batches self._create_batches_by_policy() return batches_preempt_overloaded_requests是资源不足时的兜底策略。如果当前显存块不足Scheduler 会选择挂起或抢占一些请求把显存让给优先级更高或更接近完成的请求。这个处理方式让我想起操作系统里的页面置换只不过这里的“页”变成了 KV cache 块。3.3 调度决策的关键指标调度器在决定“本轮要把哪些请求合并成一个 batch”时不是简单地把 waiting 里的请求拉出来而是要算一笔账。我摘录几个关键约束请求数上限max_num_seqs防止 batch 里请求数过多导致图执行时间不可控。token 数上限max_num_batched_tokens整批的输入 token 总数不能超过这个值。显存块预算当前空闲 KV cache 块能否覆盖新请求的max_new_tokens需求。优先级规则默认先到先服务也支持按优先级字段调整。这些约束需要组合判断即便请求数很少如果每个请求都很长token 总数也可能撞到上限。相反如果请求全部很短可能在请求数量上先撞限。写调度器的时候最容易忽略的是显存块预算因为它在单请求视角下看不出来只有把模型参数、序列长度和可用显存放到一起才算得明白。3.4 调度结果如何交给 ExecutionEngineScheduler 完成决策后产出的是一个ScheduleBatch对象里面不再直接放原始的字符串 prompt而是放“准备好执行的 token 序列 对应的采样参数 显存块索引”。它相当于一份“施工图”Engine 拿到之后不需要关心为什么是这些序列在一起执行只需要照着这份图纸把工作做掉。我特别看重这个边界。如果 Scheduler 把原始请求直接传给 Executor那么 Executor 里必然要再做一遍调度决策职责就混乱了。现在这样设计Scheduler 可以单独用一组纯 CPU 的单元测试覆盖各种队列场景不需要拉起卡做验证Engine 侧也能专注于执行和资源管理。调试的时候只要确认ScheduleBatch内容符合预期问题范围就能缩小很多。4. ExecutionEngine从调度计划到可执行任务4.1 Engine 在整条链路中的坐标ExecutionEngine 在 cannn-recipes-infer 里的定位我总结成一句话把“调度计划”翻译成“硬件能跑的任务”并在任务执行过程中管理生命周期。它不是直接调用模型的模块更接近“任务编排器”。第一次看源码时我被 ExecutionEngine 和 ModelWorker 的边界搞晕过。后来发现简单场景下 Engine 可以直接驱动模型执行但一旦涉及多卡、时序并行或者流水并行Engine 就要跟多个 Worker 通信。所以源码里更常见的做法是 Engine 维护一组 Worker 句柄按调度批次把子任务拆给不同 Worker再汇总结果。4.2 执行引擎的接口设计从调用关系看ExecutionEngine 至少暴露三个核心方法class ExecutionEngine: def __init__(self, config): self.device_ids config.device_ids self.workers self._init_workers(config) self.cache_blocks self._init_cache_blocks(config) def execute_batch(self, batch: ScheduleBatch): tasks self._split_batch_to_tasks(batch) futures [worker.submit(task) for worker, task in zip(self.workers, tasks)] return futures def collect(self, futures): results [] for f in futures: results.append(f.result()) return results_init_workers这一步值得展开。它根据传入的device_ids决定启动几个 Worker。比如你有 4 张卡tensor_parallel_size4那通常会启动 4 个 Worker每个 Worker 绑定一张卡如果tensor_parallel_size2则可能启动 2 个 Worker每个 Worker 管理两张卡的计算流程。初始化时如果显存不足或卡被占用错误会在这里提前暴露而不是等到跑 batch 时才报错这一点对线上排障很友好。4.3 与 ModelWorker 的通信任务下发与回调Engine 和 Worker 之间是典型的“生产-消费”模式。Engine 作为产出任务的消费者把每个子任务封装成一个WorkerTask里面包含 input_ids、block_tables、sampling_params 等。Worker 执行完任务后返回WorkerOutput包含生成的 token 序列、logprobs、耗时统计等。如果只看代码很容易把“Engine 调 Worker 方法”理解成同步调用。实际上在多卡场景下任务下发通常走队列Worker 之间还有可能做通信。一个有用的经验是调试多卡任务时不要只看主进程日志最好把每个 Worker 的日志分开输出否则你很难判断延迟到底卡在哪个设备上。4.4 Cache 管理与显存块分配ExecutionEngine 里还有一个常被忽视但很重要的模块Cache Manager。它负责 KV cache 块的分发、回收、拷贝。Scheduler 做预算判断时需要向 Engine 查询剩余块数Engine 执行完一轮后又会把不再使用的块标记为空闲。KV cache 块的大小通常由block_size控制常见值是 16 或 32。block_size 越大管理开销越小但内部碎片越多block_size 越小显存利用越灵活但块表管理更复杂。这块调优对离线吞吐影响很大。我之前跑一批长文本评测时把 block_size 从 32 改成 16显著降低了显存碎片批大小能往上抬整体吞吐提升了将近一成。如果你的业务序列长度波动大值得试试更小的块大小。5. ModelWorker真正碰硬件的那一环5.1 Worker 初始化加载模型、绑定设备、配置运行模式ModelWorker 是整条链路里最重的一环。初始化时它要做的事情包括class ModelWorker: def __init__(self, config, device_id): self.device_id device_id self.model self._load_model(config) self._set_device(device_id) self._prepare_cache_blocks(config)_load_model要用 CANN 环境下的模型加载接口把权重从磁盘搬到设备内存中。这里有个容易忽略的问题模型的权重复制不是一次写完就完了。分布式推理时每个 Worker 只持有自己负责的权重分片如果配置了开启图模式的算子融合还会额外编译和缓存执行图。图模式的好处是运行时省去不少 Python 算子分发开销代价是首次启动时间明显变长。你要是第一次跑脚本发现等了很久才有日志输出多半是在编译融合图。5.2 Worker 的执行循环Worker 启动后通常进入一个事件循环等待 Engine 下发任务。这个循环的设计直接决定了吞吐上限def worker_loop(self): while True: task self.task_queue.get() if task is None: break output self._execute_task(task) self.result_queue.put(output)_execute_task内部会完成把 input_ids 搬到设备、创建图输入、执行模型 forward、采样生成下一个 token、更新 KV cache、把新 token 和设备上的采样结果搬回宿主内存。如果开启流式输出这个函数还会通过回调把中间结果推给上层。我第一次读到这里时惊讶于“采样”居然也在 Worker 里做。后来想通了如果采样放回 Python 高层做每生成一个 token 就要设备到宿主拷贝一次延迟和带宽都扛不住。在设备端直接完成采样只把最终 token id 传回是性能上的必然选择。5.3 结果回传与流式输出离线推理虽然不像在线服务那样每个 token 都要马上推给用户但有些场景仍然需要流式回调。比如评测程序想要看每轮生成的进度或者长文本生成任务希望实时落盘。Worker 通过回调接口上报中间状态def _execute_task(self, task): for step in range(task.max_new_tokens): logits self.model.forward(input_ids, cache_blocks) next_token self.sample(logits, task.sampling_params) if task.stream_callback: task.stream_callback(step, next_token) if next_token EOS: breakstream_callback的设计让上层能拿到逐步生成状态又不必每步都进 Python 主线程做同步因为回调本身还是在 Worker 线程里执行的避免了额外的进程切换开销。5.4 单测与硬件隔离我个人的经验是ModelWorker 是整条链路里最难单独调试的模块因为一跑就要占卡、占显存。为了快速验证 Worker 逻辑可以在代码里加一个execution_mode开关当它等于mock时Worker 不真正调用模型只返回基于规则生成的伪结果。这样既能验证调度链路又避开了硬件依赖。但要注意mock 模式只能验证通信和状态机不能验证算子融合和显存管理。真正上卡之前专门预留一台干净的测试机器跑一遍全链路确认没有设备占用后才能放心做批量任务。6. 串起来看一次完整推理的时序与排障6.1 关键源码位置速查为了方便接着读源码我把自己梳理的“模块与可能文件/类名”的对应关系列在下面。不同版本可能略有差异但按职责去找一般都能定位到逻辑角色常见命名线索需要关注的重点函数/成员OfflineInferenceoffline_inference.pyinfer,_build_request,_assemble_outputsSchedulerscheduler.pyschedule,add_request,_preemptExecutionEngineexecution_engine.pyexecute_batch,_init_workers,collectModelWorkermodel_worker.pyworker_loop,_execute_task,load_modelCache Managercache_manager.pyallocate,free,block_tables拿到一份新代码时我习惯先打开日志文件按日志打印顺序反推调用链再回到源码确认每个模块的边界。日志里如果出现schedule和execute_batch交替出现基本说明进入了主循环。6.2 调试手法断点、日志与性能观察调试这条链路最实用的三个手段断点打在Scheduler.schedule()返回前观察ScheduleBatch里请求数和 token 数是否符合预期。日志切到 debug 级别能看到每次调度后的队列长度变化这是判断调度策略是否生效最快的方式。在 Engine 的execute_batch前后各打一个时间戳能一眼看出是“调度太慢”还是“执行太慢”。调度慢多半是 Python 层队列操作或显存计算的问题执行慢则要往模型、设备通信方向排查。还有一个压箱底的经验大量使用device to host拷贝时通过昇腾环境里的 profiling 工具看 Host 和 Device 之间是否有长时间同步等待。很多离线条慢的原因不是算力不够而是 CPU 端数据准备好了但设备端还在忙或者反过来设备端在干等 CPU 下发。用 profiling 数据对照调度时间戳基本能立刻定位瓶颈。6.3 常见问题与速查表把我在实际跑批过程中遇到的高频问题整理成一张速查表供你排查时直接对号入座现象可能原因建议处理启动很慢日志长时间无输出图模式编译或首次加载权重耗时如果是首次编译耐心等反复出现则检查缓存目录显存溢出 OOMbatch 参数配得过大或 KV cache 预算不足调小 max_num_batched_tokens或减小 block_size吞吐很低但显存没满调度批次过小设备利用率不足增大 max_num_seqs检查是否有显存碎片特定请求特别慢长序列抢占短序列资源或优先级策略不当开启抢占/超时策略把超长请求单独分桶多任务并行时结果错乱Worker 之间共享了同一份显存块检查块分配是否加锁确认 device_id 绑定正确流式回调收不到中间结果回调在线程中异常退出在 callback 里加 try/except查看 Worker 日志6.4 我的一些优化心得这条链路跑顺之后我最大的感受是优化调度参数比优化算子实现更划算。因为底层模型算子是 CANN 已经优化好的你干预空间不大但调度层的参数几乎是完全可以自己掌控的。一个简单的调优套路是先定显存上限再定 block_size然后逐步抬高 max_num_batched_tokens 和 max_num_seqs观察吞吐曲线。注意吞吐不是单调上升的batch 太大时单次 forward 耗时变长整体吞吐反而下降。我实测中曲线会有一个明显的平台期取平台期的中间点作为最终配置最稳妥。另外一个让我受益很多的做法是把每次离线推理的配置、请求数、序列长度分布、吞吐、耗时记成一张实验表。因为调度策略之间的差异只有在大量样本上才能看出来靠一两次直觉调整很容易误判。有了实验表你就能用数据说服自己和同事而不是靠“感觉这次更快”。回到开头那个问题为什么离线推理简单脚本内部要拆出这么多角色现在我可以说
返回列表