A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Misc] Make timeout passable in init_distributed_environment (#24522)
Signed-off-by: jberkhahn <jaberkha@us.ibm.com>
Jonathan Berkhahn committed
Sep 10, 2025 at 15:41 UTC
cc99baf14dacc2497d0c5ed84e076ef2c37f6a4d
1 file changed
+9
-8
vllm/distributed/parallel_state.py
+9
-8
@@ -29,6 +29,7 @@ import weakref
29
from collections import namedtuple
30
from contextlib import contextmanager, nullcontext
31
from dataclasses import dataclass
32
+from datetime import timedelta
33
from multiprocessing import shared_memory
34
from typing import Any, Callable, Optional, Union
35
from unittest.mock import patch
@@ -978,13 +979,12 @@ def set_custom_all_reduce(enable: bool):
979
_ENABLE_CUSTOM_ALL_REDUCE = enable
980
981
981
-def init_distributed_environment(
982
- world_size: int = -1,
983
- rank: int = -1,
984
- distributed_init_method: str = "env://",
985
- local_rank: int = -1,
986
- backend: str = "nccl",
987
-):
982
+def init_distributed_environment(world_size: int = -1,
983
+ rank: int = -1,
984
+ distributed_init_method: str = "env://",
985
+ local_rank: int = -1,
986
+ backend: str = "nccl",
987
+ timeout: Optional[timedelta] = None):
988
logger.debug(
989
"world_size=%d rank=%d local_rank=%d "
990
"distributed_init_method=%s backend=%s", world_size, rank, local_rank,
@@ -1020,7 +1020,8 @@ def init_distributed_environment(
1020
backend=backend,
1021
init_method=distributed_init_method,
1022
world_size=world_size,
1023
- rank=rank)
1023
+ rank=rank,
1024
+ timeout=timeout)
1025
# set the local rank
1026
# local_rank is not available in torch ProcessGroup,
1027
# see https://github.com/pytorch/pytorch/issues/122816