[Refactor] Introduce sock_send/sock_recv wrappers for zmq IPC (#29012)

Author: merrymercy

python/sglang/srt/debug_utils/dumper.py

@@ -20,6 +20,9 @@
 
 import torch
 import torch.distributed as dist
+import zmq
+
+from sglang.srt.managers.io_struct import sock_recv, sock_send
 
 # -------------------------------------- config base ------------------------------------------
 
@@ -1419,8 +1422,6 @@ def _create_zmq_rpc_broadcast(
     handler, timeout_seconds: int = 60
 ) -> Optional["_ZmqRpcBroadcast"]:
     """A general-purpose minimal RPC to support broadcasting executions to multi processes"""
-    import zmq
-
     rank = _get_rank()
     world_size = dist.get_world_size() if dist.is_initialized() else 1
 
@@ -1433,13 +1434,13 @@ def _create_zmq_rpc_broadcast(
     def serve_loop():
         while True:
             try:
-                req = sock.recv_pyobj()
+                req = sock_recv(sock)
                 result = getattr(handler, req["method"])(*req["args"], **req["kwargs"])
                 resp = {"result": result, "error": None}
             except Exception as e:
                 _log(f"[ZmqRpc] error inside handler: {e}")
                 resp = {"result": None, "error": str(e)}
-            sock.send_pyobj(resp)
+            sock_send(sock, resp)
 
     thread = threading.Thread(target=serve_loop, daemon=True)
     thread.start()
@@ -1476,14 +1477,15 @@ def __init__(self, socket, debug_name: str):
 
     def __getattr__(self, method_name: str):
         def call(*args, **kwargs):
-            self._socket.send_pyobj(
+            sock_send(
+                self._socket,
                 {
                     "method": method_name,
                     "args": args,
                     "kwargs": kwargs,
-                }
+                },
             )
-            response = self._socket.recv_pyobj()
+            response = sock_recv(self._socket)
             if response["error"]:
                 raise RuntimeError(
                     f"RPC error on {self._debug_name}: {response['error']}"

python/sglang/srt/disaggregation/encode_grpc_server.py

@@ -26,6 +26,7 @@
     handle_scheduler_receive_url_request,
     launch_encoder,
 )
+from sglang.srt.managers.io_struct import async_sock_send
 from sglang.srt.managers.schedule_batch import Modality
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import random_uuid
@@ -95,7 +96,7 @@ async def Encode(
                 "part_idx": request.part_idx,
             }
             for socket in self.send_sockets:
-                await socket.send_pyobj(request_dict)
+                await async_sock_send(socket, request_dict)
 
             # gRPC encode is image-only; encoder.encode() requires modality
             (

python/sglang/srt/disaggregation/encode_server.py

@@ -44,7 +44,14 @@
 )
 from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import initialize_dp_attention
-from sglang.srt.managers.io_struct import ProfileReq, ProfileReqInput, ProfileReqType
+from sglang.srt.managers.io_struct import (
+    ProfileReq,
+    ProfileReqInput,
+    ProfileReqType,
+    async_sock_recv,
+    async_sock_send,
+    sock_send,
+)
 from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
 from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
 from sglang.srt.model_loader import get_model
@@ -2390,13 +2397,14 @@ async def _dispatch_group(
         requests = [p.request for p in group]
         start = time.time()
         for sock in self.send_sockets:
-            sock.send_pyobj(
+            sock_send(
+                sock,
                 {
                     "type": "batch_encode",
                     "modality": modality.name,
                     "requests": requests,
                     "enter_time": start,
-                }
+                },
             )
 
         logger.info(f"Dispatching batch of {len(group)} {modality.name} requests")
@@ -2442,7 +2450,7 @@ async def _dispatch_per_request(
             req = p.request
             try:
                 for sock in self.send_sockets:
-                    sock.send_pyobj(req)
+                    sock_send(sock, req)
                 result = await self.encoder.encode_request(req, modality)
                 if not p.future.done():
                     p.future.set_result(result)
@@ -2731,7 +2739,7 @@ async def dispatch(self, request: dict) -> dict:
         )
 
         try:
-            await self.dispatch_sockets[rank].send_pyobj(request)
+            await async_sock_send(self.dispatch_sockets[rank], request)
             # An alive-but-stuck worker (NCCL deadlock etc.) wouldn't trip
             # the watchdog, so bound the wait explicitly.
             return await asyncio.wait_for(future, timeout=ENCODER_REQ_TIMEOUT)
@@ -2780,7 +2788,7 @@ async def dispatch_send(self, request: dict) -> dict:
             f"dp_rank={rank}, pending={self.pending_counts}"
         )
         try:
-            await self.dispatch_sockets[rank].send_pyobj(request)
+            await async_sock_send(self.dispatch_sockets[rank], request)
             return await asyncio.wait_for(future, timeout=ENCODER_REQ_TIMEOUT)
         except asyncio.TimeoutError:
             self.pending_futures[rank].pop(key, None)
@@ -2821,7 +2829,7 @@ async def broadcast(
                 self.req_id_to_rank[req_id] = rank
                 rank_keys.append((rank, req_id))
                 request_copy = {**request, "req_id": req_id}
-                await self.dispatch_sockets[rank].send_pyobj(request_copy)
+                await async_sock_send(self.dispatch_sockets[rank], request_copy)
                 futures.append(future)
             # Concurrent wait → total bounded by eff_timeout, not
             # dp_size × eff_timeout.
@@ -2901,7 +2909,7 @@ async def _result_listener(self) -> None:
         consecutive_errors = 0
         while True:
             try:
-                msg = await self.result_socket.recv_pyobj()
+                msg = await async_sock_recv(self.result_socket)
                 consecutive_errors = 0
             except asyncio.CancelledError:
                 raise
@@ -3051,10 +3059,10 @@ async def _dp_worker_handle_request(
             "_error_code": err_code,
         }
 
-    # pyzmq async send_pyobj isn't safe for concurrent senders.
+    # pyzmq async send isn't safe for concurrent senders.
     try:
         async with send_lock:
-            await send_sock.send_pyobj(envelope)
+            await async_sock_send(send_sock, envelope)
     except Exception:
         logger.error(
             f"DP worker {dp_rank} failed to send envelope for "
@@ -3113,7 +3121,7 @@ async def run_dp_worker(
             spawned = False
             try:
                 try:
-                    request = await recv_sock.recv_pyobj()
+                    request = await async_sock_recv(recv_sock)
                 except asyncio.CancelledError:
                     raise
                 except Exception:
@@ -3197,7 +3205,7 @@ async def run_encoder(
 ):
     encoder = MMEncoder(server_args, schedule_path, dist_init_method, rank)
     while True:
-        request = await encoder.schedule_socket.recv_pyobj()
+        request = await async_sock_recv(encoder.schedule_socket)
         if isinstance(request, ProfileReq):
             if request.type == ProfileReqType.START_PROFILE:
                 if encoder.profiler is None:
@@ -3601,7 +3609,7 @@ def start_background_send(req_id):
                     )
             else:
                 for socket in send_sockets:
-                    socket.send_pyobj(request)
+                    sock_send(socket, request)
                 nbytes, embedding_len, embedding_dim, error_msg, error_code = (
                     await encoder.encode_request(request, modality)
                 )
@@ -3822,7 +3830,7 @@ async def health_generate():
 
         # Broadcast to other TP ranks so distributed ops stay in sync
         for socket in send_sockets:
-            socket.send_pyobj(dummy_request)
+            sock_send(socket, dummy_request)
 
         # Run encode on rank 0 with timeout
         _, _, _, error_msg, _ = await asyncio.wait_for(
@@ -3900,7 +3908,7 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None):
             profile_stages=obj.profile_stages,
         )
     for socket in send_sockets:
-        socket.send_pyobj(req)
+        sock_send(socket, req)
     if encoder.profiler is None:
         encoder.profiler = EncoderProfiler(encoder.rank)
     ok, msg = encoder.profiler.start(req)
@@ -3931,7 +3939,7 @@ async def stop_profile_async():
         )
     req = ProfileReq(ProfileReqType.STOP_PROFILE)
     for socket in send_sockets:
-        socket.send_pyobj(req)
+        sock_send(socket, req)
     ok, msg = encoder.profiler.stop()
     if ok:
         return Response(content="Stop profiling.\n", status_code=200)

python/sglang/srt/elastic_ep/expert_backup_client.py

@@ -12,7 +12,7 @@
 )
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
-from sglang.srt.managers.io_struct import UpdateExpertBackupReq
+from sglang.srt.managers.io_struct import UpdateExpertBackupReq, sock_recv, sock_send
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils.network import get_local_ip_auto
 
@@ -65,15 +65,15 @@ def __init__(self, server_args: ServerArgs, model_runner):
             self.ready_sockets[i].connect(
                 f"tcp://{all_ips[i * get_world_size() // server_args.nnodes]}:{PORT_BASE + i * 2}"
             )
-            self.ready_sockets[i].send_pyobj(UpdateExpertBackupReq())
+            sock_send(self.ready_sockets[i], UpdateExpertBackupReq())
 
         self._receive_thread = threading.Thread(target=self._receive_loop, daemon=True)
         self._receive_thread.start()
 
     def _receive_loop(self):
         cnt = 0
         while cnt < self.engine_num:
-            response = self.recv_list[cnt].recv_pyobj()
+            response = sock_recv(self.recv_list[cnt])
             self.dram_map_list[response.rank] = response.weight_pointer_map
             self.session_id_list[response.rank] = response.session_id
             self.buffer_size = max(self.buffer_size, response.buffer_size)

python/sglang/srt/elastic_ep/expert_backup_manager.py

@@ -9,7 +9,7 @@
 from sglang.srt.configs.load_config import LoadConfig
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.environ import envs
-from sglang.srt.managers.io_struct import BackupDramReq
+from sglang.srt.managers.io_struct import BackupDramReq, sock_recv, sock_send
 from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
 from sglang.srt.model_loader.utils import set_default_torch_dtype
 from sglang.srt.server_args import (
@@ -62,7 +62,7 @@ def __init__(self, server_args: ServerArgs, port_args: PortArgs):
         num_ready_clients = 0
 
         while num_ready_clients < server_args.tp_size:
-            self.recv_from_expert_backup_client.recv_pyobj()
+            sock_recv(self.recv_from_expert_backup_client)
             num_ready_clients += 1
 
         back_req = BackupDramReq(
@@ -72,7 +72,7 @@ def __init__(self, server_args: ServerArgs, port_args: PortArgs):
             buffer_size=self.continuous_buffer.numel()
             * self.continuous_buffer.element_size(),
         )
-        self.send_to_expert_backup_client.send_pyobj(back_req)
+        sock_send(self.send_to_expert_backup_client, back_req)
 
         # Keep the manager subprocess alive until signals
         signal.pause()

python/sglang/srt/entrypoints/engine.py

@@ -77,6 +77,8 @@
     UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromIPCReqInput,
     UpdateWeightsFromTensorReqInput,
+    sock_recv,
+    sock_send,
 )
 from sglang.srt.managers.multi_tokenizer_mixin import (
     MultiTokenizerRouter,
@@ -1220,8 +1222,8 @@ def freeze_gc(self):
 
     def collective_rpc(self, method: str, **kwargs):
         obj = RpcReqInput(method=method, parameters=kwargs)
-        self.send_to_rpc.send_pyobj(obj)
-        recv_req = self.send_to_rpc.recv_pyobj(zmq.BLOCKY)
+        sock_send(self.send_to_rpc, obj)
+        recv_req = sock_recv(self.send_to_rpc, flags=zmq.BLOCKY)
         assert isinstance(recv_req, RpcReqOutput)
         assert recv_req.success, recv_req.message
 

python/sglang/srt/managers/communicator.py

@@ -5,7 +5,7 @@
 from collections import deque
 from typing import Deque, Generic, List, Optional, TypeVar
 
-import zmq
+from sglang.srt.managers.io_struct import sock_send
 
 T = TypeVar("T")
 
@@ -22,7 +22,7 @@ class FanOutCommunicator(Generic[T]):
     Only one request is in-flight at any time in either mode.
     """
 
-    def __init__(self, sender: zmq.Socket, fan_out: int, mode="queueing"):
+    def __init__(self, sender, fan_out: int, mode="queueing"):
         self._sender = sender
         self._fan_out = fan_out
         self._mode = mode
@@ -41,7 +41,7 @@ async def queueing_call(self, obj: T):
             assert self._result_values is None
 
         if obj is not None:
-            self._sender.send_pyobj(obj)
+            sock_send(self._sender, obj)
 
         self._result_event = asyncio.Event()
         self._result_values = []
@@ -61,7 +61,7 @@ async def watching_call(self, obj):
             self._result_event = asyncio.Event()
 
             if obj is not None:
-                self._sender.send_pyobj(obj)
+                sock_send(self._sender, obj)
 
         # Capture local refs before await -- after event fires, the first
         # awakened coroutine clears shared state; later awaiters use local refs.