| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import gc |
| 4 | import os |
| 5 | import queue |
| 6 | import signal |
| 7 | import threading |
| 8 | import time |
| 9 | from collections import defaultdict, deque |
| 10 | from collections.abc import Callable, Generator, Sequence |
| 11 | from concurrent.futures import Future |
| 12 | from contextlib import ExitStack, contextmanager |
| 13 | from enum import IntEnum |
| 14 | from functools import partial |
| 15 | from inspect import isclass, signature |
| 16 | from logging import DEBUG |
| 17 | from multiprocessing.queues import Queue |
| 18 | from typing import Any, TypeVar, cast, get_args |
| 19 | |
| 20 | import msgspec |
| 21 | import zmq |
| 22 | |
| 23 | import vllm.envs as envs |
| 24 | from vllm.config import ParallelConfig, VllmConfig |
| 25 | from vllm.config.pooler import POOLER_CONFIG_LOG_FIELDS |
| 26 | from vllm.distributed import ( |
| 27 | cleanup_dist_env_and_memory, |
| 28 | stateless_destroy_torch_distributed_process_group, |
| 29 | ) |
| 30 | from vllm.envs import enable_envs_cache |
| 31 | from vllm.logger import configure_logging, init_logger |
| 32 | from vllm.logging_utils.dump_input import dump_engine_exception |
| 33 | from vllm.lora.request import LoRARequest |
| 34 | from vllm.multimodal.cache import ( |
| 35 | MultiModalCacheMissError, |
| 36 | engine_receiver_cache_from_config, |
| 37 | ) |
| 38 | from vllm.tasks import POOLING_TASKS, SupportedTask |
| 39 | from vllm.tracing import instrument, maybe_init_worker_tracer |
| 40 | from vllm.transformers_utils.config import maybe_register_config_serialize_by_value |
| 41 | from vllm.utils import numa_utils |
| 42 | from vllm.utils.gc_utils import ( |
| 43 | freeze_gc_heap, |
| 44 | maybe_attach_gc_debug_callback, |
| 45 | ) |
| 46 | from vllm.utils.hashing import get_hash_fn_by_name |
| 47 | from vllm.utils.network_utils import make_zmq_socket |
| 48 | from vllm.utils.system_utils import decorate_logs, set_process_title |
| 49 | from vllm.v1.attention.backends.utils import resolve_kv_cache_layout |
| 50 | from vllm.v1.core.kv_cache_utils import ( |
| 51 | BlockHash, |
| 52 | generate_scheduler_kv_cache_config, |
| 53 | get_kv_cache_configs, |
| 54 | get_request_block_hasher, |
| 55 | init_none_hash, |
| 56 | resolve_kv_cache_block_sizes, |
| 57 | update_kv_cache_capacity, |
| 58 | ) |
| 59 | from vllm.v1.core.sched.interface import PauseState, SchedulerInterface |
| 60 | from vllm.v1.core.sched.output import SchedulerOutput |
| 61 | from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs |
| 62 | from vllm.v1.engine import ( |
| 63 | EEP_NOTIFICATION_CALL_ID, |
| 64 | EEPNotificationType, |
| 65 | EngineCoreOutput, |
| 66 | EngineCoreOutputs, |
| 67 | EngineCoreReadyResponse, |
| 68 | EngineCoreRequest, |
| 69 | EngineCoreRequestType, |
| 70 | FinishReason, |
| 71 | PauseMode, |
| 72 | ReconfigureDistributedRequest, |
| 73 | ReconfigureRankType, |
| 74 | UtilityOutput, |
| 75 | UtilityResult, |
| 76 | ) |
| 77 | from vllm.v1.engine.tensor_ipc import TensorIpcReceiver |
| 78 | from vllm.v1.engine.utils import ( |
| 79 | EngineHandshakeMetadata, |
| 80 | EngineZmqAddresses, |
| 81 | SignalCallback, |
| 82 | get_physical_gpu_ids_for_local_dp_rank, |
| 83 | ) |
| 84 | from vllm.v1.executor import Executor |
| 85 | from vllm.v1.fault_tolerance.engine_core_sentinel import ( |
| 86 | FT_UTILITY_METHOD, |
| 87 | EngineCoreSentinel, |
| 88 | fault_tolerant_wrapper, |
| 89 | ) |
| 90 | from vllm.v1.kv_cache_interface import ( |
| 91 | KVCacheConfig, |
| 92 | get_kv_cache_spec_kind, |
| 93 | is_full_attention_spec, |
| 94 | ) |
| 95 | from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats |
| 96 | from vllm.v1.outputs import ModelRunnerOutput |
| 97 | from vllm.v1.request import Request, RequestStatus |
| 98 | from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr |
| 99 | from vllm.v1.structured_output import StructuredOutputManager |
| 100 | from vllm.v1.utils import compute_iteration_details |
| 101 | from vllm.version import __version__ as VLLM_VERSION |
| 102 | |
| 103 | logger = init_logger(__name__) |
| 104 | |
| 105 | |
| 106 | HANDSHAKE_TIMEOUT_MINS = 5 |
| 107 | |
| 108 | _R = TypeVar("_R") # Return type for collective_rpc |
| 109 | |
| 110 | |
| 111 | class EngineCore: |
| 112 | """Inner loop of vLLM's Engine.""" |
| 113 | |
| 114 | def __init__( |
| 115 | self, |
| 116 | vllm_config: VllmConfig, |
| 117 | executor_class: type[Executor], |
| 118 | log_stats: bool, |
| 119 | executor_fail_callback: Callable | None = None, |
| 120 | include_finished_set: bool = False, |
| 121 | ): |
| 122 | # plugins need to be loaded at the engine/scheduler level too |
| 123 | from vllm.plugins import load_general_plugins |
| 124 | |
| 125 | load_general_plugins() |
| 126 | |
| 127 | self.vllm_config = vllm_config |
| 128 | if not vllm_config.parallel_config.data_parallel_rank_local: |
| 129 | logger.info( |
| 130 | "Initializing a V1 LLM engine (v%s) with config: %s", |
| 131 | VLLM_VERSION, |
| 132 | vllm_config, |
| 133 | ) |
| 134 | |
| 135 | self.log_stats = log_stats |
| 136 | # Opaque weight version supplied by the caller. |
| 137 | self._weight_version = "default" |
| 138 | |
| 139 | # Setup Model. |
| 140 | self.model_executor = executor_class(vllm_config) |
| 141 | self._pooler_config_logged = False |
| 142 | if executor_fail_callback is not None: |
| 143 | self.model_executor.register_failure_callback(executor_fail_callback) |
| 144 | |
| 145 | self.available_gpu_memory_for_kv_cache = -1 |
| 146 | |
| 147 | if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: |
| 148 | self._eep_scale_up_before_kv_init() |
| 149 | |
| 150 | # Setup KV Caches and update CacheConfig after profiling. |
| 151 | kv_cache_config = self._initialize_kv_caches(vllm_config) |
| 152 | self.structured_output_manager = StructuredOutputManager(vllm_config) |
| 153 | |
| 154 | # Setup scheduler. |
| 155 | Scheduler = vllm_config.scheduler_config.get_scheduler_cls() |
| 156 | |
| 157 | if len(kv_cache_config.kv_cache_groups) == 0: # noqa: SIM102 |
| 158 | # Encoder models without KV cache don't support |
| 159 | # chunked prefill. But do SSM models? |
| 160 | if vllm_config.scheduler_config.enable_chunked_prefill: |
| 161 | logger.warning("Disabling chunked prefill for model without KVCache") |
| 162 | vllm_config.scheduler_config.enable_chunked_prefill = False |
| 163 | |
| 164 | scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( |
| 165 | kv_cache_config, vllm_config |
| 166 | ) |
| 167 | |
| 168 | self.scheduler: SchedulerInterface = Scheduler( |
| 169 | vllm_config=vllm_config, |
| 170 | kv_cache_config=kv_cache_config, |
| 171 | structured_output_manager=self.structured_output_manager, |
| 172 | include_finished_set=include_finished_set, |
| 173 | log_stats=self.log_stats, |
| 174 | block_size=scheduler_block_size, |
| 175 | hash_block_size=hash_block_size, |
| 176 | ) |
| 177 | self._initialize_effective_attention_block_size() |
| 178 | self.use_spec_decode = vllm_config.speculative_config is not None |
| 179 | self.check_for_draft_tokens = ( |
| 180 | self.use_spec_decode or vllm_config.model_config.is_diffusion |
| 181 | ) |
| 182 | if self.scheduler.connector is not None: # type: ignore |
| 183 | self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore |
| 184 | if self.scheduler.ec_connector is not None: # type: ignore |
| 185 | self.model_executor.init_ec_output_aggregator() |
| 186 | |
| 187 | self.mm_receiver_cache = engine_receiver_cache_from_config(vllm_config) |
| 188 | |
| 189 | # If a KV connector is initialized for scheduler, we want to collect |
| 190 | # handshake metadata from all workers so the connector in the scheduler |
| 191 | # will have the full context |
| 192 | kv_connector = self.scheduler.get_kv_connector() |
| 193 | if kv_connector is not None: |
| 194 | # Collect and store KV connector xfer metadata from workers |
| 195 | # (after KV cache registration) |
| 196 | xfer_handshake_metadata = ( |
| 197 | self.model_executor.get_kv_connector_handshake_metadata() |
| 198 | ) |
| 199 | |
| 200 | if xfer_handshake_metadata: |
| 201 | # xfer_handshake_metadata is list of dicts from workers |
| 202 | # Each dict already has structure {(pp_rank, tp_rank): metadata} |
| 203 | # Merge all worker dicts into a single dict |
| 204 | content: dict[tuple[int, int], Any] = {} |
| 205 | for worker_dict in xfer_handshake_metadata: |
| 206 | if worker_dict is not None: |
| 207 | content.update(worker_dict) |
| 208 | kv_connector.set_xfer_handshake_metadata_pp_aware(content) |
| 209 | |
| 210 | # Setup batch queue for pipeline parallelism. |
| 211 | # Batch queue for scheduled batches. This enables us to asynchronously |
| 212 | # schedule and execute batches, and is required by pipeline parallelism |
| 213 | # to eliminate pipeline bubbles. |
| 214 | self.batch_queue_size = vllm_config.max_concurrent_batches |
| 215 | self.batch_queue: ( |
| 216 | deque[tuple[Future[ModelRunnerOutput], SchedulerOutput, Future[Any]]] | None |
| 217 | ) = None |
| 218 | if self.batch_queue_size > 1: |
| 219 | logger.debug("Batch queue is enabled with size %d", self.batch_queue_size) |
| 220 | self.batch_queue = deque(maxlen=self.batch_queue_size) |
| 221 | |
| 222 | self.is_mm_encoder_only = vllm_config.is_mm_encoder_only |
| 223 | self.is_pooling_model = vllm_config.model_config.runner_type == "pooling" |
| 224 | |
| 225 | self.request_block_hasher: Callable[[Request], list[BlockHash]] | None = None |
| 226 | if vllm_config.cache_config.enable_prefix_caching or kv_connector is not None: |
| 227 | caching_hash_fn = get_hash_fn_by_name( |
| 228 | vllm_config.cache_config.prefix_caching_hash_algo |
| 229 | ) |
| 230 | init_none_hash(caching_hash_fn) |
| 231 | |
| 232 | self.request_block_hasher = get_request_block_hasher( |
| 233 | hash_block_size, caching_hash_fn |
| 234 | ) |
| 235 | |
| 236 | self.step_fn = ( |
| 237 | self.step if self.batch_queue is None else self.step_with_batch_queue |
| 238 | ) |
| 239 | self.async_scheduling = vllm_config.scheduler_config.async_scheduling |
| 240 | |
| 241 | self.aborts_queue = queue.Queue[list[str]]() |
| 242 | |
| 243 | self._idle_state_callbacks: list[Callable] = [] |
| 244 | |
| 245 | # Mark the startup heap as static so that it's ignored by GC. |
| 246 | # Reduces pause times of oldest generation collections. |
| 247 | freeze_gc_heap() |
| 248 | # If enable, attach GC debugger after static variable freeze. |
| 249 | maybe_attach_gc_debug_callback() |
| 250 | # Enable environment variable cache (e.g. assume no more |
| 251 | # environment variable overrides after this point) |
| 252 | enable_envs_cache() |
| 253 | |
| 254 | @instrument(span_name="Prepare model") |
| 255 | def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: |
| 256 | start = time.time() |
| 257 | |
| 258 | # register all kvcache specs in enginecore process. |
| 259 | register_all_kvcache_specs(vllm_config) |
| 260 | |
| 261 | # Get all kv cache needed by the model |
| 262 | kv_cache_specs = self.model_executor.get_kv_cache_specs() |
| 263 | |
| 264 | # Some layers (e.g. Prefix LM attention) run non-causally and tag their |
| 265 | # KV cache spec with ``non_causal=True``. The specs are collected here in |
| 266 | # the engine-core process (the same process that builds the scheduler), |
| 267 | # so this is the multiproc-safe place to translate that layer-level |
| 268 | # signal into a scheduling policy: chunked prefill and prefix caching |
| 269 | # both assume causal attention and would corrupt non-causal prefill. |
| 270 | if any( |
| 271 | getattr(spec, "non_causal", False) |
| 272 | for worker_specs in kv_cache_specs |
| 273 | for spec in worker_specs.values() |
| 274 | ): |
| 275 | if vllm_config.scheduler_config.enable_chunked_prefill: |
| 276 | logger.info( |
| 277 | "Disabling chunked prefill: model has non-causal attention layers." |
| 278 | ) |
| 279 | vllm_config.scheduler_config.enable_chunked_prefill = False |
| 280 | if vllm_config.cache_config.enable_prefix_caching: |
| 281 | logger.info( |
| 282 | "Disabling prefix caching: model has non-causal attention layers." |
| 283 | ) |
| 284 | vllm_config.cache_config.enable_prefix_caching = False |
| 285 | |
| 286 | # Resolve the KV cache layout before memory profiling: workers that |
| 287 | # capture full cudagraphs initialize a minimal KV cache during it. |
| 288 | # Attention-free models resolve the default so layout reads never precede |
| 289 | # resolution. |
| 290 | layout = resolve_kv_cache_layout( |
| 291 | vllm_config, |
| 292 | self.model_executor.get_supported_kv_cache_layouts(), |
| 293 | [s for specs in kv_cache_specs for s in specs.values()], |
| 294 | ) |
| 295 | self.model_executor.set_kv_cache_layout(layout.name) |
| 296 | |
| 297 | has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs) |
| 298 | if has_kv_cache: |
| 299 | if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: |
| 300 | # NOTE(yongji): should already be set |
| 301 | # during _eep_scale_up_before_kv_init |
| 302 | assert self.available_gpu_memory_for_kv_cache > 0 |
| 303 | available_gpu_memory = [self.available_gpu_memory_for_kv_cache] * len( |
| 304 | kv_cache_specs |
| 305 | ) |
| 306 | else: |
| 307 | # Profiles the peak memory usage of the model to determine how |
| 308 | # much memory can be allocated for kv cache. |
| 309 | available_gpu_memory = self.model_executor.determine_available_memory() |
| 310 | self.available_gpu_memory_for_kv_cache = available_gpu_memory[0] |
| 311 | else: |
| 312 | # Attention free models don't need memory for kv cache |
| 313 | available_gpu_memory = [0] * len(kv_cache_specs) |
| 314 | |
| 315 | assert len(kv_cache_specs) == len(available_gpu_memory) |
| 316 | |
| 317 | # Track max_model_len before KV cache config to detect auto-fit changes |
| 318 | max_model_len_before = vllm_config.model_config.max_model_len |
| 319 | |
| 320 | kv_cache_configs = get_kv_cache_configs( |
| 321 | vllm_config, kv_cache_specs, available_gpu_memory |
| 322 | ) |
| 323 | for kv_cache_config in kv_cache_configs: |
| 324 | kv_cache_config.kv_cache_layout = vllm_config.cache_config.kv_cache_layout |
| 325 | |
| 326 | # If auto-fit reduced max_model_len, sync the new value to workers. |
| 327 | # This is needed because workers were spawned before memory profiling |
| 328 | # and have the original (larger) max_model_len cached. |
| 329 | max_model_len_after = vllm_config.model_config.max_model_len |
| 330 | if max_model_len_after != max_model_len_before: |
| 331 | self.collective_rpc("update_max_model_len", args=(max_model_len_after,)) |
| 332 | |
| 333 | scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) |
| 334 | vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks |
| 335 | kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups |
| 336 | if kv_cache_groups: |
| 337 | # Exclude groups that opt out of prefix caching (e.g. GLM-5.3-Flash |
| 338 | # kpool tail, a 1-block/req scratch buffer with block_size=kpool): |
| 339 | # their small block_size would otherwise drag the global block_size |
| 340 | # below the real allocator block size and desync it from mamba. |
| 341 | participating = [ |
| 342 | g.kv_cache_spec.block_size |
| 343 | for g in kv_cache_groups |
| 344 | if g.kv_cache_spec.prefix_cacheable |
| 345 | ] |
| 346 | vllm_config.cache_config.block_size = min( |
| 347 | participating |
| 348 | if participating |
| 349 | else [g.kv_cache_spec.block_size for g in kv_cache_groups] |
| 350 | ) |
| 351 | update_kv_cache_capacity(vllm_config, scheduler_kv_cache_config) |
| 352 | |
| 353 | vllm_config.validate_block_size() |
| 354 | |
| 355 | self.model_executor.initialize_from_config(kv_cache_configs) |
| 356 | if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: |
| 357 | self.model_executor.compile_or_warm_up_model() |
| 358 | |
| 359 | elapsed = time.time() - start |
| 360 | compile_time = vllm_config.compilation_config.compilation_time |
| 361 | encoder_compile_time = vllm_config.compilation_config.encoder_compilation_time |
| 362 | if encoder_compile_time > 0: |
| 363 | logger.info_once( |
| 364 | "init engine (profile, create kv cache, warmup model) took " |
| 365 | "%.2f s (compilation: %.2f s — language_model: %.2f s, " |
| 366 | "encoder: %.2f s)", |
| 367 | elapsed, |
| 368 | compile_time + encoder_compile_time, |
| 369 | compile_time, |
| 370 | encoder_compile_time, |
| 371 | ) |
| 372 | elif compile_time > 0: |
| 373 | logger.info_once( |
| 374 | "init engine (profile, create kv cache, warmup model) took " |
| 375 | "%.2f s (compilation: %.2f s)", |
| 376 | elapsed, |
| 377 | compile_time, |
| 378 | ) |
| 379 | else: |
| 380 | logger.info_once( |
| 381 | "init engine (profile, create kv cache, warmup model) took %.2f s", |
| 382 | elapsed, |
| 383 | ) |
| 384 | return scheduler_kv_cache_config |
| 385 | |
| 386 | def _initialize_effective_attention_block_size(self) -> None: |
| 387 | cache_config = self.vllm_config.cache_config |
| 388 | cache_config.effective_attention_block_size = None |
| 389 | cache_manager = getattr(self.scheduler, "kv_cache_manager", None) |
| 390 | if cache_manager is None: |
| 391 | return |
| 392 | block_sizes = { |
| 393 | manager.block_size |
| 394 | for manager in cache_manager.coordinator.single_type_managers |
| 395 | if is_full_attention_spec(manager.kv_cache_spec) |
| 396 | } |
| 397 | cache_config.effective_attention_block_size = ( |
| 398 | block_sizes.pop() if len(block_sizes) == 1 else None |
| 399 | ) |
| 400 | |
| 401 | def get_supported_tasks(self) -> tuple[SupportedTask, ...]: |
| 402 | supported_tasks = self.model_executor.supported_tasks |
| 403 | self._log_pooler_config(supported_tasks) |
| 404 | return supported_tasks |
| 405 | |
| 406 | def _log_pooler_config(self, supported_tasks: tuple[SupportedTask, ...]) -> None: |
| 407 | if self._pooler_config_logged: |
| 408 | return |
| 409 | |
| 410 | model_config = self.vllm_config.model_config |
| 411 | pooler_config = model_config.pooler_config |
| 412 | if ( |
| 413 | self.vllm_config.parallel_config.data_parallel_rank_local |
| 414 | or model_config.runner_type != "pooling" |
| 415 | or pooler_config is None |
| 416 | ): |
| 417 | return |
| 418 | |
| 419 | supported_pooling_tasks = tuple( |
| 420 | sorted(set(supported_tasks) & set(POOLING_TASKS)) |
| 421 | ) |
| 422 | if not supported_pooling_tasks: |
| 423 | return |
| 424 | |
| 425 | self._pooler_config_logged = True |
| 426 | task_set = set(supported_pooling_tasks) |
| 427 | use_activation = pooler_config.use_activation |
| 428 | if use_activation is None: |
| 429 | use_activation = True |
| 430 | sources = getattr(model_config, "_pooler_config_sources", {}) |
| 431 | pooling_type_field = ( |
| 432 | "seq_pooling_type" |
| 433 | if task_set & {"embed", "classify"} |
| 434 | else "tok_pooling_type" |
| 435 | ) |
| 436 | |
| 437 | def log_field(name: str, field: str) -> str: |
| 438 | value = ( |
| 439 | use_activation |
| 440 | if field == "use_activation" |
| 441 | else getattr(pooler_config, field) |
| 442 | ) |
| 443 | source = sources.get(field, "unknown") |
| 444 | return f"{name}={value}(source={source})" |
| 445 | |
| 446 | log_items = [("pooling_type", pooling_type_field)] |
| 447 | log_items.extend( |
| 448 | (field, field) |
| 449 | for field in POOLER_CONFIG_LOG_FIELDS |
| 450 | if field != pooling_type_field |
| 451 | ) |
| 452 | config_fields = ", ".join(log_field(name, field) for name, field in log_items) |
| 453 | |
| 454 | logger.info_once( |
| 455 | "Resolved pooling config: %s, supported_tasks=%s", |
| 456 | config_fields, |
| 457 | supported_pooling_tasks, |
| 458 | ) |
| 459 | |
| 460 | def get_kv_cache_group_metadata(self) -> list[dict[str, int | str | None]]: |
| 461 | """Return serializable KV cache metadata for external event consumers |
| 462 | (e.g., Dynamo). This function has no call sites within vLLM but needs to |
| 463 | be kept. |
| 464 | """ |
| 465 | kv_cache_config = getattr(self.scheduler, "kv_cache_config", None) |
| 466 | if kv_cache_config is None: |
| 467 | return [] |
| 468 | |
| 469 | kv_cache_manager = cast(Any, self.scheduler).kv_cache_manager |
| 470 | managers = kv_cache_manager.coordinator.single_type_managers |
| 471 | metadata: list[dict[str, int | str | None]] = [] |
| 472 | for group_idx, (group, manager) in enumerate( |
| 473 | zip(kv_cache_config.kv_cache_groups, managers, strict=True) |
| 474 | ): |
| 475 | spec = group.kv_cache_spec |
| 476 | metadata.append( |
| 477 | { |
| 478 | "group_idx": group_idx, |
| 479 | "kind": get_kv_cache_spec_kind(spec).value, |
| 480 | "block_size": manager.block_size, |
| 481 | "sliding_window": getattr(spec, "sliding_window", None), |
| 482 | } |
| 483 | ) |
| 484 | return metadata |
| 485 | |
| 486 | def add_request(self, request: Request, request_wave: int = 0): |
| 487 | """Add request to the scheduler. |
| 488 | |
| 489 | `request_wave`: indicate which wave of requests this is expected to |
| 490 | belong to in DP case |
| 491 | """ |
| 492 | # Validate the request_id type. |
| 493 | if not isinstance(request.request_id, str): |
| 494 | raise TypeError( |
| 495 | f"request_id must be a string, got {type(request.request_id)}" |
| 496 | ) |
| 497 | |
| 498 | if pooling_params := request.pooling_params: |
| 499 | supported_pooling_tasks = [ |
| 500 | task for task in self.get_supported_tasks() if task in POOLING_TASKS |
| 501 | ] |
| 502 | |
| 503 | if pooling_params.task not in supported_pooling_tasks: |
| 504 | raise ValueError( |
| 505 | f"Unsupported task: {pooling_params.task!r} " |
| 506 | f"Supported tasks: {supported_pooling_tasks}" |
| 507 | ) |
| 508 | |
| 509 | if request.kv_transfer_params is not None and ( |
| 510 | not self.scheduler.get_kv_connector() |
| 511 | ): |
| 512 | logger.warning( |
| 513 | "Got kv_transfer_params, but no KVConnector found. " |
| 514 | "Disabling KVTransfer for this request." |
| 515 | ) |
| 516 | |
| 517 | if ( |
| 518 | request.ec_transfer_params is not None |
| 519 | and self.scheduler.get_ec_connector() is None |
| 520 | ): |
| 521 | logger.warning( |
| 522 | "Got ec_transfer_params, but no ECConnector found. " |
| 523 | "Disabling ECTransfer for this request." |
| 524 | ) |
| 525 | |
| 526 | self.scheduler.add_request(request) |
| 527 | if request.abort_immediately: |
| 528 | # Immediately abort so the connector's request_finished hook runs |
| 529 | # to free any pre-admission KV-transfer resources. |
| 530 | self.abort_requests([request.request_id]) |
| 531 | |
| 532 | def abort_requests(self, request_ids: list[str]): |
| 533 | """Abort requests from the scheduler.""" |
| 534 | # TODO: The scheduler doesn't really need to know the |
| 535 | # specific finish reason, TBD whether we propagate that |
| 536 | # (i.e. client-aborted vs stop criteria met). |
| 537 | self.scheduler.finish_requests(request_ids, RequestStatus.FINISHED_ABORTED) |
| 538 | |
| 539 | @contextmanager |
| 540 | def log_error_detail(self, scheduler_output: SchedulerOutput): |
| 541 | """Execute the model and log detailed info on failure.""" |
| 542 | try: |
| 543 | yield |
| 544 | except Exception as err: |
| 545 | # We do not want to catch BaseException here since we're only |
| 546 | # interested in dumping info when the exception is due to an |
| 547 | # error from execute_model itself. |
| 548 | |
| 549 | # NOTE: This method is exception-free |
| 550 | dump_engine_exception( |
| 551 | self.vllm_config, scheduler_output, self.scheduler.make_stats() |
| 552 | ) |
| 553 | raise err |
| 554 | |
| 555 | @contextmanager |
| 556 | def capture_iteration_details( |
| 557 | self, scheduler_output: SchedulerOutput | None |
| 558 | ) -> Generator[SchedulerIterationDetails | None, None, None]: |
| 559 | enable_details = ( |
| 560 | self.vllm_config.observability_config.enable_logging_iteration_details |
| 561 | ) |
| 562 | if not self.log_stats or not enable_details: |
| 563 | yield None |
| 564 | return |
| 565 | # 0-token step: let the dummy_batch wrapper log it (avoids double-log). |
| 566 | if ( |
| 567 | scheduler_output is not None |
| 568 | and scheduler_output.total_num_scheduled_tokens == 0 |
| 569 | ): |
| 570 | yield None |
| 571 | return |
| 572 | |
| 573 | iteration_index = getattr(self, "_iteration_index", 0) |
| 574 | # scheduler_output=None marks a DP dummy iteration. |
| 575 | if scheduler_output is None: |
| 576 | iteration_details = SchedulerIterationDetails( |
| 577 | iteration_index=iteration_index, |
| 578 | num_ctx_requests=0, |
| 579 | num_ctx_tokens=0, |
| 580 | num_generation_requests=0, |
| 581 | num_generation_tokens=0, |
| 582 | elapsed_ms=0.0, |
| 583 | is_dummy=True, |
| 584 | ) |
| 585 | else: |
| 586 | details = compute_iteration_details(scheduler_output) |
| 587 | iteration_details = SchedulerIterationDetails( |
| 588 | iteration_index=iteration_index, |
| 589 | num_ctx_requests=details.num_ctx_requests, |
| 590 | num_ctx_tokens=details.num_ctx_tokens, |
| 591 | num_generation_requests=details.num_generation_requests, |
| 592 | num_generation_tokens=details.num_generation_tokens, |
| 593 | elapsed_ms=0.0, |
| 594 | num_encoder_inputs=details.num_encoder_inputs, |
| 595 | num_encoder_output_tokens=details.num_encoder_output_tokens, |
| 596 | ) |
| 597 | |
| 598 | start_time = time.monotonic() |
| 599 | yield iteration_details |
| 600 | iteration_details.elapsed_ms = (time.monotonic() - start_time) * 1000 |
| 601 | self._iteration_index = iteration_index + 1 |
| 602 | |
| 603 | def _make_iteration_details_stats( |
| 604 | self, iteration_details: SchedulerIterationDetails |
| 605 | ) -> SchedulerStats: |
| 606 | stats = self.scheduler.make_stats() or SchedulerStats() |
| 607 | stats.iteration_details = iteration_details |
| 608 | return stats |
| 609 | |
| 610 | def _attach_iteration_details( |
| 611 | self, |
| 612 | outputs: dict[int, EngineCoreOutputs], |
| 613 | iteration_details: SchedulerIterationDetails | None, |
| 614 | ) -> None: |
| 615 | if iteration_details is None: |
| 616 | return |
| 617 | |
| 618 | if (eco := next(iter(outputs.values()), None)) is None: |
| 619 | outputs[0] = eco = EngineCoreOutputs() |
| 620 | if eco.scheduler_stats is None: |
| 621 | eco.scheduler_stats = self._make_iteration_details_stats(iteration_details) |
| 622 | else: |
| 623 | eco.scheduler_stats.iteration_details = iteration_details |
| 624 | |
| 625 | def _should_throttle_prefills(self) -> bool: |
| 626 | """Whether to defer new prefills this step (DP prefill balancing). |
| 627 | Overridden by the DP engine core; never throttles otherwise.""" |
| 628 | return False |
| 629 | |
| 630 | def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: |
| 631 | """Schedule, execute, and make output. |
| 632 | |
| 633 | Returns tuple of outputs and a flag indicating whether the model |
| 634 | was executed. |
| 635 | """ |
| 636 | # Check for any requests remaining in the scheduler - unfinished, |
| 637 | # or finished and not yet removed from the batch. |
| 638 | if not self.scheduler.has_requests(): |
| 639 | return {}, False |
| 640 | scheduler_output = self.scheduler.schedule(self._should_throttle_prefills()) |
| 641 | future = self.model_executor.execute_model(scheduler_output, non_block=True) |
| 642 | grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output) |
| 643 | with ( |
| 644 | self.capture_iteration_details(scheduler_output) as iteration_details, |
| 645 | self.log_error_detail(scheduler_output), |
| 646 | ): |
| 647 | model_output = future.result() |
| 648 | if model_output is None: |
| 649 | model_output = self.model_executor.sample_tokens(grammar_output) |
| 650 | |
| 651 | # Before processing the model output, process any aborts that happened |
| 652 | # during the model execution. |
| 653 | self._process_aborts_queue() |
| 654 | engine_core_outputs = self.scheduler.update_from_output( |
| 655 | scheduler_output, model_output |
| 656 | ) |
| 657 | self._attach_iteration_details(engine_core_outputs, iteration_details) |
| 658 | |
| 659 | return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 |
| 660 | |
| 661 | def post_step(self, model_executed: bool) -> None: |
| 662 | # When using async scheduling we can't get draft token ids in advance, |
| 663 | # so we update draft token ids in the worker process and don't |
| 664 | # need to update draft token ids here. |
| 665 | if self.check_for_draft_tokens and not self.async_scheduling and model_executed: |
| 666 | draft_token_ids = self.model_executor.take_draft_token_ids() |
| 667 | if draft_token_ids is not None: |
| 668 | self.scheduler.update_draft_token_ids(draft_token_ids) |
| 669 | |
| 670 | def step_with_batch_queue( |
| 671 | self, |
| 672 | ) -> tuple[dict[int, EngineCoreOutputs] | None, bool]: |
| 673 | """Schedule and execute batches with the batch queue. |
| 674 | Note that if nothing to output in this step, None is returned. |
| 675 | |
| 676 | The execution flow is as follows: |
| 677 | 1. Try to schedule a new batch if the batch queue is not full. |
| 678 | If a new batch is scheduled, directly return an empty engine core |
| 679 | output. In other words, fulfilling the batch queue has a higher priority |
| 680 | than getting model outputs. |
| 681 | 2. If there is no new scheduled batch, meaning that the batch queue |
| 682 | is full or no other requests can be scheduled, we block until the first |
| 683 | batch in the job queue is finished. |
| 684 | 3. Update the scheduler from the output. |
| 685 | """ |
| 686 | batch_queue = self.batch_queue |
| 687 | assert batch_queue is not None |
| 688 | |
| 689 | # Try to schedule a new batch if the batch queue is not full, but |
| 690 | # the scheduler may return an empty batch if all requests are scheduled. |
| 691 | # Note that this is not blocking. |
| 692 | assert len(batch_queue) < self.batch_queue_size |
| 693 | |
| 694 | model_executed = False |
| 695 | deferred_scheduler_output = None |
| 696 | if self.scheduler.has_requests(): |
| 697 | scheduler_output = self.scheduler.schedule(self._should_throttle_prefills()) |
| 698 | with self.log_error_detail(scheduler_output): |
| 699 | exec_future = self.model_executor.execute_model( |
| 700 | scheduler_output, non_block=True |
| 701 | ) |
| 702 | if not self.is_mm_encoder_only: |
| 703 | model_executed = scheduler_output.total_num_scheduled_tokens > 0 |
| 704 | |
| 705 | if self.is_pooling_model or not model_executed: |
| 706 | # No sampling required (no requests scheduled). |
| 707 | future = cast(Future[ModelRunnerOutput], exec_future) |
| 708 | else: |
| 709 | if not scheduler_output.pending_structured_output_tokens: |
| 710 | # We aren't waiting for any tokens, get any grammar output |
| 711 | # and sample immediately. |
| 712 | grammar_output = self.scheduler.get_grammar_bitmask( |
| 713 | scheduler_output |
| 714 | ) |
| 715 | future = self.model_executor.sample_tokens( |
| 716 | grammar_output, non_block=True |
| 717 | ) |
| 718 | else: |
| 719 | # We need to defer sampling until we have processed the model output |
| 720 | # from the prior step. |
| 721 | deferred_scheduler_output = scheduler_output |
| 722 | |
| 723 | if not deferred_scheduler_output: |
| 724 | # Add this step's future to the queue. |
| 725 | batch_queue.appendleft((future, scheduler_output, exec_future)) |
| 726 | if len(batch_queue) < self.batch_queue_size and ( |
| 727 | model_executed or self.scheduler.has_requests() |
| 728 | ): |
| 729 | # Don't block on next worker response unless the queue is full |
| 730 | # or there are no more requests to schedule. |
| 731 | return None, model_executed |
| 732 | |
| 733 | elif not batch_queue: |
| 734 | # Queue is empty. We should not reach here since this method should |
| 735 | # only be called when the scheduler contains requests or the queue |
| 736 | # is non-empty. |
| 737 | return None, False |
| 738 | |
| 739 | # Block until the next result is available. |
| 740 | future, scheduler_output, exec_model_fut = batch_queue.pop() |
| 741 | with ( |
| 742 | self.capture_iteration_details(scheduler_output) as iteration_details, |
| 743 | self.log_error_detail(scheduler_output), |
| 744 | ): |
| 745 | model_output = future.result() |
| 746 | if model_output is None: |
| 747 | # None from sample_tokens() implies that the original execute_model() |
| 748 | # call failed - raise that exception. |
| 749 | exec_model_fut.result() |
| 750 | raise RuntimeError("unexpected error") |
| 751 | |
| 752 | # Before processing the model output, process any aborts that happened |
| 753 | # during the model execution. |
| 754 | self._process_aborts_queue() |
| 755 | engine_core_outputs = self.scheduler.update_from_output( |
| 756 | scheduler_output, model_output |
| 757 | ) |
| 758 | self._attach_iteration_details(engine_core_outputs, iteration_details) |
| 759 | |
| 760 | # NOTE(nick): We can either handle the deferred tasks here or save |
| 761 | # in a field and do it immediately once step_with_batch_queue is |
| 762 | # re-called. The latter slightly favors TTFT over TPOT/throughput. |
| 763 | if deferred_scheduler_output: |
| 764 | # When draft tokens are used with structured output, validate them |
| 765 | # before computing the grammar bitmask for the deferred request. |
| 766 | if self.check_for_draft_tokens: |
| 767 | draft_token_ids = self.model_executor.take_draft_token_ids() |
| 768 | if draft_token_ids is not None: |
| 769 | # Update the draft token ids in the scheduler output to |
| 770 | # filter out the invalid spec tokens, which will be padded |
| 771 | # with -1 and skipped by the grammar bitmask computation. |
| 772 | self.scheduler.update_draft_token_ids_in_output( |
| 773 | draft_token_ids, deferred_scheduler_output |
| 774 | ) |
| 775 | # We now have the tokens needed to compute the bitmask for the |
| 776 | # deferred request. Get the bitmask and call sample tokens. |
| 777 | grammar_output = self.scheduler.get_grammar_bitmask( |
| 778 | deferred_scheduler_output |
| 779 | ) |
| 780 | future = self.model_executor.sample_tokens(grammar_output, non_block=True) |
| 781 | batch_queue.appendleft((future, deferred_scheduler_output, exec_future)) |
| 782 | |
| 783 | return engine_core_outputs, model_executed |
| 784 | |
| 785 | def _process_aborts_queue(self): |
| 786 | if not self.aborts_queue.empty(): |
| 787 | request_ids = [] |
| 788 | while not self.aborts_queue.empty(): |
| 789 | ids = self.aborts_queue.get_nowait() |
| 790 | # Should be a list here, but also handle string just in case. |
| 791 | request_ids.extend((ids,) if isinstance(ids, str) else ids) |
| 792 | # More efficient to abort all as a single batch. |
| 793 | self.abort_requests(request_ids) |
| 794 | |
| 795 | def shutdown(self): |
| 796 | logger.debug_once("[shutdown] EngineCore: tearing down local resources") |
| 797 | self.structured_output_manager.clear_backend() |
| 798 | if self.model_executor: |
| 799 | self.model_executor.shutdown() |
| 800 | if self.scheduler: |
| 801 | self.scheduler.shutdown() |
| 802 | |
| 803 | # Undo the gc.freeze() from __init__ so that the objects allocated |
| 804 | # during engine startup (model weights, KV caches, etc.) become |
| 805 | # visible to the garbage collector again. Without this, deleting |
| 806 | # the engine in-process (e.g. unit tests) leaks GPU memory. |
| 807 | gc.unfreeze() |
| 808 | # Tear down distributed state initialized in this EngineCore process |
| 809 | # before it exits and release cached memory. |
| 810 | cleanup_dist_env_and_memory() |
| 811 | logger.debug_once("[shutdown] EngineCore: local resource teardown complete") |
| 812 | |
| 813 | def profile(self, is_start: bool = True, profile_prefix: str | None = None): |
| 814 | self.model_executor.profile(is_start, profile_prefix) |
| 815 | |
| 816 | def reset_mm_cache(self): |
| 817 | # NOTE: Since this is mainly for debugging, we don't attempt to |
| 818 | # re-sync the internal caches (P0 sender, P1 receiver) |
| 819 | if self.scheduler.has_unfinished_requests(): |
| 820 | logger.warning( |
| 821 | "Resetting the multi-modal cache when requests are " |
| 822 | "in progress may lead to desynced internal caches." |
| 823 | ) |
| 824 | |
| 825 | # The cache either exists in EngineCore or WorkerWrapperBase |
| 826 | if self.mm_receiver_cache is not None: |
| 827 | self.mm_receiver_cache.clear_cache() |
| 828 | |
| 829 | self.model_executor.reset_mm_cache() |
| 830 | |
| 831 | def reset_prefix_cache( |
| 832 | self, reset_running_requests: bool = False, reset_connector: bool = False |
| 833 | ) -> bool: |
| 834 | return self.scheduler.reset_prefix_cache( |
| 835 | reset_running_requests, reset_connector |
| 836 | ) |
| 837 | |
| 838 | def reset_encoder_cache(self) -> None: |
| 839 | """Reset the encoder cache to invalidate all cached encoder outputs. |
| 840 | |
| 841 | This should be called when model weights are updated to ensure |
| 842 | stale vision embeddings computed with old weights are not reused. |
| 843 | Clears both the scheduler's cache manager and the GPU model runner's cache. |
| 844 | """ |
| 845 | # NOTE: Since this is mainly for debugging, we don't attempt to |
| 846 | # re-sync the internal caches (P0 sender, P1 receiver) |
| 847 | if self.scheduler.has_unfinished_requests(): |
| 848 | logger.warning( |
| 849 | "Resetting the encoder cache when requests are " |
| 850 | "in progress may lead to desynced internal caches." |
| 851 | ) |
| 852 | |
| 853 | # Reset the scheduler's encoder cache manager (logical state) |
| 854 | self.scheduler.reset_encoder_cache() |
| 855 | # Reset the GPU model runner's encoder cache (physical storage) |
| 856 | self.model_executor.reset_encoder_cache() |
| 857 | |
| 858 | def _reset_caches( |
| 859 | self, |
| 860 | reset_running_requests: bool = True, |
| 861 | reset_connector: bool = True, |
| 862 | ) -> None: |
| 863 | # reset_connector=True so external connectors clear alongside |
| 864 | # local caches, matching the pause_generation(clear_cache=True) |
| 865 | # contract. No-op when no connector is configured. |
| 866 | if not self.reset_prefix_cache( |
| 867 | reset_running_requests=reset_running_requests, |
| 868 | reset_connector=reset_connector, |
| 869 | ): |
| 870 | raise RuntimeError("Failed to reset the KV connector cache.") |
| 871 | self.reset_mm_cache() |
| 872 | self.reset_encoder_cache() |
| 873 | |
| 874 | def _finish_pause(self, clear_cache: bool) -> None: |
| 875 | # A completed pause promises an idle device: nothing else waits on |
| 876 | # the last dummy batch an idle DP rank launches. |
| 877 | self.model_executor.collective_rpc("synchronize_device") |
| 878 | if clear_cache: |
| 879 | self._reset_caches() |
| 880 | |
| 881 | def pause_scheduler( |
| 882 | self, mode: PauseMode = "abort", clear_cache: bool = True |
| 883 | ) -> Future | None: |
| 884 | """Pause generation; behavior depends on mode. |
| 885 | |
| 886 | All pause modes queue new adds -- "abort" and "keep" skip step(); |
| 887 | "wait" allows step() so in-flight requests can drain. |
| 888 | |
| 889 | - ``abort``: Set PAUSED_NEW, abort all requests, wait for abort |
| 890 | outputs to be sent (when running with output_queue), optionally |
| 891 | clear caches, then complete the returned Future. |
| 892 | - ``wait``: Set PAUSED_NEW (queue adds, keep stepping); when drained, |
| 893 | optionally clear caches, then complete the returned Future. |
| 894 | - ``keep``: Set PAUSED_ALL; return a Future that completes when the |
| 895 | output queue is empty. |
| 896 | """ |
| 897 | if mode not in get_args(PauseMode): |
| 898 | raise ValueError(f"Invalid pause mode: {mode}") |
| 899 | if mode == "wait": |
| 900 | raise ValueError("'wait' mode can't be used in inproc-engine mode") |
| 901 | |
| 902 | if mode == "abort": |
| 903 | self.scheduler.finish_requests(None, RequestStatus.FINISHED_ABORTED) |
| 904 | |
| 905 | pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW |
| 906 | self.scheduler.set_pause_state(pause_state) |
| 907 | self._finish_pause(clear_cache) |
| 908 | |
| 909 | return None |
| 910 | |
| 911 | def resume_scheduler(self) -> None: |
| 912 | """Resume the scheduler and flush any requests queued while paused.""" |
| 913 | self.scheduler.set_pause_state(PauseState.UNPAUSED) |
| 914 | |
| 915 | def is_scheduler_paused(self) -> bool: |
| 916 | """Return whether the scheduler is in any pause state.""" |
| 917 | return self.scheduler.pause_state != PauseState.UNPAUSED |
| 918 | |
| 919 | def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None | Future: |
| 920 | """Put the engine to sleep at the specified level. |
| 921 | |
| 922 | Args: |
| 923 | level: Sleep level. |
| 924 | - Level 0: Pause scheduling only. Requests are still accepted |
| 925 | but not processed. No GPU memory changes. |
| 926 | - Level 1: Offload model weights to CPU, discard KV cache. |
| 927 | - Level 2: Discard all GPU memory. |
| 928 | mode: Pause mode - how to deal with any existing requests, see |
| 929 | documentation of pause_scheduler method. |
| 930 | |
| 931 | """ |
| 932 | # Pause scheduler before sleeping. |
| 933 | clear_prefix_cache = level >= 1 |
| 934 | pause_future = self.pause_scheduler(mode=mode, clear_cache=clear_prefix_cache) |
| 935 | if level < 1: |
| 936 | return pause_future |
| 937 | |
| 938 | # Level 1+: Delegate to executor for GPU memory management |
| 939 | model_executor = self.model_executor |
| 940 | if pause_future is None: |
| 941 | model_executor.sleep(level) |
| 942 | return None |
| 943 | |
| 944 | future = Future[Any]() |
| 945 | |
| 946 | def pause_complete(f: Future): |
| 947 | try: |
| 948 | f.result() # propagate any exception |
| 949 | future.set_result(model_executor.sleep(level)) |
| 950 | except Exception as e: |
| 951 | future.set_exception(e) |
| 952 | |
| 953 | logger.info("Waiting for in-flight requests to complete before sleeping...") |
| 954 | pause_future.add_done_callback(pause_complete) |
| 955 | return future |
| 956 | |
| 957 | def wake_up(self, tags: list[str] | None = None) -> bool: |
| 958 | """Wake up the engine from sleep. |
| 959 | |
| 960 | Args: |
| 961 | tags: Tags to wake up. Use ["scheduling"] for level 0 wake up. |
| 962 | |
| 963 | Returns: |
| 964 | Whether all executor memory is resident again (fully awake). |
| 965 | |
| 966 | """ |
| 967 | if tags is not None and "scheduling" in tags: |
| 968 | # Remove "scheduling" from tags if there are other tags to process. |
| 969 | tags = [t for t in tags if t != "scheduling"] |
| 970 | |
| 971 | if tags is None or tags: |
| 972 | self.model_executor.wake_up(tags) |
| 973 | |
| 974 | # Partial wakes intentionally keep the remaining allocations asleep. |
| 975 | # Resume scheduling only once all executor memory is resident again. |
| 976 | fully_awake = not self.model_executor.is_sleeping |
| 977 | if fully_awake: |
| 978 | self.resume_scheduler() |
| 979 | return fully_awake |
| 980 | |
| 981 | def release_kv_cache_memory(self) -> None: |
| 982 | """Discard KV cache physical memory. Requires a completed pause |
| 983 | and all executor memory to be resident. Kept requests are recomputed |
| 984 | after wake-up. |
| 985 | """ |
| 986 | if not ( |
| 987 | self.is_scheduler_paused() |
| 988 | and not self.scheduler.has_requests() |
| 989 | and not self.batch_queue |
| 990 | ): |
| 991 | raise RuntimeError( |
| 992 | "release_kv_cache_memory() requires a completed pause first" |
| 993 | ) |
| 994 | if self.model_executor.is_sleeping: |
| 995 | raise RuntimeError( |
| 996 | "release_kv_cache_memory() requires all executor memory to be resident" |
| 997 | ) |
| 998 | self._reset_caches() |
| 999 | self.model_executor.discard(("kv_cache",)) |
| 1000 | |
| 1001 | def is_sleeping(self) -> bool: |
| 1002 | """Check if engine is sleeping at any level.""" |
| 1003 | return self.is_scheduler_paused() or self.model_executor.is_sleeping |
| 1004 | |
| 1005 | def execute_dummy_batch(self): |
| 1006 | self.model_executor.execute_dummy_batch() |
| 1007 | |
| 1008 | def compute_weight_checksums(self) -> list[dict[str, str]]: |
| 1009 | return self.collective_rpc("compute_weight_checksums") |
| 1010 | |
| 1011 | def add_lora(self, lora_request: LoRARequest) -> bool: |
| 1012 | return self.model_executor.add_lora(lora_request) |
| 1013 | |
| 1014 | def remove_lora(self, lora_id: int) -> bool: |
| 1015 | return self.model_executor.remove_lora(lora_id) |
| 1016 | |
| 1017 | def list_loras(self) -> set[int]: |
| 1018 | return self.model_executor.list_loras() |
| 1019 | |
| 1020 | def pin_lora(self, lora_id: int) -> bool: |
| 1021 | return self.model_executor.pin_lora(lora_id) |
| 1022 | |
| 1023 | def save_sharded_state( |
| 1024 | self, |
| 1025 | path: str, |
| 1026 | pattern: str | None = None, |
| 1027 | max_size: int | None = None, |
| 1028 | ) -> None: |
| 1029 | self.model_executor.save_sharded_state( |
| 1030 | path=path, pattern=pattern, max_size=max_size |
| 1031 | ) |
| 1032 | |
| 1033 | def collective_rpc( |
| 1034 | self, |
| 1035 | method: str | Callable[..., _R], |
| 1036 | timeout: float | None = None, |
| 1037 | args: tuple = (), |
| 1038 | kwargs: dict[str, Any] | None = None, |
| 1039 | ) -> list[_R]: |
| 1040 | return self.model_executor.collective_rpc(method, timeout, args, kwargs) |
| 1041 | |
| 1042 | def set_weight_version(self, weight_version: str) -> None: |
| 1043 | self._weight_version = weight_version |
| 1044 | |
| 1045 | def get_weight_version(self) -> str: |
| 1046 | """Return the latest committed weight version.""" |
| 1047 | return self._weight_version |
| 1048 | |
| 1049 | def preprocess_add_request(self, request: EngineCoreRequest) -> tuple[Request, int]: |
| 1050 | """Preprocess the request. |
| 1051 | |
| 1052 | This function could be directly used in input processing thread to allow |
| 1053 | request initialization running in parallel with Model forward |
| 1054 | """ |
| 1055 | # Note on thread safety: no race condition. |
| 1056 | # `mm_receiver_cache` is reset at the end of LLMEngine init, |
| 1057 | # and will only be accessed in the input processing thread afterwards. |
| 1058 | if self.mm_receiver_cache is not None and request.mm_features: |
| 1059 | request.mm_features = self.mm_receiver_cache.get_and_update_features( |
| 1060 | request.mm_features |
| 1061 | ) |
| 1062 | |
| 1063 | req = Request.from_engine_core_request(request, self.request_block_hasher) |
| 1064 | if req.use_structured_output: |
| 1065 | # Note on thread safety: no race condition. |
| 1066 | # `grammar_init` is only invoked in input processing thread. For |
| 1067 | # `structured_output_manager`, each request is independent and |
| 1068 | # grammar compilation is async. Scheduler always checks grammar |
| 1069 | # compilation status before scheduling request. |
| 1070 | self.structured_output_manager.grammar_init(req) |
| 1071 | return req, request.current_wave |
| 1072 | |
| 1073 | def _eep_scale_up_before_kv_init(self): |
| 1074 | raise NotImplementedError |
| 1075 | |
| 1076 | def _eep_send_engine_core_notification( |
| 1077 | self, notification_type: EEPNotificationType |
| 1078 | ): |
| 1079 | raise NotImplementedError |
| 1080 | |
| 1081 | |
| 1082 | class EngineShutdownState(IntEnum): |
| 1083 | RUNNING = 0 |
| 1084 | REQUESTED = 1 |
| 1085 | SHUTTING_DOWN = 2 |
| 1086 | |
| 1087 | |
| 1088 | class EngineCoreProc(EngineCore): |
| 1089 | """ZMQ-wrapper for running EngineCore in background process.""" |
| 1090 | |
| 1091 | ENGINE_CORE_DEAD = b"ENGINE_CORE_DEAD" |
| 1092 | addresses: EngineZmqAddresses |
| 1093 | |
| 1094 | @instrument(span_name="EngineCoreProc init") |
| 1095 | def __init__( |
| 1096 | self, |
| 1097 | vllm_config: VllmConfig, |
| 1098 | local_client: bool, |
| 1099 | handshake_address: str, |
| 1100 | executor_class: type[Executor], |
| 1101 | log_stats: bool, |
| 1102 | client_handshake_address: str | None = None, |
| 1103 | tensor_queue: Queue | None = None, |
| 1104 | *, |
| 1105 | engine_index: int = 0, |
| 1106 | ): |
| 1107 | self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]() |
| 1108 | self.output_queue = queue.Queue[tuple[int, EngineCoreOutputs] | bytes]() |
| 1109 | executor_fail_callback = lambda: self.input_queue.put_nowait( |
| 1110 | (EngineCoreRequestType.EXECUTOR_FAILED, b"") |
| 1111 | ) |
| 1112 | |
| 1113 | self.engine_index = engine_index |
| 1114 | identity = self.engine_index.to_bytes(length=2, byteorder="little") |
| 1115 | self.engines_running = False |
| 1116 | self.shutdown_state = EngineShutdownState.RUNNING |
| 1117 | |
| 1118 | # Receiver for tensor IPC |
| 1119 | self.tensor_ipc_receiver: TensorIpcReceiver | None = None |
| 1120 | if tensor_queue is not None: |
| 1121 | self.tensor_ipc_receiver = TensorIpcReceiver(tensor_queue) |
| 1122 | logger.info("Using tensor IPC queue for multimodal tensor sharing") |
| 1123 | |
| 1124 | with self._perform_handshakes( |
| 1125 | handshake_address, |
| 1126 | identity, |
| 1127 | local_client, |
| 1128 | vllm_config, |
| 1129 | client_handshake_address, |
| 1130 | ) as addresses: |
| 1131 | # Set up data parallel environment. |
| 1132 | self.has_coordinator = addresses.coordinator_output is not None |
| 1133 | self.frontend_stats_publish_address = ( |
| 1134 | addresses.frontend_stats_publish_address |
| 1135 | ) |
| 1136 | logger.debug( |
| 1137 | "Has DP Coordinator: %s, stats publish address: %s", |
| 1138 | self.has_coordinator, |
| 1139 | self.frontend_stats_publish_address, |
| 1140 | ) |
| 1141 | internal_dp_balancing = ( |
| 1142 | self.has_coordinator |
| 1143 | and not vllm_config.parallel_config.data_parallel_external_lb |
| 1144 | ) |
| 1145 | # Only publish request queue stats to coordinator for "internal" |
| 1146 | # and "hybrid" LB modes. |
| 1147 | self.publish_dp_lb_stats = internal_dp_balancing |
| 1148 | self.last_counts = (0, 0) |
| 1149 | |
| 1150 | self.addresses = addresses |
| 1151 | self.process_input_queue_block = True |
| 1152 | self._init_data_parallel(vllm_config) |
| 1153 | |
| 1154 | super().__init__( |
| 1155 | vllm_config, |
| 1156 | executor_class, |
| 1157 | log_stats, |
| 1158 | executor_fail_callback, |
| 1159 | internal_dp_balancing, |
| 1160 | ) |
| 1161 | |
| 1162 | # Initialize fault tolerance settings. |
| 1163 | self.enable_fault_tolerance = ( |
| 1164 | vllm_config.parallel_config.enable_fault_tolerance |
| 1165 | ) |
| 1166 | if self.enable_fault_tolerance: |
| 1167 | self.ft_sentinel = EngineCoreSentinel( |
| 1168 | engine=self, |
| 1169 | parallel_config=vllm_config.parallel_config, |
| 1170 | ) |
| 1171 | |
| 1172 | # Background Threads and Queues for IO. These enable us to |
| 1173 | # overlap ZMQ socket IO with GPU since they release the GIL, |
| 1174 | # and to overlap some serialization/deserialization with the |
| 1175 | # model forward pass. |
| 1176 | # Threads handle Socket <-> Queues and core_busy_loop uses Queue. |
| 1177 | ready_event = threading.Event() |
| 1178 | input_thread = threading.Thread( |
| 1179 | target=self.process_input_sockets, |
| 1180 | args=( |
| 1181 | addresses.inputs, |
| 1182 | addresses.coordinator_input, |
| 1183 | identity, |
| 1184 | ready_event, |
| 1185 | ), |
| 1186 | daemon=True, |
| 1187 | ) |
| 1188 | input_thread.start() |
| 1189 | |
| 1190 | self.output_thread = threading.Thread( |
| 1191 | target=self.process_output_sockets, |
| 1192 | args=( |
| 1193 | addresses.outputs, |
| 1194 | addresses.coordinator_output, |
| 1195 | self.engine_index, |
| 1196 | ), |
| 1197 | daemon=True, |
| 1198 | ) |
| 1199 | self.output_thread.start() |
| 1200 | |
| 1201 | # Don't complete handshake until DP coordinator ready message is |
| 1202 | # received. |
| 1203 | while not ready_event.wait(timeout=10): |
| 1204 | if not input_thread.is_alive(): |
| 1205 | raise RuntimeError("Input socket thread died during startup") |
| 1206 | assert addresses.coordinator_input is not None |
| 1207 | logger.info("Waiting for READY message from DP Coordinator...") |
| 1208 | |
| 1209 | @contextmanager |
| 1210 | def _perform_handshakes( |
| 1211 | self, |
| 1212 | handshake_address: str, |
| 1213 | identity: bytes, |
| 1214 | local_client: bool, |
| 1215 | vllm_config: VllmConfig, |
| 1216 | client_handshake_address: str | None, |
| 1217 | ) -> Generator[EngineZmqAddresses, None, None]: |
| 1218 | """Perform startup handshakes. |
| 1219 | |
| 1220 | For DP=1 or offline mode, this is with the colocated front-end process. |
| 1221 | |
| 1222 | For DP>1 with internal load-balancing this is with the shared front-end |
| 1223 | process which may reside on a different node. |
| 1224 | |
| 1225 | For DP>1 with external or hybrid load-balancing, two handshakes are |
| 1226 | performed: |
| 1227 | - With the rank 0 front-end process which retrieves the |
| 1228 | DP Coordinator ZMQ addresses and DP process group address. |
| 1229 | - With the colocated front-end process which retrieves the |
| 1230 | client input/output socket addresses. |
| 1231 | with the exception of the rank 0 and colocated engines themselves which |
| 1232 | don't require the second handshake. |
| 1233 | |
| 1234 | Here, "front-end" process can mean the process containing the engine |
| 1235 | core client (which is the API server process in the case the API |
| 1236 | server is not scaled out), OR the launcher process running the |
| 1237 | run_multi_api_server() function in serve.py. |
| 1238 | """ |
| 1239 | input_ctx = zmq.Context() |
| 1240 | is_local = local_client and client_handshake_address is None |
| 1241 | headless = not local_client |
| 1242 | handshake = self._perform_handshake( |
| 1243 | input_ctx, |
| 1244 | handshake_address, |
| 1245 | identity, |
| 1246 | is_local, |
| 1247 | headless, |
| 1248 | vllm_config, |
| 1249 | vllm_config.parallel_config, |
| 1250 | ) |
| 1251 | if client_handshake_address is None: |
| 1252 | # We only need to handshake with one party. |
| 1253 | with handshake as addresses: |
| 1254 | yield addresses |
| 1255 | else: |
| 1256 | # We need to handshake with rank 0 front-end and our colocated frontend. |
| 1257 | assert local_client |
| 1258 | local_handshake = self._perform_handshake( |
| 1259 | input_ctx, client_handshake_address, identity, True, False, vllm_config |
| 1260 | ) |
| 1261 | with handshake as addresses, local_handshake as client_addresses: |
| 1262 | # 1. Obtain DP Coordinator zmq address and DP process group address |
| 1263 | # (addresses). |
| 1264 | # 2. Add front-end input/output addresses from colocated front-end |
| 1265 | # (client_addresses). |
| 1266 | addresses.inputs = client_addresses.inputs |
| 1267 | addresses.outputs = client_addresses.outputs |
| 1268 | yield addresses |
| 1269 | |
| 1270 | # Update config which may have changed from the handshake |
| 1271 | vllm_config.__post_init__() |
| 1272 | |
| 1273 | @contextmanager |
| 1274 | def _perform_handshake( |
| 1275 | self, |
| 1276 | ctx: zmq.Context, |
| 1277 | handshake_address: str, |
| 1278 | identity: bytes, |
| 1279 | local_client: bool, |
| 1280 | headless: bool, |
| 1281 | vllm_config: VllmConfig, |
| 1282 | parallel_config_to_update: ParallelConfig | None = None, |
| 1283 | ) -> Generator[EngineZmqAddresses, None, None]: |
| 1284 | with make_zmq_socket( |
| 1285 | ctx, |
| 1286 | handshake_address, |
| 1287 | zmq.DEALER, |
| 1288 | identity=identity, |
| 1289 | linger=5000, |
| 1290 | bind=False, |
| 1291 | ) as handshake_socket: |
| 1292 | # Register engine with front-end. |
| 1293 | addresses = self.startup_handshake( |
| 1294 | handshake_socket, local_client, headless, parallel_config_to_update |
| 1295 | ) |
| 1296 | yield addresses |
| 1297 | |
| 1298 | # Send ready message. |
| 1299 | ready_msg = { |
| 1300 | "status": "READY", |
| 1301 | "local": local_client, |
| 1302 | "headless": headless, |
| 1303 | } |
| 1304 | # Include config hash for DP configuration validation |
| 1305 | if vllm_config.parallel_config.data_parallel_size > 1: |
| 1306 | ready_msg["parallel_config_hash"] = ( |
| 1307 | vllm_config.parallel_config.compute_hash() |
| 1308 | ) |
| 1309 | |
| 1310 | handshake_socket.send(msgspec.msgpack.encode(ready_msg)) |
| 1311 | |
| 1312 | @staticmethod |
| 1313 | def startup_handshake( |
| 1314 | handshake_socket: zmq.Socket, |
| 1315 | local_client: bool, |
| 1316 | headless: bool, |
| 1317 | parallel_config: ParallelConfig | None = None, |
| 1318 | ) -> EngineZmqAddresses: |
| 1319 | # Send registration message. |
| 1320 | handshake_socket.send( |
| 1321 | msgspec.msgpack.encode( |
| 1322 | { |
| 1323 | "status": "HELLO", |
| 1324 | "local": local_client, |
| 1325 | "headless": headless, |
| 1326 | } |
| 1327 | ) |
| 1328 | ) |
| 1329 | |
| 1330 | # Receive initialization message. |
| 1331 | logger.debug("Waiting for init message from front-end.") |
| 1332 | if not handshake_socket.poll(timeout=HANDSHAKE_TIMEOUT_MINS * 60_000): |
| 1333 | raise RuntimeError( |
| 1334 | "Did not receive response from front-end " |
| 1335 | f"process within {HANDSHAKE_TIMEOUT_MINS} " |
| 1336 | f"minutes" |
| 1337 | ) |
| 1338 | init_bytes = handshake_socket.recv() |
| 1339 | init_message: EngineHandshakeMetadata = msgspec.msgpack.decode( |
| 1340 | init_bytes, type=EngineHandshakeMetadata |
| 1341 | ) |
| 1342 | logger.debug("Received init message: %s", init_message) |
| 1343 | |
| 1344 | if parallel_config is not None: |
| 1345 | for key, value in init_message.parallel_config.items(): |
| 1346 | setattr(parallel_config, key, value) |
| 1347 | |
| 1348 | return init_message.addresses |
| 1349 | |
| 1350 | @staticmethod |
| 1351 | def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs): |
| 1352 | """Launch EngineCore busy loop in background process.""" |
| 1353 | vllm_config: VllmConfig = kwargs["vllm_config"] |
| 1354 | if logging_config := getattr(vllm_config, "logging_config", None): |
| 1355 | configure_logging(logging_config) |
| 1356 | |
| 1357 | # Ensure we can serialize transformer config after spawning |
| 1358 | maybe_register_config_serialize_by_value() |
| 1359 | |
| 1360 | engine_core: EngineCoreProc | None = None |
| 1361 | signal_callback: SignalCallback | None = None |
| 1362 | clean_shutdown = False |
| 1363 | try: |
| 1364 | parallel_config: ParallelConfig = vllm_config.parallel_config |
| 1365 | data_parallel = parallel_config.data_parallel_size > 1 or dp_rank > 0 |
| 1366 | if data_parallel: |
| 1367 | parallel_config.data_parallel_rank_local = local_dp_rank |
| 1368 | process_title = f"EngineCore_DP{dp_rank}" |
| 1369 | else: |
| 1370 | process_title = "EngineCore" |
| 1371 | set_process_title(process_title) |
| 1372 | maybe_init_worker_tracer("vllm.engine_core", "engine_core", process_title) |
| 1373 | decorate_logs() |
| 1374 | if parallel_config.numa_bind: |
| 1375 | numa_utils.log_current_affinity_state(process_title) |
| 1376 | |
| 1377 | if data_parallel and vllm_config.kv_transfer_config is not None: |
| 1378 | # modify the engine_id and append the dp_rank to it to ensure |
| 1379 | # that the kv_transfer_config is unique for each DP rank. |
| 1380 | vllm_config.kv_transfer_config.engine_id = ( |
| 1381 | f"{vllm_config.kv_transfer_config.engine_id}_dp{dp_rank}" |
| 1382 | ) |
| 1383 | logger.debug( |
| 1384 | "Setting kv_transfer_config.engine_id to %s", |
| 1385 | vllm_config.kv_transfer_config.engine_id, |
| 1386 | ) |
| 1387 | |
| 1388 | parallel_config.data_parallel_index = dp_rank |
| 1389 | if data_parallel and vllm_config.model_config.is_moe: |
| 1390 | # Set data parallel rank for this engine process. |
| 1391 | parallel_config.data_parallel_rank = dp_rank |
| 1392 | engine_core = DPEngineCoreProc(*args, **kwargs) |
| 1393 | else: |
| 1394 | # Non-MoE DP ranks are completely independent, so treat like DP=1. |
| 1395 | # Note that parallel_config.data_parallel_index will still reflect |
| 1396 | # the original DP rank. |
| 1397 | parallel_config.reconfigure_for_independent_dp_rank() |
| 1398 | engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs) |
| 1399 | |
| 1400 | assert engine_core is not None |
| 1401 | |
| 1402 | def wakeup_engine(): |
| 1403 | # Wakes up idle engine via input_queue when shutdown is requested |
| 1404 | # Not safe in a signal handler - we may interrupt the main thread |
| 1405 | # while it is holding the non-reentrant input_queue.mutex |
| 1406 | engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None)) |
| 1407 | |
| 1408 | signal_callback = SignalCallback(wakeup_engine) |
| 1409 | |
| 1410 | def signal_handler(signum, frame): |
| 1411 | signal_name = signal.Signals(signum).name |
| 1412 | logger.info( |
| 1413 | "[shutdown] EngineCore: trigger received signal=%s", |
| 1414 | signal_name, |
| 1415 | ) |
| 1416 | engine_core.shutdown_state = EngineShutdownState.REQUESTED |
| 1417 | signal_callback.trigger() |
| 1418 | |
| 1419 | signal.signal(signal.SIGTERM, signal_handler) |
| 1420 | signal.signal(signal.SIGINT, signal_handler) |
| 1421 | |
| 1422 | engine_core.run_busy_loop() |
| 1423 | |
| 1424 | except SystemExit as e: |
| 1425 | logger.info_once("[shutdown] EngineCore: exiting busy loop") |
| 1426 | clean_shutdown = ( |
| 1427 | e.code in (None, 0) |
| 1428 | and engine_core is not None |
| 1429 | and engine_core.shutdown_state == EngineShutdownState.SHUTTING_DOWN |
| 1430 | and not engine_core.has_work() |
| 1431 | and engine_core.vllm_config.shutdown_timeout == 0 |
| 1432 | ) |
| 1433 | raise |
| 1434 | except Exception as e: |
| 1435 | if engine_core is None: |
| 1436 | logger.exception("EngineCore failed to start.") |
| 1437 | else: |
| 1438 | logger.exception("EngineCore encountered a fatal error.") |
| 1439 | engine_core._send_engine_dead() |
| 1440 | raise e |
| 1441 | finally: |
| 1442 | signal.signal(signal.SIGTERM, signal.SIG_DFL) |
| 1443 | signal.signal(signal.SIGINT, signal.SIG_DFL) |
| 1444 | if signal_callback is not None: |
| 1445 | signal_callback.stop() |
| 1446 | if engine_core is not None: |
| 1447 | engine_core.shutdown() |
| 1448 | if clean_shutdown: |
| 1449 | from vllm.platforms import current_platform |
| 1450 | |
| 1451 | if current_platform.is_rocm(): |
| 1452 | # Cleanup above already unfreezes and collects the heap. |
| 1453 | # Freeze the surviving graph to skip another slow cyclic-GC |
| 1454 | # scan during finalization; process exit reclaims it. |
| 1455 | gc.freeze() |
| 1456 | |
| 1457 | def _init_data_parallel(self, vllm_config: VllmConfig): |
| 1458 | pass |
| 1459 | |
| 1460 | def has_work(self) -> bool: |
| 1461 | """Returns true if the engine should be stepped.""" |
| 1462 | return ( |
| 1463 | self.engines_running |
| 1464 | or self.scheduler.has_requests() |
| 1465 | or bool(self.batch_queue) |
| 1466 | ) |
| 1467 | |
| 1468 | def is_running(self) -> bool: |
| 1469 | """Returns true if shutdown has not been requested.""" |
| 1470 | return self.shutdown_state == EngineShutdownState.RUNNING |
| 1471 | |
| 1472 | @fault_tolerant_wrapper |
| 1473 | def run_busy_loop(self): |
| 1474 | """Core busy loop of the EngineCore.""" |
| 1475 | while self._handle_shutdown(): |
| 1476 | # 1) Poll the input queue until there is work to do. |
| 1477 | self._process_input_queue() |
| 1478 | # Publish request counts before and after GPU step to ensure freshness. |
| 1479 | self._maybe_publish_request_counts() |
| 1480 | # 2) Step the engine core and return the outputs. |
| 1481 | self._process_engine_step() |
| 1482 | self._maybe_publish_request_counts() |
| 1483 | |
| 1484 | raise SystemExit |
| 1485 | |
| 1486 | def _maybe_publish_request_counts(self): |
| 1487 | if not self.publish_dp_lb_stats: |
| 1488 | return |
| 1489 | |
| 1490 | # Publish our request counts (if they've changed). |
| 1491 | counts = self.scheduler.get_request_counts() |
| 1492 | if counts != self.last_counts: |
| 1493 | self.last_counts = counts |
| 1494 | stats = SchedulerStats( |
| 1495 | *counts, kv_cache_usage=self.scheduler.get_kv_cache_usage() |
| 1496 | ) |
| 1497 | self.output_queue.put_nowait((-1, EngineCoreOutputs(scheduler_stats=stats))) |
| 1498 | |
| 1499 | def _process_input_queue(self): |
| 1500 | """Exits when an engine step needs to be performed.""" |
| 1501 | waited = False |
| 1502 | while not self.has_work() and self.is_running(): |
| 1503 | # Notify callbacks waiting for engine to become idle. |
| 1504 | self._notify_idle_state_callbacks() |
| 1505 | if self.input_queue.empty(): |
| 1506 | # Drain aborts queue; all aborts are also processed via input_queue. |
| 1507 | with self.aborts_queue.mutex: |
| 1508 | self.aborts_queue.queue.clear() |
| 1509 | if logger.isEnabledFor(DEBUG): |
| 1510 | logger.debug("EngineCore waiting for work.") |
| 1511 | waited = True |
| 1512 | block = self.process_input_queue_block |
| 1513 | try: |
| 1514 | req = self.input_queue.get(block=block) |
| 1515 | self._handle_client_request(*req) |
| 1516 | except queue.Empty: |
| 1517 | break |
| 1518 | if not block: |
| 1519 | break |
| 1520 | |
| 1521 | if waited: |
| 1522 | logger.debug("EngineCore loop active.") |
| 1523 | |
| 1524 | # Handle any more client requests. |
| 1525 | while not self.input_queue.empty(): |
| 1526 | req = self.input_queue.get_nowait() |
| 1527 | self._handle_client_request(*req) |
| 1528 | |
| 1529 | def _process_engine_step(self) -> bool: |
| 1530 | """Called only when there are unfinished local requests.""" |
| 1531 | # Step the engine core. |
| 1532 | outputs, model_executed = self.step_fn() |
| 1533 | # Put EngineCoreOutputs into the output queue. |
| 1534 | for output in outputs.items() if outputs else (): |
| 1535 | self.output_queue.put_nowait(output) |
| 1536 | # Post-step hook. |
| 1537 | self.post_step(model_executed) |
| 1538 | |
| 1539 | # If no model execution happened but there is still scheduler work |
| 1540 | # (e.g. WAITING_FOR_REMOTE_KVS or delayed KV connector frees), yield |
| 1541 | # the GIL briefly to allow background transfer threads to make progress. |
| 1542 | if not model_executed and self.scheduler.has_requests(): |
| 1543 | time.sleep(0.001) |
| 1544 | |
| 1545 | return model_executed |
| 1546 | |
| 1547 | def _notify_idle_state_callbacks(self) -> None: |
| 1548 | while self._idle_state_callbacks: |
| 1549 | callback = self._idle_state_callbacks.pop() |
| 1550 | callback(self) |
| 1551 | |
| 1552 | def _handle_shutdown(self) -> bool: |
| 1553 | # Check if shutdown was requested and handle it |
| 1554 | if self.shutdown_state == EngineShutdownState.RUNNING: |
| 1555 | return True |
| 1556 | |
| 1557 | if self.shutdown_state == EngineShutdownState.REQUESTED: |
| 1558 | shutdown_timeout = self.vllm_config.shutdown_timeout |
| 1559 | mode = "abort" if shutdown_timeout == 0 else "drain" |
| 1560 | |
| 1561 | logger.info( |
| 1562 | "[shutdown] EngineCore: start mode=%s timeout=%ds", |
| 1563 | mode, |
| 1564 | shutdown_timeout, |
| 1565 | ) |
| 1566 | |
| 1567 | if shutdown_timeout == 0: |
| 1568 | num_requests = self.scheduler.get_num_unfinished_requests() |
| 1569 | if num_requests > 0: |
| 1570 | logger.info( |
| 1571 | "[shutdown] EngineCore: aborting in-flight requests count=%d", |
| 1572 | num_requests, |
| 1573 | ) |
| 1574 | aborted_reqs = self.scheduler.finish_requests( |
| 1575 | None, RequestStatus.FINISHED_ABORTED |
| 1576 | ) |
| 1577 | self._send_abort_outputs(aborted_reqs) |
| 1578 | else: |
| 1579 | num_requests = self.scheduler.get_num_unfinished_requests() |
| 1580 | if num_requests > 0: |
| 1581 | logger.info( |
| 1582 | "[shutdown] EngineCore: draining in-flight requests " |
| 1583 | "count=%d timeout=%ds", |
| 1584 | num_requests, |
| 1585 | shutdown_timeout, |
| 1586 | ) |
| 1587 | |
| 1588 | self.shutdown_state = EngineShutdownState.SHUTTING_DOWN |
| 1589 | |
| 1590 | # Exit when no work remaining |
| 1591 | if not self.has_work(): |
| 1592 | logger.info( |
| 1593 | "[shutdown] EngineCore: request processing complete; " |
| 1594 | "starting resource teardown" |
| 1595 | ) |
| 1596 | return False |
| 1597 | |
| 1598 | return True |
| 1599 | |
| 1600 | def _handle_client_request( |
| 1601 | self, request_type: EngineCoreRequestType, request: Any |
| 1602 | ) -> None: |
| 1603 | """Dispatch request from client.""" |
| 1604 | if request_type == EngineCoreRequestType.WAKEUP: |
| 1605 | return |
| 1606 | elif request_type == EngineCoreRequestType.ADD: |
| 1607 | req, request_wave = request |
| 1608 | if self._reject_add_in_shutdown(req): |
| 1609 | return |
| 1610 | self.add_request(req, request_wave) |
| 1611 | elif request_type == EngineCoreRequestType.ABORT: |
| 1612 | self.abort_requests(request) |
| 1613 | elif request_type == EngineCoreRequestType.UTILITY: |
| 1614 | client_idx, call_id, method_name, args = request |
| 1615 | if self._reject_utility_in_shutdown(client_idx, call_id, method_name): |
| 1616 | return |
| 1617 | output = UtilityOutput(call_id) |
| 1618 | # Lazily look-up utility method so that failure will be handled/returned. |
| 1619 | get_result = lambda: ( |
| 1620 | (method := getattr(self, method_name)) |
| 1621 | and method(*self._convert_msgspec_args(method, args)) |
| 1622 | ) |
| 1623 | enqueue_output = lambda out: self.output_queue.put_nowait( |
| 1624 | (client_idx, EngineCoreOutputs(utility_output=out)) |
| 1625 | ) |
| 1626 | self._invoke_utility_method(method_name, get_result, output, enqueue_output) |
| 1627 | elif request_type == EngineCoreRequestType.EXECUTOR_FAILED: |
| 1628 | raise RuntimeError("Executor failed.") |
| 1629 | else: |
| 1630 | logger.error( |
| 1631 | "Unrecognized input request type encountered: %s", request_type |
| 1632 | ) |
| 1633 | |
| 1634 | def _reject_add_in_shutdown(self, request: Request) -> bool: |
| 1635 | if self.shutdown_state == EngineShutdownState.RUNNING: |
| 1636 | return False |
| 1637 | |
| 1638 | logger.debug( |
| 1639 | "[shutdown] EngineCore: rejecting new request request_id=%s", |
| 1640 | request.request_id, |
| 1641 | ) |
| 1642 | self._send_abort_outputs_to_client([request.request_id], request.client_index) |
| 1643 | return True |
| 1644 | |
| 1645 | def _reject_utility_in_shutdown( |
| 1646 | self, client_idx: int, call_id: int, method_name: str |
| 1647 | ) -> bool: |
| 1648 | if self.shutdown_state == EngineShutdownState.RUNNING: |
| 1649 | return False |
| 1650 | |
| 1651 | logger.warning( |
| 1652 | "[shutdown] EngineCore: rejecting utility call method=%s", |
| 1653 | method_name, |
| 1654 | ) |
| 1655 | output = UtilityOutput(call_id, failure_message="Server shutting down") |
| 1656 | self.output_queue.put_nowait( |
| 1657 | (client_idx, EngineCoreOutputs(utility_output=output)) |
| 1658 | ) |
| 1659 | return True |
| 1660 | |
| 1661 | @staticmethod |
| 1662 | def _invoke_utility_method( |
| 1663 | name: str, get_result: Callable, output: UtilityOutput, enqueue_output: Callable |
| 1664 | ): |
| 1665 | try: |
| 1666 | result = get_result() |
| 1667 | if isinstance(result, Future): |
| 1668 | # Defer utility output handling until future completion. |
| 1669 | callback = lambda future: EngineCoreProc._invoke_utility_method( |
| 1670 | name, future.result, output, enqueue_output |
| 1671 | ) |
| 1672 | result.add_done_callback(callback) |
| 1673 | return |
| 1674 | output.result = UtilityResult(result) |
| 1675 | except Exception as e: |
| 1676 | logger.exception("Invocation of %s method failed", name) |
| 1677 | output.failure_message = f"Call to {name} method failed: {str(e)}" |
| 1678 | enqueue_output(output) |
| 1679 | |
| 1680 | @staticmethod |
| 1681 | def _convert_msgspec_args(method, args): |
| 1682 | """If a provided arg type doesn't match corresponding target method |
| 1683 | arg type, try converting to msgspec object.""" |
| 1684 | if not args: |
| 1685 | return args |
| 1686 | arg_types = signature(method).parameters.values() |
| 1687 | assert len(args) <= len(arg_types) |
| 1688 | return tuple( |
| 1689 | msgspec.convert(v, type=p.annotation) |
| 1690 | if isclass(p.annotation) |
| 1691 | and issubclass(p.annotation, msgspec.Struct) |
| 1692 | and not isinstance(v, p.annotation) |
| 1693 | else v |
| 1694 | for v, p in zip(args, arg_types) |
| 1695 | ) |
| 1696 | |
| 1697 | def _send_engine_dead(self): |
| 1698 | """Send EngineDead status to the EngineCoreClient.""" |
| 1699 | # Put ENGINE_CORE_DEAD in the queue. |
| 1700 | self.output_queue.put_nowait(EngineCoreProc.ENGINE_CORE_DEAD) |
| 1701 | |
| 1702 | # Wait until msg sent by the daemon before shutdown. |
| 1703 | self.output_thread.join(timeout=5.0) |
| 1704 | if self.output_thread.is_alive(): |
| 1705 | logger.fatal( |
| 1706 | "vLLM shutdown signal from EngineCore failed " |
| 1707 | "to send. Please report this issue." |
| 1708 | ) |
| 1709 | |
| 1710 | def _make_ready_response(self) -> EngineCoreReadyResponse: |
| 1711 | parallel_config = self.vllm_config.parallel_config |
| 1712 | scheduler_config = self.vllm_config.scheduler_config |
| 1713 | return EngineCoreReadyResponse( |
| 1714 | max_model_len=self.vllm_config.model_config.max_model_len, |
| 1715 | num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, |
| 1716 | block_size=self.vllm_config.cache_config.block_size, |
| 1717 | mamba_block_size=self.vllm_config.cache_config.mamba_block_size, |
| 1718 | effective_attention_block_size=( |
| 1719 | self.vllm_config.cache_config.effective_attention_block_size |
| 1720 | ), |
| 1721 | dp_stats_address=self.frontend_stats_publish_address, |
| 1722 | dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), |
| 1723 | vllm_version=VLLM_VERSION, |
| 1724 | world_size=self.vllm_config.parallel_config.world_size, |
| 1725 | data_parallel_size=parallel_config.data_parallel_size, |
| 1726 | kv_cache_size_tokens=self.vllm_config.cache_config.kv_cache_size_tokens, |
| 1727 | kv_cache_max_concurrency=( |
| 1728 | self.vllm_config.cache_config.kv_cache_max_concurrency |
| 1729 | ), |
| 1730 | tensor_parallel_size=parallel_config.tensor_parallel_size, |
| 1731 | pipeline_parallel_size=parallel_config.pipeline_parallel_size, |
| 1732 | decode_context_parallel_size=parallel_config.decode_context_parallel_size, |
| 1733 | data_parallel_rank=self.engine_index, |
| 1734 | max_num_seqs=scheduler_config.max_num_seqs, |
| 1735 | max_num_batched_tokens=scheduler_config.max_num_batched_tokens, |
| 1736 | instance_id=self.vllm_config.instance_id, |
| 1737 | supports_lora=self.vllm_config.lora_config is not None, |
| 1738 | max_loras=( |
| 1739 | self.vllm_config.lora_config.max_loras |
| 1740 | if self.vllm_config.lora_config is not None |
| 1741 | else 0 |
| 1742 | ), |
| 1743 | kv_events_config=self.scheduler.get_kv_event_publisher_config(), |
| 1744 | weight_transfer_backend=( |
| 1745 | self.vllm_config.weight_transfer_config.backend |
| 1746 | if self.vllm_config.weight_transfer_config is not None |
| 1747 | else None |
| 1748 | ), |
| 1749 | enable_sleep_mode=self.vllm_config.model_config.enable_sleep_mode, |
| 1750 | supports_draft_weight_updates=( |
| 1751 | self.model_executor.supports_draft_weight_updates() |
| 1752 | if self.vllm_config.weight_transfer_config is not None |
| 1753 | else False |
| 1754 | ), |
| 1755 | ) |
| 1756 | |
| 1757 | def process_input_sockets( |
| 1758 | self, |
| 1759 | input_addresses: list[str], |
| 1760 | coord_input_address: str | None, |
| 1761 | identity: bytes, |
| 1762 | ready_event: threading.Event, |
| 1763 | ): |
| 1764 | """Input socket IO thread.""" |
| 1765 | # Msgpack serialization decoding with optional tensor IPC receiver. |
| 1766 | add_request_decoder = MsgpackDecoder( |
| 1767 | EngineCoreRequest, oob_tensor_provider=self.tensor_ipc_receiver |
| 1768 | ) |
| 1769 | generic_decoder = MsgpackDecoder(oob_tensor_provider=self.tensor_ipc_receiver) |
| 1770 | |
| 1771 | with ExitStack() as stack, zmq.Context() as ctx: |
| 1772 | input_sockets = [ |
| 1773 | stack.enter_context( |
| 1774 | make_zmq_socket( |
| 1775 | ctx, input_address, zmq.DEALER, identity=identity, bind=False |
| 1776 | ) |
| 1777 | ) |
| 1778 | for input_address in input_addresses |
| 1779 | ] |
| 1780 | if coord_input_address is None: |
| 1781 | coord_socket = None |
| 1782 | else: |
| 1783 | coord_socket = stack.enter_context( |
| 1784 | make_zmq_socket( |
| 1785 | ctx, |
| 1786 | coord_input_address, |
| 1787 | zmq.XSUB, |
| 1788 | identity=identity, |
| 1789 | bind=False, |
| 1790 | ) |
| 1791 | ) |
| 1792 | # Send subscription message to coordinator. |
| 1793 | coord_socket.send(b"\x01") |
| 1794 | |
| 1795 | # Register sockets with poller. |
| 1796 | poller = zmq.Poller() |
| 1797 | ready_response = self._make_ready_response() |
| 1798 | ready_payload = msgspec.msgpack.encode(ready_response) |
| 1799 | for input_socket in input_sockets: |
| 1800 | # Send initial message to each input socket - this is required |
| 1801 | # before the front-end ROUTER socket can send input messages |
| 1802 | # back to us. |
| 1803 | input_socket.send(ready_payload) |
| 1804 | poller.register(input_socket, zmq.POLLIN) |
| 1805 | |
| 1806 | if coord_socket is not None: |
| 1807 | # Wait for ready message from coordinator. |
| 1808 | assert coord_socket.recv() == b"READY" |
| 1809 | poller.register(coord_socket, zmq.POLLIN) |
| 1810 | |
| 1811 | ready_event.set() |
| 1812 | del ready_event |
| 1813 | while True: |
| 1814 | for input_socket, _ in poller.poll(): |
| 1815 | # (RequestType, RequestData) |
| 1816 | type_frame, *data_frames = input_socket.recv_multipart(copy=False) |
| 1817 | # NOTE(yongji): ignore READY message sent by DP coordinator |
| 1818 | # that is used to notify newly started engines |
| 1819 | if type_frame.buffer == b"READY": |
| 1820 | assert input_socket == coord_socket |
| 1821 | continue |
| 1822 | request_type = EngineCoreRequestType(bytes(type_frame.buffer)) |
| 1823 | |
| 1824 | # Deserialize the request data. |
| 1825 | request: Any |
| 1826 | if request_type == EngineCoreRequestType.ADD: |
| 1827 | req: EngineCoreRequest = add_request_decoder.decode(data_frames) |
| 1828 | try: |
| 1829 | request = self.preprocess_add_request(req) |
| 1830 | except MultiModalCacheMissError as e: |
| 1831 | # P0/P1 shadow drift -- return a retryable signal (P0 |
| 1832 | # drops the stale entry, client resends with data). |
| 1833 | self._handle_mm_cache_miss(req, e) |
| 1834 | continue |
| 1835 | except Exception: |
| 1836 | self._handle_request_preproc_error(req) |
| 1837 | continue |
| 1838 | elif request_type == EngineCoreRequestType.UTILITY: |
| 1839 | request = generic_decoder.decode(data_frames) |
| 1840 | client_idx, call_id, method, args = request |
| 1841 | if method == FT_UTILITY_METHOD: |
| 1842 | self.ft_sentinel.handle_command( |
| 1843 | client_idx, call_id, args[0] |
| 1844 | ) |
| 1845 | continue |
| 1846 | else: |
| 1847 | request = generic_decoder.decode(data_frames) |
| 1848 | |
| 1849 | if request_type == EngineCoreRequestType.ABORT: |
| 1850 | # Aborts are added to *both* queues, allows us to eagerly |
| 1851 | # process aborts while also ensuring ordering in the input |
| 1852 | # queue to avoid leaking requests. This is ok because |
| 1853 | # aborting in the scheduler is idempotent. |
| 1854 | self.aborts_queue.put_nowait(request) |
| 1855 | |
| 1856 | # Push to input queue for core busy loop. |
| 1857 | self.input_queue.put_nowait((request_type, request)) |
| 1858 | |
| 1859 | def process_output_sockets( |
| 1860 | self, output_paths: list[str], coord_output_path: str | None, engine_index: int |
| 1861 | ): |
| 1862 | """Output socket IO thread.""" |
| 1863 | # Msgpack serialization encoding. |
| 1864 | encoder = MsgpackEncoder() |
| 1865 | # Send buffers to reuse. |
| 1866 | reuse_buffers: list[bytearray] = [] |
| 1867 | # Payload buffers that can't be reused yet because zmq may still be |
| 1868 | # sending them. |
| 1869 | # Buffers of the zero-copy tensor/ndarray frames don't need tracking |
| 1870 | # here: zmq itself holds a reference to each until it's done with it. |
| 1871 | pending = deque[tuple[zmq.MessageTracker, bytearray]]() |
| 1872 | |
| 1873 | # We must set linger to ensure the ENGINE_CORE_DEAD |
| 1874 | # message is sent prior to closing the socket. |
| 1875 | with ExitStack() as stack, zmq.Context() as ctx: |
| 1876 | sockets = [ |
| 1877 | stack.enter_context( |
| 1878 | make_zmq_socket(ctx, output_path, zmq.PUSH, linger=4000) |
| 1879 | ) |
| 1880 | for output_path in output_paths |
| 1881 | ] |
| 1882 | coord_socket = ( |
| 1883 | stack.enter_context( |
| 1884 | make_zmq_socket( |
| 1885 | ctx, coord_output_path, zmq.PUSH, bind=False, linger=4000 |
| 1886 | ) |
| 1887 | ) |
| 1888 | if coord_output_path is not None |
| 1889 | else None |
| 1890 | ) |
| 1891 | max_reuse_bufs = len(sockets) + 1 |
| 1892 | |
| 1893 | while True: |
| 1894 | output = self.output_queue.get() |
| 1895 | if output == EngineCoreProc.ENGINE_CORE_DEAD: |
| 1896 | for socket in sockets: |
| 1897 | socket.send(output) |
| 1898 | break |
| 1899 | assert not isinstance(output, bytes) |
| 1900 | client_index, outputs = output |
| 1901 | outputs.engine_index = engine_index |
| 1902 | |
| 1903 | if client_index == -1: |
| 1904 | # Don't reuse buffer for coordinator message |
| 1905 | # which will be very small. |
| 1906 | assert coord_socket is not None |
| 1907 | coord_socket.send_multipart(encoder.encode(outputs)) |
| 1908 | continue |
| 1909 | |
| 1910 | # Reclaim buffers that zmq is finished with. |
| 1911 | while pending and pending[-1][0].done: |
| 1912 | reclaimed = pending.pop()[1] |
| 1913 | if len(reuse_buffers) < max_reuse_bufs: |
| 1914 | reuse_buffers.append(reclaimed) |
| 1915 | |
| 1916 | buffer = reuse_buffers.pop() if reuse_buffers else bytearray() |
| 1917 | buffers = encoder.encode_into(outputs, buffer) |
| 1918 | tracker = self._send_msg_tracking_payload( |
| 1919 | sockets[client_index], buffers |
| 1920 | ) |
| 1921 | if not tracker.done: |
| 1922 | pending.appendleft((tracker, buffer)) |
| 1923 | elif len(reuse_buffers) < max_reuse_bufs: |
| 1924 | # Limit the number of buffers to reuse. |
| 1925 | reuse_buffers.append(buffer) |
| 1926 | |
| 1927 | def _handle_mm_cache_miss( |
| 1928 | self, request: EngineCoreRequest, err: MultiModalCacheMissError |
| 1929 | ) -> None: |
| 1930 | """Return a retryable response for a P0/P1 cache-drift miss. |
| 1931 | |
| 1932 | Surfaces every drifted hash via ``EngineCoreOutput.mm_cache_miss_hashes`` so |
| 1933 | the frontend drops them from its sender cache and the client resends with |
| 1934 | data (see ``MultiModalCacheMissError``). Logged at warning, not exception, |
| 1935 | because it is expected and self-healing. |
| 1936 | """ |
| 1937 | logger.warning( |
| 1938 | "Multi-modal cache miss for request %s (mm_hashes=%s): P0/P1 cache " |
| 1939 | "drift; returning a retryable response so the items are resent with data.", |
| 1940 | request.request_id, |
| 1941 | err.mm_hashes, |
| 1942 | ) |
| 1943 | self.output_queue.put_nowait( |
| 1944 | ( |
| 1945 | request.client_index, |
| 1946 | EngineCoreOutputs( |
| 1947 | engine_index=self.engine_index, |
| 1948 | finished_requests={request.request_id}, |
| 1949 | outputs=[ |
| 1950 | EngineCoreOutput( |
| 1951 | request_id=request.request_id, |
| 1952 | new_token_ids=[], |
| 1953 | finish_reason=FinishReason.ERROR, |
| 1954 | mm_cache_miss_hashes=err.mm_hashes, |
| 1955 | ) |
| 1956 | ], |
| 1957 | ), |
| 1958 | ) |
| 1959 | ) |
| 1960 | |
| 1961 | @staticmethod |
| 1962 | def _send_msg_tracking_payload( |
| 1963 | socket: zmq.Socket, buffers: Sequence[bytestr] |
| 1964 | ) -> zmq.MessageTracker: |
| 1965 | """Send `buffers` as a zero-copy multipart message, returning a tracker |
| 1966 | for the *first* frame. |
| 1967 | |
| 1968 | Used instead of `Socket.send_multipart()` because we reuse the buffer |
| 1969 | passed to `MsgpackEncoder.encode_into()`: `send_multipart()` returns a |
| 1970 | tracker for the last frame only. |
| 1971 | """ |
| 1972 | more_flag = zmq.SNDMORE if len(buffers) > 1 else 0 |
| 1973 | tracker = socket.send(buffers[0], more_flag, copy=False, track=True) |
| 1974 | if more_flag: |
| 1975 | socket.send_multipart(buffers[1:], copy=False) |
| 1976 | return tracker |
| 1977 | |
| 1978 | def _handle_request_preproc_error(self, request: EngineCoreRequest) -> None: |
| 1979 | """Log and return a request-scoped error response for exceptions raised |
| 1980 | from the add request preprocessing in the input socket processing thread. |
| 1981 | """ |
| 1982 | logger.exception( |
| 1983 | "Unexpected error pre-processing request %s", request.request_id |
| 1984 | ) |
| 1985 | self._send_error_outputs_to_client([request.request_id], request.client_index) |
| 1986 | |
| 1987 | def pause_scheduler( |
| 1988 | self, mode: PauseMode = "abort", clear_cache: bool = True |
| 1989 | ) -> Future | None: |
| 1990 | """Pause generation; behavior depends on mode. |
| 1991 | |
| 1992 | All pause modes queue new adds -- "abort" and "keep" skip step(); |
| 1993 | "wait" allows step() so in-flight requests can drain. |
| 1994 | |
| 1995 | - ``abort``: Set PAUSED_NEW, abort all requests, wait for abort |
| 1996 | outputs to be sent (when running with output_queue), optionally |
| 1997 | clear caches, then complete the returned Future. |
| 1998 | - ``wait``: Set PAUSED_NEW (queue adds, keep stepping); when drained, |
| 1999 | optionally clear caches, then complete the returned Future. |
| 2000 | - ``keep``: Set PAUSED_ALL; return a Future that completes when the |
| 2001 | output queue is empty. |
| 2002 | """ |
| 2003 | if mode not in get_args(PauseMode): |
| 2004 | raise ValueError(f"Invalid pause mode: {mode}") |
| 2005 | |
| 2006 | def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None: |
| 2007 | try: |
| 2008 | engine._finish_pause(clear_cache) |
| 2009 | except Exception as e: |
| 2010 | future.set_exception(e) |
| 2011 | else: |
| 2012 | future.set_result(None) |
| 2013 | |
| 2014 | if mode == "abort": |
| 2015 | aborted_reqs = self.scheduler.finish_requests( |
| 2016 | None, RequestStatus.FINISHED_ABORTED |
| 2017 | ) |
| 2018 | self._send_abort_outputs(aborted_reqs) |
| 2019 | |
| 2020 | pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW |
| 2021 | self.scheduler.set_pause_state(pause_state) |
| 2022 | |
| 2023 | if self._pause_complete(): |
| 2024 | self._finish_pause(clear_cache) |
| 2025 | return None |
| 2026 | |
| 2027 | future = Future[Any]() |
| 2028 | self._idle_state_callbacks.append(partial(engine_idle_callback, future=future)) |
| 2029 | return future |
| 2030 | |
| 2031 | def _pause_complete(self) -> bool: |
| 2032 | """Returns True if the pause has fully completed and the caller can |
| 2033 | return ``None`` synchronously; False if the pause is still pending |
| 2034 | and the caller should register an idle-state callback to finish it. |
| 2035 | """ |
| 2036 | return not self.has_work() |
| 2037 | |
| 2038 | def _send_finish_outputs_to_client( |
| 2039 | self, req_ids: list[str], client_index: int, finish_reason: FinishReason |
| 2040 | ) -> None: |
| 2041 | outputs = [ |
| 2042 | EngineCoreOutput(req_id, [], finish_reason=finish_reason) |
| 2043 | for req_id in req_ids |
| 2044 | ] |
| 2045 | eco = EngineCoreOutputs(finished_requests=req_ids, outputs=outputs) |
| 2046 | self.output_queue.put_nowait((client_index, eco)) |
| 2047 | |
| 2048 | def _send_abort_outputs_to_client( |
| 2049 | self, req_ids: list[str], client_index: int |
| 2050 | ) -> None: |
| 2051 | self._send_finish_outputs_to_client(req_ids, client_index, FinishReason.ABORT) |
| 2052 | |
| 2053 | def _send_error_outputs_to_client( |
| 2054 | self, req_ids: list[str], client_index: int |
| 2055 | ) -> None: |
| 2056 | self._send_finish_outputs_to_client(req_ids, client_index, FinishReason.ERROR) |
| 2057 | |
| 2058 | def _send_abort_outputs(self, aborted_reqs: list[Request]) -> None: |
| 2059 | # TODO(nick) this will be moved inside the scheduler |
| 2060 | if aborted_reqs: |
| 2061 | # Map client_index to list of request_ids that belong to that client. |
| 2062 | by_client = defaultdict[int, set[str]](set) |
| 2063 | for request in aborted_reqs: |
| 2064 | by_client[request.client_index].add(request.request_id) |
| 2065 | for client_index, req_ids in by_client.items(): |
| 2066 | self._send_abort_outputs_to_client(list(req_ids), client_index) |
| 2067 | |
| 2068 | |
| 2069 | class DPEngineCoreProc(EngineCoreProc): |
| 2070 | """ZMQ-wrapper for running EngineCore in background process |
| 2071 | in a data parallel context.""" |
| 2072 | |
| 2073 | def __init__( |
| 2074 | self, |
| 2075 | vllm_config: VllmConfig, |
| 2076 | local_client: bool, |
| 2077 | handshake_address: str, |
| 2078 | executor_class: type[Executor], |
| 2079 | log_stats: bool, |
| 2080 | client_handshake_address: str | None = None, |
| 2081 | tensor_queue: Queue | None = None, |
| 2082 | ): |
| 2083 | assert vllm_config.model_config.is_moe, ( |
| 2084 | "DPEngineCoreProc should only be used for MoE models" |
| 2085 | ) |
| 2086 | |
| 2087 | scheduler_config = vllm_config.scheduler_config |
| 2088 | self.prefill_schedule_interval = scheduler_config.prefill_schedule_interval |
| 2089 | self.dp_sync_interval = vllm_config.parallel_config.dp_sync_interval |
| 2090 | |
| 2091 | # Counts forward-passes of the model so that we can synchronize |
| 2092 | # finished with DP peers every N steps. |
| 2093 | self.step_counter = 0 |
| 2094 | self.current_wave = 0 |
| 2095 | |
| 2096 | # Two-phase pause protocol state. When pending_pause is True, the |
| 2097 | # engine keeps stepping (dummy batches) while waiting for all DP |
| 2098 | # ranks to also set pending_pause. Once all ranks agree via |
| 2099 | # all-reduce, ignore_start_dp_wave is set so that stale |
| 2100 | # START_DP_WAVE messages cannot re-wake the engines. |
| 2101 | self.pending_pause = False |
| 2102 | self.ignore_start_dp_wave = False |
| 2103 | |
| 2104 | from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState |
| 2105 | |
| 2106 | self.eep_scaling_state: ElasticEPScalingState | None = None |
| 2107 | |
| 2108 | # Initialize the engine. |
| 2109 | dp_rank = vllm_config.parallel_config.data_parallel_rank |
| 2110 | super().__init__( |
| 2111 | vllm_config, |
| 2112 | local_client, |
| 2113 | handshake_address, |
| 2114 | executor_class, |
| 2115 | log_stats, |
| 2116 | client_handshake_address, |
| 2117 | engine_index=dp_rank, |
| 2118 | tensor_queue=tensor_queue, |
| 2119 | ) |
| 2120 | |
| 2121 | def _init_data_parallel(self, vllm_config: VllmConfig): |
| 2122 | # Configure GPUs and stateless process group for data parallel. |
| 2123 | parallel_config = vllm_config.parallel_config |
| 2124 | dp_rank = parallel_config.data_parallel_rank |
| 2125 | dp_size = parallel_config.data_parallel_size |
| 2126 | local_dp_rank = parallel_config.data_parallel_rank_local |
| 2127 | |
| 2128 | assert dp_size > 1 |
| 2129 | assert local_dp_rank is not None |
| 2130 | assert 0 <= local_dp_rank <= dp_rank < dp_size |
| 2131 | |
| 2132 | self.dp_rank = dp_rank |
| 2133 | self.dp_size = dp_size |
| 2134 | dp_group, dp_store = parallel_config.stateless_init_dp_group(return_store=True) |
| 2135 | self.dp_group, self.dp_store = dp_group, dp_store |
| 2136 | |
| 2137 | def shutdown(self): |
| 2138 | super().shutdown() |
| 2139 | if dp_group := getattr(self, "dp_group", None): |
| 2140 | stateless_destroy_torch_distributed_process_group(dp_group) |
| 2141 | |
| 2142 | def _pause_complete(self) -> bool: |
| 2143 | """Two-phase DP-aware pause. |
| 2144 | |
| 2145 | Phase 1: Set local pause state and ``pending_pause`` flag. If the |
| 2146 | engines are idle, kick-start them by setting ``engines_running`` to |
| 2147 | True so ranks enter the stepping loop and reach the all-reduce |
| 2148 | consensus checkpoint in ``_has_global_unfinished_reqs``. |
| 2149 | |
| 2150 | Phase 2 (in ``_has_global_unfinished_reqs``): Once the all-reduce |
| 2151 | confirms that **all** ranks have ``pending_pause`` set, collectively |
| 2152 | stop stepping and set ``ignore_start_dp_wave`` so that stale |
| 2153 | ``START_DP_WAVE`` messages cannot re-wake any engine. |
| 2154 | """ |
| 2155 | self.pending_pause = True |
| 2156 | self.engines_running = True |
| 2157 | |
| 2158 | return False |
| 2159 | |
| 2160 | def add_request(self, request: Request, request_wave: int = 0): |
| 2161 | super().add_request(request, request_wave) |
| 2162 | if self.has_coordinator and request_wave != self.current_wave: |
| 2163 | if request_wave > self.current_wave: |
| 2164 | self.current_wave = request_wave |
| 2165 | elif ( |
| 2166 | not self.engines_running |
| 2167 | and self.scheduler.pause_state == PauseState.UNPAUSED |
| 2168 | ): |
| 2169 | # Request received for an already-completed wave, notify |
| 2170 | # front-end that we need to start the next one. |
| 2171 | self.engines_running = True |
| 2172 | self.output_queue.put_nowait( |
| 2173 | (-1, EngineCoreOutputs(start_wave=self.current_wave)) |
| 2174 | ) |
| 2175 | |
| 2176 | def resume_scheduler(self): |
| 2177 | if self.pending_pause or (self.engines_running and self.ignore_start_dp_wave): |
| 2178 | raise RuntimeError( |
| 2179 | "resume_scheduler called while pause is still in " |
| 2180 | "flight. Wait for the pause future to resolve before " |
| 2181 | "resuming." |
| 2182 | ) |
| 2183 | if self.engines_running: |
| 2184 | logger.debug("Resume called while engines are not paused, ignoring.") |
| 2185 | return |
| 2186 | |
| 2187 | super().resume_scheduler() |
| 2188 | self.ignore_start_dp_wave = False |
| 2189 | |
| 2190 | # Barrier: wait for all DP ranks to have resumed (and cleared |
| 2191 | # ignore_start_dp_wave) before any rank starts stepping. Uses |
| 2192 | # the existing all-reduce which is safe because engines are |
| 2193 | # stopped. |
| 2194 | has_global_unfinished = ParallelConfig.has_unfinished_dp( |
| 2195 | self.dp_group, self.scheduler.has_unfinished_requests() |
| 2196 | ) |
| 2197 | |
| 2198 | if has_global_unfinished: |
| 2199 | self.engines_running = True |
| 2200 | |
| 2201 | def barrier(self): |
| 2202 | """Blocking barrier on the DP process group (test-only utility).""" |
| 2203 | import torch.distributed as dist |
| 2204 | |
| 2205 | dist.barrier(group=self.dp_group) |
| 2206 | |
| 2207 | def _handle_client_request( |
| 2208 | self, request_type: EngineCoreRequestType, request: Any |
| 2209 | ) -> None: |
| 2210 | if request_type == EngineCoreRequestType.START_DP_WAVE: |
| 2211 | if self.ignore_start_dp_wave: |
| 2212 | return |
| 2213 | new_wave, exclude_eng_index = request |
| 2214 | if exclude_eng_index != self.engine_index and ( |
| 2215 | new_wave >= self.current_wave |
| 2216 | ): |
| 2217 | self.current_wave = new_wave |
| 2218 | if not self.engines_running: |
| 2219 | logger.debug( |
| 2220 | "EngineCore starting idle loop for wave %d.", |
| 2221 | new_wave, |
| 2222 | ) |
| 2223 | self.engines_running = True |
| 2224 | else: |
| 2225 | super()._handle_client_request(request_type, request) |
| 2226 | |
| 2227 | def _maybe_publish_request_counts(self): |
| 2228 | if not self.publish_dp_lb_stats: |
| 2229 | return |
| 2230 | |
| 2231 | # Publish our request counts (if they've changed), stamped with the |
| 2232 | # lockstep-synchronized step counter and wave number. |
| 2233 | counts = self.scheduler.get_request_counts() |
| 2234 | if counts != self.last_counts: |
| 2235 | self.last_counts = counts |
| 2236 | stats = SchedulerStats( |
| 2237 | *counts, |
| 2238 | kv_cache_usage=self.scheduler.get_kv_cache_usage(), |
| 2239 | step_counter=self.step_counter, |
| 2240 | current_wave=self.current_wave, |
| 2241 | ) |
| 2242 | self.output_queue.put_nowait((-1, EngineCoreOutputs(scheduler_stats=stats))) |
| 2243 | |
| 2244 | def _should_throttle_prefills(self) -> bool: |
| 2245 | # Throttle new prefills to cadence-aligned steps for DP balancing. |
| 2246 | # step_counter is identical across DP ranks. On a fresh wave the |
| 2247 | # counter is 0, so prefills are admitted immediately after idle. |
| 2248 | return ( |
| 2249 | self.prefill_schedule_interval > 1 |
| 2250 | and self.step_counter % self.prefill_schedule_interval != 0 |
| 2251 | ) |
| 2252 | |
| 2253 | @fault_tolerant_wrapper |
| 2254 | def run_busy_loop(self): |
| 2255 | """Core busy loop of the EngineCore for data parallel case.""" |
| 2256 | # Loop until process is sent a SIGINT or SIGTERM |
| 2257 | while self._handle_shutdown(): |
| 2258 | # 1) Poll the input queue until there is work to do. |
| 2259 | was_running = self.engines_running |
| 2260 | self._process_input_queue() |
| 2261 | # Publish request counts before and after GPU step to ensure freshness. |
| 2262 | self._maybe_publish_request_counts() |
| 2263 | |
| 2264 | if self.eep_scaling_state is not None: |
| 2265 | state = self.eep_scaling_state |
| 2266 | if state.commit_requested or not state.is_ready_for_switch(): |
| 2267 | state.progress() |
| 2268 | if state.is_complete(): |
| 2269 | if state.worker_type == "removing": |
| 2270 | raise SystemExit |
| 2271 | self.process_input_queue_block = True |
| 2272 | self.eep_scaling_state = None |
| 2273 | elif not state.commit_requested and state.is_ready_for_switch(): |
| 2274 | self.process_input_queue_block = True |
| 2275 | |
| 2276 | executed = self._process_engine_step() |
| 2277 | self._maybe_publish_request_counts() |
| 2278 | |
| 2279 | local_unfinished_reqs = self.scheduler.has_unfinished_requests() |
| 2280 | if not executed: |
| 2281 | if not local_unfinished_reqs and not self.engines_running: |
| 2282 | # All engines are idle. |
| 2283 | continue |
| 2284 | |
| 2285 | # Execute a dummy pass when no ready requests ran, unless the |
| 2286 | # engine is sleeping. |
| 2287 | elif not self.model_executor.is_sleeping: |
| 2288 | with self.capture_iteration_details(None) as iteration_details: |
| 2289 | self.execute_dummy_batch() |
| 2290 | if iteration_details is not None and not self.has_coordinator: |
| 2291 | stats = self._make_iteration_details_stats(iteration_details) |
| 2292 | self.output_queue.put_nowait( |
| 2293 | (0, EngineCoreOutputs(scheduler_stats=stats)) |
| 2294 | ) |
| 2295 | |
| 2296 | # 3) All-reduce operation to determine global unfinished reqs. |
| 2297 | self.engines_running = self._has_global_unfinished_reqs( |
| 2298 | local_unfinished_reqs |
| 2299 | ) |
| 2300 | |
| 2301 | if not self.engines_running: |
| 2302 | if self.dp_rank == 0 or not self.has_coordinator: |
| 2303 | # Notify client that we are pausing the loop. |
| 2304 | logger.debug( |
| 2305 | "Wave %d finished, pausing engine loop.", self.current_wave |
| 2306 | ) |
| 2307 | # In the coordinator case, dp rank 0 sends updates to the |
| 2308 | # coordinator. Otherwise (offline spmd case), each rank |
| 2309 | # sends the update to its colocated front-end process. |
| 2310 | client_index = -1 if self.has_coordinator else 0 |
| 2311 | self.output_queue.put_nowait( |
| 2312 | ( |
| 2313 | client_index, |
| 2314 | EngineCoreOutputs(wave_complete=self.current_wave), |
| 2315 | ) |
| 2316 | ) |
| 2317 | # Increment wave count and reset step counter. |
| 2318 | self.current_wave += 1 |
| 2319 | self.step_counter = 0 |
| 2320 | elif ( |
| 2321 | not was_running |
| 2322 | and self.has_coordinator |
| 2323 | and self.dp_rank == 0 |
| 2324 | and not self.pending_pause |
| 2325 | ): |
| 2326 | # Mirror of the wave_complete notification above: the |
| 2327 | # coordinator must observe this edge too rather than assume |
| 2328 | # that a START_DP_WAVE it sent was acted upon, since a paused |
| 2329 | # engine discards it. |
| 2330 | self.output_queue.put_nowait( |
| 2331 | (-1, EngineCoreOutputs(start_wave=self.current_wave)) |
| 2332 | ) |
| 2333 | |
| 2334 | raise SystemExit |
| 2335 | |
| 2336 | def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool: |
| 2337 | # Sync step 1 too: an idle pause needs one dummy batch, not a full interval. |
| 2338 | self.step_counter += 1 |
| 2339 | if self.step_counter != 1 and self.step_counter % self.dp_sync_interval != 0: |
| 2340 | return True |
| 2341 | |
| 2342 | has_unfinished, pause_consensus = ParallelConfig.sync_dp_state( |
| 2343 | self.dp_group, |
| 2344 | has_unfinished=local_unfinished, |
| 2345 | pending_pause=self.pending_pause, |
| 2346 | ) |
| 2347 | |
| 2348 | if pause_consensus: |
| 2349 | self.ignore_start_dp_wave = True |
| 2350 | self.pending_pause = False |
| 2351 | logger.debug("DP pause consensus reached, ignoring START_DP_WAVE.") |
| 2352 | |
| 2353 | return has_unfinished |
| 2354 | |
| 2355 | def reinitialize_distributed( |
| 2356 | self, reconfig_request: ReconfigureDistributedRequest |
| 2357 | ) -> str: |
| 2358 | from copy import deepcopy |
| 2359 | |
| 2360 | from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState |
| 2361 | |
| 2362 | new_parallel_config = deepcopy(self.vllm_config.parallel_config) |
| 2363 | old_dp_size = new_parallel_config.data_parallel_size |
| 2364 | new_parallel_config.data_parallel_size = reconfig_request.new_data_parallel_size |
| 2365 | if ( |
| 2366 | reconfig_request.new_data_parallel_rank |
| 2367 | != ReconfigureRankType.KEEP_CURRENT_RANK |
| 2368 | ): |
| 2369 | new_parallel_config.data_parallel_rank = ( |
| 2370 | reconfig_request.new_data_parallel_rank |
| 2371 | ) |
| 2372 | new_parallel_config.data_parallel_master_ip = ( |
| 2373 | reconfig_request.new_data_parallel_master_ip |
| 2374 | ) |
| 2375 | new_parallel_config.data_parallel_master_port = ( |
| 2376 | reconfig_request.new_data_parallel_master_port |
| 2377 | ) |
| 2378 | new_parallel_config._data_parallel_master_port_list = ( |
| 2379 | reconfig_request.new_data_parallel_master_port_list |
| 2380 | ) |
| 2381 | new_parallel_config._coord_store_port = reconfig_request.coord_store_port |
| 2382 | |
| 2383 | is_scale_down = reconfig_request.new_data_parallel_size < old_dp_size |
| 2384 | is_shutdown = ( |
| 2385 | reconfig_request.new_data_parallel_rank |
| 2386 | == ReconfigureRankType.SHUTDOWN_CURRENT_RANK |
| 2387 | ) |
| 2388 | |
| 2389 | if self.eep_scaling_state is not None: |
| 2390 | raise RuntimeError("Elastic EP reconfiguration is already active") |
| 2391 | |
| 2392 | state = ElasticEPScalingState( |
| 2393 | model_executor=self.model_executor, |
| 2394 | engine_core=self, |
| 2395 | vllm_config=self.vllm_config, |
| 2396 | new_parallel_config=new_parallel_config, |
| 2397 | worker_type="removing" if is_shutdown else "existing", |
| 2398 | scale_type="scale_down" if is_scale_down else "scale_up", |
| 2399 | reconfig_request=reconfig_request, |
| 2400 | ) |
| 2401 | self.eep_scaling_state = state |
| 2402 | |
| 2403 | self.process_input_queue_block = False |
| 2404 | logger.info( |
| 2405 | "[Elastic EP] Received reconfiguration request and starting scaling up/down" |
| 2406 | ) |
| 2407 | return state.ready_key |
| 2408 | |
| 2409 | def commit_prepared_elastic_ep(self) -> None: |
| 2410 | state = self.eep_scaling_state |
| 2411 | if state is None or state.commit_requested or not state.is_ready_for_switch(): |
| 2412 | raise RuntimeError("No prepared Elastic EP reconfiguration is ready") |
| 2413 | state.commit_requested = True |
| 2414 | self.process_input_queue_block = False |
| 2415 | logger.info("[Elastic EP] Committing prepared reconfiguration") |
| 2416 | |
| 2417 | def _eep_send_engine_core_notification( |
| 2418 | self, notification_type: EEPNotificationType |
| 2419 | ): |
| 2420 | """Send notifications to EngineCoreClient, which can then forward |
| 2421 | the notifications to other engine core processes. It is used for: |
| 2422 | 1) In scale down: removing core engines to notify EngineCoreClient |
| 2423 | so EngineCoreClient can release their ray placement groups; |
| 2424 | 2) Both scale up/down: to notify EngineCoreClient that existing |
| 2425 | core engines have already switched to the new parallel setup. |
| 2426 | """ |
| 2427 | dp_rank = self.vllm_config.parallel_config.data_parallel_rank |
| 2428 | notification_data = (notification_type.value, dp_rank) |
| 2429 | outputs = EngineCoreOutputs( |
| 2430 | utility_output=UtilityOutput( |
| 2431 | call_id=EEP_NOTIFICATION_CALL_ID, |
| 2432 | result=UtilityResult(notification_data), |
| 2433 | ) |
| 2434 | ) |
| 2435 | outputs.engine_index = self.engine_index |
| 2436 | |
| 2437 | if hasattr(self, "output_thread") and self.output_thread.is_alive(): |
| 2438 | self.output_queue.put_nowait((0, outputs)) |
| 2439 | else: |
| 2440 | encoder = MsgpackEncoder() |
| 2441 | with ( |
| 2442 | zmq.Context() as ctx, |
| 2443 | make_zmq_socket( |
| 2444 | ctx, self.addresses.outputs[0], zmq.PUSH, linger=4000 |
| 2445 | ) as socket, |
| 2446 | ): |
| 2447 | socket.send_multipart(encoder.encode(outputs)) |
| 2448 | |
| 2449 | def _eep_scale_up_before_kv_init(self): |
| 2450 | from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState |
| 2451 | |
| 2452 | self.ignore_start_dp_wave = True |
| 2453 | state = ElasticEPScalingState( |
| 2454 | model_executor=self.model_executor, |
| 2455 | engine_core=self, |
| 2456 | vllm_config=self.vllm_config, |
| 2457 | new_parallel_config=self.vllm_config.parallel_config, |
| 2458 | worker_type="new", |
| 2459 | scale_type="scale_up", |
| 2460 | reconfig_request=None, |
| 2461 | ) |
| 2462 | if self.eep_scaling_state is not None: |
| 2463 | raise RuntimeError("Elastic EP reconfiguration is already active") |
| 2464 | self.eep_scaling_state = state |
| 2465 | state.run_pre_kv_init_states() |
| 2466 | self.process_input_queue_block = False |
| 2467 | |
| 2468 | |
| 2469 | class EngineCoreActorMixin: |
| 2470 | """Ray actor for running EngineCore in a data parallel context.""" |
| 2471 | |
| 2472 | def __init__( |
| 2473 | self, |
| 2474 | vllm_config: VllmConfig, |
| 2475 | addresses: EngineZmqAddresses, |
| 2476 | dp_rank: int = 0, |
| 2477 | local_dp_rank: int = 0, |
| 2478 | ): |
| 2479 | if logging_config := getattr(vllm_config, "logging_config", None): |
| 2480 | configure_logging(logging_config) |
| 2481 | |
| 2482 | # Initialize tracer for distributed tracing if configured. |
| 2483 | maybe_init_worker_tracer( |
| 2484 | instrumenting_module_name="vllm.engine_core", |
| 2485 | process_kind="engine_core", |
| 2486 | process_name=f"DPEngineCoreActor_DP{dp_rank}", |
| 2487 | ) |
| 2488 | |
| 2489 | self.addresses = addresses |
| 2490 | vllm_config.parallel_config.data_parallel_index = dp_rank |
| 2491 | vllm_config.parallel_config.data_parallel_rank_local = local_dp_rank |
| 2492 | |
| 2493 | self._set_nixl_side_channel_host() |
| 2494 | |
| 2495 | # Set CUDA_VISIBLE_DEVICES as early as possible in actor life cycle |
| 2496 | # NOTE: in MP we set CUDA_VISIBLE_DEVICES at process creation time, |
| 2497 | # and this cannot be done in the same way for Ray because: |
| 2498 | # 1) Ray manages life cycle of all ray workers (including |
| 2499 | # DPEngineCoreActor) |
| 2500 | # 2) Ray sets CUDA_VISIBLE_DEVICES based on num_gpus configuration |
| 2501 | # To bypass 2, we need to also set |
| 2502 | # RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES, but vLLM workers created |
| 2503 | # thereafter would have CUDA_VISIBLE_DEVICES set, which is sticky: |
| 2504 | # https://github.com/ray-project/ray/blob/e752fc319ddedd9779a0989b6d3613909bad75c9/python/ray/_private/worker.py#L456 # noqa: E501 |
| 2505 | # This is problematic because when the vLLM worker (a Ray actor) |
| 2506 | # executes a task, it indexes into the sticky CUDA_VISIBLE_DEVICES |
| 2507 | # rather than directly using the GPU ID, potentially resulting in |
| 2508 | # index out of bounds error. See: |
| 2509 | # https://github.com/ray-project/ray/pull/40461/files#diff-31e8159767361e4bc259b6d9883d9c0d5e5db780fcea4a52ead4ee3ee4a59a78R1860 # noqa: E501 |
| 2510 | # and get_accelerator_ids_for_accelerator_resource() in worker.py |
| 2511 | # of ray. |
| 2512 | self._set_visible_devices(vllm_config, local_dp_rank) |
| 2513 | |
| 2514 | @staticmethod |
| 2515 | def _set_nixl_side_channel_host(): |
| 2516 | import ray |
| 2517 | |
| 2518 | # The driver-side value is excluded from Ray actor env propagation. |
| 2519 | # Fill in an actor-local default while preserving explicit overrides. |
| 2520 | os.environ.setdefault( |
| 2521 | "VLLM_NIXL_SIDE_CHANNEL_HOST", ray.util.get_node_ip_address() |
| 2522 | ) |
| 2523 | |
| 2524 | def _set_visible_devices(self, vllm_config: VllmConfig, local_dp_rank: int): |
| 2525 | from vllm.platforms import current_platform |
| 2526 | |
| 2527 | device_control_env_var = current_platform.device_control_env_var |
| 2528 | self._set_assigned_physical_gpu_ids( |
| 2529 | vllm_config, local_dp_rank, device_control_env_var |
| 2530 | ) |
| 2531 | |
| 2532 | def _set_assigned_physical_gpu_ids( |
| 2533 | self, |
| 2534 | vllm_config: VllmConfig, |
| 2535 | local_dp_rank: int, |
| 2536 | device_control_env_var: str, |
| 2537 | ): |
| 2538 | world_size = vllm_config.parallel_config.world_size |
| 2539 | try: |
| 2540 | physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank( |
| 2541 | device_control_env_var, |
| 2542 | local_dp_rank, |
| 2543 | world_size, |
| 2544 | user_assigned_gpu_ids=( |
| 2545 | vllm_config.parallel_config.assigned_physical_gpu_ids |
| 2546 | ), |
| 2547 | ) |
| 2548 | vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids |
| 2549 | except IndexError as e: |
| 2550 | raise Exception( |
| 2551 | f"Error computing assigned_physical_gpu_ids: " |
| 2552 | f"local range: [{local_dp_rank * world_size}, " |
| 2553 | f"{(local_dp_rank + 1) * world_size}) " |
| 2554 | f'base value: "{os.getenv(device_control_env_var)}"' |
| 2555 | ) from e |
| 2556 | |
| 2557 | @contextmanager |
| 2558 | def _perform_handshakes( |
| 2559 | self, |
| 2560 | handshake_address: str, |
| 2561 | identity: bytes, |
| 2562 | local_client: bool, |
| 2563 | vllm_config: VllmConfig, |
| 2564 | client_handshake_address: str | None, |
| 2565 | ): |
| 2566 | """For Ray, we don't need to actually perform handshake. |
| 2567 | All addresses information is known before the actor creation. |
| 2568 | Therefore, we simply yield these addresses. |
| 2569 | """ |
| 2570 | yield self.addresses |
| 2571 | |
| 2572 | def wait_for_init(self): |
| 2573 | """Wait until the engine core is initialized. |
| 2574 | |
| 2575 | This is just an empty method. When ray.get() on this method |
| 2576 | (or any other method of the actor) returns, it is guaranteed |
| 2577 | that actor creation (i.e., __init__) is complete. |
| 2578 | """ |
| 2579 | pass |
| 2580 | |
| 2581 | def run(self): |
| 2582 | """Run the engine core busy loop.""" |
| 2583 | try: |
| 2584 | self.run_busy_loop() # type: ignore[attr-defined] |
| 2585 | except SystemExit: |
| 2586 | logger.debug("EngineCore exiting.") |
| 2587 | raise |
| 2588 | except Exception: |
| 2589 | logger.exception("EngineCore encountered a fatal error.") |
| 2590 | raise |
| 2591 | finally: |
| 2592 | self.shutdown() # type: ignore[attr-defined] |
| 2593 | |
| 2594 | |
| 2595 | class DPMoEEngineCoreActor(EngineCoreActorMixin, DPEngineCoreProc): |
| 2596 | """Used for MoE model data parallel cases.""" |
| 2597 | |
| 2598 | def __init__( |
| 2599 | self, |
| 2600 | vllm_config: VllmConfig, |
| 2601 | local_client: bool, |
| 2602 | addresses: EngineZmqAddresses, |
| 2603 | executor_class: type[Executor], |
| 2604 | log_stats: bool, |
| 2605 | dp_rank: int = 0, |
| 2606 | local_dp_rank: int = 0, |
| 2607 | ): |
| 2608 | vllm_config.parallel_config.data_parallel_rank = dp_rank |
| 2609 | |
| 2610 | EngineCoreActorMixin.__init__( |
| 2611 | self, vllm_config, addresses, dp_rank, local_dp_rank |
| 2612 | ) |
| 2613 | DPEngineCoreProc.__init__( |
| 2614 | self, vllm_config, local_client, "", executor_class, log_stats |
| 2615 | ) |
| 2616 | |
| 2617 | |
| 2618 | class EngineCoreActor(EngineCoreActorMixin, EngineCoreProc): |
| 2619 | """Used for non-MoE and/or non-DP cases.""" |
| 2620 | |
| 2621 | def __init__( |
| 2622 | self, |
| 2623 | vllm_config: VllmConfig, |
| 2624 | local_client: bool, |
| 2625 | addresses: EngineZmqAddresses, |
| 2626 | executor_class: type[Executor], |
| 2627 | log_stats: bool, |
| 2628 | dp_rank: int = 0, |
| 2629 | local_dp_rank: int = 0, |
| 2630 | ): |
| 2631 | vllm_config.parallel_config.reconfigure_for_independent_dp_rank() |
| 2632 | EngineCoreActorMixin.__init__( |
| 2633 | self, vllm_config, addresses, dp_rank, local_dp_rank |
| 2634 | ) |
| 2635 | EngineCoreProc.__init__( |
| 2636 | self, |
| 2637 | vllm_config, |
| 2638 | local_client, |
| 2639 | "", |
| 2640 | executor_class, |
| 2641 | log_stats, |
| 2642 | engine_index=dp_rank, |
| 2643 | ) |