[lihenga]

动机

修复 N10 / 禅道 2816528：P 已与 D 握手但尚未完成 KV 发送时取消请求，D 可能一直等到 waiting 超时。本 PR 让等待、bootstrap、inflight、运行中和 chunked 请求进入可重试的安全清理流程。

修改

• 关闭 sender 后，等待前向结果和本 rank 的源传输全部排空，再发送 DRAINED_ABORT 并释放资源；保留原取消原因并避免重复响应。
• 末块结果处理时，所有 rank 都将中止请求放入 inflight 队列，包括本地已经清理完的 rank；由原有 MIN 共识统一移出，避免某些 rank 跳过集合通信。
• 源传输排空后，先在 bootstrap 锁内原子保存通知目标、发布终态记录并清除 request_status、transfer_infos 和 failure_records 等活动状态；用最多 4096 条的终态记录处理晚到握手。TTL 为 bootstrap/waiting timeout 与 300 秒的最大值。终态握手收到通知后不再重建活动传输状态。
• bootstrap 登记与终态清理仅在锁内更新状态和保存目标快照，ZMQ 通知在锁外发送，避免 bootstrap 的慢发送阻塞调度线程获取清理锁。通知失败时 sender 保留目标快照并重试；终态晚到握手仍获通知且不重建活动状态。
• 回归使用真实 CommonKVManager 的 update_status、CommonKVSender 的 abort/clear 和 Mooncake failure_exception 状态机，取消原先无条件写回 Failed 的替身。

部署要求

P、D 必须同时升级到支持四段 DRAINED_ABORT 消息的版本。 旧 D 不支持该协议；本 PR 未增加能力协商。终态记录有 TTL 和容量上限，超出保留窗口的迟到握手仍由原超时机制收尾。

验证

指定环境：10.211.10.10 / Docker sglang0518_0920_yb，独立源码目录，未替换安装包或重启服务。

• VERIFIED：test_external_kv_abort_drain.py 27 passed、23 subtests passed。包含真实双 rank Gloo：rank 0 已排空、rank 1 未排空，两个 rank 都入队，均在第二轮统一移出；以及真实 clear 状态机后的晚到通知、TTL 到期、4096 条容量限制。
• VERIFIED：实际导入的 test_mooncake_final_chunk_poll.py 20 passed、10 subtests passed，覆盖多 rank 协议、晚到目标、发送失败后的目标保留与重试，以及真实 bootstrap 登记路径在通知被阻塞时另一线程仍能完成终态清理（活动 Failed 和终态两种情况）。
• VERIFIED：有界差异自审、定向 Ruff 和 git diff --check。完整 pre-commit 未执行。
• UNTESTED：本轮未重跑 1P1D / TP2 / overlap / 64K 请求的 GPU 压力验收，不将 CPU/Gloo 结果视为真实设备验收。本次修复 bootstrap 持锁发送造成的调度器锁等待；调度线程自己发送通知的 ZMQ 超时及排空期间复用 rid 的低优先级建议未扩展处理。

日志：/stortest/zhy_spcae/fixbug/review-pr38-20261021/review3-lock38-ast.log；/stortest/zhy_spcae/fixbug/t01-session-rpc-20261021/review3-lock38-protocol.log。

AI 辅助：Codex 实现、测试与差异自查；维护者人工 review 仍待完成。

[lzb16]

审核与实测意见

我们审核了 #38（48f0ae831c、ff6654082）：两人独立只读审核，再回代码核实。实测用的构建包含本 PR 的改动（变更行与本 PR 逐行相同），同时还包含 #37、#39、#40；对照组是不含这几个改动的上一个包。行号按 ff6654082。

结论：已握手的请求被中止后通知 D 的主路径，实测有效。还有两处建议修改：

• 多 rank 下末块中止时，各 rank 入队可能不一致（高，条件触发）；
• bootstrap 队列中止后，D 晚到的握手仍收不到通知（中）。

1. 实测

1P1D，P、D 均为 TP2；P 开 overlap 调度，chunked prefill 16K；D 的 waiting 超时设为 120 s。

场景	修复前	本 PR
并发 24 个 64K 请求，4 s 后 abort_all：被中止的 23 个请求，D 返回 500 的时间	120.36–120.54 s	4.77–5.05 s
并发 6 个 64K 请求，3 s 后按 rid 中止 1 个（此时它在 bootstrap 或 waiting 队列，D 已握手）：D 返回的时间	120.16 s	4.80 s
中止后的新请求	正常	正常
D 健康检查	全程 200	全程 200

本 PR 头上 test_external_kv_abort_drain.py 的 26 个用例全部通过。

2. 需要修改的问题

(1) 多 rank 下，末块中止按本 rank 的排空状态决定是否进 inflight 队列（高，条件触发，代码分析）

• prefill.py:1941-1958：处理末块结果时，cleanup_done 取决于本 rank 的 prefill_abort_sender_terminal(req)，也就是本 rank 之前分块的传输是否已经排空。为真的 rank 直接结束，为假的 rank 把请求放进 disagg_prefill_inflight_queue。
• prefill.py:2116-2117：inflight 队列为空的 rank 直接返回，不参加 :2123 的 poll_and_all_reduce_attn_cp_tp_group。另一个 rank 进入了集合通信，两边就会错配或挂起。后面按 all_reduce_min 统一移出的逻辑，在入口已经不一致时救不回来。
• 本 PR 之前，用户中止落到这里时走 transfer_failed，所有 rank 都入队，是一致的。本 PR 之后，任意用户取消都会走这条分支，原来只有外部 KV 加载失败才会进来。
• 触发条件：attention TP 或 CP 至少为 2；开 overlap，即末块 forward 已经发出、结果还没处理时收到中止；各 rank 的传输进度不同。我们的实测（开 overlap、TP2）两次都没有挂起，只能说明不是必然触发。
• 建议：
  • 末块分支无条件进 inflight 队列，交给 inflight 循环开头的 external_kv_cleanup_done 判断和 all_reduce_min 一致移出；或者在决定 cleanup_done 之前，对 sender 的排空状态做组内共识；
  • 补一个 2 rank 的真实 gloo 用例：一个 rank 已排空，另一个未排空。

(2) bootstrap 队列中止后，D 晚到的握手仍收不到通知（中，用真实状态机推演确认）

• scheduler.py:5830-5831：bootstrap 队列里的中止只调用 sender.abort()。
• 随后 handle_bootstrap_failure（prefill.py:2385 起）调用 failure_exception()（:2392），把这个 room 的状态清掉。
• CommonKVManager.update_status（common/conn.py:427-435）在 clear() 之后不会再写回 Failed。
• D 的握手晚到时，_notify_drained_abort_target_if_needed（mooncake/conn.py:553）查不到 Failed，不会发 DRAINED_ABORT，D 只能等 waiting 超时。我们之前测过，bootstrap 中止后 D 约 240 s 才结束。
• test_external_kv_abort_drain.py:322 的替身把 statuses[room] 无条件写回 Failed，与真实 update_status 的行为相反，所以单测发现不了这个问题。
• 建议：给已中止的 room 保留一条有界、带 TTL 的终态记录，晚到的握手据此通知，不要在握手可能到达之前就把通知依据清掉；单测改用真实的 manager 状态机。

3. 建议同时处理

• 中止后 P 侧按 room 的状态不清理。 用户中止的新清理路径 abort_external_kv_request（prefill.py:1723-1848）通知 D 之后，不调用 sender.clear()，request_status、transfer_infos、failure_records 都留着。waiting 路径每次中止还会多留一条 failure_records，持续取消会一直累积。但第 (2) 条里“晚到握手补通知”恰好依赖 Failed 条目常驻，所以不能简单 clear，建议和第 (2) 条一起，用有界的终态记录解决。
• P、D 需要同时升级。 本 PR 之后，每次用户取消已握手的请求，P 都会发 4 段的 DRAINED_ABORT 消息（mooncake/conn.py:2129）。还不支持 4 段消息的旧 D 解析时会报错，状态线程退出。建议在发布说明里写明这一点，或者做能力协商。
• 低优先级：
  • 通知是在调度线程上同步发送的。D 不可达、发送队列满时，PUSH socket 最多阻塞 30 s（common/conn.py:992，SNDTIMEO），之后每轮重试（推断）；
  • 中止是按 rid 记录的（_external_kv_abort_rids，prefill.py:1759、:2138），同一个 rid 在排空期间被新请求复用时，可能被误判为待中止。
• 整合分支： 移植到 integrate-dfs-bugfixes-20261020 之后，test_external_kv_abort_drain.py 有 9 个用例报 AttributeError: 'Combined' object has no attribute 'server_args'。原因是整合分支的 process_disagg_prefill_inflight_queue 读取了 self.server_args.enable_unified_cache_external_linker，而测试替身里没有这个字段。PR 头上没有这个问题。

[lzb16]

新提交 8bdd0f61d 复核

上次评论的 (1)–(3) 已经修好，PR 头上 test_external_kv_abort_drain.py（27 passed）和 test_mooncake_final_chunk_poll.py（19 passed）都通过了。还有一处建议，行号按 8bdd0f61d：

bootstrap 线程持有 _drained_abort_lock 时同步发送通知，可能让调度线程跟着等待。

• mooncake/conn.py:3065-3097：bootstrap 线程在 with self._drained_abort_lock: 里处理握手。如果 room 已经是终态，或者需要补发通知，就在锁内调用 _notify_drained_abort_target_if_needed，进而调用 sync_status_to_decode_endpoint（:2140），同步发送 zmq 消息。
• 这个 PUSH socket 设了 SNDTIMEO=30000（common/conn.py:992）。D 不可达、发送队列又满的时候，一次发送最多阻塞 30 s，这期间锁一直被占着。
• 调度线程处理中止时，会经 notify_drained_abort（:3504 起）在末尾调用 finish_drained_abort（:3519），后者要拿同一把锁（:566）。所以 bootstrap 线程卡在发送上时，调度线程也会被拖住，最长约 30 s。
• 上一版这处发送不在锁内，最多卡住 bootstrap 线程本身。调度线程会等这把锁，是这次加锁之后才有的。
• 建议： 锁内只做判断和登记（是否终态、要通知哪些目标、登记 transfer_infos 等），把要发的通知收集起来，出锁之后再调用 sync_status_to_decode_endpoint。重复通知可以照现在的做法，用 _drained_abort_notified_targets 之类的集合去重。