diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 44f205c8f..daf499f8a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -173,6 +173,9 @@ from sglang.srt.managers.scheduler_components.invariant_checker import ( from sglang.srt.managers.scheduler_components.kv_events_publisher import ( SchedulerKvEventsPublisher, ) +from sglang.srt.managers.scheduler_components.load_inquirer import ( + SchedulerLoadInquirer, +) from sglang.srt.managers.scheduler_components.pool_stats_observer import ( SchedulerPoolStatsObserver, ) @@ -695,6 +698,28 @@ class Scheduler( get_stats=lambda: self.stats, ) + self.load_inquirer = SchedulerLoadInquirer( + disaggregation_mode=self.disaggregation_mode, + ps=self.ps, + server_args=self.server_args, + max_total_num_tokens=self.max_total_num_tokens, + max_running_requests=self.max_running_requests, + pool_stats_observer=self.pool_stats_observer, + tp_worker=self.tp_worker, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + spec_algorithm=self.spec_algorithm, + get_running_batch=lambda: self.running_batch, + get_waiting_queue=lambda: self.waiting_queue, + get_stats=lambda: self.stats, + get_chunked_req=lambda: self.chunked_req, + get_disagg_prefill_bootstrap_queue=lambda: self.disagg_prefill_bootstrap_queue, + get_disagg_prefill_inflight_queue=lambda: self.disagg_prefill_inflight_queue, + get_disagg_decode_prealloc_queue=lambda: self.disagg_decode_prealloc_queue, + get_disagg_decode_transfer_queue=lambda: self.disagg_decode_transfer_queue, + get_spec_total_num_accept_tokens=lambda: self.spec_total_num_accept_tokens, + get_spec_total_num_forward_ct=lambda: self.spec_total_num_forward_ct, + ) + self.is_initializing = False def init_zbal_on_npu(self): @@ -1455,7 +1480,10 @@ class Scheduler( self.load_lora_adapter_from_tensors, ), (UnloadLoRAAdapterReqInput, self.unload_lora_adapter), - (GetLoadsReqInput, self.get_loads), + ( + GetLoadsReqInput, + lambda req: self.get_loads(self.load_inquirer, req), + ), (PauseGenerationReqInput, self.pause_generation), (ContinueGenerationReqInput, self.continue_generation), (DumperControlReqInput, self.handle_dumper_control), @@ -2631,11 +2659,12 @@ class Scheduler( self.running_batch.reqs, self.enable_priority_scheduling, num_pending_tokens=self._get_num_pending_tokens( + self.load_inquirer, chunk_deduct=( self.chunked_req.extend_input_len if self.chunked_req is not None else 0 - ) + ), ), ) diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py new file mode 100644 index 000000000..82cc1c9ce --- /dev/null +++ b/python/sglang/srt/managers/scheduler_components/load_inquirer.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable + +from sglang.srt.disaggregation.utils import DisaggregationMode + +if TYPE_CHECKING: + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + from sglang.srt.managers.scheduler_components.pool_stats_observer import ( + SchedulerPoolStatsObserver, + ) + from sglang.srt.managers.tp_worker import BaseTpWorker + from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator + from sglang.srt.server_args import ServerArgs + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + +logger = logging.getLogger(__name__) + + +@dataclass(kw_only=True, slots=True, frozen=True) +class SchedulerLoadInquirer: + disaggregation_mode: "DisaggregationMode" + ps: "ParallelState" + server_args: "ServerArgs" + max_total_num_tokens: int + max_running_requests: int + pool_stats_observer: "SchedulerPoolStatsObserver" + tp_worker: "BaseTpWorker" + token_to_kv_pool_allocator: "BaseTokenToKVPoolAllocator" + spec_algorithm: "SpeculativeAlgorithm" + get_running_batch: Callable + get_waiting_queue: Callable + get_stats: Callable + get_chunked_req: Callable + get_disagg_prefill_bootstrap_queue: Callable + get_disagg_prefill_inflight_queue: Callable + get_disagg_decode_prealloc_queue: Callable + get_disagg_decode_transfer_queue: Callable + get_spec_total_num_accept_tokens: Callable + get_spec_total_num_forward_ct: Callable diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 85ee86d2a..683bee419 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -1050,7 +1050,10 @@ class SchedulerOutputProcessorMixin: spec_correct_drafts_histogram = [] retraction_counts = [] output_hidden_states = None - load = self.get_loads(GetLoadsReqInput(include=["core"])) + load = self.get_loads( + self.load_inquirer, + GetLoadsReqInput(include=["core"]), + ) routed_experts = None indexer_topk = None customized_info = {} diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 263556d5c..911a4ad03 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -34,6 +34,9 @@ if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_policy import PrefillAdder from sglang.srt.managers.scheduler import EmbeddingBatchResult, Scheduler + from sglang.srt.managers.scheduler_components.load_inquirer import ( + SchedulerLoadInquirer, + ) logger = logging.getLogger(__name__) @@ -913,7 +916,10 @@ class SchedulerMetricsMixin: self.stats.token_usage / 0.9, ) - def _get_num_pending_tokens(self: Scheduler, chunk_deduct: int = 0) -> int: + @staticmethod + def _get_num_pending_tokens( + self: "SchedulerLoadInquirer", chunk_deduct: int = 0 + ) -> int: """Get the total number of tokens pending prefill. This includes tokens from waiting queue requests plus remaining tokens @@ -927,13 +933,16 @@ class SchedulerMetricsMixin: time ``prefix_indices`` is already up-to-date, so the default 0 is correct. """ - num_pending_tokens = sum(req.seqlen for req in self.waiting_queue) - if self.chunked_req is not None: - req = self.chunked_req + num_pending_tokens = sum(req.seqlen for req in self.get_waiting_queue()) + if self.get_chunked_req() is not None: + req = self.get_chunked_req() num_pending_tokens += req.seqlen - len(req.prefix_indices) - chunk_deduct return num_pending_tokens - def get_loads(self: Scheduler, req: GetLoadsReqInput = None) -> GetLoadsReqOutput: + @staticmethod + def get_loads( + self: "SchedulerLoadInquirer", req: GetLoadsReqInput = None + ) -> GetLoadsReqOutput: """ Get comprehensive load metrics for /v1/loads endpoint. @@ -949,15 +958,17 @@ class SchedulerMetricsMixin: include = set(req.include) if req.include else {"core"} include_all = "all" in include - num_running_reqs = len(self.running_batch.reqs) + num_running_reqs = len(self.get_running_batch().reqs) - waiting_queues = [self.waiting_queue] + waiting_queues = [self.get_waiting_queue()] if self.disaggregation_mode == DisaggregationMode.PREFILL: - waiting_queues.append(self.disagg_prefill_bootstrap_queue.queue) + waiting_queues.append(self.get_disagg_prefill_bootstrap_queue().queue) elif self.disaggregation_mode == DisaggregationMode.DECODE: - waiting_queues.append(self.disagg_decode_prealloc_queue.queue) - waiting_queues.append(self.disagg_decode_transfer_queue.queue) - waiting_queues.append(self.disagg_decode_prealloc_queue.retracted_queue) + waiting_queues.append(self.get_disagg_decode_prealloc_queue().queue) + waiting_queues.append(self.get_disagg_decode_transfer_queue().queue) + waiting_queues.append( + self.get_disagg_decode_prealloc_queue().retracted_queue + ) num_waiting_reqs = sum(len(queue) for queue in waiting_queues) num_used_tokens, kv_token_usage = ( @@ -985,22 +996,25 @@ class SchedulerMetricsMixin: speculative = None if include_all or "spec" in include: - if not self.spec_algorithm.is_none() and self.spec_total_num_forward_ct > 0: + if ( + not self.spec_algorithm.is_none() + and self.get_spec_total_num_forward_ct() > 0 + ): speculative = SpeculativeMetrics( accept_length=( - self.spec_total_num_accept_tokens - / self.spec_total_num_forward_ct + self.get_spec_total_num_accept_tokens() + / self.get_spec_total_num_forward_ct() ), - accept_rate=self.stats.spec_accept_rate, + accept_rate=self.get_stats().spec_accept_rate, ) lora = None if include_all or "lora" in include: - if self.enable_lora: + if self.server_args.enable_lora: lora = LoRAMetrics( - slots_used=self.stats.lora_pool_slots_used, - slots_total=self.stats.lora_pool_slots_total, - utilization=self.stats.lora_pool_utilization, + slots_used=self.get_stats().lora_pool_slots_used, + slots_total=self.get_stats().lora_pool_slots_total, + utilization=self.get_stats().lora_pool_utilization, ) disaggregation = None @@ -1014,14 +1028,14 @@ class SchedulerMetricsMixin: if self.disaggregation_mode == DisaggregationMode.PREFILL: mode_str = "prefill" - prefill_bootstrap = len(self.disagg_prefill_bootstrap_queue.queue) - prefill_inflight = len(self.disagg_prefill_inflight_queue) + prefill_bootstrap = len(self.get_disagg_prefill_bootstrap_queue().queue) + prefill_inflight = len(self.get_disagg_prefill_inflight_queue()) elif self.disaggregation_mode == DisaggregationMode.DECODE: mode_str = "decode" - decode_prealloc = len(self.disagg_decode_prealloc_queue.queue) - decode_transfer = len(self.disagg_decode_transfer_queue.queue) + decode_prealloc = len(self.get_disagg_decode_prealloc_queue().queue) + decode_transfer = len(self.get_disagg_decode_transfer_queue().queue) decode_retracted = len( - self.disagg_decode_prealloc_queue.retracted_queue + self.get_disagg_decode_prealloc_queue().retracted_queue ) disaggregation = DisaggregationMetrics( @@ -1031,17 +1045,17 @@ class SchedulerMetricsMixin: decode_prealloc_queue_reqs=decode_prealloc, decode_transfer_queue_reqs=decode_transfer, decode_retracted_queue_reqs=decode_retracted, - kv_transfer_speed_gb_s=self.stats.kv_transfer_speed_gb_s, - kv_transfer_latency_ms=self.stats.kv_transfer_latency_ms, + kv_transfer_speed_gb_s=self.get_stats().kv_transfer_speed_gb_s, + kv_transfer_latency_ms=self.get_stats().kv_transfer_latency_ms, ) queues = None if include_all or "queues" in include: queues = QueueMetrics( - waiting=len(self.waiting_queue), - grammar=self.stats.num_grammar_queue_reqs, - paused=self.stats.num_paused_reqs, - retracted=self.stats.num_retracted_reqs, + waiting=len(self.get_waiting_queue()), + grammar=self.get_stats().num_grammar_queue_reqs, + paused=self.get_stats().num_paused_reqs, + retracted=self.get_stats().num_retracted_reqs, ) return GetLoadsReqOutput( @@ -1053,9 +1067,9 @@ class SchedulerMetricsMixin: num_total_tokens=num_total_tokens, max_total_num_tokens=self.max_total_num_tokens, token_usage=round(kv_token_usage, 4), - gen_throughput=round(self.stats.gen_throughput, 2), - cache_hit_rate=round(self.stats.cache_hit_rate, 4), - utilization=round(self.stats.utilization, 4), + gen_throughput=round(self.get_stats().gen_throughput, 2), + cache_hit_rate=round(self.get_stats().cache_hit_rate, 4), + utilization=round(self.get_stats().utilization, 4), max_running_requests=self.max_running_requests, memory=memory, speculative=speculative,