Prhub

#42423 [EC Connector] Add shutdown API to EC Connector.

原始 PR 作者 omerpaz95 合并时间 2026-05-28 20:28 文件变更 5 提交数 5 评论 5 代码增减 +25 / -1

执行摘要

为 EC Connector 添加 shutdown 关闭接口

引自 PR body:'Previously, the scheduler called shutdown() on the KV connector but not the EC connector, leaving no clean teardown path for EC connector implementations that need to drain async operations on exit.'

该 PR 改动量小,设计清晰,值得快速审核合并。对于 EC connector 实现者,建议阅读 ECConnectorBase.shutdown 和 ensure_ec_transfer_shutdown 的使用方式,并在子类中覆盖 shutdown 以处理异步操作排空。

讨论亮点

Review 中 gemini-code-assist 提出两点:

  • shutdown 方法的 docstring 仅提到 scheduler 进程,应泛化以覆盖 worker 角色。
  • worker 侧(如 ModelRunner.shutdown)缺少对应的 shutdown 调用,可能导致异步 I/O 操作未完成就退出。

orozery 后来合并了 worker 侧的调用(通过 GPUWorker.shutdown 而非 ModelRunner),但 docstring 的泛化建议未被采纳。orozery 提出了类型提示的 minor nit(建议 def shutdown(self) -> None:),最终代码已采纳。

实现拆解

  1. 在 ECConnectorBase 中新增 shutdown 方法 (vllm/distributed/ec_transfer/ec_connector/base.py):添加默认的 shutdown(self) -> None 方法,子类可覆盖。
  2. 新增全局关闭函数 (vllm/distributed/ec_transfer/ec_transfer_state.py):添加 ensure_ec_transfer_shutdown(),检查全局 _EC_CONNECTOR_AGENT 状态,若存在则调用其 shutdown() 并置为 None
  3. 导出符号 (vllm/distributed/ec_transfer/__init__.py):将 ensure_ec_transfer_shutdown 加入 __all__
  4. 在 Scheduler 中调用 (vllm/v1/core/sched/scheduler.py):在 Scheduler.shutdown() 中增加 if self.ec_connector is not None: self.ec_connector.shutdown()
  5. 在 GPUWorker 中调用 (vllm/v1/worker/gpu_worker.py):导入 ensure_ec_transfer_shutdown 并在 GPUWorker.shutdown() 中调用。
文件 模块 状态 重要度
vllm/distributed/ec_transfer/ec_transfer_state.py EC 连接器 modified 6.2
vllm/distributed/ec_transfer/ec_connector/base.py EC 连接器 modified 5.5
vllm/v1/worker/gpu_worker.py 工作器 modified 5.4

关键符号

ensure_ec_transfer_shutdown ECConnectorBase.shutdown GPUWorker.shutdown Scheduler.shutdown

关键源码片段

vllm/distributed/ec_transfer/ec_transfer_state.py core-logic

新增核心全局关闭函数 ensure_ec_transfer_shutdown,是 EC connector 清理的入口。

def ensure_ec_transfer_shutdown() -> None:
    global _EC_CONNECTOR_AGENT
    if _EC_CONNECTOR_AGENT is not None:
        # 关闭连接器以排空异步操作
        _EC_CONNECTOR_AGENT.shutdown()
        # 清除全局引用以避免重复关闭
        _EC_CONNECTOR_AGENT = None
vllm/v1/worker/gpu_worker.py dependency-wiring

集成 EC connector 关闭到 GPU Worker 的 shutdown 流程,确保工作进程能清理 EC 连接。

def shutdown(self) -> None:
    # 注意:在解释器关闭期间 ensure_kv_transfer_shutdown 可能为 None
    if ensure_kv_transfer_shutdown is not None:
        ensure_kv_transfer_shutdown()
    if ensure_ec_transfer_shutdown is not None:
        # 关闭 EC 连接器,排空异步传输操作
        ensure_ec_transfer_shutdown()
    if self.profiler is not None:
        self.profiler.shutdown()
    # ... 后续清理步骤

评论区精华

shutdown 方法 docstring 仅提及 scheduler 进程 documentation

gemini-code-assist 指出 shutdown 方法的 docstring 只描述了 scheduler 场景,应涵盖 worker 角色以避免误导。

结论:未采纳,最终代码保持原有 docstring。 · unresolved

Worker 侧缺少 shutdown 调用 正确性

gemini-code-assist 指出在 worker 侧(如 ModelRunner.shutdown)缺少对应的 shutdown 调用,可能导致异步操作未完成就退出。

结论:PR 后续在 GPUWorker.shutdown 中增加了函数调用,但未在 gemini 提及的 ModelRunner 层面修改。任务中需确认是否足够。 · 已解决

风险与影响

变更本质是新增调用路径,不改变现有行为,风险较低。主要风险包括:

  • 若 EC connector 子类未正确实现 shutdown,默认空实现不会报错,可能隐藏资源泄漏。
  • docstring 未泛化,可能使新开发者误解 shutdown 的适用范围。
    建议后续子类覆盖时明确文档。

对用户无直接功能影响,但为 EC connector 实现者提供了标准的清理 hook,有助于提高分布式场景下异步操作完成率和资源释放及时性。对系统而言,进程退出路径更加对称和完善。团队可在自定义 EC connector 中覆盖 shutdown 方法执行自定义清理。

Worker 侧异步操作可能未等待 默认 shutdown 空实现可能隐藏资源泄漏

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论