Prhub

#36031 refactor(disagg): dedupe mooncake failure_exception into a mixin

原始 PR 作者 ShangmingCai 合并时间 2026-08-23 14:17 文件变更 1 提交数 1 评论 2 代码增减 +29 / -33

执行摘要

Mooncake 收发端 failure_exception 抽为共享 mixin,纯重构

PR body 明确说明:MooncakeKVSender.failure_exceptionMooncakeKVReceiver.failure_exception 在 15 行与 14 行方法体上只有一行注释不同,diff 仅报 1 行差异;两者对失败房间的收尾方式一致——锁存 Failedclear() 本地状态、在 failure_lock 下弹出记录的原因,本地无记录时按“来自另一端失败”抛出 KVTransferError。该 PR 与 #35950 组合生效:#35950 删除了 CommonKVSender / CommonKVReceiver 上的 “Fake” 占位实现,使基类恢复抽象约束,本 PR 的 mixin 恰好补齐 Mooncake 系需要的实现且不产生遮蔽。

值得快速阅读:真正的价值不在代码本身,而在两个设计决策——一是为什么用 mixin 而非基类方法(避免给形态不同的子类塞入不可达实现),二是为什么主动拒绝跨后端统一(用少量重复保护错误路径的局部性)。作者用 AST 等价性 + MRO 模拟验证纯移动重构的方法,也值得在同类改动中复制。无需精读。

讨论亮点

本 PR 没有 review 评论(review_comments_count 为 0),唯一的互动在 issue 评论中:作者发出 /rerun-group disaggregation,机器人返回三组结果——4-gpu-gb300(1 个用例)与 2-gpu-h100(5 个用例)全部通过,8-gpu-h20(5 个用例)出现失败记录且回复被截断、完整状态未展示。PR body 中作者主动展开了两个设计权衡:为什么用 mixin 而不是放到 Common 基类(Mori 形态不同、会给无 staging 的类带来不可达方法),以及为什么拒绝跨后端统一(重复换隐式契约不划算)。这些权衡未引发进一步分歧。

实现拆解

  1. 新增共享 mixin:在 python/sglang/srt/disaggregation/mooncake/conn.pyMooncakeKVSender 之前定义 MooncakeFailureExceptionMixin,把两份 failure_exception 方法体原样移入(仅统一注释措辞)。Docstring 显式列出子类必须提供的属性(conclude_stateclear()bootstrap_roomkv_mgr),把 mixin 依赖子类状态这一事实写成文档化契约。
  2. 接入两个具体类MooncakeKVSenderMooncakeKVReceiver 的基类列表分别改为 (MooncakeFailureExceptionMixin, CommonKVSender)(MooncakeFailureExceptionMixin, CommonKVReceiver),mixin 放首位保证 MRO 优先解析;同时删除两份内联副本。AscendKVSender / AscendKVReceiver 通过继承 Mooncake 类传递解析,无需改动。
  3. 验证方式:用 AST 抽取旧方法体与 mixin 方法体、去掉注释后比较——old sender == old receivernew mixin == old 均为 True,证明是移动而非重写;grep 确认两个具体类上不再残留 failure_exception;模拟四个类(含 Ascend 传递继承)的 MRO 解析均命中 mixin。PD 失败路径需要多 GPU 硬件,未在本地执行,依赖 CI 的 disagg 测试组实测。
  4. 刻意不扩展到其他后端:Mori 的 sender 调用 _finalize_failure() 而非锁存 conclude_state,默认 reason 也不同;NIXL 的 sender 会先 re-raise self._send_error 或存储的 exceptions[room] 再回退。跨三后端统一需要在三处错误路径上 duck-typing 一个适配 sender / receiver 形态的 helper,作者判断“用少量重复换一个只会在失败路径上生效的隐式契约”不划算。
  5. 测试配套:无新增单测(PR 判定不适用,失败路径需多 GPU);通过 pinned 的 ruff 0.15.1 --select=F401,F821,UP037black 26.1.0isort 7.0.0
文件 模块 状态 重要度
python/sglang/srt/disaggregation/mooncake/conn.py 失败处理 modified 7.92

关键符号

failure_exception MooncakeFailureExceptionMixin

关键源码片段

python/sglang/srt/disaggregation/mooncake/conn.py refactor

唯一改动文件:新增 `MooncakeFailureExceptionMixin` 并把 `MooncakeKVSender` / `MooncakeKVReceiver` 的 `failure_exception` 收敛到 mixin,同时调整两类的继承列表以配合 MRO。

class MooncakeFailureExceptionMixin:
    """Shared `failure_exception` for the Mooncake sender and receiver.    Both sides conclude a failed room identically: latch Failed, clear local
    state, then raise with the recorded reason -- or, when no reason was
    recorded locally, report it as propagated from another rank. Expects the
    concrete class to provide ``conclude_state``, ``clear()``,
    ``bootstrap_room`` and ``kv_mgr``.
    """
​
    def failure_exception(self):
        # 请求已在另一个 rank 失败,本地状态未锁定时显式置为 Failed
        if self.conclude_state is None:
            self.conclude_state = KVPoll.Failed
​
        # 清理本地发送 / 接收缓冲等状态
        self.clear()
​
        # 在 failure_lock 保护下取出本 room 记录的失败原因
        with self.kv_mgr.failure_lock:
            failure_reason = self.kv_mgr.failure_records.pop(self.bootstrap_room, None)
        # 本地无记录说明失败由另一端传播而来,按 propagated 上报
        is_propagated = failure_reason is None
        if is_propagated:
            failure_reason = "Failed due to an unknown reason from another rank"
        raise KVTransferError(
            self.bootstrap_room, failure_reason, is_from_another_rank=is_propagated
        )
​
​
# mixin 放在继承列表首位,保证 MRO 优先命中共享实现;
# Common 基类在 #35950 中已移除占位实现,故这里不产生任何遮蔽
class MooncakeKVSender(MooncakeFailureExceptionMixin, CommonKVSender):
    ...
​
​
class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver):
    ...

评论区精华

disaggregation CI 测试组 rerun 结果 测试

作者在 issue 评论发出 `/rerun-group disaggregation`,机器人随后汇报 4-gpu-gb300(1 个用例)、2-gpu-h100(5 个用例)、8-gpu-h20(5 个用例)三组结果。

结论:4-gpu-gb300 与 2-gpu-h100 全部通过;8-gpu-h20 组显示失败记录,且机器人回复被截断,无法确认最终状态。本 PR 为纯移动重构,理论上与失败无关,但失败路径需要多 GPU 硬件实测。 · 已处理

风险与影响

行为风险低但客观存在:失败路径只能靠多 GPU 实测,本地未执行;CI rerun 中 8-gpu-h20 组有失败记录,虽然本 PR 为纯移动重构、理论上无关,仍需留意后续 disagg 测试是否稳定。MRO 顺序敏感性:mixin 位于继承列表首位,若未来 CommonKVSender / CommonKVReceiver 重新实现 failure_exception,会被 mixin 遮蔽;若未来删除 mixin,需同步确认 Ascend 子类解析。隐式契约风险:mixin 依赖子类的 conclude_stateclear()bootstrap_roomkv_mgr,未来新增子类若缺失其一,会在失败路径抛 AttributeError(docstring 已作说明但无运行时防护)。无性能、安全、数据面兼容性影响。

用户与系统层面无外部行为变化:仅 mooncake 后端 sender / receiver 失败路径的代码组织方式变化,AscendKVSender / AscendKVReceiver 经 Mooncake 继承不受影响。团队层面,这是 disagg 清理系列的一环,消除了 mooncake 收发端失败处理的重复维护点,并为后续统一错误处理模式(mixin 组合补齐基类抽象方法)提供先例。影响范围窄、影响程度低。

失败路径仅多 GPU 可实测 CI 存在失败记录 MRO 顺序敏感性 mixin 依赖子类隐式契约

关联 Issue

#35948 refactor(disagg): hoist duplicated _handle_staging_req into a mixin
#35950 refactor(disagg): drop dead placeholder overrides in Common KV sender/receiver
#35980 refactor(disagg): hoist staging helper imports out of the bootstrap loops
#36006 refactor(disagg): register SGLANG_ENCODER_MM_LOAD_WORKERS in Envs
#36030 refactor(disagg): move _is_watermark_ready into StagingManagerMixin

完整报告

参与讨论