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

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

Browse Source
This commit is contained in:
fzyzcjy
2026-05-18 18:40:46 +08:00
committed by GitHub
parent 1213277879
commit bde932cbbb
4 changed files with 125 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+31 -2
View File
@@ -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
+43
View File
@@ -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
+4 -1
View File
@@ -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
+47 -33
View File
@@ -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,