Convert formatting to use ruff instead of yapf + isort (#26247)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
Harry Mellor
2025-10-05 15:06:22 +01:00
committed by GitHub
parent 17edd8a807
commit d6953beb91
1508 changed files with 115244 additions and 94146 deletions

View File

@@ -4,8 +4,9 @@
import asyncio
import time
from abc import ABC, abstractmethod
from collections.abc import Awaitable
from functools import cached_property
from typing import Any, Awaitable, Callable, List, Optional, Set, Union
from typing import Any, Callable, List, Optional, Set, Union
from typing_extensions import TypeVar
@@ -60,11 +61,13 @@ class ExecutorBase(ABC):
raise NotImplementedError
@abstractmethod
def collective_rpc(self,
method: Union[str, Callable[[WorkerBase], _R]],
timeout: Optional[float] = None,
args: tuple = (),
kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
def collective_rpc(
self,
method: Union[str, Callable[[WorkerBase], _R]],
timeout: Optional[float] = None,
args: tuple = (),
kwargs: Optional[dict[str, Any]] = None,
) -> list[_R]:
"""
Execute an RPC call on all workers.
@@ -109,22 +112,29 @@ class ExecutorBase(ABC):
return a, b
def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks) -> None:
"""Initialize the KV cache by invoking the underlying worker.
"""
"""Initialize the KV cache by invoking the underlying worker."""
# NOTE: This is logged in the executor because there can be >1 workers.
logger.info("# %s blocks: %d, # CPU blocks: %d",
vllm.platforms.current_platform.device_name,
num_gpu_blocks, num_cpu_blocks)
max_concurrency = (num_gpu_blocks * self.cache_config.block_size /
self.model_config.max_model_len)
logger.info("Maximum concurrency for %s tokens per request: %.2fx",
self.model_config.max_model_len, max_concurrency)
logger.info(
"# %s blocks: %d, # CPU blocks: %d",
vllm.platforms.current_platform.device_name,
num_gpu_blocks,
num_cpu_blocks,
)
max_concurrency = (
num_gpu_blocks
* self.cache_config.block_size
/ self.model_config.max_model_len
)
logger.info(
"Maximum concurrency for %s tokens per request: %.2fx",
self.model_config.max_model_len,
max_concurrency,
)
self.cache_config.num_gpu_blocks = num_gpu_blocks
self.cache_config.num_cpu_blocks = num_cpu_blocks
self.collective_rpc("initialize_cache",
args=(num_gpu_blocks, num_cpu_blocks))
self.collective_rpc("initialize_cache", args=(num_gpu_blocks, num_cpu_blocks))
@cached_property # Avoid unnecessary RPC calls
def supported_tasks(self) -> tuple[SupportedTask, ...]:
@@ -134,8 +144,7 @@ class ExecutorBase(ABC):
def execute_model(
self, execute_model_req: ExecuteModelRequest
) -> Optional[List[Union[SamplerOutput, PoolerOutput]]]:
output = self.collective_rpc("execute_model",
args=(execute_model_req, ))
output = self.collective_rpc("execute_model", args=(execute_model_req,))
return output[0]
def stop_remote_worker_execution_loop(self) -> None:
@@ -144,15 +153,15 @@ class ExecutorBase(ABC):
def add_lora(self, lora_request: LoRARequest) -> bool:
assert lora_request.lora_int_id > 0, "lora_id must be greater than 0."
return all(self.collective_rpc("add_lora", args=(lora_request, )))
return all(self.collective_rpc("add_lora", args=(lora_request,)))
def remove_lora(self, lora_id: int) -> bool:
assert lora_id > 0, "lora_id must be greater than 0."
return all(self.collective_rpc("remove_lora", args=(lora_id, )))
return all(self.collective_rpc("remove_lora", args=(lora_id,)))
def pin_lora(self, lora_id: int) -> bool:
assert lora_id > 0, "lora_id must be greater than 0."
return all(self.collective_rpc("pin_lora", args=(lora_id, )))
return all(self.collective_rpc("pin_lora", args=(lora_id,)))
def list_loras(self) -> Set[int]:
sets = self.collective_rpc("list_loras")
@@ -175,8 +184,9 @@ class ExecutorBase(ABC):
time_after_sleep = time.perf_counter()
self.sleeping_tags = {"weights", "kv_cache"}
self.is_sleeping = True
logger.info("It took %.6f seconds to fall asleep.",
time_after_sleep - time_before_sleep)
logger.info(
"It took %.6f seconds to fall asleep.", time_after_sleep - time_before_sleep
)
def wake_up(self, tags: Optional[list[str]] = None):
if not self.is_sleeping:
@@ -185,15 +195,18 @@ class ExecutorBase(ABC):
if tags:
for tag in tags:
if tag not in self.sleeping_tags:
logger.warning("Tag %s is not in sleeping tags %s", tag,
self.sleeping_tags)
logger.warning(
"Tag %s is not in sleeping tags %s", tag, self.sleeping_tags
)
return
time_before_wakeup = time.perf_counter()
self.collective_rpc("wake_up", kwargs=dict(tags=tags))
time_after_wakeup = time.perf_counter()
logger.info("It took %.6f seconds to wake up tags %s.",
time_after_wakeup - time_before_wakeup,
tags if tags is not None else self.sleeping_tags)
logger.info(
"It took %.6f seconds to wake up tags %s.",
time_after_wakeup - time_before_wakeup,
tags if tags is not None else self.sleeping_tags,
)
if tags:
for tag in tags:
self.sleeping_tags.remove(tag)
@@ -208,10 +221,10 @@ class ExecutorBase(ABC):
pattern: Optional[str] = None,
max_size: Optional[int] = None,
) -> None:
self.collective_rpc("save_sharded_state",
kwargs=dict(path=path,
pattern=pattern,
max_size=max_size))
self.collective_rpc(
"save_sharded_state",
kwargs=dict(path=path, pattern=pattern, max_size=max_size),
)
@abstractmethod
def check_health(self) -> None:
@@ -224,8 +237,8 @@ class ExecutorBase(ABC):
self.collective_rpc("shutdown")
async def execute_model_async(
self,
execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]:
self, execute_model_req: ExecuteModelRequest
) -> List[SamplerOutput]:
"""Executes one model step on the given sequences."""
output = await make_async(self.execute_model)(execute_model_req)
return output
@@ -242,7 +255,8 @@ class ExecutorBase(ABC):
def init_kv_output_aggregator(self, finished_count: Optional[int]) -> None:
"""Init KVOutputAggregator"""
self.kv_output_aggregator = KVOutputAggregator(
finished_count or self.parallel_config.world_size)
finished_count or self.parallel_config.world_size
)
class DistributedExecutorBase(ExecutorBase):
@@ -263,7 +277,8 @@ class DistributedExecutorBase(ExecutorBase):
if self.parallel_worker_tasks is None:
self.parallel_worker_tasks = self._run_workers(
"start_worker_execution_loop",
async_run_tensor_parallel_workers_only=True)
async_run_tensor_parallel_workers_only=True,
)
# Only the driver worker returns the sampling results.
driver_outputs = self._driver_execute_model(execute_model_req)
@@ -293,11 +308,13 @@ class DistributedExecutorBase(ExecutorBase):
"""
raise NotImplementedError
def collective_rpc(self,
method: Union[str, Callable],
timeout: Optional[float] = None,
args: tuple = (),
kwargs: Optional[dict[str, Any]] = None) -> list[Any]:
def collective_rpc(
self,
method: Union[str, Callable],
timeout: Optional[float] = None,
args: tuple = (),
kwargs: Optional[dict[str, Any]] = None,
) -> list[Any]:
return self._run_workers(method, *args, **(kwargs or {}))
@abstractmethod
@@ -328,12 +345,13 @@ class DistributedExecutorBase(ExecutorBase):
raise NotImplementedError
async def execute_model_async(
self,
execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]:
self, execute_model_req: ExecuteModelRequest
) -> List[SamplerOutput]:
if self.parallel_worker_tasks is None:
# Start model execution loop running in the parallel workers
self.parallel_worker_tasks = asyncio.create_task(
self._start_worker_execution_loop())
self._start_worker_execution_loop()
)
# Only the driver worker returns the sampling results.
return await self._driver_execute_model_async(execute_model_req)

View File

@@ -16,7 +16,8 @@ def encode_hook(obj: Any) -> Any:
if isinstance(obj, array):
assert obj.typecode == VLLM_TOKEN_ID_ARRAY_TYPE, (
f"vLLM array type should use '{VLLM_TOKEN_ID_ARRAY_TYPE}' type. "
f"Given array has a type code of {obj.typecode}.")
f"Given array has a type code of {obj.typecode}."
)
return obj.tobytes()
if isinstance(obj, MultiModalKwargs):
return dict(obj)

View File

@@ -11,17 +11,20 @@ import cloudpickle
import msgspec
import vllm.envs as envs
from vllm.executor.executor_base import (
DistributedExecutorBase) # yapf: disable
from vllm.executor.executor_base import DistributedExecutorBase # yapf: disable
from vllm.executor.msgspec_utils import encode_hook
from vllm.executor.ray_utils import (RayWorkerWrapper, initialize_ray_cluster,
ray)
from vllm.executor.ray_utils import RayWorkerWrapper, initialize_ray_cluster, ray
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.ray.ray_env import get_env_vars_to_copy
from vllm.sequence import ExecuteModelRequest
from vllm.utils import (_run_task_with_lock, get_distributed_init_method,
get_ip, get_open_port, make_async)
from vllm.utils import (
_run_task_with_lock,
get_distributed_init_method,
get_ip,
get_open_port,
make_async,
)
from vllm.v1.outputs import SamplerOutput
if ray is not None:
@@ -43,6 +46,7 @@ class RayWorkerMetaData:
The order of ray worker creation can be random,
and we need to reset the rank after creating all workers.
"""
worker: ActorHandle
created_rank: int
adjusted_rank: int = -1
@@ -55,7 +59,10 @@ class RayDistributedExecutor(DistributedExecutorBase):
# These env vars are worker-specific, therefore are NOT copied
# from the driver to the workers
WORKER_SPECIFIC_ENV_VARS = {
"VLLM_HOST_IP", "VLLM_HOST_PORT", "LOCAL_RANK", "CUDA_VISIBLE_DEVICES"
"VLLM_HOST_IP",
"VLLM_HOST_PORT",
"LOCAL_RANK",
"CUDA_VISIBLE_DEVICES",
}
# These non-vLLM env vars are copied from the driver to workers
@@ -86,13 +93,13 @@ class RayDistributedExecutor(DistributedExecutorBase):
self.use_ray_spmd_worker = envs.VLLM_USE_RAY_SPMD_WORKER
if self.use_ray_compiled_dag:
assert self.use_ray_spmd_worker, (
"VLLM_USE_RAY_COMPILED_DAG=1 requires "
"VLLM_USE_RAY_SPMD_WORKER=1")
"VLLM_USE_RAY_COMPILED_DAG=1 requires VLLM_USE_RAY_SPMD_WORKER=1"
)
if self.use_ray_spmd_worker:
# TODO: Support SPMD worker for non-DAG Ray executor.
assert self.use_ray_compiled_dag, (
"VLLM_USE_RAY_SPMD_WORKER=1 requires "
"VLLM_USE_RAY_COMPILED_DAG=1")
"VLLM_USE_RAY_SPMD_WORKER=1 requires VLLM_USE_RAY_COMPILED_DAG=1"
)
assert self.uses_ray
initialize_ray_cluster(self.parallel_config)
@@ -107,14 +114,12 @@ class RayDistributedExecutor(DistributedExecutorBase):
self._init_workers_ray(placement_group)
self.input_encoder = msgspec.msgpack.Encoder(enc_hook=encode_hook)
self.output_decoder = msgspec.msgpack.Decoder(
Optional[List[SamplerOutput]])
self.output_decoder = msgspec.msgpack.Decoder(Optional[List[SamplerOutput]])
self.use_v1 = envs.VLLM_USE_V1
self.pp_locks: Optional[List[asyncio.Lock]] = None
if not self.use_ray_compiled_dag:
self.driver_exec_method = make_async(
self.driver_worker.execute_method)
self.driver_exec_method = make_async(self.driver_worker.execute_method)
def shutdown(self) -> None:
if logger:
@@ -122,26 +127,29 @@ class RayDistributedExecutor(DistributedExecutorBase):
logger.info(
"Shutting down Ray distributed executor. If you see error log "
"from logging.cc regarding SIGTERM received, please ignore "
"because this is the expected termination process in Ray.")
"because this is the expected termination process in Ray."
)
if hasattr(self, "forward_dag") and self.forward_dag is not None:
self.forward_dag.teardown()
import ray
for worker in self.workers:
ray.kill(worker)
self.forward_dag = None
def _configure_ray_workers_use_nsight(self,
ray_remote_kwargs) -> Dict[str, Any]:
def _configure_ray_workers_use_nsight(self, ray_remote_kwargs) -> Dict[str, Any]:
# If nsight profiling is enabled, we need to set the profiling
# configuration for the ray workers as runtime env.
runtime_env = ray_remote_kwargs.setdefault("runtime_env", {})
runtime_env.update({
"nsight": {
"t": "cuda,cudnn,cublas",
"o": "'worker_process_%p'",
"cuda-graph-trace": "node",
runtime_env.update(
{
"nsight": {
"t": "cuda,cudnn,cublas",
"o": "'worker_process_%p'",
"cuda-graph-trace": "node",
}
}
})
)
return ray_remote_kwargs
@@ -149,8 +157,7 @@ class RayDistributedExecutor(DistributedExecutorBase):
def _get_env_vars_to_be_updated(self):
return self._env_vars_for_all_workers
def _init_workers_ray(self, placement_group: "PlacementGroup",
**ray_remote_kwargs):
def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwargs):
num_gpus = envs.VLLM_RAY_PER_WORKER_GPUS
# The driver dummy worker does not actually use any resources.
@@ -166,7 +173,8 @@ class RayDistributedExecutor(DistributedExecutorBase):
if self.parallel_config.ray_workers_use_nsight:
ray_remote_kwargs = self._configure_ray_workers_use_nsight(
ray_remote_kwargs)
ray_remote_kwargs
)
logger.info("use_ray_spmd_worker: %s", self.use_ray_spmd_worker)
@@ -174,22 +182,23 @@ class RayDistributedExecutor(DistributedExecutorBase):
bundle_indices: List[int]
if envs.VLLM_RAY_BUNDLE_INDICES:
# Use the bundle indices specified by the user.
bundle_indices = list(
map(int, envs.VLLM_RAY_BUNDLE_INDICES.split(",")))
assert len(bundle_indices) == self.parallel_config.world_size, \
("VLLM_RAY_BUNDLE_INDICES must have the same size"
f" as the world size, but got {bundle_indices=} "
f"and {self.parallel_config.world_size=}")
assert len(set(bundle_indices)) == len(bundle_indices), \
("VLLM_RAY_BUNDLE_INDICES cannot have duplicate values,"
f" but got {bundle_indices=}")
bundle_indices = list(map(int, envs.VLLM_RAY_BUNDLE_INDICES.split(",")))
assert len(bundle_indices) == self.parallel_config.world_size, (
"VLLM_RAY_BUNDLE_INDICES must have the same size"
f" as the world size, but got {bundle_indices=} "
f"and {self.parallel_config.world_size=}"
)
assert len(set(bundle_indices)) == len(bundle_indices), (
"VLLM_RAY_BUNDLE_INDICES cannot have duplicate values,"
f" but got {bundle_indices=}"
)
else:
# use the first N bundles that have GPU resources.
bundle_indices = []
for bundle_id, bundle in enumerate(placement_group.bundle_specs):
if bundle.get(current_platform.ray_device_key, 0):
bundle_indices.append(bundle_id)
bundle_indices = bundle_indices[:self.parallel_config.world_size]
bundle_indices = bundle_indices[: self.parallel_config.world_size]
worker_metadata: List[RayWorkerMetaData] = []
driver_ip = get_ip()
@@ -207,8 +216,7 @@ class RayDistributedExecutor(DistributedExecutorBase):
num_gpus=num_gpus,
scheduling_strategy=scheduling_strategy,
**ray_remote_kwargs,
)(RayWorkerWrapper).remote(vllm_config=self.vllm_config,
rpc_rank=rank)
)(RayWorkerWrapper).remote(vllm_config=self.vllm_config, rpc_rank=rank)
else:
worker = ray.remote(
num_cpus=0,
@@ -216,15 +224,15 @@ class RayDistributedExecutor(DistributedExecutorBase):
resources={current_platform.ray_device_key: num_gpus},
scheduling_strategy=scheduling_strategy,
**ray_remote_kwargs,
)(RayWorkerWrapper).remote(vllm_config=self.vllm_config,
rpc_rank=rank)
worker_metadata.append(
RayWorkerMetaData(worker=worker, created_rank=rank))
)(RayWorkerWrapper).remote(vllm_config=self.vllm_config, rpc_rank=rank)
worker_metadata.append(RayWorkerMetaData(worker=worker, created_rank=rank))
worker_ips = ray.get([
each.worker.get_node_ip.remote() # type: ignore[attr-defined]
for each in worker_metadata
])
worker_ips = ray.get(
[
each.worker.get_node_ip.remote() # type: ignore[attr-defined]
for each in worker_metadata
]
)
for each, ip in zip(worker_metadata, worker_ips):
each.ip = ip
@@ -239,7 +247,8 @@ class RayDistributedExecutor(DistributedExecutorBase):
# as the resource holder for the driver process.
self.driver_dummy_worker = worker
self.driver_worker = RayWorkerWrapper(
vllm_config=self.vllm_config, rpc_rank=0)
vllm_config=self.vllm_config, rpc_rank=0
)
worker_metadata.pop(i)
break
@@ -250,7 +259,8 @@ class RayDistributedExecutor(DistributedExecutorBase):
"Ray does not allocate any GPUs on the driver node."
f"Driver IP: {driver_ip}, worker IPs: {worker_ips}."
"Consider adjusting the Ray placement group or running "
"the driver on a GPU node.")
"the driver on a GPU node."
)
ip_counts: Dict[str, int] = {}
for ip in worker_ips:
@@ -272,15 +282,15 @@ class RayDistributedExecutor(DistributedExecutorBase):
# After sorting, the workers on the same node will be
# close to each other, and the workers on the driver
# node will be placed first.
sorted_worker_metadata = sorted(worker_metadata,
key=sort_by_driver_then_worker_ip)
sorted_worker_metadata = sorted(
worker_metadata, key=sort_by_driver_then_worker_ip
)
start_rank = 0 if self.use_ray_spmd_worker else 1
for i, item in enumerate(sorted_worker_metadata):
item.adjusted_rank = i + start_rank
self.workers = [item.worker for item in sorted_worker_metadata]
rerank_mapping = {
item.created_rank: item.adjusted_rank
for item in sorted_worker_metadata
item.created_rank: item.adjusted_rank for item in sorted_worker_metadata
}
self._run_workers("adjust_rank", rerank_mapping)
@@ -291,8 +301,8 @@ class RayDistributedExecutor(DistributedExecutorBase):
# driver_dummy_worker can be None when using ray spmd worker.
continue
worker_node_and_gpu_ids.append(
ray.get(worker.get_node_and_gpu_ids.remote()) \
) # type: ignore
ray.get(worker.get_node_and_gpu_ids.remote())
) # type: ignore
node_workers = defaultdict(list) # node id -> list of worker ranks
node_gpus = defaultdict(list) # node id -> list of gpu ids
@@ -320,20 +330,27 @@ class RayDistributedExecutor(DistributedExecutorBase):
f"{n_ips} unique IP addresses {all_ips}. Please check your"
" network configuration. If you set `VLLM_HOST_IP`"
" environment variable, make sure it is unique for"
" each node.")
" each node."
)
# Set environment variables for the driver and workers.
all_args_to_update_environment_variables = [{
current_platform.device_control_env_var:
",".join(map(str, node_gpus[node_id])),
} for (node_id, _) in worker_node_and_gpu_ids]
all_args_to_update_environment_variables = [
{
current_platform.device_control_env_var: ",".join(
map(str, node_gpus[node_id])
),
}
for (node_id, _) in worker_node_and_gpu_ids
]
# Environment variables to copy from driver to workers
env_vars_to_copy = get_env_vars_to_copy(
exclude_vars=self.WORKER_SPECIFIC_ENV_VARS,
additional_vars=set(current_platform.additional_env_vars).union(
self.ADDITIONAL_ENV_VARS),
destination="workers")
self.ADDITIONAL_ENV_VARS
),
destination="workers",
)
# Copy existing env vars to each worker's args
for args in all_args_to_update_environment_variables:
@@ -342,11 +359,11 @@ class RayDistributedExecutor(DistributedExecutorBase):
if name in os.environ:
args[name] = os.environ[name]
self._env_vars_for_all_workers = (
all_args_to_update_environment_variables)
self._env_vars_for_all_workers = all_args_to_update_environment_variables
self._run_workers("update_environment_variables",
self._get_env_vars_to_be_updated())
self._run_workers(
"update_environment_variables", self._get_env_vars_to_be_updated()
)
if len(node_gpus) == 1:
# in single node case, we don't need to get the IP address.
@@ -359,7 +376,8 @@ class RayDistributedExecutor(DistributedExecutorBase):
# the node.
driver_ip = "127.0.0.1"
distributed_init_method = get_distributed_init_method(
driver_ip, get_open_port())
driver_ip, get_open_port()
)
# Initialize the actual workers inside worker wrapper.
all_kwargs = []
@@ -377,19 +395,20 @@ class RayDistributedExecutor(DistributedExecutorBase):
self._run_workers("init_worker", all_kwargs)
self._run_workers("init_device")
self._run_workers("load_model",
max_concurrent_workers=self.parallel_config.
max_parallel_loading_workers)
self._run_workers(
"load_model",
max_concurrent_workers=self.parallel_config.max_parallel_loading_workers,
)
if self.use_ray_spmd_worker:
for pp_rank in range(self.parallel_config.pipeline_parallel_size):
self.pp_tp_workers.append([])
for tp_rank in range(
self.parallel_config.tensor_parallel_size):
for tp_rank in range(self.parallel_config.tensor_parallel_size):
# PP=2, TP=4
# pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]]
rank = (pp_rank * self.parallel_config.tensor_parallel_size
) + tp_rank
rank = (
pp_rank * self.parallel_config.tensor_parallel_size
) + tp_rank
assert len(self.pp_tp_workers[pp_rank]) == tp_rank
assert pp_rank < len(self.pp_tp_workers)
self.pp_tp_workers[pp_rank].append(self.workers[rank])
@@ -421,13 +440,13 @@ class RayDistributedExecutor(DistributedExecutorBase):
loop running in each of the remote workers.
"""
assert not self.use_ray_spmd_worker, (
"driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1")
return self.driver_worker.execute_method("execute_model",
execute_model_req)
"driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1"
)
return self.driver_worker.execute_method("execute_model", execute_model_req)
def execute_model(
self,
execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]:
self, execute_model_req: ExecuteModelRequest
) -> List[SamplerOutput]:
if not self.use_ray_spmd_worker:
return super().execute_model(execute_model_req)
@@ -470,12 +489,11 @@ class RayDistributedExecutor(DistributedExecutorBase):
del method
if self.use_ray_spmd_worker:
assert not async_run_tensor_parallel_workers_only, (
"async_run_tensor_parallel_workers_only is not supported for "
"spmd mode.")
"async_run_tensor_parallel_workers_only is not supported for spmd mode."
)
if max_concurrent_workers:
raise NotImplementedError(
"max_concurrent_workers is not supported yet.")
raise NotImplementedError("max_concurrent_workers is not supported yet.")
# Start the ray workers first.
ray_workers = self.workers
@@ -519,23 +537,27 @@ class RayDistributedExecutor(DistributedExecutorBase):
required_version = version.parse("2.43.0")
current_version = version.parse(importlib.metadata.version("ray"))
if current_version < required_version:
raise ValueError(f"Ray version {required_version} is "
f"required, but found {current_version}")
raise ValueError(
f"Ray version {required_version} is "
f"required, but found {current_version}"
)
import importlib.util
cgraph_spec = importlib.util.find_spec(
"ray.experimental.compiled_dag_ref")
cgraph_spec = importlib.util.find_spec("ray.experimental.compiled_dag_ref")
if cgraph_spec is None:
raise ValueError("Ray Compiled Graph is not installed. "
"Run `pip install ray[cgraph]` to install it.")
raise ValueError(
"Ray Compiled Graph is not installed. "
"Run `pip install ray[cgraph]` to install it."
)
cupy_spec = importlib.util.find_spec("cupy")
if (cupy_spec is None
and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "nccl"):
if cupy_spec is None and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "nccl":
raise ValueError(
"cupy is not installed but required since "
"VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE is set to 'nccl'. "
"Run `pip install ray[cgraph]` and check cupy installation.")
"Run `pip install ray[cgraph]` and check cupy installation."
)
def _compiled_ray_dag(self, enable_asyncio: bool):
assert self.parallel_config.use_ray
@@ -549,18 +571,25 @@ class RayDistributedExecutor(DistributedExecutorBase):
# ray.dag, otherwise it will not take effect.
os.environ.setdefault("RAY_CGRAPH_get_timeout", "300") # noqa: SIM112
from ray.dag import InputNode, MultiOutputNode
logger.info("RAY_CGRAPH_get_timeout is set to %s",
os.environ["RAY_CGRAPH_get_timeout"]) # noqa: SIM112
logger.info("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE = %s",
envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE)
logger.info("VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM = %s",
envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM)
logger.info(
"RAY_CGRAPH_get_timeout is set to %s", os.environ["RAY_CGRAPH_get_timeout"]
) # noqa: SIM112
logger.info(
"VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE = %s",
envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE,
)
logger.info(
"VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM = %s",
envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM,
)
channel_type = envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE
if channel_type not in ("auto", "nccl", "shm"):
raise ValueError(
"Invalid value for VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: "
f"{channel_type}. Valid values are: 'auto', 'nccl', or 'shm'.")
f"{channel_type}. Valid values are: 'auto', 'nccl', or 'shm'."
)
with InputNode() as input_data:
# Example DAG: PP=2, TP=4
@@ -585,20 +614,24 @@ class RayDistributedExecutor(DistributedExecutorBase):
# and the TP group executes in SPMD fashion.
if self.use_v1:
outputs = [
worker.execute_model_ray.
bind( # type: ignore[attr-defined]
outputs[i]) for i, worker in enumerate(tp_group)
worker.execute_model_ray.bind( # type: ignore[attr-defined]
outputs[i]
)
for i, worker in enumerate(tp_group)
]
else:
outputs = [
worker.execute_model_spmd.
bind( # type: ignore[attr-defined]
outputs[i]) for i, worker in enumerate(tp_group)
worker.execute_model_spmd.bind( # type: ignore[attr-defined]
outputs[i]
)
for i, worker in enumerate(tp_group)
]
last_pp_rank = len(self.pp_tp_workers) - 1
if (pp_rank < last_pp_rank and
envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE != "shm"):
if (
pp_rank < last_pp_rank
and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE != "shm"
):
# Specify how intermediate tensors should be passed
# between pp stages, no need to specify for the last
# pp stage or when using shared memory (the default).
@@ -612,30 +645,37 @@ class RayDistributedExecutor(DistributedExecutorBase):
if envs.VLLM_USE_RAY_WRAPPED_PP_COMM:
from ray.experimental.channel.accelerator_context import (
register_accelerator_context)
register_accelerator_context,
)
from vllm.distributed.device_communicators.ray_communicator import (
RayPPCommunicator)
register_accelerator_context(torch_module_name="cuda",
communicator_cls=RayPPCommunicator)
logger.info("Using RayPPCommunicator "
"(which wraps vLLM _PP GroupCoordinator) "
"for Ray Compiled Graph communication.")
RayPPCommunicator,
)
register_accelerator_context(
torch_module_name="cuda", communicator_cls=RayPPCommunicator
)
logger.info(
"Using RayPPCommunicator "
"(which wraps vLLM _PP GroupCoordinator) "
"for Ray Compiled Graph communication."
)
else:
logger.info("Using Ray's NCCL communicator for "
"Ray Compiled Graph communication.")
logger.info(
"Using Ray's NCCL communicator for Ray Compiled Graph communication."
)
return forward_dag.experimental_compile(
enable_asyncio=enable_asyncio,
_overlap_gpu_communication=envs.
VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM)
_overlap_gpu_communication=envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM,
)
def __del__(self):
self.shutdown()
async def execute_model_async(
self,
execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]:
self, execute_model_req: ExecuteModelRequest
) -> List[SamplerOutput]:
if not self.use_ray_spmd_worker:
return await super().execute_model_async(execute_model_req)
@@ -648,14 +688,13 @@ class RayDistributedExecutor(DistributedExecutorBase):
return self.output_decoder.decode(output)
async def _driver_execute_model_async(
self,
execute_model_req: Optional[ExecuteModelRequest] = None
self, execute_model_req: Optional[ExecuteModelRequest] = None
) -> List[SamplerOutput]:
assert not self.use_ray_spmd_worker, (
"driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1")
"driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1"
)
if not self.tp_driver_workers:
return await self.driver_exec_method("execute_model",
execute_model_req)
return await self.driver_exec_method("execute_model", execute_model_req)
if self.pp_locks is None:
# This locks each pipeline parallel stage so multiple virtual
# engines can't execute on the same stage at the same time
@@ -668,16 +707,25 @@ class RayDistributedExecutor(DistributedExecutorBase):
tasks = [
asyncio.create_task(
_run_task_with_lock(self.driver_exec_method, self.pp_locks[0],
"execute_model", execute_model_req))
_run_task_with_lock(
self.driver_exec_method,
self.pp_locks[0],
"execute_model",
execute_model_req,
)
)
]
for pp_rank, driver_worker in enumerate(self.tp_driver_workers,
start=1):
for pp_rank, driver_worker in enumerate(self.tp_driver_workers, start=1):
tasks.append(
asyncio.create_task(
_run_task_with_lock(driver_worker.execute_method.remote,
self.pp_locks[pp_rank],
"execute_model", execute_model_req)))
_run_task_with_lock(
driver_worker.execute_method.remote,
self.pp_locks[pp_rank],
"execute_model",
execute_model_req,
)
)
)
results = await asyncio.gather(*tasks)
@@ -686,7 +734,8 @@ class RayDistributedExecutor(DistributedExecutorBase):
async def _start_worker_execution_loop(self):
assert not self.use_ray_spmd_worker, (
"worker loop is disabled for VLLM_USE_RAY_SPMD_WORKER=1")
"worker loop is disabled for VLLM_USE_RAY_SPMD_WORKER=1"
)
coros = [
worker.execute_method.remote("start_worker_execution_loop")
for worker in self.non_driver_workers

View File

@@ -30,11 +30,13 @@ try:
import ray
from ray.util import placement_group_table
from ray.util.placement_group import PlacementGroup
try:
from ray._private.state import available_resources_per_node
except ImportError:
# Ray 2.9.x doesn't expose `available_resources_per_node`
from ray._private.state import state as _state
available_resources_per_node = _state._available_resources_per_node
class RayWorkerWrapper(WorkerWrapperBase):
@@ -49,8 +51,9 @@ try:
# that thread.
self.compiled_dag_cuda_device_set = False
self.input_decoder = msgspec.msgpack.Decoder(ExecuteModelRequest,
dec_hook=decode_hook)
self.input_decoder = msgspec.msgpack.Decoder(
ExecuteModelRequest, dec_hook=decode_hook
)
self.output_encoder = msgspec.msgpack.Encoder(enc_hook=encode_hook)
def get_node_ip(self) -> str:
@@ -60,16 +63,16 @@ try:
node_id = ray.get_runtime_context().get_node_id()
device_key = vllm.platforms.current_platform.ray_device_key
if not device_key:
raise RuntimeError("current platform %s does not support ray.",
vllm.platforms.current_platform.device_name)
gpu_ids = ray.get_runtime_context().get_accelerator_ids(
)[device_key]
raise RuntimeError(
"current platform %s does not support ray.",
vllm.platforms.current_platform.device_name,
)
gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]
return node_id, gpu_ids
def execute_model_spmd(
self, req_or_tuple: Union[bytes,
Tuple[bytes,
Optional[IntermediateTensors]]]
self,
req_or_tuple: Union[bytes, Tuple[bytes, Optional[IntermediateTensors]]],
) -> bytes:
"""Execute model in SPMD fashion: used only when SPMD worker and
compiled DAG are both enabled.
@@ -94,8 +97,9 @@ try:
current_platform.set_device(self.worker.device)
self.compiled_dag_cuda_device_set = True
output = self.worker._execute_model_spmd(execute_model_req,
intermediate_tensors)
output = self.worker._execute_model_spmd(
execute_model_req, intermediate_tensors
)
# Pipeline model request and output to the next pipeline stage.
if isinstance(output, IntermediateTensors):
output = serialized_req, output
@@ -121,11 +125,12 @@ try:
def execute_model_ray(
self,
scheduler_output: Union["SchedulerOutput",
Tuple["SchedulerOutput",
"IntermediateTensors"]],
) -> Union["ModelRunnerOutput", Tuple["SchedulerOutput",
"IntermediateTensors"]]:
scheduler_output: Union[
"SchedulerOutput", Tuple["SchedulerOutput", "IntermediateTensors"]
],
) -> Union[
"ModelRunnerOutput", Tuple["SchedulerOutput", "IntermediateTensors"]
]:
# This method is used by Ray Compiled Graph to execute the model,
# and it needs a special logic of self.setup_device_if_necessary()
self.setup_device_if_necessary()
@@ -135,7 +140,8 @@ try:
else:
scheduler_output, intermediate_tensors = scheduler_output, None
output = self.worker.model_runner.execute_model(
scheduler_output, intermediate_tensors)
scheduler_output, intermediate_tensors
)
if isinstance(output, IntermediateTensors):
output = scheduler_output, output
elif not get_pp_group().is_last_rank:
@@ -171,12 +177,15 @@ def ray_is_available() -> bool:
def assert_ray_available():
"""Raise an exception if Ray is not available."""
if ray is None:
raise ValueError(f"Failed to import Ray: {ray_import_err}."
"Please install Ray with `pip install ray`.")
raise ValueError(
f"Failed to import Ray: {ray_import_err}."
"Please install Ray with `pip install ray`."
)
def _verify_bundles(placement_group: "PlacementGroup",
parallel_config: ParallelConfig, device_str: str):
def _verify_bundles(
placement_group: "PlacementGroup", parallel_config: ParallelConfig, device_str: str
):
"""Verify a given placement group has bundles located in the right place.
There are 2 rules.
@@ -184,7 +193,8 @@ def _verify_bundles(placement_group: "PlacementGroup",
- Fail if driver node is not included in a placement group.
"""
assert ray.is_initialized(), (
"Ray is not initialized although distributed-executor-backend is ray.")
"Ray is not initialized although distributed-executor-backend is ray."
)
pg_data = placement_group_table(placement_group)
# bundle_idx -> node_id
bundle_to_node_ids = pg_data["bundles_to_node_id"]
@@ -217,8 +227,13 @@ def _verify_bundles(placement_group: "PlacementGroup",
"unless you have fast interconnect across nodes, like "
"Infiniband. To resolve this issue, make sure you have more "
"than %d GPUs available at each node.",
parallel_config.tensor_parallel_size, device_str, len(bundles),
device_str, node_id, parallel_config.tensor_parallel_size)
parallel_config.tensor_parallel_size,
device_str,
len(bundles),
device_str,
node_id,
parallel_config.tensor_parallel_size,
)
def _wait_until_pg_ready(current_placement_group: "PlacementGroup"):
@@ -250,7 +265,9 @@ def _wait_until_pg_ready(current_placement_group: "PlacementGroup"):
" and make sure the IP addresses used by ray cluster"
" are the same as VLLM_HOST_IP environment variable"
" specified in each node if you are running on a multi-node.",
int(time.time() - s), placement_group_specs)
int(time.time() - s),
placement_group_specs,
)
try:
ray.get(pg_ready_ref, timeout=0)
@@ -259,7 +276,8 @@ def _wait_until_pg_ready(current_placement_group: "PlacementGroup"):
"Cannot provide a placement group of "
f"{placement_group_specs=} within {PG_WAIT_TIMEOUT} seconds. See "
"`ray status` and `ray list nodes` to make sure the cluster has "
"enough resources.") from None
"enough resources."
) from None
def _wait_until_pg_removed(current_placement_group: "PlacementGroup"):
@@ -274,8 +292,9 @@ def _wait_until_pg_removed(current_placement_group: "PlacementGroup"):
# Exponential backoff for warning print.
wait_interval *= 2
logger.info(
"Waiting for removing a placement group of specs for "
"%d seconds.", int(time.time() - s))
"Waiting for removing a placement group of specs for %d seconds.",
int(time.time() - s),
)
time.sleep(wait_interval)
@@ -306,19 +325,21 @@ def initialize_ray_cluster(
except ConnectionError:
logger.warning(
"No existing RAY instance detected. "
"A new instance will be launched with current node resources.")
ray.init(address=ray_address,
num_gpus=parallel_config.world_size,
runtime_env=parallel_config.ray_runtime_env)
"A new instance will be launched with current node resources."
)
ray.init(
address=ray_address,
num_gpus=parallel_config.world_size,
runtime_env=parallel_config.ray_runtime_env,
)
else:
ray.init(address=ray_address,
runtime_env=parallel_config.ray_runtime_env)
ray.init(address=ray_address, runtime_env=parallel_config.ray_runtime_env)
device_str = current_platform.ray_device_key
if not device_str:
raise ValueError(
f"current platform {current_platform.device_name} does not "
"support ray.")
f"current platform {current_platform.device_name} does not support ray."
)
# Create or get the placement group for worker processes
if parallel_config.placement_group:
@@ -337,8 +358,8 @@ def initialize_ray_cluster(
bundle_devices = bundle.get(device_str, 0)
if bundle_devices > 1:
raise ValueError(
"Placement group bundle cannot have more than 1 "
f"{device_str}.")
f"Placement group bundle cannot have more than 1 {device_str}."
)
if bundle_devices:
device_bundles += 1
if parallel_config.world_size > device_bundles:
@@ -346,10 +367,10 @@ def initialize_ray_cluster(
f"The number of required {device_str}s exceeds the total "
f"number of available {device_str}s in the placement group. "
f"Required number of devices: {parallel_config.world_size}. "
f"Total number of devices: {device_bundles}.")
f"Total number of devices: {device_bundles}."
)
else:
logger.info("No current placement group found. "
"Creating a new placement group.")
logger.info("No current placement group found. Creating a new placement group.")
num_devices_in_cluster = ray.cluster_resources().get(device_str, 0)
# Log a warning message and delay resource allocation failure response.
# Avoid immediate rejection to allow user-initiated placement group
@@ -357,12 +378,14 @@ def initialize_ray_cluster(
if parallel_config.world_size > num_devices_in_cluster:
logger.warning(
"The number of required %ss exceeds the total "
"number of available %ss in the placement group.", device_str,
device_str)
"number of available %ss in the placement group.",
device_str,
device_str,
)
# Create a new placement group
placement_group_specs: List[Dict[str, float]] = ([{
device_str: 1.0
} for _ in range(parallel_config.world_size)])
placement_group_specs: List[Dict[str, float]] = [
{device_str: 1.0} for _ in range(parallel_config.world_size)
]
# vLLM engine is also a worker to execute model with an accelerator,
# so it requires to have the device in a current node. Check if
@@ -375,14 +398,16 @@ def initialize_ray_cluster(
f"Current node has no {device_str} available. "
f"{current_node_resource=}. vLLM engine cannot start without "
f"{device_str}. Make sure you have at least 1 {device_str} "
f"available in a node {current_node_id=} {current_ip=}.")
f"available in a node {current_node_id=} {current_ip=}."
)
# This way, at least bundle is required to be created in a current
# node.
placement_group_specs[0][f"node:{current_ip}"] = 0.001
# By default, Ray packs resources as much as possible.
current_placement_group = ray.util.placement_group(
placement_group_specs, strategy="PACK")
placement_group_specs, strategy="PACK"
)
_wait_until_pg_ready(current_placement_group)
assert current_placement_group is not None
@@ -393,6 +418,7 @@ def initialize_ray_cluster(
def get_num_tpu_nodes() -> int:
from ray._private.accelerators import TPUAcceleratorManager
cluster_resources = ray.cluster_resources()
total_tpus = int(cluster_resources["TPU"])
tpus_per_node = TPUAcceleratorManager.get_current_node_num_accelerators()

View File

@@ -14,8 +14,7 @@ from vllm.executor.executor_base import ExecutorBase
from vllm.logger import init_logger
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.cache import worker_receiver_cache_from_config
from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,
run_method)
from vllm.utils import get_distributed_init_method, get_ip, get_open_port, run_method
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
from vllm.v1.executor.utils import get_and_update_mm_cache
from vllm.v1.outputs import AsyncModelRunnerOutput
@@ -25,14 +24,11 @@ logger = init_logger(__name__)
class UniProcExecutor(ExecutorBase):
uses_ray: bool = False
def _init_executor(self) -> None:
"""Initialize the worker and load the model.
"""
self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config,
rpc_rank=0)
"""Initialize the worker and load the model."""
self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config, rpc_rank=0)
distributed_init_method, rank, local_rank = self._distributed_args()
is_driver_worker = True
kwargs = dict(
@@ -43,24 +39,24 @@ class UniProcExecutor(ExecutorBase):
is_driver_worker=is_driver_worker,
)
self.mm_receiver_cache = worker_receiver_cache_from_config(
self.vllm_config, MULTIMODAL_REGISTRY, Lock())
self.vllm_config, MULTIMODAL_REGISTRY, Lock()
)
self.async_output_thread: Optional[ThreadPoolExecutor] = None
if self.max_concurrent_batches > 1:
self.async_output_thread = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="WorkerAsyncOutput")
max_workers=1, thread_name_prefix="WorkerAsyncOutput"
)
self.collective_rpc("init_worker", args=([kwargs], ))
self.collective_rpc("init_worker", args=([kwargs],))
self.collective_rpc("init_device")
self.collective_rpc("load_model")
def _distributed_args(self) -> tuple[str, int, int]:
"""Return (distributed_init_method, rank, local_rank)."""
distributed_init_method = get_distributed_init_method(
get_ip(), get_open_port())
distributed_init_method = get_distributed_init_method(get_ip(), get_open_port())
# set local rank as the device index if specified
device_info = self.vllm_config.device_config.device.__str__().split(
":")
device_info = self.vllm_config.device_config.device.__str__().split(":")
local_rank = int(device_info[1]) if len(device_info) > 1 else 0
return distributed_init_method, 0, local_rank
@@ -68,12 +64,14 @@ class UniProcExecutor(ExecutorBase):
def max_concurrent_batches(self) -> int:
return 2 if self.scheduler_config.async_scheduling else 1
def collective_rpc(self,
method: Union[str, Callable],
timeout: Optional[float] = None,
args: Tuple = (),
kwargs: Optional[Dict] = None,
non_block: bool = False) -> List[Any]:
def collective_rpc(
self,
method: Union[str, Callable],
timeout: Optional[float] = None,
args: Tuple = (),
kwargs: Optional[Dict] = None,
non_block: bool = False,
) -> List[Any]:
if kwargs is None:
kwargs = {}
if self.mm_receiver_cache is not None and method == "execute_model":
@@ -101,10 +99,13 @@ class UniProcExecutor(ExecutorBase):
return
def reinitialize_distributed(
self, reconfig_request: ReconfigureDistributedRequest) -> None:
self, reconfig_request: ReconfigureDistributedRequest
) -> None:
self.driver_worker.reinitialize_distributed(reconfig_request)
if reconfig_request.new_data_parallel_rank == \
ReconfigureRankType.SHUTDOWN_CURRENT_RANK:
if (
reconfig_request.new_data_parallel_rank
== ReconfigureRankType.SHUTDOWN_CURRENT_RANK
):
self.shutdown()
return
@@ -132,15 +133,16 @@ class ExecutorWithExternalLauncher(UniProcExecutor):
deterministic, all the engines will generate the same outputs,
and they don't need to synchronize the states with each other.
"""
uses_ray: bool = False
def _init_executor(self) -> None:
"""Initialize the worker and load the model.
"""
"""Initialize the worker and load the model."""
if envs.VLLM_USE_V1:
assert not envs.VLLM_ENABLE_V1_MULTIPROCESSING, \
("To get deterministic execution in V1, "
"please set VLLM_ENABLE_V1_MULTIPROCESSING=0")
assert not envs.VLLM_ENABLE_V1_MULTIPROCESSING, (
"To get deterministic execution in V1, "
"please set VLLM_ENABLE_V1_MULTIPROCESSING=0"
)
super()._init_executor()
def _distributed_args(self) -> tuple[str, int, int]:
@@ -168,6 +170,7 @@ class ExecutorWithExternalLauncher(UniProcExecutor):
"""
a, b = super().determine_num_available_blocks()
from vllm.distributed.parallel_state import get_world_group
cpu_group = get_world_group().cpu_group
a_tensor = torch.tensor([a], device="cpu", dtype=torch.int64)
b_tensor = torch.tensor([b], device="cpu", dtype=torch.int64)