Prhub

#53293 [Bugfix] Set breakable graph env before Ray actor import

原始 PR 作者 alexeldeib 合并时间 2026-08-29 07:11 文件变更 4 提交数 5 评论 7 代码增减 +51 / -2

执行摘要

修复 Ray 启动顺序导致 breakable CUDA 图失效与输出错乱

PR body 明确指出根因:DSv4 automatically enables breakable CUDA graphs in VllmConfig. Native multiprocessing workers inherit that setting before importing the model, but Ray creates its worker actors first and copies the driver environment only later in initialize_worker. 由于 @eager_break_during_capture 在模型模块 import 时做决策,设置缺席时装饰器永久返回原始 attention 和 indexer 函数,Ray 后续虽启用 breakable capture,但缺少断点导致超大图与乱码输出。作者实测未修复 Ray 自动启用时抓图 0.69 GiB、颜色与数学问答 0/10;修复后 0.37 GiB、10/10。

值得精读。该 PR 展示了一个典型的 Ray 环境变量传播时序与模块导入期求值之间的隐式契约问题,修复方式最小侵入且兼顾可扩展性。建议重点关注:@eager_break_during_capture 的导入期求值语义、frozenset 白名单的扩展方式,以及 runtime_env 注入与 initialize_worker 延迟应用两种机制的适用边界。

讨论亮点
  • jeffreywang88 提出通用化建议并对比两种机制:"I'm wondering if we could make this more general. For example, maintain a list of environment variables that should always be set before actor creation, and add VLLM_USE_BREAKABLE_CUDAGRAPH to that list." 他同时指出:现有机制(#36836)由 initialize_worker 在已运行的 actor 内应用环境变量,属于延迟通用传播;本 PR 是在 actor 创建前注入,属于提前、面向导入期的传播。
  • alexeldeib 回应:"I initially did it as a frozenset and iterate over the list. but since I only found 1 that mattered, I avoided adding more." 并确认未发现其它需要同样处理的 env var。
  • jeffreywang88 最终拍板:"let's go ahead with the frozenset approach to be more general and easier to extend to other env vars going forward." alexeldeib 随即修改提交。
  • njhill 认可修复方向但请求 Ray 侧专家复核:"this looks reasonable to me but it would be good to get approval from someone more familiar with the ray side",最终 jeffreywang88 与 njhill 均 APPROVED。

实现拆解

  1. 新增白名单与注入函数:在 vllm/v1/executor/ray_env_utils.py 中定义 _RAY_WORKER_PRE_IMPORT_ENV_VARS = frozenset({"VLLM_USE_BREAKABLE_CUDAGRAPH"}),并新增 update_runtime_env_for_worker_import(runtime_env),遍历白名单,仅当变量存在于 driver os.environ 时才写入 runtime_env["env_vars"],原地修改并返回同一 dict。设计为白名单而非单变量直写,便于后续扩展其它"导入期必需"的环境变量。
  2. V2 executor 接入:在 vllm/v1/executor/ray_executor_v2.py 的 RayExecutorV2._build_runtime_env 中,写完 current_platform.ray_noset_device_env_vars 后调用 update_runtime_env_for_worker_import(runtime_env)。该函数在 actor 创建前执行,保证模型模块导入时变量可见。
  3. V1 executor 接入:在 vllm/v1/executor/ray_executor.py 的 _update_noset_device_env_vars 中追加同样的调用,保证旧 Ray executor 路径行为一致。
  4. 测试配套:tests/test_ray_env_utils.py 新增 TestImportTimePropagation 两个用例——test_breakable_cudagraph_is_available_during_actor_import 验证变量注入且保留既有 env_vars;test_other_env_vars_are_unchanged 用 patch(clear=True) 模拟仅含无关变量的 driver 环境,验证白名单外变量完全不进入 runtime_env。
  5. 演进过程:首个提交为单变量直写,经 review 后按 jeffreywang88 建议泛化为 frozenset 白名单,期间两次合入 main 保持同步,最终由 njhill 合入。
文件 模块 状态 重要度
vllm/v1/executor/ray_env_utils.py 环境工具 modified 6.73
tests/test_ray_env_utils.py 环境测试 modified 5.86
vllm/v1/executor/ray_executor_v2.py 执行器 modified 5.19
vllm/v1/executor/ray_executor.py 执行器 modified 4.67

关键符号

update_runtime_env_for_worker_import

关键源码片段

vllm/v1/executor/ray_env_utils.py core-logic

核心变更文件:新增 _RAY_WORKER_PRE_IMPORT_ENV_VARS 白名单与 update_runtime_env_for_worker_import 注入函数,是本 PR 修复逻辑的主体。

# vllm/v1/executor/ray_env_utils.py
import os
​
from vllm.ray.ray_env import RAY_NON_CARRY_OVER_ENV_VARS
​
# 需要在 Ray actor 导入 vLLM 之前就可见的 driver 环境变量白名单。
# 目前只有 breakable CUDA graph 开关一项;未来若有其它“模型模块 import
# 时才读取”的变量,直接加入该集合即可,无需改动调用方。
_RAY_WORKER_PRE_IMPORT_ENV_VARS = frozenset({"VLLM_USE_BREAKABLE_CUDAGRAPH"})
​
​
def get_driver_env_vars(worker_specific_vars: set[str]) -> dict[str, str]:
    """返回需要传播给 Ray worker 的 driver 环境变量。
​
    排除 worker 专属变量与用户配置的黑名单(RAY_NON_CARRY_OVER_ENV_VARS)。
    """
    exclude_vars = worker_specific_vars | RAY_NON_CARRY_OVER_ENV_VARS
    return {key: value for key, value in os.environ.items() if key not in exclude_vars}
​
​
def update_runtime_env_for_worker_import(runtime_env: dict) -> dict:
    """把 driver 中“导入期必需”的环境变量提前注入 Ray actor 的 runtime_env。
​
    VllmConfig 可能在 Ray 启动后才自动开启 breakable CUDA graph,而
    @eager_break_during_capture 装饰器在模型模块导入时就完成求值;
    Ray 默认在 initialize_worker 阶段才复制 driver 环境,届时为时已晚。
    因此必须在 actor 创建前把该变量放进 runtime_env["env_vars"]。
    """
    for env_var in _RAY_WORKER_PRE_IMPORT_ENV_VARS:
        if env_var in os.environ:
            env_vars = runtime_env.setdefault("env_vars", {})
            env_vars[env_var] = os.environ[env_var]
    return runtime_env
vllm/v1/executor/ray_executor_v2.py dependency-wiring

V2 Ray executor 接入点:在 _build_runtime_env 中调用注入函数,是修复生效的关键调用路径。

def _build_runtime_env(self) -> dict:
    """构建 RayWorkerProc actor 的 runtime_env。
​
    driver 环境变量在 initialize_worker 阶段以 setdefault 语义另行应用;
    但其中“导入期必需”的变量必须在此处提前注入,否则模型模块导入时
    看不到对应设置(见 update_runtime_env_for_worker_import 注释)。
    """
    base = self.parallel_config.ray_runtime_env
    runtime_env: dict = copy.deepcopy(dict(base)) if base else {}
​
    env_vars = runtime_env.setdefault("env_vars", {})
    # 设备可见性相关变量沿用既有逻辑,统一置为 "1"
    env_vars.update({v: "1" for v in current_platform.ray_noset_device_env_vars})
    # 提前注入 breakable CUDA graph 开关,保证模型导入时即可见
    update_runtime_env_for_worker_import(runtime_env)
    if self.parallel_config.ray_workers_use_nsight:
        runtime_env["nsight"] = {
            "t": "cuda,cudnn,cublas",
            "o": "'worker_process_%p'",
            "cuda-graph-trace": "node",
        }
    return runtime_env

评论区精华

将单变量直写泛化为可扩展的预导入环境变量白名单 设计

jeffreywang88 建议维护“应在 actor 创建前设置的 env var 列表”,将 VLLM_USE_BREAKABLE_CUDAGRAPH 加入其中;alexeldeib 回应最初已用 frozenset 遍历实现,但因只发现 1 个相关变量而未保留,担心过度设计。

结论:jeffreywang88 拍板采用 frozenset 方案以保持通用性和可扩展性,alexeldeib 修改提交并合入。 · 已解决

是否还有其它需要在 actor 创建前设置的 env var question

jeffreywang88 询问是否已识别其它同样受导入时序影响的变量;alexeldeib 表示仅确认 VLLM_USE_BREAKABLE_CUDAGRAPH 一项,未发现其它变量需要同样处理。

结论:未发现其它变量,白名单暂保持 1 项,但机制已具备扩展能力。 · 已解决

需要 Ray 侧专家复核 question

njhill 认为修复方向合理,但希望由更熟悉 Ray 代码的维护者确认 runtime_env 注入的时点是否正确,cc @jeffreywang88。

结论:jeffreywang88 review 后 APPROVED,njhill 随后 APPROVED 并合入。 · 已解决

风险与影响

  1. 覆盖语义:update_runtime_env_for_worker_import 使用 env_vars[env_var] = os.environ[env_var] 直接赋值而非 setdefault,若用户在 parallel_config.ray_runtime_env 中显式配置了同名变量,driver 值会覆盖用户配置,需要在使用自定义 runtime_env 时注意。
  2. 时序隐式契约:修复依赖 @eager_break_during_capture 在模型模块导入时读取环境变量的实现约定;若未来改为惰性求值或改读 VllmConfig 对象,此注入将变为冗余但无害。
  3. 测试覆盖局限:单元测试只覆盖注入函数本身;TP8 双节点 GB300 的端到端验证依赖人工执行(作者已在 PR body 中给出详细验证矩阵),CI 不会自动覆盖 Ray + breakable CUDA graph 组合场景。
  4. 回归面:改动触及 V1/V2 两个 Ray executor 的 runtime_env 构建路径,但注入逻辑幂等且仅白名单 1 项,未显式设置该变量时行为完全不变。

对用户:修复 DSv4 在 Ray 后端下自动启用 breakable CUDA graph 时的输出乱码与显存浪费(抓图 0.69 GiB 降至 0.37 GiB,颜色/数学问答从 0/10 恢复至 10/10),影响 TP 多机部署的正确性。对系统:为"导入期必需"环境变量建立了通用的前置传播机制,V1/V2 两个 Ray executor 行为对齐,且不会改变原有通用环境传播、worker 排除和节点本地变量。对团队:后续新增此类变量只需扩展 frozenset,降低同类问题再次发生的概率。

Ray 启动路径变更 依赖导入期求值时序契约 环境变量注入可能覆盖用户同名配置 端到端场景依赖人工验证

关联 Issue

#41834 [New Model][Nvidia] Add SM12x support for DeepSeek V4 Flash with essential fixes

完整报告

参与讨论