minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
fzyzcjy
2026-05-18 18:36:36 +08:00
committed by GitHub
parent c58b47bc86
commit 9fdf73f393
4 changed files with 155 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+27 -2
View File
@@ -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
+31 -1
View File
@@ -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
+82 -33
View File
@@ -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:
return self.tree_cache.session_held_tokens(self._active_pool_idxs())
@staticmethod
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:
return self.tree_cache.session_held_full_tokens(self._active_pool_idxs())
@staticmethod
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:
return self.tree_cache.session_held_swa_tokens(self._active_pool_idxs())
@staticmethod
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:
return self.tree_cache.session_held_mamba_slots(self._active_pool_idxs())
@staticmethod
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.stats.num_streaming_sessions = self._streaming_session_count()
self.stats.streaming_session_held_tokens = self._session_held_tokens()
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.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
+15 -5
View File
@@ -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.streaming_session_held_tokens = self._session_held_tokens()
self.stats.num_streaming_sessions = self.streaming_session_count(
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
)