执行摘要
此PR通过重构disaggregation接收器生命周期,修复了因DP平衡(如total_requests/total_tokens)导致的hang问题,确保decode ranks间状态同步,提升系统可靠性。
功能与动机
此变更旨在解决issue #21297报告的PD disaggregation hang问题。当prefill使用内部DP平衡时,decode可能在bootstrap_room -> prefill_dp_rank映射可用前收到请求,导致各rank独立解析pending请求,引发不一致的本地队列状态和handshake stall。PR body中引用issue #21680进行讨论,最终采用receiver生命周期重构方案。
实现拆解
关键改动按模块梳理:
- 基础层(base/conn.py):添加抽象方法
init()和send_metadata(),分离初始化和元数据发送。
- 通用层(common/conn.py):重构CommonKVReceiver,
__init__仅设置基本属性(如bootstrap_addr),init()方法解析prefill DP-rank并调用_setup_bootstrap_infos(),延迟bootstrap设置。
- 解码层(decode.py):修改
add()方法,立即创建receiver并加入队列,但延迟调用init();pending请求列表改为存储DecodeRequest对象;移除leader-based resolution逻辑。
- 具体实现层(fake/mooncake/mori/nixl/conn.py):适配新方法,移除prefill_dp_rank参数,添加
init()和send_metadata()实现。
- 测试层(test_disaggregation_dp_attention.py):新增TestDisaggregationDPAttentionTotalRequests和TestDisaggregationDPAttentionTotalTokens测试类,覆盖DP平衡路径。
评论区精华
review讨论聚焦以下要点:
- 正确性问题:ShangmingCai指出“If we only let the leader rank do the bootstrapping and broadcast to other ranks, will other ranks fail to init
prefill_info_table?”,weireweire回应修改为各rank自行运行但添加sync逻辑。
- 设计权衡:ShangmingCai建议“Maybe we should only call
decode_req.kv_receiver.abort() here... to avoid double free”,weireweire表示failure handling未在此PR解决,需后续决策。
- 演进方向:基于issue #21680讨论,从leader-based workaround转向receiver生命周期重构,确保一致性。
风险与影响
风险:
- 新生命周期分离增加复杂性,若
init()或send_metadata()实现不一致,可能导致handshake失败。
- failure处理逻辑不完整(如abort问题),可能引发资源泄漏或错误报告不及时。
- 同步机制依赖ranks间状态一致,若广播失败,可能重现hang问题。
- 测试覆盖虽增强,但可能遗漏故障场景的edge cases。
影响:
- 对用户:修复hang问题,提升使用DP平衡时的服务稳定性。
- 对系统:改善decode侧pending请求处理,防止因状态不一致导致的stall。
- 对团队:需适应新receiver生命周期,failure handling需后续工作。
关联脉络
与历史PR关联较弱,主要基于issue #21680的讨论进行重构。近期历史PR中涉及disaggregation的测试优化(如PR 21745),但未直接关联核心逻辑变更。此PR揭示了disaggregation模块中DP平衡路径的演进,强调一致性同步的重要性。
参与讨论