@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 647214f3d5
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[V0 Deprecation] Remove V0 executors (#27142)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Oct 21, 2025 at 11:09 UTC 647214f3d52e35f2275cb288aa3460db79b07e90
31 files changed +424 -1042
tests/basic_correctness/test_basic_correctness.py
+2 -4
@@ -157,11 +157,9 @@ def test_models_distributed(
157 and distributed_executor_backend == "ray"
158 and attention_backend == ""
159 and test_suite == "L4"
160 + and enable_prompt_embeds
161 ): # noqa
161 - if enable_prompt_embeds:
162 - pytest.skip("enable_prompt_embeds does not work with ray compiled dag.")
163 - monkeypatch_context.setenv("VLLM_USE_RAY_SPMD_WORKER", "1")
164 - monkeypatch_context.setenv("VLLM_USE_RAY_COMPILED_DAG", "1")
162 + pytest.skip("enable_prompt_embeds does not work with ray compiled dag.")
163
164 if attention_backend:
165 monkeypatch_context.setenv(
tests/distributed/test_multi_node_assignment.py
+1 -1
@@ -18,8 +18,8 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
18
19 from vllm import initialize_ray_cluster
20 from vllm.config import ParallelConfig
21 -from vllm.executor.ray_utils import _wait_until_pg_removed
21 from vllm.utils.network_utils import get_ip
22 +from vllm.v1.executor.ray_utils import _wait_until_pg_removed
23
24 VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
25
tests/distributed/test_pipeline_parallel.py
+1 -3
@@ -305,10 +305,8 @@ def _compare_tp(
305 common_args.extend(["--max-num-seqs", f"{max_num_seqs}"])
306
307 if distributed_backend == "ray":
308 - # For V1, test Ray Compiled Graph for all the tests
308 + # Test Ray Compiled Graph for all the tests
309 pp_env = {
310 - "VLLM_USE_RAY_COMPILED_DAG": "1",
311 - "VLLM_USE_RAY_SPMD_WORKER": "1",
310 "VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL": "1",
311 }
312 # Temporary. Currently when zeromq + SPMD is used, it does not properly
tests/model_executor/model_loader/tensorizer_loader/conftest.py
+1 -1
@@ -9,7 +9,7 @@ from vllm.distributed import cleanup_dist_env_and_memory
9 from vllm.model_executor.model_loader import tensorizer as tensorizer_mod
10 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
11 from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
12 -from vllm.v1.executor.abstract import UniProcExecutor
12 +from vllm.v1.executor import UniProcExecutor
13 from vllm.v1.worker.worker_base import WorkerWrapperBase
14
15 MODEL_REF = "facebook/opt-125m"
tests/v1/engine/test_engine_core.py
+2 -1
@@ -15,7 +15,8 @@ from vllm.platforms import current_platform
15 from vllm.utils.torch_utils import set_default_torch_num_threads
16 from vllm.v1.engine import EngineCoreRequest
17 from vllm.v1.engine.core import EngineCore
18 -from vllm.v1.executor.abstract import Executor, UniProcExecutor
18 +from vllm.v1.executor.abstract import Executor
19 +from vllm.v1.executor.uniproc_executor import UniProcExecutor
20 from vllm.v1.kv_cache_interface import KVCacheConfig
21 from vllm.v1.outputs import ModelRunnerOutput
22
tools/pre_commit/check_pickle_imports.py
+4 -4
@@ -17,8 +17,6 @@ import regex as re
17 # add to this list if absolutely necessary and after careful security review.
18 ALLOWED_FILES = {
19 # pickle
20 - "vllm/v1/serial_utils.py",
21 - "vllm/v1/executor/multiproc_executor.py",
20 "vllm/multimodal/hasher.py",
21 "vllm/transformers_utils/config.py",
22 "vllm/model_executor/models/registry.py",
@@ -38,11 +36,13 @@ ALLOWED_FILES = {
36 "benchmarks/cutlass_benchmarks/w8a8_benchmarks.py",
37 "benchmarks/cutlass_benchmarks/sparse_benchmarks.py",
38 # cloudpickle
41 - "vllm/executor/mp_distributed_executor.py",
42 - "vllm/executor/ray_distributed_executor.py",
39 + "vllm/v1/executor/multiproc_executor.py",
40 + "vllm/v1/executor/ray_executor.py",
41 "vllm/entrypoints/llm.py",
42 "vllm/utils/__init__.py",
43 "tests/utils.py",
44 + # pickle and cloudpickle
45 + "vllm/v1/serial_utils.py",
46 }
47
48 PICKLE_RE = re.compile(
vllm/__init__.py
+2 -2
@@ -21,7 +21,7 @@ MODULE_ATTRS = {
21 "AsyncLLMEngine": ".engine.async_llm_engine:AsyncLLMEngine",
22 "LLMEngine": ".engine.llm_engine:LLMEngine",
23 "LLM": ".entrypoints.llm:LLM",
24 - "initialize_ray_cluster": ".executor.ray_utils:initialize_ray_cluster",
24 + "initialize_ray_cluster": ".v1.executor.ray_utils:initialize_ray_cluster",
25 "PromptType": ".inputs:PromptType",
26 "TextPrompt": ".inputs:TextPrompt",
27 "TokensPrompt": ".inputs:TokensPrompt",
@@ -45,7 +45,6 @@ if typing.TYPE_CHECKING:
45 from vllm.engine.async_llm_engine import AsyncLLMEngine
46 from vllm.engine.llm_engine import LLMEngine
47 from vllm.entrypoints.llm import LLM
48 - from vllm.executor.ray_utils import initialize_ray_cluster
48 from vllm.inputs import PromptType, TextPrompt, TokensPrompt
49 from vllm.model_executor.models import ModelRegistry
50 from vllm.outputs import (
@@ -62,6 +61,7 @@ if typing.TYPE_CHECKING:
61 )
62 from vllm.pooling_params import PoolingParams
63 from vllm.sampling_params import SamplingParams
64 + from vllm.v1.executor.ray_utils import initialize_ray_cluster
65
66 from ._bc_linter import bc_linter_include, bc_linter_skip
67 else:
vllm/config/parallel.py
+14 -8
@@ -25,11 +25,11 @@ if TYPE_CHECKING:
25 from ray.runtime_env import RuntimeEnv
26 from ray.util.placement_group import PlacementGroup
27
28 - from vllm.executor.executor_base import ExecutorBase
28 + from vllm.v1.executor import Executor
29 else:
30 RuntimeEnv = Any
31 PlacementGroup = Any
32 - ExecutorBase = Any
32 + Executor = Any
33
34 logger = init_logger(__name__)
35
@@ -189,7 +189,7 @@ class ParallelConfig:
189 """ray distributed model workers placement group."""
190
191 distributed_executor_backend: (
192 - str | DistributedExecutorBackend | type[ExecutorBase] | None
192 + str | DistributedExecutorBackend | type[Executor] | None
193 ) = None
194 """Backend to use for distributed model
195 workers, either "ray" or "mp" (multiprocessing). If the product
@@ -511,7 +511,7 @@ class ParallelConfig:
511 # We use multiprocessing by default if world_size fits on the
512 # current node and we aren't in a ray placement group.
513
514 - from vllm.executor import ray_utils
514 + from vllm.v1.executor import ray_utils
515
516 backend: DistributedExecutorBackend = "mp"
517 ray_found = ray_utils.ray_is_available()
@@ -553,6 +553,12 @@ class ParallelConfig:
553 if self.distributed_executor_backend is None and self.world_size == 1:
554 self.distributed_executor_backend = "uni"
555
556 + if self.max_parallel_loading_workers is not None:
557 + logger.warning(
558 + "max_parallel_loading_workers is currently "
559 + "not supported and will be ignored."
560 + )
561 +
562 @property
563 def use_ray(self) -> bool:
564 return self.distributed_executor_backend == "ray" or (
@@ -563,7 +569,7 @@ class ParallelConfig:
569 @model_validator(mode="after")
570 def _verify_args(self) -> Self:
571 # Lazy import to avoid circular import
566 - from vllm.executor.executor_base import ExecutorBase
572 + from vllm.v1.executor import Executor
573
574 # Enable batch invariance settings if requested
575 if vllm_is_batch_invariant():
@@ -574,17 +580,17 @@ class ParallelConfig:
580 and not isinstance(self.distributed_executor_backend, str)
581 and not (
582 isinstance(self.distributed_executor_backend, type)
577 - and issubclass(self.distributed_executor_backend, ExecutorBase)
583 + and issubclass(self.distributed_executor_backend, Executor)
584 )
585 ):
586 raise ValueError(
587 "Unrecognized distributed executor backend "
588 f"{self.distributed_executor_backend}. Supported "
589 "values are 'ray', 'mp' 'uni', 'external_launcher', "
584 - " custom ExecutorBase subclass or its import path."
590 + " custom Executor subclass or its import path."
591 )
592 if self.use_ray:
587 - from vllm.executor import ray_utils
593 + from vllm.v1.executor import ray_utils
594
595 ray_utils.assert_ray_available()
596
vllm/config/scheduler.py
-6
@@ -107,12 +107,6 @@ class SchedulerConfig:
107 NOTE: This is not currently configurable. It will be overridden by
108 max_num_batched_tokens in case max multimodal embedding size is larger."""
109
110 - send_delta_data: bool = False
111 - """Private API. If used, scheduler sends delta data to
112 - workers instead of an entire data. It should be enabled only
113 - when SPMD worker architecture is enabled. I.e.,
114 - VLLM_USE_RAY_SPMD_WORKER=1"""
115 -
110 policy: SchedulerPolicy = "fcfs"
111 """The scheduling policy to use:\n
112 - "fcfs" means first come first served, i.e. requests are handled in order
vllm/distributed/device_communicators/tpu_communicator.py
+1 -1
@@ -31,7 +31,7 @@ if not USE_TPU_INFERENCE:
31 )
32
33 if USE_RAY:
34 - from vllm.executor import ray_utils
34 + from vllm.v1.executor import ray_utils
35
36
37 class TpuCommunicator(DeviceCommunicatorBase):
vllm/engine/arg_utils.py
+3 -4
@@ -88,12 +88,12 @@ from vllm.utils.network_utils import get_ip
88 from vllm.v1.sample.logits_processor import LogitsProcessor
89
90 if TYPE_CHECKING:
91 - from vllm.executor.executor_base import ExecutorBase
91 from vllm.model_executor.layers.quantization import QuantizationMethods
92 from vllm.model_executor.model_loader import LoadFormats
93 from vllm.usage.usage_lib import UsageContext
94 + from vllm.v1.executor import Executor
95 else:
96 - ExecutorBase = Any
96 + Executor = Any
97 QuantizationMethods = Any
98 LoadFormats = Any
99 UsageContext = Any
@@ -369,7 +369,7 @@ class EngineArgs:
369 # is intended for expert use only. The API may change without
370 # notice.
371 distributed_executor_backend: (
372 - str | DistributedExecutorBackend | type[ExecutorBase] | None
372 + str | DistributedExecutorBackend | type[Executor] | None
373 ) = ParallelConfig.distributed_executor_backend
374 # number of P/D disaggregation (or other disaggregation) workers
375 pipeline_parallel_size: int = ParallelConfig.pipeline_parallel_size
@@ -1549,7 +1549,6 @@ class EngineArgs:
1549 disable_chunked_mm_input=self.disable_chunked_mm_input,
1550 is_multimodal_model=model_config.is_multimodal_model,
1551 is_encoder_decoder=model_config.is_encoder_decoder,
1552 - send_delta_data=(envs.VLLM_USE_RAY_SPMD_WORKER and parallel_config.use_ray),
1552 policy=self.scheduling_policy,
1553 scheduler_cls=self.scheduler_cls,
1554 max_num_partial_prefills=self.max_num_partial_prefills,
vllm/entrypoints/cli/serve.py
+1 -1
@@ -26,7 +26,7 @@ from vllm.utils import (
26 from vllm.utils.network_utils import get_tcp_uri
27 from vllm.v1.engine.core import EngineCoreProc
28 from vllm.v1.engine.utils import CoreEngineProcManager, launch_core_engines
29 -from vllm.v1.executor.abstract import Executor
29 +from vllm.v1.executor import Executor
30 from vllm.v1.metrics.prometheus import setup_multiprocess_prometheus
31 from vllm.v1.utils import APIServerProcessManager, wait_for_completion_or_failure
32
vllm/envs.py
+1 -22
@@ -56,8 +56,6 @@ if TYPE_CHECKING:
56 VLLM_XLA_CHECK_RECOMPILATION: bool = False
57 VLLM_FUSED_MOE_CHUNK_SIZE: int = 64 * 1024
58 VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING: bool = True
59 - VLLM_USE_RAY_SPMD_WORKER: bool = False
60 - VLLM_USE_RAY_COMPILED_DAG: bool = False
59 VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"
60 VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
61 VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True
@@ -623,22 +621,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
621 "VLLM_CPU_MOE_PREPACK": lambda: bool(int(os.getenv("VLLM_CPU_MOE_PREPACK", "1"))),
622 # (CPU backend only) whether to use SGL kernels, optimized for small batch.
623 "VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
626 - # If the env var is set, then all workers will execute as separate
627 - # processes from the engine, and we use the same mechanism to trigger
628 - # execution on all workers.
629 - # Run vLLM with VLLM_USE_RAY_SPMD_WORKER=1 to enable it.
630 - "VLLM_USE_RAY_SPMD_WORKER": lambda: bool(
631 - int(os.getenv("VLLM_USE_RAY_SPMD_WORKER", "0"))
632 - ),
633 - # If the env var is set, it uses the Ray's Compiled Graph
634 - # (previously known as ADAG) API which optimizes the
635 - # control plane overhead.
636 - # Run vLLM with VLLM_USE_RAY_COMPILED_DAG=1 to enable it.
637 - # Note that this variable is set to 1 in V1 by default
638 - # when ray distributed executor is used.
639 - "VLLM_USE_RAY_COMPILED_DAG": lambda: bool(
640 - int(os.getenv("VLLM_USE_RAY_COMPILED_DAG", "0"))
641 - ),
624 # If the env var is set, Ray Compiled Graph uses the specified
625 # channel type to communicate between workers belonging to
626 # different pipeline-parallel stages.
@@ -646,20 +628,17 @@ environment_variables: dict[str, Callable[[], Any]] = {
628 # - "auto": use the default channel type
629 # - "nccl": use NCCL for communication
630 # - "shm": use shared memory and gRPC for communication
649 - # This flag is ignored if VLLM_USE_RAY_COMPILED_DAG is not set.
631 "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE": env_with_choices(
632 "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "auto", ["auto", "nccl", "shm"]
633 ),
634 # If the env var is set, it enables GPU communication overlap
654 - # (experimental feature) in Ray's Compiled Graph. This flag is ignored if
655 - # VLLM_USE_RAY_COMPILED_DAG is not set.
635 + # (experimental feature) in Ray's Compiled Graph.
636 "VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM": lambda: bool(
637 int(os.getenv("VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM", "0"))
638 ),
639 # If the env var is set, it uses a Ray Communicator wrapping
640 # vLLM's pipeline parallelism communicator to interact with Ray's
641 # Compiled Graph. Otherwise, it uses Ray's NCCL communicator.
662 - # This flag is ignored if VLLM_USE_RAY_COMPILED_DAG is not set.
642 "VLLM_USE_RAY_WRAPPED_PP_COMM": lambda: bool(
643 int(os.getenv("VLLM_USE_RAY_WRAPPED_PP_COMM", "1"))
644 ),
vllm/executor/__init__.py
vllm/executor/executor_base.py deleted
-393
@@ -1,393 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import asyncio
5 -import time
6 -from abc import ABC, abstractmethod
7 -from collections.abc import Awaitable, Callable
8 -from functools import cached_property
9 -from typing import Any
10 -
11 -from typing_extensions import TypeVar
12 -
13 -import vllm.platforms
14 -from vllm.config import VllmConfig
15 -from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
16 -from vllm.logger import init_logger
17 -from vllm.lora.request import LoRARequest
18 -from vllm.sequence import ExecuteModelRequest
19 -from vllm.tasks import SupportedTask
20 -from vllm.utils.async_utils import make_async
21 -from vllm.v1.outputs import SamplerOutput
22 -from vllm.v1.worker.worker_base import WorkerBase
23 -
24 -logger = init_logger(__name__)
25 -
26 -_R = TypeVar("_R", default=Any)
27 -
28 -
29 -class ExecutorBase(ABC):
30 - """Base class for all executors.
31 -
32 - An executor is responsible for executing the model on one device,
33 - or it can be a distributed executor
34 - that can execute the model on multiple devices.
35 - """
36 -
37 - uses_ray: bool # whether the executor uses Ray for orchestration.
38 - supports_pp: bool = False # whether the executor supports PP
39 -
40 - def __init__(
41 - self,
42 - vllm_config: VllmConfig,
43 - ) -> None:
44 - self.vllm_config = vllm_config
45 - self.model_config = vllm_config.model_config
46 - self.cache_config = vllm_config.cache_config
47 - self.lora_config = vllm_config.lora_config
48 - self.load_config = vllm_config.load_config
49 - self.parallel_config = vllm_config.parallel_config
50 - self.scheduler_config = vllm_config.scheduler_config
51 - self.device_config = vllm_config.device_config
52 - self.speculative_config = vllm_config.speculative_config
53 - self.observability_config = vllm_config.observability_config
54 - self._init_executor()
55 - self.is_sleeping = False
56 - self.sleeping_tags: set[str] = set()
57 - self.kv_output_aggregator: KVOutputAggregator | None = None
58 -
59 - @abstractmethod
60 - def _init_executor(self) -> None:
61 - raise NotImplementedError
62 -
63 - @abstractmethod
64 - def collective_rpc(
65 - self,
66 - method: str | Callable[[WorkerBase], _R],
67 - timeout: float | None = None,
68 - args: tuple = (),
69 - kwargs: dict[str, Any] | None = None,
70 - ) -> list[_R]:
71 - """
72 - Execute an RPC call on all workers.
73 -
74 - Args:
75 - method: Name of the worker method to execute, or a callable that
76 - is serialized and sent to all workers to execute.
77 -
78 - If the method is a callable, it should accept an additional
79 - `self` argument, in addition to the arguments passed in `args`
80 - and `kwargs`. The `self` argument will be the worker object.
81 - timeout: Maximum time in seconds to wait for execution. Raises a
82 - [`TimeoutError`][] on timeout. `None` means wait indefinitely.
83 - args: Positional arguments to pass to the worker method.
84 - kwargs: Keyword arguments to pass to the worker method.
85 -
86 - Returns:
87 - A list containing the results from each worker.
88 -
89 - Note:
90 - It is recommended to use this API to only pass control messages,
91 - and set up data-plane communication to pass data.
92 - """
93 - raise NotImplementedError
94 -
95 - def determine_num_available_blocks(self) -> tuple[int, int]:
96 - """Determine the number of available blocks for the GPU KV cache and
97 - swappable CPU KV cache.
98 -
99 - Normally, this should simply delegate to the underlying Worker. Some
100 - ExecutorBase may require modification of the result, e.g. to ensure the
101 - selected cache sizes are compatible with all workers.
102 -
103 - Returns a tuple `(num_gpu_blocks, num_cpu_blocks)`, where
104 - `num_gpu_blocks` are blocks that are "active" on the device and can be
105 - appended to.
106 - `num_cpu_blocks` refers to "swapped" blocks in CPU memory and cannot be
107 - appended to.
108 - """
109 - results = self.collective_rpc("determine_num_available_blocks")
110 - a = min([r[0] for r in results])
111 - b = min([r[1] for r in results])
112 - return a, b
113 -
114 - def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks) -> None:
115 - """Initialize the KV cache by invoking the underlying worker."""
116 - # NOTE: This is logged in the executor because there can be >1 workers.
117 - logger.info(
118 - "# %s blocks: %d, # CPU blocks: %d",
119 - vllm.platforms.current_platform.device_name,
120 - num_gpu_blocks,
121 - num_cpu_blocks,
122 - )
123 - max_concurrency = (
124 - num_gpu_blocks
125 - * self.cache_config.block_size
126 - / self.model_config.max_model_len
127 - )
128 - logger.info(
129 - "Maximum concurrency for %s tokens per request: %.2fx",
130 - self.model_config.max_model_len,
131 - max_concurrency,
132 - )
133 -
134 - self.cache_config.num_gpu_blocks = num_gpu_blocks
135 - self.cache_config.num_cpu_blocks = num_cpu_blocks
136 -
137 - self.collective_rpc("initialize_cache", args=(num_gpu_blocks, num_cpu_blocks))
138 -
139 - @cached_property # Avoid unnecessary RPC calls
140 - def supported_tasks(self) -> tuple[SupportedTask, ...]:
141 - output = self.collective_rpc("get_supported_tasks")
142 - return output[0]
143 -
144 - def execute_model(
145 - self, execute_model_req: ExecuteModelRequest
146 - ) -> list[SamplerOutput]:
147 - output = self.collective_rpc("execute_model", args=(execute_model_req,))
148 - assert output[0] is not None
149 - return output[0]
150 -
151 - def stop_remote_worker_execution_loop(self) -> None:
152 - """Releases parallel workers from model loop."""
153 - return
154 -
155 - def add_lora(self, lora_request: LoRARequest) -> bool:
156 - assert lora_request.lora_int_id > 0, "lora_id must be greater than 0."
157 - return all(self.collective_rpc("add_lora", args=(lora_request,)))
158 -
159 - def remove_lora(self, lora_id: int) -> bool:
160 - assert lora_id > 0, "lora_id must be greater than 0."
161 - return all(self.collective_rpc("remove_lora", args=(lora_id,)))
162 -
163 - def pin_lora(self, lora_id: int) -> bool:
164 - assert lora_id > 0, "lora_id must be greater than 0."
165 - return all(self.collective_rpc("pin_lora", args=(lora_id,)))
166 -
167 - def list_loras(self) -> set[int]:
168 - sets = self.collective_rpc("list_loras")
169 - for s in sets:
170 - assert s == sets[0], "All workers should have the same LORAs."
171 - return sets[0]
172 -
173 - def reset_mm_cache(self) -> None:
174 - """Reset the multi-modal cache in each worker."""
175 - self.collective_rpc("reset_mm_cache")
176 -
177 - def start_profile(self) -> None:
178 - self.collective_rpc("start_profile")
179 -
180 - def stop_profile(self) -> None:
181 - self.collective_rpc("stop_profile")
182 -
183 - def sleep(self, level: int = 1):
184 - if self.is_sleeping:
185 - logger.warning("Executor is already sleeping.")
186 - return
187 - time_before_sleep = time.perf_counter()
188 - self.collective_rpc("sleep", kwargs=dict(level=level))
189 - time_after_sleep = time.perf_counter()
190 - self.sleeping_tags = {"weights", "kv_cache"}
191 - self.is_sleeping = True
192 - logger.info(
193 - "It took %.6f seconds to fall asleep.", time_after_sleep - time_before_sleep
194 - )
195 -
196 - def wake_up(self, tags: list[str] | None = None):
197 - if not self.is_sleeping:
198 - logger.warning("Executor is not sleeping.")
199 - return
200 - if tags:
201 - for tag in tags:
202 - if tag not in self.sleeping_tags:
203 - logger.warning(
204 - "Tag %s is not in sleeping tags %s", tag, self.sleeping_tags
205 - )
206 - return
207 - time_before_wakeup = time.perf_counter()
208 - self.collective_rpc("wake_up", kwargs=dict(tags=tags))
209 - time_after_wakeup = time.perf_counter()
210 - logger.info(
211 - "It took %.6f seconds to wake up tags %s.",
212 - time_after_wakeup - time_before_wakeup,
213 - tags if tags is not None else self.sleeping_tags,
214 - )
215 - if tags:
216 - for tag in tags:
217 - self.sleeping_tags.remove(tag)
218 - else:
219 - self.sleeping_tags.clear()
220 - if not self.sleeping_tags:
221 - self.is_sleeping = False
222 -
223 - def save_sharded_state(
224 - self,
225 - path: str,
226 - pattern: str | None = None,
227 - max_size: int | None = None,
228 - ) -> None:
229 - self.collective_rpc(
230 - "save_sharded_state",
231 - kwargs=dict(path=path, pattern=pattern, max_size=max_size),
232 - )
233 -
234 - @abstractmethod
235 - def check_health(self) -> None:
236 - """Checks if the executor is healthy. If not, it should raise an
237 - exception."""
238 - raise NotImplementedError
239 -
240 - def shutdown(self) -> None:
241 - """Shutdown the executor."""
242 - self.collective_rpc("shutdown")
243 -
244 - async def execute_model_async(
245 - self, execute_model_req: ExecuteModelRequest
246 - ) -> list[SamplerOutput]:
247 - """Executes one model step on the given sequences."""
248 - output = await make_async(self.execute_model)(execute_model_req)
249 - return output
250 -
251 - async def stop_remote_worker_execution_loop_async(self) -> None:
252 - """Releases parallel workers from model loop."""
253 - return
254 -
255 - async def check_health_async(self) -> None:
256 - """Checks if the executor is healthy. If not, it should raise an
257 - exception."""
258 - self.check_health()
259 -
260 - def init_kv_output_aggregator(self, finished_count: int | None) -> None:
261 - """Init KVOutputAggregator"""
262 - self.kv_output_aggregator = KVOutputAggregator(
263 - finished_count or self.parallel_config.world_size
264 - )
265 -
266 -
267 -class DistributedExecutorBase(ExecutorBase):
268 - """Abstract superclass of distributed executor implementations."""
269 -
270 - def __init__(self, *args, **kwargs):
271 - # This is non-None when the execute model loop is running
272 - # in the parallel workers. It's a coroutine in the AsyncLLMEngine case.
273 - self.parallel_worker_tasks: Any | Awaitable[Any] | None = None
274 -
275 - super().__init__(*args, **kwargs)
276 -
277 - def execute_model(
278 - self,
279 - execute_model_req: ExecuteModelRequest,
280 - ) -> list[SamplerOutput]:
281 - # TODO: unify into collective_rpc
282 - if self.parallel_worker_tasks is None:
283 - self.parallel_worker_tasks = self._run_workers(
284 - "start_worker_execution_loop",
285 - async_run_tensor_parallel_workers_only=True,
286 - )
287 -
288 - # Only the driver worker returns the sampling results.
289 - driver_outputs = self._driver_execute_model(execute_model_req)
290 - assert driver_outputs is not None
291 - return driver_outputs
292 -
293 - def stop_remote_worker_execution_loop(self) -> None:
294 - if self.parallel_worker_tasks is None:
295 - return
296 -
297 - self._driver_execute_model(execute_model_req=None)
298 - parallel_worker_tasks = self.parallel_worker_tasks
299 - self.parallel_worker_tasks = None
300 - # Ensure that workers exit model loop cleanly
301 - # (this will raise otherwise)
302 - self._wait_for_tasks_completion(parallel_worker_tasks)
303 -
304 - @abstractmethod
305 - def _driver_execute_model(
306 - self, execute_model_req: ExecuteModelRequest | None
307 - ) -> list[SamplerOutput] | None:
308 - """Run execute_model in the driver worker.
309 -
310 - Passing None will cause the driver to stop the model execution loop
311 - running in each of the remote workers. In this case, this method
312 - returns None. Otherwise, this method returns the model output.
313 - """
314 - raise NotImplementedError
315 -
316 - def collective_rpc(
317 - self,
318 - method: str | Callable,
319 - timeout: float | None = None,
320 - args: tuple = (),
321 - kwargs: dict[str, Any] | None = None,
322 - ) -> list[Any]:
323 - return self._run_workers(method, *args, **(kwargs or {}))
324 -
325 - @abstractmethod
326 - def _run_workers(
327 - self,
328 - method: str | Callable,
329 - *args,
330 - async_run_tensor_parallel_workers_only: bool = False,
331 - max_concurrent_workers: int | None = None,
332 - **kwargs,
333 - ) -> Any:
334 - """Runs the given method on all workers.
335 -
336 - Args:
337 - async_run_tensor_parallel_workers_only: If True the method will be
338 - run only in the remote TP workers, not the driver worker.
339 - It will also be run asynchronously and return a list of futures
340 - rather than blocking on the results.
341 -
342 - # TODO: simplify and merge with collective_rpc
343 - """
344 - raise NotImplementedError
345 -
346 - @abstractmethod
347 - def _wait_for_tasks_completion(self, parallel_worker_tasks: Any) -> None:
348 - """Wait for futures returned from _run_workers() with
349 - async_run_remote_workers_only to complete."""
350 - raise NotImplementedError
351 -
352 - async def execute_model_async(
353 - self, execute_model_req: ExecuteModelRequest
354 - ) -> list[SamplerOutput]:
355 - if self.parallel_worker_tasks is None:
356 - # Start model execution loop running in the parallel workers
357 - self.parallel_worker_tasks = asyncio.create_task(
358 - self._start_worker_execution_loop()
359 - )
360 -
361 - # Only the driver worker returns the sampling results.
362 - return await self._driver_execute_model_async(execute_model_req)
363 -
364 - async def stop_remote_worker_execution_loop_async(self) -> None:
365 - if self.parallel_worker_tasks is None:
366 - return
367 -
368 - await self._driver_execute_model_async()
369 - parallel_worker_tasks = self.parallel_worker_tasks
370 - self.parallel_worker_tasks = None
371 - # Ensure that workers exit model loop cleanly
372 - # (this will raise otherwise)
373 - await parallel_worker_tasks
374 -
375 - @abstractmethod
376 - async def _driver_execute_model_async(
377 - self,
378 - execute_model_req: ExecuteModelRequest | None = None,
379 - ) -> list[SamplerOutput]:
380 - """Execute the model asynchronously in the driver worker.
381 -
382 - Passing None will cause the driver to stop the model execution
383 - loop running in each of the remote workers.
384 - """
385 - raise NotImplementedError
386 -
387 - @abstractmethod
388 - async def _start_worker_execution_loop(self):
389 - """Run execution loop on all workers. It guarantees all workers run
390 - the loop or None of them is running the loop. Loop can be stopped by
391 - `stop_remote_worker_execution_loop`.
392 - The API is idempotent (guarantee only 1 loop run at any moment)."""
393 - raise NotImplementedError
vllm/executor/msgspec_utils.py deleted
-36
@@ -1,36 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from array import array
5 -from typing import Any
6 -
7 -from vllm.multimodal.inputs import MultiModalKwargs
8 -from vllm.sequence import VLLM_TOKEN_ID_ARRAY_TYPE
9 -
10 -
11 -def encode_hook(obj: Any) -> Any:
12 - """Custom msgspec enc hook that supports array types and MultiModalKwargs.
13 -
14 - See https://jcristharif.com/msgspec/api.html#msgspec.msgpack.Encoder
15 - """
16 - if isinstance(obj, array):
17 - assert obj.typecode == VLLM_TOKEN_ID_ARRAY_TYPE, (
18 - f"vLLM array type should use '{VLLM_TOKEN_ID_ARRAY_TYPE}' type. "
19 - f"Given array has a type code of {obj.typecode}."
20 - )
21 - return obj.tobytes()
22 - if isinstance(obj, MultiModalKwargs):
23 - return dict(obj)
24 -
25 -
26 -def decode_hook(type: type, obj: Any) -> Any:
27 - """Custom msgspec dec hook that supports array types and MultiModalKwargs.
28 -
29 - See https://jcristharif.com/msgspec/api.html#msgspec.msgpack.Encoder
30 - """
31 - if type is array:
32 - deserialized = array(VLLM_TOKEN_ID_ARRAY_TYPE)
33 - deserialized.frombytes(obj)
34 - return deserialized
35 - if type is MultiModalKwargs:
36 - return MultiModalKwargs(obj)
vllm/sequence.py
-10
@@ -5,7 +5,6 @@
5 from dataclasses import dataclass
6 from typing import TYPE_CHECKING, Any
7
8 -import msgspec
8 import torch
9
10 if TYPE_CHECKING:
@@ -92,12 +91,3 @@ class IntermediateTensors:
91
92 def __repr__(self) -> str:
93 return f"IntermediateTensors(tensors={self.tensors})"
95 -
96 -
97 -class ExecuteModelRequest(
98 - msgspec.Struct,
99 - array_like=True, # type: ignore[call-arg]
100 - omit_defaults=True,
101 -): # type: ignore[call-arg]
102 - # Placeholder. Remove.
103 - pass
vllm/transformers_utils/config.py
+1 -1
@@ -943,7 +943,7 @@ def maybe_register_config_serialize_by_value() -> None:
943 cloudpickle.register_pickle_by_value(transformers_modules)
944
945 # ray vendors its own version of cloudpickle
946 - from vllm.executor.ray_utils import ray
946 + from vllm.v1.executor.ray_utils import ray
947
948 if ray:
949 ray.cloudpickle.register_pickle_by_value(transformers_modules)
vllm/v1/engine/async_llm.py
+1 -1
@@ -39,7 +39,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
39 from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector
40 from vllm.v1.engine.parallel_sampling import ParentRequest
41 from vllm.v1.engine.processor import Processor
42 -from vllm.v1.executor.abstract import Executor
42 +from vllm.v1.executor import Executor
43 from vllm.v1.metrics.loggers import (
44 StatLoggerFactory,
45 StatLoggerManager,
vllm/v1/engine/core.py
+2 -11
@@ -60,7 +60,7 @@ from vllm.v1.engine.utils import (
60 EngineZmqAddresses,
61 get_device_indices,
62 )
63 -from vllm.v1.executor.abstract import Executor
63 +from vllm.v1.executor import Executor
64 from vllm.v1.kv_cache_interface import KVCacheConfig
65 from vllm.v1.metrics.stats import SchedulerStats
66 from vllm.v1.outputs import ModelRunnerOutput
@@ -322,7 +322,6 @@ class EngineCore:
322 with self.log_error_detail(scheduler_output):
323 model_output = self.model_executor.execute_model(scheduler_output)
324
325 - assert isinstance(model_output, ModelRunnerOutput)
325 engine_core_outputs = self.scheduler.update_from_output(
326 scheduler_output, model_output
327 )
@@ -364,7 +363,7 @@ class EngineCore:
363 if self.scheduler.has_requests():
364 scheduler_output = self.scheduler.schedule()
365 future = self.model_executor.execute_model(scheduler_output, non_block=True)
367 - batch_queue.appendleft((future, scheduler_output)) # type: ignore[arg-type]
366 + batch_queue.appendleft((future, scheduler_output))
367
368 model_executed = scheduler_output.total_num_scheduled_tokens > 0
369 if (
@@ -463,14 +462,6 @@ class EngineCore:
462 ) -> list[_R]:
463 return self.model_executor.collective_rpc(method, timeout, args, kwargs)
464
466 - def save_tensorized_model(
467 - self,
468 - tensorizer_config,
469 - ) -> None:
470 - self.model_executor.save_tensorized_model(
471 - tensorizer_config=tensorizer_config,
472 - )
473 -
465 def preprocess_add_request(self, request: EngineCoreRequest) -> tuple[Request, int]:
466 """Preprocess the request.
467
vllm/v1/engine/core_client.py
+1 -1
@@ -46,7 +46,7 @@ from vllm.v1.engine.utils import (
46 CoreEngineProcManager,
47 launch_core_engines,
48 )
49 -from vllm.v1.executor.abstract import Executor
49 +from vllm.v1.executor import Executor
50 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr
51
52 logger = init_logger(__name__)
vllm/v1/engine/llm_engine.py
+1 -1
@@ -32,7 +32,7 @@ from vllm.v1.engine.core_client import EngineCoreClient
32 from vllm.v1.engine.output_processor import OutputProcessor
33 from vllm.v1.engine.parallel_sampling import ParentRequest
34 from vllm.v1.engine.processor import Processor
35 -from vllm.v1.executor.abstract import Executor
35 +from vllm.v1.executor import Executor
36 from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
37 from vllm.v1.metrics.reader import Metric, get_metrics_snapshot
38 from vllm.v1.metrics.stats import IterationStats
vllm/v1/engine/utils.py
+1 -1
@@ -23,7 +23,7 @@ from vllm.ray.ray_env import get_env_vars_to_copy
23 from vllm.utils import get_mp_context
24 from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
25 from vllm.v1.engine.coordinator import DPCoordinator
26 -from vllm.v1.executor.abstract import Executor
26 +from vllm.v1.executor import Executor
27 from vllm.v1.utils import get_engine_client_zmq_addr, shutdown
28
29 if TYPE_CHECKING:
vllm/v1/executor/__init__.py
+6
@@ -0,0 +1,6 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from .abstract import Executor
4 +from .uniproc_executor import UniProcExecutor
5 +
6 +__all__ = ["Executor", "UniProcExecutor"]
vllm/v1/executor/abstract.py
+229 -40
@@ -1,31 +1,40 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
3 +import time
4 +from abc import ABC, abstractmethod
5 from collections.abc import Callable
6 from concurrent.futures import Future
6 -from typing import Any
7 -
8 -import torch
9 -import torch.distributed as dist
7 +from functools import cached_property
8 +from typing import Literal, TypeVar, overload
9
10 from vllm.config import VllmConfig
12 -from vllm.executor.executor_base import ExecutorBase
13 -from vllm.executor.uniproc_executor import ( # noqa
14 - ExecutorWithExternalLauncher as ExecutorWithExternalLauncherV0,
15 -)
16 -from vllm.executor.uniproc_executor import UniProcExecutor as UniProcExecutorV0 # noqa
11 +from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
12 +from vllm.logger import init_logger
13 +from vllm.lora.request import LoRARequest
14 +from vllm.tasks import SupportedTask
15 from vllm.utils.import_utils import resolve_obj_by_qualname
16 from vllm.v1.core.sched.output import SchedulerOutput
17 +from vllm.v1.engine import ReconfigureDistributedRequest
18 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
19 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
20 +from vllm.v1.worker.worker_base import WorkerBase
21 +
22 +logger = init_logger(__name__)
23 +
24 +_R = TypeVar("_R")
25
26 FailureCallback = Callable[[], None]
27
28
25 -class Executor(ExecutorBase):
29 +class Executor(ABC):
30 + """Abstract base class for vLLM executors."
31 +
32 + An executor is responsible for executing the model on one device,
33 + or it can be a distributed executor that can execute the model on multiple devices.
34 """
27 - Abstract class for v1 executors, mainly define some methods for v1.
28 - For methods shared by v0 and v1, define them in ExecutorBase"""
35 +
36 + uses_ray: bool = False # whether the executor uses Ray for orchestration.
37 + supports_pp: bool = False # whether the executor supports PP
38
39 @staticmethod
40 def get_class(vllm_config: VllmConfig) -> type["Executor"]:
@@ -34,16 +43,14 @@ class Executor(ExecutorBase):
43 distributed_executor_backend = parallel_config.distributed_executor_backend
44 # distributed_executor_backend must be set in VllmConfig.__post_init__
45 if isinstance(distributed_executor_backend, type):
37 - if not issubclass(distributed_executor_backend, ExecutorBase):
46 + if not issubclass(distributed_executor_backend, Executor):
47 raise TypeError(
48 "distributed_executor_backend must be a subclass of "
40 - f"ExecutorBase. Got {distributed_executor_backend}."
49 + f"Executor. Got {distributed_executor_backend}."
50 )
51 executor_class = distributed_executor_backend
52 elif distributed_executor_backend == "ray":
44 - from vllm.v1.executor.ray_distributed_executor import ( # noqa
45 - RayDistributedExecutor,
46 - )
53 + from vllm.v1.executor.ray_executor import RayDistributedExecutor
54
55 executor_class = RayDistributedExecutor
56 elif distributed_executor_backend == "mp":
@@ -51,6 +58,8 @@ class Executor(ExecutorBase):
58
59 executor_class = MultiprocExecutor
60 elif distributed_executor_backend == "uni":
61 + from vllm.v1.executor.uniproc_executor import UniProcExecutor
62 +
63 executor_class = UniProcExecutor
64 elif distributed_executor_backend == "external_launcher":
65 # TODO: make v1 scheduling deterministic
@@ -58,10 +67,10 @@ class Executor(ExecutorBase):
67 executor_class = ExecutorWithExternalLauncher
68 elif isinstance(distributed_executor_backend, str):
69 executor_class = resolve_obj_by_qualname(distributed_executor_backend)
61 - if not issubclass(executor_class, ExecutorBase):
70 + if not issubclass(executor_class, Executor):
71 raise TypeError(
72 "distributed_executor_backend must be a subclass of "
64 - f"ExecutorBase. Got {executor_class}."
73 + f"Executor. Got {executor_class}."
74 )
75 else:
76 raise ValueError(
@@ -69,6 +78,29 @@ class Executor(ExecutorBase):
78 )
79 return executor_class
80
81 + def __init__(
82 + self,
83 + vllm_config: VllmConfig,
84 + ) -> None:
85 + self.vllm_config = vllm_config
86 + self.model_config = vllm_config.model_config
87 + self.cache_config = vllm_config.cache_config
88 + self.lora_config = vllm_config.lora_config
89 + self.load_config = vllm_config.load_config
90 + self.parallel_config = vllm_config.parallel_config
91 + self.scheduler_config = vllm_config.scheduler_config
92 + self.device_config = vllm_config.device_config
93 + self.speculative_config = vllm_config.speculative_config
94 + self.observability_config = vllm_config.observability_config
95 + self._init_executor()
96 + self.is_sleeping = False
97 + self.sleeping_tags: set[str] = set()
98 + self.kv_output_aggregator: KVOutputAggregator | None = None
99 +
100 + @abstractmethod
101 + def _init_executor(self) -> None:
102 + raise NotImplementedError
103 +
104 def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
105 """
106 Initialize the KV caches and begin the model execution loop of the
@@ -77,7 +109,7 @@ class Executor(ExecutorBase):
109 self.collective_rpc("initialize_from_config", args=(kv_cache_configs,))
110 self.collective_rpc("compile_or_warm_up_model")
111
80 - def register_failure_callback(self, callback: FailureCallback):
112 + def register_failure_callback(self, callback: FailureCallback): # noqa: B027
113 """
114 Register a function to be called if the executor enters a permanent
115 failed state.
@@ -90,22 +122,78 @@ class Executor(ExecutorBase):
122 def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
123 return self.collective_rpc("get_kv_cache_spec")
124
125 + @overload
126 + def collective_rpc(
127 + self,
128 + method: str | Callable[[WorkerBase], _R],
129 + timeout: float | None = None,
130 + args: tuple = (),
131 + kwargs: dict | None = None,
132 + non_block: Literal[False] = False,
133 + ) -> list[_R]:
134 + """
135 + Execute an RPC call on all workers.
136 +
137 + Args:
138 + method: Name of the worker method to execute, or a callable that
139 + is serialized and sent to all workers to execute.
140 +
141 + If the method is a callable, it should accept an additional
142 + `self` argument, in addition to the arguments passed in `args`
143 + and `kwargs`. The `self` argument will be the worker object.
144 + timeout: Maximum time in seconds to wait for execution. Raises a
145 + [`TimeoutError`][] on timeout. `None` means wait indefinitely.
146 + args: Positional arguments to pass to the worker method.
147 + kwargs: Keyword arguments to pass to the worker method.
148 + non_block: If `True`, returns a list of Futures instead of waiting
149 + for the results.
150 +
151 + Returns:
152 + A list containing the results from each worker.
153 +
154 + Note:
155 + It is recommended to use this API to only pass control messages,
156 + and set up data-plane communication to pass data.
157 + """
158 + pass
159 +
160 + @overload
161 def collective_rpc(
162 self,
95 - method: str | Callable,
163 + method: str | Callable[[WorkerBase], _R],
164 timeout: float | None = None,
165 args: tuple = (),
166 kwargs: dict | None = None,
99 - non_block: bool = False,
100 - ) -> list[Any]:
167 + non_block: Literal[True] = True,
168 + ) -> list[Future[_R]]:
169 + pass
170 +
171 + @abstractmethod
172 + def collective_rpc(
173 + self, method, timeout=None, args=(), kwargs=None, non_block: bool = False
174 + ):
175 raise NotImplementedError
176
177 + @overload
178 def execute_model(
179 self,
180 scheduler_output: SchedulerOutput,
106 - non_block: bool = False,
181 + non_block: Literal[False] = False,
182 + ) -> ModelRunnerOutput:
183 + pass
184 +
185 + @overload
186 + def execute_model(
187 + self,
188 + scheduler_output: SchedulerOutput,
189 + non_block: Literal[True] = True,
190 + ) -> Future[ModelRunnerOutput]:
191 + pass
192 +
193 + def execute_model(
194 + self, scheduler_output: SchedulerOutput, non_block: bool = False
195 ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
108 - output = self.collective_rpc(
196 + output = self.collective_rpc( # type: ignore[call-overload]
197 "execute_model", args=(scheduler_output,), non_block=non_block
198 )
199 return output[0]
@@ -114,7 +202,7 @@ class Executor(ExecutorBase):
202 self.collective_rpc("execute_dummy_batch")
203
204 def take_draft_token_ids(self) -> DraftTokenIds | None:
117 - output = self.collective_rpc("take_draft_token_ids")
205 + output: list[DraftTokenIds] = self.collective_rpc("take_draft_token_ids")
206 return output[0]
207
208 @property
@@ -124,19 +212,120 @@ class Executor(ExecutorBase):
212 def profile(self, is_start: bool = True):
213 self.collective_rpc("profile", args=(is_start,))
214
215 + def save_sharded_state(
216 + self,
217 + path: str,
218 + pattern: str | None = None,
219 + max_size: int | None = None,
220 + ) -> None:
221 + self.collective_rpc(
222 + "save_sharded_state",
223 + kwargs=dict(path=path, pattern=pattern, max_size=max_size),
224 + )
225 +
226 + @abstractmethod
227 + def check_health(self) -> None:
228 + """Checks if the executor is healthy. If not, it should raise an
229 + exception."""
230 + raise NotImplementedError
231
128 -class UniProcExecutor(UniProcExecutorV0, Executor):
129 - pass
232 + def shutdown(self) -> None:
233 + """Shutdown the executor."""
234 + self.collective_rpc("shutdown")
235
236 + def init_kv_output_aggregator(self, finished_count: int | None) -> None:
237 + """Init KVOutputAggregator"""
238 + self.kv_output_aggregator = KVOutputAggregator(
239 + finished_count or self.parallel_config.world_size
240 + )
241
132 -class ExecutorWithExternalLauncher(ExecutorWithExternalLauncherV0, Executor):
133 - def determine_available_memory(self) -> list[int]: # in bytes
134 - # same as determine_num_available_blocks in v0,
135 - # we need to get the min across all ranks.
136 - memory = super().determine_available_memory()
137 - from vllm.distributed.parallel_state import get_world_group
138 -
139 - cpu_group = get_world_group().cpu_group
140 - memory_tensor = torch.tensor([memory], device="cpu", dtype=torch.int64)
141 - dist.all_reduce(memory_tensor, group=cpu_group, op=dist.ReduceOp.MIN)
142 - return [memory_tensor.item()]
242 + @cached_property # Avoid unnecessary RPC calls
243 + def supported_tasks(self) -> tuple[SupportedTask, ...]:
244 + output: list[tuple[SupportedTask, ...]]
245 + output = self.collective_rpc("get_supported_tasks")
246 + return output[0]
247 +
248 + def add_lora(self, lora_request: LoRARequest) -> bool:
249 + assert lora_request.lora_int_id > 0, "lora_id must be greater than 0."
250 + return all(self.collective_rpc("add_lora", args=(lora_request,)))
251 +
252 + def remove_lora(self, lora_id: int) -> bool:
253 + assert lora_id > 0, "lora_id must be greater than 0."
254 + return all(self.collective_rpc("remove_lora", args=(lora_id,)))
255 +
256 + def pin_lora(self, lora_id: int) -> bool:
257 + assert lora_id > 0, "lora_id must be greater than 0."
258 + return all(self.collective_rpc("pin_lora", args=(lora_id,)))
259 +
260 + def list_loras(self) -> set[int]:
261 + sets: list[set[int]] = self.collective_rpc("list_loras")
262 + for s in sets:
263 + assert s == sets[0], "All workers should have the same LORAs."
264 + return sets[0]
265 +
266 + def reset_mm_cache(self) -> None:
267 + """Reset the multi-modal cache in each worker."""
268 + self.collective_rpc("reset_mm_cache")
269 +
270 + def start_profile(self) -> None:
271 + self.collective_rpc("start_profile")
272 +
273 + def stop_profile(self) -> None:
274 + self.collective_rpc("stop_profile")
275 +
276 + def sleep(self, level: int = 1):
277 + if self.is_sleeping:
278 + logger.warning("Executor is already sleeping.")
279 + return
280 + time_before_sleep = time.perf_counter()
281 + self.collective_rpc("sleep", kwargs=dict(level=level))
282 + time_after_sleep = time.perf_counter()
283 + self.sleeping_tags = {"weights", "kv_cache"}
284 + self.is_sleeping = True
285 + logger.info(
286 + "It took %.6f seconds to fall asleep.", time_after_sleep - time_before_sleep
287 + )
288 +
289 + def wake_up(self, tags: list[str] | None = None):
290 + if not self.is_sleeping:
291 + logger.warning("Executor is not sleeping.")
292 + return
293 + if tags:
294 + for tag in tags:
295 + if tag not in self.sleeping_tags:
296 + logger.warning(
297 + "Tag %s is not in sleeping tags %s", tag, self.sleeping_tags
298 + )
299 + return
300 + time_before_wakeup = time.perf_counter()
301 + self.collective_rpc("wake_up", kwargs=dict(tags=tags))
302 + time_after_wakeup = time.perf_counter()
303 + logger.info(
304 + "It took %.6f seconds to wake up tags %s.",
305 + time_after_wakeup - time_before_wakeup,
306 + tags if tags is not None else self.sleeping_tags,
307 + )
308 + if tags:
309 + for tag in tags:
310 + self.sleeping_tags.remove(tag)
311 + else:
312 + self.sleeping_tags.clear()
313 + if not self.sleeping_tags:
314 + self.is_sleeping = False
315 +
316 + def reinitialize_distributed(
317 + self, reconfig_request: ReconfigureDistributedRequest
318 + ) -> None:
319 + raise NotImplementedError
320 +
321 +
322 +from vllm.v1.executor.uniproc_executor import ( # noqa: E402
323 + ExecutorWithExternalLauncher as _ExecutorWithExternalLauncher,
324 +)
325 +from vllm.v1.executor.uniproc_executor import ( # noqa: E402
326 + UniProcExecutor as _UniProcExecutor,
327 +)
328 +
329 +# For backwards compatibility.
330 +UniProcExecutor = _UniProcExecutor
331 +ExecutorWithExternalLauncher = _ExecutorWithExternalLauncher
vllm/v1/executor/multiproc_executor.py
+2 -1
@@ -179,7 +179,7 @@ class MultiprocExecutor(Executor):
179 else:
180 self.failure_callback = callback
181
182 - def execute_model(
182 + def execute_model( # type: ignore[override]
183 self,
184 scheduler_output: SchedulerOutput,
185 non_block: bool = False,
@@ -204,6 +204,7 @@ class MultiprocExecutor(Executor):
204 )
205
206 # aggregate all workers output to a single output
207 + assert self.kv_output_aggregator is not None
208 if non_block:
209 return self.kv_output_aggregator.async_aggregate(outputs, self.output_rank)
210 return self.kv_output_aggregator.aggregate(outputs, self.output_rank)
vllm/v1/executor/ray_distributed_executor.py
+4 -107
@@ -1,111 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from concurrent.futures import Future
5 -
6 -from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
7 -from vllm.executor.ray_distributed_executor import ( # noqa
8 - RayDistributedExecutor as RayDistributedExecutorV0,
3 +from vllm.v1.executor.ray_executor import (
4 + RayDistributedExecutor as _RayDistributedExecutor,
5 )
10 -from vllm.logger import init_logger
11 -from vllm.v1.core.sched.output import SchedulerOutput
12 -from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
13 -from vllm.v1.executor.abstract import Executor
14 -from vllm.v1.outputs import ModelRunnerOutput
15 -
16 -logger = init_logger(__name__)
17 -
18 -
19 -class FutureWrapper(Future):
20 - """A wrapper around Ray output reference to meet the interface
21 - of .execute_model(): The top level (core busy loop) expects .result() api
22 - to block and return a single output.
23 -
24 - If aggregator is provided, the outputs from all workers are aggregated upon
25 - the result() call. If not only the first worker's output is returned.
26 - """
27 -
28 - def __init__(self, refs, aggregator: KVOutputAggregator | None = None):
29 - super().__init__()
30 - self.refs = refs
31 - self.aggregator = aggregator
32 -
33 - def result(self, timeout=None):
34 - if timeout is not None:
35 - raise NotImplementedError("timeout is not supported")
36 -
37 - if self.aggregator is None:
38 - return self.refs[0].get()
39 -
40 - outputs = [ref.get() for ref in self.refs]
41 - return self.aggregator.aggregate(outputs, output_rank=0)
42 -
43 -
44 -class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
45 - """Ray distributed executor using Ray Compiled Graphs."""
46 -
47 - supports_pp: bool = True
48 -
49 - def _init_executor(self) -> None:
50 - super()._init_executor()
51 -
52 - # KV connector setup
53 - self.has_connector = self.vllm_config.kv_transfer_config is not None
54 -
55 - @property
56 - def max_concurrent_batches(self) -> int:
57 - """Ray distributed executor supports pipeline parallelism,
58 - meaning that it allows PP size batches to be executed concurrently.
59 - """
60 - if self.scheduler_config.async_scheduling:
61 - return 2
62 - return self.parallel_config.pipeline_parallel_size
63 -
64 - def execute_model(
65 - self,
66 - scheduler_output: SchedulerOutput,
67 - non_block: bool = False,
68 - ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
69 - """Execute the model on the Ray workers.
70 -
71 - Args:
72 - scheduler_output: The scheduler output to execute.
73 - non_block: If True, the method will return a Future.
74 -
75 - Returns:
76 - The model runner output.
77 - """
78 - # Build the compiled DAG for the first time.
79 - if self.forward_dag is None: # type: ignore
80 - self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)
81 -
82 - refs = self.forward_dag.execute(scheduler_output) # type: ignore
83 -
84 - if not self.has_connector:
85 - # Get output only from a single worker (output_rank)
86 - # When PP is not used, we block here until the result is available.
87 - if not non_block:
88 - return refs[0].get()
89 -
90 - # When PP is used, we return a FutureWrapper immediately so that
91 - # the scheduler can yield to the next batch.
92 - return FutureWrapper(refs)
93 -
94 - # Get output from all workers when connector is present
95 - if not non_block:
96 - # Block and get results from all workers
97 - outputs = [ref.get() for ref in refs]
98 - return self.kv_output_aggregator.aggregate(outputs)
99 -
100 - # Return a future that will aggregate outputs from all workers
101 - return FutureWrapper(refs, self.kv_output_aggregator)
6
103 - def reinitialize_distributed(
104 - self, reconfig_request: ReconfigureDistributedRequest
105 - ) -> None:
106 - self._run_workers("reinitialize_distributed", reconfig_request)
107 - if (
108 - reconfig_request.new_data_parallel_rank
109 - == ReconfigureRankType.SHUTDOWN_CURRENT_RANK
110 - ):
111 - self.shutdown()
7 +# For backwards compatibility.
8 +RayDistributedExecutor = _RayDistributedExecutor
vllm/v1/executor/ray_executor.py renamed
+104 -280
@@ -1,31 +1,34 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import asyncio
4 import os
5 from collections import defaultdict
6 from collections.abc import Callable
7 +from concurrent.futures import Future
8 from dataclasses import dataclass
9 from typing import TYPE_CHECKING, Any
10
11 import cloudpickle
12 -import msgspec
12
13 import vllm.envs as envs
15 -from vllm.executor.executor_base import DistributedExecutorBase
16 -from vllm.executor.msgspec_utils import encode_hook
17 -from vllm.executor.ray_utils import RayWorkerWrapper, initialize_ray_cluster, ray
14 from vllm.logger import init_logger
15 from vllm.platforms import current_platform
16 from vllm.ray.ray_env import get_env_vars_to_copy
21 -from vllm.sequence import ExecuteModelRequest
22 -from vllm.utils.async_utils import make_async
17 from vllm.utils.network_utils import (
18 get_distributed_init_method,
19 get_ip,
20 get_open_port,
21 )
28 -from vllm.v1.outputs import SamplerOutput
22 +from vllm.v1.core.sched.output import SchedulerOutput
23 +from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
24 +from vllm.v1.executor.abstract import Executor
25 +from vllm.v1.executor.ray_utils import (
26 + FutureWrapper,
27 + RayWorkerWrapper,
28 + initialize_ray_cluster,
29 + ray,
30 +)
31 +from vllm.v1.outputs import ModelRunnerOutput
32
33 if ray is not None:
34 from ray.actor import ActorHandle
@@ -53,7 +56,7 @@ class RayWorkerMetaData:
56 ip: str = ""
57
58
56 -class RayDistributedExecutor(DistributedExecutorBase):
59 +class RayDistributedExecutor(Executor):
60 """Ray-based distributed executor"""
61
62 # These env vars are worker-specific, therefore are NOT copied
@@ -69,37 +72,14 @@ class RayDistributedExecutor(DistributedExecutorBase):
72 ADDITIONAL_ENV_VARS = {"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"}
73
74 uses_ray: bool = True
75 + supports_pp: bool = True
76
77 def _init_executor(self) -> None:
78 self.forward_dag: ray.dag.CompiledDAG | None = None
75 - if envs.VLLM_USE_V1:
76 - # V1 uses SPMD worker and compiled DAG
77 - os.environ["VLLM_USE_RAY_SPMD_WORKER"] = "1"
78 - os.environ["VLLM_USE_RAY_COMPILED_DAG"] = "1"
79 -
80 - # For TPU or XPU, avoid compiling NVIDIA's NCCL
81 - if current_platform.is_tpu() or current_platform.is_xpu():
82 - os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = "shm"
83 -
84 - # If the env var is set, it uses the Ray's compiled DAG API
85 - # which optimizes the control plane overhead.
86 - # Run vLLM with VLLM_USE_RAY_COMPILED_DAG=1 to enable it.
87 - # Currently, this requires USE_RAY_SPMD_WORKER=True.
88 - self.use_ray_compiled_dag = envs.VLLM_USE_RAY_COMPILED_DAG
89 - # If the env var is set, then we do not distinguish between the
90 - # "driver worker" vs other workers. Also, the rank 0 worker will
91 - # be executed in a remote Ray worker. Currently this requires
92 - # USE_RAY_COMPILED_DAG=True.
93 - self.use_ray_spmd_worker = envs.VLLM_USE_RAY_SPMD_WORKER
94 - if self.use_ray_compiled_dag:
95 - assert self.use_ray_spmd_worker, (
96 - "VLLM_USE_RAY_COMPILED_DAG=1 requires VLLM_USE_RAY_SPMD_WORKER=1"
97 - )
98 - if self.use_ray_spmd_worker:
99 - # TODO: Support SPMD worker for non-DAG Ray executor.
100 - assert self.use_ray_compiled_dag, (
101 - "VLLM_USE_RAY_SPMD_WORKER=1 requires VLLM_USE_RAY_COMPILED_DAG=1"
102 - )
79 +
80 + # For TPU or XPU, avoid compiling NVIDIA's NCCL
81 + if current_platform.is_tpu() or current_platform.is_xpu():
82 + os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = "shm"
83
84 assert self.uses_ray
85 initialize_ray_cluster(self.parallel_config)
@@ -113,13 +93,17 @@ class RayDistributedExecutor(DistributedExecutorBase):
93 # Create the parallel GPU workers.
94 self._init_workers_ray(placement_group)
95
116 - self.input_encoder = msgspec.msgpack.Encoder(enc_hook=encode_hook)
117 - self.output_decoder = msgspec.msgpack.Decoder(list[SamplerOutput] | None)
118 - self.use_v1 = envs.VLLM_USE_V1
96 + # KV connector setup
97 + self.has_connector = self.vllm_config.kv_transfer_config is not None
98
120 - self.pp_locks: list[asyncio.Lock] | None = None
121 - if not self.use_ray_compiled_dag:
122 - self.driver_exec_method = make_async(self.driver_worker.execute_method)
99 + @property
100 + def max_concurrent_batches(self) -> int:
101 + """Ray distributed executor supports pipeline parallelism,
102 + meaning that it allows PP size batches to be executed concurrently.
103 + """
104 + if self.scheduler_config.async_scheduling:
105 + return 2
106 + return self.parallel_config.pipeline_parallel_size
107
108 def shutdown(self) -> None:
109 if logger:
@@ -176,8 +160,6 @@ class RayDistributedExecutor(DistributedExecutorBase):
160 ray_remote_kwargs
161 )
162
179 - logger.info("use_ray_spmd_worker: %s", self.use_ray_spmd_worker)
180 -
163 # Create the workers.
164 bundle_indices: list[int]
165 if envs.VLLM_RAY_BUNDLE_INDICES:
@@ -241,30 +223,8 @@ class RayDistributedExecutor(DistributedExecutorBase):
223 for each, ip in zip(worker_metadata, worker_ips):
224 each.ip = ip
225
244 - if not self.use_ray_spmd_worker:
245 - for i, each in enumerate(worker_metadata):
246 - # find and remove the dummy worker from the list
247 - worker = each.worker
248 - worker_ip = each.ip
249 - if self.driver_dummy_worker is None and worker_ip == driver_ip:
250 - # If the worker is on the same node as the driver, we use it
251 - # as the resource holder for the driver process.
252 - self.driver_dummy_worker = worker
253 - self.driver_worker = RayWorkerWrapper(
254 - vllm_config=self.vllm_config, rpc_rank=0
255 - )
256 - worker_metadata.pop(i)
257 - break
258 -
226 logger.debug("workers: %s", worker_metadata)
227 logger.debug("driver_dummy_worker: %s", self.driver_dummy_worker)
261 - if not self.use_ray_spmd_worker and self.driver_dummy_worker is None:
262 - raise ValueError(
263 - "Ray does not allocate any GPUs on the driver node."
264 - f"Driver IP: {driver_ip}, worker IPs: {worker_ips}."
265 - "Consider adjusting the Ray placement group or running "
266 - "the driver on a GPU node."
267 - )
228
229 ip_counts: dict[str, int] = {}
230 for ip in worker_ips:
@@ -281,7 +241,7 @@ class RayDistributedExecutor(DistributedExecutorBase):
241 should be placed first.
242 """
243 ip = item.ip
284 - return (0 if ip == driver_ip else 1, ip_counts[ip], ip)
244 + return 0 if ip == driver_ip else 1, ip_counts[ip], ip
245
246 # After sorting, the workers on the same node will be
247 # close to each other, and the workers on the driver
@@ -289,14 +249,13 @@ class RayDistributedExecutor(DistributedExecutorBase):
249 sorted_worker_metadata = sorted(
250 worker_metadata, key=sort_by_driver_then_worker_ip
251 )
292 - start_rank = 0 if self.use_ray_spmd_worker else 1
252 for i, item in enumerate(sorted_worker_metadata):
294 - item.adjusted_rank = i + start_rank
253 + item.adjusted_rank = i
254 self.workers = [item.worker for item in sorted_worker_metadata]
255 rerank_mapping = {
256 item.created_rank: item.adjusted_rank for item in sorted_worker_metadata
257 }
299 - self._run_workers("adjust_rank", rerank_mapping)
258 + self.collective_rpc("adjust_rank", args=(rerank_mapping,))
259
260 # Get the set of GPU IDs used on each node.
261 worker_node_and_gpu_ids = []
@@ -365,8 +324,8 @@ class RayDistributedExecutor(DistributedExecutorBase):
324
325 self._env_vars_for_all_workers = all_args_to_update_environment_variables
326
368 - self._run_workers(
369 - "update_environment_variables", self._get_env_vars_to_be_updated()
327 + self.collective_rpc(
328 + "update_environment_variables", args=(self._get_env_vars_to_be_updated(),)
329 )
330
331 if len(node_gpus) == 1:
@@ -396,138 +355,95 @@ class RayDistributedExecutor(DistributedExecutorBase):
355 or (rank % self.parallel_config.tensor_parallel_size == 0),
356 )
357 all_kwargs.append(kwargs)
399 - self._run_workers("init_worker", all_kwargs)
358 + self.collective_rpc("init_worker", args=(all_kwargs,))
359 +
360 + self.collective_rpc("init_device")
361 + self.collective_rpc("load_model")
362 +
363 + for pp_rank in range(self.parallel_config.pipeline_parallel_size):
364 + self.pp_tp_workers.append([])
365 + for tp_rank in range(self.parallel_config.tensor_parallel_size):
366 + # PP=2, TP=4
367 + # pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]]
368 + rank = (pp_rank * self.parallel_config.tensor_parallel_size) + tp_rank
369 + assert len(self.pp_tp_workers[pp_rank]) == tp_rank
370 + assert pp_rank < len(self.pp_tp_workers)
371 + self.pp_tp_workers[pp_rank].append(self.workers[rank])
372 +
373 + def reinitialize_distributed(
374 + self, reconfig_request: ReconfigureDistributedRequest
375 + ) -> None:
376 + self.collective_rpc("reinitialize_distributed", args=(reconfig_request,))
377 + if (
378 + reconfig_request.new_data_parallel_rank
379 + == ReconfigureRankType.SHUTDOWN_CURRENT_RANK
380 + ):
381 + self.shutdown()
382 +
383 + def execute_model( # type: ignore[override]
384 + self, scheduler_output: SchedulerOutput, non_block: bool = False
385 + ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
386 + """Execute the model on the Ray workers.
387
401 - self._run_workers("init_device")
402 - self._run_workers(
403 - "load_model",
404 - max_concurrent_workers=self.parallel_config.max_parallel_loading_workers,
405 - )
388 + Args:
389 + scheduler_output: The scheduler output to execute.
390 + non_block: If True, the method will return a Future.
391
407 - if self.use_ray_spmd_worker:
408 - for pp_rank in range(self.parallel_config.pipeline_parallel_size):
409 - self.pp_tp_workers.append([])
410 - for tp_rank in range(self.parallel_config.tensor_parallel_size):
411 - # PP=2, TP=4
412 - # pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]]
413 - rank = (
414 - pp_rank * self.parallel_config.tensor_parallel_size
415 - ) + tp_rank
416 - assert len(self.pp_tp_workers[pp_rank]) == tp_rank
417 - assert pp_rank < len(self.pp_tp_workers)
418 - self.pp_tp_workers[pp_rank].append(self.workers[rank])
419 -
420 - # This is the list of workers that are rank 0 of each TP group EXCEPT
421 - # global rank 0. These are the workers that will broadcast to the
422 - # rest of the workers.
423 - self.tp_driver_workers: list[RayWorkerWrapper] = []
424 - # This is the list of workers that are not drivers and not the first
425 - # worker in a TP group. These are the workers that will be
426 - # broadcasted to.
427 - self.non_driver_workers: list[RayWorkerWrapper] = []
428 -
429 - # Enforce rank order for correct rank to return final output.
430 - for index, worker in enumerate(self.workers):
431 - # The driver worker is rank 0 and not in self.workers.
432 - rank = index + 1
433 - if rank % self.parallel_config.tensor_parallel_size == 0:
434 - self.tp_driver_workers.append(worker)
435 - else:
436 - self.non_driver_workers.append(worker)
392 + Returns:
393 + The model runner output.
394 + """
395 + # Build the compiled DAG for the first time.
396 + if self.forward_dag is None: # type: ignore
397 + self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)
398
438 - def _driver_execute_model(
439 - self, execute_model_req: ExecuteModelRequest | None
440 - ) -> list[SamplerOutput] | None:
441 - """Run execute_model in the driver worker.
399 + refs = self.forward_dag.execute(scheduler_output) # type: ignore
400
443 - Passing None will cause the driver to stop the model execution
444 - loop running in each of the remote workers.
445 - """
446 - assert not self.use_ray_spmd_worker, (
447 - "driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1"
448 - )
449 - return self.driver_worker.execute_method("execute_model", execute_model_req)
401 + if not self.has_connector:
402 + # Get output only from a single worker (output_rank)
403 + # When PP is not used, we block here until the result is available.
404 + if not non_block:
405 + return refs[0].get()
406
451 - def execute_model(
452 - self, execute_model_req: ExecuteModelRequest
453 - ) -> list[SamplerOutput]:
454 - if not self.use_ray_spmd_worker:
455 - return super().execute_model(execute_model_req)
407 + # When PP is used, we return a FutureWrapper immediately so that
408 + # the scheduler can yield to the next batch.
409 + return FutureWrapper(refs)
410
457 - if self.forward_dag is None:
458 - self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)
411 + # Get output from all workers when connector is present
412 + assert self.kv_output_aggregator is not None
413 + if not non_block:
414 + # Block and get results from all workers
415 + outputs = [ref.get() for ref in refs]
416 + return self.kv_output_aggregator.aggregate(outputs)
417
460 - if self.use_v1:
461 - serialized_data = execute_model_req
462 - else:
463 - serialized_data = self.input_encoder.encode(execute_model_req)
464 - outputs = ray.get(self.forward_dag.execute(serialized_data))
465 - output = outputs[0] if self.use_v1 else self.output_decoder.decode(outputs[0])
466 - return output
418 + # Return a future that will aggregate outputs from all workers
419 + return FutureWrapper(refs, self.kv_output_aggregator)
420
468 - def _run_workers(
421 + def collective_rpc(
422 self,
423 method: str | Callable,
471 - *args,
472 - async_run_tensor_parallel_workers_only: bool = False,
473 - max_concurrent_workers: int | None = None,
474 - **kwargs,
475 - ) -> Any:
476 - """Runs the given method on all workers. Can be used in the following
477 - ways:
478 -
479 - Args:
480 - - async_run_tensor_parallel_workers_only: If True the method will be
481 - run only in the remote TP workers, not the driver worker.
482 - It will also be run asynchronously and return a list of futures
483 - rather than blocking on the results.
484 - - args/kwargs: All workers share the same args/kwargs
485 - """
424 + timeout: float | None = None,
425 + args: tuple = (),
426 + kwargs: dict[str, Any] | None = None,
427 + non_block: bool = False,
428 + ) -> list[Any]:
429 + """Runs the given method on all workers."""
430 sent_method = method if isinstance(method, str) else cloudpickle.dumps(method)
431 del method
488 - if self.use_ray_spmd_worker:
489 - assert not async_run_tensor_parallel_workers_only, (
490 - "async_run_tensor_parallel_workers_only is not supported for spmd mode."
491 - )
492 -
493 - if max_concurrent_workers:
494 - raise NotImplementedError("max_concurrent_workers is not supported yet.")
432
496 - # Start the ray workers first.
497 - ray_workers = self.workers
498 - if async_run_tensor_parallel_workers_only:
499 - ray_workers = self.non_driver_workers
433 + if kwargs is None:
434 + kwargs = {}
435 ray_worker_outputs = [
436 worker.execute_method.remote( # type: ignore[attr-defined]
437 sent_method, *args, **kwargs
438 )
504 - for worker in ray_workers
439 + for worker in self.workers
440 ]
441
507 - if async_run_tensor_parallel_workers_only:
508 - # Just return futures
509 - return ray_worker_outputs
510 -
511 - driver_worker_output = []
512 - # In SPMD mode, the driver worker is the same as any other worker,
513 - # so we only explicitly execute on the driver worker if using a
514 - # non-SPMD worker class.
515 - if not self.use_ray_spmd_worker:
516 - # Start the driver worker after all the ray workers.
517 - driver_worker_output = [
518 - self.driver_worker.execute_method(sent_method, *args, **kwargs)
519 - ]
520 -
442 # Get the results of the ray workers.
522 - if self.workers:
523 - ray_worker_outputs = ray.get(ray_worker_outputs)
524 -
525 - return driver_worker_output + ray_worker_outputs
443 + if non_block:
444 + return [FutureWrapper((output,)) for output in ray_worker_outputs]
445
527 - def _wait_for_tasks_completion(self, parallel_worker_tasks: Any) -> None:
528 - """Wait for futures returned from _run_workers() with
529 - async_run_remote_workers_only to complete."""
530 - ray.get(parallel_worker_tasks)
446 + return ray.get(ray_worker_outputs, timeout=timeout)
447
448 def _check_ray_cgraph_installation(self):
449 import importlib.metadata
@@ -595,13 +511,6 @@ class RayDistributedExecutor(DistributedExecutorBase):
511 with InputNode() as input_data:
512 # Example DAG: PP=2, TP=4
513 #
598 - # For V0:
599 - # ExecuteModelRequest -> 0 -> (ExecuteModelReq, IntermediateTensors) -> 4 -> SamplerOutput # noqa: E501
600 - # ExecuteModelRequest -> 1 -> (ExecuteModelReq, IntermediateTensors) -> 5 -> SamplerOutput # noqa: E501
601 - # ExecuteModelRequest -> 2 -> (ExecuteModelReq, IntermediateTensors) -> 6 -> SamplerOutput # noqa: E501
602 - # ExecuteModelRequest -> 3 -> (ExecuteModelReq, IntermediateTensors) -> 7 -> SamplerOutput # noqa: E501
603 - #
604 - # For V1:
514 # SchedulerOutput -> 0 -> (SchedulerOutput, IntermediateTensors) -> 4 -> ModelRunnerOutput # noqa: E501
515 # SchedulerOutput -> 1 -> (SchedulerOutput, IntermediateTensors) -> 5 -> ModelRunnerOutput # noqa: E501
516 # SchedulerOutput -> 2 -> (SchedulerOutput, IntermediateTensors) -> 6 -> ModelRunnerOutput # noqa: E501
@@ -613,20 +522,10 @@ class RayDistributedExecutor(DistributedExecutorBase):
522 for pp_rank, tp_group in enumerate(self.pp_tp_workers):
523 # Each PP worker takes in the output of the previous PP worker,
524 # and the TP group executes in SPMD fashion.
616 - if self.use_v1:
617 - outputs = [
618 - worker.execute_model_ray.bind( # type: ignore[attr-defined]
619 - outputs[i]
620 - )
621 - for i, worker in enumerate(tp_group)
622 - ]
623 - else:
624 - outputs = [
625 - worker.execute_model_spmd.bind( # type: ignore[attr-defined]
626 - outputs[i]
627 - )
628 - for i, worker in enumerate(tp_group)
629 - ]
525 + outputs = [
526 + worker.execute_model_ray.bind(outputs[i]) # type: ignore[attr-defined]
527 + for i, worker in enumerate(tp_group)
528 + ]
529
530 last_pp_rank = len(self.pp_tp_workers) - 1
531 if (
@@ -674,82 +573,7 @@ class RayDistributedExecutor(DistributedExecutorBase):
573 def __del__(self):
574 self.shutdown()
575
677 - async def execute_model_async(
678 - self, execute_model_req: ExecuteModelRequest
679 - ) -> list[SamplerOutput]:
680 - if not self.use_ray_spmd_worker:
681 - return await super().execute_model_async(execute_model_req)
682 -
683 - if self.forward_dag is None:
684 - self.forward_dag = self._compiled_ray_dag(enable_asyncio=True)
685 -
686 - serialized_data = self.input_encoder.encode(execute_model_req)
687 - dag_future = await self.forward_dag.execute_async(serialized_data)
688 - output = await dag_future[0]
689 - return self.output_decoder.decode(output)
690 -
691 - async def _driver_execute_model_async(
692 - self, execute_model_req: ExecuteModelRequest | None = None
693 - ) -> list[SamplerOutput]:
694 - assert not self.use_ray_spmd_worker, (
695 - "driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1"
696 - )
697 - if not self.tp_driver_workers:
698 - return await self.driver_exec_method("execute_model", execute_model_req)
699 - if self.pp_locks is None:
700 - # This locks each pipeline parallel stage so multiple virtual
701 - # engines can't execute on the same stage at the same time
702 - # We create the locks here to avoid creating them in the constructor
703 - # which uses a different asyncio loop.
704 - self.pp_locks = [
705 - asyncio.Lock()
706 - for _ in range(self.parallel_config.pipeline_parallel_size)
707 - ]
708 -
709 - tasks = [
710 - asyncio.create_task(
711 - _run_task_with_lock(
712 - self.driver_exec_method,
713 - self.pp_locks[0],
714 - "execute_model",
715 - execute_model_req,
716 - )
717 - )
718 - ]
719 - for pp_rank, driver_worker in enumerate(self.tp_driver_workers, start=1):
720 - tasks.append(
721 - asyncio.create_task(
722 - _run_task_with_lock(
723 - driver_worker.execute_method.remote, # type: ignore[attr-defined]
724 - self.pp_locks[pp_rank],
725 - "execute_model",
726 - execute_model_req,
727 - )
728 - )
729 - )
730 -
731 - results = await asyncio.gather(*tasks)
732 -
733 - # Only the last PP stage has the final results.
734 - return results[-1]
735 -
736 - async def _start_worker_execution_loop(self):
737 - assert not self.use_ray_spmd_worker, (
738 - "worker loop is disabled for VLLM_USE_RAY_SPMD_WORKER=1"
739 - )
740 - coros = [
741 - worker.execute_method.remote("start_worker_execution_loop") # type: ignore[attr-defined]
742 - for worker in self.non_driver_workers
743 - ]
744 - return await asyncio.gather(*coros)
745 -
576 def check_health(self) -> None:
577 # Assume that the Ray workers are healthy.
578 # TODO: check the health of the Ray workers
579 return
750 -
751 -
752 -async def _run_task_with_lock(task: Callable, lock: asyncio.Lock, *args, **kwargs):
753 - """Utility function to run async task in a lock"""
754 - async with lock:
755 - return await task(*args, **kwargs)
vllm/v1/executor/ray_utils.py renamed
+28 -50
@@ -4,17 +4,16 @@
4 import os
5 import time
6 from collections import defaultdict
7 +from concurrent.futures import Future
8 from typing import TYPE_CHECKING, Union
9
9 -import msgspec
10 -
10 import vllm.platforms
11 from vllm.config import ParallelConfig
12 from vllm.distributed import get_pp_group
14 -from vllm.executor.msgspec_utils import decode_hook, encode_hook
13 +from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
14 from vllm.logger import init_logger
15 from vllm.platforms import current_platform
17 -from vllm.sequence import ExecuteModelRequest, IntermediateTensors
16 +from vllm.sequence import IntermediateTensors
17 from vllm.utils.network_utils import get_ip
18 from vllm.v1.outputs import AsyncModelRunnerOutput
19 from vllm.v1.worker.worker_base import WorkerWrapperBase
@@ -51,11 +50,6 @@ try:
50 # that thread.
51 self.compiled_dag_cuda_device_set = False
52
54 - self.input_decoder = msgspec.msgpack.Decoder(
55 - ExecuteModelRequest, dec_hook=decode_hook
56 - )
57 - self.output_encoder = msgspec.msgpack.Encoder(enc_hook=encode_hook)
58 -
53 def get_node_ip(self) -> str:
54 return get_ip()
55
@@ -70,47 +64,6 @@ try:
64 gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]
65 return node_id, gpu_ids
66
73 - def execute_model_spmd(
74 - self,
75 - req_or_tuple: bytes | tuple[bytes, IntermediateTensors | None],
76 - ) -> bytes:
77 - """Execute model in SPMD fashion: used only when SPMD worker and
78 - compiled DAG are both enabled.
79 -
80 - Args:
81 - req_or_tuple: A request or a tuple containing the
82 - request and intermediate tensors. Intermediate tensors are
83 - None unless if it is provided because it is > 0 pipeline
84 - stage. The request is serialized by msgspec.
85 - """
86 - if isinstance(req_or_tuple, bytes):
87 - serialized_req, intermediate_tensors = req_or_tuple, None
88 - else:
89 - serialized_req, intermediate_tensors = req_or_tuple
90 -
91 - execute_model_req = self.input_decoder.decode(serialized_req)
92 -
93 - assert self.worker is not None, "Worker is not initialized"
94 -
95 - # TODO(swang): This is needed right now because Ray Compiled Graph
96 - # executes on a background thread, so we need to reset torch's
97 - # current device.
98 - if not self.compiled_dag_cuda_device_set:
99 - assert self.worker.device is not None
100 - current_platform.set_device(self.worker.device)
101 - self.compiled_dag_cuda_device_set = True
102 -
103 - output = self.worker._execute_model_spmd( # type: ignore[attr-defined]
104 - execute_model_req, intermediate_tensors
105 - )
106 - # Pipeline model request and output to the next pipeline stage.
107 - if isinstance(output, IntermediateTensors):
108 - output = serialized_req, output
109 - else:
110 - output = self.output_encoder.encode(output)
111 -
112 - return output
113 -
67 def setup_device_if_necessary(self):
68 # TODO(swang): This is needed right now because Ray CG executes
69 # on a background thread, so we need to reset torch's current
@@ -174,6 +127,31 @@ except ImportError as e:
127 RayWorkerWrapper = None # type: ignore
128
129
130 +class FutureWrapper(Future):
131 + """A wrapper around Ray output reference to meet the interface
132 + of .execute_model(): The top level (core busy loop) expects .result() api
133 + to block and return a single output.
134 +
135 + If aggregator is provided, the outputs from all workers are aggregated upon
136 + the result() call. If not only the first worker's output is returned.
137 + """
138 +
139 + def __init__(self, refs, aggregator: KVOutputAggregator | None = None):
140 + super().__init__()
141 + self.refs = refs
142 + self.aggregator = aggregator
143 +
144 + def result(self, timeout=None):
145 + if timeout is not None:
146 + raise NotImplementedError("timeout is not supported")
147 +
148 + if self.aggregator is None:
149 + return self.refs[0].get()
150 +
151 + outputs = [ref.get() for ref in self.refs]
152 + return self.aggregator.aggregate(outputs, output_rank=0)
153 +
154 +
155 def ray_is_available() -> bool:
156 """Returns True if Ray is available."""
157 return ray is not None
vllm/v1/executor/uniproc_executor.py renamed
+11 -29
@@ -11,20 +11,18 @@ import torch
11 import torch.distributed as dist
12
13 import vllm.envs as envs
14 -from vllm.executor.executor_base import ExecutorBase
14 from vllm.logger import init_logger
15 from vllm.utils import run_method
16 from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
17 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
18 +from vllm.v1.executor.abstract import Executor
19 from vllm.v1.outputs import AsyncModelRunnerOutput
20 from vllm.v1.worker.worker_base import WorkerWrapperBase
21
22 logger = init_logger(__name__)
23
24
25 -class UniProcExecutor(ExecutorBase):
26 - uses_ray: bool = False
27 -
25 +class UniProcExecutor(Executor):
26 def _init_executor(self) -> None:
27 """Initialize the worker and load the model."""
28 self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config, rpc_rank=0)
@@ -44,9 +42,9 @@ class UniProcExecutor(ExecutorBase):
42 max_workers=1, thread_name_prefix="WorkerAsyncOutput"
43 )
44
47 - self.collective_rpc("init_worker", args=([kwargs],))
48 - self.collective_rpc("init_device")
49 - self.collective_rpc("load_model")
45 + self.driver_worker.init_worker(all_kwargs=[kwargs])
46 + self.driver_worker.init_device()
47 + self.driver_worker.load_model()
48
49 def _distributed_args(self) -> tuple[str, int, int]:
50 """Return (distributed_init_method, rank, local_rank)."""
@@ -101,16 +99,12 @@ class UniProcExecutor(ExecutorBase):
99 == ReconfigureRankType.SHUTDOWN_CURRENT_RANK
100 ):
101 self.shutdown()
104 - return
102
103 def shutdown(self) -> None:
104 if worker := self.driver_worker:
105 worker.shutdown()
106
107
111 -UniProcExecutorAsync = UniProcExecutor
112 -
113 -
108 class ExecutorWithExternalLauncher(UniProcExecutor):
109 """An executor that uses external launchers to launch engines,
110 specially designed for torchrun-compatible launchers, for
@@ -128,8 +122,6 @@ class ExecutorWithExternalLauncher(UniProcExecutor):
122 and they don't need to synchronize the states with each other.
123 """
124
131 - uses_ray: bool = False
132 -
125 def _init_executor(self) -> None:
126 """Initialize the worker and load the model."""
127 if envs.VLLM_USE_V1:
@@ -152,22 +144,12 @@ class ExecutorWithExternalLauncher(UniProcExecutor):
144 local_rank = int(os.environ["LOCAL_RANK"])
145 return distributed_init_method, rank, local_rank
146
155 - def determine_num_available_blocks(self) -> tuple[int, int]:
156 - """
157 - Determine the number of available KV blocks.
158 - Add an additional all_reduce to get the min across all ranks.
159 - Note that even if we have the same `gpu_memory_utilization` and
160 - `swap_space`, the available memory in every rank might still
161 - differ because NCCL can take different amounts of memory in
162 - different ranks. Therefore, it is necessary to test if all ranks
163 - agree on the same KV cache configuration.
164 - """
165 - a, b = super().determine_num_available_blocks()
147 + def determine_available_memory(self) -> list[int]: # in bytes
148 + # we need to get the min across all ranks.
149 + memory = super().determine_available_memory()
150 from vllm.distributed.parallel_state import get_world_group
151
152 cpu_group = get_world_group().cpu_group
169 - a_tensor = torch.tensor([a], device="cpu", dtype=torch.int64)
170 - b_tensor = torch.tensor([b], device="cpu", dtype=torch.int64)
171 - dist.all_reduce(a_tensor, group=cpu_group, op=dist.ReduceOp.MIN)
172 - dist.all_reduce(b_tensor, group=cpu_group, op=dist.ReduceOp.MIN)
173 - return a_tensor.item(), b_tensor.item()
153 + memory_tensor = torch.tensor([memory], device="cpu", dtype=torch.int64)
154 + dist.all_reduce(memory_tensor, group=cpu_group, op=dist.ReduceOp.MIN)
155 + return [memory_tensor.item()]
vllm/v1/worker/worker_base.py
-22
@@ -128,28 +128,6 @@ class WorkerBase:
128 def execute_model(self, scheduler_output: SchedulerOutput) -> ModelRunnerOutput:
129 raise NotImplementedError
130
131 - def start_worker_execution_loop(self) -> None:
132 - """Execute model loop in parallel worker.
133 -
134 - You can stop the loop by executing a driver worker with an empty output.
135 - See `stop_remote_worker_execution_loop` for more details.
136 - """
137 - raise NotImplementedError("Dead V0 code")
138 -
139 - def determine_num_available_blocks(self) -> tuple[int, int]:
140 - """Determine the number of available blocks for the GPU KV cache and
141 - swappable CPU KV cache.
142 -
143 - The implementation may run profiling or other heuristics to determine
144 - the size of caches.
145 -
146 - Returns a tuple[num_gpu_blocks, num_cpu_blocks], where num_gpu_blocks
147 - are blocks that are "active" on the device and can be appended to.
148 - num_cpu_blocks refers to "swapped" blocks in CPU memory and cannot be
149 - appended to.
150 - """
151 - raise NotImplementedError
152 -
131 def get_cache_block_size_bytes(self) -> int:
132 """Return the size of a single cache block, in bytes. Used in
133 speculative decoding.