@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4406874e7c

[Bugfix][FT] Pass stateless process-group timeouts explicitly (#58611)

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Codex <noreply@openai.com>

Tyler Michael Smith committed Sep 28, 2026 at 23:14 UTC 4406874e7c5ffca4b1758f9153e18e57f3020096

6 files changed +39 -14

tests/test_config.py

+15

@@ -5,6 +5,7 @@ import json

5 import logging

6 import os

7 from dataclasses import MISSING, Field, asdict, dataclass, field

8 +from datetime import timedelta

9 from pathlib import Path

10 from types import SimpleNamespace

11 from typing import cast

@@ -1464,6 +1465,20 @@ def test_all2all_backend_has_portable_default():

1465 assert ParallelConfig().all2all_backend == "allgather_reducescatter"

1466

1467

1468 +def test_dp_group_uses_configured_timeout_without_current_config(monkeypatch):

1469 + monkeypatch.setattr(vllm_config_module, "_current_vllm_config", None)

1470 + config = ParallelConfig(cpu_distributed_timeout_seconds=30)

1471 + with (

1472 + patch(

1473 + "vllm.distributed.utils.rendezvous",

1474 + return_value=iter([(torch.distributed.HashStore(), 0, 1)]),

1475 + ),

1476 + patch("vllm.distributed.utils.init_gloo_process_group") as init_group,

1477 + ):

1478 + config.stateless_init_dp_group()

1479 + assert init_group.call_args.kwargs["timeout"] == timedelta(seconds=30)

1480 +

1481 +

1482 @pytest.mark.parametrize(

1483 "dp_size, across_dp, expected",

1484 [(1, False, 4), (1, True, 4), (2, False, 4), (2, True, 8)],

vllm/config/parallel.py

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

4 import os

5 import socket

6 from collections.abc import Callable

7 +from datetime import timedelta

8 from typing import TYPE_CHECKING, Any, Literal, overload

10 import regex as re

@@ -613,6 +614,11 @@ class ParallelConfig:

614 """

615 return self.data_parallel_external_lb or self.data_parallel_hybrid_lb

616

617 + @property

618 + def cpu_distributed_timeout(self) -> timedelta | None:

619 + seconds = self.cpu_distributed_timeout_seconds

620 + return timedelta(seconds=seconds) if seconds is not None else None

621 +

622 def get_next_dp_init_port(self) -> int:

623 """We might need to initialize process groups in multiple

624 processes that is related to data parallelism,

@@ -694,6 +700,7 @@ class ParallelConfig:

700 backend="gloo",

701 return_store=return_store,

702 listen_socket=listen_socket,

703 + timeout=self.cpu_distributed_timeout,

704 )

705 except DistNetworkError as e:

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

vllm/distributed/stateless_coordinator.py

@@ -17,6 +17,8 @@ from vllm.distributed.parallel_state import (

17 )

18 from vllm.distributed.utils import (

19 StatelessProcessGroup,

20 + get_cpu_distributed_timeout_or_none,

21 + get_distributed_timeout_or_none,

22 stateless_destroy_torch_distributed_process_group,

23 stateless_init_torch_distributed_process_group,

24 )

@@ -104,6 +106,10 @@ class StatelessGroupCoordinator(GroupCoordinator):

106

107 backend = str(torch_distributed_backend)

108 self.backend = backend

109 + cpu_timeout = get_cpu_distributed_timeout_or_none()

110 + device_timeout = (

111 + cpu_timeout if backend == "gloo" else get_distributed_timeout_or_none()

112 + )

113 for idx, ranks in enumerate(group_ranks):

114 if self.rank in ranks:

115 self.ranks = ranks

@@ -130,6 +136,7 @@ class StatelessGroupCoordinator(GroupCoordinator):

136 backend=backend,

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

138 listen_socket=socks[0] if socks else None,

139 + timeout=device_timeout,

140 )

141 cpu_group = stateless_init_torch_distributed_process_group(

142 host=host,

@@ -139,6 +146,7 @@ class StatelessGroupCoordinator(GroupCoordinator):

146 backend="gloo",

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

148 listen_socket=socks[1] if socks else None,

149 + timeout=cpu_timeout,

150 )

151 tcp_store_group = StatelessProcessGroup.create(

152 host=host,

vllm/distributed/utils.py

+6 -11

@@ -530,8 +530,7 @@ def get_cpu_distributed_timeout_or_none() -> timedelta | None:

530 vllm_config = get_current_vllm_config_or_none()

531 if vllm_config is None:

532 return None

533 - timeout_seconds = vllm_config.parallel_config.cpu_distributed_timeout_seconds

534 - return timedelta(seconds=timeout_seconds) if timeout_seconds is not None else None

533 + return vllm_config.parallel_config.cpu_distributed_timeout

534

535

536 def get_distributed_timeout_or_none() -> timedelta | None:

@@ -582,6 +581,7 @@ def stateless_init_torch_distributed_process_group(

581 group_name: str | None = None,

582 return_store: bool = False,

583 listen_socket: socket.socket | None = None,

584 + timeout: timedelta | None = None,

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

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

587 pollute the global state. The created ProcessGroup object can be used for

@@ -617,18 +617,13 @@ def stateless_init_torch_distributed_process_group(

617 is skipped and a ``TCPStore`` server is created directly using the

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

619 between port allocation and binding.

620 +

621 + If *timeout* is None, use PyTorch's default for the backend.

622 """

623 init_method = get_tcp_uri(host, port)

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

623 - timeout = _get_default_timeout(backend)

624 - if backend == "gloo":

625 - gloo_timeout = get_cpu_distributed_timeout_or_none()

626 - if gloo_timeout is not None:

627 - timeout = gloo_timeout

628 - else:

629 - device_timeout = get_distributed_timeout_or_none()

630 - if device_timeout is not None:

631 - timeout = device_timeout

625 + if timeout is None:

626 + timeout = _get_default_timeout(backend)

627

628 if listen_socket is not None:

629 store = create_tcp_store(

vllm/v1/fault_tolerance/engine_core_sentinel.py

+2 -3

@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING

10 import msgspec

12 -from vllm.config import set_current_vllm_config

12 from vllm.distributed import stateless_destroy_torch_distributed_process_group

13 from vllm.distributed.utils import stateless_init_torch_distributed_process_group

14 from vllm.logger import init_logger

@@ -123,8 +122,7 @@ class EngineCoreSentinel:

122 engine = self.engine

123 executor = engine.model_executor

124

126 - with set_current_vllm_config(engine.vllm_config):

127 - ft_request.params.update(self._reinit_dp_group())

125 + ft_request.params.update(self._reinit_dp_group())

126 if hasattr(engine, "step_counter"):

127 engine.step_counter = 0

128

@@ -165,6 +163,7 @@ class EngineCoreSentinel:

163 parallel_config.data_parallel_size,

164 backend="gloo",

165 return_store=True,

166 + timeout=parallel_config.cpu_distributed_timeout,

167 )

168 )

169 return {"new_stateless_dp_group_ports": worker_ports}

vllm/v1/worker/sentinel/gpu_worker_sentinel.py

@@ -65,6 +65,7 @@ class WorkerSentinel:

65 self.dp_rank,

66 self.dp_size,

67 backend="gloo",

68 + timeout=self.worker.parallel_config.cpu_distributed_timeout,

69 )

71 def _clean_worker_state(self):