Prhub

#56104 [Bugfix][NIXL] Fix multiple handles xfer race

原始 PR 作者 NickLucche 合并时间 2026-09-15 12:27 文件变更 5 提交数 5 评论 27 代码增减 +233 / -151

执行摘要

NIXL 多句柄传输竞态修复:失败上报与块失效延后

PR body 给出了明确的动机:"A request can own multiple in-flight NIXL transfer handles (one per remote TP rank under heterogeneous TP). On main, the first handle to reach ERR immediately invalidates the request's blocks and reports the failure to the scheduler, while the sibling handles may still be posting RDMA writes into those same blocks. Posted transfers cannot be aborted, so once the scheduler reassigns the blocks, the in-flight writes corrupt whichever request owns them next." 即 NIXL 已投递的传输无法中止,因此只能通过「推迟失败上报与块失效」来规避对复用块的覆写;这也是后续所有级联改动的设计前提。

值得精读。重点看三处:(1) base_worker.py 的失败记账模型与 _try_release_xfer_handle 的「释放失败即保留句柄」语义,这是整个修复成立的基础;(2) get_transfer_results 中失败上报、block 失效、pending notif 撤回三件事被合并到同一次 done_recving 循环,理解这一时序才能评估后续改动;(3) push 侧 done_pushing - failed_pushing 加 _recving_metadata 过滤的双重条件,以及由 njhill 内联建议触发、随后被作者修正的 pull_worker 判空回归。评审阅读时建议顺带确认是否有超时兜底来覆盖「句柄永不终态」的情况,以及 HMA 的 TODO 是否有后续跟进项。

讨论亮点
  • push 侧连带回归(核心交锋):njhill 在 issue 评论中先预警 "I think this change may introduce some problems on the push side though, will add more detail soon.",随后在 push_worker.py 给出内联建议并说明理由:"Otherwise failed sends could be reported as succesful (done_sending) and successful sends could complete early at write handle time bypassing the consumer notif accounting." 作者回复 "Thanks for reviewing @njhill , added your push fix",并移植了 njhill 提供的测试分支(njhill/vllm commit 4143de0)中的两个用例。
  • reviewer 自己的建议引入回归:pull_worker.py 上 njhill 先提 "nit: redundant assert",作者按建议把 handle 判空折进条件写成 if handle is not None and not self._handle_failed_transfer(...);njhill 随后认错:"@NickLucche my bad I think this was my incorrect prior suggestion. This should fix the test ^^"——因为 make_prepped_xfer 在返回句柄前抛错时 handle is None,短路会让 _handle_failed_transfer 完全不被调用,setup 失败被静默吞掉。作者在 commit aa37987 中修正为 if not self._handle_failed_transfer(...): assert handle is not None。
  • 签名与风格 nit:_handle_failed_transfer 的签名被建议收敛为一行 self, req_id: str, handle: int | None, failed_req_ids: set[str] | None = None,已在 commit 63e5257 中应用。
  • 协作摩擦:njhill 表示 "I would have pushed this but don't have permission for your branch.",作者回应会查分支权限;合并前 mergify 两次报 merge conflict,最终在落后 main 52 个提交的情况下以 /ci run --allow-stale 跑 CI,njhill 评论 "Not sure what is going on with the CI here :("。

实现拆解

  1. 失败记账模型重建(vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py):_failed_inflight_recvs 与 _report_failed_recv 被删除,改为单一集合 _recv_failures,语义是「已观察到失败、等待句柄排空」。失败来源分两条通道:后台线程(_background_nixl_handshake 的 request_ready 回调、push worker 的 _on_handshake / _do_send_reg_notif)只往线程安全的 _failed_recv_reqs 队列 put,主线程在 get_transfer_results 里统一消费;主线程自身的失败(pull worker 的 setup 失败、过期读)直接写 _recv_failures 集合,避免队列开销,这也是 PR body 中 "main engine thread failures are reported on _recv_failures set to minimize overhead" 的取舍。
  2. 轮询函数纯化:_pop_done_transfers 返回值从 set[str] 变为 tuple[set[str], set[str]](无未决句柄的请求、本轮观察到失败的请求),并移除了内部的 _send_pending_recv_notifs / _report_failed_recv 副作用。失败句柄只有在 release_xfer_handle 真正成功后才退出跟踪,否则回填 in_progress 留待下轮重试;轮询抛异常也不再被当作「句柄已终态」。新增的 _try_release_xfer_handle 封装了这一语义。纯化后的函数被 pull 侧用于轮询 _recving_transfers,也被 push 侧用于轮询 _sending_transfers。
  3. 上报与失效时机前移/后移:get_transfer_results 里先把 _pop_done_transfers 返回的失败并入 _recv_failures,再逐条判断——请求元数据已不存在则静默清除;元数据存在但 _recving_metadata/_recving_transfers 中已无未决句柄,才移入 failed_recv_reqs 并合入 done_recving。block 失效从原来的 _report_failed_recv 搬到 done_recving 循环内,与 finished_recving 上报同一次 get_finished 生效;失败请求同时 pop 掉 _pending_recv_notifs,避免向 producer 发出成功通知。成功请求的 _send_pending_recv_notifs 也一并搬到该循环。
  4. pull 侧调用点适配(pull_worker.py):_read_blocks 的 transfer_setup_failed 分支改为 if not self._handle_failed_transfer(request_id, handle, self._recv_failures): assert handle is not None; self._recving_transfers[request_id].append(handle),并把日志文案改成 "Deferring failure reporting until outstanding transfers finish";_read_blocks_mixed 的两处清理路径改用 _try_release_xfer_handle,释放失败的句柄同样回填跟踪表;_is_turn2_read_expired 的过期读直接传 self._recv_failures。
  5. push 侧一致化(push_worker.py):_xfer_blocks 的 WRITE 提交失败改用 _handle_failed_transfer(释放失败则把 handle 交还调用方继续跟踪);get_transfer_results 接收新的二元返回值并做过滤 done_pushing = {req_id for req_id in done_pushing - failed_pushing if req_id in self._recving_metadata},保证失败发送绝不会作为 done_sending 上报(其 block 交给 lease/watchdog 回收),也不会在写句柄完成时就绕过 consumer notif 记账。
  6. 测试与契约调整:tests/v1/kv_connector/unit/test_nixl_connector.py 新增 recv_worker fixture、test_mixed_memory_read_failure_does_not_notify_producer、参数化的 test_recv_failure_waits_for_sibling_transfer、test_recv_failure_waits_for_unpollable_handle_to_be_released,并新增 _ScriptedXferWrapper(其 release_xfer_handle 断言生产代码永不释放 PROC 句柄);test_handshake_failure_returns_finished、test_transfer_setup_failure_returns_finished 改成「get_finished 之前 get_block_ids_with_load_errors() 为空、之后才出现失效块」,把新的延迟失效契约固化下来。tests/v1/kv_connector/unit/test_nixl_push_connector.py 新增 _pollable_worker / _make_sending_req 及三个 push 完成语义测试(移植自 njhill 的 review 分支),并把原 test_failed_push_transfer_finishes_sending 改名为 test_failed_push_transfer_does_not_evict_writer_state。
文件 模块 状态 重要度
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py 传输基类 modified 8.57
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py 推送发送端 modified 6.43
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py 拉取接收端 modified 6.37
tests/v1/kv_connector/unit/test_nixl_connector.py 连接器测试 modified 7.52
tests/v1/kv_connector/unit/test_nixl_push_connector.py 推送端测试 modified 6.96

关键符号

_pop_done_transfers _handle_failed_transfer _try_release_xfer_handle get_transfer_results _read_blocks _read_blocks_mixed get_finished _background_nixl_handshake _on_handshake

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py core-logic

修复的核心:重写失败记账(`_recv_failures` 取代 `_failed_inflight_recvs`/`_report_failed_recv`)、把 `_pop_done_transfers` 改成保留句柄的纯轮询函数、新增 `_try_release_xfer_handle`,并把失败上报与 block 失效推迟到句柄排空后的 `get_transfer_results` 内。

# vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
​
    def _try_release_xfer_handle(self, req_id: str, handle: int) -> bool:
        """Release a handle, returning False if the caller must retain it."""
        try:
            self.nixl_wrapper.release_xfer_handle(handle)
        except Exception as e:
            # 释放失败并不代表后端已经停下 DMA:句柄必须继续留在
            # _recving_transfers 里被轮询,绝不能就此把 block 交还调度器,
            # 否则仍在途的 RDMA 写入会污染下一个占用这些 block 的请求。
            self._log_failure(
                failure_type="transfer_release_failed",
                msg="Retaining handle and blocks until release succeeds",
                req_id=req_id,
                error=e,
            )
            return False
        return True
​
    def _handle_failed_transfer(
        self, req_id: str, handle: int | None, failed_req_ids: set[str] | None = None
    ) -> bool:
        """Record a failure and release its handle, returning False to retain it."""
        self.xfer_stats.record_failed_transfer()
        # 只做「记账 + 尝试释放」两件事:不再操作 _pending_recv_notifs,也不做
        # block 失效,所以 push writer 等后台线程可以直接调用而不触碰主线程状态。
        if failed_req_ids is not None:
            failed_req_ids.add(req_id)
        # handle 为 None(握手或 setup 阶段就失败)时无需保留,直接返回 True。
        return handle is None or self._try_release_xfer_handle(req_id, handle)
​
    def _pop_done_transfers(
        self, transfers: dict[str, list[int]]
    ) -> tuple[set[str], set[str]]:
        """Poll transfers, retaining handles until they can be released.
​
        Returns:
            Requests with no outstanding handles and requests with observed
            failures. Failed requests may still have outstanding handles.
        """
        done_req_ids: set[str] = set()
        failed_req_ids: set[str] = set()
        for req_id, handles in list(transfers.items()):
            in_progress = []
            for handle in handles:
                try:
                    xfer_state = self.nixl_wrapper.check_xfer_state(handle)
                    if xfer_state == "DONE":
                        res = self.nixl_wrapper.get_xfer_telemetry(handle)
                        self.xfer_stats.record_transfer(res)
                        self.nixl_wrapper.release_xfer_handle(handle)
                    elif xfer_state == "PROC":
                        in_progress.append(handle)
                    else:
                        # ERR 等终态:登记失败,但只有句柄确实释放成功才认为它
                        # 已退出跟踪;释放失败则回填,下一轮继续轮询。
                        self._log_failure(
                            failure_type="transfer_failed",
                            req_id=req_id,
                            xfer_state=xfer_state,
                        )
                        if not self._handle_failed_transfer(
                            req_id, handle, failed_req_ids
                        ):
                            in_progress.append(handle)
                except Exception as e:
                    # 轮询异常不再当作「句柄已终态」,同样先尝试释放、失败即保留。
                    self._log_failure(
                        failure_type="transfer_exception", req_id=req_id, error=e
                    )
                    if not self._handle_failed_transfer(req_id, handle, failed_req_ids):
                        in_progress.append(handle)
​
            if not in_progress:
                del transfers[req_id]
                done_req_ids.add(req_id)
            else:
                # 仍有在途(或释放失败)的句柄:请求既不完成也不上报失败。
                transfers[req_id] = in_progress
        return done_req_ids, failed_req_ids
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py core-logic

njhill 指出本改动会破坏 push 侧完成语义并给出内联修复:`get_transfer_results` 需用 `failed_pushing` 与 `_recving_metadata` 双重过滤,避免失败发送被上报为 `done_sending`、成功发送在写句柄完成时绕过 consumer notif 记账;同时把握手 / 注册通知失败改为投递 `_failed_recv_reqs` 队列。

# vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
​
    def get_transfer_results(self) -> KVConnectorTransferResults:
        # 引擎主线程来取完成情况时顺带唤醒 writer,让它有机会排空 NIXL 通知
        # (heartbeat、completion notif、迟到的 PUSH_REG),即使此前已 park。
        self._push_writer_wake.set()
​
        results = super().get_transfer_results()
        done_sending = results.finished_sending
​
        # _pop_done_transfers 会原地修改 _sending_transfers,而 writer 线程
        # 也会往里 append,所以必须在锁内轮询。
        with self._sending_transfers_lock:
            done_pushing, failed_pushing = self._pop_done_transfers(
                self._sending_transfers
            )
        # 失败的发送绝不能上报为 done_sending:上报会让 P 侧误以为 KV 已送达而
        # 提前释放 block,而 D 侧其实还在等数据;这类请求留给 lease / watchdog。
        # 同时只有仍被 _recving_metadata 跟踪的请求才允许在写句柄完成时上报,
        # 否则会绕过 consumer notif 计数提前 complete。
        done_pushing = {
            req_id
            for req_id in done_pushing - failed_pushing
            if req_id in self._recving_metadata
        }
        for req_id in done_pushing:
            self._reqs_to_send.pop(req_id, None)
            self._reqs_to_process.discard(req_id)
            self.consumer_notification_counts_by_req.pop(req_id, None)
            done_sending.add(req_id)
​
        # 通知 writer 丢弃已完成(push 成功)或已过期(lease 到期仍无 D 注册)
        # 的请求状态。
        for req_id in done_sending:
            self._evict_finished_inbox.put(req_id)
        if done_sending:
            self._push_writer_wake.set()
​
        return results

评论区精华

push 侧失败发送被误报为 done_sending、成功发送提前完成 正确性

njhill 先预警 "I think this change may introduce some problems on the push side though",随后在 `push_worker.get_transfer_results` 给出内联建议并解释:"Otherwise failed sends could be reported as succesful (done_sending) and successful sends could complete early at write handle time bypassing the consumer notif accounting." 同时附上自己分支的测试(njhill/vllm commit 4143de0)。作者回复 "added your push fix" 并移植测试。

结论:已采纳:`done_pushing - failed_pushing` 再按 `req_id in self._recving_metadata` 过滤,失败发送交 lease/watchdog 处理;补齐 push 完成语义测试,njhill 随后 APPROVED。 · 已解决

pull_worker setup 失败分支的判空与「冗余 assert」 正确性

njhill 先提 "nit: redundant assert",作者据此把 `handle is not None` 折进条件;njhill 随后认错 "@NickLucche my bad I think this was my incorrect prior suggestion. This should fix the test ^^",因为 `make_prepped_xfer` 在返回句柄前抛错时 `handle is None`,短路会导致 `_handle_failed_transfer` 根本不被调用、setup 失败被漏记。

结论:最终形态为 `if not self._handle_failed_transfer(request_id, handle, self._recv_failures): assert handle is not None; self._recving_transfers[request_id].append(handle)`,并在 commit aa37987(Co-authored-by: OpenAI Codex)修正。作者另回复会排查 fork 分支权限。 · 已解决

`_handle_failed_transfer` 签名收敛 style

njhill 给出 nit 建议,把参数列表折叠为 `self, req_id: str, handle: int | None, failed_req_ids: set[str] | None = None`,与新增的 `failed_req_ids` 出参语义保持一致。

结论:已在 commit 63e5257 中应用(commit message 中记录为 "collapse _handle_failed_transfer signature")。 · 已解决

分支权限与 CI 组织 other

njhill 表示 "I would have pushed this but don't have permission for your branch.";mergify 两次报 merge conflict;最终在落后 main 52 个提交的情况下以 `/ci run --allow-stale` 触发 CI,njhill 评论 "Not sure what is going on with the CI here :("。

结论:实现层面无影响,但暴露了 fork 协作与陈旧 CI 组合下的验证盲区,合并时未再补齐最新主干验证。 · unresolved

风险与影响

  1. 失败上报可能被无限期推迟:新逻辑要求 req_id 完全不在 _recving_transfers 中才上报失败。若某个 sibling 句柄长期停留在 PROC,或 release_xfer_handle 持续抛错(_try_release_xfer_handle 会一直返回 False),该请求既不会失效 block 也不会出现在 finished_recving,调度器只能一直等。修复前是「立即失败」,修复后是「可能挂起」,diff 中未见超时或 watchdog 兜底,依赖 NIXL 最终进入终态。
  2. 失败被静默丢弃的场景:get_transfer_results 中若 req_id not in self._recving_metadata,失败记录被直接 remove 且不加入 failed_recv_reqs。需要确认这只会发生在请求已成功结束之后,否则会出现「失败无任何上报」的静默路径。
  3. HMA 缺口被保留:if not self._is_hma_required: self._invalid_block_ids.put(...) 旁边保留了 # TODO (NickLucche) handle failed transfer for HMA.,HMA 场景下失败请求仍以 finished_recving 上报但不失效 block,属于已知未覆盖范围。
  4. 跨线程约束变隐式:_handle_failed_transfer 被刻意裁成「只记统计 + 尝试释放」,push writer 线程调用时不传 failed_req_ids 以避免跨线程写集合,后台线程一律走 _failed_recv_reqs 队列。这一约束只写在注释里,后续若有人往该函数里加回主线程状态操作(例如恢复 _pending_recv_notifs 清理),竞态会重新出现。
  5. push 测试可能空断言:_pollable_worker() / _StubWriterWorker.fresh() 的 _recving_metadata 为空字典,而 push 的 done_pushing 过滤条件包含 if req_id in self._recving_metadata;因此 test_failed_push_send_is_not_reported_done_sending、test_completed_push_send_waits_for_consumer_notif 中「不在 finished_sending」的断言即使失败过滤逻辑失效也会成立,测试对本地句柄完成路径的约束力弱于预期(正向的 consumer notif 路径仍有覆盖)。
  6. 合并流程风险:PR 在落后 main 52 个提交、CI 多次冲突重跑的情况下以 --allow-stale 合并,njhill 也表达了对 CI 状态的困惑,存在「CI 未覆盖最新主干」的验证盲区。
  • 用户/正确性:PD 分离与异构 TP 场景下,KV block 复用不再被在途 RDMA 写入污染,消除一类难以复现的「输出错乱/串味」问题,属于静默数据损坏级修复。
  • 系统行为:三个 worker(base/pull/push)的失败语义统一为「先排空句柄、再上报失败、与 get_finished 同步失效 block」;_pop_done_transfers 从带副作用的完成判定变成可复用的纯轮询函数,push 侧因此也能用它轮询 _sending_transfers。
  • 失败恢复代价:错误路径的恢复时延从「立即」变为「等到最后一个句柄终态」,在 NIXL 卡死场景下可能表现为请求长期未完成而非快速降级重算。
  • 团队协作:确立了「后台线程只投队列、主线程集中记账、释放失败即保留句柄」的约定,是一次可复用的并发状态治理范式;同时也暴露了 review suggestion 直接落地需要连带补测试的经验。
核心传输路径变更 失败上报延迟且无超时兜底 HMA 路径未失效 block push 侧测试可能空断言 落后 main 合并

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论