Prhub

#32403 [Mooncake] Fix ProcessGroup API imports

原始 PR 作者 UNIDY2002 合并时间 2026-08-03 16:16 文件变更 2 提交数 1 评论 7 代码增减 +19 / -17

执行摘要

Mooncake PG 导入路径迁移,统一从 mooncake.pg 引用

PR body 指出:"We historically imports several Mooncake ProcessGroup APIs through mooncake.ep. Mooncake keeps that path working for compatibility, but these APIs are part of the ProcessGroup interface and should be imported directly from mooncake.pg." 即这些 API 在 Mooncake 中本就属于 ProcessGroup 接口,mooncake.ep 只是为兼容保留的历史路径。为遵循上游接口规范、减少对兼容层的依赖,需要将导入统一迁移到 mooncake.pg

值得快速浏览,不建议精读。该 PR 是典型的第三方依赖接口规范跟进:API 未变、仅换导入路径,核心学习点是 _wait_for_peer_state 借机简化签名、去掉多余的模块参数传递。若团队正在使用 Mooncake 后端,建议关注 CI 中 mooncake 相关测试是否真实跑过 Elastic EP 路径,并确认依赖版本约束;若未使用 Mooncake,可忽略。

讨论亮点

该 PR 没有实质性的 review 讨论线程。唯一的审核人 ch-wan 直接 APPROVED(未填写评论)。Issue 评论中除作者多次触发 /tag-and-rerun-ci 重跑 CI 外,只有 gemini-code-assist[bot] 提示消费者版 Gemini Code Assist 已停止服务、不再产生代码评审活动。整体是一个低争议的机械改动。

实现拆解

  1. Elastic EP 恢复路径导入迁移python/sglang/srt/elastic_ep/elastic_ep.py):将 _try_recover_world_wait_for_peer_statetry_recover_ranks_join_world_groupjoin_process_groups 中原本的 from mooncake import ep as mooncake_ep 改为按需从 mooncake.pg 导入 get_peer_staterecover_ranksjoin_group,并同步调整调用方式。
  2. _wait_for_peer_state 签名简化:该函数原接收 mooncake_ep 模块对象作为第一参数,改为从函数体内直接导入 get_peer_state 后,签名简化为 (backend, ranks),删除了不必要的模块参数传递,使调用点 _wait_for_peer_state(group.device_group, local_ranks) 更简洁。
  3. 进程组创建路径导入迁移python/sglang/srt/distributed/parallel_state.py):GroupCoordinator.__init__ 中 mooncake 分支和 init_distributed_environmentbackend == "mooncake" 分支的 from mooncake.ep import MooncakeBackendOptions 一并改为 from mooncake.pg import MooncakeBackendOptions
  4. 测试与文档配套:无新增测试、文档或配置变更——PR 本身为纯导入路径调整,未改变任何行为。
文件 模块 状态 重要度
python/sglang/srt/elastic_ep/elastic_ep.py 弹性 EP modified 6.98
python/sglang/srt/distributed/parallel_state.py 并行状态 modified 5.07

关键符号

_wait_for_peer_state _try_recover_world try_recover_ranks _join_world_group join_process_groups join_scale_process_group

关键源码片段

python/sglang/srt/elastic_ep/elastic_ep.py core-logic

Elastic EP 恢复与重加入流程的核心模块,本次将 `get_peer_state`、`recover_ranks`、`join_group` 三个 ProcessGroup API 的导入从 `mooncake.ep` 全部迁移到 `mooncake.pg`,并简化了 `_wait_for_peer_state` 的函数签名,是本次改动的核心逻辑文件。

# python/sglang/srt/elastic_ep/elastic_ep.py
# 本次改动的核心:把 mooncake.ep 兼容路径统一迁移到 mooncake.pg。
# 注意 _wait_for_peer_state 不再接收模块对象作为参数,改为函数内直接导入。
​
​
def _wait_for_peer_state(backend, ranks: List[int]) -> None:
    # 从 mooncake.pg 直接导入 ProcessGroup 接口,替代原先通过
    # mooncake.ep 兼容路径间接访问的方式。
    from mooncake.pg import get_peer_state
​
    # 每 10ms 轮询一次对端状态,直到所有 rank 恢复在线。
    while not all(get_peer_state(backend, ranks)):
        time.sleep(_PEER_STATE_POLL_INTERVAL_SEC)
​
​
def _try_recover_world(global_ranks: List[int]) -> bool:
    from mooncake.pg import get_peer_state, recover_ranks
​
    world_backend = torch.distributed.group.WORLD
    # 任一 peer 不在线则直接返回 False,交由上层决定是否回退。
    if not all(get_peer_state(world_backend, global_ranks)):
        return False
​
    # 对端全部在线后执行 recover_ranks,把扩容 / 恢复的 rank 重新接回 WORLD 组。
    recover_ranks(world_backend, global_ranks)
    logger.debug("[Elastic EP][recover] WORLD recover_ranks(%s) done", global_ranks)
    return True
​
​
def try_recover_ranks(global_ranks: List[int]) -> bool:
    """Recover ranks in WORLD and every launch-time parallel group."""
    if not _try_recover_world(global_ranks):
        return False
​
    from mooncake.pg import recover_ranks
​
    for group in _iter_live_parallel_groups():
        local_ranks = _map_global_to_group_local_ranks(group.ranks, global_ranks)
        if not local_ranks:
            continue
​
        # 每个并行组都要先等 peer 上线、再 recover,
        # 依次处理 device_group 与 cpu_group,最后重建消息队列 broadcaster。
        _wait_for_peer_state(group.device_group, local_ranks)
        recover_ranks(group.device_group, local_ranks)
        _wait_for_peer_state(group.cpu_group, local_ranks)
        recover_ranks(group.cpu_group, local_ranks)
        _maybe_create_message_queue(group)
​
    _refresh_ep_members()
    return True
​
​
def _join_world_group() -> None:
    from mooncake.pg import join_group
​
    join_group(torch.distributed.group.WORLD)
​
​
def join_process_groups() -> None:
    """Rejoin WORLD and every launch-time parallel group after recovery."""
    from mooncake.pg import join_group
​
    _join_world_group()
    for group in _iter_live_parallel_groups():
        if group.world_size <= 1:
            continue
        join_group(group.device_group)
        join_group(group.cpu_group)
        _maybe_create_message_queue(group)
​
    _refresh_ep_members()
python/sglang/srt/distributed/parallel_state.py dependency-wiring

进程组创建的基础设施文件,`GroupCoordinator` 和 `init_distributed_environment` 中创建 Mooncake 后端进程组时使用的 `MooncakeBackendOptions` 导入路径从 `mooncake.ep` 迁移到 `mooncake.pg`。

# python/sglang/srt/distributed/parallel_state.py
# 创建 mooncake 后端的子并行组时,按上游新接口从 mooncake.pg 导入。for ranks in group_ranks:
    subgroup_timeout = _MODEL_PARALLEL_GROUP_TIMEOUT
    if "mooncake" in torch_distributed_backend:
        from mooncake.pg import MooncakeBackendOptions
​
        pg_active_size = len(ranks)
        if not recovered_rank and max_world_size is not None:
            # 扩容场景:组只激活部分 rank,但后端选项按 max_world_size 预留。
            assert max_world_size >= len(ranks), (
                f"max_world_size ({max_world_size}) must be >= "
                f"group size ({len(ranks)})"
            )
            pg_active_size = max_world_size
​
        pg_active_ranks = torch.zeros(
            pg_active_size, dtype=torch.int32, device=self.device
        )
        pg_active_ranks[: len(ranks)] = 1
        pg_active_ranks_cpu = torch.zeros(pg_active_size, dtype=torch.int32)
        pg_active_ranks_cpu[: len(ranks)] = 1
​
        # recovered_rank 与 max_world_size 决定是否传入扩容大小参数。
        if not recovered_rank and max_world_size is not None:
            dev_opts = MooncakeBackendOptions(
                pg_active_ranks, recovered_rank, max_world_size
            )
            cpu_opts = MooncakeBackendOptions(
                pg_active_ranks_cpu, recovered_rank, max_world_size
            )
        else:
            dev_opts = MooncakeBackendOptions(pg_active_ranks, recovered_rank)
            cpu_opts = MooncakeBackendOptions(
                pg_active_ranks_cpu, recovered_rank
            )

评论区精华

Gemini Code Assist 停止服务通知 other

gemini-code-assist[bot] 两次发布 CAUTION 提示:Google 消费者版 Gemini Code Assist 已 sunset,所有在 GitHub 上的代码评审活动已正式停止。

结论:对 PR 内容无实际影响,仅说明没有 AI 辅助评审输出。 · 已解决

CI 重跑请求 other

作者 UNIDY2002 先后 4 次发送 `/tag-and-rerun-ci` 命令触发 CI 重跑,并 @ch-wan 请求 review。

结论:CI 状态显示 PR Test (Base) 通过、PR Test (Extra) 失败,作者持续重跑;ch-wan 最终 APPROVED。 · 已解决

风险与影响

  1. 依赖版本兼容风险:改动依赖 Mooncake 新版本提供 mooncake.pg 子模块。若运行环境的 Mooncake 版本较旧、尚未提供 mooncake.pgimport 会直接抛 ModuleNotFoundError,导致 backend=mooncake 启动失败或 Elastic EP 恢复流程崩溃。需要确认部署侧 Mooncake 的最低版本要求。
  2. 缺失测试覆盖:PR 未附带任何单元测试或集成测试,mooncake.pg 导入路径的回归检测完全依赖 CI 中的 mooncake 相关用例是否存在。若 CI 未覆盖 Elastic EP 恢复/重加入路径,该改动可能静默失效。
  3. 兼容层删除风险mooncake.ep 是上游保留的兼容路径,长期看上游可能在某个版本移除该兼容层;本 PR 消除了 SGLang 对这层兼容的依赖,方向正确,但目前仍与旧版本不兼容。

影响范围集中在使用 Mooncake 作为分布式后端的场景:elastic_ep.py 中的 Elastic EP 动态扩容、故障恢复与进程组重加入流程,以及 parallel_state.py 中 mooncake 进程组(含 WORLD 组和子模型并行组)的创建。由于是纯导入路径替换,对使用 NCCL 等其他后端的用户零影响。对团队而言,这是对 Mooncake 上游接口规范的跟随,未来升级 Mooncake 时更安全;但对 mooncake 依赖版本提出了新的最低要求。

依赖版本兼容风险 缺少测试覆盖

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论