执行摘要
- 一句话:空闲循环树缓存深度校验默认关闭,修复 DP 组调度卡死
- 推荐动作:值得精读。代码量虽小,却包含两个可复用设计:一是用 EnvField 的 callable default 实现「生产默认关、CI 默认开、显式覆盖优先」的三态配置语义,让调试仪器默认不进生产路径同时不牺牲 CI 回归;二是测试通过 contextmanager 清理进程环境变量来稳定默认值断言,是环境耦合测试的典型解法。对调度器与缓存团队,建议同时关注 sanity_check 自身性能(TODO(hzh))与文档同步。
功能与动机
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」,即把调试仪器从生产路径中分离出去。
实现拆解
- 新增环境开关(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 配置文件显式注入变量。
- 在调用点加 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 内,仍每轮执行。
- 新增单元测试(test/registered/unit/managers/scheduler_components/test_invariant_checker.py):用 MagicMock 构造 SchedulerInvariantChecker,覆盖 4 个场景——非 CI 默认关闭、CI 默认开启、CI 中显式关闭、非 CI 显式开启;通过 _without_explicit_sanity_check_setting() contextmanager 在断言默认值前清掉进程环境里的显式设置,解决 Codex 指出的环境变量泄漏问题。
- 设计演进:第一版是 EnvBool(False) 纯默认关闭,review 中 ispobock 与 hzh0425 提出「至少 CI 保留」后,由 hzh0425 的 commit “run sanity check in ci” 改为与 SGLANG_IS_IN_CI 联动的延迟默认值,兼顾生产止血与 CI 回归。
关键文件:
python/sglang/srt/environ.py(模块 环境配置;类别 source;类型 core-logic;符号 _default_tree_cache_sanity_check, SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK): 定义开关语义的核心文件:新增 SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK 与延迟默认值函数 _default_tree_cache_sanity_check,实现生产关、CI 开、显式优先的三态配置。
python/sglang/srt/managers/scheduler_components/invariant_checker.py(模块 调度器;类别 source;类型 core-logic;符号 _check_tree_cache): 生产路径的实际变更落点:在 _check_tree_cache() 入口加短路 gate,是解除 DP 组 livelock 的关键一行。
test/registered/unit/managers/scheduler_components/test_invariant_checker.py(模块 单元测试;类别 test;类型 test-coverage;符号 TestCheckTreeCacheGate, _without_explicit_sanity_check_setting, _make_checker, test_disabled_by_default): 新增 4 个场景的 gate 测试,并示范如何通过 contextmanager 隔离进程环境变量以稳定默认值断言。
关键符号:_default_tree_cache_sanity_check, _check_tree_cache, TestCheckTreeCacheGate, _without_explicit_sanity_check_setting
关键源码片段
python/sglang/srt/managers/scheduler_components/invariant_checker.py
生产路径的实际变更落点:在 _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
新增 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()
评论区精华
核心交锋是默认值策略:
- 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 个测试场景与相关导入,该意见未落地。
-
默认关闭的取舍与 CI 保留 (design): 改为 callable 默认值 _default_tree_cache_sanity_check(),非 CI 默认 False、CI 默认 True,显式设置优先。
- 测试需隔离进程环境变量(Codex P2) (testing): 引入 _without_explicit_sanity_check_setting() contextmanager,用 patch.dict + clear() 清理后恢复。
- hzh0425 认为某处测试不必要 (question): 最终合并版本保留了全部 4 个测试场景与相关导入,该意见未落地。
风险与影响
- 风险:
- 生产深度校验丢失:默认关闭后,大规模部署不再执行 tree_cache.sanity_check(),树结构回归(节点计数、父子链、SWA 视图)依赖 CI 兜底,而 CI 覆盖的树规模远小于生产,隐藏问题可能上线后才暴露。
- 求值时机耦合:callable 默认值在首次读取时求值 SGLANG_IS_IN_CI,若运行环境在启动后动态改变该变量,开关行为可能不一致;SGLANG 的 EnvBool 通常早期固定,风险有限。
- 显式优先级语义:代码注释明确了「显式变量优先于 CI 默认」,但 PR 未勾选 documentation checklist,用户可能不知道这个开关的存在与语义。
- 根因未除:sanity_check() 本身的 O(节点数) 纯 Python 遍历仍在,只是默认不触发;后续新增调试校验需求时可能重蹈覆辙,建议后续将 sanity_check 改为稀疏/增量检查或编译优化。
- 测试环境耦合:若外部 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 仍为全量遍历), 文档未同步
关联脉络
- PR #37151 [Unified Cache Linker][3/N]: Add backend-independent linker core: 同属 unified cache / 基数树演进线,被 gate 的 sanity_check 正是 unified_radix_cache / unified_cache_linker 的调试仪器。
- PR #37166 fix(staging): make empty staging rings reusable: 同为 DP/PD 场景调度卡死类修复,与本 PR 同属调度器稳定性止损脉络。
- PR #37094 [mem_cache] Move
req_pool_idx into ReqKvInfo: mem_cache / 调度器内存路径重构系列的一部分,本 PR 的 gate 为这类调试校验提供可开关模式。
参与讨论