Prhub

#27145 fix(load-snapshot): avoid duplicate zmq bind in multi-tokenizer mode

原始 PR 作者 whybeyoung 合并时间 2026-06-04 09:24 文件变更 5 提交数 2 评论 8 代码增减 +157 / -19

执行摘要

修复多 tokenizer 模式 ZMQ 绑定冲突

Issue #27142 报告在 multi-tokenizer 模式下服务器启动失败,错误为 'Address already in use'。根本原因是引入 ZMQ 负载快照传输后,没有考虑到多 tokenizer 场景,每个 TokenizerWorker 都试图绑定同一个 ZMQ PULL 端点。本变更旨在确保同一节点上只有一个进程绑定 ZMQ 套接字,其他进程通过共享内存读取数据。

此 PR 值得精读,尤其是多进程通信中唯一所有者模式的实现、事件驱动的 fd 注册方式以及如何处理异步/同步上下文冲突。对于需要在多 tokenizer 环境下使用负载均衡的团队至关重要。

讨论亮点

review 中主要讨论:

  1. gemini-code-assist[bot] 指出异步 ZMQ 上下文冲突风险,建议重置 PyZMQ 上下文单例。
  2. merrymercy 要求统一 caller 命名风格,避免使用 getattr。这些建议均已在后续修正中落实。

实现拆解

按步骤描述:

  1. 在 load_snapshot.py 中添加 _tokenizer_load_snapshot_owner_caller 函数,根据 tokenizer_worker_num 返回 'MultiTokenizerRouter' 或 'TokenizerManager'。
  2. 修改 zmq_reader_owner 函数,引入 tokenizer_owner 变量,支持新的 caller 类型,更新文档字符串。
  3. 在 multi_tokenizer_mixin.py 中,MultiTokenizerRouter 初始化时根据 zmq_reader_owner 决定是否创建 ZMQ 读取器,并通过 _register_load_snapshot_reader 将读取器的 fd 注册到事件循环,实现事件驱动的轮询。
  4. 统一 caller 命名规范:将 'tokenizer' 改为 'TokenizerManager','dp_controller' 改为 'DataParallelController',涉及 data_parallel_controller.py 和 tokenizer_manager.py。
  5. 添加 TestZmqReaderOwner 单元测试类,覆盖多种配置下唯一 owner 的验证。
文件 模块 状态 重要度
python/sglang/srt/managers/load_snapshot.py 负载快照 modified 7.97
python/sglang/srt/managers/multi_tokenizer_mixin.py 多 tokenizer 路由 modified 6.97
test/registered/unit/managers/test_load_snapshot_backends.py 单元测试 modified 6.82
python/sglang/srt/managers/data_parallel_controller.py 数据并行控制 modified 4.49
python/sglang/srt/managers/tokenizer_manager.py Tokenizer 管理 modified 4.49

关键符号

_tokenizer_load_snapshot_owner_caller zmq_reader_owner _register_load_snapshot_reader fileno poll

关键源码片段

python/sglang/srt/managers/multi_tokenizer_mixin.py core-logic

在 MultiTokenizerRouter 中集成 ZMQ 读取器,实现事件驱动轮询。

# python/sglang/srt/managers/multi_tokenizer_mixin.py (MultiTokenizerRouter 部分 )def __init__(self, server_args, port_args):
    # ...
    self.load_snapshot_reader = None
    if zmq_reader_owner(server_args, 'MultiTokenizerRouter'):
        self.load_snapshot_reader = create_load_snapshot_reader(
            server_args, port_args, caller='MultiTokenizerRouter'
        )
        self._loop.call_soon_threadsafe(self._register_load_snapshot_reader)
    # ...def _register_load_snapshot_reader(self):
    assert self.load_snapshot_reader is not None
    self._loop.add_reader(
        self.load_snapshot_reader.fileno(),
        self.load_snapshot_reader.poll
    )
    self.load_snapshot_reader.poll()

评论区精华

异步 ZMQ 上下文冲突 正确性

gemini-code-assist[bot] 指出 MultiTokenizerRouter 使用的异步 ZMQ 上下文污染了全局单例,导致同步 load_snapshot_reader 收到 Future 对象,建议在创建 reader 前重置 PyZMQ 上下文单例。

结论:PR 作者确认已处理,最终代码中通过重新设计 reader 创建或上下文隔离解决了此问题。 · 已解决

统一命名规范 style

merrymercy 要求将所有 caller 名称从 'router'、'tokenizer'、'dp_controller' 改为 'MultiTokenizerRouter'、'TokenizerManager'、'DataParallelController'。

结论:已按要求修改,最终代码中使用新命名。 · 已解决

避免 getattr 正确性

merrymercy 指出不应使用 getattr 访问 server_args 字段,应直接访问以确保字段存在。

结论:已移除 getattr 调用,改用直接属性访问。 · 已解决

风险与影响

主要风险:

  1. 异步 ZMQ 上下文污染(可能已通过重置或内部上下文管理解决)。
  2. 事件驱动轮询可能引入竞态,但 zmq PULL socket 的 edge-triggered fd 特性可保证安全。
  3. 多 tokenizer 模式配置组合较多,新增的单元测试覆盖了主要场景,但仍可能存在边界情况(如 dp_size > 1 且 tokenizer_worker_num > 1 且非负载感知方法)。
  4. 与 future PR 的负载快照功能可能存在交互,需注意兼容性。

对用户:修复了 multi-tokenizer 模式下服务器启动崩溃,使该配置可用。对系统:引入了事件驱动读取机制,减少了轮询开销。对团队:统一了负载快照的所有权模型,为后续扩展奠定了基础。影响范围:仅在启用 ZMQ 传输(多节点 DP 或强制环境变量)且 tokenizer_worker_num > 1 时触发,属于非默认路径。

ZMQ 上下文冲突 事件驱动回归 多 tokenizer 边界情况

关联 Issue

#27142 [Bug] Multi-tokenizer mode crashes on startup with "Address already in use" when zmq load snapshots are enabled

完整报告

参与讨论