执行摘要
- 一句话:延迟预占用重入请求的 KV 加载直到传输完成
- 推荐动作:值得精读,展示了在异步调度框架下如何优雅处理预占用与传输作业重叠的边界情况。
功能与动机
Issue #46014 报告 OffloadingConnector 在预占用请求被重入时断言失败:assert not req_status.transfer_jobs。在异步调度中,存储作业可能已被 worker 刷新但调度器尚未消费完成输出,请求在此窗口内被重新接纳时,调度器不应立即发起加载。
实现拆解
- 在 OffloadingConnectorScheduler.get_num_new_matched_tokens 开头,清空 block_ids 后,检查 req_status.transfer_jobs 是否非空。
- 若非空,记录调试日志并返回 (None, False),告知调度器需要更多时间,稍后重试。
- 返回 None 会阻止后续的 update_state_after_alloc 中的加载路径,避免冲突。
- 新增两个测试: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(模块 调度器;类别 source;类型 core-logic;符号 get_num_new_matched_tokens): 核心修复文件,在 get_num_new_matched_tokens 入口增加传输进行中检查,避免加载冲突。
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py(模块 测试;类别 test;类型 test-coverage;符号 test_pending_transfer_defers_prefix_lookup, test_async_preempt_readmit_before_transfer_output_is_deferred): 新增两个回归测试,覆盖单元和集成场景,确保延迟行为。
关键符号: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
新增两个回归测试,覆盖单元和集成场景,确保延迟行为。
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 中的评论,随后作者调整后获得批准。无其他争议。
- 预占用重入时传输作业未完成的处理 (design): 作者实现了延迟加载的修复,通过检查 transfer_jobs 返回 None 避免冲突。
风险与影响
- 风险:改动极小(+7/-0),但返回 (None, False) 后调度器会稍后重试,可能略微增加请求响应时间。需确保调用者正确处理 None(已有处理)。
- 影响:影响范围仅限于 OffloadingConnector 调度器,修复了一个可能导致进程崩溃的竞态条件 bug,提升稳定性。对正常路径无影响。
- 风险标记:核心路径变更, 需确保调用者正确处理 None
关联脉络
参与讨论