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

Move queue-load reporting to SchedulerLoadInquirer (#25628)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-18 18:41:12 +08:00
committed by GitHub
parent bde932cbbb
commit 8357d07569
4 changed files with 169 additions and 179 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+2 -3
View File
@@ -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(
self.load_inquirer,
num_pending_tokens=self.load_inquirer._get_num_pending_tokens(
chunk_deduct=(
self.chunked_req.extend_input_len
if self.chunked_req is not None
python/sglang/srt/managers/scheduler_components/load_inquirer.py
+166
View File
@@ -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
+1 -2
View File
@@ -1050,8 +1050,7 @@ class SchedulerOutputProcessorMixin:
spec_correct_drafts_histogram = []
retraction_counts = []
output_hidden_states = None
load = self.get_loads(
self.load_inquirer,
load = self.load_inquirer.get_loads(
GetLoadsReqInput(include=["core"]),
)
routed_experts = None
python/sglang/srt/observability/scheduler_metrics_mixin.py
-174
View File
@@ -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