From 6afb5e17712e2e90b60ba8456ca893e529316869 Mon Sep 17 00:00:00 2001 From: Sam Shleifer Date: Sat, 29 Aug 2026 12:10:28 -0400 Subject: [PATCH] Gate the idle-loop tree-cache sanity check behind a default-off env (#36205) Co-authored-by: Claude Fable 5 Co-authored-by: hzh0425 --- python/sglang/srt/environ.py | 8 ++ .../scheduler_components/invariant_checker.py | 2 + .../test_invariant_checker.py | 89 +++++++++++++++++++ 3 files changed, 99 insertions(+) create mode 100644 test/registered/unit/managers/scheduler_components/test_invariant_checker.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 196a203d1..89426142f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -37,6 +37,11 @@ def _default_cache_subdir(name: str) -> str: return os.path.join(os.path.expanduser(envs.SGLANG_CACHE_DIR.get()), name) +def _default_tree_cache_sanity_check() -> bool: + """Enable the expensive tree-cache sanity check by default in CI.""" + return envs.SGLANG_IS_IN_CI.get() + + class EnvField: _allow_set_name = True @@ -464,6 +469,9 @@ class Envs: SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(True) SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY = EnvInt(0) SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE = EnvBool(True) + # The explicit environment variable still takes precedence over this CI + # default, so production remains opt-in and CI remains opt-out if needed. + SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK = EnvBool(_default_tree_cache_sanity_check) # Physical KV-page checks: committed<=allocated + no page alias. SGLANG_CHECK_KV_PAGE_INVARIANTS = EnvBool(False) SGLANG_TBO_DEBUG = EnvBool(False) diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 2cb0b9b5e..fb0af2db2 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -462,6 +462,8 @@ class SchedulerInvariantChecker: return has_leak, messages def _check_tree_cache(self): + 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()) diff --git a/test/registered/unit/managers/scheduler_components/test_invariant_checker.py b/test/registered/unit/managers/scheduler_components/test_invariant_checker.py new file mode 100644 index 000000000..d4ea9dff3 --- /dev/null +++ b/test/registered/unit/managers/scheduler_components/test_invariant_checker.py @@ -0,0 +1,89 @@ +import os +import unittest +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.environ import envs +from sglang.srt.managers.scheduler_components.invariant_checker import ( + SchedulerInvariantChecker, +) + +register_cpu_ci(est_time=4, suite="base-a-test-cpu") + + +class TestCheckTreeCacheGate(CustomTestCase): + @contextmanager + def _without_explicit_sanity_check_setting(self): + with patch.dict(os.environ, {}, clear=False): + envs.SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK.clear() + yield + + def _make_checker(self): + tree_cache = MagicMock() + tree_cache.is_tree_cache.return_value = True + tree_cache.supports_swa.return_value = True + return SchedulerInvariantChecker( + is_hybrid_swa=True, + is_hybrid_ssm=False, + disaggregation_mode=DisaggregationMode.NULL, + page_size=1, + full_tokens_per_layer=None, + swa_tokens_per_layer=None, + max_total_num_tokens=1024, + tree_cache=tree_cache, + token_to_kv_pool_allocator=MagicMock(), + req_to_token_pool=MagicMock(), + pool_stats_observer=MagicMock(), + get_last_batch=lambda: None, + get_running_batch=lambda: None, + ) + + def test_disabled_by_default(self): + 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): + 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() + + def test_explicitly_disabled_in_ci(self): + with envs.SGLANG_IS_IN_CI.override(True): + checker = self._make_checker() + + with envs.SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK.override(False): + checker._check_tree_cache() + + checker.tree_cache.sanity_check.assert_not_called() + + def test_runs_when_enabled(self): + with envs.SGLANG_IS_IN_CI.override(False): + checker = self._make_checker() + + with envs.SGLANG_ENABLE_TREE_CACHE_SANITY_CHECK.override(True): + checker._check_tree_cache() + + checker.tree_cache.sanity_check.assert_called_once() + + +if __name__ == "__main__": + unittest.main()