| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import time |
| 4 | from abc import ABC, abstractmethod |
| 5 | from collections.abc import Callable |
| 6 | from concurrent.futures import Future |
| 7 | from functools import cached_property |
| 8 | from typing import TYPE_CHECKING, Literal, TypeVar, overload |
| 9 | |
| 10 | import vllm.envs as envs |
| 11 | from vllm.config import VllmConfig |
| 12 | from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator |
| 13 | from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator |
| 14 | from vllm.distributed.kv_transfer.kv_connector.v1.base import ( |
| 15 | KVConnectorHandshakeMetadata, |
| 16 | ) |
| 17 | from vllm.logger import init_logger |
| 18 | from vllm.lora.request import LoRARequest |
| 19 | from vllm.tasks import SupportedTask |
| 20 | from vllm.tracing import instrument |
| 21 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 22 | from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput |
| 23 | from vllm.v1.engine import ReconfigureDistributedRequest |
| 24 | from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec |
| 25 | from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput |
| 26 | from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase |
| 27 | |
| 28 | if TYPE_CHECKING: |
| 29 | from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase |
| 30 | |
| 31 | logger = init_logger(__name__) |
| 32 | |
| 33 | _R = TypeVar("_R") |
| 34 | |
| 35 | # Memory-pool tags managed by sleep/wake_up/discard. |
| 36 | SLEEP_TAGS = frozenset(("weights", "kv_cache")) |
| 37 | |
| 38 | FailureCallback = Callable[[], None] |
| 39 | |
| 40 | |
| 41 | class Executor(ABC): |
| 42 | """Abstract base class for vLLM executors. |
| 43 | |
| 44 | An executor is responsible for executing the model on one device, |
| 45 | or it can be a distributed executor that can execute the model on multiple devices. |
| 46 | """ |
| 47 | |
| 48 | uses_ray: bool = False # whether the executor uses Ray for orchestration. |
| 49 | supports_pp: bool = False # whether the executor supports PP |
| 50 | |
| 51 | @staticmethod |
| 52 | def get_class(vllm_config: VllmConfig) -> type["Executor"]: |
| 53 | executor_class: type[Executor] |
| 54 | parallel_config = vllm_config.parallel_config |
| 55 | distributed_executor_backend = parallel_config.distributed_executor_backend |
| 56 | # distributed_executor_backend must be set in VllmConfig.__post_init__ |
| 57 | if isinstance(distributed_executor_backend, type): |
| 58 | if not issubclass(distributed_executor_backend, Executor): |
| 59 | raise TypeError( |
| 60 | "distributed_executor_backend must be a subclass of " |
| 61 | f"Executor. Got {distributed_executor_backend}." |
| 62 | ) |
| 63 | executor_class = distributed_executor_backend |
| 64 | elif distributed_executor_backend == "ray": |
| 65 | if envs.VLLM_USE_RAY_V2_EXECUTOR_BACKEND: |
| 66 | from vllm.v1.executor.ray_executor_v2 import RayExecutorV2 |
| 67 | |
| 68 | executor_class = RayExecutorV2 |
| 69 | else: |
| 70 | from vllm.v1.executor.ray_executor import RayDistributedExecutor |
| 71 | |
| 72 | executor_class = RayDistributedExecutor |
| 73 | elif distributed_executor_backend == "mp": |
| 74 | from vllm.v1.executor.multiproc_executor import MultiprocExecutor |
| 75 | |
| 76 | executor_class = MultiprocExecutor |
| 77 | elif distributed_executor_backend == "uni": |
| 78 | from vllm.v1.executor.uniproc_executor import UniProcExecutor |
| 79 | |
| 80 | executor_class = UniProcExecutor |
| 81 | elif distributed_executor_backend == "external_launcher": |
| 82 | # TODO: make v1 scheduling deterministic |
| 83 | # to support external launcher |
| 84 | executor_class = ExecutorWithExternalLauncher |
| 85 | elif isinstance(distributed_executor_backend, str): |
| 86 | executor_class = resolve_obj_by_qualname(distributed_executor_backend) |
| 87 | if not issubclass(executor_class, Executor): |
| 88 | raise TypeError( |
| 89 | "distributed_executor_backend must be a subclass of " |
| 90 | f"Executor. Got {executor_class}." |
| 91 | ) |
| 92 | else: |
| 93 | raise ValueError( |
| 94 | f"Unknown distributed executor backend: {distributed_executor_backend}" |
| 95 | ) |
| 96 | return executor_class |
| 97 | |
| 98 | @instrument(span_name="Executor init") |
| 99 | def __init__( |
| 100 | self, |
| 101 | vllm_config: VllmConfig, |
| 102 | ) -> None: |
| 103 | self.vllm_config = vllm_config |
| 104 | self.model_config = vllm_config.model_config |
| 105 | self.cache_config = vllm_config.cache_config |
| 106 | self.lora_config = vllm_config.lora_config |
| 107 | self.load_config = vllm_config.load_config |
| 108 | self.parallel_config = vllm_config.parallel_config |
| 109 | self.scheduler_config = vllm_config.scheduler_config |
| 110 | self.device_config = vllm_config.device_config |
| 111 | self.speculative_config = vllm_config.speculative_config |
| 112 | self.observability_config = vllm_config.observability_config |
| 113 | self._init_executor() |
| 114 | self.sleeping_tags: set[str] = set() |
| 115 | self.kv_output_aggregator: KVOutputAggregator | None = None |
| 116 | self.ec_output_aggregator: ECOutputAggregator | None = None |
| 117 | |
| 118 | @abstractmethod |
| 119 | def _init_executor(self) -> None: |
| 120 | raise NotImplementedError |
| 121 | |
| 122 | def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None: |
| 123 | """Initialize the KV caches on the underlying workers.""" |
| 124 | self.collective_rpc("initialize_from_config", args=(kv_cache_configs,)) |
| 125 | |
| 126 | def compile_or_warm_up_model(self) -> None: |
| 127 | """Compile/warm up the model and capture cudagraphs on workers.""" |
| 128 | compilation_times: list[CompilationTimes] = self.collective_rpc( |
| 129 | "compile_or_warm_up_model" |
| 130 | ) |
| 131 | # Propagate compilation time from workers back to the main process. |
| 132 | # With TP>1, compilation happens in worker processes, so the main |
| 133 | # process config is never updated. Use max across workers since they |
| 134 | # compile in parallel. |
| 135 | if compilation_times: |
| 136 | self.vllm_config.compilation_config.compilation_time = max( |
| 137 | t.language_model for t in compilation_times |
| 138 | ) |
| 139 | self.vllm_config.compilation_config.encoder_compilation_time = max( |
| 140 | t.encoder for t in compilation_times |
| 141 | ) |
| 142 | |
| 143 | def register_failure_callback(self, callback: FailureCallback): # noqa: B027 |
| 144 | """Register a function to be called if the executor enters a permanent |
| 145 | failed state. |
| 146 | """ |
| 147 | pass |
| 148 | |
| 149 | def determine_available_memory(self) -> list[int]: # in bytes |
| 150 | return self.collective_rpc("determine_available_memory") |
| 151 | |
| 152 | def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]: |
| 153 | return self.collective_rpc("get_kv_cache_spec") |
| 154 | |
| 155 | def get_supported_kv_cache_layouts(self) -> list[list[str]]: |
| 156 | """Layouts each worker's backends support, most preferred first.""" |
| 157 | return self.collective_rpc("get_supported_kv_cache_layouts") |
| 158 | |
| 159 | def set_kv_cache_layout(self, layout_name: str) -> None: |
| 160 | """Publish the resolved KV cache layout to the workers.""" |
| 161 | self.collective_rpc("set_kv_cache_layout", args=(layout_name,)) |
| 162 | |
| 163 | @overload |
| 164 | def collective_rpc( |
| 165 | self, |
| 166 | method: str | Callable[[WorkerBase], _R], |
| 167 | timeout: float | None = None, |
| 168 | args: tuple = (), |
| 169 | kwargs: dict | None = None, |
| 170 | non_block: Literal[False] = False, |
| 171 | ) -> list[_R]: |
| 172 | pass |
| 173 | |
| 174 | @overload |
| 175 | def collective_rpc( |
| 176 | self, |
| 177 | method: str | Callable[[WorkerBase], _R], |
| 178 | timeout: float | None = None, |
| 179 | args: tuple = (), |
| 180 | kwargs: dict | None = None, |
| 181 | non_block: Literal[True] = True, |
| 182 | ) -> Future[list[_R]]: |
| 183 | pass |
| 184 | |
| 185 | @abstractmethod |
| 186 | def collective_rpc( |
| 187 | self, |
| 188 | method: str | Callable[[WorkerBase], _R], |
| 189 | timeout: float | None = None, |
| 190 | args: tuple = (), |
| 191 | kwargs: dict | None = None, |
| 192 | non_block: bool = False, |
| 193 | ) -> list[_R] | Future[list[_R]]: |
| 194 | """Execute an RPC call on all workers. |
| 195 | |
| 196 | Args: |
| 197 | method: Name of the worker method to execute, or a callable that |
| 198 | is serialized and sent to all workers to execute. |
| 199 | |
| 200 | If the method is a callable, it should accept an additional |
| 201 | `self` argument, in addition to the arguments passed in `args` |
| 202 | and `kwargs`. The `self` argument will be the worker object. |
| 203 | timeout: Maximum time in seconds to wait for execution. Raises a |
| 204 | [`TimeoutError`][] on timeout. `None` means wait indefinitely. |
| 205 | args: Positional arguments to pass to the worker method. |
| 206 | kwargs: Keyword arguments to pass to the worker method. |
| 207 | non_block: If `True`, returns a list of Futures instead of waiting |
| 208 | for the results. |
| 209 | |
| 210 | Returns: |
| 211 | A list containing the results from each worker. |
| 212 | |
| 213 | Note: |
| 214 | It is recommended to use this API to only pass control messages, |
| 215 | and set up data-plane communication to pass data. |
| 216 | |
| 217 | """ |
| 218 | raise NotImplementedError |
| 219 | |
| 220 | def get_kv_connector_handshake_metadata( |
| 221 | self, |
| 222 | ) -> list[dict[tuple[int, int], KVConnectorHandshakeMetadata]]: |
| 223 | return self.collective_rpc("get_kv_connector_handshake_metadata") |
| 224 | |
| 225 | @overload |
| 226 | def execute_model( |
| 227 | self, scheduler_output: SchedulerOutput, non_block: Literal[False] = False |
| 228 | ) -> ModelRunnerOutput | None: |
| 229 | pass |
| 230 | |
| 231 | @overload |
| 232 | def execute_model( |
| 233 | self, scheduler_output: SchedulerOutput, non_block: Literal[True] = True |
| 234 | ) -> Future[ModelRunnerOutput | None]: |
| 235 | pass |
| 236 | |
| 237 | def execute_model( |
| 238 | self, scheduler_output: SchedulerOutput, non_block: bool = False |
| 239 | ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]: |
| 240 | output = self.collective_rpc( # type: ignore[call-overload] |
| 241 | "execute_model", args=(scheduler_output,), non_block=non_block |
| 242 | ) |
| 243 | return output[0] |
| 244 | |
| 245 | @overload |
| 246 | def sample_tokens( |
| 247 | self, grammar_output: GrammarOutput | None, non_block: Literal[False] = False |
| 248 | ) -> ModelRunnerOutput: |
| 249 | pass |
| 250 | |
| 251 | @overload |
| 252 | def sample_tokens( |
| 253 | self, grammar_output: GrammarOutput | None, non_block: Literal[True] = True |
| 254 | ) -> Future[ModelRunnerOutput]: |
| 255 | pass |
| 256 | |
| 257 | def sample_tokens( |
| 258 | self, grammar_output: GrammarOutput | None, non_block: bool = False |
| 259 | ) -> ModelRunnerOutput | Future[ModelRunnerOutput]: |
| 260 | output = self.collective_rpc( # type: ignore[call-overload] |
| 261 | "sample_tokens", args=(grammar_output,), non_block=non_block |
| 262 | ) |
| 263 | return output[0] |
| 264 | |
| 265 | def execute_dummy_batch(self) -> None: |
| 266 | self.collective_rpc("execute_dummy_batch") |
| 267 | |
| 268 | def take_draft_token_ids(self) -> DraftTokenIds | None: |
| 269 | output: list[DraftTokenIds] = self.collective_rpc("take_draft_token_ids") |
| 270 | return output[0] |
| 271 | |
| 272 | def profile(self, is_start: bool = True, profile_prefix: str | None = None): |
| 273 | self.collective_rpc("profile", args=(is_start, profile_prefix)) |
| 274 | |
| 275 | def save_sharded_state( |
| 276 | self, |
| 277 | path: str, |
| 278 | pattern: str | None = None, |
| 279 | max_size: int | None = None, |
| 280 | ) -> None: |
| 281 | self.collective_rpc( |
| 282 | "save_sharded_state", |
| 283 | kwargs=dict(path=path, pattern=pattern, max_size=max_size), |
| 284 | ) |
| 285 | |
| 286 | @abstractmethod |
| 287 | def check_health(self) -> None: |
| 288 | """Checks if the executor is healthy. If not, it should raise an |
| 289 | exception.""" |
| 290 | raise NotImplementedError |
| 291 | |
| 292 | def shutdown(self) -> None: |
| 293 | """Shutdown the executor.""" |
| 294 | self.collective_rpc("shutdown") |
| 295 | |
| 296 | def init_kv_output_aggregator(self, connector: "KVConnectorBase") -> None: |
| 297 | """Init KVOutputAggregator.""" |
| 298 | self.kv_output_aggregator = KVOutputAggregator.from_connector( |
| 299 | connector, self.parallel_config.world_size |
| 300 | ) |
| 301 | |
| 302 | def init_ec_output_aggregator(self) -> None: |
| 303 | self.ec_output_aggregator = ECOutputAggregator() |
| 304 | |
| 305 | @cached_property # Avoid unnecessary RPC calls |
| 306 | def supported_tasks(self) -> tuple[SupportedTask, ...]: |
| 307 | output: list[tuple[SupportedTask, ...]] |
| 308 | output = self.collective_rpc("get_supported_tasks") |
| 309 | return output[0] |
| 310 | |
| 311 | def supports_draft_weight_updates(self) -> bool: |
| 312 | worker_support: list[bool] = self.collective_rpc( |
| 313 | "supports_draft_weight_updates" |
| 314 | ) |
| 315 | return all(worker_support) |
| 316 | |
| 317 | def add_lora(self, lora_request: LoRARequest) -> bool: |
| 318 | assert lora_request.lora_int_id > 0, "lora_id must be greater than 0." |
| 319 | return all(self.collective_rpc("add_lora", args=(lora_request,))) |
| 320 | |
| 321 | def remove_lora(self, lora_id: int) -> bool: |
| 322 | assert lora_id > 0, "lora_id must be greater than 0." |
| 323 | return all(self.collective_rpc("remove_lora", args=(lora_id,))) |
| 324 | |
| 325 | def pin_lora(self, lora_id: int) -> bool: |
| 326 | assert lora_id > 0, "lora_id must be greater than 0." |
| 327 | return all(self.collective_rpc("pin_lora", args=(lora_id,))) |
| 328 | |
| 329 | def list_loras(self) -> set[int]: |
| 330 | sets: list[set[int]] = self.collective_rpc("list_loras") |
| 331 | for s in sets: |
| 332 | assert s == sets[0], "All workers should have the same LORAs." |
| 333 | return sets[0] |
| 334 | |
| 335 | def reset_mm_cache(self) -> None: |
| 336 | """Reset the multi-modal cache in each worker.""" |
| 337 | self.collective_rpc("reset_mm_cache") |
| 338 | |
| 339 | def reset_encoder_cache(self) -> None: |
| 340 | """Reset the encoder cache in each worker to clear cached encoder outputs.""" |
| 341 | self.collective_rpc("reset_encoder_cache") |
| 342 | |
| 343 | @property |
| 344 | def is_sleeping(self) -> bool: |
| 345 | return bool(self.sleeping_tags) |
| 346 | |
| 347 | def sleep(self, level: int = 1): |
| 348 | if "weights" in self.sleeping_tags: |
| 349 | logger.warning("Executor is already sleeping.") |
| 350 | return |
| 351 | time_before_sleep = time.perf_counter() |
| 352 | self.collective_rpc("sleep", kwargs=dict(level=level)) |
| 353 | time_after_sleep = time.perf_counter() |
| 354 | self.sleeping_tags |= SLEEP_TAGS |
| 355 | logger.info( |
| 356 | "It took %.6f seconds to fall asleep.", time_after_sleep - time_before_sleep |
| 357 | ) |
| 358 | |
| 359 | def wake_up(self, tags: list[str] | None = None): |
| 360 | if not self.is_sleeping: |
| 361 | logger.warning("Executor is not sleeping.") |
| 362 | return |
| 363 | if tags: |
| 364 | for tag in tags: |
| 365 | if tag not in self.sleeping_tags: |
| 366 | logger.warning( |
| 367 | "Tag %s is not in sleeping tags %s", tag, self.sleeping_tags |
| 368 | ) |
| 369 | return |
| 370 | time_before_wakeup = time.perf_counter() |
| 371 | self.collective_rpc("wake_up", kwargs=dict(tags=tags)) |
| 372 | time_after_wakeup = time.perf_counter() |
| 373 | logger.info( |
| 374 | "It took %.6f seconds to wake up tags %s.", |
| 375 | time_after_wakeup - time_before_wakeup, |
| 376 | tags if tags is not None else self.sleeping_tags, |
| 377 | ) |
| 378 | if tags: |
| 379 | for tag in tags: |
| 380 | self.sleeping_tags.remove(tag) |
| 381 | else: |
| 382 | self.sleeping_tags.clear() |
| 383 | |
| 384 | def discard(self, tags: tuple[str, ...]) -> None: |
| 385 | tags_to_discard = set(tags) - self.sleeping_tags |
| 386 | if not tags_to_discard: |
| 387 | logger.warning("Tags %s are already sleeping.", tags) |
| 388 | return |
| 389 | time_before_discard = time.perf_counter() |
| 390 | try: |
| 391 | self.collective_rpc("discard", args=(tuple(tags_to_discard),)) |
| 392 | finally: |
| 393 | self.sleeping_tags |= tags_to_discard |
| 394 | time_after_discard = time.perf_counter() |
| 395 | logger.info( |
| 396 | "It took %.6f seconds to discard tags %s.", |
| 397 | time_after_discard - time_before_discard, |
| 398 | tags_to_discard, |
| 399 | ) |
| 400 | |
| 401 | def reinitialize_distributed( |
| 402 | self, reconfig_request: ReconfigureDistributedRequest |
| 403 | ) -> None: |
| 404 | raise NotImplementedError |
| 405 | |
| 406 | @classmethod |
| 407 | def supports_async_scheduling(cls) -> bool: |
| 408 | """Whether the executor supports async scheduling.""" |
| 409 | return False |
| 410 | |
| 411 | |
| 412 | from vllm.v1.executor.uniproc_executor import ( # noqa: E402 |
| 413 | ExecutorWithExternalLauncher as _ExecutorWithExternalLauncher, |
| 414 | ) |
| 415 | from vllm.v1.executor.uniproc_executor import ( # noqa: E402 |
| 416 | UniProcExecutor as _UniProcExecutor, |
| 417 | ) |
| 418 | |
| 419 | # For backwards compatibility. |
| 420 | UniProcExecutor = _UniProcExecutor |
| 421 | ExecutorWithExternalLauncher = _ExecutorWithExternalLauncher |