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, 13:30 UTC
Use custom address for listening socket (#15988)
Signed-off-by: Jens Glaser <glaserj@ornl.gov>
jglaser committed
Apr 24, 2025 at 21:57 UTC
0d6e187e88874c39cda7409cf673f9e6546893e7
1 file changed
+21
-1
vllm/distributed/utils.py
+21
-1
@@ -7,6 +7,7 @@
7
import dataclasses
8
import datetime
9
import pickle
10
+import socket
11
import time
12
from collections import deque
13
from typing import Any, Deque, Dict, Optional, Sequence, Tuple
@@ -123,6 +124,10 @@ class StatelessProcessGroup:
124
rank: int
125
world_size: int
126
store: torch._C._distributed_c10d.Store
127
+
128
+ # stores a reference to the socket so that the file descriptor stays alive
129
+ socket: Optional[socket.socket]
130
+
131
data_expiration_seconds: int = 3600 # 1 hour
132
133
# dst rank -> counter
@@ -234,18 +239,33 @@ class StatelessProcessGroup:
239
can call `StatelessProcessGroup.create` to form a group, and then process A, B,
240
C, and D can call `StatelessProcessGroup.create` to form another group.
241
""" # noqa
242
+ launch_server = rank == 0
243
+ if launch_server:
244
+ # listen on the specified interface (instead of 0.0.0.0)
245
+ listen_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
246
+ listen_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
247
+ listen_socket.bind((host, port))
248
+ listen_socket.listen()
249
+ listen_fd = listen_socket.fileno()
250
+ else:
251
+ listen_socket = None
252
+ listen_fd = None
253
+
254
store = TCPStore(
255
host_name=host,
256
port=port,
257
world_size=world_size,
241
- is_master=(rank == 0),
258
+ is_master=launch_server,
259
timeout=datetime.timedelta(seconds=store_timeout),
260
+ use_libuv=False, # for now: github.com/pytorch/pytorch/pull/150215
261
+ master_listen_fd=listen_fd,
262
)
263
264
return StatelessProcessGroup(
265
rank=rank,
266
world_size=world_size,
267
store=store,
268
+ socket=listen_socket,
269
data_expiration_seconds=data_expiration_seconds)
270
271