diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index daf499f8a..c2c5b3036 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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( - 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 diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py index 82cc1c9ce..32acaa44e 100644 --- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py +++ b/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, + ) diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 683bee419..e7ad362c7 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/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( - self.load_inquirer, + load = self.load_inquirer.get_loads( GetLoadsReqInput(include=["core"]), ) routed_experts = None diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 911a4ad03..8a1afc243 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/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