Prhub

#36205 Gate the idle-loop tree-cache sanity check behind a default-off env

原始 PR 作者 sshleifer 合并时间 2026-08-30 00:10 文件变更 3 提交数 3 评论 9 代码增减 +99 / -0

执行摘要

空闲循环树缓存深度校验默认关闭,修复 DP 组调度卡死

PR body 描述了确定性复现的生产事故:Scheduler.on_idle 在每个空闲迭代无条件调用 _check_tree_cache() → tree_cache.sanity_check(),而该方法是纯 Python 全树遍历(自带 TODO(hzh) 注明高延迟),大树单次可达数十秒。在 8×GPU DeepSeek-V4 类部署(tp4/dp4、--enable-dp-attention、hybrid SWA)上,请求排空的 rank 进入空闲后不断重启全树 walk,其余 rank 阻塞在 broadcast_pyobj(recv_requests)与 MoE a2a 集合通信上,出现 GPUs 0%、调度器约 190% CPU 的永久楔死。py-spy 多次抓取确认 walk 每轮 idle 迭代重新开始,组内无法推进。作者结论是「The full-tree walk is debug instrumentation, so make it opt-in」,即把调试仪器从生产路径中分离出去。

值得精读。代码量虽小,却包含两个可复用设计:一是用 EnvField 的 callable default 实现「生产默认关、CI 默认开、显式覆盖优先」的三态配置语义,让调试仪器默认不进生产路径同时不牺牲 CI 回归;二是测试通过 contextmanager 清理进程环境变量来稳定默认值断言,是环境耦合测试的典型解法。对调度器与缓存团队,建议同时关注 sanity_check 自身性能(TODO(hzh))与文档同步。

讨论亮点

核心交锋是默认值策略:

  • ispobock 在 environ.py 上发问:「what do you think to turn off it by default? Maybe we need to keep it in CI at least?」
  • hzh0425 表示赞同并点名作者:「yes, we should keep in ci at least @sshleifer」,随后追加「@sshleifer could you help improve this?」
  • 最终以 commit “run sanity check in ci” 落地为 callable 默认值方案:非 CI 默认 False、CI 默认 True,显式设置覆盖。
  • Codex bot 提出 P2:test_disabled_by_default 未隔离进程环境(若运行环境已导出 SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK=1,EnvField.get() 会读到继承值导致 assert_not_called() 失败),最终通过 _without_explicit_sanity_check_setting() contextmanager 保存并清除该变量解决。
  • hzh0425 对测试文件第 11 行(DisaggregationMode 导入附近)留言「This test is not necessary」,但最终合并版本保留了全部 4 个测试场景与相关导入,该意见未落地。

实现拆解

  1. 新增环境开关(python/sglang/srt/environ.py):定义 _default_tree_cache_sanity_check() 延迟默认值函数,返回 envs.SGLANG_IS_IN_CI.get();声明 SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK = EnvBool(_default_tree_cache_sanity_check)。利用 EnvField 的 callable default 机制(_resolve_default 仅在该变量未被显式设置时才求值默认值),实现「生产(非 CI)默认关、CI 默认开、显式设置永远优先」的三态语义,无需 CI 配置文件显式注入变量。
  2. 在调用点加 gate(python/sglang/srt/managers/scheduler_components/invariant_checker.py):_check_tree_cache() 开头先判断 envs.SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK.get(),为 False 直接 return;原有 tree cache 类型 + hybrid SWA/SSM 支持判定与 sanity_check 调用保持不变。廉价的空闲池与 req-pool 泄漏检查(_check_full_pool / _check_swa_pool / _check_mamba_pool)不在 gate 内,仍每轮执行。
  3. 新增单元测试(test/registered/unit/managers/scheduler_components/test_invariant_checker.py):用 MagicMock 构造 SchedulerInvariantChecker,覆盖 4 个场景——非 CI 默认关闭、CI 默认开启、CI 中显式关闭、非 CI 显式开启;通过 _without_explicit_sanity_check_setting() contextmanager 在断言默认值前清掉进程环境里的显式设置,解决 Codex 指出的环境变量泄漏问题。
  4. 设计演进:第一版是 EnvBool(False) 纯默认关闭,review 中 ispobock 与 hzh0425 提出「至少 CI 保留」后,由 hzh0425 的 commit “run sanity check in ci” 改为与 SGLANG_IS_IN_CI 联动的延迟默认值,兼顾生产止血与 CI 回归。
文件 模块 状态 重要度
python/sglang/srt/environ.py 环境配置 modified 6.23
python/sglang/srt/managers/scheduler_components/invariant_checker.py 调度器 modified 5.1
test/registered/unit/managers/scheduler_components/test_invariant_checker.py 单元测试 added 7.04

关键符号

_default_tree_cache_sanity_check _check_tree_cache TestCheckTreeCacheGate _without_explicit_sanity_check_setting

关键源码片段

python/sglang/srt/managers/scheduler_components/invariant_checker.py core-logic

生产路径的实际变更落点:在 _check_tree_cache() 入口加短路 gate,是解除 DP 组 livelock 的关键一行。

def _check_tree_cache(self):
    # 全树一致性校验是纯 Python 调试仪器,每个空闲迭代都会命中;
    # 大树(数千条序列)单次遍历可达数十秒,DP 场景会拖死整个组,
    # 因此默认仅 CI / 显式调试开启,生产路径直接短路返回。
    if not envs.SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK.get():
        return
​
    if (
        self.tree_cache.is_tree_cache()
        and (self.is_hybrid_swa and self.tree_cache.supports_swa())
        or (self.is_hybrid_ssm and self.tree_cache.supports_mamba())
    ):
        self.tree_cache.sanity_check()
test/registered/unit/managers/scheduler_components/test_invariant_checker.py test-coverage

新增 4 个场景的 gate 测试,并示范如何通过 contextmanager 隔离进程环境变量以稳定默认值断言。

class TestCheckTreeCacheGate(CustomTestCase):
    @contextmanager
    def _without_explicit_sanity_check_setting(self):
        # 从进程环境剔除显式设置,避免外层 CI / 调试运行环境已导出的
        # SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK 干扰默认值断言;
        # clear=False 仅清除该字段,其余环境变量保持不变。
        with patch.dict(os.environ, {}, clear=False):
            envs.SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK.clear()
            yield
​
    def test_disabled_by_default(self):
        # 非 CI 且未显式设置:默认应关闭,sanity_check 不得被调用。
        with (
            envs.SGLANG_IS_IN_CI.override(False),
            self._without_explicit_sanity_check_setting(),
        ):
            checker = self._make_checker()
            checker._check_tree_cache()
            checker.tree_cache.sanity_check.assert_not_called()
​
    def test_enabled_by_default_in_ci(self):
        # CI 且未显式覆盖:默认应开启,深度一致性回归仍需生效。
        with (
            envs.SGLANG_IS_IN_CI.override(True),
            self._without_explicit_sanity_check_setting(),
        ):
            checker = self._make_checker()
            checker._check_tree_cache()
            checker.tree_cache.sanity_check.assert_called_once()

评论区精华

默认关闭的取舍与 CI 保留 设计

ispobock: what do you think to turn off it by default? Maybe we need to keep it in CI at least?;hzh0425: yes, we should keep in ci at least @sshleifer。

结论:改为 callable 默认值 _default_tree_cache_sanity_check(),非 CI 默认 False、CI 默认 True,显式设置优先。 · 已解决

测试需隔离进程环境变量(Codex P2) 测试

运行环境若导出 SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK=1,EnvField.get() 读到继承值会使 test_disabled_by_default 的 assert_not_called() 失败;建议保存并清除该变量。

结论:引入 _without_explicit_sanity_check_setting() contextmanager,用 patch.dict + clear() 清理后恢复。 · 已解决

hzh0425 认为某处测试不必要 question

hzh0425 对测试文件第 11 行(DisaggregationMode 导入附近)留言 This test is not necessary,意图无法从材料完全确定。

结论:最终合并版本保留了全部 4 个测试场景与相关导入,该意见未落地。 · unresolved

风险与影响

  1. 生产深度校验丢失:默认关闭后,大规模部署不再执行 tree_cache.sanity_check(),树结构回归(节点计数、父子链、SWA 视图)依赖 CI 兜底,而 CI 覆盖的树规模远小于生产,隐藏问题可能上线后才暴露。
  2. 求值时机耦合:callable 默认值在首次读取时求值 SGLANG_IS_IN_CI,若运行环境在启动后动态改变该变量,开关行为可能不一致;SGLANG 的 EnvBool 通常早期固定,风险有限。
  3. 显式优先级语义:代码注释明确了「显式变量优先于 CI 默认」,但 PR 未勾选 documentation checklist,用户可能不知道这个开关的存在与语义。
  4. 根因未除:sanity_check() 本身的 O(节点数) 纯 Python 遍历仍在,只是默认不触发;后续新增调试校验需求时可能重蹈覆辙,建议后续将 sanity_check 改为稀疏/增量检查或编译优化。
  5. 测试环境耦合:若外部 CI 全局导出 SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK=1,test_disabled_by_default 已通过 helper 隔离;但并行执行环境下仍需注意环境变量串扰。

影响范围集中在调度器空闲路径与调试配置:

  • 用户/部署:DP attention + EP(tp4/dp4、hybrid SWA)等会周期性整体空闲的部署直接受益,饱和 benchmark 后半段不再楔死;非 DP 部署的空闲开销也显著下降,且无需任何配置即生效。
  • CI:SGLANG_IS_IN_CI 为 True 的套件默认保持深度检查,一致性回归能力不丢失;显式设 False 可对特殊用例豁免。
  • 团队:为调试仪器与生产路径分离提供了新模式(延迟默认值 + 环境 gate),后续同类检查(如 SGLANG_ENABLE_ASYNC_ASSERT、SGLANG_CHECK_KV_PAGE_INVARIANTS)可沿用同一套路;但文档未同步,需要后续补充。
生产默认关闭深度校验 CI 默认开启 测试需隔离环境变量 根因未除(sanity_check 仍为全量遍历) 文档未同步

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论