vllm/vllm/ray/ray_env.py
Harry Mellor d6953beb91
Convert formatting to use ruff instead of yapf + isort (#26247)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2025-10-05 07:06:22 -07:00

81 lines
2.5 KiB
Python

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import os
from typing import Optional
import vllm.envs as envs
from vllm.logger import init_logger
logger = init_logger(__name__)
CONFIG_HOME = envs.VLLM_CONFIG_ROOT
# This file contains a list of env vars that should not be copied
# from the driver to the Ray workers.
RAY_NON_CARRY_OVER_ENV_VARS_FILE = os.path.join(
CONFIG_HOME, "ray_non_carry_over_env_vars.json"
)
try:
if os.path.exists(RAY_NON_CARRY_OVER_ENV_VARS_FILE):
with open(RAY_NON_CARRY_OVER_ENV_VARS_FILE) as f:
RAY_NON_CARRY_OVER_ENV_VARS = set(json.load(f))
else:
RAY_NON_CARRY_OVER_ENV_VARS = set()
except json.JSONDecodeError:
logger.warning(
"Failed to parse %s. Using an empty set for non-carry-over env vars.",
RAY_NON_CARRY_OVER_ENV_VARS_FILE,
)
RAY_NON_CARRY_OVER_ENV_VARS = set()
def get_env_vars_to_copy(
exclude_vars: Optional[set[str]] = None,
additional_vars: Optional[set[str]] = None,
destination: Optional[str] = None,
) -> set[str]:
"""
Get the environment variables to copy to downstream Ray actors.
Example use cases:
- Copy environment variables from RayDistributedExecutor to Ray workers.
- Copy environment variables from RayDPClient to Ray DPEngineCoreActor.
Args:
exclude_vars: A set of vllm defined environment variables to exclude
from copying.
additional_vars: A set of additional environment variables to copy.
If a variable is in both exclude_vars and additional_vars, it will
be excluded.
destination: The destination of the environment variables.
Returns:
A set of environment variables to copy.
"""
exclude_vars = exclude_vars or set()
additional_vars = additional_vars or set()
env_vars_to_copy = {
v
for v in set(envs.environment_variables).union(additional_vars)
if v not in exclude_vars and v not in RAY_NON_CARRY_OVER_ENV_VARS
}
to_destination = " to " + destination if destination is not None else ""
logger.info(
"RAY_NON_CARRY_OVER_ENV_VARS from config: %s", RAY_NON_CARRY_OVER_ENV_VARS
)
logger.info(
"Copying the following environment variables%s: %s",
to_destination,
[v for v in env_vars_to_copy if v in os.environ],
)
logger.info(
"If certain env vars should NOT be copied, add them to %s file",
RAY_NON_CARRY_OVER_ENV_VARS_FILE,
)
return env_vars_to_copy