From f3dce0828321c45f32235c34f468c455e74f7a1e Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 18 May 2026 18:39:24 +0800 Subject: [PATCH] Move invariant checks to SchedulerInvariantChecker and retire runtime_checker mixin (#25624) --- .../hardware_backend/mlx/scheduler_mixin.py | 4 +- python/sglang/srt/managers/scheduler.py | 16 +- .../scheduler_components/invariant_checker.py | 235 +++++++++++++++- .../scheduler_runtime_checker_mixin.py | 254 +----------------- 4 files changed, 244 insertions(+), 265 deletions(-) diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index 6e3ac9de7..7faa7e149 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py @@ -211,7 +211,7 @@ class SchedulerMlxOverlapMixin: self.cur_batch = pending_curr.batch_copy self.last_batch = pending_curr.batch_copy if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): - self.self_check_during_busy() + self.invariant_checker.self_check_during_busy() continue # 4. Chain is broken. Finalise pending_next (if any), then @@ -230,4 +230,4 @@ class SchedulerMlxOverlapMixin: self.last_batch = next_batch if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): - self.self_check_during_busy() + self.invariant_checker.self_check_during_busy() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 920790b7c..a8382d91f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -331,8 +331,7 @@ def create_scheduler_watchdog( def dump_info() -> str: if scheduler.is_initializing: return "" - _, messages = scheduler._check_all_pools( - scheduler.invariant_checker, + _, messages = scheduler.invariant_checker._check_all_pools( scheduler.pool_stats_observer.get_pool_stats(), ) return ( @@ -1523,7 +1522,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.invariant_checker) + self.invariant_checker.self_check_during_busy() @DynamicGradMode() def event_loop_overlap(self): @@ -1578,7 +1577,7 @@ class Scheduler( self.last_batch = batch if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): - self.self_check_during_busy(self.invariant_checker) + self.invariant_checker.self_check_during_busy() 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. @@ -3105,16 +3104,15 @@ class Scheduler( # memory leak check (skipped for hisparse — pool counters intentionally # diverge during host-backup, see _get_swa_token_info clamp). if not self.enable_hisparse: - has_leak, messages = self._check_all_pools( - self.invariant_checker, + has_leak, messages = self.invariant_checker._check_all_pools( self.pool_stats_observer.get_pool_stats(), ) if has_leak: - self._report_leak(self.invariant_checker, "pool", "\n".join(messages)) - self._check_req_pool(self.invariant_checker) + self.invariant_checker._report_leak("pool", "\n".join(messages)) + self.invariant_checker._check_req_pool() # tree cache sanity check - self._check_tree_cache(self.invariant_checker) + self.invariant_checker._check_tree_cache() # metrics every 30s self._maybe_log_idle_metrics() diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 0f64d17f5..e7ed22aa1 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -1,17 +1,29 @@ from __future__ import annotations import logging +import warnings from dataclasses import dataclass -from typing import Callable, Optional +from typing import ( + Callable, + List, + Optional, + Tuple, +) from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.environ import envs from sglang.srt.managers.scheduler_components.pool_stats_observer import ( + PoolStats, 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 +from sglang.srt.utils.common import ( + ceil_align, + raise_error_or_warn, +) logger = logging.getLogger(__name__) @@ -34,3 +46,224 @@ class SchedulerInvariantChecker: get_running_batch: Callable count_req_pool_leak_warnings: int = 0 count_memory_leak_warnings: int = 0 + + @staticmethod + def _check_pool_invariant( + pool_name: str, + available: int, + evictable: int, + protected: int, + session_held: int, + total: int, + uncached: int = 0, + ) -> Tuple[bool, str]: + """Check: available + evictable + protected + session_held + uncached == total.""" + total_accounted = available + evictable + protected + session_held + uncached + leak = total_accounted != total + msg = ( + f"[{pool_name}] {total=}, {available=}, {evictable=}, " + f"{protected=}, {session_held=}, {uncached=}" + ) + return leak, msg + + def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: + if self.is_hybrid_swa: + protected = self.tree_cache.full_protected_size() + session_held = self.pool_stats_observer.session_held_full_tokens() + total = self.full_tokens_per_layer + elif self.is_hybrid_ssm and self.tree_cache.supports_mamba(): + protected = self.tree_cache.full_protected_size() + session_held = self.pool_stats_observer.session_held_tokens() + total = self.token_to_kv_pool_allocator.size + else: + 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( + "full", + ps.full_available_size, + ps.full_evictable_size, + protected, + session_held, + total, + uncached, + ) + + def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: + return self._check_pool_invariant( + "swa", + ps.swa_available_size, + ps.swa_evictable_size, + self.tree_cache.swa_protected_size(), + self.pool_stats_observer.session_held_swa_tokens(), + self.swa_tokens_per_layer, + uncached, + ) + + def _check_mamba_pool(self, ps: PoolStats) -> Tuple[bool, str]: + leak, msg = self._check_pool_invariant( + "mamba", + ps.mamba_available_size, + ps.mamba_evictable_size, + self.tree_cache.mamba_protected_size(), + self.pool_stats_observer.session_held_mamba_slots(), + self.req_to_token_pool.mamba_pool.size, + ) + if leak: + # Page-level leak diagnosis for mamba + free_full_pages = set( + self.token_to_kv_pool_allocator.free_pages.tolist() + + self.token_to_kv_pool_allocator.release_pages.tolist() + ) + cached_full_pages = set(self.tree_cache.all_values_flatten().tolist()) + expected_full_pages = set( + range(1, self.token_to_kv_pool_allocator.size + 1) + ) + leaked_full_pages = ( + expected_full_pages - free_full_pages - cached_full_pages + ) + free_mamba_pages = set( + self.req_to_token_pool.mamba_pool.free_slots.tolist() + ) + cached_mamba_pages = set( + self.tree_cache.all_mamba_values_flatten().tolist() + ) + expected_mamba_pages = set(range(self.req_to_token_pool.mamba_pool.size)) + leaked_mamba_pages = ( + expected_mamba_pages - free_mamba_pages - cached_mamba_pages + ) + msg += ( + f", leaked_full_pages={leaked_full_pages or None}" + f", leaked_mamba_pages={leaked_mamba_pages or None}" + ) + return leak, msg + + def _get_total_uncached_sizes( + self, + ) -> 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. + + For full pool: uncached = allocated - cache_protected_len + For SWA pool: uncached = allocated - max(cache_protected_len, swa_evicted_seqlen) + """ + # After decode: running_batch IS last_batch (same object), count once. + # After prefill: they differ, both hold uncached tokens. + batches = [self.get_last_batch()] + if ( + self.get_running_batch() not in (None, self.get_last_batch()) + and not self.get_running_batch().is_empty() + ): + batches.append(self.get_running_batch()) + + full_uncached = 0 + swa_uncached = 0 + for batch in batches: + for req in batch.reqs: + assert req.kv_committed_freed == req.kv_overallocated_freed + if req.kv_committed_freed or req.req_pool_idx is None: + continue + + allocated_len = req.kv_allocated_len + if self.page_size > 1: + allocated_len = ceil_align(allocated_len, self.page_size) + assert req.cache_protected_len % self.page_size == 0 + + full_uncached += allocated_len - req.cache_protected_len + if self.is_hybrid_swa: + swa_uncached += allocated_len - max( + req.cache_protected_len, req.swa_evicted_seqlen + ) + + return full_uncached, swa_uncached + + def self_check_during_busy(self): + if self.get_last_batch() is None: + return + + spec_topk = self.server_args.speculative_eagle_topk or 1 + if spec_topk > 1: + warnings.warn( + "Runtime memory check (busy) is not supported when speculation topk > 1." + ) + return + + ps = self.pool_stats_observer.get_pool_stats() + full_uncached, swa_uncached = self._get_total_uncached_sizes() + + full_leak, full_msg = self._check_full_pool(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) + + if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get() > 1: + logger.info(f"[Mem Check (BUSY)] {full_msg}") + if swa_msg: + logger.info(f"[Mem Check (BUSY)] {swa_msg}") + 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): + if self.disaggregation_mode == DisaggregationMode.DECODE: + req_total_size = ( + self.req_to_token_pool.size + self.req_to_token_pool.pre_alloc_size + ) + else: + req_total_size = self.req_to_token_pool.size + + session_req_count = self.pool_stats_observer.session_held_req_count() + if len(self.req_to_token_pool.free_slots) + session_req_count != req_total_size: + msg = ( + "req_to_token_pool memory leak detected!" + f"available_size={len(self.req_to_token_pool.free_slots)}, " + f"session_held={session_req_count}, " + f"total_size={self.req_to_token_pool.size}\n" + ) + raise_error_or_warn( + self, + envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(), + "count_req_pool_leak_warnings", + msg, + ) + + def _report_leak(self, pool_name: str, token_msg: str): + msg = f"{pool_name} memory leak detected! {token_msg}" + raise_error_or_warn( + self, + envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(), + "count_memory_leak_warnings", + msg, + ) + + def _check_all_pools( + self, 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) + has_leak |= full_leak + messages.append(full_msg) + + if self.is_hybrid_swa: + swa_leak, swa_msg = self._check_swa_pool(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) + has_leak |= mamba_leak + messages.append(mamba_msg) + + return has_leak, messages + + def _check_tree_cache(self): + 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() diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 5fd1a3516..10cbe4670 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -2,261 +2,18 @@ from __future__ import annotations import logging import time -import warnings -from typing import TYPE_CHECKING, List, Tuple +from typing import TYPE_CHECKING from sglang.srt.disaggregation.utils import DisaggregationMode -from sglang.srt.environ import envs from sglang.srt.observability.metrics_collector import QueueCount -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__) class SchedulerRuntimeCheckerMixin: - @staticmethod - def _check_pool_invariant( - pool_name: str, - available: int, - evictable: int, - protected: int, - session_held: int, - total: int, - uncached: int = 0, - ) -> Tuple[bool, str]: - """Check: available + evictable + protected + session_held + uncached == total.""" - total_accounted = available + evictable + protected + session_held + uncached - leak = total_accounted != total - msg = ( - f"[{pool_name}] {total=}, {available=}, {evictable=}, " - f"{protected=}, {session_held=}, {uncached=}" - ) - return leak, msg - - @staticmethod - def _check_full_pool( - self: "SchedulerInvariantChecker", ps: PoolStats, uncached: int = 0 - ) -> Tuple[bool, str]: - if self.is_hybrid_swa: - protected = self.tree_cache.full_protected_size() - session_held = self.pool_stats_observer.session_held_full_tokens() - total = self.full_tokens_per_layer - elif self.is_hybrid_ssm and self.tree_cache.supports_mamba(): - protected = self.tree_cache.full_protected_size() - session_held = self.pool_stats_observer.session_held_tokens() - total = self.token_to_kv_pool_allocator.size - else: - protected = self.tree_cache.protected_size() - session_held = self.pool_stats_observer.session_held_tokens() - total = self.max_total_num_tokens - return SchedulerRuntimeCheckerMixin._check_pool_invariant( - "full", - ps.full_available_size, - ps.full_evictable_size, - protected, - session_held, - total, - uncached, - ) - - @staticmethod - def _check_swa_pool( - self: "SchedulerInvariantChecker", ps: PoolStats, uncached: int = 0 - ) -> Tuple[bool, str]: - return SchedulerRuntimeCheckerMixin._check_pool_invariant( - "swa", - ps.swa_available_size, - ps.swa_evictable_size, - self.tree_cache.swa_protected_size(), - self.pool_stats_observer.session_held_swa_tokens(), - self.swa_tokens_per_layer, - uncached, - ) - - @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, - self.tree_cache.mamba_protected_size(), - self.pool_stats_observer.session_held_mamba_slots(), - self.req_to_token_pool.mamba_pool.size, - ) - if leak: - # Page-level leak diagnosis for mamba - free_full_pages = set( - self.token_to_kv_pool_allocator.free_pages.tolist() - + self.token_to_kv_pool_allocator.release_pages.tolist() - ) - cached_full_pages = set(self.tree_cache.all_values_flatten().tolist()) - expected_full_pages = set( - range(1, self.token_to_kv_pool_allocator.size + 1) - ) - leaked_full_pages = ( - expected_full_pages - free_full_pages - cached_full_pages - ) - free_mamba_pages = set( - self.req_to_token_pool.mamba_pool.free_slots.tolist() - ) - cached_mamba_pages = set( - self.tree_cache.all_mamba_values_flatten().tolist() - ) - expected_mamba_pages = set(range(self.req_to_token_pool.mamba_pool.size)) - leaked_mamba_pages = ( - expected_mamba_pages - free_mamba_pages - cached_mamba_pages - ) - msg += ( - f", leaked_full_pages={leaked_full_pages or None}" - f", leaked_mamba_pages={leaked_mamba_pages or None}" - ) - return leak, msg - - @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. - - For full pool: uncached = allocated - cache_protected_len - For SWA pool: uncached = allocated - max(cache_protected_len, swa_evicted_seqlen) - """ - # After decode: running_batch IS last_batch (same object), count once. - # After prefill: they differ, both hold uncached tokens. - batches = [self.get_last_batch()] - if ( - self.get_running_batch() not in (None, self.get_last_batch()) - and not self.get_running_batch().is_empty() - ): - batches.append(self.get_running_batch()) - - full_uncached = 0 - swa_uncached = 0 - for batch in batches: - for req in batch.reqs: - assert req.kv_committed_freed == req.kv_overallocated_freed - if req.kv_committed_freed or req.req_pool_idx is None: - continue - - allocated_len = req.kv_allocated_len - if self.page_size > 1: - allocated_len = ceil_align(allocated_len, self.page_size) - assert req.cache_protected_len % self.page_size == 0 - - full_uncached += allocated_len - req.cache_protected_len - if self.is_hybrid_swa: - swa_uncached += allocated_len - max( - req.cache_protected_len, req.swa_evicted_seqlen - ) - - return full_uncached, swa_uncached - - @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 - if spec_topk > 1: - warnings.warn( - "Runtime memory check (busy) is not supported when speculation topk > 1." - ) - return - - ps = self.pool_stats_observer.get_pool_stats() - full_uncached, swa_uncached = ( - SchedulerRuntimeCheckerMixin._get_total_uncached_sizes(self) - ) - - 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 = 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}") - if swa_msg: - logger.info(f"[Mem Check (BUSY)] {swa_msg}") - assert not full_leak, f"Full Pool Mem Leak Detected! {full_msg}" - assert not swa_leak, f"SWA Pool Mem Leak Detected! {swa_msg}" - - @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 - ) - else: - req_total_size = self.req_to_token_pool.size - - session_req_count = self.pool_stats_observer.session_held_req_count() - if len(self.req_to_token_pool.free_slots) + session_req_count != req_total_size: - msg = ( - "req_to_token_pool memory leak detected!" - f"available_size={len(self.req_to_token_pool.free_slots)}, " - f"session_held={session_req_count}, " - f"total_size={self.req_to_token_pool.size}\n" - ) - raise_error_or_warn( - self, - envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(), - "count_req_pool_leak_warnings", - msg, - ) - - @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, - envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(), - "count_memory_leak_warnings", - msg, - ) - - @staticmethod - def _check_all_pools( - 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 = 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 = 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 = SchedulerRuntimeCheckerMixin._check_mamba_pool( - self, ps - ) - has_leak |= mamba_leak - messages.append(mamba_msg) - - return has_leak, messages - def _maybe_log_idle_metrics(self: Scheduler): """Collect and log metrics every 30 seconds during idle.""" if ( @@ -297,12 +54,3 @@ class SchedulerRuntimeCheckerMixin: self.disagg_decode_transfer_queue.queue, priority_enabled ) self.metrics_collector.log_stats(self.stats) - - @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()) - or (self.is_hybrid_ssm and self.tree_cache.supports_mamba()) - ): - self.tree_cache.sanity_check()