A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 04:00 UTC
[BugFix] Fix engine hanging after KV cache initialization failure (#35478)
Signed-off-by: Shiyan Deng <dsy842974287@meta.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
Shiyan Deng committed
Mar 5, 2026 at 20:58 UTC
0a208d1f549a5e35605af5b01685d64cd727b73b
2 files changed
+59
-25
vllm/v1/engine/core.py
+54
-25
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import contextlib
4
import os
5
import queue
6
import signal
@@ -117,9 +118,17 @@ class EngineCore:
118
self._eep_scale_up_before_kv_init()
119
120
# Setup KV Caches and update CacheConfig after profiling.
120
- num_gpu_blocks, num_cpu_blocks, kv_cache_config = self._initialize_kv_caches(
121
- vllm_config
122
- )
121
+ try:
122
+ num_gpu_blocks, num_cpu_blocks, kv_cache_config = (
123
+ self._initialize_kv_caches(vllm_config)
124
+ )
125
+ except Exception:
126
+ logger.exception(
127
+ "EngineCore failed during KV cache initialization; "
128
+ "shutting down executor."
129
+ )
130
+ self.model_executor.shutdown()
131
+ raise
132
133
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
134
vllm_config.cache_config.num_cpu_blocks = num_cpu_blocks
@@ -958,29 +967,49 @@ class EngineCoreProc(EngineCore):
967
addresses = self.startup_handshake(
968
handshake_socket, local_client, headless, parallel_config_to_update
969
)
961
- yield addresses
962
-
963
- # Send ready message.
964
- num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
965
- # We pass back the coordinator stats update address here for the
966
- # external LB case for our colocated front-end to use (coordinator
967
- # only runs with rank 0).
968
- dp_stats_address = self.frontend_stats_publish_address
969
-
970
- # Include config hash for DP configuration validation
971
- ready_msg = {
972
- "status": "READY",
973
- "local": local_client,
974
- "headless": headless,
975
- "num_gpu_blocks": num_gpu_blocks,
976
- "dp_stats_address": dp_stats_address,
977
- }
978
- if vllm_config.parallel_config.data_parallel_size > 1:
979
- ready_msg["parallel_config_hash"] = (
980
- vllm_config.parallel_config.compute_hash()
981
- )
970
+ exc_during_init = False
971
+ try:
972
+ yield addresses
973
+ except Exception:
974
+ exc_during_init = True
975
+ raise
976
+ finally:
977
+ if exc_during_init:
978
+ # Send FAILED status so the front-end detects init
979
+ # failure immediately via ZMQ instead of waiting for
980
+ # process sentinel (which may be delayed by cleanup).
981
+ with contextlib.suppress(Exception):
982
+ handshake_socket.send(
983
+ msgspec.msgpack.encode(
984
+ {
985
+ "status": "FAILED",
986
+ "local": local_client,
987
+ "headless": headless,
988
+ }
989
+ )
990
+ )
991
+ else:
992
+ # Send ready message.
993
+ num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
994
+ # We pass back the coordinator stats update address
995
+ # here for the external LB case for our colocated
996
+ # front-end to use (coordinator only runs with rank 0).
997
+ dp_stats_address = self.frontend_stats_publish_address
998
+
999
+ # Include config hash for DP configuration validation
1000
+ ready_msg = {
1001
+ "status": "READY",
1002
+ "local": local_client,
1003
+ "headless": headless,
1004
+ "num_gpu_blocks": num_gpu_blocks,
1005
+ "dp_stats_address": dp_stats_address,
1006
+ }
1007
+ if vllm_config.parallel_config.data_parallel_size > 1:
1008
+ ready_msg["parallel_config_hash"] = (
1009
+ vllm_config.parallel_config.compute_hash()
1010
+ )
1011
983
- handshake_socket.send(msgspec.msgpack.encode(ready_msg))
1012
+ handshake_socket.send(msgspec.msgpack.encode(ready_msg))
1013
1014
@staticmethod
1015
def startup_handshake(
vllm/v1/engine/utils.py
+5
@@ -1101,6 +1101,11 @@ def wait_for_engine_startup(
1101
1102
start_pending[0 if local else 1] -= 1
1103
engine.state = CoreEngineState.READY
1104
+ elif status == "FAILED":
1105
+ raise RuntimeError(
1106
+ f"Engine core {eng_index} reported initialization failure. "
1107
+ "See root cause above."
1108
+ )
1109
else:
1110
raise RuntimeError(
1111
f"Unexpected {status} message for "