Move queue-load reporting to SchedulerLoadInquirer (#25628)

python/sglang/srt/managers/scheduler.py

@@ -1482,7 +1482,7 @@ class Scheduler(
                 (UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
                 (
                     GetLoadsReqInput,
-                    lambda req: self.get_loads(self.load_inquirer, req),
+                    lambda req: self.load_inquirer.get_loads(req),
                 ),
                 (PauseGenerationReqInput, self.pause_generation),
                 (ContinueGenerationReqInput, self.continue_generation),
@@ -2658,8 +2658,7 @@ class Scheduler(
             adder,
             self.running_batch.reqs,
             self.enable_priority_scheduling,
-            num_pending_tokens=self._get_num_pending_tokens(
+            num_pending_tokens=self.load_inquirer._get_num_pending_tokens(
-                self.load_inquirer,
                 chunk_deduct=(
                     self.chunked_req.extend_input_len
                     if self.chunked_req is not None

python/sglang/srt/managers/scheduler_components/load_inquirer.py

@@ -1,10 +1,20 @@
 from __future__ import annotations
 
 import logging
+import time
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Callable
 
 from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.io_struct import (
+    DisaggregationMetrics,
+    GetLoadsReqInput,
+    GetLoadsReqOutput,
+    LoRAMetrics,
+    MemoryMetrics,
+    QueueMetrics,
+    SpeculativeMetrics,
+)
 
 if TYPE_CHECKING:
     from sglang.srt.distributed.parallel_state_wrapper import ParallelState
@@ -41,3 +51,159 @@ class SchedulerLoadInquirer:
     get_disagg_decode_transfer_queue: Callable
     get_spec_total_num_accept_tokens: Callable
     get_spec_total_num_forward_ct: Callable
+
+    def _get_num_pending_tokens(self, chunk_deduct: int = 0) -> int:
+        """Get the total number of tokens pending prefill.
+
+        This includes tokens from waiting queue requests plus remaining tokens
+        from the currently chunked request.
+
+        Args:
+            chunk_deduct: extra tokens to subtract from the chunked request's
+                remaining count. At batch-scheduling time the current chunk
+                has been planned but ``prefix_indices`` does not yet include it,
+                so callers pass ``extend_input_len`` here. At load-reporting
+                time ``prefix_indices`` is already up-to-date, so the default
+                0 is correct.
+        """
+        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, req: GetLoadsReqInput = None) -> GetLoadsReqOutput:
+        """
+        Get comprehensive load metrics for /v1/loads endpoint.
+
+        Args:
+            req: Request containing include list and optional dp_rank filter
+
+        Returns:
+            GetLoadsReqOutput with core metrics and optional detailed sections
+        """
+        if req is None:
+            req = GetLoadsReqInput()
+
+        include = set(req.include) if req.include else {"core"}
+        include_all = "all" in include
+
+        num_running_reqs = len(self.get_running_batch().reqs)
+
+        waiting_queues = [self.get_waiting_queue()]
+        if self.disaggregation_mode == DisaggregationMode.PREFILL:
+            waiting_queues.append(self.get_disagg_prefill_bootstrap_queue().queue)
+        elif self.disaggregation_mode == DisaggregationMode.DECODE:
+            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 = (
+            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
+        )
+
+        memory = None
+        if include_all or "memory" in include:
+            try:
+                memory = MemoryMetrics(
+                    weight_gb=round(
+                        self.tp_worker.model_runner.weight_load_mem_usage, 3
+                    ),
+                    kv_cache_gb=round(
+                        self.token_to_kv_pool_allocator.get_kvcache().mem_usage, 3
+                    ),
+                    graph_gb=round(self.tp_worker.model_runner.graph_mem_usage, 3),
+                    token_capacity=int(self.max_total_num_tokens),
+                )
+            except AttributeError as e:
+                logger.debug(f"Memory metrics not available: {e}")
+
+        speculative = None
+        if include_all or "spec" in include:
+            if (
+                not self.spec_algorithm.is_none()
+                and self.get_spec_total_num_forward_ct() > 0
+            ):
+                speculative = SpeculativeMetrics(
+                    accept_length=(
+                        self.get_spec_total_num_accept_tokens()
+                        / self.get_spec_total_num_forward_ct()
+                    ),
+                    accept_rate=self.get_stats().spec_accept_rate,
+                )
+
+        lora = None
+        if include_all or "lora" in include:
+            if self.server_args.enable_lora:
+                lora = LoRAMetrics(
+                    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
+        if include_all or "disagg" in include:
+            mode_str = "null"
+            prefill_bootstrap = 0
+            prefill_inflight = 0
+            decode_prealloc = 0
+            decode_transfer = 0
+            decode_retracted = 0
+
+            if self.disaggregation_mode == DisaggregationMode.PREFILL:
+                mode_str = "prefill"
+                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.get_disagg_decode_prealloc_queue().queue)
+                decode_transfer = len(self.get_disagg_decode_transfer_queue().queue)
+                decode_retracted = len(
+                    self.get_disagg_decode_prealloc_queue().retracted_queue
+                )
+
+            disaggregation = DisaggregationMetrics(
+                mode=mode_str,
+                prefill_bootstrap_queue_reqs=prefill_bootstrap,
+                prefill_inflight_queue_reqs=prefill_inflight,
+                decode_prealloc_queue_reqs=decode_prealloc,
+                decode_transfer_queue_reqs=decode_transfer,
+                decode_retracted_queue_reqs=decode_retracted,
+                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.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(
+            dp_rank=self.ps.dp_rank,
+            timestamp=time.time(),
+            num_running_reqs=num_running_reqs,
+            num_waiting_reqs=num_waiting_reqs,
+            num_used_tokens=num_used_tokens,
+            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.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,
+            lora=lora,
+            disaggregation=disaggregation,
+            queues=queues,
+        )

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -1050,8 +1050,7 @@ class SchedulerOutputProcessorMixin:
         spec_correct_drafts_histogram = []
         retraction_counts = []
         output_hidden_states = None
-        load = self.get_loads(
+        load = self.load_inquirer.get_loads(
-            self.load_inquirer,
             GetLoadsReqInput(include=["core"]),
         )
         routed_experts = None

python/sglang/srt/observability/scheduler_metrics_mixin.py

@@ -9,15 +9,6 @@ from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.environ import envs
-from sglang.srt.managers.io_struct import (
-    DisaggregationMetrics,
-    GetLoadsReqInput,
-    GetLoadsReqOutput,
-    LoRAMetrics,
-    MemoryMetrics,
-    QueueMetrics,
-    SpeculativeMetrics,
-)
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.observability.metrics_collector import (
@@ -34,9 +25,6 @@ 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__)
 
@@ -916,168 +904,6 @@ class SchedulerMetricsMixin:
                     self.stats.token_usage / 0.9,
                 )
 
-    @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
-        from the currently chunked request.
-
-        Args:
-            chunk_deduct: extra tokens to subtract from the chunked request's
-                remaining count. At batch-scheduling time the current chunk
-                has been planned but ``prefix_indices`` does not yet include it,
-                so callers pass ``extend_input_len`` here. At load-reporting
-                time ``prefix_indices`` is already up-to-date, so the default
-                0 is correct.
-        """
-        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
-
-    @staticmethod
-    def get_loads(
-        self: "SchedulerLoadInquirer", req: GetLoadsReqInput = None
-    ) -> GetLoadsReqOutput:
-        """
-        Get comprehensive load metrics for /v1/loads endpoint.
-
-        Args:
-            req: Request containing include list and optional dp_rank filter
-
-        Returns:
-            GetLoadsReqOutput with core metrics and optional detailed sections
-        """
-        if req is None:
-            req = GetLoadsReqInput()
-
-        include = set(req.include) if req.include else {"core"}
-        include_all = "all" in include
-
-        num_running_reqs = len(self.get_running_batch().reqs)
-
-        waiting_queues = [self.get_waiting_queue()]
-        if self.disaggregation_mode == DisaggregationMode.PREFILL:
-            waiting_queues.append(self.get_disagg_prefill_bootstrap_queue().queue)
-        elif self.disaggregation_mode == DisaggregationMode.DECODE:
-            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 = (
-            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
-        )
-
-        memory = None
-        if include_all or "memory" in include:
-            try:
-                memory = MemoryMetrics(
-                    weight_gb=round(
-                        self.tp_worker.model_runner.weight_load_mem_usage, 3
-                    ),
-                    kv_cache_gb=round(
-                        self.token_to_kv_pool_allocator.get_kvcache().mem_usage, 3
-                    ),
-                    graph_gb=round(self.tp_worker.model_runner.graph_mem_usage, 3),
-                    token_capacity=int(self.max_total_num_tokens),
-                )
-            except AttributeError as e:
-                logger.debug(f"Memory metrics not available: {e}")
-
-        speculative = None
-        if include_all or "spec" in include:
-            if (
-                not self.spec_algorithm.is_none()
-                and self.get_spec_total_num_forward_ct() > 0
-            ):
-                speculative = SpeculativeMetrics(
-                    accept_length=(
-                        self.get_spec_total_num_accept_tokens()
-                        / self.get_spec_total_num_forward_ct()
-                    ),
-                    accept_rate=self.get_stats().spec_accept_rate,
-                )
-
-        lora = None
-        if include_all or "lora" in include:
-            if self.server_args.enable_lora:
-                lora = LoRAMetrics(
-                    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
-        if include_all or "disagg" in include:
-            mode_str = "null"
-            prefill_bootstrap = 0
-            prefill_inflight = 0
-            decode_prealloc = 0
-            decode_transfer = 0
-            decode_retracted = 0
-
-            if self.disaggregation_mode == DisaggregationMode.PREFILL:
-                mode_str = "prefill"
-                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.get_disagg_decode_prealloc_queue().queue)
-                decode_transfer = len(self.get_disagg_decode_transfer_queue().queue)
-                decode_retracted = len(
-                    self.get_disagg_decode_prealloc_queue().retracted_queue
-                )
-
-            disaggregation = DisaggregationMetrics(
-                mode=mode_str,
-                prefill_bootstrap_queue_reqs=prefill_bootstrap,
-                prefill_inflight_queue_reqs=prefill_inflight,
-                decode_prealloc_queue_reqs=decode_prealloc,
-                decode_transfer_queue_reqs=decode_transfer,
-                decode_retracted_queue_reqs=decode_retracted,
-                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.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(
-            dp_rank=self.ps.dp_rank,
-            timestamp=time.time(),
-            num_running_reqs=num_running_reqs,
-            num_waiting_reqs=num_waiting_reqs,
-            num_used_tokens=num_used_tokens,
-            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.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,
-            lora=lora,
-            disaggregation=disaggregation,
-            queues=queues,
-        )
-
     def update_device_timer(self: Scheduler):
         if not ENABLE_METRICS_DEVICE_TIMER:
             return