Carve out SchedulerLoadInquirer for queue-load state (#25627)

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
-                )
+                ),
             ),
         )
 

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

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 = {}

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)
+        num_pending_tokens = sum(req.seqlen for req in self.get_waiting_queue())
-        if self.chunked_req is not None:
+        if self.get_chunked_req() is not None:
-            req = self.chunked_req
+            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.get_disagg_decode_prealloc_queue().queue)
-            waiting_queues.append(self.disagg_decode_transfer_queue.queue)
+            waiting_queues.append(self.get_disagg_decode_transfer_queue().queue)
-            waiting_queues.append(self.disagg_decode_prealloc_queue.retracted_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.get_spec_total_num_accept_tokens()
-                        / self.spec_total_num_forward_ct
+                        / 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_used=self.get_stats().lora_pool_slots_used,
-                    slots_total=self.stats.lora_pool_slots_total,
+                    slots_total=self.get_stats().lora_pool_slots_total,
-                    utilization=self.stats.lora_pool_utilization,
+                    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_bootstrap = len(self.get_disagg_prefill_bootstrap_queue().queue)
-                prefill_inflight = len(self.disagg_prefill_inflight_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_prealloc = len(self.get_disagg_decode_prealloc_queue().queue)
-                decode_transfer = len(self.disagg_decode_transfer_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_speed_gb_s=self.get_stats().kv_transfer_speed_gb_s,
-                kv_transfer_latency_ms=self.stats.kv_transfer_latency_ms,
+                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),
+                waiting=len(self.get_waiting_queue()),
-                grammar=self.stats.num_grammar_queue_reqs,
+                grammar=self.get_stats().num_grammar_queue_reqs,
-                paused=self.stats.num_paused_reqs,
+                paused=self.get_stats().num_paused_reqs,
-                retracted=self.stats.num_retracted_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),
+            gen_throughput=round(self.get_stats().gen_throughput, 2),
-            cache_hit_rate=round(self.stats.cache_hit_rate, 4),
+            cache_hit_rate=round(self.get_stats().cache_hit_rate, 4),
-            utilization=round(self.stats.utilization, 4),
+            utilization=round(self.get_stats().utilization, 4),
             max_running_requests=self.max_running_requests,
             memory=memory,
             speculative=speculative,