| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | 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], |
| 12 | ) -> dict[str, str]: |
| 13 | """Return driver env vars to propagate to Ray workers. |
| 14 | |
| 15 | Returns everything from ``os.environ`` except ``worker_specific_vars`` |
| 16 | and user-configured exclusions (``RAY_NON_CARRY_OVER_ENV_VARS``). |
| 17 | """ |
| 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 |