执行摘要

• 一句话：统一 KV offload 请求完成生命周期契约
• 推荐动作：值得精读。PR 清晰地展示了跨层生命周期契约的设计与统一过程，Review 讨论中包含大量设计权衡（时机选择、状态聚合、命名争议）。对于理解 vLLM 中多层 offloading 的架构和演进非常有帮助。建议关注 RequestState 和 _maybe_finalize_request 的实现细节，以及 reset_cache 的边界处理。

功能与动机

关联 Issue #46027 指出 on_request_finished 在两个 offloading 层间保证不一致：OffloadingManager.on_request_finished 在 GPU→CPU 存储完成后触发，而 SecondaryTierManager.on_request_finished 在 CPU→secondary 仅提交后触发，导致契约混乱。PR body 明确要求统一为“on_request_finished 在层 X = X 不会再收到该请求的 submit 调用”，完成回调可能仍会发生。

实现拆解

1. Scheduler 侧简化（OffloadingConnectorScheduler.request_finished）：恢复为无条件直接调用 manager.on_request_finished()，移除此前 #45823 引入的延迟逻辑；update_connector_output 中也不再回调该 hook。这样保证了“request_finished 后不再有新的 prepare_store/prepare_load”。
2. TieringManager 引入 RequestState：新增 RequestState dataclass（slots=True），包含 req_context、pending_primary_stores 计数、is_finished 标志和 request_level_tiers。替代此前分散的 _finished_req_contexts、_pending_primary_stores 等数据结构和特殊值。
3. 主存储计数与 secondary 通知分离：在 prepare_store 中递增 pending_primary_stores，在 complete_store 中递减，并在递减后调用 _maybe_finalize_request。on_request_finished 只设置 is_finished = True 并尝试 finalize；最终 SecondaryTierManager.on_request_finished 在 pending_primary_stores == 0 && is_finished 时才真正执行。
4. reset_cache 适配：清空未完成的 pending 计数，对已 finished 但仍有 pending 的请求立即 finalize；对未 finished 的请求保留状态，准备后续恢复。
5. 文档更新：更新 OffloadingManager 和 SecondaryTierManager 基类的 on_request_finished docstring，明确契约语义。
6. 测试覆盖：新增 5 个 tiering 测试用例（延迟通知、失败 store、零 store、reset_cache 等）并调整 scheduler 测试验证 on_request_finished 不再延迟。

关键文件：

• vllm/v1/kv_offload/tiering/manager.py（模块 卸载管理；类别 source；类型 core-logic；符号 RequestState, _maybe_finalize_request, on_request_finished, complete_store）: 核心变更文件，引入 RequestState 和 _maybe_finalize_request，实现二次层延迟 finalize 逻辑。
• vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py（模块 连接器调度；类别 source；类型 core-logic；符号 request_finished, update_connector_output, reset_cache）: Scheduler 端解除延迟，立即调用 manager.on_request_finished，简化逻辑。
• tests/v1/kv_offload/tiering/test_tiering_offloading.py（模块 卸载测试；类别 test；类型 test-coverage；符号 _start_request, test_on_request_finished_delays_secondary_until_store_submitted, test_failed_store_finalizes_finished_request, test_zero_store_request_finalizes_immediately）: 新增 5 个测试用例验证各种 finalize 场景，确保新逻辑正确。

关键符号：TieringOffloadingManager._maybe_finalize_request, TieringOffloadingManager.on_request_finished, TieringOffloadingManager.complete_store, OffloadingConnectorScheduler.request_finished, OffloadingConnectorScheduler.update_connector_output

关键源码片段

vllm/v1/kv_offload/tiering/manager.py

核心变更文件，引入 RequestState 和 _maybe_finalize_request，实现二次层延迟 finalize 逻辑。

# vllm/v1/kv_offload/tiering/manager.py (head 版本关键部分 )
@dataclass(slots=True)
class RequestState:
    req_context: ReqContext
    pending_primary_stores: int = 0 # 尚未完成 complete_store 的主存储数
    is_finished: bool = False # on_request_finished 是否已被调用
    request_level_tiers: set[SecondaryTierManager] | None = None # request 级 tier 决策
​
class TieringOffloadingManager:
    def __init__(self, ...):
        self._req_state: dict[str, RequestState] = {} # req_id -> state
        # 移除旧的 _request_level_tiers、_finished_req_contexts 等
​
    def _maybe_finalize_request(self, req_id: str) -> None:
        """若请求已 finished 且无 pending 主存储，则 finalize secondary 层。"""
        state = self._req_state[req_id]
        if not state.is_finished: # 请求尚未 finished 则跳过
            return
        if state.pending_primary_stores > 0: # 仍有待处理主存储
            return
        # 现在通知所有 secondary 层
        for tier in self.secondary_tiers:
            tier.on_request_finished(state.req_context)
        # 清理 request-level 决策
        state.request_level_tiers = None
​
    def on_request_finished(self, req_context: ReqContext) -> None:
        self.primary_tier.on_request_finished(req_context)
        state = self._req_state.get(req_context.req_id)
        if state is None:
            state = RequestState(req_context=req_context)
            self._req_state[req_context.req_id] = state
        state.is_finished = True
        self._maybe_finalize_request(req_context.req_id)
​
    def complete_store(self, keys, req_context, success):
        self.primary_tier.complete_store(keys, req_context, success)
        state = self._req_state.get(req_context.req_id)
        if state:
            # 递减计数（无论成功与否，都表示该 store 已结束）
            state.pending_primary_stores -= 1
            if success:
                # 级联到 secondary 层（递增 ref_cnt 等）
                ...
            self._maybe_finalize_request(req_context.req_id)

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

Scheduler 端解除延迟，立即调用 manager.on_request_finished，简化逻辑。

# vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py (head)
def request_finished(self, request):
    req_status = self._req_status.get(request.request_id)
    if req_status is None:
        # 未跟踪请求：直接创建 context 并完成
        req_context = _create_req_context(request)
        self.manager.on_new_request(req_context)
        self.manager.on_request_finished(req_context)
        return False, None
​
    # 统一立即通知 manager（不再等待传输完成）
    self.manager.on_request_finished(req_status.req_context)
    if not req_status.transfer_jobs:
        del self._req_status[request.request_id]
    else:
        # 保留状态以处理后续 complete_store/complete_load
        for job_id in req_status.transfer_jobs:
            ...
    return False, None
​
# update_connector_output 中不再调用 manager.on_request_finished
def update_connector_output(self, connector_output):
    ...
    req_status.transfer_jobs.remove(job_id)
    if not req_status.transfer_jobs and req_status.req.is_finished():
        del self._req_status[job_status.req_id] # 只删除状态，不触发 hook

tests/v1/kv_offload/tiering/test_tiering_offloading.py

新增 5 个测试用例验证各种 finalize 场景，确保新逻辑正确。

# tests/v1/kv_offload/tiering/test_tiering_offloading.py (head 新增测试示例 )
def test_on_request_finished_delays_secondary_until_store_submitted(self, manager_setup):
    """验证 manager 层 on_request_finished 立即触发，但 secondary 层等待 pending store 完成。"""
    blocks = to_keys(range(2))
    self._start_request()
    self.manager.prepare_store(blocks, _CTX) # pending_primary_stores=1
​
    self.manager.on_request_finished(_CTX) # 设置 is_finished=True
    # secondary 层尚未收到 on_request_finished（pending > 0）
    assert self.secondary_tier1.on_request_finished.call_count == 0
​
    self.manager.complete_store(blocks, _CTX, success=True) # pending--，触发 finalize
    assert self.secondary_tier1.on_request_finished.call_count == 1
​
def test_failed_store_finalizes_finished_request(self, manager_setup):
    """即使 store 失败，pending 计数也应递减并触发 finalize。"""
    self._start_request()
    self.manager.prepare_store(blocks, _CTX) # pending=1
    self.manager.on_request_finished(_CTX) # is_finished=True
    self.manager.complete_store(blocks, _CTX, success=False) # pending--, cascades 不被提交
    assert self.secondary_tier1.on_request_finished.call_count == 1

评论区精华

1. Scheduler 侧简化方向：orozery 建议“恢复到 #45823 之前的逻辑，无条件调用 manager.on_request_finished”，作者采纳并简化。
2. 单一 RequestState 字典：orozery 建议用单个 req_id -> RequestState 替代多个散落的数据结构，作者实现。
3. _maybe_finalize_request 命名：ronensc 提议用 _maybe_finalize_request 突出动作目的，orozery 同意并强调需添加 docstring；最终采纳。
4. reset_cache 中 pending 计数清零：orozery 发现非 finished 请求的 pending 计数在 reset 时应清零，作者修正。
5. 缺少的测试场景：orozery 指出缺少对非 finished 请求在 reset 中存活、complete_store(success=False) 时 finalize 交互、零 store 请求 finalize 的测试；作者补全。
6. pre-commit 修复：多次自动修复 mypy 警告和类型问题。

• Scheduler 侧简化方向 (design): 在 request_finished 中立即调用 manager.on_request_finished。
• 引入 RequestState 替代多个散落结构 (design): 使用 @dataclass(slots=True) RequestState 统一管理。
• 方法命名：_maybe_finish_secondary_tiers vs _maybe_finalize_request (style): 最终采用 _maybe_finalize_request 并添加 docstring。
• reset_cache 中 pending 计数清零 (correctness): reset_cache 中提前清零所有请求的 pending_primary_stores。
• 缺失的测试场景 (testing): 新增三个测试覆盖上述场景。
• complete_store 中 assert 代替条件检查 (style): 使用 assert 来检查 pending_primary_stores > 0。

风险与影响

• 风险：

  1. 竞态条件：修改了 on_request_finished 的触发时机，若 pending_primary_stores 计数不同步（如重复 complete_store），可能导致 secondary 层过早或过晚通知。代码通过 assert 和统一状态机降低风险。
  2. reset_cache 中 pending 清零影响：若 reset 时未正确清零，后续恢复时可能出现计数错误。目前已调整为对所有活跃请求清零。
  3. 向后兼容：仅影响 KV offload 层级内部接口，外部用户无感知；但使用自定义 SecondaryTierManager 的插件可能依赖旧版语义，需更新。
  4. 测试风险：新测试依赖 Python mock 和模拟 tier，可能遗漏真实硬件下的异步竞争。CI 中仅 tiering 测试通过，scheduler 测试因缺少 CUDA 扩展未完整执行。
     - 影响：用户：无直接用户可见变更，但 KV offload 行为更符合预期，减少了因生命周期不一致导致的潜在 resource leak 或级联失败。
     系统：变更集中在 KV offload 链路，尤其是多级 offloading（CPU + 存储/网络）场景。统一契约后，各层可安全地依赖 on_request_finished 释放资源。
     团队：PR 经过多轮 review 和 13 次提交迭代，设计决策清晰，代码可维护性提升。测试覆盖率显著增加，为后续重构（如移除 #45823 遗留逻辑）铺平道路。

• 风险标记：核心路径变更, 需要与 #45823 的衔接, 同步层 mock 测试可能掩盖异步竞态

关联脉络

• PR #46027 [Bug][KV Offload]: on_request_finished() has inconsistent guarantees across offloading layers: 本 PR 直接修复的 Issue，详细描述了契约不一致问题。
• PR #45823 [KV offload] Defer on_request_finished until in-flight transfers drain: 前一个 PR 引入的延迟机制，本 PR 基本上推翻了该方案，采用更清晰的分层契约。