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 15, 16:00 UTC
[V1][BugFix] Exit properly if engine core fails during startup (#16137)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Apr 7, 2025 at 15:30 UTC
7f6d47c1a2a74bd68a69103149ca729560076b9e
5 files changed
+67
-14
requirements/test.in
+1
@@ -5,6 +5,7 @@ pytest-forked
5
pytest-asyncio
6
pytest-rerunfailures
7
pytest-shard
8
+pytest-timeout
9
10
# testing utils
11
awscli
requirements/test.txt
+3
@@ -444,6 +444,7 @@ pytest==8.3.3
444
# pytest-mock
445
# pytest-rerunfailures
446
# pytest-shard
447
+ # pytest-timeout
448
pytest-asyncio==0.24.0
449
# via -r requirements/test.in
450
pytest-forked==1.6.0
@@ -454,6 +455,8 @@ pytest-rerunfailures==14.0
455
# via -r requirements/test.in
456
pytest-shard==0.1.2
457
# via -r requirements/test.in
458
+pytest-timeout==2.3.1
459
+ # via -r requirements/test.in
460
python-dateutil==2.9.0.post0
461
# via
462
# botocore
tests/v1/engine/test_engine_core_client.py
+41
@@ -3,8 +3,10 @@
3
import asyncio
4
import time
5
import uuid
6
+from threading import Thread
7
from typing import Optional
8
9
+import psutil
10
import pytest
11
from transformers import AutoTokenizer
12
@@ -245,3 +247,42 @@ async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):
247
await core_client.call_utility_async("echo", None, "help!")
248
249
assert str(e_info.value) == "Call to echo method failed: help!"
250
+
251
+
252
+@pytest.mark.timeout(10)
253
+def test_startup_failure(monkeypatch: pytest.MonkeyPatch):
254
+
255
+ with monkeypatch.context() as m, pytest.raises(Exception) as e_info:
256
+ m.setenv("VLLM_USE_V1", "1")
257
+
258
+ engine_args = EngineArgs(model=MODEL_NAME)
259
+ vllm_config = engine_args.create_engine_config(
260
+ usage_context=UsageContext.UNKNOWN_CONTEXT)
261
+ executor_class = Executor.get_class(vllm_config)
262
+
263
+ # Start another thread to wait for engine core process to start
264
+ # and kill it - simulate fatal uncaught process exit.
265
+ this_proc = psutil.Process()
266
+ children_before = set(this_proc.children())
267
+
268
+ def kill_first_child():
269
+ while True:
270
+ time.sleep(0.5)
271
+ children = set(this_proc.children()) - children_before
272
+ if children:
273
+ child = children.pop()
274
+ print("Killing child core process", child.pid)
275
+ child.kill()
276
+ break
277
+
278
+ Thread(target=kill_first_child, daemon=True).start()
279
+
280
+ _core_client = EngineCoreClient.make_client(
281
+ multiprocess_mode=True,
282
+ asyncio_mode=True,
283
+ vllm_config=vllm_config,
284
+ executor_class=executor_class,
285
+ log_stats=True,
286
+ )
287
+
288
+ assert "Engine core initialization failed" in str(e_info.value)
vllm/v1/engine/core_client.py
+14
-9
@@ -411,10 +411,21 @@ class MPClient(EngineCoreClient):
411
412
# Wait for engine core process(es) to send ready messages.
413
identities = set(eng.index for eng in self.resources.core_engines)
414
+ poller = zmq.Poller()
415
+ poller.register(sync_input_socket, zmq.POLLIN)
416
+ for eng in self.resources.core_engines:
417
+ poller.register(eng.proc_handle, zmq.POLLIN)
418
while identities:
415
- while not sync_input_socket.poll(timeout=STARTUP_POLL_PERIOD_MS):
416
- logger.info("Waiting for %d core engine proc(s) to start: %s",
417
- len(identities), identities)
419
+ events = poller.poll(STARTUP_POLL_PERIOD_MS)
420
+ if not events:
421
+ logger.debug("Waiting for %d core engine proc(s) to start: %s",
422
+ len(identities), identities)
423
+ continue
424
+ if len(events) > 1 or events[0][0] != sync_input_socket:
425
+ # One of the core processes exited.
426
+ raise RuntimeError("Engine core initialization failed. "
427
+ "See root cause above.")
428
+
429
eng_id_bytes, msg = sync_input_socket.recv_multipart()
430
eng_id = int.from_bytes(eng_id_bytes, byteorder="little")
431
if eng_id not in identities:
@@ -424,12 +435,6 @@ class MPClient(EngineCoreClient):
435
logger.info("Core engine process %d ready.", eng_id)
436
identities.discard(eng_id)
437
427
- # Double check that the process are running.
428
- for engine in self.resources.core_engines:
429
- proc = engine.proc_handle.proc
430
- if proc.exitcode is not None:
431
- raise RuntimeError(f"Engine proc {proc.name} not running")
432
-
438
def _init_core_engines(
439
self,
440
vllm_config: VllmConfig,
vllm/v1/utils.py
+8
-5
@@ -1,10 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-import multiprocessing
3
import os
4
import weakref
5
from collections import defaultdict
6
from collections.abc import Sequence
7
+from multiprocessing import Process
8
from typing import (TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar,
9
Union, overload)
10
@@ -112,20 +112,23 @@ class BackgroundProcHandle:
112
process_kwargs["output_path"] = output_path
113
114
# Run busy loop in background process.
115
- self.proc = context.Process(target=target_fn,
116
- kwargs=process_kwargs,
117
- name=process_name)
115
+ self.proc: Process = context.Process(target=target_fn,
116
+ kwargs=process_kwargs,
117
+ name=process_name)
118
self._finalizer = weakref.finalize(self, shutdown, self.proc,
119
input_path, output_path)
120
self.proc.start()
121
122
+ def fileno(self):
123
+ return self.proc.sentinel
124
+
125
def shutdown(self):
126
self._finalizer()
127
128
129
# Note(rob): shutdown function cannot be a bound method,
130
# else the gc cannot collect the object.
128
-def shutdown(proc: multiprocessing.Process, input_path: str, output_path: str):
131
+def shutdown(proc: Process, input_path: str, output_path: str):
132
# Shutdown the process.
133
if proc.is_alive():
134
proc.terminate()