手写 mini-sglang 阶段三 · 服务化：变成一个服务

消息、序列化与 ZMQ¶

前两部分的所有代码都在一个进程里运行。要变成在线服务，API Server、tokenizer、各个调度器进程之间需要传递消息。mini-sglang 的做法很朴素：每种消息是一个 dataclass，用一个 70 行的通用序列化函数变成由基本类型组成的字典，再用 msgpack 变成字节，通过 ZMQ 的几种 socket 发送。这一章把这套机制写出来。

自测：能答上来就可以跳过本章

系统里一共有哪几类消息？分别从哪个进程发往哪个进程？
张量（提示词的 token）怎样放进 msgpack？
ZMQ 的 PUSH/PULL 和 PUB/SUB 有什么区别？调度器分别在哪里用它们？
为什么 rank 0 转发给其他 rank 的是原始字节，而不是解码后的消息？

自测参考答案（先自己答，再展开对照）

三组，按接收方划分：发给 tokenizer / detokenizer 的（TokenizeMsg 来自 API Server，DetokenizeMsg 是调度器发出的新 token，还有 AbortMsg）；发给调度器的（tokenizer 发来的 UserMsg、AbortBackendMsg、ExitMsg）；发给 API Server 的（detokenizer 发来的 UserReply）。每组还有一个把多条打包在一起的 Batch...Msg，各组在自己的模块里按类名反序列化。
1 维张量存成原始字节（加上 dtype 等信息），放进 msgpack 的 bytes 字段；接收方用这些字节重新构造张量。
PUSH / PULL 是点对点的管道（消息被一个接收者取走）；PUB / SUB 是广播（每个订阅者都收到一份）。调度器用 PULL 接收上游的请求、PUSH 发出结果，rank 0 用 PUB 把收到的消息广播给其他 rank。
原样转发原始字节，保证所有 rank 看到完全相同的消息，省掉一次解码再编码，也避免了重新编码可能带来的差异。

本章要写的文件：message/utils.py、message/backend.py、message/tokenizer.py、message/frontend.py、message/__init__.py、utils/mp.py。

到这一章为止的文件：63 个（新建 6，修改 0，提前引入 0）

minisgl/
├── __init__.py
├── core.py
├── env.py
├── attention/  （4 个文件，不变）
├── distributed/  （3 个文件，不变）
├── engine/  （5 个文件，不变）
├── kernel/  （2 个文件，不变）
├── kvcache/  （5 个文件，不变）
├── layers/  （8 个文件，不变）
├── llm/  （2 个文件，不变）
├── message/
│   ├── __init__.py  ← 本章新建
│   ├── backend.py  ← 本章新建
│   ├── utils.py  ← 本章新建
│   ├── tokenizer.py  ← 本章新建
│   ├── frontend.py  ← 本章新建
├── models/  （7 个文件，不变）
├── moe/  （2 个文件，不变）
├── scheduler/  （9 个文件，不变）
├── utils/
│   ├── __init__.py
│   ├── misc.py
│   ├── hf.py
│   ├── torch_utils.py
│   ├── registry.py
│   ├── logger.py
│   ├── device.py
│   ├── mp.py  ← 本章新建

这一步的 main：examples/ch12_message.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。

三组消息¶

按接收方分成三组，每组有一个基类：

基类	接收方	消息
`BaseTokenizerMsg`	tokenizer / detokenizer	`TokenizeMsg`（文本或对话 + 采样参数）、`DetokenizeMsg`（一个新 token + 是否结束）、`AbortMsg`
`BaseBackendMsg`	调度器	`UserMsg`（token 张量 + 采样参数）、`AbortBackendMsg`、`ExitMsg`
`BaseFrontendMsg`	API Server	`UserReply`（增量文本 + 是否结束）

每组还有一个 Batch...Msg，把多条消息打包成一条发送，减少消息数。

minisgl/message/backend.py

"""发往调度器（后端）的消息。"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List

import torch
from minisgl.core import SamplingParams

from .utils import deserialize_type, serialize_type


@dataclass
class BaseBackendMsg:
    def encoder(self) -> Dict:
        return serialize_type(self)

    @staticmethod
    def decoder(json: Dict) -> BaseBackendMsg:
        return deserialize_type(globals(), json)


@dataclass
class BatchBackendMsg(BaseBackendMsg):
    data: List[BaseBackendMsg]


@dataclass
class ExitMsg(BaseBackendMsg):
    pass


@dataclass
class UserMsg(BaseBackendMsg):
    uid: int
    input_ids: torch.Tensor  # CPU 上的 1 维 int32 张量
    sampling_params: SamplingParams


@dataclass
class AbortBackendMsg(BaseBackendMsg):
    uid: int

注意 decoder 里的 globals()：反序列化时按类名在定义消息的那个模块的全局命名空间里找类。这就是为什么三组消息分在三个文件里、各自有自己的 decoder——每个接收方只认识发给自己的那组消息。

通用序列化¶

minisgl/message/utils.py

"""消息的序列化：dataclass ↔ 由基本类型组成的 dict，再交给 msgpack 变成字节。

每个对象序列化成 {"__type__": 类名, 字段...}；1 维张量存成原始字节加 dtype。
反序列化时按类名在给定的类表（通常是定义消息的模块的 globals()）里找到类，递归还原字段。
"""

from __future__ import annotations

from typing import Any, Dict, Type

import numpy as np
import torch


def _serialize_any(value: Any) -> Any:
    if isinstance(value, dict):
        return {k: _serialize_any(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return type(value)(_serialize_any(v) for v in value)
    if isinstance(value, (int, float, str, type(None), bool, bytes)):
        return value
    return serialize_type(value)


def serialize_type(obj: Any) -> Dict:
    if isinstance(obj, torch.Tensor):
        assert obj.dim() == 1, "only 1D tensors are supported"
        return {"__type__": "Tensor", "buffer": obj.numpy().tobytes(), "dtype": str(obj.dtype)}
    result: Dict[str, Any] = {"__type__": type(obj).__name__}
    for k, v in obj.__dict__.items():
        result[k] = _serialize_any(v)
    return result


def _deserialize_any(cls_map: Dict[str, Type], data: Any) -> Any:
    if isinstance(data, dict):
        if "__type__" in data:
            return deserialize_type(cls_map, data)
        return {k: _deserialize_any(cls_map, v) for k, v in data.items()}
    if isinstance(data, (list, tuple)):
        return type(data)(_deserialize_any(cls_map, d) for d in data)
    return data


def deserialize_type(cls_map: Dict[str, Type], data: Dict) -> Any:
    type_name = data["__type__"]
    if type_name == "Tensor":
        np_dtype = getattr(np, data["dtype"].replace("torch.", ""))
        return torch.from_numpy(np.frombuffer(data["buffer"], dtype=np_dtype).copy())
    kwargs = {k: _deserialize_any(cls_map, v) for k, v in data.items() if k != "__type__"}
    return cls_map[type_name](**kwargs)

规则很简单：任何对象变成 {"__type__": 类名, 字段名: 值...}，递归处理字段；基本类型原样保留；1 维张量变成原始字节加 dtype 字符串。反序列化反过来：看到 __type__ 就按类名构造对象。

examples/ch12_message.py

"""第 12 章：一条消息从对象到字节再回到对象；ZMQ 的 PUSH/PULL 与 PUB/SUB。"""

import threading

import msgpack
import torch
from minisgl.core import SamplingParams
from minisgl.message import BaseBackendMsg, BatchBackendMsg, UserMsg
from minisgl.utils import ZmqPubQueue, ZmqPullQueue, ZmqPushQueue, ZmqSubQueue

msg = BatchBackendMsg(data=[
    UserMsg(uid=7, input_ids=torch.tensor([9707, 11, 1879], dtype=torch.int32),
            sampling_params=SamplingParams(temperature=0.6, max_tokens=64)),
])
d = msg.encoder()
print("encoder() 得到的字典:")
print(" ", {k: v for k, v in d.items() if k != "data"}, "data[0] =")
inner = dict(d["data"][0])
print("   ", inner)
raw = msgpack.packb(d, use_bin_type=True)
print(f"msgpack 之后 {len(raw)} 字节")
back = BaseBackendMsg.decoder(msgpack.unpackb(raw, raw=False))
print("解码回来:", back.data[0].uid, back.data[0].input_ids, back.data[0].sampling_params)

# PUSH/PULL：一对一的队列；PUB/SUB：一对多的广播
addr = "ipc:///tmp/minisgl_example_pushpull"
pull = ZmqPullQueue(addr, create=True, decoder=BaseBackendMsg.decoder)
push = ZmqPushQueue(addr, create=False, encoder=BaseBackendMsg.encoder)
threading.Thread(target=lambda: [push.put(UserMsg(uid=i, input_ids=torch.tensor([i], dtype=torch.int32),
                                                  sampling_params=SamplingParams())) for i in range(3)]).start()
print("PULL 收到 uid:", [pull.get().uid for _ in range(3)], " 队列已空:", pull.empty())

paddr = "ipc:///tmp/minisgl_example_pubsub"
pub = ZmqPubQueue(paddr, create=True, encoder=BaseBackendMsg.encoder)
subs = [ZmqSubQueue(paddr, create=False, decoder=BaseBackendMsg.decoder) for _ in range(2)]
pub.wait_for_subscribers(2)  # 订阅生效之前发布的消息会被丢弃，所以先等两个订阅者到齐
pub.put_raw(raw)  # rank 0 把收到的原始字节原样转发，不用重新编码
print("两个 SUB 都收到了:", [s.get().data[0].uid for s in subs])
for q in [pull, push, pub, *subs]:
    q.stop()

运行结果

encoder() 得到的字典:
  {'__type__': 'BatchBackendMsg'} data[0] =
    {'__type__': 'UserMsg', 'uid': 7, 'input_ids': {'__type__': 'Tensor', 'buffer': b'\xeb%\x00\x00\x0b\x00\x00\x00W\x07\x00\x00', 'dtype': 'torch.int32'}, 'sampling_params': {'__type__': 'SamplingParams', 'temperature': 0.6, 'top_k': -1, 'top_p': 1.0, 'ignore_eos': False, 'max_tokens': 64}}
msgpack 之后 229 字节
解码回来: 7 tensor([9707,   11, 1879], dtype=torch.int32) SamplingParams(temperature=0.6, top_k=-1, top_p=1.0, ignore_eos=False, max_tokens=64)
PULL 收到 uid: [0, 1, 2]  队列已空: True
两个 SUB 都收到了: [7, 7]

一条带 3 个 token 的请求消息大约 200 多字节，其中大部分是字段名。字段名在每条消息里重复出现，这是通用序列化的代价；换来的是新增一种消息只需要写一个 dataclass，不用写任何编解码代码。提示词的 token 以原始字节存放（int32 每个 4 字节），一个 8000 token 的提示词也只有 32 KB。

ZMQ 队列¶

minisgl/utils/mp.py · ZmqPushQueue

class ZmqPushQueue(Generic[T]):
    def __init__(self, addr: str, create: bool, encoder: Callable[[T], Dict]):
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.PUSH)
        self.socket.bind(addr) if create else self.socket.connect(addr)
        self.encoder = encoder

    def put(self, obj: T) -> None:
        self.socket.send(msgpack.packb(self.encoder(obj), use_bin_type=True), copy=False)

    def stop(self) -> None:
        self.socket.close(linger=0)
        self.context.term()

utils/mp.py 把 ZMQ socket 包成五种队列，都是"编码器 + msgpack + socket"的薄封装：

PUSH / PULL：一对一（或多对一）的管道，消息按顺序到达、不会丢失。tokenizer → 调度器、调度器 → detokenizer、detokenizer → API Server 都是这种；
异步版本：API Server 运行在 asyncio 事件循环里，用 zmq.asyncio 的 socket，await 收发不阻塞事件循环；
PUB / SUB：一对多的广播，rank 0 调度器用它把消息转发给其他 rank。

create=True 的一方 bind（占住地址），另一方 connect。地址是 ipc:///tmp/minisgl_N.pid=启动进程的 PID，同一台机器上起多个服务也不会冲突。

PUB/SUB 有一个经典陷阱——"慢订阅者"（slow joiner）：订阅在连接建立之后才异步生效，在此之前发布的消息会被直接丢弃。如果 rank 0 的第一条广播赶在 rank 1 的订阅生效之前，rank 1 就会永远等下去。官方靠启动时序躲开了它（调度器初始化完还要等 tokenizer、API Server 启动，时间很充裕），但这只是时间差，不是保证。我们把发布端换成 XPUB：它能收到"有订阅者订阅了"的通知，wait_for_subscribers(n) 等 n 个订阅者都到齐再返回：

minisgl/utils/mp.py · ZmqPubQueue

XPUB_VERBOSE 让每个订阅者的订阅通知都送达（默认会把相同主题的重复订阅合 an>

class ZmqPubQueue(Generic[T]): """广播队列。用 XPUB 而不是 PUB：XPUB 能收到"有订阅者订阅了"的通知，于是可以等订阅者到齐再发。 class="sd">    PUB/SUB 的一个经典陷阱（"慢订阅者"）：订阅在连接建立后才异步生效，在此之前发布的消息会被直接丢弃。 class="sd">    """ def __init__(self, addr: str, create: bool, encoder: Callable[[T], Dict]): self.context = zmq.Context() self.socket = self.context.socket(zmq.XPUB) self.socket.setsockopt(zmq.XPUB_VERBOSE, 1)  # 每个订阅者的订阅通知都要，不去重 self.socket.bind(addr) if create else self.socket.connect(addr) self.encoder = encoder def wait_for_subscribers(self, n: int) -> None: for _ in range(n): self.socket.recv()  # 订阅通知：b"\x01" + 主题 def put_raw(self, raw: bytes) -> None: self.socket.send(raw, copy=False) def put(self, obj: T) -> None: self.socket.send(msgpack.packb(self.encoder(obj), use_bin_type=True), copy=False) def stop(self) -> None: self.socket.close(linger=0) self.context.term() 并成一条）。
为什么转发原始字节¶
rank 0 从 tokenizer 收到消息后，要原样转发给其他 rank。ZmqPullQueue 提供了 get_raw() 和 decode() 两个方法：rank 0 先取原始字节、原样 put_raw 给其他 rank，再自己解码。这样省掉了一次"解码再编码"，更重要的是保证所有 rank 收到的是逐字节相同的消息——各 rank 必须做出完全相同的调度决策（第 14 章）。

官方实现

序列化：message/utils.py · serialize_type（第 20～35 行）
队列：utils/mp.py · ZmqPushQueue（第 12～30 行）
官方的 tests/misc/test_serialize.py 测试了序列化的往返

我们的队列在 stop() 时使用 close(linger=0)：官方用默认的 close()，如果还有未送出的消息，进程退出时可能会卡住等待。


与官方的差异：XPUB
官方的 ZmqPubQueue 用普通的 PUB socket，没有等待订阅者。本书第 14 章的演示程序（主进程启动两个 rank 后立即发消息）用官方的写法，在本书的机器上每次都会卡死：rank 0 收到第一条消息后广播、然后在"广播条数"的集合通信里等 rank 1，而 rank 1 永远收不到那条被丢弃的广播。换成 XPUB 并等待订阅者后就不会了。tests/test_ch12_message.py 里有对应的测试。

测试¶
tests/test_ch12_message.py · test_zmq_push_pull_over_ipcdef test_zmq_push_pull_over_ipc(ipc_dir):
    addr = f"ipc://{ipc_dir}/q"
    pull = ZmqPullQueue(addr, create=True, decoder=BaseTokenizerMsg.decoder)
    push = ZmqPushQueue(addr, create=False, encoder=BaseTokenizerMsg.encoder)
    msgs = [TokenizeMsg(uid=1, text=[{"role": "user", "content": "你好"}], sampling_params=SamplingParams()),
            DetokenizeMsg(uid=1, next_token=42, finished=True)]
    threading.Thread(target=lambda: [push.put(m) for m in msgs]).start()
    got = [pull.get(), pull.get()]
    assert got == msgs and pull.empty()
    push.stop()
    pull.stop()


ipc_dir 而不是 pytest 的 tmp_path
Unix domain socket 的路径存在 sockaddr_un.sun_path 里，长度有硬上限：Linux 108 字节，macOS 只有 104。pytest 的 tmp_path 在 macOS 上形如 /private/var/folders/9k/3m_.../T/pytest-of-you/pytest-12/test_zmq_push_pull_over_ip0/，光目录就一百多字节，拿它拼 socket 路径会被 ZMQ 直接拒掉（ipc path ... is longer than 103 characters）。所以 conftest.py 里的 ipc_dir 固定在 /tmp 下开一个短目录。真实的推理引擎也绕不开这一条——本书的 ipc:///tmp/minisgl_0.pid=... 之所以这么短，正是为了留出余量。


怎么讲清楚
讲进程间通信：三组消息——API Server 发给 tokenizer（请求）、tokenizer 发给调度器（token 化后的请求）、调度器经 detokenizer 发回 API Server（结果）；通用序列化把对象变成 {"__type__": 类名, ...} 再用 msgpack 编码，一维张量直接存原始字节，比 pickle 快也更安全。ZMQ 的 PUSH/PULL 是点对点管道，PUB/SUB 是广播（订阅者要先连上，否则会丢早期消息）；rank 0 收到请求后把原始字节原样广播给其他 rank，保证所有 rank 看到完全相同的消息。追问"为什么不用 NCCL 传控制消息"：控制消息小而不规则，NCCL 只适合张量。

练习¶

序列化只支持 1 维张量。如果要传一个 [batch, seq] 的二维张量，需要改哪几行？
通用序列化每条消息都带着字段名。设计一种更紧凑的格式（例如按字段顺序存成列表），它的代价是什么？
为什么 API Server 用异步版本的队列，而调度器和 tokenizer 用同步版本？


参考答案

序列化时额外存 shape（list(obj.shape)），反序列化时 np.frombuffer(...).reshape(shape)；去掉 dim() == 1 的断言。
收发双方必须对字段顺序达成一致，加字段、删字段都要同步修改两边，且版本不同的进程无法通信；可读性和调试也变差。msgpack 已经很紧凑，瓶颈通常不在这里。
API Server 在一个事件循环里同时处理成百上千个 HTTP 连接，任何阻塞调用都会卡住所有连接；调度器和 tokenizer 是单线程的循环，本来就在"收消息 → 处理"之间轮转，同步调用更简单。


小结¶

 三组消息分别发往 tokenizer、调度器、API Server，每组按类名在自己的模块里反序列化。
 通用序列化：对象 → {"__type__": 类名, ...} → msgpack 字节；1 维张量存原始字节。
 PUSH/PULL 做点对点管道，PUB/SUB 做 rank 0 到其他 rank 的广播；订阅者必须先连上。
 rank 0 原样转发原始字节，保证所有 rank 看到相同的消息。


本章练习题（1 道）
在浏览器里直接写代码、跑测试；也可以在本地用 python practice/judge.py 判题。

196. 进程间消息的序列化（中等）









  
  






                
  
  
    
    发现错误或有更好的写法？在 GitHub 上编辑此页 · 提交问题