feat(zmq): teach the direct-backend path to speak SGLang

bindings/python/src/lib.rs

@@ -753,6 +753,7 @@ impl Router {
                 match self.backend {
                     BackendType::Vllm => Some(worker::RuntimeType::Vllm),
                     BackendType::Tokenspeed => Some(worker::RuntimeType::TokenSpeed),
+                    BackendType::Sglang => Some(worker::RuntimeType::Sglang),
                     _ => None,
                 }
             } else {

bindings/python/src/smg/_sglang_zmq_launcher.py

@@ -0,0 +1,121 @@
+"""Headless SGLang scheduler launcher for the SMG direct-ZMQ backend.
+
+On the direct-ZMQ path SMG plays the SGLang TokenizerManager role itself: it
+binds the PUSH input and PULL output ``ipc://`` sockets, sends tokenized
+generate requests, and receives batched token outputs straight off the wire.
+This module therefore launches *only* the scheduler process(es) — no
+TokenizerManager, no DetokenizerManager — wired to SMG's two sockets through a
+hand-built ``PortArgs``, with ``skip_tokenizer_init`` so the scheduler runs
+tokenizer-free and routes its outputs back over the tokenizer socket (which SMG
+owns).
+
+Unmodified SGLang exposes no CLI flag or env var to point a bare scheduler at
+externally-chosen sockets, so SMG drives SGLang's own ``PortArgs`` /
+``Engine._launch_scheduler_processes`` primitives directly. Every argument other
+than the two SMG-owned endpoints (``--smg-input-ipc`` / ``--smg-output-ipc``) is
+a native SGLang server argument forwarded verbatim.
+
+The scheduler defaults to pickle over ZMQ; SMG only decodes msgpack, so
+``SGLANG_USE_PICKLE_IPC`` is forced off before SGLang is imported (the flag is
+snapshotted at import time).
+
+Invoked as a subprocess by the ``sglang`` launcher in ``serve.py``.
+"""
+
+import os
+
+# Must precede any sglang import: io_struct snapshots this at module load to
+# choose msgpack vs pickle for the ZMQ wire, and SMG only speaks msgpack. Hard
+# override (not setdefault): an inherited "1" would silently pickle the wire.
+os.environ["SGLANG_USE_PICKLE_IPC"] = "0"
+
+import logging  # noqa: E402
+import sys  # noqa: E402
+
+logger = logging.getLogger("smg.sglang_zmq_launcher")
+
+_INPUT_FLAG = "--smg-input-ipc"
+_OUTPUT_FLAG = "--smg-output-ipc"
+
+
+def _extract_flag(argv: list[str], flag: str) -> tuple[str, list[str]]:
+    """Pop ``flag VALUE`` or ``flag=VALUE`` from argv; return (value, rest).
+
+    These are SMG-owned endpoints, not SGLang args, so they are stripped before
+    the remainder is handed to SGLang's own parser.
+    """
+    for i, token in enumerate(argv):
+        if token == flag:
+            if i + 1 >= len(argv):
+                raise SystemExit(f"{flag} requires a value")
+            return argv[i + 1], argv[:i] + argv[i + 2 :]
+        if token.startswith(flag + "="):
+            return token.split("=", 1)[1], argv[:i] + argv[i + 1 :]
+    raise SystemExit(f"{flag} is required")
+
+
+def main(argv: list[str] | None = None) -> None:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    # scheduler_input_ipc_name: SMG PUSHes requests here, the scheduler PULLs.
+    input_ipc, argv = _extract_flag(argv, _INPUT_FLAG)
+    # tokenizer_ipc_name: under skip_tokenizer_init the scheduler PUSHes outputs
+    # here, and SMG PULLs them.
+    output_ipc, argv = _extract_flag(argv, _OUTPUT_FLAG)
+
+    from sglang.srt.entrypoints.engine import (
+        Engine,
+        _set_envs_and_config,
+        _set_gc,
+    )
+    from sglang.srt.managers.scheduler import run_scheduler_process
+    from sglang.srt.plugins import load_plugins
+    from sglang.srt.server_args import PortArgs, prepare_server_args
+    from sglang.srt.utils import configure_logger
+
+    # SMG tokenizes upstream and drives the tokenizer<->scheduler ZMQ wire
+    # directly, so the scheduler must run without its own tokenizer. Pass this as
+    # a CLI flag: SGLang resolves and freezes server_args in prepare_server_args,
+    # so it can no longer be set afterward.
+    if "--skip-tokenizer-init" not in argv:
+        argv.append("--skip-tokenizer-init")
+    server_args = prepare_server_args(argv)
+
+    # Mirror the head of Engine._launch_subprocesses so the scheduler subprocess
+    # inherits the same logging, env/config, plugins, and mp start method it
+    # would under a normal SGLang launch.
+    configure_logger(server_args)
+    _set_envs_and_config(server_args)
+    load_plugins()
+    server_args.check_server_args()
+    _set_gc(server_args)
+
+    # Reuse SGLang's own port derivation (nccl port, instance id, the unused
+    # detokenizer/rpc/metrics ipc names), then repoint the two SMG-owned sockets
+    # at the endpoints SMG binds.
+    port_args = PortArgs.init_new(server_args)
+    port_args.scheduler_input_ipc_name = input_ipc
+    port_args.tokenizer_ipc_name = output_ipc
+
+    result, procs = Engine._launch_scheduler_processes(
+        server_args, port_args, run_scheduler_process
+    )
+    # Block until every scheduler rank has finished loading; the SMG router
+    # gates request admission on its own ZMQ readiness, so once the schedulers
+    # are up this process just keeps them alive.
+    result.wait_for_ready()
+    logger.info(
+        "SGLang scheduler(s) ready on input=%s output=%s; SMG owns the tokenizer ZMQ wire",
+        input_ipc,
+        output_ipc,
+    )
+    result.block_until_scheduler_exits()
+
+    # A scheduler exiting means the engine is gone; propagate a non-zero status
+    # (rather than a silent exit 0) so the parent launcher sees the failure
+    # instead of leaving the router pushing to a dead socket.
+    if any(proc.exitcode for proc in procs):
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

bindings/python/src/smg/serve.py

@@ -37,6 +37,17 @@ def _zmq_ipc_url(port: int) -> str:
     return f"ipc://{_ZMQ_SOCKET_DIR}/engine-{port}"
 
 
+def _zmq_data_addresses(ipc_url: str) -> tuple[str, str]:
+    """The ipc:// data-plane sockets SMG binds for a worker: (input, output).
+
+    Mirrors `zmq_socket_addresses` in model_gateway/src/worker/worker.rs: SMG
+    binds `<path>-in.sock` for requests and `<path>-out.sock` for outputs. An
+    engine that connects directly (no TCP handshake, e.g. SGLang) must dial these
+    exact endpoints.
+    """
+    return f"{ipc_url}-in.sock", f"{ipc_url}-out.sock"
+
+
 def _zmq_handshake_port(ipc_url: str) -> int:
     """The tcp handshake port SMG derives from an ipc:// URL.
 
@@ -189,11 +200,23 @@ class SglangWorkerLauncher(WorkerLauncher):
     """Launcher for sglang inference workers."""
 
     def _get_tp_size(self, args: argparse.Namespace) -> int:
-        return getattr(args, "tensor_parallel_size", 1)
+        # sglang's ServerArgs registers --tp-size with dest `tp_size` (the alias
+        # --tensor-parallel-size shares that dest), so read tp_size first, then
+        # fall back to the alias and finally a single-GPU default. The default
+        # applies only when neither attribute is set; an explicit value (even a
+        # non-positive one) is preserved so gpu_env can reject it.
+        if hasattr(args, "tp_size"):
+            return int(args.tp_size)
+        if hasattr(args, "tensor_parallel_size"):
+            return int(args.tensor_parallel_size)
+        return 1
 
     def build_command(
         self, args: argparse.Namespace, backend_args: list[str], host: str, port: int
     ) -> list[str]:
+        if getattr(args, "connection_mode", "grpc") == "zmq":
+            return self._build_zmq_command(args, backend_args, port)
+
         cmd = [
             sys.executable,
             "-m",
@@ -216,6 +239,37 @@ def build_command(
         cmd.extend(self._filter_backend_args(backend_args, ["--model-path", "--host", "--port"]))
         return cmd
 
+    def _build_zmq_command(
+        self, args: argparse.Namespace, backend_args: list[str], port: int
+    ) -> list[str]:
+        """Launch a headless SGLang scheduler wired to SMG's ZMQ sockets.
+
+        Unmodified SGLang exposes no CLI to point a bare scheduler at
+        externally-chosen sockets, so this dispatches the `_sglang_zmq_launcher`
+        module, which drives SGLang's own scheduler-launch primitives on the two
+        ipc endpoints SMG binds. Each worker is one standalone scheduler; running
+        several is dense data parallelism as N independent ZMQ workers.
+        """
+        input_ipc, output_ipc = _zmq_data_addresses(_zmq_ipc_url(port))
+        cmd = [
+            sys.executable,
+            "-m",
+            "smg._sglang_zmq_launcher",
+            "--smg-input-ipc",
+            input_ipc,
+            "--smg-output-ipc",
+            output_ipc,
+            "--model-path",
+            getattr(args, "model_path", ""),
+        ]
+        cmd.extend(
+            self._filter_backend_args(
+                backend_args,
+                ["--smg-input-ipc", "--smg-output-ipc", "--model-path", "--host", "--port"],
+            )
+        )
+        return cmd
+
 
 class VllmWorkerLauncher(WorkerLauncher):
     """Launcher for vLLM inference workers."""
@@ -654,7 +708,7 @@ def add_serve_args(parser: argparse.ArgumentParser) -> None:
         help=(
             "Connection mode for workers (default: grpc). Note: trtllm only "
             "supports grpc, tokenspeed only supports zmq, and zmq is otherwise "
-            "only supported for the vllm backend"
+            "only supported for the vllm and sglang backends"
         ),
     )
     # Router host/port - may be overridden by backend (e.g. sglang)
@@ -735,10 +789,14 @@ def parse_serve_args(
 
     # ZMQ direct-backend is a same-host engine connection; only vLLM EngineCore
     # and TokenSpeed speak a supported ZMQ wire protocol.
-    if serve_router_args.connection_mode == "zmq" and backend not in ("vllm", "tokenspeed"):
+    if serve_router_args.connection_mode == "zmq" and backend not in (
+        "vllm",
+        "tokenspeed",
+        "sglang",
+    ):
         pre_parser.error(
-            "connection-mode zmq is only supported for the vllm and tokenspeed "
-            f"backends, not {backend}"
+            "connection-mode zmq is only supported for the vllm, tokenspeed, and "
+            f"sglang backends, not {backend}"
         )
 
     # Pass 2: full parser with backend-specific args; resolve so backend can override
@@ -760,6 +818,12 @@ def parse_serve_args(
     else:
         args = parser.parse_args(argv)
 
+    # Some backends declare data parallelism under their own dest (sglang uses
+    # dp_size), which resolves over the serve-level --data-parallel-size. Give
+    # the orchestrator one canonical field regardless of backend.
+    if not hasattr(args, "data_parallel_size"):
+        args.data_parallel_size = getattr(args, "dp_size", 1)
+
     return backend, args, backend_args
 
 

crates/engine_zmq_client/examples/live_probe.rs

@@ -65,13 +65,17 @@ async fn main() {
     .expect("handshake with engine");
 
     let engine = &transport.engines[0];
+    let ready = engine
+        .ready_response
+        .as_ref()
+        .expect("handshake engine registers a ready response");
     eprintln!(
         "[probe] engine registered: vllm_version={} max_model_len={} num_gpu_blocks={} block_size={} dtype={} engine_index={:?}",
-        engine.ready_response.vllm_version,
-        engine.ready_response.max_model_len,
-        engine.ready_response.num_gpu_blocks,
-        engine.ready_response.block_size,
-        engine.ready_response.dtype.as_str(),
+        ready.vllm_version,
+        ready.max_model_len,
+        ready.num_gpu_blocks,
+        ready.block_size,
+        ready.dtype.as_str(),
         engine.engine_id.engine_index(),
     );