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, 03:30 UTC
[DP] Support external DP Load Balancer mode (#19790)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Jul 2, 2025 at 18:21 UTC
657f2f301a431542a731719fa8c6326deacc317d
11 files changed
+1250
-783
.buildkite/test-pipeline.yaml
+5
-1
@@ -155,6 +155,7 @@ steps:
155
- examples/offline_inference/rlhf_colocate.py
156
- tests/examples/offline_inference/data_parallel.py
157
- tests/v1/test_async_llm_dp.py
158
+ - tests/v1/test_external_lb_dp.py
159
- tests/v1/engine/test_engine_core_client.py
160
commands:
161
# test with tp=2 and external_dp=2
@@ -163,8 +164,9 @@ steps:
164
# test with tp=2 and pp=2
165
- PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
166
# test with internal dp
166
- - python3 ../examples/offline_inference/data_parallel.py
167
+ - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
168
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
169
+ - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_external_lb_dp.py
170
- pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
171
- pytest -v -s distributed/test_utils.py
172
- pytest -v -s compile/test_basic_correctness.py
@@ -682,10 +684,12 @@ steps:
684
- vllm/worker/model_runner.py
685
- entrypoints/llm/test_collective_rpc.py
686
- tests/v1/test_async_llm_dp.py
687
+ - tests/v1/test_external_lb_dp.py
688
- tests/v1/entrypoints/openai/test_multi_api_servers.py
689
- vllm/v1/engine/
690
commands:
691
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
692
+ - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_external_lb_dp.py
693
- DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
694
- pytest -v -s entrypoints/llm/test_collective_rpc.py
695
- pytest -v -s ./compile/test_basic_correctness.py
tests/v1/engine/test_engine_core_client.py
+3
-3
@@ -26,8 +26,8 @@ from vllm.v1.engine import EngineCoreRequest
26
from vllm.v1.engine.core import EngineCore
27
from vllm.v1.engine.core_client import (AsyncMPClient, EngineCoreClient,
28
SyncMPClient)
29
+from vllm.v1.engine.utils import CoreEngineProcManager
30
from vllm.v1.executor.abstract import Executor
30
-from vllm.v1.utils import CoreEngineProcManager
31
32
from ...distributed.conftest import MockSubscriber
33
from ...utils import create_new_process_for_each_test
@@ -563,7 +563,7 @@ def test_engine_core_proc_instantiation_cuda_empty(
563
m.setenv("VLLM_USE_V1", "1")
564
m.setenv("CUDA_VISIBLE_DEVICES", "") # No CUDA devices
565
566
- from vllm.v1.utils import EngineZmqAddresses
566
+ from vllm.v1.engine.utils import EngineZmqAddresses
567
568
def mock_startup_handshake(self, handshake_socket, on_head_node,
569
parallel_config):
@@ -580,7 +580,7 @@ def test_engine_core_proc_instantiation_cuda_empty(
580
trust_remote_code=True).create_engine_config()
581
engine_core_proc = EngineCoreProc(
582
vllm_config=vllm_config,
583
- on_head_node=True,
583
+ local_client=True,
584
handshake_address="tcp://127.0.0.1:12345",
585
executor_class=mock_executor_class,
586
log_stats=False,
tests/v1/test_external_lb_dp.py
new
+312
@@ -0,0 +1,312 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import asyncio
4
+import os
5
+import threading
6
+import time
7
+from contextlib import AsyncExitStack
8
+
9
+import openai # use the official client for correctness check
10
+import pytest
11
+import pytest_asyncio
12
+
13
+from tests.utils import RemoteOpenAIServer
14
+from vllm.platforms import Platform
15
+
16
+MODEL_NAME = "ibm-research/PowerMoE-3b"
17
+
18
+# Number of data parallel ranks for external LB testing
19
+DP_SIZE = int(os.getenv("DP_SIZE", "2"))
20
+# Default tensor parallell size to use
21
+TP_SIZE = int(os.getenv("TP_SIZE", "1"))
22
+
23
+
24
+class ExternalLBServerManager:
25
+ """Manages data parallel vLLM server instances for external
26
+ load balancer testing."""
27
+
28
+ def __init__(self,
29
+ model_name: str,
30
+ dp_size: int,
31
+ api_server_count: int,
32
+ base_server_args: list,
33
+ tp_size: int = TP_SIZE):
34
+ self.model_name = model_name
35
+ self.dp_size = dp_size
36
+ self.tp_size = tp_size
37
+ self.api_server_count = api_server_count
38
+ self.base_server_args = base_server_args
39
+ self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
40
+ self.server_threads: list[threading.Thread] = []
41
+
42
+ def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
43
+ """Start all server instances for external LB mode."""
44
+ for rank in range(self.dp_size):
45
+ # Create server args for this specific rank
46
+ server_args = self.base_server_args.copy()
47
+
48
+ # Add external LB specific arguments
49
+ server_args.extend([
50
+ "--data-parallel-size",
51
+ str(self.dp_size),
52
+ "--data-parallel-rank",
53
+ str(rank),
54
+ "--data-parallel-size-local",
55
+ "1",
56
+ "--tensor-parallel-size",
57
+ str(self.tp_size),
58
+ "--port",
59
+ str(8000 + rank), # Different port for each rank
60
+ "--api-server-count",
61
+ str(self.api_server_count),
62
+ ])
63
+
64
+ # Use a thread to start each server to allow parallel initialization
65
+ def start_server(r: int, sargs: list[str]):
66
+ try:
67
+ # Start the server
68
+ server = RemoteOpenAIServer(
69
+ self.model_name,
70
+ sargs,
71
+ auto_port=False,
72
+ env_dict={
73
+ "CUDA_VISIBLE_DEVICES":
74
+ ",".join(
75
+ str(Platform.device_id_to_physical_device_id(
76
+ i))
77
+ for i in range(r * TP_SIZE, (r + 1) * TP_SIZE))
78
+ })
79
+ server.__enter__()
80
+ print(f"Server rank {r} started successfully with "
81
+ f"{self.api_server_count} API servers")
82
+ self.servers.append((server, sargs))
83
+ except Exception as e:
84
+ print(f"Failed to start server rank {r}: {e}")
85
+ raise
86
+
87
+ thread = threading.Thread(target=start_server,
88
+ args=(rank, server_args))
89
+ thread.start()
90
+
91
+ self.server_threads.append(thread)
92
+
93
+ # Wait for all servers to start
94
+ for thread in self.server_threads:
95
+ thread.join()
96
+
97
+ # Give servers additional time to fully initialize and coordinate
98
+ time.sleep(2)
99
+
100
+ if len(self.servers) != self.dp_size:
101
+ raise Exception("Servers failed to start")
102
+
103
+ return self.servers
104
+
105
+ def __exit__(self, exc_type, exc_val, exc_tb):
106
+ """Stop all server instances."""
107
+ while self.servers:
108
+ try:
109
+ self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
110
+ except Exception as e:
111
+ print(f"Error stopping server: {e}")
112
+
113
+
114
+@pytest.fixture(scope="module")
115
+def default_server_args():
116
+ return [
117
+ # use half precision for speed and memory savings in CI environment
118
+ "--dtype",
119
+ "bfloat16",
120
+ "--max-model-len",
121
+ "2048",
122
+ "--max-num-seqs",
123
+ "128",
124
+ "--enforce-eager",
125
+ ]
126
+
127
+
128
+@pytest.fixture(scope="module", params=[1, 4])
129
+def servers(request, default_server_args):
130
+ api_server_count = request.param
131
+ with ExternalLBServerManager(MODEL_NAME, DP_SIZE, api_server_count,
132
+ default_server_args) as server_list:
133
+ yield server_list
134
+
135
+
136
+@pytest_asyncio.fixture
137
+async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
138
+ # Create a client for each server
139
+ async with AsyncExitStack() as stack:
140
+ yield [
141
+ await stack.enter_async_context(server.get_async_client())
142
+ for server, _ in servers
143
+ ]
144
+
145
+
146
+@pytest.mark.asyncio
147
+@pytest.mark.parametrize(
148
+ "model_name",
149
+ [MODEL_NAME],
150
+)
151
+async def test_external_lb_single_completion(clients: list[
152
+ openai.AsyncOpenAI], servers: list[tuple[RemoteOpenAIServer, list[str]]],
153
+ model_name: str) -> None:
154
+
155
+ async def make_request(client: openai.AsyncOpenAI):
156
+ completion = await client.completions.create(
157
+ model=model_name,
158
+ prompt="Hello, my name is",
159
+ max_tokens=10,
160
+ temperature=1.0)
161
+
162
+ assert completion.id is not None
163
+ assert completion.choices is not None and len(completion.choices) == 1
164
+
165
+ choice = completion.choices[0]
166
+ # The exact number of tokens can vary slightly with temperature=1.0,
167
+ # so we check for a reasonable minimum length.
168
+ assert len(choice.text) >= 1
169
+ # Finish reason might not always be 'length' if the model finishes early
170
+ # or due to other reasons, especially with high temperature.
171
+ # So, we'll accept 'length' or 'stop'.
172
+ assert choice.finish_reason in ("length", "stop")
173
+
174
+ # Token counts can also vary, so we check they are positive.
175
+ assert completion.usage.completion_tokens > 0
176
+ assert completion.usage.prompt_tokens > 0
177
+ assert completion.usage.total_tokens > 0
178
+ return completion
179
+
180
+ # Test single request to each server
181
+ for i, client in enumerate(clients):
182
+ result = await make_request(client)
183
+ assert result is not None
184
+ print(f"Server {i} handled single completion request successfully")
185
+
186
+ await asyncio.sleep(0.5)
187
+
188
+ # Send requests to all servers in round-robin fashion
189
+ num_requests_per_server = 25 # Total 50 requests across 2 servers
190
+ all_tasks = []
191
+
192
+ for i, client in enumerate(clients):
193
+ tasks = [make_request(client) for _ in range(num_requests_per_server)]
194
+ all_tasks.extend(tasks)
195
+
196
+ results = await asyncio.gather(*all_tasks)
197
+ assert len(results) == num_requests_per_server * len(clients)
198
+ assert all(completion is not None for completion in results)
199
+
200
+ await asyncio.sleep(0.5)
201
+
202
+ # Second burst of requests
203
+ all_tasks = []
204
+ for i, client in enumerate(clients):
205
+ tasks = [make_request(client) for _ in range(num_requests_per_server)]
206
+ all_tasks.extend(tasks)
207
+
208
+ results = await asyncio.gather(*all_tasks)
209
+ assert len(results) == num_requests_per_server * len(clients)
210
+ assert all(completion is not None for completion in results)
211
+
212
+ _, server_args = servers[0]
213
+ api_server_count = (
214
+ server_args.count('--api-server-count')
215
+ and server_args[server_args.index('--api-server-count') + 1] or 1)
216
+ print(
217
+ f"Successfully completed external LB test with {len(clients)} servers "
218
+ f"(API server count: {api_server_count})")
219
+
220
+
221
+@pytest.mark.asyncio
222
+@pytest.mark.parametrize(
223
+ "model_name",
224
+ [MODEL_NAME],
225
+)
226
+async def test_external_lb_completion_streaming(clients: list[
227
+ openai.AsyncOpenAI], servers: list[tuple[RemoteOpenAIServer, list[str]]],
228
+ model_name: str) -> None:
229
+ prompt = "What is an LLM?"
230
+
231
+ async def make_streaming_request(client: openai.AsyncOpenAI):
232
+ # Perform a non-streaming request to get the expected full output
233
+ single_completion = await client.completions.create(
234
+ model=model_name,
235
+ prompt=prompt,
236
+ max_tokens=5,
237
+ temperature=0.0,
238
+ )
239
+ single_output = single_completion.choices[0].text
240
+
241
+ # Perform the streaming request
242
+ stream = await client.completions.create(model=model_name,
243
+ prompt=prompt,
244
+ max_tokens=5,
245
+ temperature=0.0,
246
+ stream=True)
247
+ chunks: list[str] = []
248
+ finish_reason_count = 0
249
+ last_chunk = None
250
+ async for chunk in stream:
251
+ chunks.append(chunk.choices[0].text)
252
+ if chunk.choices[0].finish_reason is not None:
253
+ finish_reason_count += 1
254
+ last_chunk = chunk # Keep track of the last chunk
255
+
256
+ # finish reason should only return in the last block for OpenAI API
257
+ assert finish_reason_count == 1, (
258
+ "Finish reason should appear exactly once.")
259
+ assert last_chunk is not None, (
260
+ "Stream should have yielded at least one chunk.")
261
+ assert last_chunk.choices[
262
+ 0].finish_reason == "length", "Finish reason should be 'length'."
263
+ # Check that the combined text matches the non-streamed version.
264
+ assert "".join(
265
+ chunks
266
+ ) == single_output, "Streamed output should match non-streamed output."
267
+ return True # Indicate success for this request
268
+
269
+ # Test single request to each server
270
+ for i, client in enumerate(clients):
271
+ result = await make_streaming_request(client)
272
+ assert result is not None
273
+ print(f"Server {i} handled single streaming request successfully")
274
+
275
+ await asyncio.sleep(0.5)
276
+
277
+ # Send streaming requests to all servers in round-robin fashion
278
+ num_requests_per_server = 25 # Total 50 requests across 2 servers
279
+ all_tasks = []
280
+
281
+ for i, client in enumerate(clients):
282
+ tasks = [
283
+ make_streaming_request(client)
284
+ for _ in range(num_requests_per_server)
285
+ ]
286
+ all_tasks.extend(tasks)
287
+
288
+ results = await asyncio.gather(*all_tasks)
289
+ assert len(results) == num_requests_per_server * len(clients)
290
+ assert all(results), "Not all streaming requests completed successfully."
291
+
292
+ await asyncio.sleep(0.5)
293
+
294
+ # Second burst of streaming requests
295
+ all_tasks = []
296
+ for i, client in enumerate(clients):
297
+ tasks = [
298
+ make_streaming_request(client)
299
+ for _ in range(num_requests_per_server)
300
+ ]
301
+ all_tasks.extend(tasks)
302
+
303
+ results = await asyncio.gather(*all_tasks)
304
+ assert len(results) == num_requests_per_server * len(clients)
305
+ assert all(results), "Not all streaming requests completed successfully."
306
+
307
+ _, server_args = servers[0]
308
+ api_server_count = (
309
+ server_args.count('--api-server-count')
310
+ and server_args[server_args.index('--api-server-count') + 1] or 1)
311
+ print(f"Successfully completed external LB streaming test with "
312
+ f"{len(clients)} servers (API server count: {api_server_count})")
vllm/config.py
+13
@@ -1784,6 +1784,10 @@ class ParallelConfig:
1784
"""Port of the data parallel master."""
1785
data_parallel_backend: str = "mp"
1786
"""Backend to use for data parallel, either "mp" or "ray"."""
1787
+ data_parallel_external_lb: bool = False
1788
+ """Whether to use "external" DP LB mode. Applies only to online serving
1789
+ and when data_parallel_size > 0. Set implicitly when
1790
+ data_parallel_rank is provided explicitly to vllm serve."""
1791
enable_expert_parallel: bool = False
1792
"""Use expert parallelism instead of tensor parallelism for MoE layers."""
1793
enable_eplb: bool = False
@@ -1953,6 +1957,11 @@ class ParallelConfig:
1957
if self.data_parallel_size > 1 or self.data_parallel_size_local == 0:
1958
# Data parallel was specified in the engine args.
1959
self.data_parallel_master_port = get_open_port()
1960
+
1961
+ if not (0 <= self.data_parallel_rank < self.data_parallel_size):
1962
+ raise ValueError(
1963
+ f"data_parallel_rank ({self.data_parallel_rank})"
1964
+ f" must be in the range [0, {self.data_parallel_size})")
1965
else:
1966
# Otherwise fall back to env vars (e.g. for offline SPMD case).
1967
self.data_parallel_size = envs.VLLM_DP_SIZE
@@ -1961,6 +1970,10 @@ class ParallelConfig:
1970
self.data_parallel_master_ip = envs.VLLM_DP_MASTER_IP
1971
self.data_parallel_master_port = envs.VLLM_DP_MASTER_PORT
1972
1973
+ if self.data_parallel_external_lb:
1974
+ raise ValueError("data_parallel_external_lb can only "
1975
+ "be set when data_parallel_size > 1")
1976
+
1977
if self.distributed_executor_backend == "external_launcher":
1978
import os
1979
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
vllm/engine/arg_utils.py
+21
-7
@@ -318,6 +318,7 @@ class EngineArgs:
318
pipeline_parallel_size: int = ParallelConfig.pipeline_parallel_size
319
tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
320
data_parallel_size: int = ParallelConfig.data_parallel_size
321
+ data_parallel_rank: Optional[int] = None
322
data_parallel_size_local: Optional[int] = None
323
data_parallel_address: Optional[str] = None
324
data_parallel_rpc_port: Optional[int] = None
@@ -655,6 +656,12 @@ class EngineArgs:
656
**parallel_kwargs["tensor_parallel_size"])
657
parallel_group.add_argument("--data-parallel-size", "-dp",
658
**parallel_kwargs["data_parallel_size"])
659
+ parallel_group.add_argument(
660
+ '--data-parallel-rank',
661
+ '-dpn',
662
+ type=int,
663
+ help='Data parallel rank of this instance. '
664
+ 'When set, enables external load balancer mode.')
665
parallel_group.add_argument('--data-parallel-size-local',
666
'-dpl',
667
type=int,
@@ -1126,10 +1133,17 @@ class EngineArgs:
1133
# but we should not do this here.
1134
placement_group = ray.util.get_current_placement_group()
1135
1129
- # Local DP size defaults to global DP size if not set.
1130
- data_parallel_size_local = self.data_parallel_size if (
1131
- self.data_parallel_size_local
1132
- is None) else self.data_parallel_size_local
1136
+ data_parallel_external_lb = self.data_parallel_rank is not None
1137
+ if data_parallel_external_lb:
1138
+ assert self.data_parallel_size_local in (1, None), (
1139
+ "data_parallel_size_local must be 1 when data_parallel_rank "
1140
+ "is set")
1141
+ data_parallel_size_local = 1
1142
+ elif self.data_parallel_size_local is not None:
1143
+ data_parallel_size_local = self.data_parallel_size_local
1144
+ else:
1145
+ # Local DP size defaults to global DP size if not set.
1146
+ data_parallel_size_local = self.data_parallel_size
1147
1148
# DP address, used in multi-node case for torch distributed group
1149
# and ZMQ sockets.
@@ -1154,16 +1168,16 @@ class EngineArgs:
1168
self.data_parallel_rpc_port
1169
is not None) else ParallelConfig.data_parallel_rpc_port
1170
1157
- data_parallel_backend = self.data_parallel_backend
1158
-
1171
parallel_config = ParallelConfig(
1172
pipeline_parallel_size=self.pipeline_parallel_size,
1173
tensor_parallel_size=self.tensor_parallel_size,
1174
data_parallel_size=self.data_parallel_size,
1175
+ data_parallel_rank=self.data_parallel_rank or 0,
1176
+ data_parallel_external_lb=data_parallel_external_lb,
1177
data_parallel_size_local=data_parallel_size_local,
1178
data_parallel_master_ip=data_parallel_address,
1179
data_parallel_rpc_port=data_parallel_rpc_port,
1166
- data_parallel_backend=data_parallel_backend,
1180
+ data_parallel_backend=self.data_parallel_backend,
1181
enable_expert_parallel=self.enable_expert_parallel,
1182
enable_eplb=self.enable_eplb,
1183
num_redundant_experts=self.num_redundant_experts,
vllm/entrypoints/cli/serve.py
+59
-118
@@ -5,9 +5,9 @@ import argparse
5
import os
6
import signal
7
import sys
8
+from typing import Optional
9
10
import uvloop
10
-import zmq
11
12
import vllm
13
import vllm.envs as envs
@@ -21,17 +21,13 @@ from vllm.entrypoints.utils import (VLLM_SUBCMD_PARSER_EPILOG,
21
from vllm.executor.multiproc_worker_utils import _add_prefix
22
from vllm.logger import init_logger
23
from vllm.usage.usage_lib import UsageContext
24
-from vllm.utils import FlexibleArgumentParser, get_tcp_uri, zmq_socket_ctx
25
-from vllm.v1.engine.coordinator import DPCoordinator
24
+from vllm.utils import FlexibleArgumentParser, get_tcp_uri
25
from vllm.v1.engine.core import EngineCoreProc
27
-from vllm.v1.engine.core_client import CoreEngineProcManager
26
+from vllm.v1.engine.utils import CoreEngineProcManager, launch_core_engines
27
from vllm.v1.executor.abstract import Executor
28
from vllm.v1.metrics.prometheus import setup_multiprocess_prometheus
30
-from vllm.v1.utils import (APIServerProcessManager, CoreEngine,
31
- CoreEngineActorManager, EngineZmqAddresses,
32
- get_engine_client_zmq_addr,
33
- wait_for_completion_or_failure,
34
- wait_for_engine_startup)
29
+from vllm.v1.utils import (APIServerProcessManager,
30
+ wait_for_completion_or_failure)
31
32
logger = init_logger(__name__)
33
@@ -48,11 +44,15 @@ class ServeSubcommand(CLISubcommand):
44
45
if args.headless or args.api_server_count < 1:
46
run_headless(args)
51
- elif args.api_server_count > 1:
52
- run_multi_api_server(args)
47
else:
54
- # Single API server (this process).
55
- uvloop.run(run_server(args))
48
+ if args.data_parallel_start_rank:
49
+ raise ValueError("data_parallel_start_rank is only "
50
+ "applicable in headless mode")
51
+ if args.api_server_count > 1:
52
+ run_multi_api_server(args)
53
+ else:
54
+ # Single API server (this process).
55
+ uvloop.run(run_server(args))
56
57
def validate(self, args: argparse.Namespace) -> None:
58
validate_parsed_serve_args(args)
@@ -121,14 +121,19 @@ def run_headless(args: argparse.Namespace):
121
122
parallel_config = vllm_config.parallel_config
123
local_engine_count = parallel_config.data_parallel_size_local
124
- host = parallel_config.data_parallel_master_ip
125
- port = engine_args.data_parallel_rpc_port # add to config too
126
- handshake_address = get_tcp_uri(host, port)
124
125
if local_engine_count <= 0:
126
raise ValueError("data_parallel_size_local must be > 0 in "
127
"headless mode")
128
129
+ if parallel_config.data_parallel_rank is not None:
130
+ raise ValueError("data_parallel_rank is not applicable in "
131
+ "headless mode")
132
+
133
+ host = parallel_config.data_parallel_master_ip
134
+ port = engine_args.data_parallel_rpc_port # add to config too
135
+ handshake_address = get_tcp_uri(host, port)
136
+
137
# Catch SIGTERM and SIGINT to allow graceful shutdown.
138
def signal_handler(signum, frame):
139
logger.debug("Received %d signal.", signum)
@@ -148,7 +153,7 @@ def run_headless(args: argparse.Namespace):
153
start_index=args.data_parallel_start_rank,
154
local_start_index=0,
155
vllm_config=vllm_config,
151
- on_head_node=False,
156
+ local_client=False,
157
handshake_address=handshake_address,
158
executor_class=Executor.get_class(vllm_config),
159
log_stats=not engine_args.disable_log_stats,
@@ -192,117 +197,53 @@ def run_multi_api_server(args: argparse.Namespace):
197
" api_server_count > 1")
198
model_config.disable_mm_preprocessor_cache = True
199
200
+ executor_class = Executor.get_class(vllm_config)
201
+ log_stats = not engine_args.disable_log_stats
202
+
203
parallel_config = vllm_config.parallel_config
204
+ dp_rank = parallel_config.data_parallel_rank
205
+ external_dp_lb = parallel_config.data_parallel_external_lb
206
+ assert external_dp_lb or dp_rank == 0
207
197
- assert parallel_config.data_parallel_rank == 0
208
+ api_server_manager: Optional[APIServerProcessManager] = None
209
199
- dp_size = parallel_config.data_parallel_size
200
- local_engine_count = parallel_config.data_parallel_size_local
201
- host = parallel_config.data_parallel_master_ip
202
- local_only = local_engine_count == dp_size
203
-
204
- # Set up input and output addresses.
205
- input_addresses = [
206
- get_engine_client_zmq_addr(local_only, host)
207
- for _ in range(num_api_servers)
208
- ]
209
- output_addresses = [
210
- get_engine_client_zmq_addr(local_only, host)
211
- for _ in range(num_api_servers)
212
- ]
213
-
214
- addresses = EngineZmqAddresses(
215
- inputs=input_addresses,
216
- outputs=output_addresses,
217
- )
210
+ with launch_core_engines(vllm_config, executor_class, log_stats,
211
+ num_api_servers) as (local_engine_manager,
212
+ coordinator, addresses):
213
219
- # Set up coordinator for dp > 1.
220
- coordinator = None
221
- stats_update_address = None
222
- if dp_size > 1:
223
- coordinator = DPCoordinator(parallel_config)
224
- addresses.coordinator_input, addresses.coordinator_output = (
225
- coordinator.get_engine_socket_addresses())
226
- stats_update_address = coordinator.get_stats_publish_address()
227
- logger.info("Started DP Coordinator process (PID: %d)",
228
- coordinator.proc.pid)
229
-
230
- if parallel_config.data_parallel_backend == "ray":
231
- logger.info("Starting ray-based data parallel backend")
232
-
233
- engine_actor_manager = CoreEngineActorManager(
234
- vllm_config=vllm_config,
235
- addresses=addresses,
236
- executor_class=Executor.get_class(vllm_config),
237
- log_stats=not engine_args.disable_log_stats,
238
- )
239
- # Start API servers using the manager
240
- api_server_manager = APIServerProcessManager(
214
+ # Construct common args for the APIServerProcessManager up-front.
215
+ api_server_manager_kwargs = dict(
216
target_server_fn=run_api_server_worker_proc,
217
listen_address=listen_address,
218
sock=sock,
219
args=args,
220
num_servers=num_api_servers,
246
- input_addresses=input_addresses,
247
- output_addresses=output_addresses,
248
- stats_update_address=stats_update_address)
249
-
250
- wait_for_completion_or_failure(api_server_manager=api_server_manager,
251
- engine_manager=engine_actor_manager,
252
- coordinator=coordinator)
253
- return
254
-
255
- handshake_address = get_engine_client_zmq_addr(
256
- local_only, host, parallel_config.data_parallel_rpc_port)
257
-
258
- with zmq_socket_ctx(handshake_address, zmq.ROUTER,
259
- bind=True) as handshake_socket:
260
-
261
- # Start local engines.
262
- if not local_engine_count:
263
- local_engine_manager = None
264
- else:
265
- local_engine_manager = CoreEngineProcManager(
266
- EngineCoreProc.run_engine_core,
267
- vllm_config=vllm_config,
268
- executor_class=Executor.get_class(vllm_config),
269
- log_stats=not engine_args.disable_log_stats,
270
- handshake_address=handshake_address,
271
- on_head_node=True,
272
- local_engine_count=local_engine_count,
273
- start_index=0,
274
- local_start_index=0)
275
-
276
- # Start API servers using the manager
221
+ input_addresses=addresses.inputs,
222
+ output_addresses=addresses.outputs,
223
+ stats_update_address=coordinator.get_stats_publish_address()
224
+ if coordinator else None)
225
+
226
+ # For dp ranks > 0 in external DP LB mode, we must delay the
227
+ # start of the API servers until the local engine is started
228
+ # (after the launcher context manager exits),
229
+ # since we get the front-end stats update address from the coordinator
230
+ # via the handshake with the local engine.
231
+ if dp_rank == 0 or not external_dp_lb:
232
+ # Start API servers using the manager.
233
+ api_server_manager = APIServerProcessManager(
234
+ **api_server_manager_kwargs)
235
+
236
+ # Start API servers now if they weren't already started.
237
+ if api_server_manager is None:
238
+ api_server_manager_kwargs["stats_update_address"] = (
239
+ addresses.frontend_stats_publish_address)
240
api_server_manager = APIServerProcessManager(
278
- target_server_fn=run_api_server_worker_proc,
279
- listen_address=listen_address,
280
- sock=sock,
281
- args=args,
282
- num_servers=num_api_servers,
283
- input_addresses=input_addresses,
284
- output_addresses=output_addresses,
285
- stats_update_address=stats_update_address)
286
-
287
- # Wait for engine handshakes to complete.
288
- core_engines = [
289
- CoreEngine(index=i, local=(i < local_engine_count))
290
- for i in range(dp_size)
291
- ]
292
- wait_for_engine_startup(
293
- handshake_socket,
294
- addresses,
295
- core_engines,
296
- parallel_config,
297
- vllm_config.cache_config,
298
- local_engine_manager,
299
- coordinator.proc if coordinator else None,
300
- )
301
-
302
- # Wait for API servers
303
- wait_for_completion_or_failure(api_server_manager=api_server_manager,
304
- engine_manager=local_engine_manager,
305
- coordinator=coordinator)
241
+ **api_server_manager_kwargs)
242
+
243
+ # Wait for API servers
244
+ wait_for_completion_or_failure(api_server_manager=api_server_manager,
245
+ engine_manager=local_engine_manager,
246
+ coordinator=coordinator)
247
248
249
def run_api_server_worker_proc(listen_address,
vllm/v1/engine/coordinator.py
+35
-13
@@ -10,7 +10,7 @@ import zmq
10
11
from vllm.config import ParallelConfig
12
from vllm.logger import init_logger
13
-from vllm.utils import get_mp_context, get_open_zmq_ipc_path, make_zmq_socket
13
+from vllm.utils import get_mp_context, make_zmq_socket
14
from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType
15
from vllm.v1.serial_utils import MsgpackDecoder
16
from vllm.v1.utils import get_engine_client_zmq_addr, shutdown
@@ -48,20 +48,33 @@ class DPCoordinator:
48
49
Engines will move into running state when receiving a new request or
50
START_DP_WAVE message.
51
+
52
+ Note that when deployed in External LB mode, no stats will be published by
53
+ the engines and thus updates will only be sent to front-ends when the
54
+ request wave / running state changes.
55
"""
56
57
def __init__(self, parallel_config: ParallelConfig):
58
55
- # Assume coordinator is colocated with front-end procs.
56
- front_publish_address = get_open_zmq_ipc_path()
57
-
59
dp_size = parallel_config.data_parallel_size
60
assert dp_size > 1, "Coordinator only used for data parallel"
61
61
- local_only = dp_size == parallel_config.data_parallel_size_local
62
host = parallel_config.data_parallel_master_ip
63
- back_publish_address = get_engine_client_zmq_addr(local_only, host)
64
- back_output_address = get_engine_client_zmq_addr(local_only, host)
63
+ external_lb = parallel_config.data_parallel_external_lb
64
+
65
+ # Assume coordinator is colocated with front-end procs when not in
66
+ # external DP LB mode.
67
+ front_publish_address = get_engine_client_zmq_addr(
68
+ local_only=not external_lb, host=host)
69
+
70
+ local_only_eng = dp_size == parallel_config.data_parallel_size_local
71
+ back_publish_address = get_engine_client_zmq_addr(local_only_eng, host)
72
+ back_output_address = get_engine_client_zmq_addr(local_only_eng, host)
73
+
74
+ # When in external LB mode, load stats aren't published, only changes
75
+ # to request wave / running state, so we don't need to rate-limit the
76
+ # updates to the front-end proc(s).
77
+ min_stats_update_interval_ms = 0 if external_lb else 100
78
79
context = get_mp_context()
80
self.proc: multiprocessing.Process = context.Process(
@@ -72,6 +85,7 @@ class DPCoordinator:
85
"front_publish_address": front_publish_address,
86
"back_output_address": back_output_address,
87
"back_publish_address": back_publish_address,
88
+ "min_stats_update_interval_ms": min_stats_update_interval_ms,
89
},
90
daemon=True)
91
self.proc.start()
@@ -100,12 +114,16 @@ class EngineState:
114
115
class CoordinatorProc:
116
103
- def __init__(self, engine_count: int):
117
+ def __init__(self,
118
+ engine_count: int,
119
+ min_stats_update_interval_ms: int = 100):
120
121
self.ctx = zmq.Context()
122
123
self.engines = [EngineState() for _ in range(engine_count)]
124
125
+ self.stats_update_interval_ms = min_stats_update_interval_ms
126
+
127
self.current_wave = 0
128
self.engines_running = False
129
self.stats_changed = False
@@ -116,8 +134,11 @@ class CoordinatorProc:
134
front_publish_address: str,
135
back_output_address: str,
136
back_publish_address: str,
137
+ min_stats_update_interval_ms: int = 100,
138
):
120
- coordinator = CoordinatorProc(engine_count=engine_count)
139
+ coordinator = CoordinatorProc(
140
+ engine_count=engine_count,
141
+ min_stats_update_interval_ms=min_stats_update_interval_ms)
142
try:
143
coordinator.process_input_socket(
144
front_publish_address,
@@ -156,9 +177,10 @@ class CoordinatorProc:
177
last_publish_time = 0
178
while True:
179
elapsed = int(time.time() * 1000) - last_publish_time
159
- # Send at 100 ms interval if the stats have changed,
160
- # or otherwise every 3 seconds.
161
- wait_for = 100 if self.stats_changed else 3000
180
+ # Send at stats_update_interval_ms interval if the stats have
181
+ # changed, or otherwise every 4 seconds.
182
+ wait_for = (self.stats_update_interval_ms
183
+ if self.stats_changed else 4000)
184
events = poller.poll(timeout=max(0, wait_for - elapsed))
185
if not events:
186
# Poller timeout - publish current stats to front-ends.
@@ -174,7 +196,7 @@ class CoordinatorProc:
196
197
if publish_front in events:
198
buffer = publish_front.recv()
177
- if buffer == b'\x01':
199
+ if buffer in (b'\x01', b'\x00'):
200
# Ignore subscription messages.
201
continue
202
vllm/v1/engine/core.py
+113
-36
@@ -34,6 +34,7 @@ from vllm.v1.core.sched.scheduler import Scheduler as V1Scheduler
34
from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
35
EngineCoreRequestType, UtilityOutput)
36
from vllm.v1.engine.mm_input_cache import MirroredProcessingCache
37
+from vllm.v1.engine.utils import EngineHandshakeMetadata, EngineZmqAddresses
38
from vllm.v1.executor.abstract import Executor
39
from vllm.v1.kv_cache_interface import KVCacheConfig
40
from vllm.v1.metrics.stats import SchedulerStats
@@ -41,7 +42,6 @@ from vllm.v1.outputs import ModelRunnerOutput
42
from vllm.v1.request import Request, RequestStatus
43
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
44
from vllm.v1.structured_output import StructuredOutputManager
44
-from vllm.v1.utils import EngineHandshakeMetadata, EngineZmqAddresses
45
from vllm.version import __version__ as VLLM_VERSION
46
47
logger = init_logger(__name__)
@@ -367,10 +367,11 @@ class EngineCoreProc(EngineCore):
367
def __init__(
368
self,
369
vllm_config: VllmConfig,
370
- on_head_node: bool,
370
+ local_client: bool,
371
handshake_address: str,
372
executor_class: type[Executor],
373
log_stats: bool,
374
+ client_handshake_address: Optional[str] = None,
375
engine_index: int = 0,
376
):
377
self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]()
@@ -383,12 +384,21 @@ class EngineCoreProc(EngineCore):
384
identity = self.engine_index.to_bytes(length=2, byteorder="little")
385
self.engines_running = False
386
386
- with self._perform_handshake(handshake_address, identity, on_head_node,
387
- vllm_config) as addresses:
387
+ with self._perform_handshakes(handshake_address, identity,
388
+ local_client, vllm_config,
389
+ client_handshake_address) as addresses:
390
self.client_count = len(addresses.outputs)
391
392
# Set up data parallel environment.
393
self.has_coordinator = addresses.coordinator_output is not None
394
+ self.frontend_stats_publish_address = (
395
+ addresses.frontend_stats_publish_address)
396
+ # Only publish request queue stats to coordinator for "internal"
397
+ # LB mode.
398
+ self.publish_dp_lb_stats = (
399
+ self.has_coordinator
400
+ and not vllm_config.parallel_config.data_parallel_external_lb)
401
+
402
self._init_data_parallel(vllm_config)
403
404
super().__init__(vllm_config, executor_class, log_stats,
@@ -414,45 +424,102 @@ class EngineCoreProc(EngineCore):
424
self.output_thread.start()
425
426
@contextmanager
417
- def _perform_handshake(
418
- self, handshake_address: str, identity: bytes, on_head_node: bool,
419
- vllm_config: VllmConfig
427
+ def _perform_handshakes(
428
+ self,
429
+ handshake_address: str,
430
+ identity: bytes,
431
+ local_client: bool,
432
+ vllm_config: VllmConfig,
433
+ client_handshake_address: Optional[str],
434
) -> Generator[EngineZmqAddresses, None, None]:
435
+ """
436
+ Perform startup handshakes.
437
+
438
+ For DP=1 or offline mode, this is with the colocated front-end process.
439
+
440
+ For DP>1 with internal loadbalancing this is with the shared front-end
441
+ process which may reside on a different node.
442
+
443
+ For DP>1 with external loadbalancing, two handshakes are performed:
444
+ - With the rank 0 front-end process which retrieves the
445
+ DP Coordinator ZMQ addresses and DP process group address.
446
+ - With the colocated front-end process which retrieves the
447
+ client input/output socket addresses.
448
+ with the exception of the rank 0 engine itself which doesn't require
449
+ the second handshake.
450
+
451
+ Here, "front-end" process can mean the process containing the engine
452
+ core client (which is the API server process in the case the API
453
+ server is not scaled out), OR the launcher process running the
454
+ run_multi_api_server() function in serve.py.
455
+ """
456
input_ctx = zmq.Context()
422
- with make_zmq_socket(input_ctx,
457
+ is_local = local_client and client_handshake_address is None
458
+ handshake = self._perform_handshake(input_ctx, handshake_address,
459
+ identity, is_local, vllm_config,
460
+ vllm_config.parallel_config)
461
+ if client_handshake_address is None:
462
+ with handshake as addresses:
463
+ yield addresses
464
+ else:
465
+ local_handshake = self._perform_handshake(
466
+ input_ctx, client_handshake_address, identity, local_client,
467
+ vllm_config)
468
+ with handshake as addresses, local_handshake as client_addresses:
469
+ addresses.inputs = client_addresses.inputs
470
+ addresses.outputs = client_addresses.outputs
471
+ yield addresses
472
+
473
+ # Update config which may have changed from the handshake
474
+ vllm_config.__post_init__()
475
+
476
+ @contextmanager
477
+ def _perform_handshake(
478
+ self,
479
+ ctx: zmq.Context,
480
+ handshake_address: str,
481
+ identity: bytes,
482
+ local_client: bool,
483
+ vllm_config: VllmConfig,
484
+ parallel_config_to_update: Optional[ParallelConfig] = None,
485
+ ) -> Generator[EngineZmqAddresses, None, None]:
486
+ with make_zmq_socket(ctx,
487
handshake_address,
488
zmq.DEALER,
489
identity=identity,
490
linger=5000,
491
bind=False) as handshake_socket:
492
# Register engine with front-end.
429
- addresses = self.startup_handshake(handshake_socket, on_head_node,
430
- vllm_config.parallel_config)
431
-
432
- # Update config which may have changed from the handshake
433
- vllm_config.__post_init__()
434
-
493
+ addresses = self.startup_handshake(handshake_socket, local_client,
494
+ parallel_config_to_update)
495
yield addresses
496
497
# Send ready message.
498
num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
499
+ # We pass back the coordinator stats update address here for the
500
+ # external LB case for our colocated front-end to use (coordinator
501
+ # only runs with rank 0).
502
+ dp_stats_address = self.frontend_stats_publish_address
503
handshake_socket.send(
504
msgspec.msgpack.encode({
505
"status": "READY",
442
- "local": on_head_node,
506
+ "local": local_client,
507
"num_gpu_blocks": num_gpu_blocks,
508
+ "dp_stats_address": dp_stats_address,
509
}))
510
511
@staticmethod
512
def startup_handshake(
448
- handshake_socket: zmq.Socket, on_head_node: bool,
449
- parallel_config: ParallelConfig) -> EngineZmqAddresses:
513
+ handshake_socket: zmq.Socket,
514
+ local_client: bool,
515
+ parallel_config: Optional[ParallelConfig] = None,
516
+ ) -> EngineZmqAddresses:
517
518
# Send registration message.
519
handshake_socket.send(
520
msgspec.msgpack.encode({
521
"status": "HELLO",
455
- "local": on_head_node,
522
+ "local": local_client,
523
}))
524
525
# Receive initialization message.
@@ -466,9 +533,9 @@ class EngineCoreProc(EngineCore):
533
init_bytes, type=EngineHandshakeMetadata)
534
logger.debug("Received init message: %s", init_message)
535
469
- received_parallel_config = init_message.parallel_config
470
- for key, value in received_parallel_config.items():
471
- setattr(parallel_config, key, value)
536
+ if parallel_config is not None:
537
+ for key, value in init_message.parallel_config.items():
538
+ setattr(parallel_config, key, value)
539
540
return init_message.addresses
541
@@ -749,12 +816,12 @@ class DPEngineCoreProc(EngineCoreProc):
816
def __init__(
817
self,
818
vllm_config: VllmConfig,
752
- on_head_node: bool,
819
+ local_client: bool,
820
handshake_address: str,
821
executor_class: type[Executor],
822
log_stats: bool,
823
+ client_handshake_address: Optional[str] = None,
824
):
757
-
825
self._decorate_logs()
826
827
# Counts forward-passes of the model so that we can synchronize
@@ -765,8 +832,9 @@ class DPEngineCoreProc(EngineCoreProc):
832
833
# Initialize the engine.
834
dp_rank = vllm_config.parallel_config.data_parallel_rank
768
- super().__init__(vllm_config, on_head_node, handshake_address,
769
- executor_class, log_stats, dp_rank)
835
+ super().__init__(vllm_config, local_client, handshake_address,
836
+ executor_class, log_stats, client_handshake_address,
837
+ dp_rank)
838
839
def _decorate_logs(self):
840
# Add process-specific prefix to stdout and stderr before
@@ -799,10 +867,18 @@ class DPEngineCoreProc(EngineCoreProc):
867
from vllm.platforms import current_platform
868
device_control_env_var = current_platform.device_control_env_var
869
world_size = vllm_config.parallel_config.world_size
802
- os.environ[device_control_env_var] = ",".join(
803
- str(current_platform.device_id_to_physical_device_id(i))
804
- for i in range(local_dp_rank * world_size, (local_dp_rank + 1) *
805
- world_size))
870
+ # Set CUDA_VISIBLE_DEVICES or equivalent.
871
+ try:
872
+ os.environ[device_control_env_var] = ",".join(
873
+ str(current_platform.device_id_to_physical_device_id(i))
874
+ for i in range(local_dp_rank *
875
+ world_size, (local_dp_rank + 1) * world_size))
876
+ except IndexError as e:
877
+ raise Exception(
878
+ f"Error setting {device_control_env_var}: "
879
+ f"local range: [{local_dp_rank * world_size}, "
880
+ f"{(local_dp_rank + 1) * world_size}) "
881
+ f"base value: \"{os.getenv(device_control_env_var)}\"") from e
882
883
self.dp_rank = dp_rank
884
self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()
@@ -839,7 +915,7 @@ class DPEngineCoreProc(EngineCoreProc):
915
super()._handle_client_request(request_type, request)
916
917
def _maybe_publish_request_counts(self):
842
- if not self.has_coordinator:
918
+ if not self.publish_dp_lb_stats:
919
return
920
921
# Publish our request counts (if they've changed).
@@ -892,9 +968,9 @@ class DPEngineCoreProc(EngineCoreProc):
968
969
def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
970
895
- # Optimization - only perform finish-sync all-reduce every 24 steps.
971
+ # Optimization - only perform finish-sync all-reduce every 32 steps.
972
self.counter += 1
897
- if self.counter != 24:
973
+ if self.counter != 32:
974
return True
975
self.counter = 0
976
@@ -910,7 +986,7 @@ class DPEngineCoreActor(DPEngineCoreProc):
986
def __init__(
987
self,
988
vllm_config: VllmConfig,
913
- on_head_node: bool,
989
+ local_client: bool,
990
addresses: EngineZmqAddresses,
991
executor_class: type[Executor],
992
log_stats: bool,
@@ -927,15 +1003,16 @@ class DPEngineCoreActor(DPEngineCoreProc):
1003
# data parallel groups.
1004
del os.environ['CUDA_VISIBLE_DEVICES']
1005
930
- super().__init__(vllm_config, on_head_node, "", executor_class,
1006
+ super().__init__(vllm_config, local_client, "", executor_class,
1007
log_stats)
1008
1009
def _decorate_logs(self):
1010
pass
1011
1012
@contextmanager
937
- def _perform_handshake(self, handshake_address: str, identity: bytes,
938
- on_head_node: bool, vllm_config: VllmConfig):
1013
+ def _perform_handshakes(self, handshake_address: str, identity: bytes,
1014
+ local_client: bool, vllm_config: VllmConfig,
1015
+ client_handshake_address: Optional[str]):
1016
"""
1017
For Ray, we don't need to actually perform handshake.
1018
All addresses information is known before the actor creation.
vllm/v1/engine/core_client.py
+126
-222
@@ -7,7 +7,7 @@ import sys
7
import uuid
8
import weakref
9
from abc import ABC, abstractmethod
10
-from collections import deque
10
+from collections import defaultdict, deque
11
from collections.abc import Awaitable, Sequence
12
from concurrent.futures import Future
13
from dataclasses import dataclass
@@ -21,18 +21,16 @@ import zmq.asyncio
21
from vllm.config import VllmConfig
22
from vllm.logger import init_logger
23
from vllm.lora.request import LoRARequest
24
-from vllm.utils import (get_open_zmq_inproc_path, make_zmq_socket,
25
- zmq_socket_ctx)
24
+from vllm.utils import get_open_zmq_inproc_path, make_zmq_socket
25
from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
26
EngineCoreRequestType, UtilityOutput)
27
from vllm.v1.engine.coordinator import DPCoordinator
28
from vllm.v1.engine.core import EngineCore, EngineCoreProc
29
from vllm.v1.engine.exceptions import EngineDeadError
30
+from vllm.v1.engine.utils import (CoreEngineActorManager,
31
+ CoreEngineProcManager, launch_core_engines)
32
from vllm.v1.executor.abstract import Executor
33
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr
33
-from vllm.v1.utils import (CoreEngine, CoreEngineActorManager,
34
- CoreEngineProcManager, EngineZmqAddresses,
35
- get_engine_client_zmq_addr, wait_for_engine_startup)
34
35
logger = init_logger(__name__)
36
@@ -40,6 +38,8 @@ AnyFuture = Union[asyncio.Future[Any], Future[Any]]
38
39
_R = TypeVar('_R') # Return type for collective_rpc
40
41
+EngineIdentity = bytes
42
+
43
44
class EngineCoreClient(ABC):
45
"""
@@ -84,14 +84,16 @@ class EngineCoreClient(ABC):
84
client_addresses: Optional[dict[str, str]] = None,
85
client_index: int = 0,
86
) -> "MPClient":
87
- if vllm_config.parallel_config.data_parallel_size > 1:
88
- if vllm_config.parallel_config.data_parallel_backend == "ray":
89
- return RayDPClient(vllm_config, executor_class, log_stats,
90
- client_addresses, client_index)
91
- return DPAsyncMPClient(vllm_config, executor_class, log_stats,
92
- client_addresses, client_index)
93
- return AsyncMPClient(vllm_config, executor_class, log_stats,
94
- client_addresses, client_index)
87
+ parallel_config = vllm_config.parallel_config
88
+ client_args = (vllm_config, executor_class, log_stats,
89
+ client_addresses, client_index)
90
+ if parallel_config.data_parallel_size > 1:
91
+ if parallel_config.data_parallel_external_lb:
92
+ # External load balancer - client per DP rank.
93
+ return DPAsyncMPClient(*client_args)
94
+ # Internal load balancer - client balances to all DP ranks.
95
+ return DPLBAsyncMPClient(*client_args)
96
+ return AsyncMPClient(*client_args)
97
98
@abstractmethod
99
def shutdown(self):
@@ -386,42 +388,32 @@ class MPClient(EngineCoreClient):
388
self._finalizer = weakref.finalize(self, self.resources)
389
success = False
390
try:
389
- parallel_config = vllm_config.parallel_config
390
- local_engine_count = parallel_config.data_parallel_size_local
391
- local_start_index = parallel_config.data_parallel_rank_local
392
- dp_size = parallel_config.data_parallel_size
393
- dp_rank = parallel_config.data_parallel_rank
394
-
391
# State used for data parallel.
392
self.engines_running = False
393
398
- # SPMD mode is where there is an LLM instance per DP rank and
399
- # one core engine per LLM, see
400
- # examples/offline_inference/data_parallel.py.
401
- spmd_mode = local_start_index is not None
402
- if spmd_mode:
403
- assert local_engine_count == 1
404
- self.core_engines = [CoreEngine(index=dp_rank, local=True)]
405
- else:
406
- assert dp_rank == 0
407
- local_start_index = 0
408
- self.core_engines = [
409
- CoreEngine(index=i, local=(i < local_engine_count))
410
- for i in range(dp_size)
411
- ]
412
-
413
- local_only = spmd_mode or local_engine_count == dp_size
414
-
394
self.stats_update_address: Optional[str] = None
395
if client_addresses is not None:
396
+ # Engines are managed externally to this client.
397
input_address = client_addresses["input_address"]
398
output_address = client_addresses["output_address"]
399
self.stats_update_address = client_addresses.get(
400
"stats_update_address")
401
else:
422
- host = parallel_config.data_parallel_master_ip
423
- input_address = get_engine_client_zmq_addr(local_only, host)
424
- output_address = get_engine_client_zmq_addr(local_only, host)
402
+ # Engines are managed by this client.
403
+ with launch_core_engines(vllm_config, executor_class,
404
+ log_stats) as (engine_manager,
405
+ coordinator,
406
+ addresses):
407
+ self.resources.coordinator = coordinator
408
+ self.resources.engine_manager = engine_manager
409
+
410
+ (input_address, ) = addresses.inputs
411
+ (output_address, ) = addresses.outputs
412
+ self.stats_update_address = (
413
+ addresses.frontend_stats_publish_address)
414
+ if coordinator is not None:
415
+ assert self.stats_update_address == (
416
+ coordinator.get_stats_publish_address())
417
418
# Create input and output sockets.
419
self.input_socket = self.resources.input_socket = make_zmq_socket(
@@ -429,18 +421,24 @@ class MPClient(EngineCoreClient):
421
self.resources.output_socket = make_zmq_socket(
422
self.ctx, output_address, zmq.PULL)
423
432
- if client_addresses is None:
433
- self._init_engines_direct(vllm_config, local_only,
434
- local_start_index, input_address,
435
- output_address, executor_class,
436
- log_stats)
437
- coordinator = self.resources.coordinator
438
- if coordinator:
439
- self.stats_update_address = (
440
- coordinator.get_stats_publish_address())
424
+ parallel_config = vllm_config.parallel_config
425
+ dp_size = parallel_config.data_parallel_size
426
+ dp_rank = parallel_config.data_parallel_rank
427
+ external_dp_lb = parallel_config.data_parallel_external_lb
428
+
429
+ offline_mode = parallel_config.data_parallel_rank_local is not None
430
+ engine_ranks = [dp_rank] if (offline_mode
431
+ or external_dp_lb) else range(dp_size)
432
+ assert parallel_config.data_parallel_size_local <= len(
433
+ engine_ranks)
434
+
435
+ # ZMQ identity of each engine that this client will talk to.
436
+ self.core_engines: list[EngineIdentity] = [
437
+ index.to_bytes(2, "little") for index in engine_ranks
438
+ ]
439
440
# Wait for ready messages from each engine on the input socket.
443
- identities = set(e.identity for e in self.core_engines)
441
+ identities = set(self.core_engines)
442
sync_input_socket = zmq.Socket.shadow(self.input_socket)
443
while identities:
444
if not sync_input_socket.poll(timeout=600_000):
@@ -449,7 +447,7 @@ class MPClient(EngineCoreClient):
447
identity, _ = sync_input_socket.recv_multipart()
448
identities.remove(identity)
449
452
- self.core_engine = self.core_engines[0]
450
+ self.core_engine: EngineIdentity = self.core_engines[0]
451
self.utility_results: dict[int, AnyFuture] = {}
452
453
# Request objects which may contain pytorch-allocated tensors
@@ -462,73 +460,6 @@ class MPClient(EngineCoreClient):
460
if not success:
461
self._finalizer()
462
465
- def _init_engines_direct(self, vllm_config: VllmConfig, local_only: bool,
466
- local_start_index: int, input_address: str,
467
- output_address: str,
468
- executor_class: type[Executor], log_stats: bool):
469
- """Self-contained client mode, launch engine and coordinator process
470
- as needed."""
471
-
472
- parallel_config = vllm_config.parallel_config
473
- local_engine_count = parallel_config.data_parallel_size_local
474
- start_index = parallel_config.data_parallel_rank
475
- host = parallel_config.data_parallel_master_ip
476
-
477
- if len(self.core_engines) > 1:
478
- self.resources.coordinator = DPCoordinator(parallel_config)
479
-
480
- handshake_address = get_engine_client_zmq_addr(
481
- local_only, host, parallel_config.data_parallel_rpc_port)
482
-
483
- with zmq_socket_ctx(handshake_address, zmq.ROUTER,
484
- bind=True) as handshake_socket:
485
-
486
- # Start local engines.
487
- if local_engine_count:
488
- # In server mode, start_index and local_start_index will
489
- # both be 0.
490
- self.resources.engine_manager = CoreEngineProcManager(
491
- EngineCoreProc.run_engine_core,
492
- vllm_config=vllm_config,
493
- executor_class=executor_class,
494
- log_stats=log_stats,
495
- handshake_address=handshake_address,
496
- on_head_node=True,
497
- local_engine_count=local_engine_count,
498
- start_index=start_index,
499
- local_start_index=local_start_index)
500
-
501
- # Wait for engine core process(es) to start.
502
- self._wait_for_engine_startup(handshake_socket, input_address,
503
- output_address)
504
-
505
- def _wait_for_engine_startup(self, handshake_socket: zmq.Socket,
506
- input_address: str, output_address: str):
507
- addresses = EngineZmqAddresses(
508
- inputs=[input_address],
509
- outputs=[output_address],
510
- )
511
-
512
- coordinator = self.resources.coordinator
513
- if coordinator is not None:
514
- addresses.coordinator_input, addresses.coordinator_output = (
515
- coordinator.get_engine_socket_addresses())
516
-
517
- proc_manager = self.resources.engine_manager
518
- assert isinstance(proc_manager, (type(None), CoreEngineProcManager)), (
519
- "_wait_for_engine_startup should only be "
520
- "called with CoreEngineProcManager")
521
-
522
- wait_for_engine_startup(
523
- handshake_socket,
524
- addresses,
525
- self.core_engines,
526
- self.vllm_config.parallel_config,
527
- self.vllm_config.cache_config,
528
- proc_manager,
529
- coordinator.proc if coordinator else None,
530
- )
531
-
463
def shutdown(self):
464
# Terminate background resources.
465
self._finalizer()
@@ -583,7 +514,6 @@ class SyncMPClient(MPClient):
514
# a ref to the client which prevents gc.
515
ctx = self.ctx
516
out_socket = self.resources.output_socket
586
- assert out_socket is not None
517
decoder = self.decoder
518
utility_results = self.utility_results
519
outputs_queue = self.outputs_queue
@@ -593,6 +523,7 @@ class SyncMPClient(MPClient):
523
resources.shutdown_path = shutdown_path
524
525
def process_outputs_socket():
526
+ assert isinstance(out_socket, zmq.Socket)
527
shutdown_socket = ctx.socket(zmq.PAIR)
528
try:
529
shutdown_socket.bind(shutdown_path)
@@ -609,7 +540,7 @@ class SyncMPClient(MPClient):
540
541
frames = out_socket.recv_multipart(copy=False)
542
resources.validate_alive(frames)
612
- outputs = decoder.decode(frames)
543
+ outputs: EngineCoreOutputs = decoder.decode(frames)
544
if outputs.utility_output:
545
_process_utility_output(outputs.utility_output,
546
utility_results)
@@ -646,7 +577,7 @@ class SyncMPClient(MPClient):
577
self.ensure_alive()
578
self.free_pending_messages()
579
# (Identity, RequestType, SerializedRequest)
649
- msg = (self.core_engine.identity, request_type.value,
580
+ msg = (self.core_engine, request_type.value,
581
*self.encoder.encode(request))
582
583
if len(msg) <= 3:
@@ -812,7 +743,7 @@ class AsyncMPClient(MPClient):
743
def _send_input(self,
744
request_type: EngineCoreRequestType,
745
request: Any,
815
- engine: Optional[CoreEngine] = None) -> Awaitable[Any]:
746
+ engine: Optional[EngineIdentity] = None) -> Awaitable[Any]:
747
if engine is None:
748
engine = self.core_engine
749
@@ -820,7 +751,7 @@ class AsyncMPClient(MPClient):
751
return self._send_input_message(message, engine, request)
752
753
def _send_input_message(self, message: tuple[bytestr,
823
- ...], engine: CoreEngine,
754
+ ...], engine: EngineIdentity,
755
objects: Any) -> Awaitable[Any]:
756
"""
757
objects is a reference to retain until zmq is finished with the
@@ -829,7 +760,7 @@ class AsyncMPClient(MPClient):
760
self.ensure_alive()
761
self.free_pending_messages()
762
832
- msg = (engine.identity, ) + message
763
+ msg = (engine, ) + message
764
if not objects or len(msg) <= 3:
765
# No auxiliary buffers => no tensor backing buffers in request.
766
return self.input_socket.send_multipart(msg, copy=False)
@@ -850,7 +781,7 @@ class AsyncMPClient(MPClient):
781
engine=self.core_engine)
782
783
async def _call_utility_async(self, method: str, *args,
853
- engine: CoreEngine) -> Any:
784
+ engine: EngineIdentity) -> Any:
785
call_id = uuid.uuid1().int >> 64
786
future = asyncio.get_running_loop().create_future()
787
self.utility_results[call_id] = future
@@ -921,7 +852,7 @@ class AsyncMPClient(MPClient):
852
853
class DPAsyncMPClient(AsyncMPClient):
854
"""Asyncio-compatible client for multi-proc, multi-engine (data parallel)
924
- EngineCore."""
855
+ EngineCore. Assumes external load-balancing by default."""
856
857
def __init__(self,
858
vllm_config: VllmConfig,
@@ -930,15 +861,12 @@ class DPAsyncMPClient(AsyncMPClient):
861
client_addresses: Optional[dict[str, str]] = None,
862
client_index: int = 0):
863
self.current_wave = 0
933
- # To route aborts to the correct engine.
934
- self.reqs_in_flight: dict[str, CoreEngine] = {}
864
865
super().__init__(vllm_config, executor_class, log_stats,
866
client_addresses, client_index)
867
939
- assert len(self.core_engines) > 1
940
-
868
# List of [waiting, running] pair per engine.
869
+ # Used only by DPLBAsyncMPClient subclass.
870
self.lb_engines: list[list[int]] = []
871
872
self.first_req_sock_addr = get_open_zmq_inproc_path()
@@ -969,6 +897,8 @@ class DPAsyncMPClient(AsyncMPClient):
897
self.first_req_sock_addr,
898
zmq.PAIR,
899
bind=False) as first_req_rcv_socket:
900
+ assert isinstance(socket, zmq.asyncio.Socket)
901
+ assert isinstance(first_req_rcv_socket, zmq.asyncio.Socket)
902
# Send subscription message.
903
await socket.send(b'\x01')
904
@@ -1012,64 +942,85 @@ class DPAsyncMPClient(AsyncMPClient):
942
resources.stats_update_task = asyncio.create_task(
943
run_engine_stats_update_task())
944
1015
- def get_core_engine_for_request(self,
1016
- dp_rank: Optional[int] = None
1017
- ) -> CoreEngine:
1018
- if dp_rank is not None:
1019
- # engines are already in rank order
1020
- return self.core_engines[dp_rank]
1021
-
1022
- if not self.lb_engines:
1023
- return self.core_engines[0]
1024
- # TODO use P2C alg for larger DP sizes
1025
- num_engines = len(self.lb_engines)
1026
- min_counts = [sys.maxsize, sys.maxsize]
1027
- eng_index = 0
1028
- for i in range(num_engines):
1029
- # Start from client_index to help with balancing when engines
1030
- # are empty.
1031
- idx = (self.client_index + i) % num_engines
1032
- counts = self.lb_engines[idx]
1033
- if counts < min_counts:
1034
- min_counts = counts
1035
- eng_index = idx
1036
- # Adjust local counts for better balancing between stats updates
1037
- # from the coordinator (which happen every 100ms).
1038
- if min_counts[0]:
1039
- min_counts[0] += 1
1040
- else:
1041
- min_counts[1] += 1
1042
- return self.core_engines[eng_index]
1043
-
1044
- async def call_utility_async(self, method: str, *args) -> Any:
1045
- # Only the result from the first engine is returned.
1046
- return (await asyncio.gather(*[
1047
- self._call_utility_async(method, *args, engine=engine)
1048
- for engine in self.core_engines
1049
- ]))[0]
1050
-
945
async def add_request_async(self, request: EngineCoreRequest) -> None:
946
self._ensure_stats_update_task()
947
948
request.current_wave = self.current_wave
949
request.client_index = self.client_index
950
1057
- chosen_engine = self.get_core_engine_for_request(
1058
- request.data_parallel_rank)
1059
- self.reqs_in_flight[request.request_id] = chosen_engine
1060
-
951
+ chosen_engine = self.get_core_engine_for_request(request)
952
to_await = self._send_input(EngineCoreRequestType.ADD, request,
953
chosen_engine)
954
if not self.engines_running:
955
# Notify coordinator that we're sending a request
1065
- await self.first_req_send_socket.send(chosen_engine.identity)
956
+ await self.first_req_send_socket.send(chosen_engine)
957
958
await to_await
959
960
self._ensure_output_queue_task()
961
962
+ def get_core_engine_for_request(self, request: EngineCoreRequest):
963
+ return self.core_engine
964
+
965
+
966
+class DPLBAsyncMPClient(DPAsyncMPClient):
967
+ """Asyncio-compatible client for multi-proc, multi-engine (data parallel)
968
+ EngineCore. Load-balances between multiple engine processes."""
969
+
970
+ def __init__(self,
971
+ vllm_config: VllmConfig,
972
+ executor_class: type[Executor],
973
+ log_stats: bool,
974
+ client_addresses: Optional[dict[str, str]] = None,
975
+ client_index: int = 0):
976
+
977
+ # To route aborts to the correct engine.
978
+ self.reqs_in_flight: dict[str, EngineIdentity] = {}
979
+
980
+ super().__init__(vllm_config, executor_class, log_stats,
981
+ client_addresses, client_index)
982
+
983
+ assert len(self.core_engines) > 1
984
+
985
+ def get_core_engine_for_request(
986
+ self, request: EngineCoreRequest) -> EngineIdentity:
987
+ # Engines are in rank order.
988
+ if (eng_index := request.data_parallel_rank) is None:
989
+ if not self.lb_engines:
990
+ return self.core_engine
991
+ # TODO use P2C alg for larger DP sizes
992
+ num_engines = len(self.lb_engines)
993
+ min_counts = [sys.maxsize, sys.maxsize]
994
+ eng_index = 0
995
+ for i in range(num_engines):
996
+ # Start from client_index to help with balancing when engines
997
+ # are empty.
998
+ idx = (self.client_index + i) % num_engines
999
+ counts = self.lb_engines[idx]
1000
+ if counts < min_counts:
1001
+ min_counts = counts
1002
+ eng_index = idx
1003
+ # Adjust local counts for better balancing between stats updates
1004
+ # from the coordinator (which happen every 100ms).
1005
+ if min_counts[0]:
1006
+ min_counts[0] += 1
1007
+ else:
1008
+ min_counts[1] += 1
1009
+
1010
+ chosen_engine = self.core_engines[eng_index]
1011
+ # Record which engine is chosen for this request, to handle aborts.
1012
+ self.reqs_in_flight[request.request_id] = chosen_engine
1013
+ return chosen_engine
1014
+
1015
+ async def call_utility_async(self, method: str, *args) -> Any:
1016
+ # Only the result from the first engine is returned.
1017
+ return (await asyncio.gather(*[
1018
+ self._call_utility_async(method, *args, engine=engine)
1019
+ for engine in self.core_engines
1020
+ ]))[0]
1021
+
1022
@staticmethod
1072
- async def process_engine_outputs(self: "DPAsyncMPClient",
1023
+ async def process_engine_outputs(self: "DPLBAsyncMPClient",
1024
outputs: EngineCoreOutputs):
1025
if outputs.finished_requests and self.reqs_in_flight:
1026
for req_id in outputs.finished_requests:
@@ -1085,61 +1036,14 @@ class DPAsyncMPClient(AsyncMPClient):
1036
await self._abort_requests(request_ids, engine)
1037
return
1038
1088
- by_engine: dict[CoreEngine, list[str]] = {}
1039
+ by_engine = defaultdict[EngineIdentity, list[str]](list)
1040
for req_id in request_ids:
1041
if engine := self.reqs_in_flight.get(req_id):
1091
- by_engine.setdefault(engine, []).append(req_id)
1042
+ by_engine[engine].append(req_id)
1043
for engine, req_ids in by_engine.items():
1044
await self._abort_requests(req_ids, engine)
1045
1046
async def _abort_requests(self, request_ids: list[str],
1096
- engine: CoreEngine) -> None:
1047
+ engine: EngineIdentity) -> None:
1048
await self._send_input(EngineCoreRequestType.ABORT, request_ids,
1049
engine)
1099
-
1100
-
1101
-class RayDPClient(DPAsyncMPClient):
1102
- """
1103
- Ray-based client for multi-proc, multi-engine (data parallel)
1104
- EngineCore.
1105
- """
1106
-
1107
- def __init__(
1108
- self,
1109
- vllm_config: VllmConfig,
1110
- executor_class: type[Executor],
1111
- log_stats: bool,
1112
- client_addresses: Optional[dict[str, str]] = None,
1113
- client_index: int = 0,
1114
- ):
1115
- super().__init__(vllm_config, executor_class, log_stats,
1116
- client_addresses, client_index)
1117
-
1118
- def _init_engines_direct(self, vllm_config: VllmConfig, local_only: bool,
1119
- local_start_index: int, input_address: str,
1120
- output_address: str,
1121
- executor_class: type[Executor], log_stats: bool):
1122
- """Self-contained client mode, launch engine and coordinator process
1123
- as needed."""
1124
-
1125
- parallel_config = vllm_config.parallel_config
1126
- assert parallel_config.data_parallel_rank == 0
1127
- assert local_start_index == 0
1128
-
1129
- addresses = EngineZmqAddresses(
1130
- inputs=[input_address],
1131
- outputs=[output_address],
1132
- )
1133
-
1134
- if len(self.core_engines) > 1:
1135
- coordinator = DPCoordinator(parallel_config)
1136
- self.resources.coordinator = coordinator
1137
- addresses.coordinator_input, addresses.coordinator_output = (
1138
- coordinator.get_engine_socket_addresses())
1139
-
1140
- # Start all engines.
1141
- self.resources.engine_manager = CoreEngineActorManager(
1142
- vllm_config=vllm_config,
1143
- addresses=addresses,
1144
- executor_class=executor_class,
1145
- log_stats=log_stats)
vllm/v1/engine/utils.py
new
+546
@@ -0,0 +1,546 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import contextlib
5
+import weakref
6
+from collections.abc import Iterator
7
+from dataclasses import dataclass
8
+from enum import Enum, auto
9
+from multiprocessing import Process, connection
10
+from multiprocessing.process import BaseProcess
11
+from typing import TYPE_CHECKING, Callable, Optional, Union
12
+
13
+import msgspec
14
+import zmq
15
+
16
+from vllm.config import CacheConfig, ParallelConfig, VllmConfig
17
+from vllm.logger import init_logger
18
+from vllm.utils import get_mp_context, get_open_zmq_ipc_path, zmq_socket_ctx
19
+from vllm.v1.engine.coordinator import DPCoordinator
20
+from vllm.v1.executor.abstract import Executor
21
+from vllm.v1.utils import get_engine_client_zmq_addr, shutdown
22
+
23
+if TYPE_CHECKING:
24
+ from ray.util.placement_group import PlacementGroup
25
+
26
+logger = init_logger(__name__)
27
+
28
+STARTUP_POLL_PERIOD_MS = 10000
29
+
30
+
31
+class CoreEngineState(Enum):
32
+ NEW = auto()
33
+ CONNECTED = auto()
34
+ READY = auto()
35
+
36
+
37
+class CoreEngine:
38
+ """One per data parallel rank, used to track state during handshaking."""
39
+
40
+ def __init__(self, index: int = 0, local: bool = True):
41
+ self.local = local
42
+ self.identity = index.to_bytes(2, "little")
43
+
44
+ self.state = CoreEngineState.NEW
45
+
46
+
47
+@dataclass
48
+class EngineZmqAddresses:
49
+ # ZMQ input socket addresses for each front-end client (requests)
50
+ inputs: list[str]
51
+ # ZMQ output socket addresses for each front-end client (responses)
52
+ outputs: list[str]
53
+ # ZMQ input socket address of DP coordinator if applicable
54
+ coordinator_input: Optional[str] = None
55
+ # ZMQ output socket address of DP coordinator if applicable
56
+ coordinator_output: Optional[str] = None
57
+ # ZMQ socket for front-end to connect to DP coordinator.
58
+ # Not used by engine, just relayed to front-end in handshake response.
59
+ # Only required for external DP LB case.
60
+ frontend_stats_publish_address: Optional[str] = None
61
+
62
+
63
+@dataclass
64
+class EngineHandshakeMetadata:
65
+ """Metadata sent to each engine process during startup handshake,
66
+ including addresses of the front-end ZMQ queues that they should
67
+ connect to.
68
+ """
69
+ addresses: EngineZmqAddresses
70
+ parallel_config: dict[str, Union[int, str]]
71
+
72
+
73
+class CoreEngineProcManager:
74
+ """
75
+ Utility class to handle creation, readiness, and shutdown
76
+ of background processes used by the AsyncLLM and LLMEngine.
77
+ """
78
+
79
+ def __init__(
80
+ self,
81
+ target_fn: Callable,
82
+ local_engine_count: int,
83
+ start_index: int,
84
+ local_start_index: int,
85
+ vllm_config: VllmConfig,
86
+ local_client: bool,
87
+ handshake_address: str,
88
+ executor_class: type[Executor],
89
+ log_stats: bool,
90
+ client_handshake_address: Optional[str] = None,
91
+ ):
92
+ context = get_mp_context()
93
+ common_kwargs = {
94
+ "vllm_config": vllm_config,
95
+ "local_client": local_client,
96
+ "handshake_address": handshake_address,
97
+ "executor_class": executor_class,
98
+ "log_stats": log_stats,
99
+ }
100
+
101
+ if client_handshake_address:
102
+ common_kwargs[
103
+ "client_handshake_address"] = client_handshake_address
104
+
105
+ self.processes: list[BaseProcess] = []
106
+ for index in range(local_engine_count):
107
+ local_index = local_start_index + index
108
+ global_index = start_index + index
109
+ # Start EngineCore in background process.
110
+ self.processes.append(
111
+ context.Process(target=target_fn,
112
+ name=f"EngineCore_{global_index}",
113
+ kwargs=common_kwargs | {
114
+ "dp_rank": global_index,
115
+ "local_dp_rank": local_index,
116
+ }))
117
+
118
+ self._finalizer = weakref.finalize(self, shutdown, self.processes)
119
+ try:
120
+ for proc in self.processes:
121
+ proc.start()
122
+ finally:
123
+ # Kill other procs if not all are running.
124
+ if self.finished_procs():
125
+ self.close()
126
+
127
+ def close(self):
128
+ """Shutdown all procs."""
129
+ self._finalizer()
130
+
131
+ def join_first(self):
132
+ """Wait for any process to exit."""
133
+ connection.wait(proc.sentinel for proc in self.processes)
134
+
135
+ def sentinels(self) -> list:
136
+ return [proc.sentinel for proc in self.processes]
137
+
138
+ def finished_procs(self) -> dict[str, int]:
139
+ """Returns dict of proc name -> exit code for any finished procs."""
140
+ return {
141
+ proc.name: proc.exitcode
142
+ for proc in self.processes if proc.exitcode is not None
143
+ }
144
+
145
+
146
+class CoreEngineActorManager:
147
+ """
148
+ Utility class to handle creation, readiness, and shutdown
149
+ of core engine Ray actors used by the AsyncLLM and LLMEngine.
150
+
151
+ Different from CoreEngineProcManager, this class manages
152
+ core engines for both local and remote nodes.
153
+ """
154
+
155
+ def __init__(
156
+ self,
157
+ vllm_config: VllmConfig,
158
+ addresses: EngineZmqAddresses,
159
+ executor_class: type[Executor],
160
+ log_stats: bool,
161
+ placement_groups: Optional[list["PlacementGroup"]] = None,
162
+ local_dp_ranks: Optional[list[int]] = None,
163
+ ):
164
+ import copy
165
+
166
+ import ray
167
+ from ray.util.scheduling_strategies import (
168
+ PlacementGroupSchedulingStrategy)
169
+
170
+ from vllm.v1.engine.core import DPEngineCoreActor
171
+
172
+ self.local_engine_actors: list[ray.ActorHandle] = []
173
+ self.remote_engine_actors: list[ray.ActorHandle] = []
174
+ dp_size = vllm_config.parallel_config.data_parallel_size
175
+ local_engine_count = \
176
+ vllm_config.parallel_config.data_parallel_size_local
177
+ world_size = vllm_config.parallel_config.world_size
178
+
179
+ if ray.is_initialized():
180
+ logger.info(
181
+ "Ray is already initialized. Skipping Ray initialization.")
182
+ else:
183
+ ray.init()
184
+
185
+ if placement_groups is not None:
186
+ assert local_dp_ranks is not None, (
187
+ "local_dp_ranks must be provided if "
188
+ "placement_groups is provided")
189
+ assert len(placement_groups) == len(local_dp_ranks), (
190
+ "placement_groups and local_dp_ranks must "
191
+ "have the same length")
192
+ logger.info("Using provided placement groups")
193
+ # TODO(rui): validate passed-in placement groups
194
+ self.created_placement_groups = []
195
+ else:
196
+ placement_groups, local_dp_ranks = \
197
+ CoreEngineActorManager.create_dp_placement_groups(vllm_config)
198
+ self.created_placement_groups = placement_groups
199
+ assert len(placement_groups) == dp_size, (
200
+ "Number of placement groups must match data parallel size")
201
+
202
+ refs = []
203
+ for index in range(dp_size):
204
+ local_index = local_dp_ranks[index]
205
+ dp_vllm_config = copy.deepcopy(vllm_config)
206
+ pg = placement_groups[index]
207
+ dp_vllm_config.parallel_config.placement_group = pg
208
+ local_client = index < local_engine_count
209
+ actor = ray.remote(DPEngineCoreActor).options(
210
+ scheduling_strategy=PlacementGroupSchedulingStrategy(
211
+ placement_group=pg,
212
+ placement_group_bundle_index=world_size,
213
+ )).remote(vllm_config=dp_vllm_config,
214
+ executor_class=executor_class,
215
+ log_stats=log_stats,
216
+ local_client=local_client,
217
+ addresses=addresses,
218
+ dp_rank=index,
219
+ local_dp_rank=local_index)
220
+ if local_client:
221
+ self.local_engine_actors.append(actor)
222
+ else:
223
+ self.remote_engine_actors.append(actor)
224
+ refs.append(actor.wait_for_init.remote())
225
+
226
+ ray.get(refs)
227
+ self.run_refs = []
228
+ for actor in self.local_engine_actors + self.remote_engine_actors:
229
+ self.run_refs.append(actor.run.remote())
230
+
231
+ @staticmethod
232
+ def create_dp_placement_groups(
233
+ vllm_config: VllmConfig
234
+ ) -> tuple[list["PlacementGroup"], list[int]]:
235
+
236
+ import ray
237
+ from ray._private.state import available_resources_per_node
238
+ from ray.util.state import list_nodes
239
+
240
+ logger.info("Creating placement groups for data parallel")
241
+ dp_master_ip = \
242
+ vllm_config.parallel_config.data_parallel_master_ip
243
+ dp_size = vllm_config.parallel_config.data_parallel_size
244
+ local_engine_count = \
245
+ vllm_config.parallel_config.data_parallel_size_local
246
+
247
+ nodes = sorted(list_nodes(),
248
+ key=lambda node: node.node_ip != dp_master_ip)
249
+ assert nodes[0].node_ip == dp_master_ip, (
250
+ "The first node must be the head node")
251
+ assert len(nodes) == 1 or nodes[1].node_ip != dp_master_ip, (
252
+ "There can only be one head node")
253
+
254
+ available_resources = available_resources_per_node()
255
+ world_size = vllm_config.parallel_config.world_size
256
+ placement_groups: list[PlacementGroup] = []
257
+ local_dp_ranks: list[int] = []
258
+
259
+ for node in nodes:
260
+ node_ip = node.node_ip
261
+ node_resources = available_resources[node.node_id]
262
+ # For now, each DP rank can only be assigned to one node
263
+ # TODO(rui): support allocating a single DP rank
264
+ # to multiple nodes
265
+ available_engine_count = int(node_resources["GPU"]) // world_size
266
+ if node_ip == dp_master_ip:
267
+ assert available_engine_count >= local_engine_count, (
268
+ "Not enough resources to allocate DP ranks "
269
+ f"on DP master node {node_ip}")
270
+ for i in range(local_engine_count):
271
+ bundles = [{
272
+ "GPU": 1.0,
273
+ "node:" + dp_master_ip: 0.001
274
+ }] * world_size + [{
275
+ "CPU": 1.0
276
+ }]
277
+ pg = ray.util.placement_group(
278
+ name=f"dp_rank_{len(placement_groups)}",
279
+ strategy="STRICT_PACK",
280
+ bundles=bundles,
281
+ )
282
+ placement_groups.append(pg)
283
+ local_dp_ranks.append(i)
284
+ else:
285
+ for i in range(available_engine_count):
286
+ if len(placement_groups) == dp_size:
287
+ break
288
+ bundles = [{"GPU": 1.0}] * world_size + [{"CPU": 1.0}]
289
+ pg = ray.util.placement_group(
290
+ name=f"dp_rank_{len(placement_groups)}",
291
+ strategy="STRICT_PACK",
292
+ bundles=bundles,
293
+ )
294
+ placement_groups.append(pg)
295
+ local_dp_ranks.append(i)
296
+ return placement_groups, local_dp_ranks
297
+
298
+ def get_run_refs(self):
299
+ return self.run_refs
300
+
301
+ def close(self):
302
+ import ray
303
+ for actor in self.local_engine_actors + self.remote_engine_actors:
304
+ ray.kill(actor)
305
+ for pg in self.created_placement_groups:
306
+ ray.util.remove_placement_group(pg)
307
+
308
+
309
+@contextlib.contextmanager
310
+def launch_core_engines(
311
+ vllm_config: VllmConfig,
312
+ executor_class: type[Executor],
313
+ log_stats: bool,
314
+ num_api_servers: int = 1,
315
+) -> Iterator[tuple[
316
+ Optional[Union[CoreEngineProcManager, CoreEngineActorManager]],
317
+ Optional[DPCoordinator],
318
+ EngineZmqAddresses,
319
+]]:
320
+ """Launch engine and DP coordinator processes as needed."""
321
+
322
+ parallel_config = vllm_config.parallel_config
323
+ dp_size = parallel_config.data_parallel_size
324
+ local_engine_count = parallel_config.data_parallel_size_local
325
+ local_start_index = parallel_config.data_parallel_rank_local
326
+ dp_rank = parallel_config.data_parallel_rank
327
+ host = parallel_config.data_parallel_master_ip
328
+ external_dp_lb = parallel_config.data_parallel_external_lb
329
+
330
+ # In offline mode there is an LLM instance per DP rank and
331
+ # one core engine per LLM, see
332
+ # examples/offline_inference/data_parallel.py.
333
+ offline_mode = local_start_index is not None
334
+
335
+ # client_local_only = True for cases where this front-end
336
+ # sends requests only to colocated engines.
337
+ client_local_only = offline_mode or external_dp_lb or (local_engine_count
338
+ == dp_size)
339
+
340
+ # Set up input and output addresses.
341
+ addresses = EngineZmqAddresses(
342
+ inputs=[
343
+ get_engine_client_zmq_addr(client_local_only, host)
344
+ for _ in range(num_api_servers)
345
+ ],
346
+ outputs=[
347
+ get_engine_client_zmq_addr(client_local_only, host)
348
+ for _ in range(num_api_servers)
349
+ ],
350
+ )
351
+
352
+ # Run the DP Coordinator process with rank 0 when in
353
+ # online DP mode.
354
+ run_coordinator = dp_size > 1 and not offline_mode and dp_rank == 0
355
+
356
+ if run_coordinator:
357
+ coordinator = DPCoordinator(parallel_config)
358
+
359
+ addresses.coordinator_input, addresses.coordinator_output = (
360
+ coordinator.get_engine_socket_addresses())
361
+ addresses.frontend_stats_publish_address = (
362
+ coordinator.get_stats_publish_address())
363
+
364
+ logger.info("Started DP Coordinator process (PID: %d)",
365
+ coordinator.proc.pid)
366
+ else:
367
+ coordinator = None
368
+
369
+ if parallel_config.data_parallel_backend == "ray":
370
+ logger.info("Starting ray-based data parallel backend")
371
+
372
+ engine_actor_manager = CoreEngineActorManager(
373
+ vllm_config=vllm_config,
374
+ addresses=addresses,
375
+ executor_class=executor_class,
376
+ log_stats=log_stats,
377
+ )
378
+
379
+ yield engine_actor_manager, coordinator, addresses
380
+ return
381
+
382
+ if offline_mode or (external_dp_lb and dp_rank > 0):
383
+ assert local_engine_count == 1
384
+ engines_to_handshake = [CoreEngine(index=dp_rank, local=True)]
385
+ else:
386
+ engines_to_handshake = [
387
+ CoreEngine(index=i, local=(i < local_engine_count))
388
+ for i in range(dp_size)
389
+ ]
390
+
391
+ # Whether the started engines will handshake only with co-located
392
+ # front-end processes. In external_dp_lb mode, ranks > 0 handshake with
393
+ # their co-located frontend and also the rank 0 front-end, and hence this
394
+ # will be False.
395
+ handshake_local_only = offline_mode or local_engine_count == dp_size
396
+
397
+ handshake_address = get_engine_client_zmq_addr(
398
+ handshake_local_only, host, parallel_config.data_parallel_rpc_port)
399
+
400
+ if external_dp_lb and dp_rank > 0:
401
+ assert not handshake_local_only
402
+ local_handshake_address = get_open_zmq_ipc_path()
403
+ client_handshake_address = local_handshake_address
404
+ else:
405
+ local_handshake_address = handshake_address
406
+ client_handshake_address = None
407
+
408
+ with zmq_socket_ctx(local_handshake_address, zmq.ROUTER,
409
+ bind=True) as handshake_socket:
410
+
411
+ from vllm.v1.engine.core import EngineCoreProc
412
+
413
+ # Start local engines.
414
+ if local_engine_count:
415
+ # In server mode, start_index and local_start_index will
416
+ # both be 0.
417
+ local_engine_manager = CoreEngineProcManager(
418
+ EngineCoreProc.run_engine_core,
419
+ vllm_config=vllm_config,
420
+ executor_class=executor_class,
421
+ log_stats=log_stats,
422
+ handshake_address=handshake_address,
423
+ client_handshake_address=client_handshake_address,
424
+ local_client=True,
425
+ local_engine_count=local_engine_count,
426
+ start_index=dp_rank,
427
+ local_start_index=local_start_index or 0)
428
+ else:
429
+ local_engine_manager = None
430
+
431
+ yield local_engine_manager, coordinator, addresses
432
+
433
+ # Now wait for engines to start.
434
+ wait_for_engine_startup(
435
+ handshake_socket,
436
+ addresses,
437
+ engines_to_handshake,
438
+ parallel_config,
439
+ vllm_config.cache_config,
440
+ local_engine_manager,
441
+ coordinator.proc if coordinator else None,
442
+ )
443
+
444
+
445
+def wait_for_engine_startup(
446
+ handshake_socket: zmq.Socket,
447
+ addresses: EngineZmqAddresses,
448
+ core_engines: list[CoreEngine],
449
+ parallel_config: ParallelConfig,
450
+ cache_config: CacheConfig,
451
+ proc_manager: Optional[CoreEngineProcManager],
452
+ coord_process: Optional[Process],
453
+):
454
+ # Wait for engine core process(es) to send ready messages.
455
+ local_count = parallel_config.data_parallel_size_local
456
+ remote_count = len(core_engines) - local_count
457
+ # [local, remote] counts
458
+ conn_pending, start_pending = [local_count, remote_count], [0, 0]
459
+ poller = zmq.Poller()
460
+ poller.register(handshake_socket, zmq.POLLIN)
461
+
462
+ if proc_manager is not None:
463
+ for sentinel in proc_manager.sentinels():
464
+ poller.register(sentinel, zmq.POLLIN)
465
+ if coord_process is not None:
466
+ poller.register(coord_process.sentinel, zmq.POLLIN)
467
+ while any(conn_pending) or any(start_pending):
468
+ events = poller.poll(STARTUP_POLL_PERIOD_MS)
469
+ if not events:
470
+ if any(conn_pending):
471
+ logger.debug(
472
+ "Waiting for %d local, %d remote core engine proc(s) "
473
+ "to connect.", *conn_pending)
474
+ if any(start_pending):
475
+ logger.debug(
476
+ "Waiting for %d local, %d remote core engine proc(s) "
477
+ "to start.", *start_pending)
478
+ continue
479
+ if len(events) > 1 or events[0][0] != handshake_socket:
480
+ # One of the local core processes exited.
481
+ finished = proc_manager.finished_procs() if proc_manager else {}
482
+ if coord_process is not None and coord_process.exitcode is not None:
483
+ finished[coord_process.name] = coord_process.exitcode
484
+ raise RuntimeError("Engine core initialization failed. "
485
+ "See root cause above. "
486
+ f"Failed core proc(s): {finished}")
487
+
488
+ # Receive HELLO and READY messages from the input socket.
489
+ eng_identity, ready_msg_bytes = handshake_socket.recv_multipart()
490
+ eng_index = int.from_bytes(eng_identity, "little")
491
+ engine = next((e for e in core_engines if e.identity == eng_identity),
492
+ None)
493
+ if engine is None:
494
+ raise RuntimeError(f"Message from engine with unexpected data "
495
+ f"parallel rank: {eng_index}")
496
+ msg = msgspec.msgpack.decode(ready_msg_bytes)
497
+ status, local = msg["status"], msg["local"]
498
+ if local != engine.local:
499
+ raise RuntimeError(f"{status} message from "
500
+ f"{'local' if local else 'remote'} "
501
+ f"engine {eng_index}, expected it to be "
502
+ f"{'local' if engine.local else 'remote'}")
503
+
504
+ if status == "HELLO" and engine.state == CoreEngineState.NEW:
505
+
506
+ # Send init message with DP config info.
507
+ init_message = msgspec.msgpack.encode(
508
+ EngineHandshakeMetadata(
509
+ addresses=addresses,
510
+ parallel_config={
511
+ "data_parallel_master_ip":
512
+ parallel_config.data_parallel_master_ip,
513
+ "data_parallel_master_port":
514
+ parallel_config.data_parallel_master_port,
515
+ "data_parallel_size":
516
+ parallel_config.data_parallel_size,
517
+ }))
518
+ handshake_socket.send_multipart((eng_identity, init_message),
519
+ copy=False)
520
+ conn_pending[0 if local else 1] -= 1
521
+ start_pending[0 if local else 1] += 1
522
+ engine.state = CoreEngineState.CONNECTED
523
+ elif status == "READY" and engine.state == CoreEngineState.CONNECTED:
524
+ # Setup KV cache config with initialization state from
525
+ # engine core process. Sum values from all engines in DP case.
526
+ num_gpu_blocks = cache_config.num_gpu_blocks or 0
527
+ num_gpu_blocks += msg["num_gpu_blocks"]
528
+ cache_config.num_gpu_blocks = num_gpu_blocks
529
+
530
+ # In external DP LB mode, the coordinator address that the
531
+ # front-end procs connect to is obtained from rank 0 via
532
+ # one of the engine handshakes, and passed to the local
533
+ # front-end process in the response from the other.
534
+ if addresses.frontend_stats_publish_address is None:
535
+ addresses.frontend_stats_publish_address = msg.get(
536
+ "dp_stats_address")
537
+
538
+ start_pending[0 if local else 1] -= 1
539
+ engine.state = CoreEngineState.READY
540
+ else:
541
+ raise RuntimeError(f"Unexpected {status} message for "
542
+ f"{'local' if local else 'remote'} engine "
543
+ f"{eng_index} in {engine.state} state.")
544
+
545
+ logger.debug("%s from %s core engine process %s.", status,
546
+ "local" if local else "remote", eng_index)
vllm/v1/utils.py
+17
-383
@@ -1,44 +1,35 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
import argparse
4
import multiprocessing
5
import time
6
import weakref
7
from collections import defaultdict
8
from collections.abc import Sequence
10
-from dataclasses import dataclass
11
-from enum import Enum, auto
12
-from multiprocessing import Process, connection
9
+from multiprocessing import connection
10
from multiprocessing.process import BaseProcess
11
from typing import (TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar,
12
Union, overload)
13
17
-import msgspec
14
import torch
19
-import zmq
15
21
-from vllm.config import CacheConfig, ParallelConfig, VllmConfig
16
from vllm.logger import init_logger
17
from vllm.model_executor.models.utils import extract_layer_index
18
from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled,
19
usage_message)
26
-from vllm.utils import (get_mp_context, get_open_port, get_open_zmq_ipc_path,
27
- get_tcp_uri, kill_process_tree)
28
-from vllm.v1.executor.abstract import Executor
20
+from vllm.utils import (get_open_port, get_open_zmq_ipc_path, get_tcp_uri,
21
+ kill_process_tree)
22
23
if TYPE_CHECKING:
31
- from ray.util.placement_group import PlacementGroup
32
-
24
from vllm.attention.layer import Attention
25
from vllm.v1.engine.coordinator import DPCoordinator
26
+ from vllm.v1.engine.utils import (CoreEngineActorManager,
27
+ CoreEngineProcManager)
28
29
logger = init_logger(__name__)
30
31
T = TypeVar("T")
32
40
-STARTUP_POLL_PERIOD_MS = 10000
41
-
33
34
class ConstantList(Generic[T], Sequence):
35
@@ -111,47 +102,16 @@ class ConstantList(Generic[T], Sequence):
102
def get_engine_client_zmq_addr(local_only: bool,
103
host: str,
104
port: int = 0) -> str:
114
- return get_open_zmq_ipc_path() if local_only else (get_tcp_uri(
115
- host, port or get_open_port()))
116
-
117
-
118
-class CoreEngineState(Enum):
119
- NEW = auto()
120
- CONNECTED = auto()
121
- READY = auto()
122
-
123
-
124
-class CoreEngine:
125
- """One per data parallel rank."""
126
-
127
- def __init__(self, index: int = 0, local: bool = True):
128
- self.local = local
129
- self.index = index
130
- self.identity = index.to_bytes(2, "little")
131
-
132
- self.state = CoreEngineState.NEW
105
+ """Assign a new ZMQ socket address.
106
107
+ If local_only is True, participants are colocated and so a unique IPC
108
+ address will be returned.
109
135
-@dataclass
136
-class EngineZmqAddresses:
137
- # ZMQ input socket addresses for each front-end client (requests)
138
- inputs: list[str]
139
- # ZMQ output socket addresses for each front-end client (responses)
140
- outputs: list[str]
141
- # ZMQ input socket address of DP coordinator if applicable
142
- coordinator_input: Optional[str] = None
143
- # ZMQ output socket address of DP coordinator if applicable
144
- coordinator_output: Optional[str] = None
110
+ Otherwise, the provided host and port will be used to construct a TCP
111
+ address (port == 0 means assign an available port)."""
112
146
-
147
-@dataclass
148
-class EngineHandshakeMetadata:
149
- """Metadata sent to each engine process during startup handshake,
150
- including addresses of the front-end ZMQ queues that they should
151
- connect to.
152
- """
153
- addresses: EngineZmqAddresses
154
- parallel_config: dict[str, Union[int, str]]
113
+ return get_open_zmq_ipc_path() if local_only else (get_tcp_uri(
114
+ host, port or get_open_port()))
115
116
117
class APIServerProcessManager:
@@ -219,339 +179,10 @@ class APIServerProcessManager:
179
self._finalizer()
180
181
222
-class CoreEngineProcManager:
223
- """
224
- Utility class to handle creation, readiness, and shutdown
225
- of background processes used by the AsyncLLM and LLMEngine.
226
- """
227
-
228
- def __init__(
229
- self,
230
- target_fn: Callable,
231
- local_engine_count: int,
232
- start_index: int,
233
- local_start_index: int,
234
- vllm_config: VllmConfig,
235
- on_head_node: bool,
236
- handshake_address: str,
237
- executor_class: type[Executor],
238
- log_stats: bool,
239
- ):
240
- context = get_mp_context()
241
- common_kwargs = {
242
- "vllm_config": vllm_config,
243
- "on_head_node": on_head_node,
244
- "handshake_address": handshake_address,
245
- "executor_class": executor_class,
246
- "log_stats": log_stats,
247
- }
248
-
249
- self.processes: list[BaseProcess] = []
250
- for index in range(local_engine_count):
251
- local_index = local_start_index + index
252
- global_index = start_index + index
253
- # Start EngineCore in background process.
254
- self.processes.append(
255
- context.Process(target=target_fn,
256
- name=f"EngineCore_{global_index}",
257
- kwargs=common_kwargs | {
258
- "dp_rank": global_index,
259
- "local_dp_rank": local_index,
260
- }))
261
-
262
- self._finalizer = weakref.finalize(self, shutdown, self.processes)
263
- try:
264
- for proc in self.processes:
265
- proc.start()
266
- finally:
267
- # Kill other procs if not all are running.
268
- if self.finished_procs():
269
- self.close()
270
-
271
- def close(self):
272
- """Shutdown all procs."""
273
- self._finalizer()
274
-
275
- def join_first(self):
276
- """Wait for any process to exit."""
277
- connection.wait(proc.sentinel for proc in self.processes)
278
-
279
- def sentinels(self) -> list:
280
- return [proc.sentinel for proc in self.processes]
281
-
282
- def finished_procs(self) -> dict[str, int]:
283
- """Returns dict of proc name -> exit code for any finished procs."""
284
- return {
285
- proc.name: proc.exitcode
286
- for proc in self.processes if proc.exitcode is not None
287
- }
288
-
289
-
290
-class CoreEngineActorManager:
291
- """
292
- Utility class to handle creation, readiness, and shutdown
293
- of core engine Ray actors used by the AsyncLLM and LLMEngine.
294
-
295
- Different from CoreEngineProcManager, this class manages
296
- core engines for both local and remote nodes.
297
- """
298
-
299
- def __init__(
300
- self,
301
- vllm_config: VllmConfig,
302
- addresses: EngineZmqAddresses,
303
- executor_class: type[Executor],
304
- log_stats: bool,
305
- placement_groups: Optional[list["PlacementGroup"]] = None,
306
- local_dp_ranks: Optional[list[int]] = None,
307
- ):
308
- import copy
309
-
310
- import ray
311
- from ray.util.scheduling_strategies import (
312
- PlacementGroupSchedulingStrategy)
313
-
314
- from vllm.v1.engine.core import DPEngineCoreActor
315
-
316
- self.local_engine_actors: list[ray.ActorHandle] = []
317
- self.remote_engine_actors: list[ray.ActorHandle] = []
318
- dp_size = vllm_config.parallel_config.data_parallel_size
319
- local_engine_count = \
320
- vllm_config.parallel_config.data_parallel_size_local
321
- world_size = vllm_config.parallel_config.world_size
322
-
323
- if ray.is_initialized():
324
- logger.info(
325
- "Ray is already initialized. Skipping Ray initialization.")
326
- else:
327
- ray.init()
328
-
329
- if placement_groups is not None:
330
- assert local_dp_ranks is not None, (
331
- "local_dp_ranks must be provided if "
332
- "placement_groups is provided")
333
- assert len(placement_groups) == len(local_dp_ranks), (
334
- "placement_groups and local_dp_ranks must "
335
- "have the same length")
336
- logger.info("Using provided placement groups")
337
- # TODO(rui): validate passed-in placement groups
338
- self.created_placement_groups = []
339
- else:
340
- placement_groups, local_dp_ranks = \
341
- CoreEngineActorManager.create_dp_placement_groups(vllm_config)
342
- self.created_placement_groups = placement_groups
343
- assert len(placement_groups) == dp_size, (
344
- "Number of placement groups must match data parallel size")
345
-
346
- refs = []
347
- for index in range(dp_size):
348
- local_index = local_dp_ranks[index]
349
- dp_vllm_config = copy.deepcopy(vllm_config)
350
- pg = placement_groups[index]
351
- dp_vllm_config.parallel_config.placement_group = pg
352
- on_head_node = index < local_engine_count
353
- actor = ray.remote(DPEngineCoreActor).options(
354
- scheduling_strategy=PlacementGroupSchedulingStrategy(
355
- placement_group=pg,
356
- placement_group_bundle_index=world_size,
357
- )).remote(vllm_config=dp_vllm_config,
358
- executor_class=executor_class,
359
- log_stats=log_stats,
360
- on_head_node=on_head_node,
361
- addresses=addresses,
362
- dp_rank=index,
363
- local_dp_rank=local_index)
364
- if on_head_node:
365
- self.local_engine_actors.append(actor)
366
- else:
367
- self.remote_engine_actors.append(actor)
368
- refs.append(actor.wait_for_init.remote())
369
-
370
- ray.get(refs)
371
- self.run_refs = []
372
- for actor in self.local_engine_actors + self.remote_engine_actors:
373
- self.run_refs.append(actor.run.remote())
374
-
375
- @staticmethod
376
- def create_dp_placement_groups(
377
- vllm_config: VllmConfig
378
- ) -> tuple[list["PlacementGroup"], list[int]]:
379
-
380
- import ray
381
- from ray._private.state import available_resources_per_node
382
- from ray.util.state import list_nodes
383
-
384
- logger.info("Creating placement groups for data parallel")
385
- dp_master_ip = \
386
- vllm_config.parallel_config.data_parallel_master_ip
387
- dp_size = vllm_config.parallel_config.data_parallel_size
388
- local_engine_count = \
389
- vllm_config.parallel_config.data_parallel_size_local
390
-
391
- nodes = list_nodes()
392
- nodes = sorted(list_nodes(),
393
- key=lambda node: node.node_ip != dp_master_ip)
394
- assert nodes[0].node_ip == dp_master_ip, (
395
- "The first node must be the head node")
396
- assert len(nodes) == 1 or nodes[1].node_ip != dp_master_ip, (
397
- "There can only be one head node")
398
-
399
- available_resources = available_resources_per_node()
400
- world_size = vllm_config.parallel_config.world_size
401
- placement_groups: list[PlacementGroup] = []
402
- local_dp_ranks: list[int] = []
403
-
404
- for node in nodes:
405
- node_ip = node.node_ip
406
- node_resources = available_resources[node.node_id]
407
- # For now, each DP rank can only be assigned to one node
408
- # TODO(rui): support allocating a single DP rank
409
- # to multiple nodes
410
- available_engine_count = int(node_resources["GPU"]) // world_size
411
- if node_ip == dp_master_ip:
412
- assert available_engine_count >= local_engine_count, (
413
- "Not enough resources to allocate DP ranks "
414
- f"on DP master node {node_ip}")
415
- for i in range(local_engine_count):
416
- bundles = [{
417
- "GPU": 1.0,
418
- "node:" + dp_master_ip: 0.001
419
- }] * world_size + [{
420
- "CPU": 1.0
421
- }]
422
- pg = ray.util.placement_group(
423
- name=f"dp_rank_{len(placement_groups)}",
424
- strategy="STRICT_PACK",
425
- bundles=bundles,
426
- )
427
- placement_groups.append(pg)
428
- local_dp_ranks.append(i)
429
- else:
430
- for i in range(available_engine_count):
431
- if len(placement_groups) == dp_size:
432
- break
433
- bundles = [{"GPU": 1.0}] * world_size + [{"CPU": 1.0}]
434
- pg = ray.util.placement_group(
435
- name=f"dp_rank_{len(placement_groups)}",
436
- strategy="STRICT_PACK",
437
- bundles=bundles,
438
- )
439
- placement_groups.append(pg)
440
- local_dp_ranks.append(i)
441
- return placement_groups, local_dp_ranks
442
-
443
- def get_run_refs(self):
444
- return self.run_refs
445
-
446
- def close(self):
447
- import ray
448
- for actor in self.local_engine_actors + self.remote_engine_actors:
449
- ray.kill(actor)
450
- for pg in self.created_placement_groups:
451
- ray.util.remove_placement_group(pg)
452
-
453
-
454
-def wait_for_engine_startup(
455
- handshake_socket: zmq.Socket,
456
- addresses: EngineZmqAddresses,
457
- core_engines: list[CoreEngine],
458
- parallel_config: ParallelConfig,
459
- cache_config: CacheConfig,
460
- proc_manager: Optional[CoreEngineProcManager],
461
- coord_process: Optional[Process],
462
-):
463
-
464
- # Wait for engine core process(es) to send ready messages.
465
- local_count = parallel_config.data_parallel_size_local
466
- remote_count = len(core_engines) - local_count
467
- # [local, remote] counts
468
- conn_pending, start_pending = [local_count, remote_count], [0, 0]
469
- poller = zmq.Poller()
470
- poller.register(handshake_socket, zmq.POLLIN)
471
-
472
- if proc_manager is not None:
473
- for sentinel in proc_manager.sentinels():
474
- poller.register(sentinel, zmq.POLLIN)
475
- if coord_process is not None:
476
- poller.register(coord_process.sentinel, zmq.POLLIN)
477
- while any(conn_pending) or any(start_pending):
478
- events = poller.poll(STARTUP_POLL_PERIOD_MS)
479
- if not events:
480
- if any(conn_pending):
481
- logger.debug(
482
- "Waiting for %d local, %d remote core engine proc(s) "
483
- "to connect.", *conn_pending)
484
- if any(start_pending):
485
- logger.debug(
486
- "Waiting for %d local, %d remote core engine proc(s) "
487
- "to start.", *start_pending)
488
- continue
489
- if len(events) > 1 or events[0][0] != handshake_socket:
490
- # One of the local core processes exited.
491
- finished = proc_manager.finished_procs() if proc_manager else {}
492
- if coord_process is not None and coord_process.exitcode is not None:
493
- finished[coord_process.name] = coord_process.exitcode
494
- raise RuntimeError("Engine core initialization failed. "
495
- "See root cause above. "
496
- f"Failed core proc(s): {finished}")
497
-
498
- # Receive HELLO and READY messages from the input socket.
499
- eng_identity, ready_msg_bytes = handshake_socket.recv_multipart()
500
- eng_index = int.from_bytes(eng_identity, "little")
501
- engine = next((e for e in core_engines if e.identity == eng_identity),
502
- None)
503
- if engine is None:
504
- raise RuntimeError(f"Message from engine with unexpected data "
505
- f"parallel rank: {eng_index}")
506
- msg = msgspec.msgpack.decode(ready_msg_bytes)
507
- status, local = msg["status"], msg["local"]
508
- if local != engine.local:
509
- raise RuntimeError(f"{status} message from "
510
- f"{'local' if local else 'remote'} "
511
- f"engine {eng_index}, expected it to be "
512
- f"{'local' if engine.local else 'remote'}")
513
-
514
- if status == "HELLO" and engine.state == CoreEngineState.NEW:
515
-
516
- # Send init message with DP config info.
517
- init_message = msgspec.msgpack.encode(
518
- EngineHandshakeMetadata(
519
- addresses=addresses,
520
- parallel_config={
521
- "data_parallel_master_ip":
522
- parallel_config.data_parallel_master_ip,
523
- "data_parallel_master_port":
524
- parallel_config.data_parallel_master_port,
525
- "data_parallel_size":
526
- parallel_config.data_parallel_size,
527
- }))
528
- handshake_socket.send_multipart((eng_identity, init_message),
529
- copy=False)
530
- conn_pending[0 if local else 1] -= 1
531
- start_pending[0 if local else 1] += 1
532
- engine.state = CoreEngineState.CONNECTED
533
- elif status == "READY" and (engine.state == CoreEngineState.CONNECTED):
534
- # Setup KV cache config with initialization state from
535
- # engine core process. Sum values from all engines in DP case.
536
- num_gpu_blocks = cache_config.num_gpu_blocks or 0
537
- num_gpu_blocks += msg["num_gpu_blocks"]
538
- cache_config.num_gpu_blocks = num_gpu_blocks
539
-
540
- start_pending[0 if local else 1] -= 1
541
- engine.state = CoreEngineState.READY
542
- else:
543
- raise RuntimeError(f"Unexpected {status} message for "
544
- f"{'local' if local else 'remote'} engine "
545
- f"{eng_index} in {engine.state} state.")
546
-
547
- logger.debug("%s from %s core engine process %s.", status,
548
- "local" if local else "remote", eng_index)
549
-
550
-
182
def wait_for_completion_or_failure(
183
api_server_manager: APIServerProcessManager,
553
- engine_manager: Optional[Union[CoreEngineProcManager,
554
- CoreEngineActorManager]] = None,
184
+ engine_manager: Optional[Union["CoreEngineProcManager",
185
+ "CoreEngineActorManager"]] = None,
186
coordinator: Optional["DPCoordinator"] = None) -> None:
187
"""Wait for all processes to complete or detect if any fail.
188
@@ -565,6 +196,9 @@ def wait_for_completion_or_failure(
196
coordinator: The coordinator for data parallel.
197
"""
198
199
+ from vllm.v1.engine.utils import (CoreEngineActorManager,
200
+ CoreEngineProcManager)
201
+
202
try:
203
logger.info("Waiting for API servers to complete ...")
204
# Create a mapping of sentinels to their corresponding processes