| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import asyncio |
| 4 | import os |
| 5 | import socket |
| 6 | import time |
| 7 | import warnings |
| 8 | from collections.abc import AsyncGenerator, Iterable, Mapping |
| 9 | from copy import copy |
| 10 | from typing import Any |
| 11 | |
| 12 | import vllm.envs as envs |
| 13 | from vllm import TokensPrompt |
| 14 | from vllm.config import VllmConfig |
| 15 | from vllm.config.kv_events import KVEventsConfig |
| 16 | from vllm.distributed.weight_transfer.base import ( |
| 17 | WeightTransferInitRequest, |
| 18 | WeightTransferUpdateRequest, |
| 19 | ) |
| 20 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 21 | from vllm.engine.protocol import EngineClient, StreamingInput |
| 22 | from vllm.entrypoints.serve.elastic_ep.middleware import set_scaling_elastic_ep |
| 23 | from vllm.exceptions import ( |
| 24 | GracefulHTTPError, |
| 25 | MaxQueuedTokensError, |
| 26 | QueueOverflowError, |
| 27 | VLLMClientError, |
| 28 | VLLMValidationError, |
| 29 | ) |
| 30 | from vllm.inputs import EngineInput, PromptType |
| 31 | from vllm.logger import configure_logging_if_needed, init_logger |
| 32 | from vllm.lora.request import LoRARequest |
| 33 | from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry |
| 34 | from vllm.outputs import STREAM_FINISHED, PoolingRequestOutput, RequestOutput |
| 35 | from vllm.pooling_params import PoolingParams |
| 36 | from vllm.profiler.wrapper import TorchProfilerWrapper |
| 37 | from vllm.renderers import renderer_from_config |
| 38 | from vllm.renderers.inputs.preprocess import extract_prompt_components |
| 39 | from vllm.sampling_params import RequestOutputKind, SamplingParams |
| 40 | from vllm.tasks import SupportedTask |
| 41 | from vllm.tokenizers import TokenizerLike |
| 42 | from vllm.tracing import init_tracer |
| 43 | from vllm.transformers_utils.config import maybe_register_config_serialize_by_value |
| 44 | from vllm.usage.usage_lib import UsageContext |
| 45 | from vllm.utils.async_utils import cancel_task_threadsafe |
| 46 | from vllm.utils.collection_utils import as_list |
| 47 | from vllm.v1.engine import EngineCoreRequest, PauseMode |
| 48 | from vllm.v1.engine.admission_control import SharedAdmissionStats |
| 49 | from vllm.v1.engine.core_client import EngineCoreClient |
| 50 | from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError |
| 51 | from vllm.v1.engine.input_processor import InputProcessor |
| 52 | from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector |
| 53 | from vllm.v1.engine.parallel_sampling import ParentRequest |
| 54 | from vllm.v1.executor import Executor |
| 55 | from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult |
| 56 | from vllm.v1.kv_hints import KvHintsEnvelope |
| 57 | from vllm.v1.metrics.loggers import ( |
| 58 | StatLoggerFactory, |
| 59 | StatLoggerManager, |
| 60 | load_stat_logger_plugin_factories, |
| 61 | ) |
| 62 | from vllm.v1.metrics.prometheus import shutdown_prometheus |
| 63 | from vllm.v1.metrics.stats import IterationStats |
| 64 | |
| 65 | logger = init_logger(__name__) |
| 66 | |
| 67 | |
| 68 | class InputStreamError(Exception): |
| 69 | """Wrapper for errors from the input stream generator. |
| 70 | |
| 71 | This is used to propagate errors from the user's input generator |
| 72 | without wrapping them in EngineGenerateError. |
| 73 | """ |
| 74 | |
| 75 | def __init__(self, cause: Exception): |
| 76 | self.cause = cause |
| 77 | super().__init__(str(cause)) |
| 78 | |
| 79 | |
| 80 | class AsyncLLM(EngineClient): |
| 81 | """An asynchronous wrapper for the vLLM engine.""" |
| 82 | |
| 83 | def __init__( |
| 84 | self, |
| 85 | vllm_config: VllmConfig, |
| 86 | executor_class: type[Executor], |
| 87 | log_stats: bool, |
| 88 | usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, |
| 89 | mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, |
| 90 | log_requests: bool = True, |
| 91 | start_engine_loop: bool = True, |
| 92 | stat_loggers: list[StatLoggerFactory] | None = None, |
| 93 | aggregate_engine_logging: bool = False, |
| 94 | client_addresses: dict[str, Any] | None = None, |
| 95 | client_count: int = 1, |
| 96 | client_index: int = 0, |
| 97 | profiler: TorchProfilerWrapper | None = None, |
| 98 | ) -> None: |
| 99 | """Create an AsyncLLM. |
| 100 | |
| 101 | Args: |
| 102 | vllm_config: global configuration. |
| 103 | executor_class: an Executor impl, e.g. MultiprocExecutor. |
| 104 | log_stats: Whether to log stats. |
| 105 | usage_context: Usage context of the LLM. |
| 106 | mm_registry: Multi-modal registry. |
| 107 | log_requests: Whether to log requests. |
| 108 | start_engine_loop: Whether to start the engine loop. |
| 109 | client_addresses: ZMQ addresses of an externally launched engine, |
| 110 | when not launching one in-process. |
| 111 | client_count: Number of API-server clients sharing the engine. |
| 112 | client_index: Index of this client within ``client_count``. |
| 113 | aggregate_engine_logging: Whether to aggregate per-engine logs |
| 114 | into a single stat logger. |
| 115 | stat_loggers: customized stat loggers for the engine. |
| 116 | If not provided, default stat loggers will be used. |
| 117 | PLEASE BE AWARE THAT STAT LOGGER IS NOT STABLE |
| 118 | IN V1, AND ITS BASE CLASS INTERFACE MIGHT CHANGE. |
| 119 | profiler: Torch profiler wrapper used to trace the frontend. |
| 120 | |
| 121 | Returns: |
| 122 | None |
| 123 | |
| 124 | """ |
| 125 | configure_logging_if_needed(vllm_config.logging_config) |
| 126 | |
| 127 | # Ensure we can serialize custom transformer configs |
| 128 | maybe_register_config_serialize_by_value() |
| 129 | |
| 130 | self.vllm_config = vllm_config |
| 131 | self._elastic_ep_lock = asyncio.Lock() |
| 132 | self.model_config = vllm_config.model_config |
| 133 | self.scheduler_config = vllm_config.scheduler_config |
| 134 | self.observability_config = vllm_config.observability_config |
| 135 | |
| 136 | tracing_endpoint = self.observability_config.otlp_traces_endpoint |
| 137 | if tracing_endpoint is not None: |
| 138 | init_tracer("vllm.llm_engine", tracing_endpoint) |
| 139 | |
| 140 | self.log_requests = log_requests |
| 141 | |
| 142 | custom_stat_loggers = list(stat_loggers or []) |
| 143 | custom_stat_loggers.extend(load_stat_logger_plugin_factories()) |
| 144 | |
| 145 | has_custom_loggers = bool(custom_stat_loggers) |
| 146 | self.log_stats = log_stats or has_custom_loggers |
| 147 | if not log_stats and has_custom_loggers: |
| 148 | logger.info( |
| 149 | "AsyncLLM created with log_stats=False, " |
| 150 | "but custom stat loggers were found; " |
| 151 | "enabling logging without default stat loggers." |
| 152 | ) |
| 153 | |
| 154 | self.renderer = renderer = renderer_from_config(self.vllm_config) |
| 155 | |
| 156 | # Convert EngineInput --> EngineCoreRequest. |
| 157 | self.input_processor = InputProcessor(self.vllm_config, renderer) |
| 158 | |
| 159 | self.admission_stats = ( |
| 160 | SharedAdmissionStats(client_addresses, client_count, client_index) |
| 161 | if client_addresses is not None |
| 162 | and "mp_admission_counters" in client_addresses |
| 163 | else None |
| 164 | ) |
| 165 | |
| 166 | # Converts EngineCoreOutputs --> RequestOutput. |
| 167 | self.output_processor = OutputProcessor( |
| 168 | renderer.tokenizer, |
| 169 | log_stats=self.log_stats, |
| 170 | stream_interval=self.vllm_config.scheduler_config.stream_interval, |
| 171 | tracing_enabled=tracing_endpoint is not None, |
| 172 | admission_stats=self.admission_stats, |
| 173 | ) |
| 174 | |
| 175 | # EngineCore (starts the engine in background process). |
| 176 | # Hand the renderer to the client so it can start the frontend MM |
| 177 | # warmup only after engine-core fork (the why is in |
| 178 | # BaseRenderer.start_mm_warmup_in_background). The warmup is joined |
| 179 | # by reset_mm_cache / warmup / shutdown. |
| 180 | self.engine_core = EngineCoreClient.make_async_mp_client( |
| 181 | vllm_config=vllm_config, |
| 182 | executor_class=executor_class, |
| 183 | log_stats=self.log_stats, |
| 184 | client_addresses=client_addresses, |
| 185 | client_count=client_count, |
| 186 | client_index=client_index, |
| 187 | renderer=renderer, |
| 188 | ) |
| 189 | |
| 190 | # Loggers. |
| 191 | self.logger_manager: StatLoggerManager | None = None |
| 192 | if self.log_stats: |
| 193 | self.logger_manager = StatLoggerManager( |
| 194 | vllm_config=vllm_config, |
| 195 | engine_idxs=self.engine_core.engine_ranks_managed, |
| 196 | custom_stat_loggers=custom_stat_loggers, |
| 197 | enable_default_loggers=log_stats, |
| 198 | client_count=client_count, |
| 199 | aggregate_engine_logging=aggregate_engine_logging, |
| 200 | ) |
| 201 | self.logger_manager.log_engine_initialized() |
| 202 | |
| 203 | self._client_count = client_count |
| 204 | |
| 205 | self.output_handler: asyncio.Task | None = None |
| 206 | try: |
| 207 | # Start output handler eagerly if we are in the asyncio eventloop. |
| 208 | asyncio.get_running_loop() |
| 209 | self._run_output_handler() |
| 210 | except RuntimeError: |
| 211 | pass |
| 212 | |
| 213 | self.profiler = profiler |
| 214 | configured_activities = vllm_config.profiler_config.torch_profiler_activities |
| 215 | if ( |
| 216 | vllm_config.profiler_config.profiler == "torch" |
| 217 | and not vllm_config.profiler_config.ignore_frontend |
| 218 | and (configured_activities is None or "CPU" in configured_activities) |
| 219 | ): |
| 220 | profiler_dir = vllm_config.profiler_config.torch_profiler_dir |
| 221 | logger.info( |
| 222 | "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", # noqa: E501 |
| 223 | profiler_dir, |
| 224 | ) |
| 225 | worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm" |
| 226 | self.profiler = TorchProfilerWrapper( |
| 227 | vllm_config.profiler_config, |
| 228 | worker_name=worker_name, |
| 229 | local_rank=0, |
| 230 | activities=["CPU"], |
| 231 | ) |
| 232 | |
| 233 | @classmethod |
| 234 | def from_vllm_config( |
| 235 | cls, |
| 236 | vllm_config: VllmConfig, |
| 237 | start_engine_loop: bool = True, |
| 238 | usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, |
| 239 | stat_loggers: list[StatLoggerFactory] | None = None, |
| 240 | enable_log_requests: bool = False, |
| 241 | aggregate_engine_logging: bool = False, |
| 242 | disable_log_stats: bool = False, |
| 243 | client_addresses: dict[str, Any] | None = None, |
| 244 | client_count: int = 1, |
| 245 | client_index: int = 0, |
| 246 | ) -> "AsyncLLM": |
| 247 | # Create the LLMEngine. |
| 248 | return cls( |
| 249 | vllm_config=vllm_config, |
| 250 | executor_class=Executor.get_class(vllm_config), |
| 251 | start_engine_loop=start_engine_loop, |
| 252 | stat_loggers=stat_loggers, |
| 253 | log_requests=enable_log_requests, |
| 254 | log_stats=not disable_log_stats, |
| 255 | aggregate_engine_logging=aggregate_engine_logging, |
| 256 | usage_context=usage_context, |
| 257 | client_addresses=client_addresses, |
| 258 | client_count=client_count, |
| 259 | client_index=client_index, |
| 260 | ) |
| 261 | |
| 262 | @classmethod |
| 263 | def from_engine_args( |
| 264 | cls, |
| 265 | engine_args: AsyncEngineArgs, |
| 266 | start_engine_loop: bool = True, |
| 267 | usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, |
| 268 | stat_loggers: list[StatLoggerFactory] | None = None, |
| 269 | ) -> "AsyncLLM": |
| 270 | """Create an AsyncLLM from the EngineArgs.""" |
| 271 | # Create the engine configs. |
| 272 | vllm_config = engine_args.create_engine_config(usage_context) |
| 273 | executor_class = Executor.get_class(vllm_config) |
| 274 | |
| 275 | # Create the AsyncLLM. |
| 276 | return cls( |
| 277 | vllm_config=vllm_config, |
| 278 | executor_class=executor_class, |
| 279 | log_requests=engine_args.enable_log_requests, |
| 280 | log_stats=not engine_args.disable_log_stats, |
| 281 | start_engine_loop=start_engine_loop, |
| 282 | usage_context=usage_context, |
| 283 | stat_loggers=stat_loggers, |
| 284 | ) |
| 285 | |
| 286 | def __del__(self): |
| 287 | self.shutdown() |
| 288 | |
| 289 | def shutdown(self, timeout: float | None = None) -> None: |
| 290 | """Shutdown, cleaning up the background proc and IPC.""" |
| 291 | shutdown_prometheus() |
| 292 | |
| 293 | if renderer := getattr(self, "renderer", None): |
| 294 | renderer.shutdown() |
| 295 | |
| 296 | if engine_core := getattr(self, "engine_core", None): |
| 297 | engine_core.shutdown(timeout=timeout) |
| 298 | |
| 299 | handler = getattr(self, "output_handler", None) |
| 300 | if handler is not None: |
| 301 | cancel_task_threadsafe(handler) |
| 302 | |
| 303 | def get_num_unfinished_requests(self) -> int: |
| 304 | return self.output_processor.get_num_unfinished_requests() |
| 305 | |
| 306 | def get_num_queued_tokens(self) -> int: |
| 307 | return self.output_processor.get_num_queued_tokens() |
| 308 | |
| 309 | def check_admission(self, n: int = 1, request_id: str | None = None) -> None: |
| 310 | """Reject the request if it would exceed queue limits. |
| 311 | |
| 312 | Both limits return HTTP 503 (Service Unavailable) so that load |
| 313 | balancers and client SDKs retry on a different instance. |
| 314 | |
| 315 | - ``max_num_queued_reqs``: approximate global cap on the number of |
| 316 | unfinished requests (waiting + running). A request with ``n > 1`` |
| 317 | counts as ``n`` slots. |
| 318 | - ``max_num_queued_tokens``: TTFT QoS — cap on the total prompt |
| 319 | tokens of requests still in prefill. |
| 320 | |
| 321 | Note: ``get_num_queued_tokens`` uses ``prompt_len`` for all prefilling requests. |
| 322 | Chunked prefill progress and prefix-cache hits are not subtracted because the |
| 323 | scheduler's ``num_computed_tokens`` and ``num_cached_tokens`` are only |
| 324 | propagated to the API server after prefill completes. The overestimation is |
| 325 | conservative — earlier rejection, preserving TTFT targets. |
| 326 | |
| 327 | Args: |
| 328 | n: Number of sequences the request will occupy. |
| 329 | request_id: Request id, used for logging only. |
| 330 | |
| 331 | Raises: |
| 332 | QueueOverflowError: If ``max_num_queued_reqs`` would be exceeded. |
| 333 | MaxQueuedTokensError: If ``max_num_queued_tokens`` would be exceeded. |
| 334 | |
| 335 | """ |
| 336 | max_num_reqs = self.scheduler_config.max_num_queued_reqs |
| 337 | if max_num_reqs is not None: |
| 338 | current_requests = ( |
| 339 | self.admission_stats.get_num_requests() |
| 340 | if self.admission_stats is not None |
| 341 | else self.get_num_unfinished_requests() |
| 342 | ) |
| 343 | if current_requests + n > max_num_reqs: |
| 344 | logger.info( |
| 345 | "Request queue full - rejecting request %s " |
| 346 | "(current=%d, n=%d, max=%d).", |
| 347 | request_id, |
| 348 | current_requests, |
| 349 | n, |
| 350 | max_num_reqs, |
| 351 | ) |
| 352 | raise QueueOverflowError() |
| 353 | |
| 354 | max_queued_tokens = self.scheduler_config.max_num_queued_tokens |
| 355 | if max_queued_tokens is not None: |
| 356 | current_tokens = self.get_num_queued_tokens() |
| 357 | if current_tokens >= max_queued_tokens: |
| 358 | logger.info( |
| 359 | "Max queued tokens reached - rejecting request %s " |
| 360 | "(current_tokens=%d, max=%d).", |
| 361 | request_id, |
| 362 | current_tokens, |
| 363 | max_queued_tokens, |
| 364 | ) |
| 365 | raise MaxQueuedTokensError() |
| 366 | |
| 367 | async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: |
| 368 | if not hasattr(self, "_supported_tasks"): |
| 369 | # Cache the result |
| 370 | self._supported_tasks = await self.engine_core.get_supported_tasks_async() |
| 371 | |
| 372 | return self._supported_tasks |
| 373 | |
| 374 | async def add_request( |
| 375 | self, |
| 376 | request_id: str, |
| 377 | prompt: EngineCoreRequest |
| 378 | | PromptType |
| 379 | | EngineInput |
| 380 | | AsyncGenerator[StreamingInput, None], |
| 381 | params: SamplingParams | PoolingParams, |
| 382 | arrival_time: float | None = None, |
| 383 | lora_request: LoRARequest | None = None, |
| 384 | tokenization_kwargs: dict[str, Any] | None = None, |
| 385 | trace_headers: Mapping[str, str] | None = None, |
| 386 | priority: int = 0, |
| 387 | data_parallel_rank: int | None = None, |
| 388 | session_id: str | None = None, |
| 389 | prompt_text: str | None = None, |
| 390 | reasoning_ended: bool | None = None, |
| 391 | reasoning_parser_kwargs: dict[str, Any] | None = None, |
| 392 | kv_hints: KvHintsEnvelope | None = None, |
| 393 | ) -> RequestOutputCollector: |
| 394 | """Add new request to the AsyncLLM.""" |
| 395 | if self.errored: |
| 396 | raise EngineDeadError() |
| 397 | |
| 398 | is_pooling = isinstance(params, PoolingParams) |
| 399 | |
| 400 | if ( |
| 401 | self.vllm_config.cache_config.kv_sharing_fast_prefill |
| 402 | and not is_pooling |
| 403 | and params.prompt_logprobs |
| 404 | ): |
| 405 | raise VLLMValidationError( |
| 406 | "--kv-sharing-fast-prefill produces incorrect logprobs for " |
| 407 | "prompt tokens, please disable it when the requests need " |
| 408 | "prompt logprobs" |
| 409 | ) |
| 410 | |
| 411 | if isinstance(prompt, AsyncGenerator): |
| 412 | if reasoning_ended is not None or reasoning_parser_kwargs is not None: |
| 413 | raise NotImplementedError |
| 414 | |
| 415 | # Streaming input case. |
| 416 | return await self._add_streaming_input_request( |
| 417 | request_id, |
| 418 | prompt, |
| 419 | params, |
| 420 | arrival_time, |
| 421 | lora_request, |
| 422 | tokenization_kwargs, |
| 423 | trace_headers, |
| 424 | priority, |
| 425 | data_parallel_rank, |
| 426 | session_id, |
| 427 | kv_hints, |
| 428 | ) |
| 429 | |
| 430 | # Convert Input --> Request. |
| 431 | if isinstance(prompt, EngineCoreRequest): |
| 432 | logger.warning_once( |
| 433 | "Passing EngineCoreRequest to AsyncLLM.generate() and .add_requests() " |
| 434 | "is deprecated and will be removed in the future. You should " |
| 435 | "instead pass the outputs of Renderer.render_cmpl() or " |
| 436 | "Renderer.render_chat()." |
| 437 | ) |
| 438 | |
| 439 | request = prompt |
| 440 | if request_id != request.request_id: |
| 441 | logger.warning_once( |
| 442 | "AsyncLLM.add_request() was passed a request_id parameter that " |
| 443 | "does not match the EngineCoreRequest.request_id attribute. The " |
| 444 | "latter will be used, and the former will be ignored." |
| 445 | ) |
| 446 | else: |
| 447 | if isinstance(prompt, dict) and "type" in prompt: |
| 448 | # Rendered EngineInput; no blocking preprocessing needed. |
| 449 | request = self.input_processor.process_inputs( |
| 450 | request_id, |
| 451 | prompt, |
| 452 | params, |
| 453 | supported_tasks=await self.get_supported_tasks(), |
| 454 | arrival_time=arrival_time, |
| 455 | lora_request=lora_request, |
| 456 | tokenization_kwargs=tokenization_kwargs, |
| 457 | trace_headers=trace_headers, |
| 458 | priority=priority, |
| 459 | data_parallel_rank=data_parallel_rank, |
| 460 | session_id=session_id, |
| 461 | kv_hints=kv_hints, |
| 462 | ) |
| 463 | else: |
| 464 | # Raw prompts require tokenization and possibly multimodal |
| 465 | # processing, which must not block the event loop. |
| 466 | request = await self.input_processor.process_inputs_async( |
| 467 | request_id, |
| 468 | prompt, |
| 469 | params, |
| 470 | supported_tasks=await self.get_supported_tasks(), |
| 471 | arrival_time=arrival_time, |
| 472 | lora_request=lora_request, |
| 473 | tokenization_kwargs=tokenization_kwargs, |
| 474 | trace_headers=trace_headers, |
| 475 | priority=priority, |
| 476 | data_parallel_rank=data_parallel_rank, |
| 477 | session_id=session_id, |
| 478 | kv_hints=kv_hints, |
| 479 | ) |
| 480 | prompt_text, _, _ = extract_prompt_components(self.model_config, prompt) |
| 481 | |
| 482 | if reasoning_ended is not None: |
| 483 | request.reasoning_ended = reasoning_ended |
| 484 | if reasoning_parser_kwargs is not None: |
| 485 | request.reasoning_parser_kwargs = reasoning_parser_kwargs |
| 486 | |
| 487 | self.input_processor.assign_request_id(request) |
| 488 | |
| 489 | # We start the output_handler on the first call to add_request() so |
| 490 | # we can call __init__ before the event loop, which enables us |
| 491 | # to handle startup failure gracefully in the OpenAI server. |
| 492 | self._run_output_handler() |
| 493 | |
| 494 | # Create a new output collector for the request. |
| 495 | queue = RequestOutputCollector(params.output_kind, request.request_id) |
| 496 | |
| 497 | # Use cloned params that may have been updated in process_inputs() |
| 498 | params = request.params |
| 499 | |
| 500 | if is_pooling or params.n == 1: |
| 501 | await self._add_request(request, prompt_text, None, 0, queue) |
| 502 | return queue |
| 503 | |
| 504 | parent_params = params |
| 505 | assert isinstance(parent_params, SamplingParams) |
| 506 | |
| 507 | # Fan out child requests (for n>1). |
| 508 | parent_request = ParentRequest(request) |
| 509 | child_requests: list[EngineCoreRequest] = [] |
| 510 | for idx in range(parent_params.n): |
| 511 | request_id, child_params = parent_request.get_child_info(idx) |
| 512 | child_request = request if idx == parent_params.n - 1 else copy(request) |
| 513 | child_request.request_id = request_id |
| 514 | child_request.sampling_params = child_params |
| 515 | child_requests.append(child_request) |
| 516 | |
| 517 | self.check_admission(parent_params.n, parent_request.request_id) |
| 518 | try: |
| 519 | for idx, child_request in enumerate(child_requests): |
| 520 | self.output_processor.add_request( |
| 521 | child_request, prompt_text, parent_request, idx, queue |
| 522 | ) |
| 523 | |
| 524 | for child_request in child_requests: |
| 525 | await self.engine_core.add_request_async(child_request) |
| 526 | if self.log_requests: |
| 527 | logger.info("Added request %s.", child_request.request_id) |
| 528 | except BaseException: |
| 529 | await self.abort(parent_request.request_id, internal=True) |
| 530 | raise |
| 531 | return queue |
| 532 | |
| 533 | async def _add_request( |
| 534 | self, |
| 535 | request: EngineCoreRequest, |
| 536 | prompt: str | None, |
| 537 | parent_req: ParentRequest | None, |
| 538 | index: int, |
| 539 | queue: RequestOutputCollector, |
| 540 | ): |
| 541 | if parent_req is None and not self.output_processor.has_request( |
| 542 | request.request_id |
| 543 | ): |
| 544 | self.check_admission(request_id=request.request_id) |
| 545 | |
| 546 | # Register locally before the first await so concurrent tasks see this request. |
| 547 | self.output_processor.add_request(request, prompt, parent_req, index, queue) |
| 548 | |
| 549 | # Add the EngineCoreRequest to EngineCore (separate process). |
| 550 | await self.engine_core.add_request_async(request) |
| 551 | |
| 552 | if self.log_requests: |
| 553 | logger.info("Added request %s.", request.request_id) |
| 554 | |
| 555 | async def _add_streaming_input_request( |
| 556 | self, |
| 557 | request_id: str, |
| 558 | input_stream: AsyncGenerator[StreamingInput, None], |
| 559 | sampling_params: SamplingParams | PoolingParams, |
| 560 | arrival_time: float | None = None, |
| 561 | lora_request: LoRARequest | None = None, |
| 562 | tokenization_kwargs: dict[str, Any] | None = None, |
| 563 | trace_headers: Mapping[str, str] | None = None, |
| 564 | priority: int = 0, |
| 565 | data_parallel_rank: int | None = None, |
| 566 | session_id: str | None = None, |
| 567 | kv_hints: KvHintsEnvelope | None = None, |
| 568 | ) -> RequestOutputCollector: |
| 569 | self._validate_streaming_input_sampling_params(sampling_params) |
| 570 | |
| 571 | inputs = dict( |
| 572 | supported_tasks=await self.get_supported_tasks(), |
| 573 | arrival_time=arrival_time, |
| 574 | lora_request=lora_request, |
| 575 | tokenization_kwargs=tokenization_kwargs, |
| 576 | trace_headers=trace_headers, |
| 577 | priority=priority, |
| 578 | data_parallel_rank=data_parallel_rank, |
| 579 | session_id=session_id, |
| 580 | kv_hints=kv_hints, |
| 581 | ) |
| 582 | |
| 583 | if not sampling_params.skip_clone: |
| 584 | sampling_params = sampling_params.clone() |
| 585 | sampling_params.skip_clone = True |
| 586 | |
| 587 | # Create request for validation, also used as the finished signal |
| 588 | # once the input stream is closed. |
| 589 | final_req = self.input_processor.process_inputs( |
| 590 | request_id=request_id, |
| 591 | prompt=TokensPrompt(prompt_token_ids=[0]), |
| 592 | params=sampling_params, |
| 593 | **inputs, # type: ignore[arg-type] |
| 594 | ) |
| 595 | self.input_processor.assign_request_id(final_req) |
| 596 | internal_req_id = final_req.request_id |
| 597 | |
| 598 | queue = RequestOutputCollector(sampling_params.output_kind, internal_req_id) |
| 599 | |
| 600 | async def handle_inputs(): |
| 601 | cancelled = False |
| 602 | try: |
| 603 | async for input_chunk in input_stream: |
| 604 | sp = input_chunk.sampling_params |
| 605 | if sp: |
| 606 | self._validate_streaming_input_sampling_params(sp) |
| 607 | else: |
| 608 | sp = sampling_params |
| 609 | # TODO(nick): Avoid re-validating reused sampling parameters |
| 610 | req = self.input_processor.process_inputs( |
| 611 | request_id=internal_req_id, |
| 612 | prompt=input_chunk.prompt, |
| 613 | params=sp, |
| 614 | resumable=True, |
| 615 | **inputs, # type: ignore[arg-type] |
| 616 | ) |
| 617 | req.external_req_id = request_id |
| 618 | if req.prompt_embeds is not None: |
| 619 | raise VLLMValidationError( |
| 620 | "prompt_embeds not supported for streaming inputs" |
| 621 | ) |
| 622 | prompt_text, _, _ = extract_prompt_components( |
| 623 | self.model_config, input_chunk.prompt |
| 624 | ) |
| 625 | await self._add_request(req, prompt_text, None, 0, queue) |
| 626 | except (asyncio.CancelledError, GeneratorExit): |
| 627 | cancelled = True |
| 628 | except Exception as error: |
| 629 | # Wrap in InputStreamError so generate() can propagate it |
| 630 | # without wrapping in EngineGenerateError. |
| 631 | queue.put(InputStreamError(error)) |
| 632 | finally: |
| 633 | queue._input_stream_task = None |
| 634 | if not cancelled: |
| 635 | # Send empty final request to indicate that inputs have |
| 636 | # finished. Don't send if cancelled (session was aborted). |
| 637 | await self._add_request(final_req, None, None, 0, queue) |
| 638 | |
| 639 | # Ensure output handler is running. |
| 640 | self._run_output_handler() |
| 641 | |
| 642 | queue._input_stream_task = asyncio.create_task(handle_inputs()) |
| 643 | return queue |
| 644 | |
| 645 | @staticmethod |
| 646 | def _validate_streaming_input_sampling_params( |
| 647 | params: SamplingParams | PoolingParams, |
| 648 | ): |
| 649 | if ( |
| 650 | not isinstance(params, SamplingParams) |
| 651 | or params.n > 1 |
| 652 | or params.output_kind == RequestOutputKind.FINAL_ONLY |
| 653 | or params.stop |
| 654 | ): |
| 655 | raise VLLMValidationError( |
| 656 | "Input streaming not currently supported " |
| 657 | "for pooling models, n > 1, request_kind = FINAL_ONLY " |
| 658 | "or with stop strings." |
| 659 | ) |
| 660 | |
| 661 | # TODO: we should support multiple prompts in one call, as you |
| 662 | # can do with LLM.generate. So that for multi-prompt completion |
| 663 | # requests we don't need to send multiple messages to core proc, |
| 664 | # and so we don't need multiple streams which then get |
| 665 | # re-multiplexed in the API server anyhow. |
| 666 | async def generate( |
| 667 | self, |
| 668 | prompt: EngineCoreRequest |
| 669 | | PromptType |
| 670 | | EngineInput |
| 671 | | AsyncGenerator[StreamingInput, None], |
| 672 | sampling_params: SamplingParams, |
| 673 | request_id: str, |
| 674 | *, |
| 675 | prompt_text: str | None = None, |
| 676 | lora_request: LoRARequest | None = None, |
| 677 | tokenization_kwargs: dict[str, Any] | None = None, |
| 678 | trace_headers: Mapping[str, str] | None = None, |
| 679 | priority: int = 0, |
| 680 | data_parallel_rank: int | None = None, |
| 681 | session_id: str | None = None, |
| 682 | kv_hints: KvHintsEnvelope | None = None, |
| 683 | reasoning_ended: bool | None = None, |
| 684 | reasoning_parser_kwargs: dict[str, Any] | None = None, |
| 685 | ) -> AsyncGenerator[RequestOutput, None]: |
| 686 | """Main function called by the API server to kick off a request |
| 687 | * 1) Making an AsyncStream corresponding to the Request. |
| 688 | * 2) Processing the Input. |
| 689 | * 3) Adding the Request to the Detokenizer. |
| 690 | * 4) Adding the Request to the EngineCore (separate process). |
| 691 | |
| 692 | A separate output_handler loop runs in a background AsyncIO task, |
| 693 | pulling outputs from EngineCore and putting them into the |
| 694 | per-request AsyncStream. |
| 695 | |
| 696 | The caller of generate() iterates the returned AsyncGenerator, |
| 697 | returning the RequestOutput back to the caller. |
| 698 | |
| 699 | Note: |
| 700 | Passing a raw prompt string directly to this method is deprecated. |
| 701 | Advanced power-users can manually bypass the raw-prompt fallback |
| 702 | path using the Engine's underlying Renderer pipeline: |
| 703 | |
| 704 | >>> from vllm.inputs import parse_model_prompt |
| 705 | >>> parsed = parse_model_prompt(self.model_config, "Prompt text") |
| 706 | >>> params = self.renderer.default_cmpl_tok_params |
| 707 | >>> (engine_input,) = self.renderer.render_cmpl([parsed], params) |
| 708 | >>> gen = self.generate(engine_input, sampling_params, request_id) |
| 709 | |
| 710 | """ |
| 711 | q: RequestOutputCollector | None = None |
| 712 | try: |
| 713 | q = await self.add_request( |
| 714 | request_id, |
| 715 | prompt, |
| 716 | sampling_params, |
| 717 | lora_request=lora_request, |
| 718 | tokenization_kwargs=tokenization_kwargs, |
| 719 | trace_headers=trace_headers, |
| 720 | priority=priority, |
| 721 | data_parallel_rank=data_parallel_rank, |
| 722 | session_id=session_id, |
| 723 | kv_hints=kv_hints, |
| 724 | prompt_text=prompt_text, |
| 725 | reasoning_ended=reasoning_ended, |
| 726 | reasoning_parser_kwargs=reasoning_parser_kwargs, |
| 727 | ) |
| 728 | |
| 729 | # The output_handler task pushes items into the queue. |
| 730 | # This task pulls from the queue and yields to caller. |
| 731 | finished = False |
| 732 | while not finished: |
| 733 | # Note: drain queue without await if possible (avoids |
| 734 | # task switching under load which helps performance). |
| 735 | out = q.get_nowait() or await q.get() |
| 736 | |
| 737 | # Note: both OutputProcessor and EngineCore handle their |
| 738 | # own request cleanup based on finished. |
| 739 | assert isinstance(out, RequestOutput) |
| 740 | finished = out.finished |
| 741 | if out is not STREAM_FINISHED: |
| 742 | yield out |
| 743 | |
| 744 | # If the request is disconnected by the client, generate() |
| 745 | # is cancelled or the generator is garbage collected. So, |
| 746 | # we abort the request if we end up here. |
| 747 | except (asyncio.CancelledError, GeneratorExit): |
| 748 | if q is not None: |
| 749 | await self.abort(q.request_id, internal=True) |
| 750 | if self.log_requests: |
| 751 | logger.info("Request %s aborted.", request_id) |
| 752 | raise |
| 753 | |
| 754 | # Engine is dead. Do not abort since we shut down. |
| 755 | except EngineDeadError: |
| 756 | if self.log_requests: |
| 757 | logger.info("Request %s failed (engine dead).", request_id) |
| 758 | raise |
| 759 | |
| 760 | # Request validation error or admission control rejection. |
| 761 | except (VLLMClientError, GracefulHTTPError) as e: |
| 762 | if self.log_requests: |
| 763 | logger.info("Request %s failed (bad request): %s.", request_id, e) |
| 764 | raise |
| 765 | |
| 766 | # Error from input stream generator - propagate directly. |
| 767 | except InputStreamError as e: |
| 768 | if q is not None: |
| 769 | await self.abort(q.request_id, internal=True) |
| 770 | if self.log_requests: |
| 771 | logger.info("Request %s failed (input error): %s.", request_id, e) |
| 772 | raise e.cause from e |
| 773 | |
| 774 | # Unexpected error in the generate() task (possibly recoverable). |
| 775 | except Exception as e: |
| 776 | if q is not None: |
| 777 | await self.abort(q.request_id, internal=True) |
| 778 | if self.log_requests: |
| 779 | try: |
| 780 | s = f"{e.__class__.__name__}: {e}" |
| 781 | except Exception as e2: |
| 782 | s = ( |
| 783 | f"{e.__class__.__name__}: " |
| 784 | "error during printing an exception of class" |
| 785 | + e2.__class__.__name__ |
| 786 | ) |
| 787 | logger.info("Request %s failed due to %s.", request_id, s) |
| 788 | raise EngineGenerateError() from e |
| 789 | finally: |
| 790 | if q is not None: |
| 791 | q.close() |
| 792 | |
| 793 | def _run_output_handler(self): |
| 794 | """Background loop: pulls from EngineCore and pushes to AsyncStreams.""" |
| 795 | if self.output_handler is not None: |
| 796 | return |
| 797 | |
| 798 | # Ensure that the task doesn't have a circular ref back to the AsyncLLM |
| 799 | # object, or else it won't be garbage collected and cleaned up properly. |
| 800 | engine_core = self.engine_core |
| 801 | output_processor = self.output_processor |
| 802 | log_stats = self.log_stats |
| 803 | # We use a mutable list for logger_manager so that it can be updated |
| 804 | # during elastic EP scaling (see scale_elastic_ep) without creating |
| 805 | # a circular reference via self. |
| 806 | self._logger_ref = [self.logger_manager] |
| 807 | logger_ref = self._logger_ref |
| 808 | renderer = self.renderer |
| 809 | # P0 multi-modal sender ("shadow") cache; None for text-only models. |
| 810 | mm_processor_cache = renderer.mm_processor_cache |
| 811 | chunk_size = envs.VLLM_V1_OUTPUT_PROC_CHUNK_SIZE |
| 812 | |
| 813 | async def output_handler(): |
| 814 | try: |
| 815 | while True: |
| 816 | # 1) Pull EngineCoreOutputs from the EngineCore. |
| 817 | outputs = await engine_core.get_output_async() |
| 818 | num_outputs = len(outputs.outputs) |
| 819 | |
| 820 | iteration_stats = ( |
| 821 | IterationStats() if (log_stats and num_outputs) else None |
| 822 | ) |
| 823 | |
| 824 | # Split outputs into chunks of at most |
| 825 | # VLLM_V1_OUTPUT_PROC_CHUNK_SIZE, so that we don't block the |
| 826 | # event loop for too long. |
| 827 | engine_core_outputs = outputs.outputs |
| 828 | for start in range(0, num_outputs, chunk_size): |
| 829 | end = start + chunk_size |
| 830 | outputs_slice = engine_core_outputs[start:end] |
| 831 | # 2) Process EngineCoreOutputs. |
| 832 | processed_outputs = output_processor.process_outputs( |
| 833 | outputs_slice, outputs.timestamp, iteration_stats |
| 834 | ) |
| 835 | # NOTE: RequestOutputs are pushed to their queues. |
| 836 | assert not processed_outputs.request_outputs |
| 837 | |
| 838 | # 2b) Recover from P0/P1 cache drift: the engine flags hashes |
| 839 | # it couldn't find (mm_cache_miss_hashes); drop them from the |
| 840 | # P0 shadow so the client's retry resends the data and |
| 841 | # repopulates P1. Hot-path no-op (field is None otherwise). |
| 842 | if mm_processor_cache is not None: |
| 843 | for eco in outputs_slice: |
| 844 | if eco.mm_cache_miss_hashes: |
| 845 | for mm_hash in eco.mm_cache_miss_hashes: |
| 846 | mm_processor_cache.invalidate(mm_hash) |
| 847 | |
| 848 | # Allow other asyncio tasks to run between chunks |
| 849 | if end < num_outputs: |
| 850 | await asyncio.sleep(0) |
| 851 | |
| 852 | # 3) Abort any reqs that finished due to stop strings. |
| 853 | if processed_outputs.reqs_to_abort: |
| 854 | await engine_core.abort_requests_async( |
| 855 | processed_outputs.reqs_to_abort |
| 856 | ) |
| 857 | |
| 858 | output_processor.update_scheduler_stats(outputs.scheduler_stats) |
| 859 | |
| 860 | # 4) Logging. |
| 861 | # TODO(rob): make into a coroutine and launch it in |
| 862 | # background thread once Prometheus overhead is non-trivial. |
| 863 | if logger_ref[0]: |
| 864 | logger_ref[0].record( |
| 865 | engine_idx=outputs.engine_index, |
| 866 | scheduler_stats=outputs.scheduler_stats, |
| 867 | iteration_stats=iteration_stats, |
| 868 | mm_cache_stats=renderer.stat_mm_cache(), |
| 869 | ) |
| 870 | except Exception as e: |
| 871 | logger.exception("AsyncLLM output_handler failed.") |
| 872 | output_processor.propagate_error(e) |
| 873 | |
| 874 | self.output_handler = asyncio.create_task(output_handler()) |
| 875 | |
| 876 | async def abort( |
| 877 | self, request_id: str | Iterable[str], internal: bool = False |
| 878 | ) -> None: |
| 879 | """Abort RequestId in OutputProcessor and EngineCore.""" |
| 880 | request_ids = ( |
| 881 | (request_id,) if isinstance(request_id, str) else as_list(request_id) |
| 882 | ) |
| 883 | all_request_ids = self.output_processor.abort_requests(request_ids, internal) |
| 884 | await self.engine_core.abort_requests_async(all_request_ids) |
| 885 | |
| 886 | if self.log_requests: |
| 887 | logger.info("Aborted request(s) %s.", ",".join(request_ids)) |
| 888 | |
| 889 | async def notify_kv_transfer_request_rejected( |
| 890 | self, |
| 891 | request_id: str, |
| 892 | kv_transfer_params: dict[str, Any], |
| 893 | *, |
| 894 | data_parallel_rank: int | None = None, |
| 895 | ) -> None: |
| 896 | """Submit a pre-aborted request so the connector's request_finished |
| 897 | hook runs to free any pre-admission KV-transfer resources (e.g. NIXL |
| 898 | prefill blocks pinned on the P node).""" |
| 899 | request = EngineCoreRequest( |
| 900 | request_id=request_id, |
| 901 | prompt_token_ids=[0], |
| 902 | mm_features=None, |
| 903 | sampling_params=SamplingParams( |
| 904 | max_tokens=1, |
| 905 | extra_args={"kv_transfer_params": dict(kv_transfer_params)}, |
| 906 | ), |
| 907 | pooling_params=None, |
| 908 | arrival_time=time.time(), |
| 909 | lora_request=None, |
| 910 | cache_salt=None, |
| 911 | data_parallel_rank=data_parallel_rank, |
| 912 | abort_immediately=True, |
| 913 | ) |
| 914 | await self.engine_core.add_request_async(request) |
| 915 | |
| 916 | async def pause_generation( |
| 917 | self, |
| 918 | *, |
| 919 | mode: PauseMode = "abort", |
| 920 | wait_for_inflight_requests: bool | None = None, |
| 921 | clear_cache: bool = True, |
| 922 | ) -> None: |
| 923 | """Pause generation to allow model weight updates. |
| 924 | |
| 925 | All mode handling (abort / wait / keep) and cache clearing is done |
| 926 | in the engine. New generation/encoding requests will not be scheduled |
| 927 | until resume is called. |
| 928 | |
| 929 | Args: |
| 930 | mode: How to handle in-flight requests: |
| 931 | - ``"abort"``: Abort all in-flight requests immediately |
| 932 | (default). |
| 933 | - ``"wait"``: Wait for in-flight requests to complete. |
| 934 | - ``"keep"``: Freeze requests in queue; they resume on |
| 935 | :meth:`resume_generation`. |
| 936 | wait_for_inflight_requests: DEPRECATED: use mode argument. |
| 937 | clear_cache: Whether to clear KV cache and prefix cache after |
| 938 | draining. Set to ``False`` to preserve cache for faster resume. |
| 939 | |
| 940 | """ |
| 941 | if wait_for_inflight_requests: |
| 942 | warnings.warn( |
| 943 | "The `wait_for_inflight_requests` parameter in " |
| 944 | "`AsyncLLM.pause_generation()` is deprecated. " |
| 945 | "Please use `mode` argument instead.", |
| 946 | DeprecationWarning, |
| 947 | stacklevel=2, |
| 948 | ) |
| 949 | mode = "wait" |
| 950 | if clear_cache: |
| 951 | await self.renderer.clear_mm_cache_async() |
| 952 | await self.engine_core.pause_scheduler_async(mode=mode, clear_cache=clear_cache) |
| 953 | # Small sleep to help ensure that final outputs from any in-flight requests are |
| 954 | # returned prior to this method returning. These outputs come out of the engine |
| 955 | # prior to the wait-for-idle completion event, but involve additional async |
| 956 | # tasks in output processing. |
| 957 | # Note that this is not required for correctness, just more intuitive ordering |
| 958 | # of events from caller's pov. |
| 959 | await asyncio.sleep(0.02) |
| 960 | |
| 961 | async def resume_generation(self) -> None: |
| 962 | """Resume generation after :meth:`pause_generation`.""" |
| 963 | await self.engine_core.resume_scheduler_async() |
| 964 | |
| 965 | async def is_paused(self) -> bool: |
| 966 | """Return whether the engine is currently paused.""" |
| 967 | return await self.engine_core.is_scheduler_paused_async() |
| 968 | |
| 969 | async def encode( |
| 970 | self, |
| 971 | prompt: PromptType | EngineInput, |
| 972 | pooling_params: PoolingParams, |
| 973 | request_id: str, |
| 974 | lora_request: LoRARequest | None = None, |
| 975 | trace_headers: Mapping[str, str] | None = None, |
| 976 | priority: int = 0, |
| 977 | tokenization_kwargs: dict[str, Any] | None = None, |
| 978 | reasoning_ended: bool | None = None, |
| 979 | ) -> AsyncGenerator[PoolingRequestOutput, None]: |
| 980 | """Main function called by the API server to kick off a request |
| 981 | * 1) Making an AsyncStream corresponding to the Request. |
| 982 | * 2) Processing the Input. |
| 983 | * 3) Adding the Request to the EngineCore (separate process). |
| 984 | |
| 985 | A separate output_handler loop runs in a background AsyncIO task, |
| 986 | pulling outputs from EngineCore and putting them into the |
| 987 | per-request AsyncStream. |
| 988 | |
| 989 | The caller of generate() iterates the returned AsyncGenerator, |
| 990 | returning the RequestOutput back to the caller. |
| 991 | """ |
| 992 | q: RequestOutputCollector | None = None |
| 993 | try: |
| 994 | q = await self.add_request( |
| 995 | request_id, |
| 996 | prompt, |
| 997 | pooling_params, |
| 998 | lora_request=lora_request, |
| 999 | tokenization_kwargs=tokenization_kwargs, |
| 1000 | trace_headers=trace_headers, |
| 1001 | priority=priority, |
| 1002 | reasoning_ended=reasoning_ended, |
| 1003 | ) |
| 1004 | |
| 1005 | # The output_handler task pushes items into the queue. |
| 1006 | # This task pulls from the queue and yields to caller. |
| 1007 | finished = False |
| 1008 | while not finished: |
| 1009 | # Note: drain queue without await if possible (avoids |
| 1010 | # task switching under load which helps performance). |
| 1011 | out = q.get_nowait() or await q.get() |
| 1012 | assert isinstance(out, PoolingRequestOutput) |
| 1013 | # Note: both OutputProcessor and EngineCore handle their |
| 1014 | # own request cleanup based on finished. |
| 1015 | finished = out.finished |
| 1016 | yield out |
| 1017 | |
| 1018 | # If the request is disconnected by the client, generate() |
| 1019 | # is cancelled. So, we abort the request if we end up here. |
| 1020 | except asyncio.CancelledError: |
| 1021 | if q is not None: |
| 1022 | await self.abort(q.request_id, internal=True) |
| 1023 | if self.log_requests: |
| 1024 | logger.info("Request %s aborted.", request_id) |
| 1025 | raise |
| 1026 | |
| 1027 | # Engine is dead. Do not abort since we shut down. |
| 1028 | except EngineDeadError: |
| 1029 | if self.log_requests: |
| 1030 | logger.info("Request %s failed (engine dead).", request_id) |
| 1031 | raise |
| 1032 | |
| 1033 | # Request validation error or admission control rejection. |
| 1034 | except (VLLMClientError, GracefulHTTPError): |
| 1035 | if self.log_requests: |
| 1036 | logger.info("Request %s failed (bad request).", request_id) |
| 1037 | raise |
| 1038 | |
| 1039 | # Unexpected error in the generate() task (possibly recoverable). |
| 1040 | except Exception as e: |
| 1041 | if q is not None: |
| 1042 | await self.abort(q.request_id, internal=True) |
| 1043 | if self.log_requests: |
| 1044 | logger.info("Request %s failed.", request_id) |
| 1045 | raise EngineGenerateError() from e |
| 1046 | finally: |
| 1047 | if q is not None: |
| 1048 | q.close() |
| 1049 | |
| 1050 | @property |
| 1051 | def tokenizer(self) -> TokenizerLike | None: |
| 1052 | return self.renderer.tokenizer |
| 1053 | |
| 1054 | def get_tokenizer(self) -> TokenizerLike: |
| 1055 | return self.renderer.get_tokenizer() |
| 1056 | |
| 1057 | async def is_tracing_enabled(self) -> bool: |
| 1058 | return self.observability_config.otlp_traces_endpoint is not None |
| 1059 | |
| 1060 | async def do_log_stats(self) -> None: |
| 1061 | if self.logger_manager: |
| 1062 | self.logger_manager.log() |
| 1063 | |
| 1064 | async def check_health(self) -> None: |
| 1065 | logger.debug("Called check_health.") |
| 1066 | if self.errored: |
| 1067 | raise self.dead_error |
| 1068 | |
| 1069 | async def start_profile(self, profile_prefix: str | None = None) -> None: |
| 1070 | coros = [self.engine_core.profile_async(True, profile_prefix)] |
| 1071 | if self.profiler is not None: |
| 1072 | coros.append(asyncio.to_thread(self.profiler.start)) |
| 1073 | await asyncio.gather(*coros) |
| 1074 | |
| 1075 | async def stop_profile(self) -> None: |
| 1076 | coros = [self.engine_core.profile_async(False)] |
| 1077 | if self.profiler is not None: |
| 1078 | coros.append(asyncio.to_thread(self.profiler.stop)) |
| 1079 | await asyncio.gather(*coros) |
| 1080 | |
| 1081 | async def reset_mm_cache(self) -> None: |
| 1082 | # Join the background MM warmup first: the mm_processor_cache is not |
| 1083 | # safe for concurrent access with its apply/clear. |
| 1084 | self.renderer._join_mm_warmup() |
| 1085 | await self.renderer.clear_mm_cache_async() |
| 1086 | await self.engine_core.reset_mm_cache_async() |
| 1087 | |
| 1088 | async def reset_prefix_cache( |
| 1089 | self, reset_running_requests: bool = False, reset_connector: bool = False |
| 1090 | ) -> bool: |
| 1091 | return await self.engine_core.reset_prefix_cache_async( |
| 1092 | reset_running_requests, reset_connector |
| 1093 | ) |
| 1094 | |
| 1095 | async def reset_encoder_cache(self) -> None: |
| 1096 | await self.engine_core.reset_encoder_cache_async() |
| 1097 | |
| 1098 | async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None: |
| 1099 | if level >= 1: |
| 1100 | await self.renderer.clear_mm_cache_async() |
| 1101 | await self.engine_core.sleep_async(level, mode) |
| 1102 | |
| 1103 | if self.logger_manager is not None: |
| 1104 | self.logger_manager.record_sleep_state(1, level) |
| 1105 | |
| 1106 | async def release_kv_cache_memory(self) -> None: |
| 1107 | await self.renderer.clear_mm_cache_async() |
| 1108 | await self.engine_core.release_kv_cache_memory_async() |
| 1109 | |
| 1110 | if self.logger_manager is not None: |
| 1111 | self.logger_manager.record_sleep_state(1, 0) |
| 1112 | |
| 1113 | async def wake_up(self, tags: list[str] | None = None) -> bool: |
| 1114 | fully_awake = await self.engine_core.wake_up_async(tags) |
| 1115 | |
| 1116 | if self.logger_manager is not None and fully_awake: |
| 1117 | self.logger_manager.record_sleep_state(0, 0) |
| 1118 | |
| 1119 | return fully_awake |
| 1120 | |
| 1121 | async def checkpoint_prepare(self) -> None: |
| 1122 | await self.collective_rpc("checkpoint_prepare") |
| 1123 | |
| 1124 | async def checkpoint_restore(self) -> None: |
| 1125 | await self.collective_rpc("checkpoint_restore") |
| 1126 | |
| 1127 | async def is_sleeping(self) -> bool: |
| 1128 | return await self.engine_core.is_sleeping_async() |
| 1129 | |
| 1130 | async def compute_weight_checksums(self) -> list[dict[str, str]]: |
| 1131 | return await self.engine_core.compute_weight_checksums_async() |
| 1132 | |
| 1133 | async def add_lora(self, lora_request: LoRARequest) -> bool: |
| 1134 | """Load a new LoRA adapter into the engine for future requests.""" |
| 1135 | return await self.engine_core.add_lora_async(lora_request) |
| 1136 | |
| 1137 | async def remove_lora(self, lora_id: int) -> bool: |
| 1138 | """Remove an already loaded LoRA adapter.""" |
| 1139 | return await self.engine_core.remove_lora_async(lora_id) |
| 1140 | |
| 1141 | async def list_loras(self) -> set[int]: |
| 1142 | """List all registered adapters.""" |
| 1143 | return await self.engine_core.list_loras_async() |
| 1144 | |
| 1145 | async def pin_lora(self, lora_id: int) -> bool: |
| 1146 | """Prevent an adapter from being evicted.""" |
| 1147 | return await self.engine_core.pin_lora_async(lora_id) |
| 1148 | |
| 1149 | async def collective_rpc( |
| 1150 | self, |
| 1151 | method: str, |
| 1152 | timeout: float | None = None, |
| 1153 | args: tuple = (), |
| 1154 | kwargs: dict | None = None, |
| 1155 | ): |
| 1156 | """Perform a collective RPC call to the given path.""" |
| 1157 | return await self.engine_core.collective_rpc_async( |
| 1158 | method, timeout, args, kwargs |
| 1159 | ) |
| 1160 | |
| 1161 | async def wait_for_requests_to_drain(self, drain_timeout: int = 300): |
| 1162 | """Wait for all requests to be drained.""" |
| 1163 | start_time = time.time() |
| 1164 | while time.time() - start_time < drain_timeout: |
| 1165 | if not self.engine_core.dp_engines_running(): |
| 1166 | logger.info("Engines are idle, requests have been drained") |
| 1167 | return |
| 1168 | |
| 1169 | logger.info("Engines are still running, waiting for requests to drain...") |
| 1170 | await asyncio.sleep(1) # Wait 1 second before checking again |
| 1171 | |
| 1172 | raise TimeoutError( |
| 1173 | f"Timeout reached after {drain_timeout} seconds " |
| 1174 | "waiting for requests to drain." |
| 1175 | ) |
| 1176 | |
| 1177 | async def _drain_requests_for_elastic_ep(self, drain_timeout: int) -> None: |
| 1178 | try: |
| 1179 | logger.info( |
| 1180 | "VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, " |
| 1181 | "waiting for requests to drain before scaling" |
| 1182 | ) |
| 1183 | await self.wait_for_requests_to_drain(drain_timeout) |
| 1184 | except BaseException: |
| 1185 | set_scaling_elastic_ep(False) |
| 1186 | raise |
| 1187 | |
| 1188 | async def scale_elastic_ep( |
| 1189 | self, new_data_parallel_size: int, drain_timeout: int = 300 |
| 1190 | ): |
| 1191 | """Scale the elastic EP data parallel size.""" |
| 1192 | async with self._elastic_ep_lock: |
| 1193 | await self._scale_elastic_ep(new_data_parallel_size, drain_timeout) |
| 1194 | |
| 1195 | async def _scale_elastic_ep( |
| 1196 | self, new_data_parallel_size: int, drain_timeout: int |
| 1197 | ) -> None: |
| 1198 | old_data_parallel_size = self.vllm_config.parallel_config.data_parallel_size |
| 1199 | if old_data_parallel_size == new_data_parallel_size: |
| 1200 | logger.info( |
| 1201 | "Data parallel size is already %s, skipping scale", |
| 1202 | new_data_parallel_size, |
| 1203 | ) |
| 1204 | return |
| 1205 | |
| 1206 | await self.engine_core.prepare_elastic_ep(new_data_parallel_size) |
| 1207 | |
| 1208 | # recreate stat loggers |
| 1209 | if new_data_parallel_size > old_data_parallel_size and self.log_stats: |
| 1210 | # TODO(rob): fix this after talking with Ray team. |
| 1211 | # This resets all the prometheus metrics since we |
| 1212 | # unregister during initialization. Need to understand |
| 1213 | # the intended behavior here better. |
| 1214 | self.logger_manager = StatLoggerManager( |
| 1215 | vllm_config=self.vllm_config, |
| 1216 | engine_idxs=list(range(new_data_parallel_size)), |
| 1217 | custom_stat_loggers=None, |
| 1218 | ) |
| 1219 | # Update the mutable ref so output_handler picks up the |
| 1220 | # new logger without creating a circular reference via self. |
| 1221 | if hasattr(self, "_logger_ref"): |
| 1222 | self._logger_ref[0] = self.logger_manager |
| 1223 | self.logger_manager.log_engine_initialized() |
| 1224 | |
| 1225 | set_scaling_elastic_ep(True) |
| 1226 | if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS: |
| 1227 | await self._drain_requests_for_elastic_ep(drain_timeout) |
| 1228 | |
| 1229 | await self.engine_core.commit_elastic_ep() |
| 1230 | self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size |
| 1231 | set_scaling_elastic_ep(False) |
| 1232 | |
| 1233 | async def handle_fault( |
| 1234 | self, fault_tolerance_request: FaultToleranceRequest |
| 1235 | ) -> FaultToleranceResult: |
| 1236 | """Send fault tolerance instruction to the engine.""" |
| 1237 | return await self.engine_core.handle_fault(fault_tolerance_request) |
| 1238 | |
| 1239 | async def get_status(self): |
| 1240 | return await self.engine_core.get_status() |
| 1241 | |
| 1242 | @property |
| 1243 | def is_running(self) -> bool: |
| 1244 | # Is None before the loop is started. |
| 1245 | return self.output_handler is None or not self.output_handler.done() |
| 1246 | |
| 1247 | @property |
| 1248 | def is_stopped(self) -> bool: |
| 1249 | return self.errored |
| 1250 | |
| 1251 | @property |
| 1252 | def errored(self) -> bool: |
| 1253 | return self.engine_core.resources.engine_dead or not self.is_running |
| 1254 | |
| 1255 | @property |
| 1256 | def dead_error(self) -> BaseException: |
| 1257 | return EngineDeadError() |
| 1258 | |
| 1259 | async def init_weight_transfer_engine( |
| 1260 | self, request: WeightTransferInitRequest |
| 1261 | ) -> None: |
| 1262 | """Initialize weight transfer for RL training. |
| 1263 | |
| 1264 | Args: |
| 1265 | request: Weight transfer initialization request with backend-specific info |
| 1266 | |
| 1267 | """ |
| 1268 | await self.collective_rpc( |
| 1269 | "init_weight_transfer_engine", kwargs={"init_info": request.init_info} |
| 1270 | ) |
| 1271 | |
| 1272 | async def start_weight_update(self) -> None: |
| 1273 | """Start a new weight update.""" |
| 1274 | await self.collective_rpc("start_weight_update") |
| 1275 | |
| 1276 | async def start_draft_weight_update(self) -> None: |
| 1277 | """Start a new weight update targeting the speculative draft model.""" |
| 1278 | await self.collective_rpc("start_draft_weight_update") |
| 1279 | |
| 1280 | async def update_weights(self, request: WeightTransferUpdateRequest) -> None: |
| 1281 | """Batched weight update for RL training. |
| 1282 | |
| 1283 | Args: |
| 1284 | request: Weight update request with backend-specific update info |
| 1285 | |
| 1286 | """ |
| 1287 | await self.collective_rpc( |
| 1288 | "update_weights", kwargs={"update_info": request.update_info} |
| 1289 | ) |
| 1290 | |
| 1291 | async def finish_weight_update(self, weight_version: str | None = None) -> None: |
| 1292 | """Finish the weight update and set its version if provided.""" |
| 1293 | await self.collective_rpc("finish_weight_update") |
| 1294 | if weight_version is not None: |
| 1295 | await self.update_weight_version(weight_version) |
| 1296 | |
| 1297 | async def update_weight_version(self, new_version: str) -> None: |
| 1298 | """Set the weight version without updating weights.""" |
| 1299 | await self.engine_core.set_weight_version_async(new_version) |
| 1300 | |
| 1301 | async def get_weight_version(self) -> str: |
| 1302 | """Return the latest committed weight version.""" |
| 1303 | return await self.engine_core.get_weight_version_async() |
| 1304 | |
| 1305 | def get_kv_event_sources(self) -> dict[int, KVEventsConfig]: |
| 1306 | return self.engine_core.get_kv_event_sources() |