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, 07:15 UTC
[V1] Support DP with Ray (#18779)
Rui Qiao committed
Jun 2, 2025 at 21:15 UTC
bdce64f2365b39335141f8efcb3a0a8ecc559153
10 files changed
+539
-108
requirements/test.in
+1
-1
@@ -17,7 +17,7 @@ vector_quantize_pytorch # required for minicpmo_26 test
17
vocos # required for minicpmo_26 test
18
peft
19
pqdm
20
-ray[cgraph]>=2.43.0, !=2.44.* # Ray Compiled Graph, required by pipeline parallelism tests
20
+ray[cgraph,default]>=2.43.0, !=2.44.* # Ray Compiled Graph, required by pipeline parallelism tests
21
sentence-transformers # required for embedding tests
22
soundfile # required for audio tests
23
jiwer # required for audio tests
requirements/test.txt
+50
@@ -10,9 +10,13 @@ aiohappyeyeballs==2.4.3
10
# via aiohttp
11
aiohttp==3.10.11
12
# via
13
+ # aiohttp-cors
14
# datasets
15
# fsspec
16
# lm-eval
17
+ # ray
18
+aiohttp-cors==0.8.1
19
+ # via ray
20
aiosignal==1.3.1
21
# via
22
# aiohttp
@@ -57,6 +61,8 @@ bounded-pool-executor==0.0.3
61
# via pqdm
62
buildkite-test-collector==0.1.9
63
# via -r requirements/test.in
64
+cachetools==5.5.2
65
+ # via google-auth
66
certifi==2024.8.30
67
# via
68
# httpcore
@@ -81,6 +87,8 @@ colorama==0.4.6
87
# sacrebleu
88
# schemathesis
89
# tqdm-multiprocess
90
+colorful==0.5.6
91
+ # via ray
92
contourpy==1.3.0
93
# via matplotlib
94
cramjam==2.9.0
@@ -108,6 +116,8 @@ dill==0.3.8
116
# evaluate
117
# lm-eval
118
# multiprocess
119
+distlib==0.3.9
120
+ # via virtualenv
121
dnspython==2.7.0
122
# via email-validator
123
docopt==0.6.2
@@ -143,6 +153,7 @@ filelock==3.16.1
153
# ray
154
# torch
155
# transformers
156
+ # virtualenv
157
fonttools==4.54.1
158
# via matplotlib
159
fqdn==1.5.1
@@ -165,8 +176,16 @@ genai-perf==0.0.8
176
# via -r requirements/test.in
177
genson==1.3.0
178
# via datamodel-code-generator
179
+google-api-core==2.24.2
180
+ # via opencensus
181
+google-auth==2.40.2
182
+ # via google-api-core
183
+googleapis-common-protos==1.70.0
184
+ # via google-api-core
185
graphql-core==3.2.6
186
# via hypothesis-graphql
187
+grpcio==1.71.0
188
+ # via ray
189
h11==0.14.0
190
# via httpcore
191
harfile==0.3.0
@@ -392,6 +411,10 @@ nvidia-nvjitlink-cu12==12.8.61
411
# torch
412
nvidia-nvtx-cu12==12.8.55
413
# via torch
414
+opencensus==0.11.4
415
+ # via ray
416
+opencensus-context==0.1.3
417
+ # via opencensus
418
opencv-python-headless==4.11.0.86
419
# via
420
# -r requirements/test.in
@@ -445,6 +468,7 @@ platformdirs==4.3.6
468
# via
469
# black
470
# pooch
471
+ # virtualenv
472
plotly==5.24.1
473
# via genai-perf
474
pluggy==1.5.0
@@ -457,10 +481,17 @@ portalocker==2.10.1
481
# via sacrebleu
482
pqdm==0.2.0
483
# via -r requirements/test.in
484
+prometheus-client==0.22.0
485
+ # via ray
486
propcache==0.2.0
487
# via yarl
488
+proto-plus==1.26.1
489
+ # via google-api-core
490
protobuf==5.28.3
491
# via
492
+ # google-api-core
493
+ # googleapis-common-protos
494
+ # proto-plus
495
# ray
496
# tensorizer
497
psutil==6.1.0
@@ -470,10 +501,18 @@ psutil==6.1.0
501
# tensorizer
502
py==1.11.0
503
# via pytest-forked
504
+py-spy==0.4.0
505
+ # via ray
506
pyarrow==18.0.0
507
# via
508
# datasets
509
# genai-perf
510
+pyasn1==0.6.1
511
+ # via
512
+ # pyasn1-modules
513
+ # rsa
514
+pyasn1-modules==0.4.2
515
+ # via google-auth
516
pybind11==2.13.6
517
# via lm-eval
518
pycparser==2.22
@@ -486,6 +525,7 @@ pydantic==2.11.5
525
# datamodel-code-generator
526
# mistral-common
527
# mteb
528
+ # ray
529
pydantic-core==2.33.2
530
# via pydantic
531
pygments==2.18.0
@@ -573,6 +613,7 @@ requests==2.32.3
613
# buildkite-test-collector
614
# datasets
615
# evaluate
616
+ # google-api-core
617
# huggingface-hub
618
# lm-eval
619
# mistral-common
@@ -601,6 +642,8 @@ rpds-py==0.20.1
642
# via
643
# jsonschema
644
# referencing
645
+rsa==4.9.1
646
+ # via google-auth
647
runai-model-streamer==0.11.0
648
# via -r requirements/test.in
649
runai-model-streamer-s3==0.11.0
@@ -648,9 +691,12 @@ shellingham==1.5.4
691
six==1.16.0
692
# via
693
# junit-xml
694
+ # opencensus
695
# python-dateutil
696
# rfc3339-validator
697
# rouge-score
698
+smart-open==7.1.0
699
+ # via ray
700
sniffio==1.3.1
701
# via
702
# anyio
@@ -801,6 +847,8 @@ urllib3==2.2.3
847
# tritonclient
848
vector-quantize-pytorch==1.21.2
849
# via -r requirements/test.in
850
+virtualenv==20.31.2
851
+ # via ray
852
vocos==0.1.0
853
# via -r requirements/test.in
854
webcolors==24.11.1
@@ -809,6 +857,8 @@ werkzeug==3.1.3
857
# via schemathesis
858
word2number==1.1
859
# via lm-eval
860
+wrapt==1.17.2
861
+ # via smart-open
862
xxhash==3.5.0
863
# via
864
# datasets
tests/v1/test_async_llm_dp.py
+10
-3
@@ -59,14 +59,22 @@ async def generate(engine: AsyncLLM,
59
60
61
@pytest.mark.parametrize(
62
- "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])
62
+ "output_kind",
63
+ [
64
+ RequestOutputKind.DELTA,
65
+ RequestOutputKind.FINAL_ONLY,
66
+ ],
67
+)
68
+@pytest.mark.parametrize("data_parallel_backend", ["mp", "ray"])
69
@pytest.mark.asyncio
64
-async def test_load(output_kind: RequestOutputKind):
70
+async def test_load(output_kind: RequestOutputKind,
71
+ data_parallel_backend: str):
72
73
with ExitStack() as after:
74
75
prompt = "This is a test of data parallel"
76
77
+ engine_args.data_parallel_backend = data_parallel_backend
78
engine = AsyncLLM.from_engine_args(engine_args)
79
after.callback(engine.shutdown)
80
@@ -82,7 +90,6 @@ async def test_load(output_kind: RequestOutputKind):
90
asyncio.create_task(
91
generate(engine, request_id, prompt, output_kind,
92
NUM_EXPECTED_TOKENS)))
85
-
93
# Confirm that we got all the EXPECTED tokens from the requests.
94
done, pending = await asyncio.wait(tasks,
95
return_when=asyncio.FIRST_EXCEPTION)
vllm/config.py
+6
@@ -1742,6 +1742,8 @@ class ParallelConfig:
1742
"""Port for data parallel messaging."""
1743
data_parallel_master_port: int = 29500
1744
"""Port of the data parallel master."""
1745
+ data_parallel_backend: str = "mp"
1746
+ """Backend to use for data parallel, either "mp" or "ray"."""
1747
enable_expert_parallel: bool = False
1748
"""Use expert parallelism instead of tensor parallelism for MoE layers."""
1749
max_parallel_loading_workers: Optional[int] = None
@@ -1911,6 +1913,10 @@ class ParallelConfig:
1913
"please install Ray with `pip install "
1914
"ray`.") from ray_utils.ray_import_err
1915
backend = "ray"
1916
+ elif self.data_parallel_backend == "ray":
1917
+ logger.info("Using ray distributed inference because "
1918
+ "data_parallel_backend is ray")
1919
+ backend = "ray"
1920
elif ray_found:
1921
if self.placement_group:
1922
backend = "ray"
vllm/engine/arg_utils.py
+25
-4
@@ -39,7 +39,7 @@ from vllm.test_utils import MODEL_WEIGHTS_S3_BUCKET, MODELS_ON_S3
39
from vllm.transformers_utils.utils import check_gguf_file
40
from vllm.usage.usage_lib import UsageContext
41
from vllm.utils import (STR_DUAL_CHUNK_FLASH_ATTN_VAL, FlexibleArgumentParser,
42
- GiB_bytes, is_in_ray_actor)
42
+ GiB_bytes, get_ip, is_in_ray_actor)
43
44
# yapf: enable
45
@@ -292,6 +292,7 @@ class EngineArgs:
292
data_parallel_size_local: Optional[int] = None
293
data_parallel_address: Optional[str] = None
294
data_parallel_rpc_port: Optional[int] = None
295
+ data_parallel_backend: str = ParallelConfig.data_parallel_backend
296
enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
297
max_parallel_loading_workers: Optional[
298
int] = ParallelConfig.max_parallel_loading_workers
@@ -624,6 +625,12 @@ class EngineArgs:
625
type=int,
626
help='Port for data parallel RPC '
627
'communication.')
628
+ parallel_group.add_argument('--data-parallel-backend',
629
+ '-dpb',
630
+ type=str,
631
+ default='mp',
632
+ help='Backend for data parallel, either '
633
+ '"mp" or "ray".')
634
parallel_group.add_argument(
635
"--enable-expert-parallel",
636
**parallel_kwargs["enable_expert_parallel"])
@@ -1059,9 +1066,20 @@ class EngineArgs:
1066
1067
# DP address, used in multi-node case for torch distributed group
1068
# and ZMQ sockets.
1062
- data_parallel_address = self.data_parallel_address if (
1063
- self.data_parallel_address
1064
- is not None) else ParallelConfig.data_parallel_master_ip
1069
+ if self.data_parallel_address is None:
1070
+ if self.data_parallel_backend == "ray":
1071
+ host_ip = get_ip()
1072
+ logger.info(
1073
+ "Using host IP %s as ray-based data parallel address",
1074
+ host_ip)
1075
+ data_parallel_address = host_ip
1076
+ else:
1077
+ assert self.data_parallel_backend == "mp", (
1078
+ "data_parallel_backend can only be ray or mp, got %s",
1079
+ self.data_parallel_backend)
1080
+ data_parallel_address = ParallelConfig.data_parallel_master_ip
1081
+ else:
1082
+ data_parallel_address = self.data_parallel_address
1083
1084
# This port is only used when there are remote data parallel engines,
1085
# otherwise the local IPC transport is used.
@@ -1069,6 +1087,8 @@ class EngineArgs:
1087
self.data_parallel_rpc_port
1088
is not None) else ParallelConfig.data_parallel_rpc_port
1089
1090
+ data_parallel_backend = self.data_parallel_backend
1091
+
1092
parallel_config = ParallelConfig(
1093
pipeline_parallel_size=self.pipeline_parallel_size,
1094
tensor_parallel_size=self.tensor_parallel_size,
@@ -1076,6 +1096,7 @@ class EngineArgs:
1096
data_parallel_size_local=data_parallel_size_local,
1097
data_parallel_master_ip=data_parallel_address,
1098
data_parallel_rpc_port=data_parallel_rpc_port,
1099
+ data_parallel_backend=data_parallel_backend,
1100
enable_expert_parallel=self.enable_expert_parallel,
1101
max_parallel_loading_workers=self.max_parallel_loading_workers,
1102
disable_custom_all_reduce=self.disable_custom_all_reduce,
vllm/entrypoints/cli/serve.py
+30
-5
@@ -27,7 +27,8 @@ from vllm.v1.engine.core_client import CoreEngineProcManager
27
from vllm.v1.executor.abstract import Executor
28
from vllm.v1.metrics.prometheus import setup_multiprocess_prometheus
29
from vllm.v1.utils import (APIServerProcessManager, CoreEngine,
30
- EngineZmqAddresses, get_engine_client_zmq_addr,
30
+ CoreEngineActorManager, EngineZmqAddresses,
31
+ get_engine_client_zmq_addr,
32
wait_for_completion_or_failure,
33
wait_for_engine_startup)
34
@@ -229,6 +230,31 @@ def run_multi_api_server(args: argparse.Namespace):
230
logger.info("Started DP Coordinator process (PID: %d)",
231
coordinator.proc.pid)
232
233
+ if parallel_config.data_parallel_backend == "ray":
234
+ logger.info("Starting ray-based data parallel backend")
235
+
236
+ engine_actor_manager = CoreEngineActorManager(
237
+ vllm_config=vllm_config,
238
+ addresses=addresses,
239
+ executor_class=Executor.get_class(vllm_config),
240
+ log_stats=not engine_args.disable_log_stats,
241
+ )
242
+ # Start API servers using the manager
243
+ api_server_manager = APIServerProcessManager(
244
+ target_server_fn=run_api_server_worker_proc,
245
+ listen_address=listen_address,
246
+ sock=sock,
247
+ args=args,
248
+ num_servers=num_api_servers,
249
+ input_addresses=input_addresses,
250
+ output_addresses=output_addresses,
251
+ stats_update_address=stats_update_address)
252
+
253
+ wait_for_completion_or_failure(api_server_manager=api_server_manager,
254
+ engine_manager=engine_actor_manager,
255
+ coordinator=coordinator)
256
+ return
257
+
258
handshake_address = get_engine_client_zmq_addr(
259
local_only, host, parallel_config.data_parallel_rpc_port)
260
@@ -277,10 +303,9 @@ def run_multi_api_server(args: argparse.Namespace):
303
)
304
305
# Wait for API servers
280
- wait_for_completion_or_failure(
281
- api_server_manager=api_server_manager,
282
- local_engine_manager=local_engine_manager,
283
- coordinator=coordinator)
306
+ wait_for_completion_or_failure(api_server_manager=api_server_manager,
307
+ engine_manager=local_engine_manager,
308
+ coordinator=coordinator)
309
310
311
def run_api_server_worker_proc(listen_address,
vllm/v1/engine/async_llm.py
+9
-4
@@ -27,7 +27,8 @@ from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
27
from vllm.usage.usage_lib import UsageContext
28
from vllm.utils import Device, cdiv
29
from vllm.v1.engine import EngineCoreRequest
30
-from vllm.v1.engine.core_client import AsyncMPClient, DPAsyncMPClient
30
+from vllm.v1.engine.core_client import (AsyncMPClient, DPAsyncMPClient,
31
+ RayDPClient)
32
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
33
from vllm.v1.engine.output_processor import (OutputProcessor,
34
RequestOutputCollector)
@@ -119,9 +120,13 @@ class AsyncLLM(EngineClient):
120
log_stats=self.log_stats)
121
122
# EngineCore (starts the engine in background process).
122
- core_client_class = AsyncMPClient if (
123
- vllm_config.parallel_config.data_parallel_size
124
- == 1) else DPAsyncMPClient
123
+ core_client_class: type[AsyncMPClient]
124
+ if vllm_config.parallel_config.data_parallel_size == 1:
125
+ core_client_class = AsyncMPClient
126
+ elif vllm_config.parallel_config.data_parallel_backend == "ray":
127
+ core_client_class = RayDPClient
128
+ else:
129
+ core_client_class = DPAsyncMPClient
130
131
self.engine_core = core_client_class(
132
vllm_config=vllm_config,
vllm/v1/engine/core.py
+123
-45
@@ -6,8 +6,9 @@ import sys
6
import threading
7
import time
8
from collections import deque
9
+from collections.abc import Generator
10
from concurrent.futures import Future
10
-from contextlib import ExitStack
11
+from contextlib import ExitStack, contextmanager
12
from inspect import isclass, signature
13
from logging import DEBUG
14
from typing import Any, Callable, Optional, TypeVar, Union
@@ -367,60 +368,35 @@ class EngineCoreProc(EngineCore):
368
log_stats: bool,
369
engine_index: int = 0,
370
):
370
- input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]()
371
-
372
- executor_fail_callback = lambda: input_queue.put_nowait(
371
+ self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]()
372
+ self.output_queue = queue.Queue[Union[tuple[int, EngineCoreOutputs],
373
+ bytes]]()
374
+ executor_fail_callback = lambda: self.input_queue.put_nowait(
375
(EngineCoreRequestType.EXECUTOR_FAILED, b''))
376
375
- # Create input socket.
376
- input_ctx = zmq.Context()
377
- identity = engine_index.to_bytes(length=2, byteorder="little")
378
- with make_zmq_socket(input_ctx,
379
- handshake_address,
380
- zmq.DEALER,
381
- identity=identity,
382
- linger=5000,
383
- bind=False) as handshake_socket:
377
+ self.engine_index = engine_index
378
+ identity = self.engine_index.to_bytes(length=2, byteorder="little")
379
+ self.engines_running = False
380
385
- # Register engine with front-end.
386
- addresses = self.startup_handshake(handshake_socket, on_head_node,
387
- vllm_config.parallel_config)
381
+ with self._perform_handshake(handshake_address, identity, on_head_node,
382
+ vllm_config) as addresses:
383
self.client_count = len(addresses.outputs)
384
390
- # Update config which may have changed from the handshake.
391
- vllm_config.__post_init__()
392
-
385
# Set up data parallel environment.
386
self.has_coordinator = addresses.coordinator_output is not None
387
self._init_data_parallel(vllm_config)
388
397
- # Initialize engine core and model.
389
super().__init__(vllm_config, executor_class, log_stats,
390
executor_fail_callback)
391
401
- self.engine_index = engine_index
402
- self.step_fn = (self.step if self.batch_queue is None else
403
- self.step_with_batch_queue)
404
- self.engines_running = False
405
- self.last_counts = (0, 0)
406
-
407
- # Send ready message.
408
- num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
409
- handshake_socket.send(
410
- msgspec.msgpack.encode({
411
- "status": "READY",
412
- "local": on_head_node,
413
- "num_gpu_blocks": num_gpu_blocks,
414
- }))
392
+ self.step_fn = (self.step if self.batch_queue is None else
393
+ self.step_with_batch_queue)
394
395
# Background Threads and Queues for IO. These enable us to
396
# overlap ZMQ socket IO with GPU since they release the GIL,
397
# and to overlap some serialization/deserialization with the
398
# model forward pass.
399
# Threads handle Socket <-> Queues and core_busy_loop uses Queue.
421
- self.input_queue = input_queue
422
- self.output_queue = queue.Queue[Union[tuple[int, EngineCoreOutputs],
423
- bytes]]()
400
threading.Thread(target=self.process_input_sockets,
401
args=(addresses.inputs, addresses.coordinator_input,
402
identity),
@@ -428,10 +404,40 @@ class EngineCoreProc(EngineCore):
404
self.output_thread = threading.Thread(
405
target=self.process_output_sockets,
406
args=(addresses.outputs, addresses.coordinator_output,
431
- engine_index),
407
+ self.engine_index),
408
daemon=True)
409
self.output_thread.start()
410
411
+ @contextmanager
412
+ def _perform_handshake(
413
+ self, handshake_address: str, identity: bytes, on_head_node: bool,
414
+ vllm_config: VllmConfig
415
+ ) -> Generator[EngineZmqAddresses, None, None]:
416
+ input_ctx = zmq.Context()
417
+ with make_zmq_socket(input_ctx,
418
+ handshake_address,
419
+ zmq.DEALER,
420
+ identity=identity,
421
+ linger=5000,
422
+ bind=False) as handshake_socket:
423
+ # Register engine with front-end.
424
+ addresses = self.startup_handshake(handshake_socket, on_head_node,
425
+ vllm_config.parallel_config)
426
+
427
+ # Update config which may have changed from the handshake
428
+ vllm_config.__post_init__()
429
+
430
+ yield addresses
431
+
432
+ # Send ready message.
433
+ num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
434
+ handshake_socket.send(
435
+ msgspec.msgpack.encode({
436
+ "status": "READY",
437
+ "local": on_head_node,
438
+ "num_gpu_blocks": num_gpu_blocks,
439
+ }))
440
+
441
@staticmethod
442
def startup_handshake(
443
handshake_socket: zmq.Socket, on_head_node: bool,
@@ -743,24 +749,29 @@ class DPEngineCoreProc(EngineCoreProc):
749
executor_class: type[Executor],
750
log_stats: bool,
751
):
746
- # Add process-specific prefix to stdout and stderr before
747
- # we initialize the engine.
748
- from multiprocessing import current_process
749
- process_name = current_process().name
750
- pid = os.getpid()
751
- _add_prefix(sys.stdout, process_name, pid)
752
- _add_prefix(sys.stderr, process_name, pid)
752
+
753
+ self._decorate_logs()
754
755
# Counts forward-passes of the model so that we can synchronize
756
# finished with DP peers every N steps.
757
self.counter = 0
758
self.current_wave = 0
759
+ self.last_counts = (0, 0)
760
761
# Initialize the engine.
762
dp_rank = vllm_config.parallel_config.data_parallel_rank
763
super().__init__(vllm_config, on_head_node, handshake_address,
764
executor_class, log_stats, dp_rank)
765
766
+ def _decorate_logs(self):
767
+ # Add process-specific prefix to stdout and stderr before
768
+ # we initialize the engine.
769
+ from multiprocessing import current_process
770
+ process_name = current_process().name
771
+ pid = os.getpid()
772
+ _add_prefix(sys.stdout, process_name, pid)
773
+ _add_prefix(sys.stderr, process_name, pid)
774
+
775
def _init_data_parallel(self, vllm_config: VllmConfig):
776
777
# Configure GPUs and stateless process group for data parallel.
@@ -880,3 +891,70 @@ class DPEngineCoreProc(EngineCoreProc):
891
892
return ParallelConfig.has_unfinished_dp(self.dp_group,
893
local_unfinished)
894
+
895
+
896
+class DPEngineCoreActor(DPEngineCoreProc):
897
+ """
898
+ Ray actor for running EngineCore in a data parallel context
899
+ """
900
+
901
+ def __init__(
902
+ self,
903
+ vllm_config: VllmConfig,
904
+ on_head_node: bool,
905
+ addresses: EngineZmqAddresses,
906
+ executor_class: type[Executor],
907
+ log_stats: bool,
908
+ dp_rank: int = 0,
909
+ local_dp_rank: int = 0,
910
+ ):
911
+ self.addresses = addresses
912
+ vllm_config.parallel_config.data_parallel_rank = dp_rank
913
+ vllm_config.parallel_config.data_parallel_rank_local = \
914
+ local_dp_rank
915
+
916
+ # Ray sets CUDA_VISIBLE_DEVICES to empty string,
917
+ # we clean this up to be able to properly initialize
918
+ # data parallel groups.
919
+ del os.environ['CUDA_VISIBLE_DEVICES']
920
+
921
+ super().__init__(vllm_config, on_head_node, "", executor_class,
922
+ log_stats)
923
+
924
+ def _decorate_logs(self):
925
+ pass
926
+
927
+ @contextmanager
928
+ def _perform_handshake(self, handshake_address: str, identity: bytes,
929
+ on_head_node: bool, vllm_config: VllmConfig):
930
+ """
931
+ For Ray, we don't need to actually perform handshake.
932
+ All addresses information is known before the actor creation.
933
+ Therefore, we simply yield these addresses.
934
+ """
935
+ yield self.addresses
936
+
937
+ def wait_for_init(self):
938
+ """
939
+ Wait until the engine core is initialized.
940
+
941
+ This is just an empty method. When ray.get() on this method
942
+ (or any other method of the actor) returns, it is guaranteed
943
+ that actor creation (i.e., __init__) is complete.
944
+ """
945
+ pass
946
+
947
+ def run(self):
948
+ """
949
+ Run the engine core busy loop.
950
+ """
951
+ try:
952
+ self.run_busy_loop()
953
+ except SystemExit:
954
+ logger.debug("EngineCore exiting.")
955
+ raise
956
+ except Exception:
957
+ logger.exception("EngineCore encountered a fatal error.")
958
+ raise
959
+ finally:
960
+ self.shutdown()
vllm/v1/engine/core_client.py
+65
-9
@@ -29,9 +29,9 @@ from vllm.v1.engine.core import EngineCore, EngineCoreProc
29
from vllm.v1.engine.exceptions import EngineDeadError
30
from vllm.v1.executor.abstract import Executor
31
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr
32
-from vllm.v1.utils import (CoreEngine, CoreEngineProcManager,
33
- EngineZmqAddresses, get_engine_client_zmq_addr,
34
- wait_for_engine_startup)
32
+from vllm.v1.utils import (CoreEngine, CoreEngineActorManager,
33
+ CoreEngineProcManager, EngineZmqAddresses,
34
+ get_engine_client_zmq_addr, wait_for_engine_startup)
35
36
logger = init_logger(__name__)
37
@@ -68,6 +68,8 @@ class EngineCoreClient(ABC):
68
69
if multiprocess_mode and asyncio_mode:
70
if vllm_config.parallel_config.data_parallel_size > 1:
71
+ if vllm_config.parallel_config.data_parallel_backend == "ray":
72
+ return RayDPClient(vllm_config, executor_class, log_stats)
73
return DPAsyncMPClient(vllm_config, executor_class, log_stats)
74
75
return AsyncMPClient(vllm_config, executor_class, log_stats)
@@ -273,7 +275,10 @@ class BackgroundResources:
275
circular reference back to the client object."""
276
277
ctx: Union[zmq.Context]
276
- local_engine_manager: Optional[CoreEngineProcManager] = None
278
+ # If CoreEngineProcManager, it manages local engines;
279
+ # if CoreEngineActorManager, it manages all engines.
280
+ engine_manager: Optional[Union[CoreEngineProcManager,
281
+ CoreEngineActorManager]] = None
282
coordinator: Optional[DPCoordinator] = None
283
output_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None
284
input_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None
@@ -290,8 +295,8 @@ class BackgroundResources:
295
"""Clean up background resources."""
296
297
self.engine_dead = True
293
- if self.local_engine_manager is not None:
294
- self.local_engine_manager.close()
298
+ if self.engine_manager is not None:
299
+ self.engine_manager.close()
300
if self.coordinator is not None:
301
self.coordinator.close()
302
@@ -457,7 +462,7 @@ class MPClient(EngineCoreClient):
462
if local_engine_count:
463
# In server mode, start_index and local_start_index will
464
# both be 0.
460
- self.resources.local_engine_manager = CoreEngineProcManager(
465
+ self.resources.engine_manager = CoreEngineProcManager(
466
EngineCoreProc.run_engine_core,
467
vllm_config=vllm_config,
468
executor_class=executor_class,
@@ -484,13 +489,18 @@ class MPClient(EngineCoreClient):
489
addresses.coordinator_input, addresses.coordinator_output = (
490
coordinator.get_engine_socket_addresses())
491
492
+ proc_manager = self.resources.engine_manager
493
+ assert isinstance(proc_manager, (type(None), CoreEngineProcManager)), (
494
+ "_wait_for_engine_startup should only be "
495
+ "called with CoreEngineProcManager")
496
+
497
wait_for_engine_startup(
498
handshake_socket,
499
addresses,
500
self.core_engines,
501
self.vllm_config.parallel_config,
502
self.vllm_config.cache_config,
493
- self.resources.local_engine_manager,
503
+ proc_manager,
504
coordinator.proc if coordinator else None,
505
)
506
@@ -887,7 +897,6 @@ class DPAsyncMPClient(AsyncMPClient):
897
log_stats: bool,
898
client_addresses: Optional[dict[str, str]] = None,
899
client_index: int = 0):
890
-
900
self.current_wave = 0
901
self.engines_running = False
902
# To route aborts to the correct engine.
@@ -1050,3 +1059,50 @@ class DPAsyncMPClient(AsyncMPClient):
1059
if not self.resources.engine_dead:
1060
await self._send_input(EngineCoreRequestType.ABORT, request_ids,
1061
engine)
1062
+
1063
+
1064
+class RayDPClient(DPAsyncMPClient):
1065
+ """
1066
+ Ray-based client for multi-proc, multi-engine (data parallel)
1067
+ EngineCore.
1068
+ """
1069
+
1070
+ def __init__(
1071
+ self,
1072
+ vllm_config: VllmConfig,
1073
+ executor_class: type[Executor],
1074
+ log_stats: bool,
1075
+ client_addresses: Optional[dict[str, str]] = None,
1076
+ client_index: int = 0,
1077
+ ):
1078
+ super().__init__(vllm_config, executor_class, log_stats,
1079
+ client_addresses, client_index)
1080
+
1081
+ def _init_engines_direct(self, vllm_config: VllmConfig, local_only: bool,
1082
+ local_start_index: int, input_address: str,
1083
+ output_address: str,
1084
+ executor_class: type[Executor], log_stats: bool):
1085
+ """Self-contained client mode, launch engine and coordinator process
1086
+ as needed."""
1087
+
1088
+ parallel_config = vllm_config.parallel_config
1089
+ assert parallel_config.data_parallel_rank == 0
1090
+ assert local_start_index == 0
1091
+
1092
+ addresses = EngineZmqAddresses(
1093
+ inputs=[input_address],
1094
+ outputs=[output_address],
1095
+ )
1096
+
1097
+ if len(self.core_engines) > 1:
1098
+ coordinator = DPCoordinator(parallel_config)
1099
+ self.resources.coordinator = coordinator
1100
+ addresses.coordinator_input, addresses.coordinator_output = (
1101
+ coordinator.get_engine_socket_addresses())
1102
+
1103
+ # Start all engines.
1104
+ self.resources.engine_manager = CoreEngineActorManager(
1105
+ vllm_config=vllm_config,
1106
+ addresses=addresses,
1107
+ executor_class=executor_class,
1108
+ log_stats=log_stats)
vllm/v1/utils.py
+220
-37
@@ -27,6 +27,8 @@ from vllm.utils import (get_mp_context, get_open_port, get_open_zmq_ipc_path,
27
from vllm.v1.executor.abstract import Executor
28
29
if TYPE_CHECKING:
30
+ from ray.util.placement_group import PlacementGroup
31
+
32
from vllm.attention.layer import Attention
33
from vllm.v1.engine.coordinator import DPCoordinator
34
@@ -112,6 +114,45 @@ def get_engine_client_zmq_addr(local_only: bool,
114
host, port or get_open_port()))
115
116
117
+class CoreEngineState(Enum):
118
+ NEW = auto()
119
+ CONNECTED = auto()
120
+ READY = auto()
121
+
122
+
123
+class CoreEngine:
124
+ """One per data parallel rank."""
125
+
126
+ def __init__(self, index: int = 0, local: bool = True):
127
+ self.local = local
128
+ self.index = index
129
+ self.identity = index.to_bytes(2, "little")
130
+
131
+ self.state = CoreEngineState.NEW
132
+
133
+
134
+@dataclass
135
+class EngineZmqAddresses:
136
+ # ZMQ input socket addresses for each front-end client (requests)
137
+ inputs: list[str]
138
+ # ZMQ output socket addresses for each front-end client (responses)
139
+ outputs: list[str]
140
+ # ZMQ input socket address of DP coordinator if applicable
141
+ coordinator_input: Optional[str] = None
142
+ # ZMQ output socket address of DP coordinator if applicable
143
+ coordinator_output: Optional[str] = None
144
+
145
+
146
+@dataclass
147
+class EngineHandshakeMetadata:
148
+ """Metadata sent to each engine process during startup handshake,
149
+ including addresses of the front-end ZMQ queues that they should
150
+ connect to.
151
+ """
152
+ addresses: EngineZmqAddresses
153
+ parallel_config: dict[str, Union[int, str]]
154
+
155
+
156
class APIServerProcessManager:
157
"""Manages a group of API server processes.
158
@@ -245,43 +286,168 @@ class CoreEngineProcManager:
286
}
287
288
248
-class CoreEngineState(Enum):
249
- NEW = auto()
250
- CONNECTED = auto()
251
- READY = auto()
252
-
289
+class CoreEngineActorManager:
290
+ """
291
+ Utility class to handle creation, readiness, and shutdown
292
+ of core engine Ray actors used by the AsyncLLM and LLMEngine.
293
254
-class CoreEngine:
255
- """One per data parallel rank."""
294
+ Different from CoreEngineProcManager, this class manages
295
+ core engines for both local and remote nodes.
296
+ """
297
257
- def __init__(self, index: int = 0, local: bool = True):
258
- self.local = local
259
- self.index = index
260
- self.identity = index.to_bytes(2, "little")
298
+ def __init__(
299
+ self,
300
+ vllm_config: VllmConfig,
301
+ addresses: EngineZmqAddresses,
302
+ executor_class: type[Executor],
303
+ log_stats: bool,
304
+ placement_groups: Optional[list["PlacementGroup"]] = None,
305
+ local_dp_ranks: Optional[list[int]] = None,
306
+ ):
307
+ import copy
308
262
- self.state = CoreEngineState.NEW
309
+ import ray
310
+ from ray.util.scheduling_strategies import (
311
+ PlacementGroupSchedulingStrategy)
312
313
+ from vllm.v1.engine.core import DPEngineCoreActor
314
265
-@dataclass
266
-class EngineZmqAddresses:
267
- # ZMQ input socket addresses for each front-end client (requests)
268
- inputs: list[str]
269
- # ZMQ output socket addresses for each front-end client (responses)
270
- outputs: list[str]
271
- # ZMQ input socket address of DP coordinator if applicable
272
- coordinator_input: Optional[str] = None
273
- # ZMQ output socket address of DP coordinator if applicable
274
- coordinator_output: Optional[str] = None
315
+ self.local_engine_actors: list[ray.ActorHandle] = []
316
+ self.remote_engine_actors: list[ray.ActorHandle] = []
317
+ dp_size = vllm_config.parallel_config.data_parallel_size
318
+ local_engine_count = \
319
+ vllm_config.parallel_config.data_parallel_size_local
320
+ world_size = vllm_config.parallel_config.world_size
321
322
+ if ray.is_initialized():
323
+ logger.info(
324
+ "Ray is already initialized. Skipping Ray initialization.")
325
+ else:
326
+ ray.init()
327
+
328
+ if placement_groups is not None:
329
+ assert local_dp_ranks is not None, (
330
+ "local_dp_ranks must be provided if "
331
+ "placement_groups is provided")
332
+ assert len(placement_groups) == len(local_dp_ranks), (
333
+ "placement_groups and local_dp_ranks must "
334
+ "have the same length")
335
+ logger.info("Using provided placement groups")
336
+ # TODO(rui): validate passed-in placement groups
337
+ self.created_placement_groups = []
338
+ else:
339
+ placement_groups, local_dp_ranks = \
340
+ CoreEngineActorManager.create_dp_placement_groups(vllm_config)
341
+ self.created_placement_groups = placement_groups
342
+ assert len(placement_groups) == dp_size, (
343
+ "Number of placement groups must match data parallel size")
344
+
345
+ refs = []
346
+ for index in range(dp_size):
347
+ local_index = local_dp_ranks[index]
348
+ dp_vllm_config = copy.deepcopy(vllm_config)
349
+ pg = placement_groups[index]
350
+ dp_vllm_config.parallel_config.placement_group = pg
351
+ on_head_node = index < local_engine_count
352
+ actor = ray.remote(DPEngineCoreActor).options(
353
+ scheduling_strategy=PlacementGroupSchedulingStrategy(
354
+ placement_group=pg,
355
+ placement_group_bundle_index=world_size,
356
+ )).remote(vllm_config=dp_vllm_config,
357
+ executor_class=executor_class,
358
+ log_stats=log_stats,
359
+ on_head_node=on_head_node,
360
+ addresses=addresses,
361
+ dp_rank=index,
362
+ local_dp_rank=local_index)
363
+ if on_head_node:
364
+ self.local_engine_actors.append(actor)
365
+ else:
366
+ self.remote_engine_actors.append(actor)
367
+ refs.append(actor.wait_for_init.remote())
368
+
369
+ ray.get(refs)
370
+ self.run_refs = []
371
+ for actor in self.local_engine_actors + self.remote_engine_actors:
372
+ self.run_refs.append(actor.run.remote())
373
+
374
+ @staticmethod
375
+ def create_dp_placement_groups(
376
+ vllm_config: VllmConfig
377
+ ) -> tuple[list["PlacementGroup"], list[int]]:
378
+
379
+ import ray
380
+ from ray._private.state import available_resources_per_node
381
+ from ray.util.state import list_nodes
382
+
383
+ logger.info("Creating placement groups for data parallel")
384
+ dp_master_ip = \
385
+ vllm_config.parallel_config.data_parallel_master_ip
386
+ dp_size = vllm_config.parallel_config.data_parallel_size
387
+ local_engine_count = \
388
+ vllm_config.parallel_config.data_parallel_size_local
389
+
390
+ nodes = list_nodes()
391
+ nodes = sorted(list_nodes(),
392
+ key=lambda node: node.node_ip != dp_master_ip)
393
+ assert nodes[0].node_ip == dp_master_ip, (
394
+ "The first node must be the head node")
395
+ assert len(nodes) == 1 or nodes[1].node_ip != dp_master_ip, (
396
+ "There can only be one head node")
397
+
398
+ available_resources = available_resources_per_node()
399
+ world_size = vllm_config.parallel_config.world_size
400
+ placement_groups: list[PlacementGroup] = []
401
+ local_dp_ranks: list[int] = []
402
+
403
+ for node in nodes:
404
+ node_ip = node.node_ip
405
+ node_resources = available_resources[node.node_id]
406
+ # For now, each DP rank can only be assigned to one node
407
+ # TODO(rui): support allocating a single DP rank
408
+ # to multiple nodes
409
+ available_engine_count = node_resources["GPU"] // world_size
410
+ if node_ip == dp_master_ip:
411
+ assert available_engine_count >= local_engine_count, (
412
+ "Not enough resources to allocate DP ranks "
413
+ f"on DP master node {node_ip}")
414
+ for i in range(local_engine_count):
415
+ bundles = [{
416
+ "GPU": 1.0,
417
+ "node:" + dp_master_ip: 0.001
418
+ }] * world_size + [{
419
+ "CPU": 1.0
420
+ }]
421
+ pg = ray.util.placement_group(
422
+ name=f"dp_rank_{len(placement_groups)}",
423
+ strategy="STRICT_PACK",
424
+ bundles=bundles,
425
+ )
426
+ placement_groups.append(pg)
427
+ local_dp_ranks.append(i)
428
+ else:
429
+ for i in range(available_engine_count):
430
+ if len(placement_groups) == dp_size:
431
+ break
432
+ bundles = [{"GPU": 1.0}] * world_size + [{"CPU": 1.0}]
433
+ pg = ray.util.placement_group(
434
+ name=f"dp_rank_{len(placement_groups)}",
435
+ strategy="STRICT_PACK",
436
+ bundles=bundles,
437
+ )
438
+ placement_groups.append(pg)
439
+ local_dp_ranks.append(i)
440
+ return placement_groups, local_dp_ranks
441
+
442
+ def get_run_refs(self):
443
+ return self.run_refs
444
277
-@dataclass
278
-class EngineHandshakeMetadata:
279
- """Metadata sent to each engine process during startup handshake,
280
- including addresses of the front-end ZMQ queues that they should
281
- connect to.
282
- """
283
- addresses: EngineZmqAddresses
284
- parallel_config: dict[str, Union[int, str]]
445
+ def close(self):
446
+ import ray
447
+ for actor in self.local_engine_actors + self.remote_engine_actors:
448
+ ray.kill(actor)
449
+ for pg in self.created_placement_groups:
450
+ ray.util.remove_placement_group(pg)
451
452
453
def wait_for_engine_startup(
@@ -383,11 +549,19 @@ def wait_for_engine_startup(
549
550
def wait_for_completion_or_failure(
551
api_server_manager: APIServerProcessManager,
386
- local_engine_manager: Optional[CoreEngineProcManager] = None,
552
+ engine_manager: Optional[Union[CoreEngineProcManager,
553
+ CoreEngineActorManager]] = None,
554
coordinator: Optional["DPCoordinator"] = None) -> None:
555
"""Wait for all processes to complete or detect if any fail.
556
557
Raises an exception if any process exits with a non-zero status.
558
+
559
+ Args:
560
+ api_server_manager: The manager for API servers.
561
+ engine_manager: The manager for engine processes.
562
+ If CoreEngineProcManager, it manages local engines;
563
+ if CoreEngineActorManager, it manages all engines.
564
+ coordinator: The coordinator for data parallel.
565
"""
566
567
try:
@@ -402,14 +576,18 @@ def wait_for_completion_or_failure(
576
if coordinator:
577
sentinel_to_proc[coordinator.proc.sentinel] = coordinator.proc
578
405
- if local_engine_manager:
406
- for proc in local_engine_manager.processes:
579
+ actor_run_refs = []
580
+ if isinstance(engine_manager, CoreEngineProcManager):
581
+ for proc in engine_manager.processes:
582
sentinel_to_proc[proc.sentinel] = proc
583
+ elif isinstance(engine_manager, CoreEngineActorManager):
584
+ actor_run_refs = engine_manager.get_run_refs()
585
586
# Check if any process terminates
410
- while sentinel_to_proc:
587
+ while sentinel_to_proc or actor_run_refs:
588
# Wait for any process to terminate
412
- ready_sentinels: list[Any] = connection.wait(sentinel_to_proc)
589
+ ready_sentinels: list[Any] = connection.wait(sentinel_to_proc,
590
+ timeout=5)
591
592
# Process any terminated processes
593
for sentinel in ready_sentinels:
@@ -420,6 +598,11 @@ def wait_for_completion_or_failure(
598
raise RuntimeError(
599
f"Process {proc.name} (PID: {proc.pid}) "
600
f"died with exit code {proc.exitcode}")
601
+
602
+ if actor_run_refs:
603
+ import ray
604
+ _, actor_run_refs = ray.wait(actor_run_refs, timeout=5)
605
+
606
except KeyboardInterrupt:
607
logger.info("Received KeyboardInterrupt, shutting down API servers...")
608
except Exception as e:
@@ -431,8 +614,8 @@ def wait_for_completion_or_failure(
614
api_server_manager.close()
615
if coordinator:
616
coordinator.close()
434
- if local_engine_manager:
435
- local_engine_manager.close()
617
+ if engine_manager:
618
+ engine_manager.close()
619
620
621
# Note(rob): shutdown function cannot be a bound method,