Prhub

#41735 File system secondary tier implemented in python

原始 PR 作者 rshavitt 合并时间 2026-05-25 02:14 文件变更 9 提交数 13 评论 111 代码增减 +980 / -4

执行摘要

新增纯 Python 文件系统 KV 二级卸载层

扩展多级卸载能力,提供一种基于本地文件系统的持久化二级存储选项,使 KV 缓存可以卸载到磁盘,降低 GPU 内存压力。PR body 明确说明添加文件系统 secondary tier。

推荐精读 DualQueueThreadPool 和 JobState 的完成追踪设计,以及 FileSystemTierManager 对 SecondaryTierManager 接口的实现。合并后可在文档中添加 fs_python 配置说明。

讨论亮点
  • 文件拆分:orozery 要求将单文件拆为 manager/thread_pool/io,去下划线,作者采纳。
  • 完成追踪设计:orozery 建议 JobState 和完成队列内置于线程池,作者重写。
  • O_DIRECT 兼容:orozery 指出 macOS 缺失,作者用 getattr 兼容。
  • shutdown 行为:orozery 主张丢弃所有队列,作者实施。
  • 测试分层:orozery 建议单元测试与 e2e 分离,作者调整。
  • 磁盘空间:NUABO 询问 eviction,orozery 回复依赖外部清理。

实现拆解

  1. FileMappervllm/v1/kv_offload/file_mapper.py):将 OffloadKey 映射为文件路径,通过哈希散列到三级子目录,支持 parallel_agnostic 模式。
  2. I/O 原语vllm/v1/kv_offload/tiering/fs/io.py):store_block 写临时文件后原子 rename,load_block 用 os.readv 读入 memoryview,均尝试 O_DIRECT。
  3. DualQueueThreadPooltiering/fs/thread_pool.py):双队列(load/store)双优先线程组,JobState 追踪子任务完成状态。
  4. FileSystemTierManagertiering/fs/manager.py):继承 SecondaryTierManager,组合 FileMapper 和线程池,submit_store/submit_load 批量入队 callable,get_finished 轮询完成队列。
  5. 注册与测试:工厂类注册 'fs_python',配置自动推导 gpu_blocks_per_file。新增两个测试文件覆盖单元功能,扩展现有连接器集成测试。
文件 模块 状态 重要度
vllm/v1/kv_offload/tiering/fs/thread_pool.py 线程池 added 9.04
vllm/v1/kv_offload/tiering/fs/manager.py 管理层 added 9.02
vllm/v1/kv_offload/file_mapper.py 文件映射 added 8.94
vllm/v1/kv_offload/tiering/fs/io.py I/O 层 added 8.84
tests/v1/kv_offload/test_fs_tier.py 测试 added 8.14
tests/v1/kv_offload/test_file_mapper.py 测试 added 7.83
vllm/v1/kv_offload/tiering/factory.py 工厂 modified 5.03

关键符号

JobState.__init__ JobState.task_done DualQueueThreadPool.__init__ DualQueueThreadPool.enqueue_load DualQueueThreadPool.enqueue_store DualQueueThreadPool.get_finished FileSystemTierManager.__init__ FileSystemTierManager.submit_store FileSystemTierManager.submit_load FileSystemTierManager.get_finished FileSystemTierManager.lookup FileSystemTierManager.shutdown FileMapper.__init__ FileMapper.from_offloading_spec FileMapper.get_file_name FileMapper._compute_base_path store_block load_block

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

将初始单文件拆分为多个模块 设计

orozery 要求将初始的单一文件拆分为 manager.py、thread_pool.py、io.py,并去掉类名的下划线前缀。

结论:作者全部采纳,重构为多文件结构。 · 已解决

JobState 位置与完成队列设计 设计

orozery 建议将 JobState 移到 thread_pool 中,并在池内维护完成队列,使 manager 的 get_finished 变得简单。

结论:作者按建议实现,thread_pool 内部维护 _finished_q,manager 直接委托。 · 已解决

O_DIRECT 兼容 macOS 性能

orozery 指出 os.O_DIRECT 在 macOS 上不存在,建议使用 getattr 兼容。

结论:作者添加 getattr(os, 'O_DIRECT', 0) 确保跨平台。 · 已解决

shutdown 时队列行为 正确性

讨论停止时应丢弃 load 还是 store 队列,orozery 认为最佳做法是丢弃所有(best effort)。

结论:实现清空 load 和 store 队列,并 notify_all 使线程退出。 · 已解决

测试分层(unit vs e2e) 测试

orozery 建议纯单元测试放在 test_fs_tier.py,端到端测试放在 test_offloading_connector.py。

结论:作者调整测试文件,保留单元测试并移除原 E2E 类,在连接器测试中添加集成测试。 · 已解决

风险与影响

  • 性能风险:文件系统 I/O 可能在高吞吐下成为瓶颈,O_DIRECT 在 macOS 不可用。
  • 兼容风险:O_DIRECT 仅 Linux,macOS 自动降级。
  • 磁盘空间:无内置 eviction,写满将报错。
  • 临时文件残留:异常 crash 可能残留 .tmp 文件。
  • 路径安全:路径从哈希构建,风险低。
  • 用户影响:可通过配置 type: 'fs_python' 启用,无需额外依赖。
  • 系统影响:新增一个纯 Python 二级存储实现,不影响现有卸载流程。
  • 测试影响:新增 ~400 行测试,覆盖单元和集成场景。
磁盘 I/O 性能瓶颈 O_DIRECT 兼容性 无磁盘 eviction 多线程竞态 临时文件残留

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论