Introduce SchedulerInvariantChecker to own invariant-check state (#25623)

python/sglang/srt/managers/scheduler.py

@@ -167,6 +167,9 @@ from sglang.srt.managers.schedule_policy import (
 from sglang.srt.managers.scheduler_components.dp_attn import (
     SchedulerDPAttnAdapter,
 )
+from sglang.srt.managers.scheduler_components.invariant_checker import (
+    SchedulerInvariantChecker,
+)
 from sglang.srt.managers.scheduler_components.pool_stats_observer import (
     SchedulerPoolStatsObserver,
 )
@@ -329,7 +332,8 @@ def create_scheduler_watchdog(
         if scheduler.is_initializing:
             return ""
         _, messages = scheduler._check_all_pools(
-            scheduler.pool_stats_observer.get_pool_stats()
+            scheduler.invariant_checker,
+            scheduler.pool_stats_observer.get_pool_stats(),
         )
         return (
             f"{scheduler.cur_batch.batch_size()=}\n"
@@ -654,6 +658,23 @@ class Scheduler(
             get_running_batch=lambda: self.running_batch,
         )
 
+        self.invariant_checker = SchedulerInvariantChecker(
+            is_hybrid_swa=self.is_hybrid_swa,
+            is_hybrid_ssm=self.is_hybrid_ssm,
+            disaggregation_mode=self.disaggregation_mode,
+            page_size=self.page_size,
+            full_tokens_per_layer=self.full_tokens_per_layer,
+            swa_tokens_per_layer=self.swa_tokens_per_layer,
+            max_total_num_tokens=self.max_total_num_tokens,
+            server_args=self.server_args,
+            tree_cache=self.tree_cache,
+            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            req_to_token_pool=self.req_to_token_pool,
+            pool_stats_observer=self.pool_stats_observer,
+            get_last_batch=lambda: self.last_batch,
+            get_running_batch=lambda: self.running_batch,
+        )
+
         self.is_initializing = False
 
     def init_zbal_on_npu(self):
@@ -1502,7 +1523,7 @@ class Scheduler(
             # Update last_batch
             self.last_batch = batch
             if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
-                self.self_check_during_busy()
+                self.self_check_during_busy(self.invariant_checker)
 
     @DynamicGradMode()
     def event_loop_overlap(self):
@@ -1557,7 +1578,7 @@ class Scheduler(
             self.last_batch = batch
 
             if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
-                self.self_check_during_busy()
+                self.self_check_during_busy(self.invariant_checker)
 
     def is_disable_overlap_for_batch(self, batch: ScheduleBatch) -> bool:
         # For two consecutive prefill batches, we disable overlap to improve the TTFT of the first batch.
@@ -3085,14 +3106,15 @@ class Scheduler(
         # diverge during host-backup, see _get_swa_token_info clamp).
         if not self.enable_hisparse:
             has_leak, messages = self._check_all_pools(
-                self.pool_stats_observer.get_pool_stats()
+                self.invariant_checker,
+                self.pool_stats_observer.get_pool_stats(),
             )
             if has_leak:
-                self._report_leak("pool", "\n".join(messages))
-            self._check_req_pool()
+                self._report_leak(self.invariant_checker, "pool", "\n".join(messages))
+            self._check_req_pool(self.invariant_checker)
 
         # tree cache sanity check
-        self._check_tree_cache()
+        self._check_tree_cache(self.invariant_checker)
 
         # metrics every 30s
         self._maybe_log_idle_metrics()

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.scheduler_components.pool_stats_observer import (
+    SchedulerPoolStatsObserver,
+)
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(kw_only=True, slots=True)
+class SchedulerInvariantChecker:
+    is_hybrid_swa: bool
+    is_hybrid_ssm: bool
+    disaggregation_mode: DisaggregationMode
+    page_size: int
+    full_tokens_per_layer: Optional[int]
+    swa_tokens_per_layer: Optional[int]
+    max_total_num_tokens: int
+    server_args: ServerArgs
+    tree_cache: BasePrefixCache
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    req_to_token_pool: ReqToTokenPool
+    pool_stats_observer: SchedulerPoolStatsObserver
+    get_last_batch: Callable
+    get_running_batch: Callable
+    count_req_pool_leak_warnings: int = 0
+    count_memory_leak_warnings: int = 0

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -12,6 +12,9 @@ from sglang.srt.utils.common import ceil_align, raise_error_or_warn
 
 if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import Scheduler
+    from sglang.srt.managers.scheduler_components.invariant_checker import (
+        SchedulerInvariantChecker,
+    )
 
 logger = logging.getLogger(__name__)
 
@@ -36,8 +39,9 @@ class SchedulerRuntimeCheckerMixin:
         )
         return leak, msg
 
+    @staticmethod
     def _check_full_pool(
-        self: Scheduler, ps: PoolStats, uncached: int = 0
+        self: "SchedulerInvariantChecker", ps: PoolStats, uncached: int = 0
     ) -> Tuple[bool, str]:
         if self.is_hybrid_swa:
             protected = self.tree_cache.full_protected_size()
@@ -51,7 +55,7 @@ class SchedulerRuntimeCheckerMixin:
             protected = self.tree_cache.protected_size()
             session_held = self.pool_stats_observer.session_held_tokens()
             total = self.max_total_num_tokens
-        return self._check_pool_invariant(
+        return SchedulerRuntimeCheckerMixin._check_pool_invariant(
             "full",
             ps.full_available_size,
             ps.full_evictable_size,
@@ -61,10 +65,11 @@ class SchedulerRuntimeCheckerMixin:
             uncached,
         )
 
+    @staticmethod
     def _check_swa_pool(
-        self: Scheduler, ps: PoolStats, uncached: int = 0
+        self: "SchedulerInvariantChecker", ps: PoolStats, uncached: int = 0
     ) -> Tuple[bool, str]:
-        return self._check_pool_invariant(
+        return SchedulerRuntimeCheckerMixin._check_pool_invariant(
             "swa",
             ps.swa_available_size,
             ps.swa_evictable_size,
@@ -74,8 +79,11 @@ class SchedulerRuntimeCheckerMixin:
             uncached,
         )
 
-    def _check_mamba_pool(self: Scheduler, ps: PoolStats) -> Tuple[bool, str]:
-        leak, msg = self._check_pool_invariant(
+    @staticmethod
+    def _check_mamba_pool(
+        self: "SchedulerInvariantChecker", ps: PoolStats
+    ) -> Tuple[bool, str]:
+        leak, msg = SchedulerRuntimeCheckerMixin._check_pool_invariant(
             "mamba",
             ps.mamba_available_size,
             ps.mamba_evictable_size,
@@ -112,7 +120,10 @@ class SchedulerRuntimeCheckerMixin:
             )
         return leak, msg
 
-    def _get_total_uncached_sizes(self: Scheduler) -> Tuple[int, int]:
+    @staticmethod
+    def _get_total_uncached_sizes(
+        self: "SchedulerInvariantChecker",
+    ) -> Tuple[int, int]:
         """Sum uncached tokens for full and SWA pools across all active batches.
 
         Returns (full_uncached, swa_uncached). For non-SWA models, swa_uncached is 0.
@@ -122,12 +133,12 @@ class SchedulerRuntimeCheckerMixin:
         """
         # After decode: running_batch IS last_batch (same object), count once.
         # After prefill: they differ, both hold uncached tokens.
-        batches = [self.last_batch]
+        batches = [self.get_last_batch()]
         if (
-            self.running_batch not in (None, self.last_batch)
-            and not self.running_batch.is_empty()
+            self.get_running_batch() not in (None, self.get_last_batch())
+            and not self.get_running_batch().is_empty()
         ):
-            batches.append(self.running_batch)
+            batches.append(self.get_running_batch())
 
         full_uncached = 0
         swa_uncached = 0
@@ -150,8 +161,9 @@ class SchedulerRuntimeCheckerMixin:
 
         return full_uncached, swa_uncached
 
-    def self_check_during_busy(self: Scheduler):
-        if self.last_batch is None:
+    @staticmethod
+    def self_check_during_busy(self: "SchedulerInvariantChecker"):
+        if self.get_last_batch() is None:
             return
 
         spec_topk = self.server_args.speculative_eagle_topk or 1
@@ -162,13 +174,19 @@ class SchedulerRuntimeCheckerMixin:
             return
 
         ps = self.pool_stats_observer.get_pool_stats()
-        full_uncached, swa_uncached = self._get_total_uncached_sizes()
+        full_uncached, swa_uncached = (
+            SchedulerRuntimeCheckerMixin._get_total_uncached_sizes(self)
+        )
 
-        full_leak, full_msg = self._check_full_pool(ps, uncached=full_uncached)
+        full_leak, full_msg = SchedulerRuntimeCheckerMixin._check_full_pool(
+            self, ps, uncached=full_uncached
+        )
 
         swa_leak, swa_msg = False, ""
         if self.is_hybrid_swa:
-            swa_leak, swa_msg = self._check_swa_pool(ps, uncached=swa_uncached)
+            swa_leak, swa_msg = SchedulerRuntimeCheckerMixin._check_swa_pool(
+                self, ps, uncached=swa_uncached
+            )
 
         if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get() > 1:
             logger.info(f"[Mem Check (BUSY)] {full_msg}")
@@ -177,7 +195,8 @@ class SchedulerRuntimeCheckerMixin:
         assert not full_leak, f"Full Pool Mem Leak Detected! {full_msg}"
         assert not swa_leak, f"SWA Pool Mem Leak Detected! {swa_msg}"
 
-    def _check_req_pool(self: Scheduler):
+    @staticmethod
+    def _check_req_pool(self: "SchedulerInvariantChecker"):
         if self.disaggregation_mode == DisaggregationMode.DECODE:
             req_total_size = (
                 self.req_to_token_pool.size + self.req_to_token_pool.pre_alloc_size
@@ -200,7 +219,8 @@ class SchedulerRuntimeCheckerMixin:
                 msg,
             )
 
-    def _report_leak(self: Scheduler, pool_name: str, token_msg: str):
+    @staticmethod
+    def _report_leak(self: "SchedulerInvariantChecker", pool_name: str, token_msg: str):
         msg = f"{pool_name} memory leak detected! {token_msg}"
         raise_error_or_warn(
             self,
@@ -209,24 +229,29 @@ class SchedulerRuntimeCheckerMixin:
             msg,
         )
 
+    @staticmethod
     def _check_all_pools(
-        self: Scheduler, ps: PoolStats, uncached: int = 0
+        self: "SchedulerInvariantChecker", ps: PoolStats, uncached: int = 0
     ) -> Tuple[bool, List[str]]:
         """Check memory invariant across all pools. Returns (has_leak, messages)."""
         has_leak = False
         messages = []
 
-        full_leak, full_msg = self._check_full_pool(ps, uncached=uncached)
+        full_leak, full_msg = SchedulerRuntimeCheckerMixin._check_full_pool(
+            self, ps, uncached=uncached
+        )
         has_leak |= full_leak
         messages.append(full_msg)
 
         if self.is_hybrid_swa:
-            swa_leak, swa_msg = self._check_swa_pool(ps)
+            swa_leak, swa_msg = SchedulerRuntimeCheckerMixin._check_swa_pool(self, ps)
             has_leak |= swa_leak
             messages.append(swa_msg)
 
         if self.is_hybrid_ssm and self.tree_cache.supports_mamba():
-            mamba_leak, mamba_msg = self._check_mamba_pool(ps)
+            mamba_leak, mamba_msg = SchedulerRuntimeCheckerMixin._check_mamba_pool(
+                self, ps
+            )
             has_leak |= mamba_leak
             messages.append(mamba_msg)
 
@@ -273,7 +298,8 @@ class SchedulerRuntimeCheckerMixin:
             )
         self.metrics_collector.log_stats(self.stats)
 
-    def _check_tree_cache(self: Scheduler):
+    @staticmethod
+    def _check_tree_cache(self: "SchedulerInvariantChecker"):
         if (
             self.tree_cache.is_tree_cache()
             and (self.is_hybrid_swa and self.tree_cache.supports_swa())