Add SchedulerPoolStatsObserver and route pool-stats state through it (#25619)

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.pool_stats_observer import (
+    SchedulerPoolStatsObserver,
+)
 from sglang.srt.managers.scheduler_components.profiler_manager import (
     SchedulerProfilerManager,
 )
@@ -611,6 +614,22 @@ class Scheduler(
             get_require_mlp_sync=lambda: self.require_mlp_sync,
         )
 
+        self.pool_stats_observer = SchedulerPoolStatsObserver(
+            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,
+            session_controller=self.session_controller,
+            hisparse_coordinator=self.hisparse_coordinator,
+            is_hybrid_swa=self.is_hybrid_swa,
+            is_hybrid_ssm=self.is_hybrid_ssm,
+            enable_hisparse=self.enable_hisparse,
+            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,
+            get_last_batch=lambda: self.last_batch,
+            get_running_batch=lambda: self.running_batch,
+        )
+
         self.is_initializing = False
 
     def init_zbal_on_npu(self):
@@ -2324,7 +2343,9 @@ class Scheduler(
         prefill_delayer_single_pass = None
         if self.prefill_delayer:
             # Get max usage across all pools for prefill delay decision
-            max_pool_usage = self.get_pool_stats().get_max_pool_usage()
+            max_pool_usage = self.get_pool_stats(
+                self.pool_stats_observer,
+            ).get_max_pool_usage()
             prefill_delayer_single_pass = PrefillDelayerSinglePassExecutor(
                 self.prefill_delayer, token_usage=max_pool_usage
             )
@@ -3039,7 +3060,11 @@ 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.get_pool_stats())
+            has_leak, messages = self._check_all_pools(
+                self.get_pool_stats(
+                    self.pool_stats_observer,
+                )
+            )
             if has_leak:
                 self._report_leak("pool", "\n".join(messages))
             self._check_req_pool()

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

@@ -1,7 +1,20 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import List, Optional, Tuple
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    List,
+    Optional,
+    Tuple,
+)
+
+if TYPE_CHECKING:
+    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
 
 
 class SchedulerStats: ...  # type: ignore[no-redef]
@@ -123,3 +136,20 @@ class PoolStats:
         stats.kv_available_tokens = self.full_available_size
         stats.kv_evictable_tokens = self.full_evictable_size
         stats.kv_used_tokens = self.full_num_used
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerPoolStatsObserver:
+    tree_cache: "BasePrefixCache"
+    token_to_kv_pool_allocator: "BaseTokenToKVPoolAllocator"
+    req_to_token_pool: "ReqToTokenPool"
+    session_controller: Any
+    hisparse_coordinator: Any
+    is_hybrid_swa: bool
+    is_hybrid_ssm: bool
+    enable_hisparse: bool
+    full_tokens_per_layer: Any
+    swa_tokens_per_layer: Any
+    max_total_num_tokens: int
+    get_last_batch: Callable
+    get_running_batch: Callable

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -15,26 +15,31 @@ from sglang.srt.utils.watchdog import WatchdogRaw
 
 if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import Scheduler
+    from sglang.srt.managers.scheduler_components.pool_stats_observer import (
+        SchedulerPoolStatsObserver,
+    )
 
 logger = logging.getLogger(__name__)
 
 
 class SchedulerRuntimeCheckerMixin:
-    def _streaming_session_count(self: Scheduler) -> int:
+    @staticmethod
+    def streaming_session_count(self: "SchedulerPoolStatsObserver") -> int:
         return sum(
             1
             for session in self.session_controller.sessions.values()
             if session.streaming
         )
 
-    def _active_pool_idxs(self: Scheduler) -> set:
+    @staticmethod
+    def active_pool_idxs(self: "SchedulerPoolStatsObserver") -> set:
         """Pool idxs currently owned by reqs in last_batch / running_batch.
 
         Used to decide which session slots' KV is owned by batch reqs
         (and thus counted via uncached_size, not session_held).
         """
         idxs = set()
-        for batch in [self.last_batch, self.running_batch]:
+        for batch in [self.get_last_batch(), self.get_running_batch()]:
             if batch is None or batch.is_empty():
                 continue
             for req in batch.reqs:
@@ -42,35 +47,51 @@ class SchedulerRuntimeCheckerMixin:
                     idxs.add(req.req_pool_idx)
         return idxs
 
-    def _session_held_tokens(self: Scheduler) -> int:
+    @staticmethod
-        return self.tree_cache.session_held_tokens(self._active_pool_idxs())
+    def session_held_tokens(self: "SchedulerPoolStatsObserver") -> int:
+        return self.tree_cache.session_held_tokens(
+            SchedulerRuntimeCheckerMixin.active_pool_idxs(self)
+        )
 
-    def _session_held_full_tokens(self: Scheduler) -> int:
+    @staticmethod
-        return self.tree_cache.session_held_full_tokens(self._active_pool_idxs())
+    def session_held_full_tokens(self: "SchedulerPoolStatsObserver") -> int:
+        return self.tree_cache.session_held_full_tokens(
+            SchedulerRuntimeCheckerMixin.active_pool_idxs(self)
+        )
 
-    def _session_held_swa_tokens(self: Scheduler) -> int:
+    @staticmethod
-        return self.tree_cache.session_held_swa_tokens(self._active_pool_idxs())
+    def session_held_swa_tokens(self: "SchedulerPoolStatsObserver") -> int:
+        return self.tree_cache.session_held_swa_tokens(
+            SchedulerRuntimeCheckerMixin.active_pool_idxs(self)
+        )
 
-    def _session_held_req_count(self: Scheduler) -> int:
+    @staticmethod
+    def session_held_req_count(self: "SchedulerPoolStatsObserver") -> int:
         return self.tree_cache.session_held_req_count()
 
-    def _session_held_mamba_slots(self: Scheduler) -> int:
+    @staticmethod
-        return self.tree_cache.session_held_mamba_slots(self._active_pool_idxs())
+    def session_held_mamba_slots(self: "SchedulerPoolStatsObserver") -> int:
+        return self.tree_cache.session_held_mamba_slots(
+            SchedulerRuntimeCheckerMixin.active_pool_idxs(self)
+        )
 
-    def get_pool_stats(self: Scheduler) -> PoolStats:
+    @staticmethod
+    def get_pool_stats(self: "SchedulerPoolStatsObserver") -> PoolStats:
         if self.is_hybrid_swa:
-            pool_stats = self._get_swa_token_info()
+            pool_stats = SchedulerRuntimeCheckerMixin._get_swa_token_info(self)
         elif self.is_hybrid_ssm:
-            pool_stats = self._get_mamba_token_info()
+            pool_stats = SchedulerRuntimeCheckerMixin._get_mamba_token_info(self)
         else:
-            pool_stats = self._get_token_info()
+            pool_stats = SchedulerRuntimeCheckerMixin._get_token_info(self)
 
         if self.enable_hisparse:
-            pool_stats = self._get_hisparse_token_info(pool_stats)
+            pool_stats = SchedulerRuntimeCheckerMixin._get_hisparse_token_info(
+                self, pool_stats
+            )
 
         # swa + ssm can coexist: overlay mamba fields onto swa stats
         if self.is_hybrid_ssm:
-            mamba_stats = self._get_mamba_token_info()
+            mamba_stats = SchedulerRuntimeCheckerMixin._get_mamba_token_info(self)
             pool_stats.is_hybrid_ssm = True
             pool_stats.mamba_num_used = mamba_stats.mamba_num_used
             pool_stats.mamba_usage = mamba_stats.mamba_usage
@@ -79,7 +100,8 @@ class SchedulerRuntimeCheckerMixin:
 
         return pool_stats
 
-    def _get_token_info(self: Scheduler) -> PoolStats:
+    @staticmethod
+    def _get_token_info(self: "SchedulerPoolStatsObserver") -> PoolStats:
         available_size = self.token_to_kv_pool_allocator.available_size()
         evictable_size = self.tree_cache.evictable_size()
         num_used = self.max_total_num_tokens - (available_size + evictable_size)
@@ -91,7 +113,10 @@ class SchedulerRuntimeCheckerMixin:
             full_evictable_size=evictable_size,
         )
 
-    def _get_hisparse_token_info(self: Scheduler, pool_stats: PoolStats) -> PoolStats:
+    @staticmethod
+    def _get_hisparse_token_info(
+        self: "SchedulerPoolStatsObserver", pool_stats: PoolStats
+    ) -> PoolStats:
         if self.enable_hisparse and self.hisparse_coordinator is not None:
             h = self.hisparse_coordinator.get_token_stats()
             return dataclasses.replace(
@@ -104,7 +129,8 @@ class SchedulerRuntimeCheckerMixin:
             )
         return pool_stats
 
-    def _get_mamba_token_info(self: Scheduler):
+    @staticmethod
+    def _get_mamba_token_info(self: "SchedulerPoolStatsObserver"):
         is_mamba_radix_cache = (
             self.tree_cache.supports_mamba() and self.tree_cache.is_tree_cache()
         )
@@ -137,7 +163,8 @@ class SchedulerRuntimeCheckerMixin:
             mamba_evictable_size=mamba_evictable_size,
         )
 
-    def _get_swa_token_info(self: Scheduler) -> PoolStats:
+    @staticmethod
+    def _get_swa_token_info(self: "SchedulerPoolStatsObserver") -> PoolStats:
         full_available_size = self.token_to_kv_pool_allocator.full_available_size()
         full_evictable_size = self.tree_cache.full_evictable_size()
         swa_available_size = self.token_to_kv_pool_allocator.swa_available_size()
@@ -194,15 +221,21 @@ class SchedulerRuntimeCheckerMixin:
     ) -> Tuple[bool, str]:
         if self.is_hybrid_swa:
             protected = self.tree_cache.full_protected_size()
-            session_held = self._session_held_full_tokens()
+            session_held = self.session_held_full_tokens(
+                self.pool_stats_observer,
+            )
             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._session_held_tokens()
+            session_held = self.session_held_tokens(
+                self.pool_stats_observer,
+            )
             total = self.token_to_kv_pool_allocator.size
         else:
             protected = self.tree_cache.protected_size()
-            session_held = self._session_held_tokens()
+            session_held = self.session_held_tokens(
+                self.pool_stats_observer,
+            )
             total = self.max_total_num_tokens
         return self._check_pool_invariant(
             "full",
@@ -222,7 +255,9 @@ class SchedulerRuntimeCheckerMixin:
             ps.swa_available_size,
             ps.swa_evictable_size,
             self.tree_cache.swa_protected_size(),
-            self._session_held_swa_tokens(),
+            self.session_held_swa_tokens(
+                self.pool_stats_observer,
+            ),
             self.swa_tokens_per_layer,
             uncached,
         )
@@ -233,7 +268,9 @@ class SchedulerRuntimeCheckerMixin:
             ps.mamba_available_size,
             ps.mamba_evictable_size,
             self.tree_cache.mamba_protected_size(),
-            self._session_held_mamba_slots(),
+            self.session_held_mamba_slots(
+                self.pool_stats_observer,
+            ),
             self.req_to_token_pool.mamba_pool.size,
         )
         if leak:
@@ -314,7 +351,9 @@ class SchedulerRuntimeCheckerMixin:
             )
             return
 
-        ps = self.get_pool_stats()
+        ps = self.get_pool_stats(
+            self.pool_stats_observer,
+        )
         full_uncached, swa_uncached = self._get_total_uncached_sizes()
 
         full_leak, full_msg = self._check_full_pool(ps, uncached=full_uncached)
@@ -338,7 +377,9 @@ class SchedulerRuntimeCheckerMixin:
         else:
             req_total_size = self.req_to_token_pool.size
 
-        session_req_count = self._session_held_req_count()
+        session_req_count = self.session_held_req_count(
+            self.pool_stats_observer,
+        )
         if len(self.req_to_token_pool.free_slots) + session_req_count != req_total_size:
             msg = (
                 "req_to_token_pool memory leak detected!"
@@ -393,9 +434,13 @@ class SchedulerRuntimeCheckerMixin:
         ):
             return
 
-        self.get_pool_stats().update_scheduler_stats(self.stats)
+        self.get_pool_stats(self.pool_stats_observer).update_scheduler_stats(self.stats)
-        self.stats.num_streaming_sessions = self._streaming_session_count()
+        self.stats.num_streaming_sessions = self.streaming_session_count(
-        self.stats.streaming_session_held_tokens = self._session_held_tokens()
+            self.pool_stats_observer,
+        )
+        self.stats.streaming_session_held_tokens = self.session_held_tokens(
+            self.pool_stats_observer,
+        )
 
         priority_enabled = self.enable_priority_scheduling
         self.stats.num_running_reqs = QueueCount.from_reqs(
@@ -437,7 +482,11 @@ def create_scheduler_watchdog(
     def dump_info() -> str:
         if scheduler.is_initializing:
             return ""
-        _, messages = scheduler._check_all_pools(scheduler.get_pool_stats())
+        _, messages = scheduler._check_all_pools(
+            scheduler.get_pool_stats(
+                scheduler.pool_stats_observer,
+            )
+        )
         return (
             f"{scheduler.cur_batch.batch_size()=}\n"
             f"{scheduler.cur_batch.reqs=}\n" + "\n".join(messages)

python/sglang/srt/observability/scheduler_metrics_mixin.py

@@ -499,7 +499,9 @@ class SchedulerMetricsMixin:
             prefill_stats.log_input_tokens / gap_latency if gap_latency > 0 else 0.0
         )
 
-        pool_stats = self.get_pool_stats()
+        pool_stats = self.get_pool_stats(
+            self.pool_stats_observer,
+        )
         token_usage_msg = ", ".join(pool_stats.get_prefill_usage_msg_parts()) + ", "
 
         self.stats.new_token_ratio = prefill_stats.new_token_ratio
@@ -660,7 +662,9 @@ class SchedulerMetricsMixin:
         self.num_generated_tokens = 0
         num_running_reqs = len(batch.reqs)
 
-        pool_stats = self.get_pool_stats()
+        pool_stats = self.get_pool_stats(
+            self.pool_stats_observer,
+        )
         token_usage_msg = ", ".join(pool_stats.get_decode_usage_msg_parts()) + ", "
 
         if RECORD_STEP_TIME:
@@ -779,8 +783,12 @@ class SchedulerMetricsMixin:
                 )
 
             # Streaming session metrics
-            self.stats.num_streaming_sessions = self._streaming_session_count()
+            self.stats.num_streaming_sessions = self.streaming_session_count(
-            self.stats.streaming_session_held_tokens = self._session_held_tokens()
+                self.pool_stats_observer,
+            )
+            self.stats.streaming_session_held_tokens = self.session_held_tokens(
+                self.pool_stats_observer,
+            )
 
             # Routing key metrics
             # (to reduce the overhead, we only compute this when all requests have routing_key)
@@ -1011,7 +1019,9 @@ class SchedulerMetricsMixin:
             waiting_queues.append(self.disagg_decode_prealloc_queue.retracted_queue)
 
         num_waiting_reqs = sum(len(queue) for queue in waiting_queues)
-        num_used_tokens, kv_token_usage = self.get_pool_stats().get_kv_token_stats()
+        num_used_tokens, kv_token_usage = self.get_pool_stats(
+            self.pool_stats_observer,
+        ).get_kv_token_stats()
         num_total_tokens = num_used_tokens + sum(
             req.seqlen for queue in waiting_queues for req in queue
         )