From ee392a1e144f11f43b958618cdac33f55461ed37 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 18 May 2026 18:37:50 +0800 Subject: [PATCH] Move pool-stats sampling to SchedulerPoolStatsObserver (#25621) --- python/sglang/srt/managers/scheduler.py | 10 +- .../pool_stats_observer.py | 150 ++++++++++++ .../scheduler_runtime_checker_mixin.py | 221 ++---------------- .../observability/scheduler_metrics_mixin.py | 22 +- 4 files changed, 176 insertions(+), 227 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2892ba071..0f2899f1c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2343,9 +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( - self.pool_stats_observer, - ).get_max_pool_usage() + max_pool_usage = ( + self.pool_stats_observer.get_pool_stats().get_max_pool_usage() + ) prefill_delayer_single_pass = PrefillDelayerSinglePassExecutor( self.prefill_delayer, token_usage=max_pool_usage ) @@ -3061,9 +3061,7 @@ 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.get_pool_stats( - self.pool_stats_observer, - ) + self.pool_stats_observer.get_pool_stats() ) if has_leak: self._report_leak("pool", "\n".join(messages)) diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 986d473da..a6ed75228 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -153,3 +153,153 @@ class SchedulerPoolStatsObserver: max_total_num_tokens: int get_last_batch: Callable get_running_batch: Callable + + def streaming_session_count(self) -> int: + return sum( + 1 + for session in self.session_controller.sessions.values() + if session.streaming + ) + + def active_pool_idxs(self) -> 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.get_last_batch(), self.get_running_batch()]: + if batch is None or batch.is_empty(): + continue + for req in batch.reqs: + if req.req_pool_idx is not None: + idxs.add(req.req_pool_idx) + return idxs + + def session_held_tokens(self) -> int: + return self.tree_cache.session_held_tokens(self.active_pool_idxs()) + + def session_held_full_tokens(self) -> int: + return self.tree_cache.session_held_full_tokens(self.active_pool_idxs()) + + def session_held_swa_tokens(self) -> int: + return self.tree_cache.session_held_swa_tokens(self.active_pool_idxs()) + + def session_held_req_count(self) -> int: + return self.tree_cache.session_held_req_count() + + def session_held_mamba_slots(self) -> int: + return self.tree_cache.session_held_mamba_slots(self.active_pool_idxs()) + + def get_pool_stats(self) -> PoolStats: + if self.is_hybrid_swa: + pool_stats = self._get_swa_token_info() + elif self.is_hybrid_ssm: + pool_stats = self._get_mamba_token_info() + else: + pool_stats = self._get_token_info() + + if self.enable_hisparse: + pool_stats = self._get_hisparse_token_info(pool_stats) + + # swa + ssm can coexist: overlay mamba fields onto swa stats + if self.is_hybrid_ssm: + mamba_stats = self._get_mamba_token_info() + pool_stats.is_hybrid_ssm = True + pool_stats.mamba_num_used = mamba_stats.mamba_num_used + pool_stats.mamba_usage = mamba_stats.mamba_usage + pool_stats.mamba_available_size = mamba_stats.mamba_available_size + pool_stats.mamba_evictable_size = mamba_stats.mamba_evictable_size + + return pool_stats + + def _get_token_info(self) -> 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) + token_usage = num_used / self.max_total_num_tokens + return PoolStats( + full_num_used=num_used, + full_token_usage=token_usage, + full_available_size=available_size, + full_evictable_size=evictable_size, + ) + + def _get_hisparse_token_info(self, 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( + pool_stats, + is_hisparse=True, + hisparse_device_tokens=h.device_tokens, + hisparse_device_token_usage=h.device_token_usage, + hisparse_host_tokens=h.host_tokens, + hisparse_host_token_usage=h.host_token_usage, + ) + return pool_stats + + def _get_mamba_token_info(self): + is_mamba_radix_cache = ( + self.tree_cache.supports_mamba() and self.tree_cache.is_tree_cache() + ) + full_available_size = self.token_to_kv_pool_allocator.available_size() + full_evictable_size = ( + self.tree_cache.full_evictable_size() if is_mamba_radix_cache else 0 + ) + mamba_available_size = self.req_to_token_pool.mamba_pool.available_size() + mamba_evictable_size = ( + self.tree_cache.mamba_evictable_size() if is_mamba_radix_cache else 0 + ) + full_num_used = self.token_to_kv_pool_allocator.size - ( + full_available_size + full_evictable_size + ) + mamba_num_used = self.req_to_token_pool.mamba_pool.size - ( + mamba_available_size + mamba_evictable_size + ) + full_token_usage = full_num_used / self.token_to_kv_pool_allocator.size + mamba_usage = mamba_num_used / self.req_to_token_pool.mamba_pool.size + + return PoolStats( + is_hybrid_ssm=True, + full_num_used=full_num_used, + full_token_usage=full_token_usage, + full_available_size=full_available_size, + full_evictable_size=full_evictable_size, + mamba_num_used=mamba_num_used, + mamba_usage=mamba_usage, + mamba_available_size=mamba_available_size, + mamba_evictable_size=mamba_evictable_size, + ) + + def _get_swa_token_info(self) -> 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() + swa_evictable_size = self.tree_cache.swa_evictable_size() + full_num_used = self.full_tokens_per_layer - ( + full_available_size + full_evictable_size + ) + swa_num_used = self.swa_tokens_per_layer - ( + swa_available_size + swa_evictable_size + ) + # FIXME(hisparse): host-backup transiently over-releases the device pool + # counter, producing negative full_num_used / swa_num_used. We clamp to 0 + # to keep token_usage / leak checks sane, but the underlying accounting + # bug should be fixed so the clamp can go away. + if self.enable_hisparse: + full_num_used = max(0, full_num_used) + swa_num_used = max(0, swa_num_used) + full_token_usage = full_num_used / self.full_tokens_per_layer + swa_token_usage = swa_num_used / self.swa_tokens_per_layer + + return PoolStats( + is_hybrid_swa=True, + full_num_used=full_num_used, + full_token_usage=full_token_usage, + full_available_size=full_available_size, + full_evictable_size=full_evictable_size, + swa_num_used=swa_num_used, + swa_token_usage=swa_token_usage, + swa_available_size=swa_available_size, + swa_evictable_size=swa_evictable_size, + ) diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index f343318e7..c14c7abb6 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -1,6 +1,5 @@ from __future__ import annotations -import dataclasses import logging import time import warnings @@ -8,195 +7,17 @@ from typing import TYPE_CHECKING, List, 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 from sglang.srt.observability.metrics_collector import QueueCount from sglang.srt.utils.common import ceil_align, raise_error_or_warn 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: - @staticmethod - def streaming_session_count(self: "SchedulerPoolStatsObserver") -> int: - return sum( - 1 - for session in self.session_controller.sessions.values() - if session.streaming - ) - - @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.get_last_batch(), self.get_running_batch()]: - if batch is None or batch.is_empty(): - continue - for req in batch.reqs: - if req.req_pool_idx is not None: - idxs.add(req.req_pool_idx) - return idxs - - @staticmethod - def session_held_tokens(self: "SchedulerPoolStatsObserver") -> int: - return self.tree_cache.session_held_tokens( - SchedulerRuntimeCheckerMixin.active_pool_idxs(self) - ) - - @staticmethod - def session_held_full_tokens(self: "SchedulerPoolStatsObserver") -> int: - return self.tree_cache.session_held_full_tokens( - SchedulerRuntimeCheckerMixin.active_pool_idxs(self) - ) - - @staticmethod - def session_held_swa_tokens(self: "SchedulerPoolStatsObserver") -> int: - return self.tree_cache.session_held_swa_tokens( - SchedulerRuntimeCheckerMixin.active_pool_idxs(self) - ) - - @staticmethod - def session_held_req_count(self: "SchedulerPoolStatsObserver") -> int: - return self.tree_cache.session_held_req_count() - - @staticmethod - def session_held_mamba_slots(self: "SchedulerPoolStatsObserver") -> int: - return self.tree_cache.session_held_mamba_slots( - SchedulerRuntimeCheckerMixin.active_pool_idxs(self) - ) - - @staticmethod - def get_pool_stats(self: "SchedulerPoolStatsObserver") -> PoolStats: - if self.is_hybrid_swa: - pool_stats = SchedulerRuntimeCheckerMixin._get_swa_token_info(self) - elif self.is_hybrid_ssm: - pool_stats = SchedulerRuntimeCheckerMixin._get_mamba_token_info(self) - else: - pool_stats = SchedulerRuntimeCheckerMixin._get_token_info(self) - - if self.enable_hisparse: - 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 = 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 - pool_stats.mamba_available_size = mamba_stats.mamba_available_size - pool_stats.mamba_evictable_size = mamba_stats.mamba_evictable_size - - return pool_stats - - @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) - token_usage = num_used / self.max_total_num_tokens - return PoolStats( - full_num_used=num_used, - full_token_usage=token_usage, - full_available_size=available_size, - full_evictable_size=evictable_size, - ) - - @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( - pool_stats, - is_hisparse=True, - hisparse_device_tokens=h.device_tokens, - hisparse_device_token_usage=h.device_token_usage, - hisparse_host_tokens=h.host_tokens, - hisparse_host_token_usage=h.host_token_usage, - ) - return pool_stats - - @staticmethod - def _get_mamba_token_info(self: "SchedulerPoolStatsObserver"): - is_mamba_radix_cache = ( - self.tree_cache.supports_mamba() and self.tree_cache.is_tree_cache() - ) - full_available_size = self.token_to_kv_pool_allocator.available_size() - full_evictable_size = ( - self.tree_cache.full_evictable_size() if is_mamba_radix_cache else 0 - ) - mamba_available_size = self.req_to_token_pool.mamba_pool.available_size() - mamba_evictable_size = ( - self.tree_cache.mamba_evictable_size() if is_mamba_radix_cache else 0 - ) - full_num_used = self.token_to_kv_pool_allocator.size - ( - full_available_size + full_evictable_size - ) - mamba_num_used = self.req_to_token_pool.mamba_pool.size - ( - mamba_available_size + mamba_evictable_size - ) - full_token_usage = full_num_used / self.token_to_kv_pool_allocator.size - mamba_usage = mamba_num_used / self.req_to_token_pool.mamba_pool.size - - return PoolStats( - is_hybrid_ssm=True, - full_num_used=full_num_used, - full_token_usage=full_token_usage, - full_available_size=full_available_size, - full_evictable_size=full_evictable_size, - mamba_num_used=mamba_num_used, - mamba_usage=mamba_usage, - mamba_available_size=mamba_available_size, - mamba_evictable_size=mamba_evictable_size, - ) - - @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() - swa_evictable_size = self.tree_cache.swa_evictable_size() - full_num_used = self.full_tokens_per_layer - ( - full_available_size + full_evictable_size - ) - swa_num_used = self.swa_tokens_per_layer - ( - swa_available_size + swa_evictable_size - ) - # FIXME(hisparse): host-backup transiently over-releases the device pool - # counter, producing negative full_num_used / swa_num_used. We clamp to 0 - # to keep token_usage / leak checks sane, but the underlying accounting - # bug should be fixed so the clamp can go away. - if self.enable_hisparse: - full_num_used = max(0, full_num_used) - swa_num_used = max(0, swa_num_used) - full_token_usage = full_num_used / self.full_tokens_per_layer - swa_token_usage = swa_num_used / self.swa_tokens_per_layer - - return PoolStats( - is_hybrid_swa=True, - full_num_used=full_num_used, - full_token_usage=full_token_usage, - full_available_size=full_available_size, - full_evictable_size=full_evictable_size, - swa_num_used=swa_num_used, - swa_token_usage=swa_token_usage, - swa_available_size=swa_available_size, - swa_evictable_size=swa_evictable_size, - ) - @staticmethod def _check_pool_invariant( pool_name: str, @@ -221,21 +42,15 @@ class SchedulerRuntimeCheckerMixin: ) -> Tuple[bool, str]: if self.is_hybrid_swa: protected = self.tree_cache.full_protected_size() - session_held = self.session_held_full_tokens( - self.pool_stats_observer, - ) + 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.session_held_tokens( - self.pool_stats_observer, - ) + 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.session_held_tokens( - self.pool_stats_observer, - ) + session_held = self.pool_stats_observer.session_held_tokens() total = self.max_total_num_tokens return self._check_pool_invariant( "full", @@ -255,9 +70,7 @@ class SchedulerRuntimeCheckerMixin: ps.swa_available_size, ps.swa_evictable_size, self.tree_cache.swa_protected_size(), - self.session_held_swa_tokens( - self.pool_stats_observer, - ), + self.pool_stats_observer.session_held_swa_tokens(), self.swa_tokens_per_layer, uncached, ) @@ -268,9 +81,7 @@ class SchedulerRuntimeCheckerMixin: ps.mamba_available_size, ps.mamba_evictable_size, self.tree_cache.mamba_protected_size(), - self.session_held_mamba_slots( - self.pool_stats_observer, - ), + self.pool_stats_observer.session_held_mamba_slots(), self.req_to_token_pool.mamba_pool.size, ) if leak: @@ -351,9 +162,7 @@ class SchedulerRuntimeCheckerMixin: ) return - ps = self.get_pool_stats( - self.pool_stats_observer, - ) + 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) @@ -377,9 +186,7 @@ class SchedulerRuntimeCheckerMixin: else: req_total_size = self.req_to_token_pool.size - session_req_count = self.session_held_req_count( - self.pool_stats_observer, - ) + 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!" @@ -434,12 +241,12 @@ class SchedulerRuntimeCheckerMixin: ): return - self.get_pool_stats(self.pool_stats_observer).update_scheduler_stats(self.stats) - self.stats.num_streaming_sessions = self.streaming_session_count( - self.pool_stats_observer, + self.pool_stats_observer.get_pool_stats().update_scheduler_stats(self.stats) + self.stats.num_streaming_sessions = ( + self.pool_stats_observer.streaming_session_count() ) - self.stats.streaming_session_held_tokens = self.session_held_tokens( - self.pool_stats_observer, + self.stats.streaming_session_held_tokens = ( + self.pool_stats_observer.session_held_tokens() ) priority_enabled = self.enable_priority_scheduling @@ -483,9 +290,7 @@ def create_scheduler_watchdog( if scheduler.is_initializing: return "" _, messages = scheduler._check_all_pools( - scheduler.get_pool_stats( - scheduler.pool_stats_observer, - ) + scheduler.pool_stats_observer.get_pool_stats() ) return ( f"{scheduler.cur_batch.batch_size()=}\n" diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index ce91ca89e..c8758dd36 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -499,9 +499,7 @@ class SchedulerMetricsMixin: prefill_stats.log_input_tokens / gap_latency if gap_latency > 0 else 0.0 ) - pool_stats = self.get_pool_stats( - self.pool_stats_observer, - ) + pool_stats = self.pool_stats_observer.get_pool_stats() token_usage_msg = ", ".join(pool_stats.get_prefill_usage_msg_parts()) + ", " self.stats.new_token_ratio = prefill_stats.new_token_ratio @@ -662,9 +660,7 @@ class SchedulerMetricsMixin: self.num_generated_tokens = 0 num_running_reqs = len(batch.reqs) - pool_stats = self.get_pool_stats( - self.pool_stats_observer, - ) + pool_stats = self.pool_stats_observer.get_pool_stats() token_usage_msg = ", ".join(pool_stats.get_decode_usage_msg_parts()) + ", " if RECORD_STEP_TIME: @@ -783,11 +779,11 @@ class SchedulerMetricsMixin: ) # Streaming session metrics - self.stats.num_streaming_sessions = self.streaming_session_count( - self.pool_stats_observer, + self.stats.num_streaming_sessions = ( + self.pool_stats_observer.streaming_session_count() ) - self.stats.streaming_session_held_tokens = self.session_held_tokens( - self.pool_stats_observer, + self.stats.streaming_session_held_tokens = ( + self.pool_stats_observer.session_held_tokens() ) # Routing key metrics @@ -1019,9 +1015,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( - self.pool_stats_observer, - ).get_kv_token_stats() + num_used_tokens, kv_token_usage = ( + self.pool_stats_observer.get_pool_stats().get_kv_token_stats() + ) num_total_tokens = num_used_tokens + sum( req.seqlen for queue in waiting_queues for req in queue )