pickle to msgpack migration

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

@@ -42,9 +42,15 @@
     init_distributed_environment,
     initialize_model_parallel,
 )
+from sglang.srt.entrypoints.http_api_specs import ProfileReqInputSpec
 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,
+    ProfileReqType,
+    async_sock_recv,
+    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 +2396,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 +2449,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)
@@ -2978,7 +2985,7 @@ async def _dp_worker_handle_profile(
         req = (
             ProfileReq(**obj)
             if obj is not None
-            else ProfileReq(ProfileReqType.START_PROFILE)
+            else ProfileReq(req_type=ProfileReqType.START_PROFILE)
         )
         if enc.profiler is None:
             enc.profiler = EncoderProfiler(dp_rank)
@@ -3197,9 +3204,9 @@ 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 request.req_type == ProfileReqType.START_PROFILE:
                 if encoder.profiler is None:
                     encoder.profiler = EncoderProfiler(encoder.rank)
                 encoder.profiler.start(request)
@@ -3601,7 +3608,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 +3829,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(
@@ -3854,12 +3861,12 @@ async def health_generate():
 
 
 @app.api_route("/start_profile", methods=["GET", "POST"])
-async def start_profile_async(obj: Optional[ProfileReqInput] = None):
+async def start_profile_async(obj: Optional[ProfileReqInputSpec] = None):
     if dp_dispatcher is not None:
         profile_req = None
         if obj is not None:
             profile_req = {
-                "type": ProfileReqType.START_PROFILE,
+                "req_type": ProfileReqType.START_PROFILE,
                 "output_dir": obj.output_dir,
                 "start_step": obj.start_step,
                 "num_steps": obj.num_steps,
@@ -3883,10 +3890,10 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None):
         return Response(content="encoder not ready\n", status_code=503)
     req = None
     if obj is None:
-        req = ProfileReq(ProfileReqType.START_PROFILE)
+        req = ProfileReq(req_type=ProfileReqType.START_PROFILE)
     else:
         req = ProfileReq(
-            type=ProfileReqType.START_PROFILE,
+            req_type=ProfileReqType.START_PROFILE,
             output_dir=obj.output_dir,
             start_step=obj.start_step,
             num_steps=obj.num_steps,
@@ -3900,7 +3907,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)
@@ -3929,9 +3936,9 @@ async def stop_profile_async():
         return Response(
             content="profiling not initialized\n", status_code=HTTPStatus.BAD_REQUEST
         )
-    req = ProfileReq(ProfileReqType.STOP_PROFILE)
+    req = ProfileReq(req_type=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

@@ -80,6 +80,8 @@
     UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromIPCReqInput,
     UpdateWeightsFromTensorReqInput,
+    sock_recv,
+    sock_send,
 )
 from sglang.srt.managers.multi_tokenizer_mixin import (
     MultiTokenizerRouter,
@@ -861,8 +863,8 @@ def _launch_subprocesses(
         for p in detoken_procs:
             scheduler_init_result.all_child_pids.append(p.pid)
 
-        # Init tokenizer manager first, as the bootstrap server is initialized here
         if server_args.tokenizer_worker_num == 1:
+            # Launch a single tokenizer manager without the router
             tokenizer_manager, template_manager = init_tokenizer_manager_func(
                 server_args, port_args
             )
@@ -1223,8 +1225,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/entrypoints/grpc_server.py

@@ -105,7 +105,7 @@ async def start_profile_handler(request):
             record_shapes = (record_shapes is not False) and env_record_shapes
 
             req = ProfileReq(
-                type=ProfileReqType.START_PROFILE,
+                req_type=ProfileReqType.START_PROFILE,
                 output_dir=body.get("output_dir"),
                 start_step=body.get("start_step"),
                 num_steps=body.get("num_steps"),
@@ -134,7 +134,7 @@ async def start_profile_handler(request):
 
     async def stop_profile_handler(request):
         try:
-            req = ProfileReq(type=ProfileReqType.STOP_PROFILE)
+            req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE)
             results = await request_manager.send_communicator_req(
                 req, "profile_communicator", timeout=600.0
             )