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

This commit is contained in:
Lianmin Zheng
2026-06-23 15:54:36 -07:00
committed by GitHub
parent ecab3f322e
commit 34dd9c28ca
22 changed files with 258 additions and 154 deletions
@@ -12,6 +12,7 @@ import zmq
from sglang.srt.entrypoints.http_server import launch_server
from sglang.srt.environ import envs
from sglang.srt.managers.io_struct import sock_recv, sock_send
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils.network import get_free_port, get_zmq_socket_on_host
from sglang.test.scripted_runtime.io_struct import (
@@ -89,7 +90,7 @@ class ScriptedHttpServer:
raise RuntimeError(f"ScriptedHttpServer is dirty: {self._dirty}")
fn_path = f"{script_fn.__module__}:{script_fn.__qualname__}"
self._socket.send_pyobj(RunScript(fn_path=fn_path, args=args))
sock_send(self._socket, RunScript(fn_path=fn_path, args=args))
if not self._socket.poll(int(timeout_s * 1000)):
if not self._server_process.is_alive():
@@ -98,7 +99,7 @@ class ScriptedHttpServer:
self._dirty = f"script {fn_path!r} timed out after {timeout_s}s"
raise TimeoutError(self._dirty)
reply = self._socket.recv_pyobj()
reply = sock_recv(self._socket)
match reply:
case ScriptFailed(traceback=tb):
raise AssertionError(f"scripted-runtime script failed:\n{tb}")
@@ -116,7 +117,7 @@ class ScriptedHttpServer:
fatal_error: Optional[OutOfBandError] = None
try:
try:
self._socket.send_pyobj(Shutdown())
sock_send(self._socket, Shutdown())
except zmq.ZMQError:
pass
@@ -139,7 +140,7 @@ class ScriptedHttpServer:
f"{LISTENER_ACCEPT_TIMEOUT_S}s"
)
ready = self._socket.recv_pyobj()
ready = sock_recv(self._socket)
if not isinstance(ready, HookReady):
raise RuntimeError(
f"ScriptedHttpServer: expected HookReady handshake, got {ready!r}"