Ace EldeibandCodex
9662ab0835
Signed-off-by: Ace Eldeib <aeldeib@coreweave.com> Co-authored-by: Codex <codex@openai.com>
34 lines
1.2 KiB
Python
34 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import os
|
|
|
|
from vllm.ray.ray_env import RAY_NON_CARRY_OVER_ENV_VARS
|
|
|
|
_RAY_WORKER_PRE_IMPORT_ENV_VARS = frozenset({"VLLM_USE_BREAKABLE_CUDAGRAPH"})
|
|
|
|
|
|
def get_driver_env_vars(
|
|
worker_specific_vars: set[str],
|
|
) -> dict[str, str]:
|
|
"""Return driver env vars to propagate to Ray workers.
|
|
|
|
Returns everything from ``os.environ`` except ``worker_specific_vars``
|
|
and user-configured exclusions (``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:
|
|
"""Expose driver environment variables needed before Ray actor import.
|
|
|
|
VllmConfig can auto-enable breakable CUDA graphs after Ray starts, while
|
|
eager_break_during_capture resolves the setting during model import.
|
|
"""
|
|
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
|