[Bugfix] Set breakable graph env before Ray actor import (#53293)
Signed-off-by: Ace Eldeib <aeldeib@coreweave.com> Co-authored-by: Codex <codex@openai.com>
Ace Eldeib committed
Aug 28, 2026 at 19:11 UTC
9662ab0835e9eac28ac7d95d4b25ecb7140b7bf3
4 files changed
+51
-2
tests/test_ray_env_utils.py
+27
-1
@@ -5,7 +5,10 @@
5
import os
6
from unittest.mock import patch
7
8
-from vllm.v1.executor.ray_env_utils import get_driver_env_vars
8
+from vllm.v1.executor.ray_env_utils import (
9
+ get_driver_env_vars,
10
+ update_runtime_env_for_worker_import,
11
+)
12
13
WORKER_VARS: set[str] = {
14
"CUDA_VISIBLE_DEVICES",
@@ -49,3 +52,26 @@ class TestExclusion:
52
)
53
def test_non_carry_over_blacklist(self):
54
assert "LMCACHE_LOCAL_CPU" not in get_driver_env_vars(WORKER_VARS)
55
+
56
+
57
+class TestImportTimePropagation:
58
+ @patch.dict(os.environ, {"VLLM_USE_BREAKABLE_CUDAGRAPH": "1"}, clear=False)
59
+ def test_breakable_cudagraph_is_available_during_actor_import(self):
60
+ runtime_env = {"env_vars": {"EXISTING": "value"}}
61
+
62
+ result = update_runtime_env_for_worker_import(runtime_env)
63
+
64
+ assert result is runtime_env
65
+ assert result["env_vars"] == {
66
+ "EXISTING": "value",
67
+ "VLLM_USE_BREAKABLE_CUDAGRAPH": "1",
68
+ }
69
+
70
+ @patch.dict(os.environ, {"OTHER_ENV_VAR": "value"}, clear=True)
71
+ def test_other_env_vars_are_unchanged(self):
72
+ runtime_env = {"env_vars": {"EXISTING": "value"}}
73
+
74
+ result = update_runtime_env_for_worker_import(runtime_env)
75
+
76
+ assert result is runtime_env
77
+ assert result["env_vars"] == {"EXISTING": "value"}
vllm/v1/executor/ray_env_utils.py
+15
@@ -4,6 +4,8 @@ import os
4
5
from vllm.ray.ray_env import RAY_NON_CARRY_OVER_ENV_VARS
6
7
+_RAY_WORKER_PRE_IMPORT_ENV_VARS = frozenset({"VLLM_USE_BREAKABLE_CUDAGRAPH"})
8
+
9
10
def get_driver_env_vars(
11
worker_specific_vars: set[str],
@@ -16,3 +18,16 @@ def get_driver_env_vars(
18
exclude_vars = worker_specific_vars | RAY_NON_CARRY_OVER_ENV_VARS
19
20
return {key: value for key, value in os.environ.items() if key not in exclude_vars}
21
+
22
+
23
+def update_runtime_env_for_worker_import(runtime_env: dict) -> dict:
24
+ """Expose driver environment variables needed before Ray actor import.
25
+
26
+ VllmConfig can auto-enable breakable CUDA graphs after Ray starts, while
27
+ eager_break_during_capture resolves the setting during model import.
28
+ """
29
+ for env_var in _RAY_WORKER_PRE_IMPORT_ENV_VARS:
30
+ if env_var in os.environ:
31
+ env_vars = runtime_env.setdefault("env_vars", {})
32
+ env_vars[env_var] = os.environ[env_var]
33
+ return runtime_env
vllm/v1/executor/ray_executor.py
+4
@@ -22,6 +22,9 @@ from vllm.utils.network_utils import (
22
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
23
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
24
from vllm.v1.executor.abstract import Executor
25
+from vllm.v1.executor.ray_env_utils import (
26
+ update_runtime_env_for_worker_import,
27
+)
28
from vllm.v1.executor.ray_utils import (
29
WORKER_SPECIFIC_ENV_VARS,
30
FutureWrapper,
@@ -147,6 +150,7 @@ class RayDistributedExecutor(Executor):
150
env_vars.update(
151
{env_var: "1" for env_var in current_platform.ray_noset_device_env_vars}
152
)
153
+ update_runtime_env_for_worker_import(runtime_env)
154
return ray_remote_kwargs
155
156
# child class could overwrite this to return actual env vars.
vllm/v1/executor/ray_executor_v2.py
+5
-1
@@ -26,7 +26,10 @@ from vllm.v1.executor.multiproc_executor import (
26
MultiprocExecutor,
27
WorkerProc,
28
)
29
-from vllm.v1.executor.ray_env_utils import get_driver_env_vars
29
+from vllm.v1.executor.ray_env_utils import (
30
+ get_driver_env_vars,
31
+ update_runtime_env_for_worker_import,
32
+)
33
from vllm.v1.executor.ray_utils import (
34
WORKER_SPECIFIC_ENV_VARS,
35
build_actor_name,
@@ -258,6 +261,7 @@ class RayExecutorV2(MultiprocExecutor):
261
262
env_vars = runtime_env.setdefault("env_vars", {})
263
env_vars.update({v: "1" for v in current_platform.ray_noset_device_env_vars})
264
+ update_runtime_env_for_worker_import(runtime_env)
265
if self.parallel_config.ray_workers_use_nsight:
266
runtime_env["nsight"] = {
267
"t": "cuda,cudnn,cublas",