EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Set ZMQ buffer size heuristic (#1801)

Browse Source
This commit is contained in:
Liangsheng Yin
2024-10-25 23:15:56 -07:00
committed by GitHub
parent 1e8903414a
commit a628dd8e31
1 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/utils.py
View File

@@ -724,14 +724,22 @@ def first_rank_print(*args, **kwargs):
def get_zmq_socket(context: zmq.Context, socket_type: zmq.SocketType, endpoint: str): def get_zmq_socket(context: zmq.Context, socket_type: zmq.SocketType, endpoint: str):
mem = psutil.virtual_memory()
total_mem = mem.total / 1024**3
available_mem = mem.available / 1024**3
if total_mem > 32 and available_mem > 16:
buf_size = int(0.5 * 1024**3)
else:
buf_size = -1
socket = context.socket(socket_type) socket = context.socket(socket_type)
if socket_type == zmq.PUSH: if socket_type == zmq.PUSH:
socket.setsockopt(zmq.SNDHWM, 0) socket.setsockopt(zmq.SNDHWM, 0)
socket.setsockopt(zmq.SNDBUF, 100000000) socket.setsockopt(zmq.SNDBUF, buf_size)
socket.connect(f"ipc://{endpoint}") socket.connect(f"ipc://{endpoint}")
elif socket_type == zmq.PULL: elif socket_type == zmq.PULL:
socket.setsockopt(zmq.RCVHWM, 0) socket.setsockopt(zmq.RCVHWM, 0)
socket.setsockopt(zmq.RCVBUF, 100000000) socket.setsockopt(zmq.RCVBUF, buf_size)
socket.bind(f"ipc://{endpoint}") socket.bind(f"ipc://{endpoint}")
else: else:
raise ValueError(f"Unsupported socket type: {socket_type}") raise ValueError(f"Unsupported socket type: {socket_type}")