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

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

elastic_ep: Fix stateless group port races (#36330)

Signed-off-by: Itay Alroy <ialroy@nvidia.com>

Itay Alroy committed Mar 18, 2026 at 16:36 UTC de1a86b7dea68dffdfbeda77c8407c587ee90542

12 files changed +224 -225

.buildkite/test_areas/expert_parallelism.yaml

+1 -2

@@ -24,8 +24,7 @@ steps:

25 - label: Elastic EP Scaling Test

26 timeout_in_minutes: 20

27 - device: b200

28 - optional: true

27 + device: h100

28 working_dir: "/vllm-workspace/tests"

29 num_devices: 4

30 source_file_dependencies:

vllm/config/parallel.py

+31 -85

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

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import os

5 +import socket

6 from collections.abc import Callable

7 from typing import TYPE_CHECKING, Any, Literal, overload

@@ -266,33 +267,9 @@ class ParallelConfig:

267 Set to be private as it's not intended to be configured by users.

268 """

269

269 - _stateless_dp_group_port_list: list[list[int]] = Field(default_factory=list)

270 - """List of open ports for stateless DP groups when enable_elastic_ep is True.

271 - Set to be private as it's not intended to be configured by users.

272 - It is a list of list[int], with each inner list contains a set of 3 ports

273 - to be used for setting up the stateless CPU/device/TCPStore groups

274 - in StatelessGroupCoordinator. The number of inner lists is equal to

275 - the number of DP groups,

276 - i.e., len(self._stateless_dp_group_port_list) == world_size_across_dp // dp_size,

277 - and len(self._stateless_dp_group_port_list[i]) == 3 for all i.

278 - """

279 -

280 - _stateless_ep_group_port_list: list[list[int]] = Field(default_factory=list)

281 - """List of open ports for stateless EP groups when enable_elastic_ep is True.

282 - Set to be private as it's not intended to be configured by users.

283 - len(self._stateless_ep_group_port_list) == world_size_across_dp // ep_size,

284 - """

285 -

286 - _stateless_eplb_group_port_list: list[list[int]] = Field(default_factory=list)

287 - """List of open ports for stateless EPLB groups when enable_elastic_ep is True.

288 - Same topology as EP but separate NCCL communicator to avoid deadlocks.

289 - """

290 -

291 - _stateless_world_group_port_list: list[list[int]] = Field(default_factory=list)

292 - """List of open ports for stateless world group when enable_elastic_ep is True.

293 - Set to be private as it's not intended to be configured by users.

294 - len(self._stateless_world_group_port_list) == 1,

295 - """

270 + _coord_store_port: int = 0

271 + """Port of the coordination TCPStore. Can be set by the API server; workers

272 + connect as clients to exchange self-picked group ports at runtime."""

273

274 decode_context_parallel_size: int = 1

275 """Number of decode context parallel groups, because the world size does

@@ -465,65 +442,32 @@ class ParallelConfig:

442

443 return answer

444

468 - def allocate_elastic_ep_ports(self) -> None:

469 - """Allocate all ports for elastic EP (stateless groups + DP master).

445 + def _pick_stateless_dp_port(self) -> tuple[int, socket.socket | None]:

446 + """Return ``(port, listen_socket)`` for DP group init.

447

471 - Must be called AFTER ray.init() so that ports claimed by Ray's

472 - idle worker pool are already in use and won't be returned by

473 - get_open_ports_list().

448 + With a coord store, rank 0 binds a socket and publishes the port;

449 + others read it. Without one, pops a pre-allocated port and

450 + returns ``listen_socket=None``.

451 """

475 - if not self.enable_elastic_ep:

476 - return

477 - if self._stateless_world_group_port_list:

478 - return

479 -

480 - num_world_groups = 1

481 - dp_size = self.data_parallel_size

482 - ep_size = self.data_parallel_size * self.world_size_across_dp

483 - num_dp_groups = max(1, self.world_size_across_dp // dp_size)

484 - num_ep_groups = max(1, self.world_size_across_dp // ep_size)

485 - num_eplb_groups = num_ep_groups

486 - total_stateless_ports = (

487 - num_world_groups + num_dp_groups + num_ep_groups + num_eplb_groups

488 - ) * 3

489 - num_dp_master_ports = 5

490 -

491 - all_ports = get_open_ports_list(total_stateless_ports + num_dp_master_ports)

492 -

493 - self._data_parallel_master_port_list = all_ports[-num_dp_master_ports:]

494 - self.data_parallel_master_port = self._data_parallel_master_port_list.pop()

495 - all_ports = all_ports[:-num_dp_master_ports]

496 -

497 - self._stateless_world_group_port_list = [

498 - all_ports[i : i + 3] for i in range(0, num_world_groups * 3, 3)

499 - ]

500 - start_idx = num_world_groups * 3

501 - self._stateless_dp_group_port_list = [

502 - all_ports[i : i + 3]

503 - for i in range(start_idx, start_idx + num_dp_groups * 3, 3)

504 - ]

505 - start_idx += num_dp_groups * 3

506 - self._stateless_ep_group_port_list = [

507 - all_ports[i : i + 3]

508 - for i in range(start_idx, start_idx + num_ep_groups * 3, 3)

509 - ]

510 - start_idx += num_ep_groups * 3

511 - self._stateless_eplb_group_port_list = [

512 - all_ports[i : i + 3]

513 - for i in range(start_idx, start_idx + num_eplb_groups * 3, 3)

514 - ]

515 -

516 - def get_next_stateless_world_group_port(self) -> list[int]:

517 - return self._stateless_world_group_port_list.pop()

518 -

519 - def get_next_stateless_dp_group_port(self) -> list[int]:

520 - return self._stateless_dp_group_port_list.pop()

521 -

522 - def get_next_stateless_ep_group_port(self) -> list[int]:

523 - return self._stateless_ep_group_port_list.pop()

524 -

525 - def get_next_stateless_eplb_group_port(self) -> list[int]:

526 - return self._stateless_eplb_group_port_list.pop()

452 + if not self._coord_store_port:

453 + return self.get_next_dp_init_port(), None

454 +

455 + from vllm.distributed.utils import get_cached_tcp_store_client

456 +

457 + store = get_cached_tcp_store_client(

458 + self.data_parallel_master_ip, self._coord_store_port

459 + )

460 +

461 + key = "dp_master_port"

462 + if self.data_parallel_rank == 0:

463 + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

464 + s.bind((self.data_parallel_master_ip, 0))

465 + s.listen()

466 + port = s.getsockname()[1]

467 + store.set(key, str(port).encode())

468 + return port, s

469 + else:

470 + return int(store.get(key).decode()), None

471

472 @overload

473 def stateless_init_dp_group(

@@ -553,14 +497,16 @@ class ParallelConfig:

497 last_exc: Exception | None = None

498 for _ in range(max_retries):

499 try:

500 + port, listen_socket = self._pick_stateless_dp_port()

501 # use gloo since the engine process might not have cuda device

502 return stateless_init_torch_distributed_process_group(

503 self.data_parallel_master_ip,

559 - self.get_next_dp_init_port(),

504 + port,

505 self.data_parallel_rank,

506 self.data_parallel_size,

507 backend="gloo",

508 return_store=return_store,

509 + listen_socket=listen_socket,

510 )

511 except DistNetworkError as e:

512 # We only want to retry when the root cause is EADDRINUSE.

vllm/distributed/elastic_ep/elastic_execute.py

+2 -4

@@ -162,10 +162,8 @@ class ElasticEPScalingExecutor:

162 new_dp_size=new_dp_size,

163 new_world_size_across_dp=new_world_size_across_dp,

164 master_ip=reconfig_request.new_data_parallel_master_ip,

165 - world_group_ports=reconfig_request.new_stateless_world_group_port_list,

166 - dp_group_ports=reconfig_request.new_stateless_dp_group_port_list,

167 - ep_group_ports=reconfig_request.new_stateless_ep_group_port_list,

168 - eplb_group_ports=reconfig_request.new_stateless_eplb_group_port_list,

165 + coord_store_port=reconfig_request.coord_store_port,

166 + enable_eplb=updated_config.parallel_config.enable_eplb,

167 )

168 self.worker.model_runner.eep_eplb_suppressed = True

169 standby_ep_group = get_standby_ep_group()

vllm/distributed/elastic_ep/elastic_state.py

+1 -12

@@ -563,15 +563,4 @@ class ElasticEPScalingState:

563 parallel_config._data_parallel_master_port_list = (

564 reconfig_request.new_data_parallel_master_port_list

565 )

566 - parallel_config._stateless_world_group_port_list = (

567 - reconfig_request.new_stateless_world_group_port_list

568 - )

569 - parallel_config._stateless_dp_group_port_list = (

570 - reconfig_request.new_stateless_dp_group_port_list

571 - )

572 - parallel_config._stateless_ep_group_port_list = (

573 - reconfig_request.new_stateless_ep_group_port_list

574 - )

575 - parallel_config._stateless_eplb_group_port_list = (

576 - reconfig_request.new_stateless_eplb_group_port_list

577 - )

566 + parallel_config._coord_store_port = reconfig_request.coord_store_port

vllm/distributed/elastic_ep/standby_state.py

+15 -9

@@ -38,10 +38,8 @@ def create_standby_groups(

38 new_dp_size: int,

39 new_world_size_across_dp: int,

40 master_ip: str,

41 - world_group_ports: list[list[int]],

42 - dp_group_ports: list[list[int]],

43 - ep_group_ports: list[list[int]],

44 - eplb_group_ports: list[list[int]] | None = None,

41 + coord_store_port: int,

42 + enable_eplb: bool = True,

43 backend: str | None = None,

44 ) -> None:

45 global \

@@ -51,19 +49,23 @@ def create_standby_groups(

49 _STANDBY_EP, \

50 _STANDBY_EPLB

52 + from vllm.distributed.utils import get_cached_tcp_store_client

53 +

54 assert new_world_size_across_dp == torch.distributed.get_world_size() * new_dp_size

55 world_group = get_world_group()

56 assert isinstance(world_group, StatelessGroupCoordinator)

57 backend = backend or world_group.backend

59 + coord_store = get_cached_tcp_store_client(master_ip, coord_store_port)

60 +

61 standby_world_ranks = [list(range(new_world_size_across_dp))]

62 _STANDBY_WORLD = _init_stateless_group(

63 standby_world_ranks,

64 "world",

63 - world_group_ports,

65 master_ip,

66 backend,

67 use_device_communicator=False,

68 + coord_store=coord_store,

69 )

70 _STANDBY_WORLD_NODE_COUNT = _node_count(_STANDBY_WORLD.tcp_store_group)

@@ -76,7 +78,7 @@ def create_standby_groups(

78 standby_dp_ranks = all_ranks.transpose(1, 3).reshape(-1, new_dp_size).unbind(0)

79 standby_dp_ranks = [x.tolist() for x in standby_dp_ranks]

80 _STANDBY_DP = _init_stateless_group(

79 - standby_dp_ranks, "dp", dp_group_ports, master_ip, backend

81 + standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store

82 )

84 standby_ep_ranks = (

@@ -84,12 +86,16 @@ def create_standby_groups(

86 )

87 standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]

88 _STANDBY_EP = _init_stateless_group(

87 - standby_ep_ranks, "ep", ep_group_ports, master_ip, backend

89 + standby_ep_ranks, "ep", master_ip, backend, coord_store=coord_store

90 )

90 - if eplb_group_ports is not None:

92 + if enable_eplb:

93 _STANDBY_EPLB = _init_stateless_group(

92 - standby_ep_ranks, "eplb", eplb_group_ports, master_ip, backend

94 + standby_ep_ranks,

95 + "eplb",

96 + master_ip,

97 + backend,

98 + coord_store=coord_store,

99 )

100

101

vllm/distributed/parallel_state.py

+20 -21

@@ -40,13 +40,16 @@ import torch

40 import torch.distributed

41 import torch.distributed._functional_collectives as funcol

42 import torch.distributed._symmetric_memory

43 -from torch.distributed import Backend, ProcessGroup

43 +from torch.distributed import Backend, ProcessGroup, Store

45 import vllm.envs as envs

46 from vllm.distributed.device_communicators.base_device_communicator import (

47 DeviceCommunicatorBase,

48 )

49 -from vllm.distributed.utils import StatelessProcessGroup

49 +from vllm.distributed.utils import (

50 + StatelessProcessGroup,

51 + get_cached_tcp_store_client,

52 +)

53 from vllm.logger import init_logger

54 from vllm.utils.import_utils import resolve_obj_by_qualname

55 from vllm.utils.network_utils import get_distributed_init_method

@@ -1164,9 +1167,9 @@ def init_model_parallel_group(

1167 def _init_stateless_group(

1168 group_ranks: list[list[int]],

1169 group_name: str,

1167 - group_ports: list[list[int]],

1170 host: str,

1171 backend: str,

1172 + coord_store: Store,

1173 use_device_communicator: bool = True,

1174 ) -> "StatelessGroupCoordinator":

1175 """Create a StatelessGroupCoordinator with the given parameters."""

@@ -1180,7 +1183,7 @@ def _init_stateless_group(

1183 use_device_communicator=use_device_communicator,

1184 group_name=group_name,

1185 host=host,

1183 - group_ports=group_ports,

1186 + coord_store=coord_store,

1187 global_rank=world.rank,

1188 global_world_size=world.world_size,

1189 )

@@ -1321,7 +1324,9 @@ def _init_elastic_ep_world(

1324 group_ranks = [all_ranks[i : i + 1] for i in range(global_world_size)]

1325 if global_rank in all_ranks:

1326 group_ranks = [all_ranks]

1324 - group_ports = [parallel_config.get_next_stateless_world_group_port()]

1327 + coord_store = get_cached_tcp_store_client(

1328 + parallel_config.data_parallel_master_ip, parallel_config._coord_store_port

1329 + )

1330 world = StatelessGroupCoordinator(

1331 group_ranks=group_ranks,

1332 local_rank=local_rank,

@@ -1329,7 +1334,7 @@ def _init_elastic_ep_world(

1334 use_device_communicator=False,

1335 group_name="world",

1336 host=parallel_config.data_parallel_master_ip,

1332 - group_ports=group_ports,

1337 + coord_store=coord_store,

1338 global_rank=global_rank,

1339 global_world_size=global_world_size,

1340 )

@@ -1513,7 +1518,13 @@ def initialize_model_parallel(

1518 config = get_current_vllm_config()

1519 data_parallel_size = config.parallel_config.data_parallel_size

1520 enable_elastic_ep = config.parallel_config.enable_elastic_ep

1521 + parallel_config = config.parallel_config

1522 + coord_store: Store | None = None

1523 if enable_elastic_ep:

1524 + coord_store = get_cached_tcp_store_client(

1525 + parallel_config.data_parallel_master_ip,

1526 + parallel_config._coord_store_port,

1527 + )

1528 # Use stateless world group for global information

1529 world_size = get_world_group().world_size

1530 rank = get_world_group().rank

@@ -1633,16 +1644,12 @@ def initialize_model_parallel(

1644 group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0)

1645 group_ranks = [x.tolist() for x in group_ranks]

1646 if enable_elastic_ep:

1636 - parallel_config = config.parallel_config

1637 - dp_ports = [

1638 - parallel_config.get_next_stateless_dp_group_port() for _ in group_ranks

1639 - ]

1647 _DP = _init_stateless_group(

1648 group_ranks,

1649 "dp",

1643 - dp_ports,

1650 parallel_config.data_parallel_master_ip,

1651 backend,

1652 + coord_store=coord_store,

1653 )

1654 else:

1655 _DP = init_model_parallel_group(

@@ -1665,16 +1672,12 @@ def initialize_model_parallel(

1672 )

1673 group_ranks = [x.tolist() for x in group_ranks]

1674 if enable_elastic_ep:

1668 - parallel_config = config.parallel_config

1669 - ep_ports = [

1670 - parallel_config.get_next_stateless_ep_group_port() for _ in group_ranks

1671 - ]

1675 _EP = _init_stateless_group(

1676 group_ranks,

1677 "ep",

1675 - ep_ports,

1678 parallel_config.data_parallel_master_ip,

1679 backend,

1680 + coord_store=coord_store,

1681 )

1682 else:

1683 _EP = init_model_parallel_group(

@@ -1693,16 +1696,12 @@ def initialize_model_parallel(

1696 and config.parallel_config.enable_eplb

1697 ):

1698 if enable_elastic_ep:

1696 - eplb_ports = [

1697 - parallel_config.get_next_stateless_eplb_group_port()

1698 - for _ in group_ranks

1699 - ]

1699 _EPLB = _init_stateless_group(

1700 group_ranks,

1701 "eplb",

1703 - eplb_ports,

1702 parallel_config.data_parallel_master_ip,

1703 backend,

1704 + coord_store=coord_store,

1705 )

1706 else:

1707 _EPLB = init_model_parallel_group(

vllm/distributed/stateless_coordinator.py

+50 -7

@@ -1,9 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import socket

4 +import struct

5 from typing import Any, Optional

7 import torch

6 -from torch.distributed import Backend, ProcessGroup

8 +from torch.distributed import Backend, ProcessGroup, Store

10 from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator

11 from vllm.distributed.parallel_state import (

@@ -23,6 +25,38 @@ from vllm.utils.import_utils import resolve_obj_by_qualname

26 logger = init_logger(__name__)

28 +_PORTS_FMT = "!3I"

29 +

30 +

31 +def _allocate_group_ports(

32 + key: str,

33 + host: str,

34 + coord_store: Store,

35 +) -> tuple[list[int], list[socket.socket]]:

36 + """Bind 3 sockets and publish the ports to *coord_store*.

37 +

38 + Called by rank 0 only. Returns ``(ports, sockets)`` with the

39 + sockets still open.

40 + """

41 + socks: list[socket.socket] = []

42 + ports: list[int] = []

43 + for _ in range(3):

44 + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

45 + s.bind((host, 0))

46 + s.listen()

47 + socks.append(s)

48 + ports.append(s.getsockname()[1])

49 + coord_store.set(key, struct.pack(_PORTS_FMT, *ports))

50 + return ports, socks

51 +

52 +

53 +def _fetch_group_ports(key: str, coord_store: Store) -> list[int]:

54 + """Read 3 ports published by rank 0 from *coord_store*.

55 +

56 + Blocks until the key is available.

57 + """

58 + return list(struct.unpack(_PORTS_FMT, coord_store.get(key)))

59 +

61 class StatelessGroupCoordinator(GroupCoordinator):

62 """

@@ -39,10 +73,10 @@ class StatelessGroupCoordinator(GroupCoordinator):

73 local_rank: int,

74 torch_distributed_backend: str | Backend,

75 use_device_communicator: bool,

76 + coord_store: Store,

77 use_message_queue_broadcaster: bool = False,

78 group_name: str | None = None,

79 host: str = "127.0.0.1",

45 - group_ports: list[list[int]] | None = None,

80 global_rank: int = 0,

81 global_world_size: int = 1,

82 ):

@@ -61,17 +95,23 @@ class StatelessGroupCoordinator(GroupCoordinator):

96 backend = str(torch_distributed_backend)

97 self.backend = backend

64 - assert group_ports is not None, "group_ports is not provided"

98 for idx, ranks in enumerate(group_ranks):

99 if self.rank in ranks:

100 self.ranks = ranks

101 self.world_size = len(ranks)

102 self.rank_in_group = ranks.index(self.rank)

103

71 - ports = group_ports[idx]

72 - device_port = ports[0]

73 - cpu_port = ports[1]

74 - tcp_store_port = ports[2]

104 + key = f"{group_name}_{idx}"

105 + if self.rank_in_group == 0:

106 + ports, socks = _allocate_group_ports(

107 + key,

108 + host,

109 + coord_store,

110 + )

111 + else:

112 + ports = _fetch_group_ports(key, coord_store)

113 + socks = []

114 + device_port, cpu_port, tcp_store_port = ports

115

116 device_group = stateless_init_torch_distributed_process_group(

117 host=host,

@@ -80,6 +120,7 @@ class StatelessGroupCoordinator(GroupCoordinator):

120 world_size=self.world_size,

121 backend=backend,

122 group_name=f"{self.unique_name}_device",

123 + listen_socket=socks[0] if socks else None,

124 )

125 cpu_group = stateless_init_torch_distributed_process_group(

126 host=host,

@@ -88,12 +129,14 @@ class StatelessGroupCoordinator(GroupCoordinator):

129 world_size=self.world_size,

130 backend="gloo",

131 group_name=f"{self.unique_name}_cpu",

132 + listen_socket=socks[1] if socks else None,

133 )

134 tcp_store_group = StatelessProcessGroup.create(

135 host=host,

136 port=tcp_store_port,

137 rank=self.rank_in_group,

138 world_size=self.world_size,

139 + listen_socket=socks[2] if socks else None,

140 )

141

142 self_device_group = device_group

vllm/distributed/utils.py

+60 -18

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

6 # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py

8 import dataclasses

9 +import functools

10 import os

11 import pickle

12 import socket

@@ -139,6 +140,29 @@ def get_pp_indices(

140 return (start_layer, end_layer)

141

142

143 +def create_tcp_store(

144 + host: str,

145 + port: int,

146 + listen_socket: socket.socket | None = None,

147 + **kwargs: Any,

148 +) -> TCPStore:

149 + """Create a TCPStore, optionally taking ownership of ``listen_socket``."""

150 + if listen_socket is None:

151 + return TCPStore(host_name=host, port=port, **kwargs)

152 +

153 + listen_fd = listen_socket.detach()

154 + try:

155 + return TCPStore(

156 + host_name=host,

157 + port=port,

158 + master_listen_fd=listen_fd,

159 + **kwargs,

160 + )

161 + except Exception:

162 + socket.close(listen_fd)

163 + raise

164 +

165 +

166 @dataclasses.dataclass

167 class StatelessProcessGroup:

168 """A dataclass to hold a metadata store, and the rank, world_size of the

@@ -150,9 +174,6 @@ class StatelessProcessGroup:

174 world_size: int

175 store: torch._C._distributed_c10d.Store

176

153 - # stores a reference to the socket so that the file descriptor stays alive

154 - socket: socket.socket | None

155 -

177 data_expiration_seconds: int = 3600 # 1 hour

178

179 # dst rank -> counter

@@ -419,6 +440,7 @@ class StatelessProcessGroup:

440 world_size: int,

441 data_expiration_seconds: int = 3600,

442 store_timeout: int = 300,

443 + listen_socket: socket.socket | None = None,

444 ) -> "StatelessProcessGroup":

445 """A replacement for `torch.distributed.init_process_group` that does not

446 pollute the global state.

@@ -436,36 +458,39 @@ class StatelessProcessGroup:

458 C, and D can call `StatelessProcessGroup.create` to form another group.

459 """ # noqa

460 launch_server = rank == 0

439 - if launch_server:

440 - # listen on the specified interface (instead of 0.0.0.0)

461 + if launch_server and listen_socket is None:

462 listen_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

463 listen_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

464 listen_socket.bind((host, port))

465 listen_socket.listen()

445 - listen_fd = listen_socket.fileno()

446 - else:

447 - listen_socket = None

448 - listen_fd = None

449 -

450 - store = TCPStore(

451 - host_name=host,

452 - port=port,

466 + store = create_tcp_store(

467 + host,

468 + port,

469 + listen_socket=listen_socket,

470 world_size=world_size,

471 is_master=launch_server,

472 timeout=timedelta(seconds=store_timeout),

473 use_libuv=False, # for now: github.com/pytorch/pytorch/pull/150215

457 - master_listen_fd=listen_fd,

474 )

475

476 return StatelessProcessGroup(

477 rank=rank,

478 world_size=world_size,

479 store=store,

464 - socket=listen_socket,

480 data_expiration_seconds=data_expiration_seconds,

481 )

482

483

484 +@functools.lru_cache(maxsize=1)

485 +def get_cached_tcp_store_client(host: str, port: int) -> TCPStore:

486 + """Return a cached TCPStore client.

487 +

488 + Cached so that every call with the same ``(host, port)`` reuses the

489 + same connection. A new ``(host, port)`` evicts the old entry.

490 + """

491 + return TCPStore(host, port, is_master=False, wait_for_workers=False)

492 +

493 +

494 def init_gloo_process_group(

495 prefix_store: PrefixStore,

496 group_rank: int,

@@ -504,6 +529,7 @@ def stateless_init_torch_distributed_process_group(

529 backend: str,

530 group_name: str | None = None,

531 return_store: bool = False,

532 + listen_socket: socket.socket | None = None,

533 ) -> ProcessGroup | tuple[ProcessGroup, Store]:

534 """

535 A replacement for `torch.distributed.init_process_group` that does not

@@ -535,14 +561,30 @@ def stateless_init_torch_distributed_process_group(

561 are the same as process 1 and 5, the main communication channel is

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

563 channel is formed with process 9 and 10.

564 +

565 + When *listen_socket* is provided, the rendezvous step

566 + is skipped and a ``TCPStore`` server is created directly using the

567 + pre-bound socket. This is useful for eliminating TOCTOU races

568 + between port allocation and binding.

569 """

570 init_method = get_tcp_uri(host, port)

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

572 timeout = _get_default_timeout(backend)

573

543 - store, rank, world_size = next(

544 - rendezvous(init_method, rank, world_size, timeout=timeout)

545 - )

574 + if listen_socket is not None:

575 + store = create_tcp_store(

576 + host,

577 + port,

578 + listen_socket=listen_socket,

579 + world_size=world_size,

580 + is_master=True,

581 + timeout=timeout,

582 + multi_tenant=True,

583 + )

584 + else:

585 + store, rank, world_size = next(

586 + rendezvous(init_method, rank, world_size, timeout=timeout)

587 + )

588 store.set_timeout(timeout)

589

590 group_rank = rank

vllm/v1/engine/__init__.py

+1 -4

@@ -237,10 +237,7 @@ class ReconfigureDistributedRequest(msgspec.Struct):

237 new_data_parallel_master_ip: str

238 new_data_parallel_master_port: int

239 new_data_parallel_master_port_list: list[int]

240 - new_stateless_world_group_port_list: list[list[int]]

241 - new_stateless_dp_group_port_list: list[list[int]]

242 - new_stateless_ep_group_port_list: list[list[int]]

243 - new_stateless_eplb_group_port_list: list[list[int]]

240 + coord_store_port: int

241

242

243 class ReconfigureRankType(enum.IntEnum):

vllm/v1/engine/core.py

@@ -1767,6 +1767,7 @@ class DPEngineCoreProc(EngineCoreProc):

1767 new_parallel_config._data_parallel_master_port_list = (

1768 reconfig_request.new_data_parallel_master_port_list

1769 )

1770 + new_parallel_config._coord_store_port = reconfig_request.coord_store_port

1771

1772 is_scale_down = reconfig_request.new_data_parallel_size < old_dp_size

1773 is_shutdown = (

vllm/v1/engine/core_client.py

+28 -62

@@ -455,56 +455,6 @@ class ElasticScalingCache:

455 pending_notifications: dict[EEPNotificationType, set[int]]

456

457

458 -def allocate_stateless_group_ports(parallel_config, new_data_parallel_size: int):

459 - """

460 - Allocate stateless group ports for elastic EP.

461 - """

462 - from vllm.utils.network_utils import get_open_ports_list

463 -

464 - assert parallel_config.enable_elastic_ep, "Elastic EP must be enabled"

465 - world_size = parallel_config.world_size

466 - new_world_size_across_dp = world_size * new_data_parallel_size

467 - num_world_groups = 1

468 - num_dp_groups = max(1, new_world_size_across_dp // new_data_parallel_size)

469 - num_ep_groups = max(

470 - 1,

471 - new_world_size_across_dp

472 - // (new_data_parallel_size * parallel_config.tensor_parallel_size),

473 - )

474 - num_eplb_groups = num_ep_groups

475 - total_ports_needed = (

476 - num_world_groups + num_dp_groups + num_ep_groups + num_eplb_groups

477 - ) * 3 + 5

478 - all_ports = get_open_ports_list(total_ports_needed)

479 - new_data_parallel_master_port_list = all_ports[-5:]

480 - all_ports = all_ports[:-5]

481 - new_stateless_world_group_port_list = [

482 - all_ports[i : i + 3] for i in range(0, num_world_groups * 3, 3)

483 - ]

484 - start_idx = num_world_groups * 3

485 - new_stateless_dp_group_port_list = [

486 - all_ports[i : i + 3] for i in range(start_idx, start_idx + num_dp_groups * 3, 3)

487 - ]

488 - start_idx += num_dp_groups * 3

489 - new_stateless_ep_group_port_list = [

490 - all_ports[i : i + 3] for i in range(start_idx, start_idx + num_ep_groups * 3, 3)

491 - ]

492 - start_idx += num_ep_groups * 3

493 - new_stateless_eplb_group_port_list = [

494 - all_ports[i : i + 3]

495 - for i in range(start_idx, start_idx + num_eplb_groups * 3, 3)

496 - ]

497 -

498 - parallel_config._stateless_world_group_port_list = (

499 - new_stateless_world_group_port_list

500 - )

501 - parallel_config._stateless_dp_group_port_list = new_stateless_dp_group_port_list

502 - parallel_config._stateless_ep_group_port_list = new_stateless_ep_group_port_list

503 - parallel_config._stateless_eplb_group_port_list = new_stateless_eplb_group_port_list

504 - parallel_config.data_parallel_master_port = new_data_parallel_master_port_list.pop()

505 - parallel_config._data_parallel_master_port_list = new_data_parallel_master_port_list

506 -

507 -

458 class MPClient(EngineCoreClient):

459 """

460 MPClient: base client for multi-proc EngineCore.

@@ -1541,6 +1491,28 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1491 self._ensure_output_queue_task()

1492 await future

1493

1494 + def _setup_elastic_ep_reconfig_bootstrap(self) -> tuple[str, int]:

1495 + from vllm.distributed.utils import create_tcp_store

1496 + from vllm.utils.network_utils import get_open_ports_list

1497 +

1498 + parallel_config = self.vllm_config.parallel_config

1499 + parallel_config._data_parallel_master_port_list = get_open_ports_list(5)

1500 + parallel_config.data_parallel_master_port = (

1501 + parallel_config._data_parallel_master_port_list.pop()

1502 + )

1503 +

1504 + ip = parallel_config.data_parallel_master_ip

1505 + store = create_tcp_store(

1506 + ip,

1507 + 0,

1508 + is_master=True,

1509 + world_size=-1,

1510 + wait_for_workers=False,

1511 + )

1512 + parallel_config._coord_store_port = store.port

1513 + self._coord_store = store

1514 + return ip, store.port

1515 +

1516 async def _scale_up_elastic_ep(

1517 self, cur_data_parallel_size: int, new_data_parallel_size: int

1518 ) -> None:

@@ -1555,7 +1527,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1527 )

1528

1529 parallel_config = self.vllm_config.parallel_config

1558 - allocate_stateless_group_ports(parallel_config, new_data_parallel_size)

1530 + ip, coord_store_port = self._setup_elastic_ep_reconfig_bootstrap()

1531

1532 # Phase 1: Send reconfig messages to existing engines

1533 reconfig_futures = []

@@ -1564,13 +1536,10 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1536 new_data_parallel_size=new_data_parallel_size,

1537 new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK,

1538 new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK,

1567 - new_data_parallel_master_ip=parallel_config.data_parallel_master_ip,

1539 + new_data_parallel_master_ip=ip,

1540 new_data_parallel_master_port=parallel_config.data_parallel_master_port,

1541 new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list,

1570 - new_stateless_world_group_port_list=parallel_config._stateless_world_group_port_list,

1571 - new_stateless_dp_group_port_list=parallel_config._stateless_dp_group_port_list,

1572 - new_stateless_ep_group_port_list=parallel_config._stateless_ep_group_port_list,

1573 - new_stateless_eplb_group_port_list=parallel_config._stateless_eplb_group_port_list,

1542 + coord_store_port=coord_store_port,

1543 )

1544 coro = self._call_utility_async(

1545 "reinitialize_distributed", reconfig_request, engine=engine

@@ -1650,7 +1619,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1619 )

1620

1621 parallel_config = self.vllm_config.parallel_config

1653 - allocate_stateless_group_ports(parallel_config, new_data_parallel_size)

1622 + ip, coord_store_port = self._setup_elastic_ep_reconfig_bootstrap()

1623

1624 reconfig_futures = []

1625 for cur_dp_rank, engine in enumerate(self.core_engines):

@@ -1658,13 +1627,10 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1627 new_data_parallel_size=new_data_parallel_size,

1628 new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK,

1629 new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK,

1661 - new_data_parallel_master_ip=parallel_config.data_parallel_master_ip,

1630 + new_data_parallel_master_ip=ip,

1631 new_data_parallel_master_port=parallel_config.data_parallel_master_port,

1632 new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list,

1664 - new_stateless_world_group_port_list=parallel_config._stateless_world_group_port_list,

1665 - new_stateless_dp_group_port_list=parallel_config._stateless_dp_group_port_list,

1666 - new_stateless_ep_group_port_list=parallel_config._stateless_ep_group_port_list,

1667 - new_stateless_eplb_group_port_list=parallel_config._stateless_eplb_group_port_list,

1633 + coord_store_port=coord_store_port,

1634 )

1635 if cur_dp_rank >= new_data_parallel_size:

1636 reconfig_request.new_data_parallel_rank = (

vllm/v1/engine/utils.py

+14 -1

@@ -301,7 +301,20 @@ class CoreEngineActorManager:

301 else:

302 ray.init()

303

304 - vllm_config.parallel_config.allocate_elastic_ep_ports()

304 + parallel_config = vllm_config.parallel_config

305 + if parallel_config.enable_elastic_ep:

306 + from vllm.distributed.utils import create_tcp_store

307 +

308 + ip = parallel_config.data_parallel_master_ip

309 + store = create_tcp_store(

310 + ip,

311 + 0,

312 + is_master=True,

313 + world_size=-1,

314 + wait_for_workers=False,

315 + )

316 + parallel_config._coord_store_port = store.port

317 + self._coord_store = store

318

319 if placement_groups is not None:

320 assert local_dp_ranks is not None, (