@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 15dac210f0

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 15, 17:15 UTC

[V1] AsyncLLM data parallel (#13923)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Mar 27, 2025 at 16:14 UTC 15dac210f0e6b907f191911917238273042552ed

18 files changed +726 -160

.buildkite/test-pipeline.yaml

@@ -135,12 +135,14 @@ steps:

135 - examples/offline_inference/rlhf.py

136 - examples/offline_inference/rlhf_colocate.py

137 - tests/examples/offline_inference/data_parallel.py

138 + - tests/v1/test_async_llm_dp.py

139 commands:

140 # test with tp=2 and external_dp=2

141 - VLLM_USE_V1=0 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py

142 - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py

143 # test with internal dp

144 - python3 ../examples/offline_inference/data_parallel.py

145 + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py

146 - pytest -v -s distributed/test_utils.py

147 - pytest -v -s compile/test_basic_correctness.py

148 - pytest -v -s distributed/test_pynccl.py

@@ -514,7 +516,10 @@ steps:

516 - vllm/worker/worker.py

517 - vllm/worker/model_runner.py

518 - entrypoints/llm/test_collective_rpc.py

519 + - tests/v1/test_async_llm_dp.py

520 + - vllm/v1/engine/

521 commands:

522 + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py

523 - VLLM_ENABLE_V1_MULTIPROCESSING=0 pytest -v -s entrypoints/llm/test_collective_rpc.py

524 - pytest -v -s ./compile/test_basic_correctness.py

525 - pytest -v -s ./compile/test_wrapper.py

examples/offline_inference/data_parallel.py

+15 -7

@@ -28,6 +28,7 @@ Multi-node:

28 --master-port=13345

29 """

30 import os

31 +from time import sleep

33 from vllm import LLM, SamplingParams

34 from vllm.utils import get_open_port

@@ -36,14 +37,13 @@ from vllm.utils import get_open_port

37 def main(model, dp_size, local_dp_rank, global_dp_rank, dp_master_ip,

38 dp_master_port, GPUs_per_dp_rank):

39 os.environ["VLLM_DP_RANK"] = str(global_dp_rank)

40 + os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)

41 os.environ["VLLM_DP_SIZE"] = str(dp_size)

42 os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip

43 os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)

42 - # set devices for each dp_rank

43 - os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(

44 - str(i)

45 - for i in range(local_dp_rank * GPUs_per_dp_rank, (local_dp_rank + 1) *

46 - GPUs_per_dp_rank))

44 +

45 + # CUDA_VISIBLE_DEVICES for each DP rank is set automatically inside the

46 + # engine processes.

48 # Sample prompts.

49 prompts = [

@@ -90,6 +90,9 @@ def main(model, dp_size, local_dp_rank, global_dp_rank, dp_master_ip,

90 print(f"DP rank {global_dp_rank}, Prompt: {prompt!r}, "

91 f"Generated text: {generated_text!r}")

93 + # Give engines time to pause their processing loops before exiting.

94 + sleep(1)

95 +

97 if __name__ == "__main__":

98 import argparse

@@ -152,8 +155,13 @@ if __name__ == "__main__":

155 procs.append(proc)

156 exit_code = 0

157 for proc in procs:

155 - proc.join()

156 - if proc.exitcode:

158 + proc.join(timeout=300)

159 + if proc.exitcode is None:

160 + print(f"Killing process {proc.pid} that "

161 + f"didn't stop within 5 minutes.")

162 + proc.kill()

163 + exit_code = 1

164 + elif proc.exitcode:

165 exit_code = proc.exitcode

166

167 exit(exit_code)

tests/v1/engine/test_engine_core_client.py

+4 -4

@@ -167,11 +167,11 @@ def test_engine_core_client(monkeypatch: pytest.MonkeyPatch,

167

168 core_client: SyncMPClient = client

169

170 - result = core_client._call_utility("echo", "testarg")

170 + result = core_client.call_utility("echo", "testarg")

171 assert result == "testarg"

172

173 with pytest.raises(Exception) as e_info:

174 - core_client._call_utility("echo", None, "help!")

174 + core_client.call_utility("echo", None, "help!")

175

176 assert str(e_info.value) == "Call to echo method failed: help!"

177

@@ -238,10 +238,10 @@ async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):

238

239 core_client: AsyncMPClient = client

240

241 - result = await core_client._call_utility_async("echo", "testarg")

241 + result = await core_client.call_utility_async("echo", "testarg")

242 assert result == "testarg"

243

244 with pytest.raises(Exception) as e_info:

245 - await core_client._call_utility_async("echo", None, "help!")

245 + await core_client.call_utility_async("echo", None, "help!")

246

247 assert str(e_info.value) == "Call to echo method failed: help!"

tests/v1/test_async_llm_dp.py new

+109

@@ -0,0 +1,109 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +

3 +import asyncio

4 +import os

5 +from contextlib import ExitStack

6 +from typing import Optional

7 +

8 +import pytest

9 +

10 +from vllm import SamplingParams

11 +from vllm.engine.arg_utils import AsyncEngineArgs

12 +from vllm.inputs import PromptType

13 +from vllm.platforms import current_platform

14 +from vllm.sampling_params import RequestOutputKind

15 +from vllm.v1.engine.async_llm import AsyncLLM

16 +from vllm.v1.engine.core_client import DPAsyncMPClient

17 +

18 +engine_args = AsyncEngineArgs(

19 + model="ibm-research/PowerMoE-3b",

20 + enforce_eager=True,

21 + disable_log_requests=True,

22 + tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),

23 + data_parallel_size=int(os.getenv("DP_SIZE", 2)),

24 +)

25 +

26 +if not current_platform.supports_v1(engine_args.create_model_config()):

27 + pytest.skip(reason="Requires V1-supporting platform.",

28 + allow_module_level=True)

29 +

30 +

31 +async def generate(engine: AsyncLLM,

32 + request_id: str,

33 + prompt: PromptType,

34 + output_kind: RequestOutputKind,

35 + max_tokens: int,

36 + prompt_logprobs: Optional[int] = None) -> tuple[int, str]:

37 + # Ensure generate doesn't complete too fast for cancellation test.

38 + await asyncio.sleep(0.2)

39 +

40 + count = 0

41 + sampling_params = SamplingParams(max_tokens=max_tokens,

42 + ignore_eos=True,

43 + output_kind=output_kind,

44 + temperature=0,

45 + prompt_logprobs=prompt_logprobs)

46 + async for out in engine.generate(request_id=request_id,

47 + prompt=prompt,

48 + sampling_params=sampling_params):

49 +

50 + num_tokens = len(out.outputs[0].token_ids)

51 + if output_kind == RequestOutputKind.DELTA:

52 + count += num_tokens

53 + else:

54 + count = num_tokens

55 +

56 + await asyncio.sleep(0.)

57 +

58 + return count, request_id

59 +

60 +

61 +@pytest.mark.parametrize(

62 + "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])

63 +@pytest.mark.asyncio

64 +async def test_load(output_kind: RequestOutputKind):

65 +

66 + with ExitStack() as after:

67 +

68 + prompt = "This is a test of data parallel"

69 +

70 + engine = AsyncLLM.from_engine_args(engine_args)

71 + after.callback(engine.shutdown)

72 +

73 + NUM_REQUESTS = 100

74 + NUM_EXPECTED_TOKENS = 10

75 +

76 + request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]

77 +

78 + # Create concurrent requests.

79 + tasks = []

80 + for request_id in request_ids:

81 + tasks.append(

82 + asyncio.create_task(

83 + generate(engine, request_id, prompt, output_kind,

84 + NUM_EXPECTED_TOKENS)))

85 +

86 + # Confirm that we got all the EXPECTED tokens from the requests.

87 + done, pending = await asyncio.wait(tasks,

88 + return_when=asyncio.FIRST_EXCEPTION)

89 + for task in pending:

90 + task.cancel()

91 + for task in done:

92 + num_generated_tokens, request_id = await task

93 + assert num_generated_tokens == NUM_EXPECTED_TOKENS, (

94 + f"{request_id} generated {num_generated_tokens} but "

95 + f"expected {NUM_EXPECTED_TOKENS}")

96 +

97 + assert not engine.output_processor.has_unfinished_requests()

98 +

99 + # testing internals here which may break

100 + core_client: DPAsyncMPClient = engine.engine_core

101 + # the engines only synchronize stopping every N steps so

102 + # allow a small amount of time here.

103 + for _ in range(10):

104 + if core_client.num_engines_running == 0:

105 + break

106 + await asyncio.sleep(0.5)

107 +

108 + assert core_client.num_engines_running == 0

109 + assert not core_client.reqs_in_flight

vllm/config.py

+16 -5

@@ -40,7 +40,8 @@ from vllm.transformers_utils.config import (

40 from vllm.transformers_utils.s3_utils import S3Model

41 from vllm.transformers_utils.utils import is_s3, maybe_model_redirect

42 from vllm.utils import (GiB_bytes, LayerBlockType, cuda_device_count_stateless,

43 - get_cpu_memory, random_uuid, resolve_obj_by_qualname)

43 + get_cpu_memory, get_open_port, random_uuid,

44 + resolve_obj_by_qualname)

46 if TYPE_CHECKING:

47 from ray.util.placement_group import PlacementGroup

@@ -1389,6 +1390,8 @@ class ParallelConfig:

1390 tensor_parallel_size: int = 1 # Number of tensor parallel groups.

1391 data_parallel_size: int = 1 # Number of data parallel groups.

1392 data_parallel_rank: int = 0 # Rank of the data parallel group.

1393 + # Local rank of the data parallel group, defaults to global rank.

1394 + data_parallel_rank_local: Optional[int] = None

1395 # IP of the data parallel master.

1396 data_parallel_master_ip: str = "127.0.0.1"

1397 data_parallel_master_port: int = 29500 # Port of the data parallel master.

@@ -1493,10 +1496,18 @@ class ParallelConfig:

1496 self.world_size = self.pipeline_parallel_size * \

1497 self.tensor_parallel_size

1498

1496 - self.data_parallel_size = envs.VLLM_DP_SIZE

1497 - self.data_parallel_rank = envs.VLLM_DP_RANK

1498 - self.data_parallel_master_ip = envs.VLLM_DP_MASTER_IP

1499 - self.data_parallel_master_port = envs.VLLM_DP_MASTER_PORT

1499 + if self.data_parallel_size > 1:

1500 + # Data parallel was specified in the engine args.

1501 + self.data_parallel_master_port = get_open_port()

1502 + # TODO multi-node

1503 + else:

1504 + # Otherwise fall back to env vars (e.g. for offline SPMD case).

1505 + self.data_parallel_size = envs.VLLM_DP_SIZE

1506 + self.data_parallel_rank = envs.VLLM_DP_RANK

1507 + self.data_parallel_rank_local = envs.VLLM_DP_RANK_LOCAL

1508 + self.data_parallel_master_ip = envs.VLLM_DP_MASTER_IP

1509 + self.data_parallel_master_port = envs.VLLM_DP_MASTER_PORT

1510 +

1511 self.world_size_across_dp = self.world_size * self.data_parallel_size

1512

1513 if self.distributed_executor_backend == "external_launcher":

vllm/distributed/utils.py

+12

@@ -15,6 +15,8 @@ import torch

15 from torch.distributed import ProcessGroup, TCPStore

16 from torch.distributed.distributed_c10d import (Backend, PrefixStore,

17 _get_default_timeout,

18 + _shutdown_backend,

19 + _unregister_process_group,

20 is_nccl_available)

21 from torch.distributed.rendezvous import rendezvous

@@ -333,3 +335,13 @@ def stateless_init_torch_distributed_process_group(

335 pg._register_backend(device, backend_type, backend_class)

336

337 return pg

338 +

339 +

340 +def stateless_destroy_torch_distributed_process_group(

341 + pg: ProcessGroup) -> None:

342 + """

343 + Destroy ProcessGroup returned by

344 + stateless_init_torch_distributed_process_group().

345 + """

346 + _shutdown_backend(pg)

347 + _unregister_process_group(pg.group_name)

vllm/engine/arg_utils.py

+10

@@ -114,6 +114,7 @@ class EngineArgs:

114 # number of P/D disaggregation (or other disaggregation) workers

115 pipeline_parallel_size: int = 1

116 tensor_parallel_size: int = 1

117 + data_parallel_size: int = 1

118 enable_expert_parallel: bool = False

119 max_parallel_loading_workers: Optional[int] = None

120 block_size: Optional[int] = None

@@ -442,6 +443,14 @@ class EngineArgs:

443 type=int,

444 default=EngineArgs.tensor_parallel_size,

445 help='Number of tensor parallel replicas.')

446 + parser.add_argument('--data-parallel-size',

447 + '-dp',

448 + type=int,

449 + default=EngineArgs.data_parallel_size,

450 + help='Number of data parallel replicas. '

451 + 'MoE layers will be sharded according to the '

452 + 'product of the tensor-parallel-size and '

453 + 'data-parallel-size.')

454 parser.add_argument(

455 '--enable-expert-parallel',

456 action='store_true',

@@ -1359,6 +1368,7 @@ class EngineArgs:

1368 parallel_config = ParallelConfig(

1369 pipeline_parallel_size=self.pipeline_parallel_size,

1370 tensor_parallel_size=self.tensor_parallel_size,

1371 + data_parallel_size=self.data_parallel_size,

1372 enable_expert_parallel=self.enable_expert_parallel,

1373 max_parallel_loading_workers=self.max_parallel_loading_workers,

1374 disable_custom_all_reduce=self.disable_custom_all_reduce,

vllm/envs.py

@@ -2,6 +2,7 @@

3 import hashlib

4 import os

5 +import sys

6 import tempfile

7 from typing import TYPE_CHECKING, Any, Callable, Optional

@@ -95,6 +96,7 @@ if TYPE_CHECKING:

96 VLLM_CUDART_SO_PATH: Optional[str] = None

97 VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH: bool = True

98 VLLM_DP_RANK: int = 0

99 + VLLM_DP_RANK_LOCAL: int = -1

100 VLLM_DP_SIZE: int = 1

101 VLLM_DP_MASTER_IP: str = ""

102 VLLM_DP_MASTER_PORT: int = 0

@@ -625,6 +627,12 @@ environment_variables: dict[str, Callable[[], Any]] = {

627 "VLLM_DP_RANK":

628 lambda: int(os.getenv("VLLM_DP_RANK", "0")),

629

630 + # Rank of the process in the data parallel setting.

631 + # Defaults to VLLM_DP_RANK when not set.

632 + "VLLM_DP_RANK_LOCAL":

633 + lambda: int(

634 + os.getenv("VLLM_DP_RANK_LOCAL", sys.modules[__name__].VLLM_DP_RANK)),

635 +

636 # World size of the data parallel setting

637 "VLLM_DP_SIZE":

638 lambda: int(os.getenv("VLLM_DP_SIZE", "1")),

vllm/utils.py

+9 -5

@@ -578,7 +578,7 @@ def get_open_port() -> int:

578 dp_port = envs.VLLM_DP_MASTER_PORT

579 while True:

580 port = _get_open_port()

581 - if port >= dp_port and port < dp_port + 10:

581 + if dp_port <= port < dp_port + 10:

582 continue

583 return port

584 return _get_open_port()

@@ -2176,11 +2176,11 @@ def make_zmq_socket(

2176 if socket_type == zmq.constants.PULL:

2177 socket.setsockopt(zmq.constants.RCVHWM, 0)

2178 socket.setsockopt(zmq.constants.RCVBUF, buf_size)

2179 - socket.connect(path)

2179 + socket.bind(path)

2180 elif socket_type == zmq.constants.PUSH:

2181 socket.setsockopt(zmq.constants.SNDHWM, 0)

2182 socket.setsockopt(zmq.constants.SNDBUF, buf_size)

2183 - socket.bind(path)

2183 + socket.connect(path)

2184 else:

2185 raise ValueError(f"Unknown Socket Type: {socket_type}")

2186

@@ -2188,7 +2188,11 @@ def make_zmq_socket(

2188

2189

2190 @contextlib.contextmanager

2191 -def zmq_socket_ctx(path: str, socket_type: Any) -> Iterator[zmq.Socket]:

2191 +def zmq_socket_ctx(

2192 + path: str,

2193 + socket_type: Any,

2194 + linger: int = 0,

2195 +) -> Iterator[zmq.Socket]:

2196 """Context manager for a ZMQ socket"""

2197

2198 ctx = zmq.Context() # type: ignore[attr-defined]

@@ -2199,7 +2203,7 @@ def zmq_socket_ctx(path: str, socket_type: Any) -> Iterator[zmq.Socket]:

2203 logger.debug("Got Keyboard Interrupt.")

2204

2205 finally:

2202 - ctx.destroy(linger=0)

2206 + ctx.destroy(linger=linger)

2207

2208

2209 def is_in_ray_actor():

vllm/v1/core/sched/scheduler.py

+15 -2

@@ -37,9 +37,10 @@ class Scheduler(SchedulerInterface):

37 cache_config: CacheConfig,

38 lora_config: Optional[LoRAConfig],

39 speculative_config: Optional[SpeculativeConfig],

40 - log_stats: bool,

40 structured_output_manager: StructuredOutputManager,

41 mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,

42 + include_finished_set: bool = False,

43 + log_stats: bool = False,

44 ) -> None:

45 self.scheduler_config = scheduler_config

46 self.cache_config = cache_config

@@ -48,6 +49,12 @@ class Scheduler(SchedulerInterface):

49 self.log_stats = log_stats

50 self.structured_output_manager = structured_output_manager

52 + # include_finished_set controls whether a separate set of finished

53 + # request ids should be included in the EngineCoreOutputs returned

54 + # by update_from_outputs(). This is currently used in the multi-engine

55 + # case to track request lifetimes efficiently.

56 + self.include_finished_set = include_finished_set

57 +

58 # Scheduling constraints.

59 self.max_num_running_reqs = self.scheduler_config.max_num_seqs

60 self.max_num_scheduled_tokens = \

@@ -663,10 +670,16 @@ class Scheduler(SchedulerInterface):

670 new_running.append(request)

671

672 self.running = new_running

666 - return EngineCoreOutputs(

673 + engine_core_outputs = EngineCoreOutputs(

674 outputs=outputs,

675 scheduler_stats=self.make_stats(),

676 )

677 + if self.include_finished_set:

678 + #TODO currently sending duplicates here, improve this

679 + engine_core_outputs.finished_requests = (

680 + scheduler_output.finished_req_ids | self.finished_req_ids)

681 +

682 + return engine_core_outputs

683

684 def add_request(self, request: Request) -> None:

685 self.waiting.append(request)

vllm/v1/engine/__init__.py

+8 -1

@@ -128,12 +128,18 @@ class EngineCoreOutputs(

128 #NOTE(Nick): We could consider ways to make this more compact,

129 # e.g. columnwise layout

130

131 + engine_index: int = 0

132 +

133 # [num_reqs]

134 outputs: list[EngineCoreOutput] = []

135 scheduler_stats: Optional[SchedulerStats] = None

136 timestamp: float = 0.0

137

138 utility_output: Optional[UtilityOutput] = None

139 + finished_requests: Optional[set[str]] = None

140 +

141 + # In DP case, used to signal that the engine is paused.

142 + engine_paused: bool = False

143

144 def __post_init__(self):

145 if self.timestamp == 0.0:

@@ -147,4 +153,5 @@ class EngineCoreRequestType(enum.Enum):

153 """

154 ADD = b'\x00'

155 ABORT = b'\x01'

150 - UTILITY = b'\x02'

156 + START_DP = b'\x02'

157 + UTILITY = b'\x03'

vllm/v1/engine/async_llm.py

+16 -7

@@ -66,11 +66,17 @@ class AsyncLLM(EngineClient):

67 self.log_requests = log_requests

68 self.log_stats = log_stats

69 - self.stat_loggers: list[StatLoggerBase] = []

69 +

70 + # Set up stat loggers; independent set for each DP rank.

71 + self.stat_loggers: list[list[StatLoggerBase]] = []

72 if self.log_stats:

71 - if logger.isEnabledFor(logging.INFO):

72 - self.stat_loggers.append(LoggingStatLogger())

73 - self.stat_loggers.append(PrometheusStatLogger(vllm_config))

73 + for i in range(vllm_config.parallel_config.data_parallel_size):

74 + loggers: list[StatLoggerBase] = []

75 + if logger.isEnabledFor(logging.INFO):

76 + loggers.append(LoggingStatLogger(engine_index=i))

77 + loggers.append(

78 + PrometheusStatLogger(vllm_config, engine_index=i))

79 + self.stat_loggers.append(loggers)

81 # Tokenizer (+ ensure liveness if running in another process).

82 self.tokenizer = init_tokenizer_from_configs(

@@ -329,6 +335,7 @@ class AsyncLLM(EngineClient):

335 # TODO(rob): make into a coroutine and launch it in

336 # background thread once Prometheus overhead is non-trivial.

337 self._record_stats(

338 + engine_index=outputs.engine_index,

339 scheduler_stats=outputs.scheduler_stats,

340 iteration_stats=iteration_stats,

341 )

@@ -350,12 +357,13 @@ class AsyncLLM(EngineClient):

357 self,

358 scheduler_stats: Optional[SchedulerStats],

359 iteration_stats: Optional[IterationStats],

360 + engine_index: int = 0,

361 ):

362 if not self.log_stats:

363 return

364

365 assert scheduler_stats is not None

358 - for stat_logger in self.stat_loggers:

366 + for stat_logger in self.stat_loggers[engine_index]:

367 stat_logger.record(scheduler_stats=scheduler_stats,

368 iteration_stats=iteration_stats)

369

@@ -393,8 +401,9 @@ class AsyncLLM(EngineClient):

401 scheduler_outputs=None,

402 model_output=None,

403 ) -> None:

396 - for stat_logger in self.stat_loggers:

397 - stat_logger.log()

404 + for loggers in self.stat_loggers:

405 + for stat_logger in loggers:

406 + stat_logger.log()

407

408 async def check_health(self) -> None:

409 logger.debug("Called check_health.")

vllm/v1/engine/core.py

+183 -31

@@ -1,12 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -

2 +import os

3 import queue

4 import signal

5 +import sys

6 import threading

7 import time

8 from concurrent.futures import Future

9 from inspect import isclass, signature

9 -from multiprocessing.connection import Connection

10 +from logging import DEBUG

11 from typing import Any, Optional

13 import msgspec

@@ -14,7 +15,9 @@ import psutil

15 import zmq

16 import zmq.asyncio

17 -from vllm.config import VllmConfig

18 +from vllm.config import ParallelConfig, VllmConfig

19 +from vllm.distributed import stateless_destroy_torch_distributed_process_group

20 +from vllm.executor.multiproc_worker_utils import _add_prefix

21 from vllm.logger import init_logger

22 from vllm.lora.request import LoRARequest

23 from vllm.transformers_utils.config import (

@@ -91,6 +94,8 @@ class EngineCore:

94 cache_config=vllm_config.cache_config,

95 lora_config=vllm_config.lora_config,

96 speculative_config=vllm_config.speculative_config,

97 + include_finished_set=vllm_config.parallel_config.data_parallel_size

98 + > 1,

99 log_stats=self.log_stats,

100 structured_output_manager=self.structured_output_manager,

101 )

@@ -283,10 +288,10 @@ class EngineCoreProc(EngineCore):

288 self,

289 input_path: str,

290 output_path: str,

286 - ready_pipe: Connection,

291 vllm_config: VllmConfig,

292 executor_class: type[Executor],

293 log_stats: bool,

294 + engine_index: int = 0,

295 ):

296 super().__init__(vllm_config, executor_class, log_stats)

297

@@ -302,14 +307,20 @@ class EngineCoreProc(EngineCore):

307 args=(input_path, ),

308 daemon=True).start()

309 threading.Thread(target=self.process_output_socket,

305 - args=(output_path, ),

310 + args=(output_path, engine_index),

311 daemon=True).start()

312

308 - # Send Readiness signal to EngineClient.

309 - ready_pipe.send({"status": "READY"})

313 + self.global_unfinished_reqs = False

314 +

315 + self.step_fn = (self.step if self.batch_queue is None else

316 + self.step_with_batch_queue)

317

318 @staticmethod

312 - def run_engine_core(*args, **kwargs):

319 + def run_engine_core(*args,

320 + dp_rank: int = 0,

321 + local_dp_rank: int = 0,

322 + ready_pipe,

323 + **kwargs):

324 """Launch EngineCore busy loop in background process."""

325

326 # Signal handler used for graceful termination.

@@ -331,9 +342,21 @@ class EngineCoreProc(EngineCore):

342 signal.signal(signal.SIGINT, signal_handler)

343

344 parent_process = psutil.Process().parent()

334 - engine_core = None

345 + engine_core: Optional[EngineCoreProc] = None

346 try:

336 - engine_core = EngineCoreProc(*args, **kwargs)

347 + parallel_config: ParallelConfig = kwargs[

348 + "vllm_config"].parallel_config

349 + if parallel_config.data_parallel_size > 1:

350 + # Set data parallel rank for this engine process.

351 + parallel_config.data_parallel_rank = dp_rank

352 + parallel_config.data_parallel_rank_local = local_dp_rank

353 + engine_core = DPEngineCoreProc(*args, **kwargs)

354 + else:

355 + engine_core = EngineCoreProc(*args, **kwargs)

356 +

357 + # Send Readiness signal to EngineClient.

358 + ready_pipe.send({"status": "READY"})

359 +

360 engine_core.run_busy_loop()

361

362 except SystemExit:

@@ -351,28 +374,44 @@ class EngineCoreProc(EngineCore):

374 def run_busy_loop(self):

375 """Core busy loop of the EngineCore."""

376

354 - step_fn = (self.step

355 - if self.batch_queue is None else self.step_with_batch_queue)

356 -

377 # Loop until process is sent a SIGINT or SIGTERM

378 while True:

379 # 1) Poll the input queue until there is work to do.

360 - while not self.scheduler.has_requests():

361 - logger.debug("EngineCore busy loop waiting.")

362 - req = self.input_queue.get()

363 - self._handle_client_request(*req)

364 -

365 - # 2) Handle any new client requests.

366 - while not self.input_queue.empty():

367 - req = self.input_queue.get_nowait()

368 - self._handle_client_request(*req)

369 -

370 - # 3) Step the engine core.

371 - outputs = step_fn()

372 -

373 - # 4) Put EngineCoreOutputs into the output queue.

374 - if outputs is not None:

375 - self.output_queue.put_nowait(outputs)

380 + self._process_input_queue()

381 + # 2) Step the engine core and return the outputs.

382 + self._process_engine_step()

383 +

384 + def _process_input_queue(self):

385 + """Exits when an engine step needs to be performed."""

386 +

387 + waited = False

388 + while not self.global_unfinished_reqs and not (

389 + self.scheduler.has_requests()):

390 + if logger.isEnabledFor(DEBUG) and self.input_queue.empty():

391 + logger.debug("EngineCore waiting for work.")

392 + waited = True

393 + req = self.input_queue.get()

394 + self._handle_client_request(*req)

395 +

396 + if waited:

397 + logger.debug(

398 + "EngineCore loop active - local unfinished: %s, finished: %s.",

399 + self.scheduler.has_unfinished_requests(),

400 + self.scheduler.has_finished_requests())

401 +

402 + # Handle any more client requests.

403 + while not self.input_queue.empty():

404 + req = self.input_queue.get_nowait()

405 + self._handle_client_request(*req)

406 +

407 + def _process_engine_step(self):

408 + """Called only when there are unfinished local requests."""

409 +

410 + # Step the engine core.

411 + outputs = self.step_fn()

412 + # Put EngineCoreOutputs into the output queue.

413 + if outputs is not None:

414 + self.output_queue.put_nowait(outputs)

415

416 def _handle_client_request(self, request_type: EngineCoreRequestType,

417 request: Any) -> None:

@@ -382,6 +421,10 @@ class EngineCoreProc(EngineCore):

421 self.add_request(request)

422 elif request_type == EngineCoreRequestType.ABORT:

423 self.abort_requests(request)

424 + elif request_type == EngineCoreRequestType.START_DP:

425 + if not self.global_unfinished_reqs:

426 + logger.debug("EngineCore starting idle loop.")

427 + self.global_unfinished_reqs = True

428 elif request_type == EngineCoreRequestType.UTILITY:

429 call_id, method_name, args = request

430 output = UtilityOutput(call_id)

@@ -432,7 +475,7 @@ class EngineCoreProc(EngineCore):

475 # Push to input queue for core busy loop.

476 self.input_queue.put_nowait((request_type, request))

477

435 - def process_output_socket(self, output_path: str):

478 + def process_output_socket(self, output_path: str, engine_index: int):

479 """Output socket IO thread."""

480

481 # Msgpack serialization encoding.

@@ -443,5 +486,114 @@ class EngineCoreProc(EngineCore):

486 with zmq_socket_ctx(output_path, zmq.constants.PUSH) as socket:

487 while True:

488 outputs = self.output_queue.get()

489 + outputs.engine_index = engine_index

490 encoder.encode_into(outputs, buffer)

447 - socket.send_multipart((buffer, ), copy=False)

491 + socket.send(buffer, copy=False)

492 +

493 +

494 +ENGINE_PAUSED_OUTPUTS = EngineCoreOutputs(engine_paused=True)

495 +

496 +

497 +class DPEngineCoreProc(EngineCoreProc):

498 + """ZMQ-wrapper for running EngineCore in background process

499 + in a data parallel context."""

500 +

501 + def __init__(

502 + self,

503 + input_path: str,

504 + output_path: str,

505 + vllm_config: VllmConfig,

506 + executor_class: type[Executor],

507 + log_stats: bool,

508 + ):

509 + # Add process-specific prefix to stdout and stderr before

510 + # we initialize the engine.

511 + from multiprocessing import current_process

512 + process_name = current_process().name

513 + pid = os.getpid()

514 + _add_prefix(sys.stdout, process_name, pid)

515 + _add_prefix(sys.stderr, process_name, pid)

516 +

517 + dp_size = vllm_config.parallel_config.data_parallel_size

518 + dp_rank = vllm_config.parallel_config.data_parallel_rank

519 + local_dp_rank = vllm_config.parallel_config.data_parallel_rank_local

520 +

521 + assert dp_size > 1

522 + assert 0 <= local_dp_rank <= dp_rank < dp_size

523 +

524 + from vllm.platforms import current_platform

525 + if current_platform.is_cuda_alike():

526 + from vllm.platforms.cuda import device_id_to_physical_device_id

527 + tp_size = vllm_config.parallel_config.tensor_parallel_size

528 + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(

529 + str(device_id_to_physical_device_id(i))

530 + for i in range(local_dp_rank * tp_size, (local_dp_rank + 1) *

531 + tp_size))

532 +

533 + self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()

534 +

535 + # Initialize the engine after setting up environment.

536 + super().__init__(input_path, output_path, vllm_config, executor_class,

537 + log_stats, dp_rank)

538 +

539 + # Counts forward-passes of the model so that we can synchronize

540 + # finished with DP peers every N steps.

541 + self.counter = 0

542 +

543 + def shutdown(self):

544 + super().shutdown()

545 + if dp_group := getattr(self, "dp_group", None):

546 + stateless_destroy_torch_distributed_process_group(dp_group)

547 +

548 + def run_busy_loop(self):

549 + """Core busy loop of the EngineCore for data parallel case."""

550 +

551 + # Loop until process is sent a SIGINT or SIGTERM

552 + while True:

553 + # 1) Poll the input queue until there is work to do.

554 + self._process_input_queue()

555 +

556 + local_unfinished_reqs = self.scheduler.has_unfinished_requests()

557 +

558 + if local_unfinished_reqs:

559 + # 2) Step the engine core.

560 + self._process_engine_step()

561 +

562 + # Check if we have now finished all requests.

563 + local_unfinished_reqs = (

564 + self.scheduler.has_unfinished_requests())

565 + else:

566 + if self.scheduler.has_finished_requests():

567 + # There are no unfinished requests, but there are some

568 + # finished requests remaining to be removed from the

569 + # batch state. This engine step won't perform a forward

570 + # pass but will flush the finished requests to ensure

571 + # up-to-date state is returned in the engine outputs.

572 + self._process_engine_step()

573 +

574 + if not self.global_unfinished_reqs:

575 + # All engines are idle.

576 + continue

577 +

578 + # There must be unfinished requests in DP peers, run a

579 + # dummy forward pass.

580 + self.execute_dummy_batch()

581 +

582 + # 3) All-reduce operation to determine global unfinished reqs.

583 + self.global_unfinished_reqs = self._has_global_unfinished_reqs(

584 + local_unfinished_reqs)

585 +

586 + if not self.global_unfinished_reqs:

587 + # Notify client that we are pausing the loop.

588 + self.output_queue.put_nowait(ENGINE_PAUSED_OUTPUTS)

589 +

590 + def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:

591 +

592 + # Optimization - only perform finish-sync all-reduce every 16 steps.

593 + self.counter += 1

594 + if self.counter != 16:

595 + return True

596 + self.counter = 0

597 +

598 + return ParallelConfig.has_unfinished_dp(self.dp_group,

599 + local_unfinished)

vllm/v1/engine/core_client.py

+266 -68

@@ -8,10 +8,11 @@ import threading

8 import uuid

9 import weakref

10 from abc import ABC, abstractmethod

11 +from collections.abc import Awaitable, Sequence

12 from concurrent.futures import Future

12 -from dataclasses import dataclass

13 +from dataclasses import dataclass, field

14 from threading import Thread

14 -from typing import Any, Optional, Union

15 +from typing import Any, Callable, Optional, Union

17 import zmq

18 import zmq.asyncio

@@ -60,6 +61,9 @@ class EngineCoreClient(ABC):

61 "is not currently supported.")

63 if multiprocess_mode and asyncio_mode:

64 + if vllm_config.parallel_config.data_parallel_size > 1:

65 + return DPAsyncMPClient(vllm_config, executor_class, log_stats)

66 +

67 return AsyncMPClient(vllm_config, executor_class, log_stats)

69 if multiprocess_mode and not asyncio_mode:

@@ -207,28 +211,74 @@ class InprocClient(EngineCoreClient):

211 return self.engine_core.pin_lora(lora_id)

212

213

214 +class CoreEngine:

215 + """One per data parallel rank."""

216 +

217 + def __init__(

218 + self,

219 + vllm_config: VllmConfig,

220 + executor_class: type[Executor],

221 + log_stats: bool,

222 + ctx: Union[zmq.Context, zmq.asyncio.Context],

223 + output_path: str,

224 + index: int = 0,

225 + local_dp_rank: int = 0,

226 + ):

227 + # Paths and sockets for IPC.

228 + input_path = get_open_zmq_ipc_path()

229 + self.input_socket = make_zmq_socket(ctx, input_path,

230 + zmq.constants.PUSH)

231 + try:

232 + # Start EngineCore in background process.

233 + self.proc_handle = BackgroundProcHandle(

234 + input_path=input_path,

235 + output_path=output_path,

236 + process_name=f"EngineCore_{index}",

237 + target_fn=EngineCoreProc.run_engine_core,

238 + process_kwargs={

239 + "vllm_config": vllm_config,

240 + "dp_rank": index,

241 + "local_dp_rank": local_dp_rank,

242 + "executor_class": executor_class,

243 + "log_stats": log_stats,

244 + })

245 +

246 + self.num_reqs_in_flight = 0

247 + finally:

248 + if not hasattr(self, "num_reqs_in_flight"):

249 + # Ensure socket is closed if process fails to start.

250 + self.close()

251 +

252 + def send_multipart(self, msg_parts: Sequence):

253 + return self.input_socket.send_multipart(msg_parts, copy=False)

254 +

255 + def close(self):

256 + if proc_handle := getattr(self, "proc_handle", None):

257 + proc_handle.shutdown()

258 + if socket := getattr(self, "input_socket", None):

259 + socket.close(linger=0)

260 +

261 +

262 @dataclass

263 class BackgroundResources:

264 """Used as a finalizer for clean shutdown, avoiding

265 circular reference back to the client object."""

266

215 - ctx: zmq.Context

267 + ctx: Union[zmq.Context]

268 + core_engines: list[CoreEngine] = field(default_factory=list)

269 output_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None

217 - input_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None

218 - proc_handle: Optional[BackgroundProcHandle] = None

270 shutdown_path: Optional[str] = None

271

272 def __call__(self):

273 """Clean up background resources."""

274

224 - if self.proc_handle is not None:

225 - self.proc_handle.shutdown()

275 + for core_engine in self.core_engines:

276 + core_engine.close()

277 +

278 # ZMQ context termination can hang if the sockets

279 # aren't explicitly closed first.

280 if self.output_socket is not None:

281 self.output_socket.close(linger=0)

230 - if self.input_socket is not None:

231 - self.input_socket.close(linger=0)

282 if self.shutdown_path is not None:

283 # We must ensure that the sync output socket is

284 # closed cleanly in its own thread.

@@ -284,7 +334,7 @@ class MPClient(EngineCoreClient):

334 self.decoder = MsgpackDecoder(EngineCoreOutputs)

335

336 # ZMQ setup.

287 - sync_ctx = zmq.Context()

337 + sync_ctx = zmq.Context(io_threads=2)

338 self.ctx = zmq.asyncio.Context(sync_ctx) if asyncio_mode else sync_ctx

339

340 # This will ensure resources created so far are closed

@@ -293,28 +343,38 @@ class MPClient(EngineCoreClient):

343 self.resources = BackgroundResources(ctx=sync_ctx)

344 self._finalizer = weakref.finalize(self, self.resources)

345

296 - # Paths for IPC.

346 + # Paths and sockets for IPC.

347 self.output_path = get_open_zmq_ipc_path()

298 - input_path = get_open_zmq_ipc_path()

348

300 - # Start EngineCore in background process.

301 - self.resources.proc_handle = BackgroundProcHandle(

302 - input_path=input_path,

303 - output_path=self.output_path,

304 - process_name="EngineCore",

305 - target_fn=EngineCoreProc.run_engine_core,

306 - process_kwargs={

307 - "vllm_config": vllm_config,

308 - "executor_class": executor_class,

309 - "log_stats": log_stats,

310 - })

311 -

312 - # Create input socket.

313 - self.resources.input_socket = make_zmq_socket(self.ctx, input_path,

314 - zmq.constants.PUSH)

315 - self.input_socket = self.resources.input_socket

349 + new_core_engine = lambda index, local_dp_rank=None: CoreEngine(

350 + vllm_config, executor_class, log_stats, self.ctx, self.output_path,

351 + index, local_dp_rank)

352 +

353 + # Start engine core process(es).

354 + self._init_core_engines(vllm_config, new_core_engine,

355 + self.resources.core_engines)

356 +

357 + # Wait for engine core process(es) to start.

358 + for engine in self.resources.core_engines:

359 + engine.proc_handle.wait_for_startup()

360 +

361 self.utility_results: dict[int, AnyFuture] = {}

362

363 + def _init_core_engines(

364 + self,

365 + vllm_config: VllmConfig,

366 + new_core_engine: Callable[[int, Optional[int]], CoreEngine],

367 + core_engines: list[CoreEngine],

368 + ) -> None:

369 +

370 + # Default case - single core engine.

371 + dp_rank = vllm_config.parallel_config.data_parallel_rank

372 + local_dp_rank = vllm_config.parallel_config.data_parallel_rank_local

373 + core_engine = new_core_engine(

374 + dp_rank, local_dp_rank if local_dp_rank is not None else dp_rank)

375 + core_engines.append(core_engine)

376 + self.core_engine = core_engine

377 +

378 def shutdown(self):

379 self._finalizer()

380

@@ -370,7 +430,7 @@ class SyncMPClient(MPClient):

430 # shutdown signal, exit thread.

431 break

432

373 - (frame, ) = out_socket.recv_multipart(copy=False)

433 + frame = out_socket.recv(copy=False)

434 outputs = decoder.decode(frame.buffer)

435 if outputs.utility_output:

436 _process_utility_output(outputs.utility_output,

@@ -391,18 +451,15 @@ class SyncMPClient(MPClient):

451 def get_output(self) -> EngineCoreOutputs:

452 return self.outputs_queue.get()

453

394 - def _send_input(self, request_type: EngineCoreRequestType,

395 - request: Any) -> None:

396 -

454 + def _send_input(self, request_type: EngineCoreRequestType, request: Any):

455 # (RequestType, SerializedRequest)

456 msg = (request_type.value, self.encoder.encode(request))

399 - self.input_socket.send_multipart(msg, copy=False)

457 + self.core_engine.send_multipart(msg)

458

401 - def _call_utility(self, method: str, *args) -> Any:

459 + def call_utility(self, method: str, *args) -> Any:

460 call_id = uuid.uuid1().int >> 64

461 future: Future[Any] = Future()

462 self.utility_results[call_id] = future

405 -

463 self._send_input(EngineCoreRequestType.UTILITY,

464 (call_id, method, args))

465

@@ -419,34 +476,34 @@ class SyncMPClient(MPClient):

476 self._send_input(EngineCoreRequestType.ABORT, request_ids)

477

478 def profile(self, is_start: bool = True) -> None:

422 - self._call_utility("profile", is_start)

479 + self.call_utility("profile", is_start)

480

481 def reset_prefix_cache(self) -> None:

425 - self._call_utility("reset_prefix_cache")

482 + self.call_utility("reset_prefix_cache")

483

484 def add_lora(self, lora_request: LoRARequest) -> bool:

428 - return self._call_utility("add_lora", lora_request)

485 + return self.call_utility("add_lora", lora_request)

486

487 def remove_lora(self, lora_id: int) -> bool:

431 - return self._call_utility("remove_lora", lora_id)

488 + return self.call_utility("remove_lora", lora_id)

489

490 def list_loras(self) -> set[int]:

434 - return self._call_utility("list_loras")

491 + return self.call_utility("list_loras")

492

493 def pin_lora(self, lora_id: int) -> bool:

437 - return self._call_utility("pin_lora", lora_id)

494 + return self.call_utility("pin_lora", lora_id)

495

496 def sleep(self, level: int = 1) -> None:

440 - self._call_utility("sleep", level)

497 + self.call_utility("sleep", level)

498

499 def wake_up(self) -> None:

443 - self._call_utility("wake_up")

500 + self.call_utility("wake_up")

501

502 def is_sleeping(self) -> bool:

446 - return self._call_utility("is_sleeping")

503 + return self.call_utility("is_sleeping")

504

505 def execute_dummy_batch(self) -> None:

449 - self._call_utility("execute_dummy_batch")

506 + self.call_utility("execute_dummy_batch")

507

508

509 class AsyncMPClient(MPClient):

@@ -464,13 +521,21 @@ class AsyncMPClient(MPClient):

521 self.outputs_queue: Optional[asyncio.Queue[EngineCoreOutputs]] = None

522 self.queue_task: Optional[asyncio.Task] = None

523

467 - async def _start_output_queue_task(self):

524 + self.outputs_handler: Optional[Callable[

525 + [AsyncMPClient, EngineCoreOutputs], Awaitable[None]]] = None

526 +

527 + def _ensure_output_queue_task(self):

528 + if self.outputs_queue is not None:

529 + return

530 +

531 # Perform IO in separate task to parallelize as much as possible.

532 # Avoid task having direct reference back to the client.

533 self.outputs_queue = asyncio.Queue()

534 decoder = self.decoder

535 utility_results = self.utility_results

536 outputs_queue = self.outputs_queue

537 + output_handler = self.outputs_handler

538 + _self_ref = weakref.ref(self) if output_handler else None

539 output_path = self.output_path

540 output_socket = make_zmq_socket(self.ctx, output_path,

541 zmq.constants.PULL)

@@ -483,34 +548,52 @@ class AsyncMPClient(MPClient):

548 if outputs.utility_output:

549 _process_utility_output(outputs.utility_output,

550 utility_results)

486 - else:

551 + continue

552 +

553 + if output_handler is not None:

554 + assert _self_ref is not None

555 + _self = _self_ref()

556 + if not _self:

557 + # Client has been garbage collected, abort.

558 + return

559 + await output_handler(_self, outputs)

560 +

561 + if outputs.outputs or outputs.scheduler_stats:

562 outputs_queue.put_nowait(outputs)

563

564 self.queue_task = asyncio.create_task(process_outputs_socket(),

565 name="EngineCoreOutputQueueTask")

566

567 async def get_output_async(self) -> EngineCoreOutputs:

493 - if self.outputs_queue is None:

494 - await self._start_output_queue_task()

495 - assert self.outputs_queue is not None

568 + self._ensure_output_queue_task()

569 + assert self.outputs_queue is not None

570 return await self.outputs_queue.get()

571

572 async def _send_input(self, request_type: EngineCoreRequestType,

573 request: Any) -> None:

574 + await self.core_engine.send_multipart(

575 + (request_type.value, self.encoder.encode(request)))

576

501 - msg = (request_type.value, self.encoder.encode(request))

502 - await self.input_socket.send_multipart(msg, copy=False)

577 + self._ensure_output_queue_task()

578

504 - if self.outputs_queue is None:

505 - await self._start_output_queue_task()

579 + async def call_utility_async(self, method: str, *args) -> Any:

580 + return await self._call_utility_async(method,

581 + *args,

582 + engine=self.core_engine)

583

507 - async def _call_utility_async(self, method: str, *args) -> Any:

584 + async def _call_utility_async(

585 + self,

586 + method: str,

587 + *args,

588 + engine: CoreEngine,

589 + ) -> Any:

590 call_id = uuid.uuid1().int >> 64

591 future = asyncio.get_running_loop().create_future()

592 self.utility_results[call_id] = future

511 - await self._send_input(EngineCoreRequestType.UTILITY,

512 - (call_id, method, args))

513 -

593 + message = (EngineCoreRequestType.UTILITY.value,

594 + self.encoder.encode((call_id, method, args)))

595 + await engine.send_multipart(message)

596 + self._ensure_output_queue_task()

597 return await future

598

599 async def add_request_async(self, request: EngineCoreRequest) -> None:

@@ -524,31 +607,146 @@ class AsyncMPClient(MPClient):

607 await self._send_input(EngineCoreRequestType.ABORT, request_ids)

608

609 async def profile_async(self, is_start: bool = True) -> None:

527 - await self._call_utility_async("profile", is_start)

610 + await self.call_utility_async("profile", is_start)

611

612 async def reset_prefix_cache_async(self) -> None:

530 - await self._call_utility_async("reset_prefix_cache")

613 + await self.call_utility_async("reset_prefix_cache")

614

615 async def sleep_async(self, level: int = 1) -> None:

533 - await self._call_utility_async("sleep", level)

616 + await self.call_utility_async("sleep", level)

617

618 async def wake_up_async(self) -> None:

536 - await self._call_utility_async("wake_up")

619 + await self.call_utility_async("wake_up")

620

621 async def is_sleeping_async(self) -> bool:

539 - return await self._call_utility_async("is_sleeping")

622 + return await self.call_utility_async("is_sleeping")

623

624 async def execute_dummy_batch_async(self) -> None:

542 - await self._call_utility_async("execute_dummy_batch")

625 + await self.call_utility_async("execute_dummy_batch")

626

627 async def add_lora_async(self, lora_request: LoRARequest) -> bool:

545 - return await self._call_utility_async("add_lora", lora_request)

628 + return await self.call_utility_async("add_lora", lora_request)

629

630 async def remove_lora_async(self, lora_id: int) -> bool:

548 - return await self._call_utility_async("remove_lora", lora_id)

631 + return await self.call_utility_async("remove_lora", lora_id)

632

633 async def list_loras_async(self) -> set[int]:

551 - return await self._call_utility_async("list_loras")

634 + return await self.call_utility_async("list_loras")

635

636 async def pin_lora_async(self, lora_id: int) -> bool:

554 - return await self._call_utility_async("pin_lora", lora_id)

637 + return await self.call_utility_async("pin_lora", lora_id)

638 +

639 +

640 +class DPAsyncMPClient(AsyncMPClient):

641 + """Asyncio-compatible client for multi-proc, multi-engine (data parallel)

642 + EngineCore."""

643 +

644 + def __init__(self, vllm_config: VllmConfig, executor_class: type[Executor],

645 + log_stats: bool):

646 + super().__init__(vllm_config, executor_class, log_stats)

647 +

648 + assert len(self.core_engines) > 1

649 +

650 + # Control message used for triggering dp idle mode loop.

651 + self.start_dp_msg = (EngineCoreRequestType.START_DP.value,

652 + self.encoder.encode(None))

653 +

654 + self.num_engines_running = 0

655 + self.reqs_in_flight: dict[str, CoreEngine] = {}

656 +

657 + self.outputs_handler = DPAsyncMPClient.process_engine_outputs # type: ignore[assignment]

658 +

659 + def _init_core_engines(

660 + self,

661 + vllm_config: VllmConfig,

662 + new_core_engine: Callable[[int, Optional[int]], CoreEngine],

663 + core_engines: list[CoreEngine],

664 + ) -> None:

665 +

666 + # Launch a core engine for each data parallel rank.

667 + dp_size = vllm_config.parallel_config.data_parallel_size

668 + for i in range(dp_size):

669 + # Multi-node not yet supported so local_dp_rank == dp_rank.

670 + core_engines.append(new_core_engine(i, i))

671 +

672 + self.core_engines = core_engines

673 +

674 + async def call_utility_async(self, method: str, *args) -> Any:

675 + # Only the result from the first engine is returned.

676 + return (await asyncio.gather(*[

677 + self._call_utility_async(method, *args, engine=engine)

678 + for engine in self.core_engines

679 + ]))[0]

680 +

681 + async def add_request_async(self, request: EngineCoreRequest) -> None:

682 + # NOTE: text prompt is not needed in the core engine as it has been

683 + # tokenized.

684 + request.prompt = None

685 +

686 + msg = (EngineCoreRequestType.ADD.value, self.encoder.encode(request))

687 +

688 + chosen_engine = self.get_core_engine_for_request()

689 + self.reqs_in_flight[request.request_id] = chosen_engine

690 + chosen_engine.num_reqs_in_flight += 1

691 + if self.num_engines_running >= len(self.core_engines):

692 + await chosen_engine.send_multipart(msg)

693 + else:

694 + # Send request to chosen engine and dp start loop

695 + # control message to all other engines.

696 + self.num_engines_running += len(self.core_engines)

697 + await asyncio.gather(*[

698 + engine.send_multipart(msg if engine is

699 + chosen_engine else self.start_dp_msg)

700 + for engine in self.core_engines

701 + ])

702 +

703 + self._ensure_output_queue_task()

704 +

705 + def get_core_engine_for_request(self) -> CoreEngine:

706 + return min(self.core_engines, key=lambda e: e.num_reqs_in_flight)

707 +

708 + @staticmethod

709 + async def process_engine_outputs(self: "DPAsyncMPClient",

710 + outputs: EngineCoreOutputs):

711 + if self.reqs_in_flight:

712 + for req_id in outputs.finished_requests or ():

713 + if engine := self.reqs_in_flight.pop(req_id, None):

714 + engine.num_reqs_in_flight -= 1

715 +

716 + if outputs.engine_paused:

717 + assert self.num_engines_running >= 1

718 + self.num_engines_running -= 1

719 + if not self.num_engines_running and self.reqs_in_flight:

720 + # If there are requests in flight here, they must have

721 + # been sent after the engines paused. We must make

722 + # sure to start the other engines:

723 + self.num_engines_running = len(self.core_engines)

724 + coros = [

725 + engine.send_multipart(self.start_dp_msg)

726 + for engine in self.core_engines

727 + if not engine.num_reqs_in_flight

728 + ]

729 + if coros:

730 + await asyncio.gather(*coros)

731 +

732 + async def abort_requests_async(self, request_ids: list[str]) -> None:

733 + if not request_ids:

734 + return

735 +

736 + if len(request_ids) == 1:

737 + # Fast-path common case.

738 + if engine := self.reqs_in_flight.get(request_ids[0]):

739 + await self._abort_requests(request_ids, engine)

740 + return

741 +

742 + by_engine: dict[CoreEngine, list[str]] = {}

743 + for req_id in request_ids:

744 + if engine := self.reqs_in_flight.get(req_id):

745 + by_engine.setdefault(engine, []).append(req_id)

746 + for engine, req_ids in by_engine.items():

747 + await self._abort_requests(req_ids, engine)

748 +

749 + async def _abort_requests(self, request_ids: list[str],

750 + engine: CoreEngine) -> None:

751 + await engine.send_multipart((EngineCoreRequestType.ABORT.value,

752 + self.encoder.encode(request_ids)))

vllm/v1/engine/llm_engine.py

+12 -5

@@ -8,6 +8,7 @@ from typing_extensions import TypeVar

9 import vllm.envs as envs

10 from vllm.config import ParallelConfig, VllmConfig

11 +from vllm.distributed import stateless_destroy_torch_distributed_process_group

12 from vllm.engine.arg_utils import EngineArgs

13 from vllm.engine.metrics_types import StatLoggerBase

14 from vllm.inputs import INPUT_REGISTRY, InputRegistry, PromptType

@@ -60,11 +61,13 @@ class LLMEngine:

61 self.cache_config = vllm_config.cache_config

63 # important: init dp group before init the engine_core

63 - self.parallel_config = vllm_config.parallel_config

64 - self.dp_enabled = self.parallel_config.data_parallel_size > 1 # noqa

64 + # In the decoupled engine case this is handled in EngineCoreProc.

65 + parallel_config = vllm_config.parallel_config

66 + if not multiprocess_mode and parallel_config.data_parallel_size > 1:

67 + self.dp_group = parallel_config.stateless_init_dp_group()

68 + else:

69 + self.dp_group = None

70 self.should_execute_dummy_batch = False

66 - if self.dp_enabled:

67 - self.dp_group = self.parallel_config.stateless_init_dp_group()

72 # Tokenizer (+ ensure liveness if running in another process).

73 self.tokenizer = init_tokenizer_from_configs(

@@ -148,7 +151,7 @@ class LLMEngine:

151

152 def has_unfinished_requests(self) -> bool:

153 has_unfinished = self.output_processor.has_unfinished_requests()

151 - if not self.dp_enabled:

154 + if self.dp_group is None:

155 return has_unfinished

156 return self.has_unfinished_requests_dp(has_unfinished)

157

@@ -280,3 +283,7 @@ class LLMEngine:

283 def pin_lora(self, lora_id: int) -> bool:

284 """Prevent an adapter from being evicted."""

285 return self.engine_core.pin_lora(lora_id)

286 +

287 + def __del__(self):

288 + if dp_group := getattr(self, "dp_group", None):

289 + stateless_destroy_torch_distributed_process_group(dp_group)

vllm/v1/executor/multiproc_executor.py

+21 -17

@@ -235,7 +235,10 @@ class WorkerProc:

235 worker_response_mq_handle = self.worker_response_mq.export_handle()

236

237 # Send Readiness signal to EngineCore process.

238 - with zmq_socket_ctx(ready_path, zmq.constants.PUSH) as ready_socket:

238 + # Set linger here because we want to ensure the message has

239 + # been sent before the context is closed.

240 + with zmq_socket_ctx(ready_path, zmq.constants.PUSH,

241 + linger=10000) as ready_socket:

242 payload = pickle.dumps(worker_response_mq_handle,

243 protocol=pickle.HIGHEST_PROTOCOL)

244 ready_socket.send_string(WorkerProc.READY_STR)

@@ -270,11 +273,13 @@ class WorkerProc:

273 proc = context.Process(target=WorkerProc.worker_main,

274 kwargs=process_kwargs,

275 daemon=True)

273 - proc.start()

276

275 - # Wait for startup

276 - worker_response_mq_handle = WorkerProc.wait_for_startup(

277 - proc, ready_path)

277 + with zmq_socket_ctx(ready_path, zmq.constants.PULL) as ready_socket:

278 + proc.start()

279 +

280 + # Wait for startup

281 + worker_response_mq_handle = WorkerProc.wait_for_startup(

282 + proc, ready_socket)

283

284 worker_response_mq = MessageQueue.create_from_handle(

285 worker_response_mq_handle, 0)

@@ -337,23 +342,22 @@ class WorkerProc:

342 @staticmethod

343 def wait_for_startup(

344 proc: BaseProcess,

340 - ready_path: str,

345 + ready_socket: zmq.Socket,

346 ) -> Optional[Handle]:

347 """Wait until the Worker is ready."""

343 - with zmq_socket_ctx(ready_path, zmq.constants.PULL) as socket:

348

345 - # Wait for Worker to send READY.

346 - while socket.poll(timeout=POLLING_TIMEOUT_MS) == 0:

347 - logger.debug("Waiting for WorkerProc to startup.")

349 + # Wait for Worker to send READY.

350 + while ready_socket.poll(timeout=POLLING_TIMEOUT_MS) == 0:

351 + logger.debug("Waiting for WorkerProc to startup.")

352

349 - if not proc.is_alive():

350 - raise RuntimeError("WorkerProc failed to start.")

353 + if not proc.is_alive():

354 + raise RuntimeError("WorkerProc failed to start.")

355

352 - message = socket.recv_string()

353 - assert message == WorkerProc.READY_STR

354 - handle_frame = socket.recv(copy=False)

355 - handle = pickle.loads(handle_frame.buffer)

356 - return handle

356 + message = ready_socket.recv_string()

357 + assert message == WorkerProc.READY_STR

358 + handle_frame = ready_socket.recv(copy=False)

359 + handle = pickle.loads(handle_frame.buffer)

360 + return handle

361

362 class ResponseStatus(Enum):

363 SUCCESS = auto()

vllm/v1/metrics/loggers.py

+10 -4

@@ -31,7 +31,8 @@ class StatLoggerBase(ABC):

32 class LoggingStatLogger(StatLoggerBase):

34 - def __init__(self):

34 + def __init__(self, engine_index: int = 0):

35 + self.engine_index = engine_index

36 self._reset(time.monotonic())

37 self.last_scheduler_stats = SchedulerStats()

38 # Prefix cache metrics. This cannot be reset.

@@ -78,11 +79,13 @@ class LoggingStatLogger(StatLoggerBase):

80 # Format and print output.

81 logger.info(

82 + "Engine %03d: "

83 "Avg prompt throughput: %.1f tokens/s, "

84 "Avg generation throughput: %.1f tokens/s, "

85 "Running: %d reqs, Waiting: %d reqs, "

86 "GPU KV cache usage: %.1f%%, "

87 "Prefix cache hit rate: %.1f%%",

88 + self.engine_index,

89 prompt_throughput,

90 generation_throughput,

91 scheduler_stats.num_running_reqs,

@@ -94,7 +97,7 @@ class LoggingStatLogger(StatLoggerBase):

98 class PrometheusStatLogger(StatLoggerBase):

97 - def __init__(self, vllm_config: VllmConfig):

100 + def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):

101 self._unregister_vllm_metrics()

102

103 # Use this flag to hide metrics that were deprecated in

@@ -102,8 +105,11 @@ class PrometheusStatLogger(StatLoggerBase):

105 self.show_hidden_metrics = \

106 vllm_config.observability_config.show_hidden_metrics

107

105 - labelnames = ["model_name"]

106 - labelvalues = [vllm_config.model_config.served_model_name]

108 + labelnames = ["model_name", "engine"]

109 + labelvalues = [

110 + vllm_config.model_config.served_model_name,

111 + str(engine_index)

112 + ]

113

114 max_model_len = vllm_config.model_config.max_model_len

115

vllm/v1/utils.py

+7 -4

@@ -105,7 +105,7 @@ class BackgroundProcHandle:

105 process_kwargs: dict[Any, Any],

106 ):

107 context = get_mp_context()

108 - reader, writer = context.Pipe(duplex=False)

108 + self.reader, writer = context.Pipe(duplex=False)

109

110 assert ("ready_pipe" not in process_kwargs

111 and "input_path" not in process_kwargs

@@ -115,14 +115,17 @@ class BackgroundProcHandle:

115 process_kwargs["output_path"] = output_path

116

117 # Run busy loop in background process.

118 - self.proc = context.Process(target=target_fn, kwargs=process_kwargs)

118 + self.proc = context.Process(target=target_fn,

119 + kwargs=process_kwargs,

120 + name=process_name)

121 self._finalizer = weakref.finalize(self, shutdown, self.proc,

122 input_path, output_path)

123 self.proc.start()

124

125 + def wait_for_startup(self):

126 # Wait for startup.

124 - if reader.recv()["status"] != "READY":

125 - raise RuntimeError(f"{process_name} initialization failed. "

127 + if self.reader.recv()["status"] != "READY":

128 + raise RuntimeError(f"{self.proc.name} initialization failed. "

129 "See root cause above.")

130

131 def shutdown(self):