From c58b47bc86d08b5a07ee6bf18eb16a64e728c000 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 18 May 2026 18:35:27 +0800 Subject: [PATCH] Move PoolStats dataclass to scheduler_components.pool_stats_observer (#25618) --- .../pool_stats_observer.py | 125 ++++++++++++++++++ .../scheduler_runtime_checker_mixin.py | 122 +---------------- .../test_scheduler_pause_generation.py | 2 +- 3 files changed, 128 insertions(+), 121 deletions(-) create mode 100644 python/sglang/srt/managers/scheduler_components/pool_stats_observer.py diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py new file mode 100644 index 000000000..331e22a47 --- /dev/null +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import dataclasses +from typing import List, Optional, Tuple + + +class SchedulerStats: ... # type: ignore[no-redef] + + +@dataclasses.dataclass +class PoolStats: + # For full pools (required) + full_num_used: int + full_token_usage: float + full_available_size: int + full_evictable_size: int + + is_hybrid_swa: bool = False + is_hybrid_ssm: bool = False + is_hisparse: bool = False + + # For hybrid-swa pools + swa_num_used: Optional[int] = None + swa_token_usage: Optional[float] = None + swa_available_size: Optional[int] = None + swa_evictable_size: Optional[int] = None + + # For mamba pools + mamba_num_used: Optional[int] = None + mamba_usage: Optional[float] = None + mamba_available_size: Optional[int] = None + mamba_evictable_size: Optional[int] = None + + # HiSparse device/host breakdown for decode logs (plain KV pool only) + hisparse_device_tokens: Optional[int] = None + hisparse_device_token_usage: Optional[float] = None + hisparse_host_tokens: Optional[int] = None + hisparse_host_token_usage: Optional[float] = None + + def get_kv_token_stats(self) -> Tuple[int, float]: + # NOTE: mamba pool is not included in the "token usage" calculation. + if self.is_hybrid_swa: + num_used = max(self.full_num_used, self.swa_num_used) + token_usage = max(self.full_token_usage, self.swa_token_usage) + else: + num_used = self.full_num_used + token_usage = self.full_token_usage + + return num_used, token_usage + + def get_max_pool_usage(self) -> float: + usage = self.full_token_usage + if self.is_hybrid_swa: + usage = max(usage, self.swa_token_usage) + if self.is_hybrid_ssm: + usage = max(usage, self.mamba_usage) + assert usage is not None and usage >= 0, f"{usage=} is not valid" + return usage + + def get_prefill_usage_msg_parts(self) -> List[str]: + parts = [] + if self.is_hybrid_swa: + parts += [ + f"full token usage: {self.full_token_usage:.2f}", + f"swa token usage: {self.swa_token_usage:.2f}", + ] + if self.is_hybrid_ssm: + if not self.is_hybrid_swa: + parts.append(f"full token usage: {self.full_token_usage:.2f}") + parts.append(f"mamba usage: {self.mamba_usage:.2f}") + if not parts: + parts.append(f"token usage: {self.full_token_usage:.2f}") + return parts + + def get_decode_usage_msg_parts(self) -> List[str]: + parts = [] + if self.is_hybrid_swa: + parts += [ + f"#full token: {self.full_num_used}", + f"full token usage: {self.full_token_usage:.2f}", + f"#swa token: {self.swa_num_used}", + f"swa token usage: {self.swa_token_usage:.2f}", + ] + if self.is_hybrid_ssm: + if not self.is_hybrid_swa: + parts += [ + f"#full token: {self.full_num_used}", + f"full token usage: {self.full_token_usage:.2f}", + ] + parts += [ + f"mamba num: {self.mamba_num_used}", + f"mamba usage: {self.mamba_usage:.2f}", + ] + if self.is_hisparse: + parts += [ + f"#gpu token: {self.hisparse_device_tokens}", + f"gpu token usage: {self.hisparse_device_token_usage:.2f}", + f"#cpu token: {self.hisparse_host_tokens}", + f"cpu token usage: {self.hisparse_host_token_usage:.2f}", + ] + if not parts: + parts.append( + f"#token: {self.full_num_used}, token usage: {self.full_token_usage:.2f}" + ) + return parts + + def update_scheduler_stats(self, stats: SchedulerStats) -> None: + """Update pool-related fields on SchedulerStats.""" + num_used, _ = self.get_kv_token_stats() + stats.num_used_tokens = num_used + stats.token_usage = round(self.get_max_pool_usage(), 2) + stats.full_token_usage = self.full_token_usage + if self.is_hybrid_swa: + stats.swa_token_usage = self.swa_token_usage + stats.swa_available_tokens = self.swa_available_size + stats.swa_evictable_tokens = self.swa_evictable_size + stats.swa_used_tokens = self.swa_num_used + if self.is_hybrid_ssm: + stats.mamba_usage = self.mamba_usage + stats.mamba_available_tokens = self.mamba_available_size + stats.mamba_evictable_tokens = self.mamba_evictable_size + stats.mamba_used_tokens = self.mamba_num_used + stats.kv_available_tokens = self.full_available_size + stats.kv_evictable_tokens = self.full_evictable_size + stats.kv_used_tokens = self.full_num_used diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 9d81c6e3a..419bca683 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -4,139 +4,21 @@ import dataclasses import logging import time import warnings -from typing import TYPE_CHECKING, List, Optional, Tuple +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.observability.metrics_collector import SchedulerStats logger = logging.getLogger(__name__) -@dataclasses.dataclass -class PoolStats: - # For full pools (required) - full_num_used: int - full_token_usage: float - full_available_size: int - full_evictable_size: int - - is_hybrid_swa: bool = False - is_hybrid_ssm: bool = False - is_hisparse: bool = False - - # For hybrid-swa pools - swa_num_used: Optional[int] = None - swa_token_usage: Optional[float] = None - swa_available_size: Optional[int] = None - swa_evictable_size: Optional[int] = None - - # For mamba pools - mamba_num_used: Optional[int] = None - mamba_usage: Optional[float] = None - mamba_available_size: Optional[int] = None - mamba_evictable_size: Optional[int] = None - - # HiSparse device/host breakdown for decode logs (plain KV pool only) - hisparse_device_tokens: Optional[int] = None - hisparse_device_token_usage: Optional[float] = None - hisparse_host_tokens: Optional[int] = None - hisparse_host_token_usage: Optional[float] = None - - def get_kv_token_stats(self) -> Tuple[int, float]: - # NOTE: mamba pool is not included in the "token usage" calculation. - if self.is_hybrid_swa: - num_used = max(self.full_num_used, self.swa_num_used) - token_usage = max(self.full_token_usage, self.swa_token_usage) - else: - num_used = self.full_num_used - token_usage = self.full_token_usage - - return num_used, token_usage - - def get_max_pool_usage(self) -> float: - usage = self.full_token_usage - if self.is_hybrid_swa: - usage = max(usage, self.swa_token_usage) - if self.is_hybrid_ssm: - usage = max(usage, self.mamba_usage) - assert usage is not None and usage >= 0, f"{usage=} is not valid" - return usage - - def get_prefill_usage_msg_parts(self) -> List[str]: - parts = [] - if self.is_hybrid_swa: - parts += [ - f"full token usage: {self.full_token_usage:.2f}", - f"swa token usage: {self.swa_token_usage:.2f}", - ] - if self.is_hybrid_ssm: - if not self.is_hybrid_swa: - parts.append(f"full token usage: {self.full_token_usage:.2f}") - parts.append(f"mamba usage: {self.mamba_usage:.2f}") - if not parts: - parts.append(f"token usage: {self.full_token_usage:.2f}") - return parts - - def get_decode_usage_msg_parts(self) -> List[str]: - parts = [] - if self.is_hybrid_swa: - parts += [ - f"#full token: {self.full_num_used}", - f"full token usage: {self.full_token_usage:.2f}", - f"#swa token: {self.swa_num_used}", - f"swa token usage: {self.swa_token_usage:.2f}", - ] - if self.is_hybrid_ssm: - if not self.is_hybrid_swa: - parts += [ - f"#full token: {self.full_num_used}", - f"full token usage: {self.full_token_usage:.2f}", - ] - parts += [ - f"mamba num: {self.mamba_num_used}", - f"mamba usage: {self.mamba_usage:.2f}", - ] - if self.is_hisparse: - parts += [ - f"#gpu token: {self.hisparse_device_tokens}", - f"gpu token usage: {self.hisparse_device_token_usage:.2f}", - f"#cpu token: {self.hisparse_host_tokens}", - f"cpu token usage: {self.hisparse_host_token_usage:.2f}", - ] - if not parts: - parts.append( - f"#token: {self.full_num_used}, token usage: {self.full_token_usage:.2f}" - ) - return parts - - def update_scheduler_stats(self, stats: SchedulerStats) -> None: - """Update pool-related fields on SchedulerStats.""" - num_used, _ = self.get_kv_token_stats() - stats.num_used_tokens = num_used - stats.token_usage = round(self.get_max_pool_usage(), 2) - stats.full_token_usage = self.full_token_usage - if self.is_hybrid_swa: - stats.swa_token_usage = self.swa_token_usage - stats.swa_available_tokens = self.swa_available_size - stats.swa_evictable_tokens = self.swa_evictable_size - stats.swa_used_tokens = self.swa_num_used - if self.is_hybrid_ssm: - stats.mamba_usage = self.mamba_usage - stats.mamba_available_tokens = self.mamba_available_size - stats.mamba_evictable_tokens = self.mamba_evictable_size - stats.mamba_used_tokens = self.mamba_num_used - stats.kv_available_tokens = self.full_available_size - stats.kv_evictable_tokens = self.full_evictable_size - stats.kv_used_tokens = self.full_num_used - - class SchedulerRuntimeCheckerMixin: def _streaming_session_count(self: Scheduler) -> int: return sum( diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index c6200584c..2769f8a7a 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -9,7 +9,7 @@ maybe_stub_sgl_kernel() from sglang.srt.managers.io_struct import PauseGenerationReqInput from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_runtime_checker_mixin import PoolStats +from sglang.srt.managers.scheduler_components.pool_stats_observer import PoolStats register_cpu_ci(est_time=15, suite="base-a-test-cpu")