From b40f605fdeee49a275f5c0d54f0f249de8cfdafe Mon Sep 17 00:00:00 2001 From: Jimmy <29097382+jinmingyi1998@users.noreply.github.com> Date: Wed, 12 Nov 2025 18:05:16 +0800 Subject: [PATCH] fix(tcp-port): replace bind_server_socket to get_zmq_socket(Port conflict) (#11961) Co-authored-by: wangchao --- .../sglang/srt/disaggregation/common/conn.py | 19 +++++++-------- python/sglang/srt/disaggregation/kv_events.py | 6 +++++ .../srt/disaggregation/mooncake/conn.py | 4 ---- python/sglang/srt/disaggregation/nixl/conn.py | 7 ++---- .../device_communicators/shm_broadcast.py | 7 +++--- python/sglang/srt/utils/common.py | 23 +++++++++++++++++++ 6 files changed, 44 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 1a66c83ae..636ff72db 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -32,8 +32,8 @@ from sglang.srt.layers.dp_attention import ( from sglang.srt.server_args import ServerArgs from sglang.srt.utils import ( format_tcp_address, - get_free_port, get_local_ip_auto, + get_zmq_socket_on_host, is_valid_ipv6_address, maybe_wrap_ipv6_address, ) @@ -68,11 +68,16 @@ class CommonKVManager(BaseKVManager): ) self.pp_size = server_args.pp_size self.pp_rank = self.kv_args.pp_rank - self.rank_port = get_free_port() self.local_ip = get_local_ip_auto() - self.server_socket = zmq.Context().socket(zmq.PULL) - if is_valid_ipv6_address(self.local_ip): - self.server_socket.setsockopt(zmq.IPV6, 1) + + # bind zmq socket + context = zmq.Context() + zmq_bind_host = maybe_wrap_ipv6_address(self.local_ip) + self.rank_port, self.server_socket = get_zmq_socket_on_host( + context, zmq.PULL, host=zmq_bind_host + ) + logger.debug(f"kv manager bind to {zmq_bind_host}:{self.rank_port}") + self.request_status: Dict[int, KVPoll] = {} if self.disaggregation_mode == DisaggregationMode.PREFILL: @@ -92,9 +97,6 @@ class CommonKVManager(BaseKVManager): f"Unsupported DisaggregationMode: {self.disaggregation_mode}" ) - def _bind_server_socket(self): - self.server_socket.bind(format_tcp_address(self.local_ip, self.rank_port)) - def _register_to_bootstrap(self): """Register KVSender to bootstrap server via HTTP POST.""" if self.dist_init_addr: @@ -177,7 +179,6 @@ class CommonKVManager(BaseKVManager): class CommonKVSender(BaseKVSender): - def __init__( self, mgr: BaseKVManager, diff --git a/python/sglang/srt/disaggregation/kv_events.py b/python/sglang/srt/disaggregation/kv_events.py index f0a3a4357..22c7aeeb3 100644 --- a/python/sglang/srt/disaggregation/kv_events.py +++ b/python/sglang/srt/disaggregation/kv_events.py @@ -238,6 +238,9 @@ class ZmqEventPublisher(EventPublisher): or self._endpoint.startswith("ipc://") or self._endpoint.startswith("inproc://") ): + logger.debug( + f"ZmqEventPublisher socket publisher_endpoint bind to {self._endpoint}" + ) self._pub.bind(self._endpoint) else: self._pub.connect(self._endpoint) @@ -248,6 +251,9 @@ class ZmqEventPublisher(EventPublisher): # 3) works in our non‑blocking poll loop alongside PUB if self._replay_endpoint is not None: self._replay = self._ctx.socket(zmq.ROUTER) + logger.debug( + f"ZmqEventPublisher socket replay_endpoint bind to {self._replay_endpoint}" + ) self._replay.bind(self._replay_endpoint) def _publisher_thread(self) -> None: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 0e1f844a5..1a7670f59 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -849,8 +849,6 @@ class MooncakeKVManager(CommonKVManager): ) def start_prefill_thread(self): - self._bind_server_socket() - def bootstrap_thread(): """This thread recvs pre-alloc notification from the decode engine""" # KVPoll.Bootstrapping -> KVPoll.WaitingForInput @@ -887,8 +885,6 @@ class MooncakeKVManager(CommonKVManager): threading.Thread(target=bootstrap_thread).start() def start_decode_thread(self): - self._bind_server_socket() - def decode_thread(): while True: msg = self.server_socket.recv_multipart() diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 5cabbf677..c61c3e980 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -83,8 +83,8 @@ class KVArgsRegisterInfo: dst_port=int(msg[2].decode("ascii")), agent_name=msg[3].decode("ascii"), agent_metadata=msg[4], - dst_kv_ptrs=list(struct.unpack(f"{len(msg[5])//8}Q", msg[5])), - dst_aux_ptrs=list(struct.unpack(f"{len(msg[6])//8}Q", msg[6])), + dst_kv_ptrs=list(struct.unpack(f"{len(msg[5]) // 8}Q", msg[5])), + dst_aux_ptrs=list(struct.unpack(f"{len(msg[6]) // 8}Q", msg[6])), gpu_id=int(msg[7].decode("ascii")), decode_tp_size=int(msg[8].decode("ascii")), decode_tp_rank=int(msg[9].decode("ascii")), @@ -647,8 +647,6 @@ class NixlKVManager(CommonKVManager): return self.transfer_statuses[room].is_done() def _start_bootstrap_thread(self): - self._bind_server_socket() - def bootstrap_thread(): """This thread recvs transfer info from the decode engine""" while True: @@ -687,7 +685,6 @@ class NixlKVManager(CommonKVManager): class NixlKVSender(CommonKVSender): - def __init__( self, mgr: NixlKVManager, diff --git a/python/sglang/srt/distributed/device_communicators/shm_broadcast.py b/python/sglang/srt/distributed/device_communicators/shm_broadcast.py index e956a2592..39788fe30 100644 --- a/python/sglang/srt/distributed/device_communicators/shm_broadcast.py +++ b/python/sglang/srt/distributed/device_communicators/shm_broadcast.py @@ -215,7 +215,6 @@ class MessageQueue: socket_addr = f"tcp://127.0.0.1:{local_subscribe_port}" logger.debug("Binding to %s", socket_addr) self.local_socket.bind(socket_addr) - self.current_idx = 0 else: @@ -232,9 +231,9 @@ class MessageQueue: remote_subscribe_port = get_open_port() if is_valid_ipv6_address(connect_ip): self.remote_socket.setsockopt(IPV6, 1) - self.remote_socket.bind( - format_tcp_address(connect_ip, remote_subscribe_port) - ) + address = format_tcp_address(connect_ip, remote_subscribe_port) + logger.debug(f"class MessageQueue: Binding remote socket to {address=}") + self.remote_socket.bind(address) else: remote_subscribe_port = None diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 3283d5834..3a0c02c2a 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1359,6 +1359,29 @@ def get_zmq_socket( return socket +def get_zmq_socket_on_host( + context: zmq.Context, + socket_type: zmq.SocketType, + host: Optional[str] = None, +) -> Tuple[int, zmq.Socket]: + """Create and configure a ZeroMQ socket. + + Args: + context: ZeroMQ context to create the socket from. + socket_type: Type of ZeroMQ socket to create. + host: Optional host to bind/connect to, without "tcp://" prefix. If None, binds to "tcp://*". + + Returns: + Tuple of (port, socket) where port is the randomly assigned TCP port. + """ + socket = context.socket(socket_type) + # Bind to random TCP port + config_socket(socket, socket_type) + bind_host = f"tcp://{host}" if host else "tcp://*" + port = socket.bind_to_random_port(bind_host) + return port, socket + + def config_socket(socket, socket_type: zmq.SocketType): mem = psutil.virtual_memory() total_mem = mem.total / 1024**3