@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 55aa7af994

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, 11:00 UTC

[V1] DP scale-out (2/N): Decouple engine process management and comms (#15977)

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

Nick Hill committed May 13, 2025 at 10:48 UTC 55aa7af9947244bdc427bfae01e750b1c157bd8f

10 files changed +525 -252

tests/async_engine/test_async_llm_engine.py

+1 -1

@@ -41,7 +41,7 @@ class MockEngine:

41 self.abort_request_calls = 0

42 self.request_id = None

43 # Ugly, remove dependency when possible

44 - self.parallel_config = ParallelConfig(1, 1, False)

44 + self.parallel_config = ParallelConfig()

45 self.model_config = MockModelConfig()

47 async def step_async(self, virtual_engine):

tests/v1/engine/test_engine_core_client.py

+8 -7

@@ -18,9 +18,10 @@ from vllm.platforms import current_platform

18 from vllm.usage.usage_lib import UsageContext

19 from vllm.v1.engine import EngineCoreRequest

20 from vllm.v1.engine.core import EngineCore

21 -from vllm.v1.engine.core_client import (AsyncMPClient, CoreEngine,

22 - EngineCoreClient, SyncMPClient)

21 +from vllm.v1.engine.core_client import (AsyncMPClient, EngineCoreClient,

22 + SyncMPClient)

23 from vllm.v1.executor.abstract import Executor

24 +from vllm.v1.utils import CoreEngineProcManager

26 from ...distributed.conftest import MockSubscriber

27 from ...utils import create_new_process_for_each_test

@@ -348,13 +349,13 @@ def test_startup_failure(monkeypatch: pytest.MonkeyPatch):

349

350 # Monkey-patch to extract core process pid while it's starting.

351 core_proc_pid = [None]

351 - ce_ctor = CoreEngine.__init__

352 + cepm_ctor = CoreEngineProcManager.__init__

353

353 - def patched_ce_ctor(self, *args, **kwargs):

354 - ce_ctor(self, *args, **kwargs)

355 - core_proc_pid[0] = self.proc_handle.proc.pid

354 + def patched_cepm_ctor(self: CoreEngineProcManager, *args, **kwargs):

355 + cepm_ctor(self, *args, **kwargs)

356 + core_proc_pid[0] = self.processes[0].pid

357

357 - m.setattr(CoreEngine, "__init__", patched_ce_ctor)

358 + m.setattr(CoreEngineProcManager, "__init__", patched_cepm_ctor)

359

360 t = time.time()

361 engine_args = EngineArgs(model=MODEL_NAME)

vllm/config.py

+19 -22

@@ -1668,25 +1668,17 @@ class ParallelConfig:

1668 data_parallel_size: int = 1

1669 """Number of data parallel groups. MoE layers will be sharded according to

1670 the product of the tensor parallel size and data parallel size."""

1671 + data_parallel_size_local: int = 1

1672 + """Number of local data parallel groups."""

1673 data_parallel_rank: int = 0

1674 """Rank of the data parallel group."""

1673 - _data_parallel_rank_local: Optional[int] = field(default=None, init=False)

1674 - """Private field to store the local rank of the data parallel group."""

1675 -

1676 - @property

1677 - def data_parallel_rank_local(self) -> int:

1678 - """Local rank of the data parallel group, defaults to global rank."""

1679 - if self._data_parallel_rank_local is None:

1680 - return self.data_parallel_rank

1681 - return self._data_parallel_rank_local

1682 -

1683 - @data_parallel_rank_local.setter

1684 - def data_parallel_rank_local(self, value: int) -> None:

1685 - """Set the local rank of the data parallel group."""

1686 - self._data_parallel_rank_local = value

1687 -

1675 + data_parallel_rank_local: Optional[int] = None

1676 + """Local rank of the data parallel group,

1677 + set only in SPMD mode."""

1678 data_parallel_master_ip: str = "127.0.0.1"

1679 """IP of the data parallel master."""

1680 + data_parallel_rpc_port: int = 29550

1681 + """Port for data parallel messaging."""

1682 data_parallel_master_port: int = 29500

1683 """Port of the data parallel master."""

1684 enable_expert_parallel: bool = False

@@ -1734,13 +1726,16 @@ class ParallelConfig:

1726

1727 world_size: int = field(init=False)

1728 """world_size is TPxPP, it affects the number of workers we create."""

1737 - world_size_across_dp: int = field(init=False)

1738 - """world_size_across_dp is TPxPPxDP, it is the size of the world

1739 - including data parallelism."""

1729

1730 rank: int = 0

1731 """Global rank in distributed setup."""

1732

1733 + @property

1734 + def world_size_across_dp(self) -> int:

1735 + """world_size_across_dp is TPxPPxDP, it is the size of the world

1736 + including data parallelism."""

1737 + return self.world_size * self.data_parallel_size

1738 +

1739 def get_next_dp_init_port(self) -> int:

1740 """

1741 We might need to initialize process groups in multiple

@@ -1800,10 +1795,14 @@ class ParallelConfig:

1795 self.world_size = self.pipeline_parallel_size * \

1796 self.tensor_parallel_size

1797

1803 - if self.data_parallel_size > 1:

1798 + if self.data_parallel_size_local > self.data_parallel_size:

1799 + raise ValueError(

1800 + f"data_parallel_size_local ({self.data_parallel_size_local}) "

1801 + f"must be <= data_parallel_size ({self.data_parallel_size})")

1802 +

1803 + if self.data_parallel_size > 1 or self.data_parallel_size_local == 0:

1804 # Data parallel was specified in the engine args.

1805 self.data_parallel_master_port = get_open_port()

1806 - # TODO multi-node

1806 else:

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

1808 self.data_parallel_size = envs.VLLM_DP_SIZE

@@ -1812,8 +1811,6 @@ class ParallelConfig:

1811 self.data_parallel_master_ip = envs.VLLM_DP_MASTER_IP

1812 self.data_parallel_master_port = envs.VLLM_DP_MASTER_PORT

1813

1815 - self.world_size_across_dp = self.world_size * self.data_parallel_size

1816 -

1814 if self.distributed_executor_backend == "external_launcher":

1815 import os

1816 os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

vllm/distributed/utils.py

+2 -1

@@ -22,6 +22,7 @@ from torch.distributed.rendezvous import rendezvous

23 import vllm.envs as envs

24 from vllm.logger import init_logger

25 +from vllm.utils import get_tcp_uri

27 logger = init_logger(__name__)

@@ -303,7 +304,7 @@ def stateless_init_torch_distributed_process_group(

304 always formed with process 1, 2, ..., 8, and the additional communication

305 channel is formed with process 9 and 10.

306 """

306 - init_method = f"tcp://{host}:{port}"

307 + init_method = get_tcp_uri(host, port)

308 backend = Backend(backend) # it is basically string

309 timeout = _get_default_timeout(backend)

310

vllm/engine/arg_utils.py

+38

@@ -283,6 +283,9 @@ class EngineArgs:

283 pipeline_parallel_size: int = ParallelConfig.pipeline_parallel_size

284 tensor_parallel_size: int = ParallelConfig.tensor_parallel_size

285 data_parallel_size: int = ParallelConfig.data_parallel_size

286 + data_parallel_size_local: Optional[int] = None

287 + data_parallel_address: Optional[str] = None

288 + data_parallel_rpc_port: Optional[int] = None

289 enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel

290 max_parallel_loading_workers: Optional[

291 int] = ParallelConfig.max_parallel_loading_workers

@@ -596,6 +599,21 @@ class EngineArgs:

599 **parallel_kwargs["tensor_parallel_size"])

600 parallel_group.add_argument("--data-parallel-size", "-dp",

601 **parallel_kwargs["data_parallel_size"])

602 + parallel_group.add_argument('--data-parallel-size-local',

603 + '-dpl',

604 + type=int,

605 + help='Number of data parallel replicas '

606 + 'to run on this node.')

607 + parallel_group.add_argument('--data-parallel-address',

608 + '-dpa',

609 + type=str,

610 + help='Address of data parallel cluster '

611 + 'head-node.')

612 + parallel_group.add_argument('--data-parallel-rpc-port',

613 + '-dpp',

614 + type=int,

615 + help='Port for data parallel RPC '

616 + 'communication.')

617 parallel_group.add_argument(

618 "--enable-expert-parallel",

619 **parallel_kwargs["enable_expert_parallel"])

@@ -1019,10 +1037,30 @@ class EngineArgs:

1037 # but we should not do this here.

1038 placement_group = ray.util.get_current_placement_group()

1039

1040 + # Local DP size defaults to global DP size if not set.

1041 + data_parallel_size_local = self.data_parallel_size if (

1042 + self.data_parallel_size_local

1043 + is None) else self.data_parallel_size_local

1044 +

1045 + # DP address, used in multi-node case for torch distributed group

1046 + # and ZMQ sockets.

1047 + data_parallel_address = self.data_parallel_address if (

1048 + self.data_parallel_address

1049 + is not None) else ParallelConfig.data_parallel_master_ip

1050 +

1051 + # This port is only used when there are remote data parallel engines,

1052 + # otherwise the local IPC transport is used.

1053 + data_parallel_rpc_port = self.data_parallel_rpc_port if (

1054 + self.data_parallel_rpc_port

1055 + is not None) else ParallelConfig.data_parallel_rpc_port

1056 +

1057 parallel_config = ParallelConfig(

1058 pipeline_parallel_size=self.pipeline_parallel_size,

1059 tensor_parallel_size=self.tensor_parallel_size,

1060 data_parallel_size=self.data_parallel_size,

1061 + data_parallel_size_local=data_parallel_size_local,

1062 + data_parallel_master_ip=data_parallel_address,

1063 + data_parallel_rpc_port=data_parallel_rpc_port,

1064 enable_expert_parallel=self.enable_expert_parallel,

1065 max_parallel_loading_workers=self.max_parallel_loading_workers,

1066 disable_custom_all_reduce=self.disable_custom_all_reduce,

vllm/entrypoints/cli/serve.py

+79 -2

@@ -1,14 +1,24 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import argparse

4 +import signal

6 import uvloop

8 +import vllm.envs as envs

9 +from vllm import AsyncEngineArgs

10 from vllm.entrypoints.cli.types import CLISubcommand

11 from vllm.entrypoints.openai.api_server import run_server

12 from vllm.entrypoints.openai.cli_args import (make_arg_parser,

13 validate_parsed_serve_args)

11 -from vllm.utils import FlexibleArgumentParser

14 +from vllm.logger import init_logger

15 +from vllm.usage.usage_lib import UsageContext

16 +from vllm.utils import FlexibleArgumentParser, get_tcp_uri

17 +from vllm.v1.engine.core import EngineCoreProc

18 +from vllm.v1.engine.core_client import CoreEngineProcManager

19 +from vllm.v1.executor.abstract import Executor

20 +

21 +logger = init_logger(__name__)

24 class ServeSubcommand(CLISubcommand):

@@ -24,7 +34,10 @@ class ServeSubcommand(CLISubcommand):

34 if hasattr(args, 'model_tag') and args.model_tag is not None:

35 args.model = args.model_tag

27 - uvloop.run(run_server(args))

37 + if args.headless:

38 + run_headless(args)

39 + else:

40 + uvloop.run(run_server(args))

42 def validate(self, args: argparse.Namespace) -> None:

43 validate_parsed_serve_args(args)

@@ -42,6 +55,18 @@ class ServeSubcommand(CLISubcommand):

55 nargs='?',

56 help="The model tag to serve "

57 "(optional if specified in config)")

58 + serve_parser.add_argument(

59 + "--headless",

60 + action='store_true',

61 + default=False,

62 + help="Run in headless mode. See multi-node data parallel "

63 + "documentation for more details.")

64 + serve_parser.add_argument(

65 + '--data-parallel-start-rank',

66 + '-dpr',

67 + type=int,

68 + default=0,

69 + help='Starting data parallel rank for secondary nodes.')

70 serve_parser.add_argument(

71 "--config",

72 type=str,

@@ -57,3 +82,55 @@ class ServeSubcommand(CLISubcommand):

83 def cmd_init() -> list[CLISubcommand]:

84 return [ServeSubcommand()]

85 +

86 +

87 +def run_headless(args: argparse.Namespace):

88 +

89 + # Create the EngineConfig.

90 + engine_args = AsyncEngineArgs.from_cli_args(args)

91 + usage_context = UsageContext.OPENAI_API_SERVER

92 + vllm_config = engine_args.create_engine_config(usage_context=usage_context)

93 +

94 + if not envs.VLLM_USE_V1:

95 + raise RuntimeError("Headless mode is only supported for V1")

96 +

97 + parallel_config = vllm_config.parallel_config

98 + local_engine_count = parallel_config.data_parallel_size_local

99 + host = parallel_config.data_parallel_master_ip

100 + port = engine_args.data_parallel_rpc_port # add to config too

101 + input_address = get_tcp_uri(host, port)

102 +

103 + if local_engine_count <= 0:

104 + raise RuntimeError("data_parallel_size_local must be > 0 in "

105 + "headless mode")

106 +

107 + # Catch SIGTERM and SIGINT to allow graceful shutdown.

108 + def signal_handler(signum, frame):

109 + logger.debug("Received %d signal.", signum)

110 + raise SystemExit

111 +

112 + signal.signal(signal.SIGTERM, signal_handler)

113 + signal.signal(signal.SIGINT, signal_handler)

114 +

115 + logger.info(

116 + "Launching %d data parallel engine(s) in headless mode, "

117 + "with head node address %s.", local_engine_count, input_address)

118 +

119 + # Create the engines.

120 + engine_manager = CoreEngineProcManager(

121 + target_fn=EngineCoreProc.run_engine_core,

122 + local_engine_count=local_engine_count,

123 + start_index=args.data_parallel_start_rank,

124 + local_start_index=0,

125 + vllm_config=vllm_config,

126 + on_head_node=False,

127 + input_address=input_address,

128 + executor_class=Executor.get_class(vllm_config),

129 + log_stats=not engine_args.disable_log_stats,

130 + )

131 +

132 + try:

133 + engine_manager.join_first()

134 + finally:

135 + logger.info("Shutting down.")

136 + engine_manager.close()

vllm/utils.py

@@ -613,6 +613,10 @@ def is_valid_ipv6_address(address: str) -> bool:

613

614

615 def get_distributed_init_method(ip: str, port: int) -> str:

616 + return get_tcp_uri(ip, port)

617 +

618 +

619 +def get_tcp_uri(ip: str, port: int) -> str:

620 # Brackets are not permitted in ipv4 addresses,

621 # see https://github.com/python/cpython/issues/103848

622 return f"tcp://[{ip}]:{port}" if ":" in ip else f"tcp://{ip}:{port}"

vllm/v1/engine/core.py

+122 -67

@@ -1,5 +1,4 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -import json

2 import os

3 import queue

4 import signal

@@ -23,7 +22,7 @@ from vllm.logging_utils.dump_input import dump_engine_exception

22 from vllm.lora.request import LoRARequest

23 from vllm.transformers_utils.config import (

24 maybe_register_config_serialize_by_value)

26 -from vllm.utils import resolve_obj_by_qualname, zmq_socket_ctx

25 +from vllm.utils import make_zmq_socket, resolve_obj_by_qualname, zmq_socket_ctx

26 from vllm.v1.core.kv_cache_utils import (get_kv_cache_config,

27 unify_kv_cache_configs)

28 from vllm.v1.core.sched.interface import SchedulerInterface

@@ -43,6 +42,7 @@ from vllm.version import __version__ as VLLM_VERSION

42 logger = init_logger(__name__)

44 POLLING_TIMEOUT_S = 2.5

45 +HANDSHAKE_TIMEOUT_MINS = 5

47 _R = TypeVar('_R') # Return type for collective_rpc

@@ -348,9 +348,9 @@ class EngineCoreProc(EngineCore):

348

349 def __init__(

350 self,

351 - input_path: str,

352 - output_path: str,

351 vllm_config: VllmConfig,

352 + on_head_node: bool,

353 + input_address: str,

354 executor_class: type[Executor],

355 log_stats: bool,

356 engine_index: int = 0,

@@ -360,28 +360,91 @@ class EngineCoreProc(EngineCore):

360 executor_fail_callback = lambda: input_queue.put_nowait(

361 (EngineCoreRequestType.EXECUTOR_FAILED, b''))

362

363 - super().__init__(vllm_config, executor_class, log_stats,

364 - executor_fail_callback)

365 -

366 - self.step_fn = (self.step if self.batch_queue is None else

367 - self.step_with_batch_queue)

368 - self.engines_running = False

369 -

370 - # Background Threads and Queues for IO. These enable us to

371 - # overlap ZMQ socket IO with GPU since they release the GIL,

372 - # and to overlap some serialization/deserialization with the

373 - # model forward pass.

374 - # Threads handle Socket <-> Queues and core_busy_loop uses Queue.

375 - self.input_queue = input_queue

376 - self.output_queue = queue.Queue[Union[EngineCoreOutputs, bytes]]()

377 - threading.Thread(target=self.process_input_socket,

378 - args=(input_path, engine_index),

379 - daemon=True).start()

380 - self.output_thread = threading.Thread(

381 - target=self.process_output_socket,

382 - args=(output_path, engine_index),

383 - daemon=True)

384 - self.output_thread.start()

363 + # Create input socket.

364 + input_ctx = zmq.Context()

365 + identity = engine_index.to_bytes(length=2, byteorder="little")

366 + input_socket = make_zmq_socket(input_ctx,

367 + input_address,

368 + zmq.DEALER,

369 + identity=identity,

370 + bind=False)

371 + try:

372 + # Register engine with front-end.

373 + output_address = self.startup_handshake(

374 + input_socket, on_head_node, vllm_config.parallel_config)

375 +

376 + # Update config which may have changed from the handshake.

377 + vllm_config.__post_init__()

378 +

379 + # Set up data parallel environment.

380 + self._init_data_parallel(vllm_config)

381 +

382 + # Initialize engine core and model.

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

384 + executor_fail_callback)

385 +

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

387 + self.step_with_batch_queue)

388 + self.engines_running = False

389 +

390 + # Send ready message.

391 + num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks

392 + input_socket.send(

393 + msgspec.msgpack.encode({

394 + "status": "READY",

395 + "local": on_head_node,

396 + "num_gpu_blocks": num_gpu_blocks,

397 + }))

398 +

399 + # Background Threads and Queues for IO. These enable us to

400 + # overlap ZMQ socket IO with GPU since they release the GIL,

401 + # and to overlap some serialization/deserialization with the

402 + # model forward pass.

403 + # Threads handle Socket <-> Queues and core_busy_loop uses Queue.

404 + self.input_queue = input_queue

405 + self.output_queue = queue.Queue[Union[EngineCoreOutputs, bytes]]()

406 + threading.Thread(target=self.process_input_socket,

407 + args=(input_socket, ),

408 + daemon=True).start()

409 + input_socket = None

410 + self.output_thread = threading.Thread(

411 + target=self.process_output_socket,

412 + args=(output_address, engine_index),

413 + daemon=True)

414 + self.output_thread.start()

415 + finally:

416 + if input_socket is not None:

417 + input_socket.close(linger=0)

418 +

419 + @staticmethod

420 + def startup_handshake(input_socket: zmq.Socket, on_head_node: bool,

421 + parallel_config: ParallelConfig) -> str:

422 +

423 + # Send registration message.

424 + input_socket.send(

425 + msgspec.msgpack.encode({

426 + "status": "HELLO",

427 + "local": on_head_node,

428 + }))

429 +

430 + # Receive initialization message.

431 + logger.info("Waiting for init message from front-end.")

432 + if not input_socket.poll(timeout=HANDSHAKE_TIMEOUT_MINS * 60 * 1000):

433 + raise RuntimeError("Did not receive response from front-end "

434 + f"process within {HANDSHAKE_TIMEOUT_MINS} "

435 + f"minutes")

436 + init_bytes = input_socket.recv()

437 + init_message = msgspec.msgpack.decode(init_bytes)

438 + logger.debug("Received init message: %s", init_message)

439 +

440 + output_socket_address = init_message["output_socket_address"]

441 + #TBD(nick) maybe replace IP with configured head node address

442 +

443 + received_parallel_config = init_message["parallel_config"]

444 + for key, value in received_parallel_config.items():

445 + setattr(parallel_config, key, value)

446 +

447 + return output_socket_address

448

449 @staticmethod

450 def run_engine_core(*args,

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

475 try:

476 parallel_config: ParallelConfig = kwargs[

477 "vllm_config"].parallel_config

415 - if parallel_config.data_parallel_size > 1:

478 + if parallel_config.data_parallel_size > 1 or dp_rank > 0:

479 # Set data parallel rank for this engine process.

480 parallel_config.data_parallel_rank = dp_rank

481 parallel_config.data_parallel_rank_local = local_dp_rank

@@ -436,6 +499,9 @@ class EngineCoreProc(EngineCore):

499 if engine_core is not None:

500 engine_core.shutdown()

501

502 + def _init_data_parallel(self, vllm_config: VllmConfig):

503 + pass

504 +

505 def run_busy_loop(self):

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

507

@@ -527,40 +593,25 @@ class EngineCoreProc(EngineCore):

593 logger.fatal("vLLM shutdown signal from EngineCore failed "

594 "to send. Please report this issue.")

595

530 - def process_input_socket(self, input_path: str, engine_index: int):

596 + def process_input_socket(self, input_socket: zmq.Socket):

597 """Input socket IO thread."""

598

599 # Msgpack serialization decoding.

600 add_request_decoder = MsgpackDecoder(EngineCoreRequest)

601 generic_decoder = MsgpackDecoder()

536 - identity = engine_index.to_bytes(length=2, byteorder="little")

602

538 - with zmq_socket_ctx(input_path,

539 - zmq.DEALER,

540 - identity=identity,

541 - bind=False) as socket:

542 -

543 - # Send ready message to front-end once input socket is connected.

544 - message_dict = {

545 - 'type': 'READY',

546 - 'num_gpu_blocks': self.vllm_config.cache_config.num_gpu_blocks,

547 - }

548 - message = json.dumps(message_dict).encode('utf-8')

549 - socket.send(message)

550 -

551 - while True:

552 - # (RequestType, RequestData)

553 - type_frame, *data_frames = socket.recv_multipart(copy=False)

554 - request_type = EngineCoreRequestType(bytes(type_frame.buffer))

603 + while True:

604 + # (RequestType, RequestData)

605 + type_frame, *data_frames = input_socket.recv_multipart(copy=False)

606 + request_type = EngineCoreRequestType(bytes(type_frame.buffer))

607

556 - # Deserialize the request data.

557 - decoder = add_request_decoder if (

558 - request_type

559 - == EngineCoreRequestType.ADD) else generic_decoder

560 - request = decoder.decode(data_frames)

608 + # Deserialize the request data.

609 + decoder = add_request_decoder if (

610 + request_type == EngineCoreRequestType.ADD) else generic_decoder

611 + request = decoder.decode(data_frames)

612

562 - # Push to input queue for core busy loop.

563 - self.input_queue.put_nowait((request_type, request))

613 + # Push to input queue for core busy loop.

614 + self.input_queue.put_nowait((request_type, request))

615

616 def process_output_socket(self, output_path: str, engine_index: int):

617 """Output socket IO thread."""

@@ -609,9 +660,9 @@ class DPEngineCoreProc(EngineCoreProc):

660

661 def __init__(

662 self,

612 - input_path: str,

613 - output_path: str,

663 vllm_config: VllmConfig,

664 + on_head_node: bool,

665 + input_address: str,

666 executor_class: type[Executor],

667 log_stats: bool,

668 ):

@@ -623,8 +674,20 @@ class DPEngineCoreProc(EngineCoreProc):

674 _add_prefix(sys.stdout, process_name, pid)

675 _add_prefix(sys.stderr, process_name, pid)

676

626 - dp_size = vllm_config.parallel_config.data_parallel_size

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

678 + # finished with DP peers every N steps.

679 + self.counter = 0

680 +

681 + # Initialize the engine.

682 + dp_rank = vllm_config.parallel_config.data_parallel_rank

683 + super().__init__(vllm_config, on_head_node, input_address,

684 + executor_class, log_stats, dp_rank)

685 +

686 + def _init_data_parallel(self, vllm_config: VllmConfig):

687 +

688 + # Configure GPUs and stateless process group for data parallel.

689 dp_rank = vllm_config.parallel_config.data_parallel_rank

690 + dp_size = vllm_config.parallel_config.data_parallel_size

691 local_dp_rank = vllm_config.parallel_config.data_parallel_rank_local

692

693 assert dp_size > 1

@@ -632,24 +695,16 @@ class DPEngineCoreProc(EngineCoreProc):

695

696 from vllm.platforms import current_platform

697 device_control_env_var = current_platform.device_control_env_var

635 - tp_size = vllm_config.parallel_config.tensor_parallel_size

698 + world_size = vllm_config.parallel_config.world_size

699 os.environ[device_control_env_var] = ",".join(

700 str(current_platform.device_id_to_physical_device_id(i))

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

639 - tp_size))

701 + for i in range(local_dp_rank * world_size, (local_dp_rank + 1) *

702 + world_size))

703

704 self.local_dp_rank = local_dp_rank

705 self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()

706 self.current_wave = 0

707

645 - # Initialize the engine after setting up environment.

646 - super().__init__(input_path, output_path, vllm_config, executor_class,

647 - log_stats, dp_rank)

648 -

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

650 - # finished with DP peers every N steps.

651 - self.counter = 0

652 -

708 def shutdown(self):

709 super().shutdown()

710 if dp_group := getattr(self, "dp_group", None):

vllm/v1/engine/core_client.py

+174 -119

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 import asyncio

3 import contextlib

4 -import json

4 import queue

5 import uuid

6 import weakref

@@ -9,25 +8,27 @@ from abc import ABC, abstractmethod

8 from collections import deque

9 from collections.abc import Awaitable, Sequence

10 from concurrent.futures import Future

12 -from dataclasses import dataclass, field

11 +from dataclasses import dataclass

12 +from enum import Enum, auto

13 from threading import Thread

14 from typing import Any, Callable, Optional, TypeVar, Union

16 +import msgspec

17 import zmq

18 import zmq.asyncio

19 -from vllm.config import VllmConfig

20 +from vllm.config import ParallelConfig, VllmConfig

21 from vllm.logger import init_logger

22 from vllm.lora.request import LoRARequest

22 -from vllm.utils import (get_open_zmq_inproc_path, get_open_zmq_ipc_path,

23 - make_zmq_socket)

23 +from vllm.utils import (get_open_port, get_open_zmq_inproc_path,

24 + get_open_zmq_ipc_path, get_tcp_uri, make_zmq_socket)

25 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,

26 EngineCoreRequestType, UtilityOutput)

27 from vllm.v1.engine.core import EngineCore, EngineCoreProc

28 from vllm.v1.engine.exceptions import EngineDeadError

29 from vllm.v1.executor.abstract import Executor

30 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr

30 -from vllm.v1.utils import BackgroundProcHandle

31 +from vllm.v1.utils import CoreEngineProcManager

33 logger = init_logger(__name__)

@@ -264,45 +265,22 @@ class InprocClient(EngineCoreClient):

265 return self.engine_core.collective_rpc(method, timeout, args, kwargs)

266

267

268 +class CoreEngineState(Enum):

269 + NEW = auto()

270 + CONNECTED = auto()

271 + READY = auto()

272 +

273 +

274 class CoreEngine:

275 """One per data parallel rank."""

276

270 - def __init__(

271 - self,

272 - vllm_config: VllmConfig,

273 - executor_class: type[Executor],

274 - log_stats: bool,

275 - input_path: str,

276 - output_path: str,

277 - index: int = 0,

278 - local_dp_rank: int = 0,

279 - ):

277 + def __init__(self, index: int = 0, local: bool = True):

278 + self.local = local

279 self.index = index

280 self.identity = index.to_bytes(length=2, byteorder="little")

282 - try:

283 - # Start EngineCore in background process.

284 - self.proc_handle = BackgroundProcHandle(

285 - input_path=input_path,

286 - output_path=output_path,

287 - process_name=f"EngineCore_{index}",

288 - target_fn=EngineCoreProc.run_engine_core,

289 - process_kwargs={

290 - "vllm_config": vllm_config,

291 - "dp_rank": index,

292 - "local_dp_rank": local_dp_rank,

293 - "executor_class": executor_class,

294 - "log_stats": log_stats,

295 - })

296 -

297 - self.num_reqs_in_flight = 0

298 - finally:

299 - if not hasattr(self, "num_reqs_in_flight"):

300 - # Ensure socket is closed if process fails to start.

301 - self.close()

281

303 - def close(self):

304 - if proc_handle := getattr(self, "proc_handle", None):

305 - proc_handle.shutdown()

282 + self.state = CoreEngineState.NEW

283 + self.num_reqs_in_flight = 0

284

285

286 @dataclass

@@ -311,7 +289,7 @@ class BackgroundResources:

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

290

291 ctx: Union[zmq.Context]

314 - core_engines: list[CoreEngine] = field(default_factory=list)

292 + local_engine_manager: Optional[CoreEngineProcManager] = None

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

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

295 output_queue_task: Optional[asyncio.Task] = None

@@ -325,8 +303,8 @@ class BackgroundResources:

303 """Clean up background resources."""

304

305 self.engine_dead = True

328 - for core_engine in self.core_engines:

329 - core_engine.close()

306 + if self.local_engine_manager is not None:

307 + self.local_engine_manager.close()

308

309 if self.output_queue_task is not None:

310 self.output_queue_task.cancel()

@@ -388,25 +366,56 @@ class MPClient(EngineCoreClient):

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

367 success = False

368 try:

391 - # Paths and sockets for IPC.

392 - self.output_path = get_open_zmq_ipc_path()

393 - input_path = get_open_zmq_ipc_path()

394 - self.input_socket = make_zmq_socket(self.ctx,

395 - input_path,

396 - zmq.ROUTER,

397 - bind=True)

398 - self.resources.input_socket = self.input_socket

399 -

400 - new_core_engine = lambda index, local_dp_rank=None: CoreEngine(

401 - vllm_config, executor_class, log_stats, input_path, self.

402 - output_path, index, local_dp_rank)

403 -

404 - # Start engine core process(es).

405 - self._init_core_engines(vllm_config, new_core_engine,

406 - self.resources.core_engines)

369 + parallel_config = vllm_config.parallel_config

370 + local_engine_count = parallel_config.data_parallel_size_local

371 + start_index = parallel_config.data_parallel_rank

372 + local_start_index = parallel_config.data_parallel_rank_local

373 +

374 + # SPMD mode is where there is an LLM instance per DP rank and

375 + # one core engine per LLM, see

376 + # examples/offline_inference/data_parallel.py.

377 + spmd_mode = local_start_index is not None

378 + if spmd_mode:

379 + assert local_engine_count == 1

380 + self.core_engines = [

381 + CoreEngine(index=local_start_index, local=True)

382 + ]

383 + else:

384 + assert start_index == 0

385 + local_start_index = 0

386 + self.core_engines = [

387 + CoreEngine(index=i, local=(i < local_engine_count))

388 + for i in range(parallel_config.data_parallel_size)

389 + ]

390 +

391 + input_address, output_address = self._get_zmq_addresses(

392 + parallel_config, spmd_mode)

393 +

394 + # Create input and output sockets.

395 + self.input_socket = self.resources.input_socket = make_zmq_socket(

396 + self.ctx, input_address, zmq.ROUTER, bind=True)

397 +

398 + self.resources.output_socket = make_zmq_socket(

399 + self.ctx, output_address, zmq.constants.PULL)

400 + # Start local engines.

401 + if local_engine_count:

402 + # In server mode, start_index and local_start_index will

403 + # both be 0.

404 + self.resources.local_engine_manager = CoreEngineProcManager(

405 + EngineCoreProc.run_engine_core,

406 + vllm_config=vllm_config,

407 + executor_class=executor_class,

408 + log_stats=log_stats,

409 + input_address=input_address,

410 + on_head_node=True,

411 + local_engine_count=local_engine_count,

412 + start_index=start_index,

413 + local_start_index=local_start_index)

414 +

415 + self.core_engine = self.core_engines[0]

416

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

409 - self._wait_for_engine_startup()

418 + self._wait_for_engine_startup(output_address, parallel_config)

419

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

421

@@ -420,56 +429,116 @@ class MPClient(EngineCoreClient):

429 if not success:

430 self._finalizer()

431

423 - def _wait_for_engine_startup(self):

432 + @staticmethod

433 + def _get_zmq_addresses(parallel_config: ParallelConfig,

434 + spmd_mode: bool) -> tuple[str, str]:

435 + """Returns (input_address, output_address)."""

436 + dp_size = parallel_config.data_parallel_size

437 + local_engine_count = parallel_config.data_parallel_size_local

438 +

439 + if local_engine_count == dp_size or spmd_mode:

440 + input_address = get_open_zmq_ipc_path()

441 + output_address = get_open_zmq_ipc_path()

442 + else:

443 + host = parallel_config.data_parallel_master_ip

444 + input_port = parallel_config.data_parallel_rpc_port

445 + output_port = get_open_port()

446 + input_address = get_tcp_uri(host, input_port)

447 + output_address = get_tcp_uri(host, output_port)

448 +

449 + return input_address, output_address

450 +

451 + def _wait_for_engine_startup(self, output_address: str,

452 + parallel_config: ParallelConfig):

453 # Get a sync handle to the socket which can be sync or async.

454 sync_input_socket = zmq.Socket.shadow(self.input_socket)

455

456 # Wait for engine core process(es) to send ready messages.

428 - identities = set(eng.index for eng in self.resources.core_engines)

457 + local_count = parallel_config.data_parallel_size_local

458 + remote_count = len(self.core_engines) - local_count

459 + # [local, remote] counts

460 + conn_pending, start_pending = [local_count, remote_count], [0, 0]

461 +

462 poller = zmq.Poller()

463 poller.register(sync_input_socket, zmq.POLLIN)

431 - for eng in self.resources.core_engines:

432 - poller.register(eng.proc_handle, zmq.POLLIN)

433 - while identities:

464 + proc_manager = self.resources.local_engine_manager

465 + if proc_manager is not None:

466 + for sentinel in proc_manager.sentinels():

467 + poller.register(sentinel, zmq.POLLIN)

468 + while any(conn_pending) or any(start_pending):

469 events = poller.poll(STARTUP_POLL_PERIOD_MS)

470 if not events:

436 - logger.debug("Waiting for %d core engine proc(s) to start: %s",

437 - len(identities), identities)

471 + if any(conn_pending):

472 + logger.debug(

473 + "Waiting for %d local, %d remote core engine proc(s) "

474 + "to connect.", *conn_pending)

475 + if any(start_pending):

476 + logger.debug(

477 + "Waiting for %d local, %d remote core engine proc(s) "

478 + "to start.", *start_pending)

479 continue

480 if len(events) > 1 or events[0][0] != sync_input_socket:

440 - # One of the core processes exited.

481 + # One of the local core processes exited.

482 + finished = proc_manager.finished_procs(

483 + ) if proc_manager else {}

484 raise RuntimeError("Engine core initialization failed. "

442 - "See root cause above.")

443 -

444 - eng_id_bytes, data = sync_input_socket.recv_multipart()

445 - eng_id = int.from_bytes(eng_id_bytes, byteorder="little")

446 - if eng_id not in identities:

447 - raise RuntimeError(f"Unexpected or duplicate engine: {eng_id}")

448 - message_dict = json.loads(data.decode('utf-8'))

449 - if message_dict['type'] != 'READY':

450 - raise RuntimeError(f"Engine {eng_id} failed: {data.decode()}")

451 - logger.info("Core engine process %d ready.", eng_id)

452 - identities.discard(eng_id)

453 - # Setup KV cache config with initialization state from

454 - # engine core process. Sum values from all engines in DP case.

455 - num_gpu_blocks = self.vllm_config.cache_config.num_gpu_blocks or 0

456 - num_gpu_blocks += message_dict['num_gpu_blocks']

457 - self.vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks

458 -

459 - def _init_core_engines(

460 - self,

461 - vllm_config: VllmConfig,

462 - new_core_engine: Callable[[int, Optional[int]], CoreEngine],

463 - core_engines: list[CoreEngine],

464 - ) -> None:

465 -

466 - # Default case - single core engine.

467 - core_engine = new_core_engine(

468 - vllm_config.parallel_config.data_parallel_rank,

469 - vllm_config.parallel_config.data_parallel_rank_local,

470 - )

471 - core_engines.append(core_engine)

472 - self.core_engine = core_engine

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 = sync_input_socket.recv_multipart()

490 + eng_index = int.from_bytes(eng_identity, byteorder="little")

491 + engine = next(

492 + (e for e in self.core_engines if e.identity == eng_identity),

493 + None)

494 + if engine is None:

495 + raise RuntimeError(f"Message from engine with unexpected data "

496 + f"parallel rank: {eng_index}")

497 + msg = msgspec.msgpack.decode(ready_msg_bytes)

498 + status, local = msg["status"], msg["local"]

499 + if local != engine.local:

500 + raise RuntimeError(f"{status} message from "

501 + f"{'local' if local else 'remote'} "

502 + f"engine {eng_index}, expected it to be "

503 + f"{'local' if engine.local else 'remote'}")

504 +

505 + if status == "HELLO" and engine.state == CoreEngineState.NEW:

506 +

507 + # Send init message with DP config info.

508 + init_message = self.encoder.encode({

509 + "output_socket_address": output_address,

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 + })

519 + sync_input_socket.send_multipart((eng_identity, *init_message),

520 + copy=False)

521 + conn_pending[0 if local else 1] -= 1

522 + start_pending[0 if local else 1] += 1

523 + engine.state = CoreEngineState.CONNECTED

524 + elif status == "READY" and (engine.state

525 + == CoreEngineState.CONNECTED):

526 + # Setup KV cache config with initialization state from

527 + # engine core process. Sum values from all engines in DP case.

528 + cache_config = self.vllm_config.cache_config

529 + num_gpu_blocks = cache_config.num_gpu_blocks or 0

530 + num_gpu_blocks += msg['num_gpu_blocks']

531 + cache_config.num_gpu_blocks = num_gpu_blocks

532 +

533 + start_pending[0 if local else 1] -= 1

534 + engine.state = CoreEngineState.READY

535 + else:

536 + raise RuntimeError(f"Unexpected {status} message for "

537 + f"{'local' if local else 'remote'} engine "

538 + f"{eng_index} in {engine.state} state.")

539 +

540 + logger.debug("%s from %s core engine process %s.", status,

541 + "local" if local else "remote", eng_index)

542

543 def shutdown(self):

544 # Terminate background resources.

@@ -520,7 +589,8 @@ class SyncMPClient(MPClient):

589 # Ensure that the outputs socket processing thread does not have

590 # a ref to the client which prevents gc.

591 ctx = self.ctx

523 - output_path = self.output_path

592 + out_socket = self.resources.output_socket

593 + assert out_socket is not None

594 decoder = self.decoder

595 utility_results = self.utility_results

596 outputs_queue = self.outputs_queue

@@ -531,7 +601,6 @@ class SyncMPClient(MPClient):

601

602 def process_outputs_socket():

603 shutdown_socket = ctx.socket(zmq.PAIR)

534 - out_socket = make_zmq_socket(ctx, output_path, zmq.constants.PULL)

604 try:

605 shutdown_socket.bind(shutdown_path)

606 poller = zmq.Poller()

@@ -566,6 +635,9 @@ class SyncMPClient(MPClient):

635 daemon=True)

636 self.output_queue_thread.start()

637

638 + # The thread takes on responsibility for closing the socket.

639 + self.resources.output_socket = None

640 +

641 def get_output(self) -> EngineCoreOutputs:

642 # If an exception arises in process_outputs_socket task,

643 # it is forwarded to the outputs_queue so we can raise it

@@ -693,10 +765,8 @@ class AsyncMPClient(MPClient):

765 self.__class__,

766 "process_engine_outputs", None)

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

696 - output_path = self.output_path

697 - output_socket = make_zmq_socket(self.ctx, output_path,

698 - zmq.constants.PULL)

699 - resources.output_socket = output_socket

768 + output_socket = resources.output_socket

769 + assert output_socket is not None

770

771 async def process_outputs_socket():

772 try:

@@ -861,21 +931,6 @@ class DPAsyncMPClient(AsyncMPClient):

931

932 assert len(self.core_engines) > 1

933

864 - def _init_core_engines(

865 - self,

866 - vllm_config: VllmConfig,

867 - new_core_engine: Callable[[int, Optional[int]], CoreEngine],

868 - core_engines: list[CoreEngine],

869 - ) -> None:

870 -

871 - # Launch a core engine for each data parallel rank.

872 - dp_size = vllm_config.parallel_config.data_parallel_size

873 - for i in range(dp_size):

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

875 - core_engines.append(new_core_engine(i, i))

876 -

877 - self.core_engines = core_engines

878 -

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

935 # Only the result from the first engine is returned.

936 return (await asyncio.gather(*[

vllm/v1/utils.py

+78 -33

@@ -1,20 +1,23 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import os

4 +import time

5 import weakref

6 from collections import defaultdict

7 from collections.abc import Sequence

7 -from multiprocessing import Process

8 -from typing import (TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar,

9 - Union, overload)

8 +from multiprocessing import Process, connection

9 +from typing import (TYPE_CHECKING, Callable, Generic, Optional, TypeVar, Union,

10 + overload)

12 import torch

14 +from vllm.config import VllmConfig

15 from vllm.logger import init_logger

16 from vllm.model_executor.models.utils import extract_layer_index

17 from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled,

18 usage_message)

19 from vllm.utils import get_mp_context, kill_process_tree

20 +from vllm.v1.executor.abstract import Executor

22 if TYPE_CHECKING:

23 from vllm.attention.layer import Attention

@@ -92,7 +95,7 @@ class ConstantList(Generic[T], Sequence):

95 return f"ConstantList({self._x})"

95 -class BackgroundProcHandle:

98 +class CoreEngineProcManager:

99 """

100 Utility class to handle creation, readiness, and shutdown

101 of background processes used by the AsyncLLM and LLMEngine.

@@ -100,49 +103,91 @@ class BackgroundProcHandle:

103

104 def __init__(

105 self,

103 - input_path: str,

104 - output_path: str,

105 - process_name: str,

106 target_fn: Callable,

107 - process_kwargs: dict[Any, Any],

107 + local_engine_count: int,

108 + start_index: int,

109 + local_start_index: int,

110 + vllm_config: VllmConfig,

111 + on_head_node: bool,

112 + input_address: str,

113 + executor_class: type[Executor],

114 + log_stats: bool,

115 ):

116 context = get_mp_context()

117 + common_kwargs = {

118 + "vllm_config": vllm_config,

119 + "on_head_node": on_head_node,

120 + "input_address": input_address,

121 + "executor_class": executor_class,

122 + "log_stats": log_stats,

123 + }

124 +

125 + self.processes: list[Process] = []

126 + for index in range(local_engine_count):

127 + local_index = local_start_index + index

128 + global_index = start_index + index

129 + # Start EngineCore in background process.

130 + self.processes.append(

131 + context.Process(target=target_fn,

132 + name=f"EngineCore_{global_index}",

133 + kwargs=common_kwargs | {

134 + "dp_rank": global_index,

135 + "local_dp_rank": local_index,

136 + }))

137 +

138 + self._finalizer = weakref.finalize(self, shutdown, self.processes,

139 + input_address)

140 + try:

141 + for proc in self.processes:

142 + proc.start()

143 + finally:

144 + # Kill other procs if not all are running.

145 + if self.finished_procs():

146 + self.close()

147 +

148 + def close(self):

149 + """Shutdown all procs."""

150 + self._finalizer()

151

111 - assert ("input_path" not in process_kwargs

112 - and "output_path" not in process_kwargs)

113 - process_kwargs["input_path"] = input_path

114 - process_kwargs["output_path"] = output_path

115 -

116 - # Run busy loop in background process.

117 - self.proc: Process = context.Process(target=target_fn,

118 - kwargs=process_kwargs,

119 - name=process_name)

120 - self._finalizer = weakref.finalize(self, shutdown, self.proc,

121 - input_path, output_path)

122 - self.proc.start()

152 + def join_first(self):

153 + """Wait for any process to exit."""

154 + connection.wait(proc.sentinel for proc in self.processes)

155

124 - def fileno(self):

125 - return self.proc.sentinel

156 + def sentinels(self) -> list:

157 + return [proc.sentinel for proc in self.processes]

158

127 - def shutdown(self):

128 - self._finalizer()

159 + def finished_procs(self) -> dict[str, int]:

160 + """Returns dict of proc name -> exit code for any finished procs."""

161 + return {

162 + proc.name: proc.exitcode

163 + for proc in self.processes if proc.exitcode is not None

164 + }

165

166

167 # Note(rob): shutdown function cannot be a bound method,

132 -# else the gc cannot collect the object.

133 -def shutdown(proc: Process, input_path: str, output_path: str):

168 +# else the gc cannot collect the objedecoupct.

169 +def shutdown(procs: list[Process], input_address: str):

170 # Shutdown the process.

135 - if proc.is_alive():

136 - proc.terminate()

137 - proc.join(5)

138 -

171 + for proc in procs:

172 + if proc.is_alive():

173 + proc.terminate()

174 +

175 + # Allow 5 seconds for remaining procs to terminate.

176 + deadline = time.monotonic() + 5

177 + for proc in procs:

178 + remaining = deadline - time.monotonic()

179 + if remaining <= 0:

180 + break

181 + if proc.is_alive():

182 + proc.join(remaining)

183 +

184 + for proc in procs:

185 if proc.is_alive() and (pid := proc.pid) is not None:

186 kill_process_tree(pid)

187

188 # Remove zmq ipc socket files.

143 - ipc_sockets = [output_path, input_path]

144 - for ipc_socket in ipc_sockets:

145 - socket_file = ipc_socket.replace("ipc://", "")

189 + if input_address.startswith("ipc://"):

190 + socket_file = input_address[len("ipc://"):]

191 if os and os.path.exists(socket_file):

192 os.remove(socket_file)

193