执行摘要

• 一句话：为 on_schedule_end 传递调度上下文
• 推荐动作：值得精读。该 PR 展示了在大型推理框架中如何安全地扩展核心钩子参数：通过引入 NamedTuple 显式传递上下文、使用 Collection 接口保证迭代安全、全量更新所有实现以维持一致性。其设计讨论（dataclass vs NamedTuple、Iterable vs Collection）对类似场景有参考价值。

功能与动机

在 KV offload 场景中，不同调度步骤可能需要针对新请求或抢占请求做出差异化处理（如立即刷新 store 或调整优先级）。原先 on_schedule_end() 不携带上下文，管理器无法区分当前步骤的调度事件。通过 ScheduleEndContext 明确传入 new_req_ids 和 preempted_req_ids，使各 tier 管理器能够智能响应特定调度决策。

实现拆解

实现步骤如下：

1. 定义调度上下文数据结构（vllm/v1/kv_offload/base.py）：新增 ScheduleEndContext 作为 NamedTuple，包含两个 Collection[str] 字段 —— new_req_ids 和 preempted_req_ids。字段类型选择 Collection 以同时支持 list、tuple、set 等可重复消费的集合。

2. 修改核心钩子签名（base.py 中 OffloadingManager.on_schedule_end()）：为方法增加 context: ScheduleEndContext 参数，并更新 docstring。

3. 更新二级 Tier 抽象基类（tiering/base.py）：SecondaryTierManager.on_schedule_end() 同样增加 context 参数。

4. 更新所有具体实现：修改 tiering/manager.py、fs/manager.py、obj/manager.py、p2p/manager.py 中的对应 on_schedule_end 方法，使其正确接受并传递 ScheduleEndContext。在 tiering/manager.py 中还将 context 依次传递给所有 secondary_tiers。

5. 调度侧构造与传递（vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py）：在 OffloadingConnectorScheduler.build_connector_meta() 中，根据 SchedulerOutput 提供的 scheduled_new_reqs 和 preempted_req_ids 构造 ScheduleEndContext 实例，然后调用 self.manager.on_schedule_end(schedule_end_context)。

6. 测试配套更新：更新 tests/v1/kv_offload/tiering/p2p/test_manager.py、tests/v1/kv_offload/tiering/test_tiering_offloading.py、tests/v1/kv_offload/tiering/test_fs_tier.py、tests/v1/kv_offload/tiering/test_obj_tier.py 等测试文件，适配新的方法签名。

关键文件：

• vllm/v1/kv_offload/base.py（模块 离载基类；类别 source；类型 core-logic；符号 ScheduleEndContext, on_schedule_end）: 核心抽象层，定义 ScheduleEndContext 并修改 on_schedule_end 接口签名，是整个变更的入口。
• vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py（模块 离载调度；类别 source；类型 core-logic）: 调度器侧构造 ScheduleEndContext 并调用 on_schedule_end，是变更的实际使用点。
• vllm/v1/kv_offload/tiering/manager.py（模块 分层管理；类别 source；类型 core-logic；符号 on_schedule_end）: 主要管理器 CPUPrimaryTierOffloadingManager 实现 on_schedule_end，并在内部将 context 转发给所有二级 tier。
• vllm/v1/kv_offload/tiering/base.py（模块 二级离载；类别 source；类型 core-logic；符号 on_schedule_end）: SecondaryTierManager 抽象基类，定义 on_schedule_end 接口，保证所有派生类一致性。
• vllm/v1/kv_offload/tiering/fs/manager.py（模块 文件离载；类别 source；类型 core-logic；符号 on_schedule_end）: 文件系统二级 tier 实现，同步更新 on_schedule_end 签名并调用 flush。
• vllm/v1/kv_offload/tiering/obj/manager.py（模块 对象离载；类别 source；类型 core-logic；符号 on_schedule_end）: 对象存储二级 tier 实现，同步更新 on_schedule_end 签名并调用 flush。
• vllm/v1/kv_offload/tiering/p2p/manager.py（模块 点对点离载；类别 source；类型 core-logic；符号 on_schedule_end）: 点对点二级 tier 实现，同步更新 on_schedule_end 签名。
• tests/v1/kv_offload/tiering/p2p/test_manager.py（模块 P2P 测试；类别 test；类型 test-coverage）: P2P 管理器测试，更新测试代码适配新签名，验证基本功能不退化。
• tests/v1/kv_offload/tiering/test_tiering_offloading.py（模块 分层测试；类别 test；类型 test-coverage）: 分层 offloading 管理器测试，适配新签名，保证级联转发逻辑正确。
• tests/v1/kv_offload/tiering/test_fs_tier.py（模块 文件离载测试；类别 test；类型 test-coverage）: 文件系统 tier 测试，适配新签名。
• tests/v1/kv_offload/tiering/test_obj_tier.py（模块 对象离载测试；类别 test；类型 test-coverage）: 对象存储 tier 测试，适配新签名。

关键符号：on_schedule_end, ScheduleEndContext, build_connector_meta

关键源码片段

vllm/v1/kv_offload/base.py

核心抽象层，定义 ScheduleEndContext 并修改 on_schedule_end 接口签名，是整个变更的入口。

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
​
from typing import NamedTuple
from collections.abc import Collection
​
class ScheduleEndContext(NamedTuple):
    # Per-step scheduling info passed to on_schedule_end()
    # 包含本轮调度中新请求和抢占请求的 ID 集合 .
    # 本轮首次调度的请求 ID (request IDs scheduled for the first time this step).
    new_req_ids: Collection[str]
    # 本轮被抢占的请求 ID (request IDs preempted this step).
    preempted_req_ids: Collection[str]
​
​
class OffloadingManager(ABC):
​
    def on_schedule_end(self, context: ScheduleEndContext) -> None:
        # Called once at the end of each scheduler step.
        # Managers may override this to flush deferred work accumulated
        # during the step (e.g., batched promotions).
        return

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

调度器侧构造 ScheduleEndContext 并调用 on_schedule_end，是变更的实际使用点。

def build_connector_meta(
    self, scheduler_output: SchedulerOutput
) -> KVConnectorMetadata:
    self._update_req_states(scheduler_output)
​
    # 构造 ScheduleEndContext，从 scheduler_output 中提取新请求 ID 和抢占请求 ID.
    schedule_end_context = ScheduleEndContext(
        new_req_ids=[
            req.req_id for req in scheduler_output.scheduled_new_reqs
        ],
        preempted_req_ids=scheduler_output.preempted_req_ids or (),
    )
    self.manager.on_schedule_end(schedule_end_context)
​
    # 处理被抢占请求的 job 刷出
    for req_id in scheduler_output.preempted_req_ids or ():
        req_status = self._req_status.get(req_id)
        if req_status is None or not req_status.transfer_jobs:
            continue
        any_jid = next(iter(req_status.transfer_jobs))
        assert self._jobs[any_jid].is_store
        self._current_batch_jobs_to_flush.update(req_status.transfer_jobs)
​
    # ... 后续逻辑不变 ...

评论区精华

Review 中的核心讨论包括：

• 设计选择：orozery 提议将 ScheduleEndContext 从 dataclass 改为 NamedTuple，ronensc 采纳（commit c283e7eb）。
• 类型安全：orozery 担心 Iterable 在多 tier 遍历时会被耗尽，建议使用 Collection 或 Sequence；ronensc 指出 preempted_req_ids 来源是 set，与 Sequence 不兼容，最终采用 Collection。
• 性能微调：关于 new_req_ids 应使用 list 还是 tuple，ronensc 给出 microbenchmark 数据，证明 list 构造开销更小，且单步请求数量有限，最终保留 list。

• 合并冲突修复：orozery 发现冲突解决时误删了 from enum import auto，ronensc 确认并提交修复。

  • ScheduleEndContext 类型：dataclass vs NamedTuple (design): 使用 NamedTuple。
  • 字段类型：Iterable vs Collection vs Sequence (design): 使用 Collection[str] 兼容 set/list/tuple。
  • 数据结构性能：list vs tuple (performance): 使用 list 构造 new_req_ids。
  • 合并冲突导致 auto 导入丢失 (other): 恢复 import auto。

风险与影响

• 风险：
  1. 接口兼容性风险：所有继承 OffloadingManager 或 SecondaryTierManager 的自定义子类必须同步更新 on_schedule_end 签名，否则运行时触发 TypeError。内部实现已全部更新，但仓库外部署可能遗漏。
  2. 测试覆盖不足：现有测试仅验证签名正确性，未充分验证 ScheduleEndContext 内容在各管理器中的实际消费逻辑，存在回归盲点。
  3. 性能风险：每步调度构造 ScheduleEndContext 对象和列表，在高并发下可能产生短暂内存分配，但相对整体 offload 开销可忽略。
• 影响：
  • 对用户：KV offload 功能用户无直接行为变化，但为未来更智能的缓存策略提供了原料。
  • 对系统：offload 管理器现可感知新请求和抢占，有望提升内存利用率和降低延迟。
  • 对团队：所有 offload 模块统一更新，降低后续二次开发成本。
  • 风险标记：接口变更, 测试覆盖不足, 自定义实现兼容性

关联脉络

• 暂无明显关联 PR