Prhub

#46231 [Bugfix] Defer offload reads while transfers are pending

原始 PR 作者 Palaiologos1453 合并时间 2026-06-21 19:30 文件变更 2 提交数 2 评论 3 代码增减 +106 / -4

执行摘要

延迟预占用重入请求的 KV 加载直到传输完成

Issue #46014 报告 OffloadingConnector 在预占用请求被重入时断言失败:assert not req_status.transfer_jobs。在异步调度中,存储作业可能已被 worker 刷新但调度器尚未消费完成输出,请求在此窗口内被重新接纳时,调度器不应立即发起加载。

值得精读,展示了在异步调度框架下如何优雅处理预占用与传输作业重叠的边界情况。

讨论亮点

Reviewer @orozery 在第一次审核中要求查看 Issue #46014 中的评论,随后作者调整后获得批准。无其他争议。

实现拆解

  1. 在 OffloadingConnectorScheduler.get_num_new_matched_tokens 开头,清空 block_ids 后,检查 req_status.transfer_jobs 是否非空。
  2. 若非空,记录调试日志并返回 (None, False),告知调度器需要更多时间,稍后重试。
  3. 返回 None 会阻止后续的 update_state_after_alloc 中的加载路径,避免冲突。
  4. 新增两个测试:test_pending_transfer_defers_prefix_lookup(纯单元)和 test_async_preempt_readmit_before_transfer_output_is_deferred(集成,使用 request_runner),验证延迟行为。
文件 模块 状态 重要度
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py 调度器 modified 6.13
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py 测试 modified 6.28

关键符号

get_num_new_matched_tokens test_pending_transfer_defers_prefix_lookup test_async_preempt_readmit_before_transfer_output_is_deferred

关键源码片段

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py test-coverage

新增两个回归测试,覆盖单元和集成场景,确保延迟行为。

def test_pending_transfer_defers_prefix_lookup():
    """验证当请求存在进行中的 store 时,lookup 被推迟且 manager.lookup 不被调用。"""
    scheduler = object.__new__(OffloadingConnectorScheduler)
    scheduler.manager = MagicMock(spec=OffloadingManager)
​
    # 模拟一个请求,具有 transfer_jobs 集
    request = SimpleNamespace(request_id="req-0")
    group_state = SimpleNamespace(block_ids=[1, 2, 3])
    req_status = SimpleNamespace(
        group_states=[group_state],
        transfer_jobs={123}, # 非空的传输作业集
    )
    scheduler._req_status = {request.request_id: req_status}
​
    # 执行 get_num_new_matched_tokens
    matched_tokens, is_async = scheduler.get_num_new_matched_tokens(
        request, num_computed_tokens=0,
    )
​
    # 验证返回 None 和 False,block_ids 被清空,manager.lookup 未被调用
    assert matched_tokens is None
    assert is_async is False
    assert group_state.block_ids == []
    scheduler.manager.lookup.assert_not_called()

评论区精华

预占用重入时传输作业未完成的处理 设计

Reviewer @orozery 要求查看 Issue #46014 中的评论("Please see my comment on the issue"),作者随后提交改动后获得批准。

结论:作者实现了延迟加载的修复,通过检查 transfer_jobs 返回 None 避免冲突。 · 已解决

风险与影响

改动极小(+7/-0),但返回 (None, False) 后调度器会稍后重试,可能略微增加请求响应时间。需确保调用者正确处理 None(已有处理)。

影响范围仅限于 OffloadingConnector 调度器,修复了一个可能导致进程崩溃的竞态条件 bug,提升稳定性。对正常路径无影响。

核心路径变更 需确保调用者正确处理 None

关联 Issue

#46014 [Bug]: OffloadingConnector crashes with `assert not req_status.transfer_jobs` when a preempted request is re-admitted

完整报告

参与讨论