Prhub

#31144 [PD] Prevent decode scheduler from blocking on ZMQ sends to a stalled prefill peer

原始 PR 作者 SovietPower 合并时间 2026-07-25 12:53 文件变更 6 提交数 7 评论 19 代码增减 +188 / -116

执行摘要

修复 decode 在 prefill 崩溃后因 ZMQ 无超时而永久挂起

当 prefill 实例因 CUDA OOM 等崩溃时,decode 实例的调度器主循环可能在 ZMQ send_multipart 上永久阻塞,且心跳线程也会因共享锁死锁,导致服务器停止响应请求,即使 prefill 重启也无法恢复(PR #31144 正文)。根源是 ZMQ PUSH socket 未设置 SNDTIMEO。

本 PR 值得所有从事 PD 部署的开发者精读,展示了分布式系统中经典死锁的分析步骤与最小化修复方案的设计权衡。

讨论亮点

关键讨论包括:ShangmingCai 要求最小化改动只保留三个核心点(SNDTIMEO、环境变量、最小 try/except),SovietPower 采纳并精简代码;zmq.DONTWAIT 改为全局 SNDTIMEO 并引入独立环境变量;ronhuafeng 通过 fault injection 发现 _register_kv_args 失败后返回 True 导致后续请求误用连接池索引,SovietPower 修正为返回 bool 并调整缓存顺序;ShangmingCai 提醒在 except 块中设置 conclude_state,已补充。

实现拆解

  1. 公共层基础修改:在 python/sglang/srt/disaggregation/common/conn.py_connect 类方法中,为 PUSH socket 设置 zmq.SNDTIMEO,值来自新增环境变量 SGLANG_DISAGGREGATION_ZMQ_SEND_TIMEOUT(默认 1 秒)。此设置影响所有 PUSH socket(send_metadata、_register_kv_args、_send_abort_notification),从根本上防止无限阻塞。
  2. 注册与发送的异常处理:在 Mooncake、Mori、nixl 三个后端的 _register_kv_argssend_metadata 方法中,将原有的无保护 send_multipart 包裹 try/except zmq.ZMQError,失败时记录错误、设置 conclude_state = KVPoll.Failedupdate_status(KVPoll.Failed),然后返回 False
  3. 失败不污染连接池:在 common/conn.py_setup_bootstrap_infos 中,将 _register_kv_args() 调用提前到缓存 connection_pool 之前,且仅当返回 True 时才缓存,失败时直接返回。避免后续请求复用指向死 prefill 的过时连接。
  4. 环境变量与文档:在 python/sglang/srt/environ.py 添加 SGLANG_DISAGGREGATION_ZMQ_SEND_TIMEOUT = EnvInt(1),并在 docs_new/docs/references/environment_variables.mdx 中记录说明。
  5. 测试覆盖:本次改动未新增专属测试,但通过现有 CI 中 disaggregation 测试套件(如 test_disaggregation_basic.py)覆盖。
文件 模块 状态 重要度
python/sglang/srt/disaggregation/common/conn.py 通信层 modified 7.0
python/sglang/srt/disaggregation/mooncake/conn.py 通信层 modified 7.63
python/sglang/srt/disaggregation/mori/conn.py 通信层 modified 7.46
python/sglang/srt/disaggregation/nixl/conn.py 通信层 modified 7.51
python/sglang/srt/environ.py 配置 modified 4.15

关键符号

_register_kv_args send_metadata _connect _setup_bootstrap_infos

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

ZMQ 发送超时策略:DONTWAIT vs SNDTIMEO 设计

ShangmingCai 认为 zmq.DONTWAIT 过于激进,SovietPower 改为 socket 全局 SNDTIMEO 并新增独立环境变量 SGLANG_DISAGGREGATION_ZMQ_SEND_TIMEOUT(默认 1s)。

结论:采用 SNDTIMEO 且默认 1s,用户可根据网络情况调大。 · 已解决

PR 范围控制:最小化改动 设计

ShangmingCai 要求只保留三个核心点(SNDTIMEO、环境变量、最小 try/except),移除非必要调整。

结论:SovietPower 精简 PR,仅保留核心三处修改。 · 已解决

_register_kv_args 返回值与连接池缓存时序 正确性

ronhuafeng 通过 fault injection 发现 _register_kv_args 失败后仍返回 True,导致后续请求误用连接池索引,再次触发 send_metadata 死锁。

结论:改为返回 bool;公共层 _setup_bootstrap_infos 先调用注册,成功后才缓存连接池。同时 except 块中设置 conclude_state 和 update_status(Failed)。 · 已解决

_register_kv_args except 块缺省 conclude_state 正确性

ShangmingCai 在 mooncake/conn.py 第 1952 行建议在 except 中也设置 self.conclude_state = KVPoll.Failed。

结论:SovietPower 确认已添加。 · 已解决

风险与影响

超时阈值风险:默认 1 秒在极端网络延迟或大消息场景下可能误判,但正常发送毫秒级,用户可调大环境变量。行为变更:_register_kv_args 失败后不再继续,修正了之前的隐含 bug。缺少专项测试:未新增单元测试覆盖超时恢复路径,回归风险依赖集成测试。未来 ZMQ 配置再次修改时需保持统一原则。

用户影响:所有使用 PD 部署的用户受益于更高稳定性,decode 实例不再因 prefill 意外崩溃而永久不可用。系统影响:调度器主循环不再被 ZMQ 发送阻塞,心跳线程可正常清理,实例自恢复时间从分钟级降至秒级。团队影响:统一 ZMQ 超时处理模式,降低后续排查和维护成本。

核心调度路径变更 缺专项测试覆盖 依赖超时默认值

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论