@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d5b981f8b1
Fork 0

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 14, 23:45 UTC

[DP] Internal Load Balancing Per Node [`one-pod-per-node`] (#21238)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Nick Hill <nhill@redhat.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

Robert Shaw committed Jul 23, 2025 at 23:57 UTC d5b981f8b1de31a54d55f9c0ead977dbf4d6b987
12 files changed +486 -45
.buildkite/test-pipeline.yaml
+2
@@ -166,6 +166,7 @@ steps:
166 - tests/v1/test_async_llm_dp.py
167 - tests/v1/test_external_lb_dp.py
168 - tests/v1/test_internal_lb_dp.py
169 + - tests/v1/test_hybrid_lb_dp.py
170 - tests/v1/engine/test_engine_core_client.py
171 commands:
172 # test with tp=2 and external_dp=2
@@ -178,6 +179,7 @@ steps:
179 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
180 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_external_lb_dp.py
181 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/test_internal_lb_dp.py
182 + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/test_hybrid_lb_dp.py
183 - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
184 - pytest -v -s distributed/test_utils.py
185 - pytest -v -s compile/test_basic_correctness.py
tests/v1/engine/test_engine_core_client.py
+2 -2
@@ -565,8 +565,8 @@ def test_engine_core_proc_instantiation_cuda_empty(
565
566 from vllm.v1.engine.utils import EngineZmqAddresses
567
568 - def mock_startup_handshake(self, handshake_socket, on_head_node,
569 - parallel_config):
568 + def mock_startup_handshake(self, handshake_socket, local_client,
569 + headless, parallel_config):
570 return EngineZmqAddresses(inputs=["tcp://127.0.0.1:5555"],
571 outputs=["tcp://127.0.0.1:5556"],
572 coordinator_input=None,
tests/v1/test_hybrid_lb_dp.py new
+352
@@ -0,0 +1,352 @@
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 tests.v1.test_utils import check_request_balancing
15 +from vllm.platforms import Platform
16 +
17 +MODEL_NAME = "ibm-research/PowerMoE-3b"
18 +
19 +# Number of data parallel ranks for hybrid LB testing (4 total)
20 +DP_SIZE = int(os.getenv("DP_SIZE", "4"))
21 +# Default tensor parallel size to use
22 +TP_SIZE = int(os.getenv("TP_SIZE", "1"))
23 +
24 +# Number of nodes (2 nodes, each with 2 DP ranks)
25 +NUM_NODES = 2
26 +DP_SIZE_LOCAL = DP_SIZE // NUM_NODES # 2 ranks per node
27 +
28 +
29 +class HybridLBServerManager:
30 + """Manages hybrid data parallel vLLM server instances where each node
31 + runs a single logical API server that balances requests only to the
32 + DP engines running on that same node."""
33 +
34 + def __init__(self,
35 + model_name: str,
36 + dp_size: int,
37 + api_server_count: int,
38 + base_server_args: list,
39 + dp_size_local: int = DP_SIZE_LOCAL,
40 + tp_size: int = TP_SIZE):
41 + self.model_name = model_name
42 + self.dp_size = dp_size
43 + self.dp_size_local = dp_size_local
44 + self.tp_size = tp_size
45 + self.api_server_count = api_server_count
46 + self.base_server_args = base_server_args
47 + self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
48 + self.server_threads: list[threading.Thread] = []
49 + self.num_nodes = dp_size // dp_size_local
50 +
51 + def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
52 + """Start all server instances for hybrid LB mode."""
53 + for node_id in range(self.num_nodes):
54 + # Create server args for this specific node
55 + server_args = self.base_server_args.copy()
56 +
57 + # Calculate start rank for this node
58 + start_rank = node_id * self.dp_size_local
59 +
60 + # Add hybrid LB specific arguments
61 + server_args.extend([
62 + "--data-parallel-size",
63 + str(self.dp_size),
64 + "--data-parallel-size-local",
65 + str(self.dp_size_local),
66 + "--data-parallel-start-rank",
67 + str(start_rank),
68 + "--data-parallel-hybrid-lb", # Enable hybrid LB mode
69 + "--tensor-parallel-size",
70 + str(self.tp_size),
71 + "--port",
72 + str(8000 + node_id), # Different port for each node
73 + "--api-server-count",
74 + str(self.api_server_count),
75 + "--data-parallel-address",
76 + "127.0.0.1",
77 + "--data-parallel-rpc-port",
78 + "13345",
79 + ])
80 +
81 + # Use a thread to start each server to allow parallel initialization
82 + def start_server(node: int, sargs: list[str]):
83 + try:
84 + # Calculate GPU devices for this node
85 + gpus_per_node = self.dp_size_local * self.tp_size
86 + gpu_start = node * gpus_per_node
87 + gpu_end = gpu_start + gpus_per_node
88 +
89 + # Start the server
90 + server = RemoteOpenAIServer(
91 + self.model_name,
92 + sargs,
93 + auto_port=False,
94 + env_dict={
95 + "CUDA_VISIBLE_DEVICES":
96 + ",".join(
97 + str(Platform.device_id_to_physical_device_id(
98 + i)) for i in range(gpu_start, gpu_end))
99 + })
100 + server.__enter__()
101 + print(f"Hybrid LB node {node} started successfully with "
102 + f"{self.dp_size_local} local DP ranks and "
103 + f"{self.api_server_count} API servers")
104 + self.servers.append((server, sargs))
105 + except Exception as e:
106 + print(f"Failed to start hybrid LB node {node}: {e}")
107 + raise
108 +
109 + thread = threading.Thread(target=start_server,
110 + args=(node_id, server_args))
111 + thread.start()
112 +
113 + self.server_threads.append(thread)
114 +
115 + # Wait for all servers to start
116 + for thread in self.server_threads:
117 + thread.join()
118 +
119 + # Give servers additional time to fully initialize and coordinate
120 + time.sleep(3)
121 +
122 + if len(self.servers) != self.num_nodes:
123 + raise Exception("Servers failed to start")
124 +
125 + return self.servers
126 +
127 + def __exit__(self, exc_type, exc_val, exc_tb):
128 + """Stop all server instances."""
129 + while self.servers:
130 + try:
131 + self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
132 + except Exception as e:
133 + print(f"Error stopping server: {e}")
134 +
135 +
136 +@pytest.fixture(scope="module")
137 +def default_server_args():
138 + return [
139 + # use half precision for speed and memory savings in CI environment
140 + "--dtype",
141 + "bfloat16",
142 + "--max-model-len",
143 + "2048",
144 + "--max-num-seqs",
145 + "128",
146 + "--enforce-eager",
147 + ]
148 +
149 +
150 +@pytest.fixture(scope="module", params=[1]) # Only 1 API server for now
151 +def servers(request, default_server_args):
152 + api_server_count = request.param
153 + with HybridLBServerManager(MODEL_NAME, DP_SIZE, api_server_count,
154 + default_server_args, DP_SIZE_LOCAL,
155 + TP_SIZE) as server_list:
156 + yield server_list
157 +
158 +
159 +@pytest_asyncio.fixture
160 +async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
161 + # Create a client for each node (each node has its own API endpoint)
162 + async with AsyncExitStack() as stack:
163 + yield [
164 + await stack.enter_async_context(server.get_async_client())
165 + for server, _ in servers
166 + ]
167 +
168 +
169 +@pytest.mark.asyncio
170 +@pytest.mark.parametrize(
171 + "model_name",
172 + [MODEL_NAME],
173 +)
174 +async def test_hybrid_lb_completion(clients: list[openai.AsyncOpenAI],
175 + servers: list[tuple[RemoteOpenAIServer,
176 + list[str]]],
177 + model_name: str) -> None:
178 +
179 + async def make_request(client: openai.AsyncOpenAI):
180 + completion = await client.completions.create(
181 + model=model_name,
182 + prompt="Hello, my name is",
183 + max_tokens=10,
184 + temperature=1.0)
185 +
186 + assert completion.id is not None
187 + assert completion.choices is not None and len(completion.choices) == 1
188 +
189 + choice = completion.choices[0]
190 + # The exact number of tokens can vary slightly with temperature=1.0,
191 + # so we check for a reasonable minimum length.
192 + assert len(choice.text) >= 1
193 + # Finish reason might not always be 'length' if the model finishes early
194 + # or due to other reasons, especially with high temperature.
195 + # So, we'll accept 'length' or 'stop'.
196 + assert choice.finish_reason in ("length", "stop")
197 +
198 + # Token counts can also vary, so we check they are positive.
199 + assert completion.usage.completion_tokens > 0
200 + assert completion.usage.prompt_tokens > 0
201 + assert completion.usage.total_tokens > 0
202 + return completion
203 +
204 + # Test single request to each node
205 + for i, client in enumerate(clients):
206 + result = await make_request(client)
207 + assert result is not None
208 + print(
209 + f"Hybrid LB node {i} handled single completion request successfully"
210 + )
211 +
212 + await asyncio.sleep(0.5)
213 +
214 + # Send requests to all nodes - each should balance within its local DP ranks
215 + num_requests_per_node = 25 # Total 50 requests across 2 nodes
216 + all_tasks = []
217 +
218 + for i, client in enumerate(clients):
219 + tasks = [make_request(client) for _ in range(num_requests_per_node)]
220 + all_tasks.extend(tasks)
221 +
222 + results = await asyncio.gather(*all_tasks)
223 + assert len(results) == num_requests_per_node * len(clients)
224 + assert all(completion is not None for completion in results)
225 +
226 + await asyncio.sleep(0.5)
227 +
228 + # Second burst of requests
229 + all_tasks = []
230 + for i, client in enumerate(clients):
231 + tasks = [make_request(client) for _ in range(num_requests_per_node)]
232 + all_tasks.extend(tasks)
233 +
234 + results = await asyncio.gather(*all_tasks)
235 + assert len(results) == num_requests_per_node * len(clients)
236 + assert all(completion is not None for completion in results)
237 +
238 + _, server_args = servers[0]
239 + api_server_count = (
240 + server_args.count('--api-server-count')
241 + and server_args[server_args.index('--api-server-count') + 1] or 1)
242 + print(
243 + f"Successfully completed hybrid LB test with {len(clients)} nodes "
244 + f"({DP_SIZE_LOCAL} DP ranks each, API server count: {api_server_count})"
245 + )
246 +
247 + # Check request balancing within each node
248 + for i, (server, _) in enumerate(servers):
249 + print(f"Checking request balancing for node {i}")
250 + check_request_balancing(server, DP_SIZE_LOCAL)
251 +
252 +
253 +@pytest.mark.asyncio
254 +@pytest.mark.parametrize(
255 + "model_name",
256 + [MODEL_NAME],
257 +)
258 +async def test_hybrid_lb_completion_streaming(clients: list[
259 + openai.AsyncOpenAI], servers: list[tuple[RemoteOpenAIServer, list[str]]],
260 + model_name: str) -> None:
261 + prompt = "What is an LLM?"
262 +
263 + async def make_streaming_request(client: openai.AsyncOpenAI):
264 + # Perform a non-streaming request to get the expected full output
265 + single_completion = await client.completions.create(
266 + model=model_name,
267 + prompt=prompt,
268 + max_tokens=5,
269 + temperature=0.0,
270 + )
271 + single_output = single_completion.choices[0].text
272 +
273 + # Perform the streaming request
274 + stream = await client.completions.create(model=model_name,
275 + prompt=prompt,
276 + max_tokens=5,
277 + temperature=0.0,
278 + stream=True)
279 + chunks: list[str] = []
280 + finish_reason_count = 0
281 + last_chunk = None
282 + async for chunk in stream:
283 + chunks.append(chunk.choices[0].text)
284 + if chunk.choices[0].finish_reason is not None:
285 + finish_reason_count += 1
286 + last_chunk = chunk # Keep track of the last chunk
287 +
288 + # finish reason should only return in the last block for OpenAI API
289 + assert finish_reason_count == 1, (
290 + "Finish reason should appear exactly once.")
291 + assert last_chunk is not None, (
292 + "Stream should have yielded at least one chunk.")
293 + assert last_chunk.choices[
294 + 0].finish_reason == "length", "Finish reason should be 'length'."
295 + # Check that the combined text matches the non-streamed version.
296 + assert "".join(
297 + chunks
298 + ) == single_output, "Streamed output should match non-streamed output."
299 + return True # Indicate success for this request
300 +
301 + # Test single request to each node
302 + for i, client in enumerate(clients):
303 + result = await make_streaming_request(client)
304 + assert result is not None
305 + print(
306 + f"Hybrid LB node {i} handled single streaming request successfully"
307 + )
308 +
309 + await asyncio.sleep(0.5)
310 +
311 + # Send streaming requests to all nodes
312 + num_requests_per_node = 25 # Total 50 requests across 2 nodes
313 + all_tasks = []
314 +
315 + for i, client in enumerate(clients):
316 + tasks = [
317 + make_streaming_request(client)
318 + for _ in range(num_requests_per_node)
319 + ]
320 + all_tasks.extend(tasks)
321 +
322 + results = await asyncio.gather(*all_tasks)
323 + assert len(results) == num_requests_per_node * len(clients)
324 + assert all(results), "Not all streaming requests completed successfully."
325 +
326 + await asyncio.sleep(0.5)
327 +
328 + # Second burst of streaming requests
329 + all_tasks = []
330 + for i, client in enumerate(clients):
331 + tasks = [
332 + make_streaming_request(client)
333 + for _ in range(num_requests_per_node)
334 + ]
335 + all_tasks.extend(tasks)
336 +
337 + results = await asyncio.gather(*all_tasks)
338 + assert len(results) == num_requests_per_node * len(clients)
339 + assert all(results), "Not all streaming requests completed successfully."
340 +
341 + _, server_args = servers[0]
342 + api_server_count = (
343 + server_args.count('--api-server-count')
344 + and server_args[server_args.index('--api-server-count') + 1] or 1)
345 + print(f"Successfully completed hybrid LB streaming test with "
346 + f"{len(clients)} nodes ({DP_SIZE_LOCAL} DP ranks each, "
347 + f"API server count: {api_server_count})")
348 +
349 + # Check request balancing within each node
350 + for i, (server, _) in enumerate(servers):
351 + print(f"Checking streaming request balancing for node {i}")
352 + check_request_balancing(server, DP_SIZE_LOCAL)
vllm/config.py
+10 -2
@@ -1908,8 +1908,16 @@ class ParallelConfig:
1908 """Backend to use for data parallel, either "mp" or "ray"."""
1909 data_parallel_external_lb: bool = False
1910 """Whether to use "external" DP LB mode. Applies only to online serving
1911 - and when data_parallel_size > 0. Set implicitly when
1912 - data_parallel_rank is provided explicitly to vllm serve."""
1911 + and when data_parallel_size > 0. This is useful for a "one-pod-per-rank"
1912 + wide-EP setup in Kuberentes. Set implicitly when --data-parallel-rank
1913 + is provided explicitly to vllm serve."""
1914 + data_parallel_hybrid_lb: bool = False
1915 + """Whether to use "hybrid" DP LB mode. Applies only to online serving
1916 + and when data_parallel_size > 0. Enables running an AsyncLLM
1917 + and API server on a "per-node" basis where vLLM load balances
1918 + between local data parallel ranks, but an external LB balances
1919 + between vLLM nodes/replicas. Set explicitly in conjunction with
1920 + --data-parallel-start-rank."""
1921 enable_expert_parallel: bool = False
1922 """Use expert parallelism instead of tensor parallelism for MoE layers."""
1923 enable_eplb: bool = False
vllm/engine/arg_utils.py
+38
@@ -295,9 +295,11 @@ class EngineArgs:
295 tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
296 data_parallel_size: int = ParallelConfig.data_parallel_size
297 data_parallel_rank: Optional[int] = None
298 + data_parallel_start_rank: Optional[int] = None
299 data_parallel_size_local: Optional[int] = None
300 data_parallel_address: Optional[str] = None
301 data_parallel_rpc_port: Optional[int] = None
302 + data_parallel_hybrid_lb: bool = False
303 data_parallel_backend: str = ParallelConfig.data_parallel_backend
304 enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
305 enable_eplb: bool = ParallelConfig.enable_eplb
@@ -604,6 +606,11 @@ class EngineArgs:
606 type=int,
607 help='Data parallel rank of this instance. '
608 'When set, enables external load balancer mode.')
609 + parallel_group.add_argument('--data-parallel-start-rank',
610 + '-dpr',
611 + type=int,
612 + help='Starting data parallel rank '
613 + 'for secondary nodes.')
614 parallel_group.add_argument('--data-parallel-size-local',
615 '-dpl',
616 type=int,
@@ -625,6 +632,9 @@ class EngineArgs:
632 default='mp',
633 help='Backend for data parallel, either '
634 '"mp" or "ray".')
635 + parallel_group.add_argument(
636 + "--data-parallel-hybrid-lb",
637 + **parallel_kwargs["data_parallel_hybrid_lb"])
638 parallel_group.add_argument(
639 "--enable-expert-parallel",
640 **parallel_kwargs["enable_expert_parallel"])
@@ -972,6 +982,7 @@ class EngineArgs:
982 def create_engine_config(
983 self,
984 usage_context: Optional[UsageContext] = None,
985 + headless: bool = False,
986 ) -> VllmConfig:
987 """
988 Create the VllmConfig.
@@ -1060,15 +1071,41 @@ class EngineArgs:
1071 # but we should not do this here.
1072 placement_group = ray.util.get_current_placement_group()
1073
1074 + assert not headless or not self.data_parallel_hybrid_lb, (
1075 + "data_parallel_hybrid_lb is not applicable in "
1076 + "headless mode")
1077 +
1078 data_parallel_external_lb = self.data_parallel_rank is not None
1079 + # Local DP rank = 1, use pure-external LB.
1080 if data_parallel_external_lb:
1081 assert self.data_parallel_size_local in (1, None), (
1082 "data_parallel_size_local must be 1 when data_parallel_rank "
1083 "is set")
1084 data_parallel_size_local = 1
1085 + # Use full external lb if we have local_size of 1.
1086 + self.data_parallel_hybrid_lb = False
1087 elif self.data_parallel_size_local is not None:
1088 data_parallel_size_local = self.data_parallel_size_local
1089 +
1090 + if self.data_parallel_start_rank and not headless:
1091 + # Infer hybrid LB mode.
1092 + self.data_parallel_hybrid_lb = True
1093 +
1094 + if self.data_parallel_hybrid_lb and data_parallel_size_local == 1:
1095 + # Use full external lb if we have local_size of 1.
1096 + data_parallel_external_lb = True
1097 + self.data_parallel_hybrid_lb = False
1098 +
1099 + if data_parallel_size_local == self.data_parallel_size:
1100 + # Disable hybrid LB mode if set for a single node
1101 + self.data_parallel_hybrid_lb = False
1102 +
1103 + self.data_parallel_rank = self.data_parallel_start_rank or 0
1104 else:
1105 + assert not self.data_parallel_hybrid_lb, (
1106 + "data_parallel_size_local must be set to use "
1107 + "data_parallel_hybrid_lb.")
1108 +
1109 # Local DP size defaults to global DP size if not set.
1110 data_parallel_size_local = self.data_parallel_size
1111
@@ -1125,6 +1162,7 @@ class EngineArgs:
1162 data_parallel_master_ip=data_parallel_address,
1163 data_parallel_rpc_port=data_parallel_rpc_port,
1164 data_parallel_backend=self.data_parallel_backend,
1165 + data_parallel_hybrid_lb=self.data_parallel_hybrid_lb,
1166 enable_expert_parallel=self.enable_expert_parallel,
1167 enable_eplb=self.enable_eplb,
1168 num_redundant_experts=self.num_redundant_experts,
vllm/entrypoints/cli/serve.py
+10 -9
@@ -45,11 +45,6 @@ class ServeSubcommand(CLISubcommand):
45 if args.headless or args.api_server_count < 1:
46 run_headless(args)
47 else:
48 - if args.data_parallel_start_rank:
49 - raise ValueError(
50 - "data_parallel_start_rank is only applicable "
51 - "in headless mode. "
52 - "Add --headless flag to enable headless mode.")
48 if args.api_server_count > 1:
49 run_multi_api_server(args)
50 else:
@@ -86,13 +81,14 @@ def run_headless(args: argparse.Namespace):
81 # Create the EngineConfig.
82 engine_args = vllm.AsyncEngineArgs.from_cli_args(args)
83 usage_context = UsageContext.OPENAI_API_SERVER
89 - vllm_config = engine_args.create_engine_config(usage_context=usage_context)
84 + vllm_config = engine_args.create_engine_config(usage_context=usage_context,
85 + headless=True)
86
87 if not envs.VLLM_USE_V1:
88 raise ValueError("Headless mode is only supported for V1")
89
94 - if engine_args.data_parallel_rank is not None:
95 - raise ValueError("data_parallel_rank is not applicable in "
90 + if engine_args.data_parallel_hybrid_lb:
91 + raise ValueError("data_parallel_hybrid_lb is not applicable in "
92 "headless mode")
93
94 parallel_config = vllm_config.parallel_config
@@ -122,7 +118,7 @@ def run_headless(args: argparse.Namespace):
118 engine_manager = CoreEngineProcManager(
119 target_fn=EngineCoreProc.run_engine_core,
120 local_engine_count=local_engine_count,
125 - start_index=args.data_parallel_start_rank,
121 + start_index=vllm_config.parallel_config.data_parallel_rank,
122 local_start_index=0,
123 vllm_config=vllm_config,
124 local_client=False,
@@ -169,6 +165,11 @@ def run_multi_api_server(args: argparse.Namespace):
165 " api_server_count > 1")
166 model_config.disable_mm_preprocessor_cache = True
167
168 + if vllm_config.parallel_config.data_parallel_hybrid_lb:
169 + raise NotImplementedError(
170 + "Hybrid load balancing with --api-server-count > 0"
171 + "is not yet supported.")
172 +
173 executor_class = Executor.get_class(vllm_config)
174 log_stats = not engine_args.disable_log_stats
175
vllm/entrypoints/openai/cli_args.py
-7
@@ -222,13 +222,6 @@ def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
222 default=False,
223 help="Run in headless mode. See multi-node data parallel "
224 "documentation for more details.")
225 - parser.add_argument(
226 - "--data-parallel-start-rank",
227 - "-dpr",
228 - type=int,
229 - default=0,
230 - help="Starting data parallel rank for secondary nodes. "
231 - "Requires --headless.")
225 parser.add_argument("--api-server-count",
226 "-asc",
227 type=int,
vllm/v1/engine/async_llm.py
+1 -1
@@ -127,7 +127,7 @@ class AsyncLLM(EngineClient):
127 if self.log_stats:
128 self.logger_manager = StatLoggerManager(
129 vllm_config=vllm_config,
130 - engine_idxs=self.engine_core.engine_ranks,
130 + engine_idxs=self.engine_core.engine_ranks_managed,
131 custom_stat_loggers=stat_loggers,
132 )
133 self.logger_manager.log_engine_initialized()
vllm/v1/engine/coordinator.py
+3 -2
@@ -61,11 +61,12 @@ class DPCoordinator:
61
62 host = parallel_config.data_parallel_master_ip
63 external_lb = parallel_config.data_parallel_external_lb
64 + hybrid_lb = parallel_config.data_parallel_hybrid_lb
65
66 # Assume coordinator is colocated with front-end procs when not in
66 - # external DP LB mode.
67 + # either external or hybrid DP LB mode.
68 front_publish_address = get_engine_client_zmq_addr(
68 - local_only=not external_lb, host=host)
69 + local_only=not external_lb and not hybrid_lb, host=host)
70
71 local_only_eng = dp_size == parallel_config.data_parallel_size_local
72 back_publish_address = get_engine_client_zmq_addr(local_only_eng, host)
vllm/v1/engine/core.py
+14 -5
@@ -467,13 +467,14 @@ class EngineCoreProc(EngineCore):
467 For DP>1 with internal loadbalancing this is with the shared front-end
468 process which may reside on a different node.
469
470 - For DP>1 with external loadbalancing, two handshakes are performed:
470 + For DP>1 with external or hybrid loadbalancing, two handshakes are
471 + performed:
472 - With the rank 0 front-end process which retrieves the
473 DP Coordinator ZMQ addresses and DP process group address.
474 - With the colocated front-end process which retrieves the
475 client input/output socket addresses.
475 - with the exception of the rank 0 engine itself which doesn't require
476 - the second handshake.
476 + with the exception of the rank 0 and colocated engines themselves which
477 + don't require the second handshake.
478
479 Here, "front-end" process can mean the process containing the engine
480 core client (which is the API server process in the case the API
@@ -482,15 +483,18 @@ class EngineCoreProc(EngineCore):
483 """
484 input_ctx = zmq.Context()
485 is_local = local_client and client_handshake_address is None
486 + headless = not local_client
487 handshake = self._perform_handshake(input_ctx, handshake_address,
486 - identity, is_local, vllm_config,
488 + identity, is_local, headless,
489 + vllm_config,
490 vllm_config.parallel_config)
491 if client_handshake_address is None:
492 with handshake as addresses:
493 yield addresses
494 else:
495 + assert local_client
496 local_handshake = self._perform_handshake(
493 - input_ctx, client_handshake_address, identity, local_client,
497 + input_ctx, client_handshake_address, identity, True, False,
498 vllm_config)
499 with handshake as addresses, local_handshake as client_addresses:
500 addresses.inputs = client_addresses.inputs
@@ -507,6 +511,7 @@ class EngineCoreProc(EngineCore):
511 handshake_address: str,
512 identity: bytes,
513 local_client: bool,
514 + headless: bool,
515 vllm_config: VllmConfig,
516 parallel_config_to_update: Optional[ParallelConfig] = None,
517 ) -> Generator[EngineZmqAddresses, None, None]:
@@ -518,6 +523,7 @@ class EngineCoreProc(EngineCore):
523 bind=False) as handshake_socket:
524 # Register engine with front-end.
525 addresses = self.startup_handshake(handshake_socket, local_client,
526 + headless,
527 parallel_config_to_update)
528 yield addresses
529
@@ -531,6 +537,7 @@ class EngineCoreProc(EngineCore):
537 msgspec.msgpack.encode({
538 "status": "READY",
539 "local": local_client,
540 + "headless": headless,
541 "num_gpu_blocks": num_gpu_blocks,
542 "dp_stats_address": dp_stats_address,
543 }))
@@ -539,6 +546,7 @@ class EngineCoreProc(EngineCore):
546 def startup_handshake(
547 handshake_socket: zmq.Socket,
548 local_client: bool,
549 + headless: bool,
550 parallel_config: Optional[ParallelConfig] = None,
551 ) -> EngineZmqAddresses:
552
@@ -547,6 +555,7 @@ class EngineCoreProc(EngineCore):
555 msgspec.msgpack.encode({
556 "status": "HELLO",
557 "local": local_client,
558 + "headless": headless,
559 }))
560
561 # Receive initialization message.
vllm/v1/engine/core_client.py
+19 -8
@@ -429,18 +429,23 @@ class MPClient(EngineCoreClient):
429 parallel_config = vllm_config.parallel_config
430 dp_size = parallel_config.data_parallel_size
431 dp_rank = parallel_config.data_parallel_rank
432 - external_dp_lb = parallel_config.data_parallel_external_lb
433 -
432 + dp_local_size = parallel_config.data_parallel_size_local
433 offline_mode = parallel_config.data_parallel_rank_local is not None
435 - self.engine_ranks = ([dp_rank] if
436 - (offline_mode or external_dp_lb) else list(
437 - range(dp_size)))
434 + # Client manages local+remote EngineCores in pure internal LB case.
435 + # Client manages local EngineCores in hybrid and external LB case.
436 + local_engines_only = (parallel_config.data_parallel_hybrid_lb
437 + or parallel_config.data_parallel_external_lb)
438 +
439 + num_ranks = dp_local_size if local_engines_only else dp_size
440 + self.engine_ranks_managed = [dp_rank] if offline_mode else list(
441 + range(dp_rank, dp_rank + num_ranks))
442 assert parallel_config.data_parallel_size_local <= len(
439 - self.engine_ranks)
443 + self.engine_ranks_managed)
444
445 # ZMQ identity of each engine that this client will talk to.
446 self.core_engines: list[EngineIdentity] = [
443 - index.to_bytes(2, "little") for index in self.engine_ranks
447 + rank.to_bytes(2, "little")
448 + for rank in self.engine_ranks_managed
449 ]
450
451 # Wait for ready messages from each engine on the input socket.
@@ -895,6 +900,12 @@ class DPAsyncMPClient(AsyncMPClient):
900 return
901
902 assert self.stats_update_address is not None
903 + assert len(self.engine_ranks_managed) > 0
904 + # NOTE: running and waiting counts are all global from
905 + # the Coordinator include all global EngineCores. This
906 + # slice includes just the cores managed by this client.
907 + count_slice = slice(self.engine_ranks_managed[0],
908 + self.engine_ranks_managed[-1] + 1)
909
910 async def run_engine_stats_update_task():
911 with make_zmq_socket(self.ctx, self.stats_update_address,
@@ -959,7 +970,7 @@ class DPAsyncMPClient(AsyncMPClient):
970 counts, wave, running = msgspec.msgpack.decode(buf)
971 self.current_wave = wave
972 self.engines_running = running
962 - self.lb_engines = counts
973 + self.lb_engines = counts[count_slice]
974
975 resources.stats_update_task = asyncio.create_task(
976 run_engine_stats_update_task())
vllm/v1/engine/utils.py
+35 -9
@@ -544,7 +544,8 @@ def launch_core_engines(
544 local_start_index = parallel_config.data_parallel_rank_local
545 dp_rank = parallel_config.data_parallel_rank
546 host = parallel_config.data_parallel_master_ip
547 - external_dp_lb = parallel_config.data_parallel_external_lb
547 + local_engines_only = (parallel_config.data_parallel_hybrid_lb
548 + or parallel_config.data_parallel_external_lb)
549
550 # In offline mode there is an LLM instance per DP rank and
551 # one core engine per LLM, see
@@ -553,8 +554,8 @@ def launch_core_engines(
554
555 # client_local_only = True for cases where this front-end
556 # sends requests only to colocated engines.
556 - client_local_only = offline_mode or external_dp_lb or (local_engine_count
557 - == dp_size)
557 + client_local_only = (offline_mode or local_engines_only
558 + or (local_engine_count == dp_size))
559
560 # Set up input and output addresses.
561 addresses = EngineZmqAddresses(
@@ -598,14 +599,27 @@ def launch_core_engines(
599 yield engine_actor_manager, coordinator, addresses
600 return
601
601 - if offline_mode or (external_dp_lb and dp_rank > 0):
602 + if offline_mode:
603 assert local_engine_count == 1
604 engines_to_handshake = [CoreEngine(index=dp_rank, local=True)]
604 - else:
605 + elif dp_rank == 0:
606 + # Rank 0 holds Coordinator, so it handshakes with all Cores
607 + # in both external dplb and internal dplb mode.
608 + # Note this also covers the case where we have zero local engines
609 + # and rank 0 is headless.
610 engines_to_handshake = [
611 CoreEngine(index=i, local=(i < local_engine_count))
612 for i in range(dp_size)
613 ]
614 + else:
615 + # Rank > 0 handshakes with just the local cores it is managing.
616 + assert local_engines_only, (
617 + "Attempting to launch core_engines from dp_rank > 0, but "
618 + "found internal DPLB, which is incompatible.")
619 + engines_to_handshake = [
620 + CoreEngine(index=i, local=True)
621 + for i in range(dp_rank, dp_rank + local_engine_count)
622 + ]
623
624 # Whether the started engines will handshake only with co-located
625 # front-end processes. In external_dp_lb mode, ranks > 0 handshake with
@@ -616,7 +630,7 @@ def launch_core_engines(
630 handshake_address = get_engine_client_zmq_addr(
631 handshake_local_only, host, parallel_config.data_parallel_rpc_port)
632
619 - if external_dp_lb and dp_rank > 0:
633 + if local_engines_only and dp_rank > 0:
634 assert not handshake_local_only
635 local_handshake_address = get_open_zmq_ipc_path()
636 client_handshake_address = local_handshake_address
@@ -631,8 +645,6 @@ def launch_core_engines(
645
646 # Start local engines.
647 if local_engine_count:
634 - # In server mode, start_index and local_start_index will
635 - # both be 0.
648 local_engine_manager = CoreEngineProcManager(
649 EngineCoreProc.run_engine_core,
650 vllm_config=vllm_config,
@@ -678,6 +690,9 @@ def wait_for_engine_startup(
690 poller = zmq.Poller()
691 poller.register(handshake_socket, zmq.POLLIN)
692
693 + remote_should_be_headless = not parallel_config.data_parallel_hybrid_lb \
694 + and not parallel_config.data_parallel_external_lb
695 +
696 if proc_manager is not None:
697 for sentinel in proc_manager.sentinels():
698 poller.register(sentinel, zmq.POLLIN)
@@ -713,13 +728,24 @@ def wait_for_engine_startup(
728 raise RuntimeError(f"Message from engine with unexpected data "
729 f"parallel rank: {eng_index}")
730 msg = msgspec.msgpack.decode(ready_msg_bytes)
716 - status, local = msg["status"], msg["local"]
731 + status, local, headless = msg["status"], msg["local"], msg["headless"]
732 if local != engine.local:
733 raise RuntimeError(f"{status} message from "
734 f"{'local' if local else 'remote'} "
735 f"engine {eng_index}, expected it to be "
736 f"{'local' if engine.local else 'remote'}")
737
738 + # Remote engines must be headless iff we aren't in hybrid dp lb mode.
739 + if not local and headless != remote_should_be_headless:
740 + if headless:
741 + raise RuntimeError(f"Remote engine {eng_index} must not use "
742 + f"--headless in external or hybrid dp lb "
743 + f"mode")
744 + else:
745 + raise RuntimeError(f"Remote engine {eng_index} must use "
746 + f"--headless unless in external or hybrid "
747 + f"dp lb mode")
748 +
749 if status == "HELLO" and engine.state == CoreEngineState.NEW:
750
751 # Send init message with DP config info.