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

[Metrics] Export scheduler stage wall time (#37636)

Browse Source 
Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com>
This commit is contained in:
Jialin Ouyang
2026-09-04 10:45:53 -07:00
committed by GitHub
co-authored by Pranjal Shankhdhar
parent 07199fa220
commit 010dc955be
17 changed files with 447 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+7 -2
View File
@@ -96,6 +96,11 @@ from sglang.srt.observability.req_time_stats import (
set_schedule_time_batch,
set_time_batch,
)
from sglang.srt.observability.scheduler_stage_metrics import (
SCHEDULER_STAGE_GET_NEXT_BATCH,
SCHEDULER_STAGE_PROCESS_QUEUE,
scheduler_stage_method,
)
from sglang.srt.runtime_context import (
get_disagg,
get_memory,
@@ -103,7 +108,6 @@ from sglang.srt.runtime_context import (
)
from sglang.srt.utils import ceil_align, get_num_new_pages, is_npu
from sglang.srt.utils.network import NetworkAddress
from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
logger = logging.getLogger(__name__)
@@ -2581,7 +2585,7 @@ class SchedulerDisaggregationDecodeMixin:
return GenerationBatchResult()
@scheduler_nvtx_method("scheduler.get_next_batch_to_run")
@scheduler_stage_method(SCHEDULER_STAGE_GET_NEXT_BATCH)
def get_next_disagg_decode_batch_to_run(
self: Scheduler, running_batch: ScheduleBatch
) -> NextBatchPlan:
@@ -2689,6 +2693,7 @@ class SchedulerDisaggregationDecodeMixin:
return new_batch
@scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE)
def process_decode_queue(self: Scheduler):
if self.enable_decode_hicache:
self.tree_cache.check_hicache_events()
python/sglang/srt/disaggregation/prefill.py
+12 -2
View File
@@ -72,13 +72,18 @@ from sglang.srt.mem_cache.common import (
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
from sglang.srt.observability.req_time_stats import set_schedule_time_batch
from sglang.srt.observability.scheduler_stage_metrics import (
SCHEDULER_STAGE_GET_NEXT_BATCH,
SCHEDULER_STAGE_PROCESS_QUEUE,
SchedulerStageMetricsRecorder,
scheduler_stage_method,
)
from sglang.srt.runtime_context import (
get_disagg,
get_parallel,
get_schedule,
)
from sglang.srt.utils import is_npu
from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
if TYPE_CHECKING:
from torch.distributed import ProcessGroup
@@ -147,6 +152,7 @@ class PrefillBootstrapQueue:
gloo_group: ProcessGroup,
max_total_num_tokens: int,
scheduler: Scheduler,
scheduler_stage_metrics: SchedulerStageMetricsRecorder,
pp_rank: int,
pp_size: int,
transfer_backend: TransferBackend,
@@ -165,6 +171,7 @@ class PrefillBootstrapQueue:
self.queue: List[Req] = []
self.gloo_group = gloo_group
self.scheduler = scheduler
self.scheduler_stage_metrics = scheduler_stage_metrics
self.max_total_num_tokens = (
self.scheduler.tp_worker.model_runner.effective_max_total_num_tokens
)
@@ -407,6 +414,7 @@ class PrefillBootstrapQueue:
"""
req.sampling_params.max_new_tokens = 1
@scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE)
def pop_bootstrapped(
self,
return_failed_reqs: bool = False,
@@ -524,6 +532,7 @@ class SchedulerDisaggregationPrefillMixin:
if room is not None and room in kv_mgr.transfer_infos:
prefetch(room)
@scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE)
def resolve_waiting_queue_bootstrap(self: Scheduler) -> None:
"""Resolve bootstrap status for waiting prefill requests before admission.
@@ -565,7 +574,7 @@ class SchedulerDisaggregationPrefillMixin:
for req in self.waiting_queue
)
@scheduler_nvtx_method("scheduler.get_next_batch_to_run")
@scheduler_stage_method(SCHEDULER_STAGE_GET_NEXT_BATCH)
def get_next_disagg_prefill_batch_to_run(
self: Scheduler,
running_batch: ScheduleBatch,
@@ -875,6 +884,7 @@ class SchedulerDisaggregationPrefillMixin:
dp_cooperation_info=batch.dp_cooperation_info,
)
@scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE)
def process_disagg_prefill_inflight_queue(
self: Scheduler, rids_to_check: Optional[List[str]] = None
) -> List[Req]:
python/sglang/srt/managers/scheduler.py
+36 -20
View File
@@ -292,6 +292,15 @@ from sglang.srt.observability.req_time_stats import (
set_schedule_time_batch,
set_time_batch,
)
from sglang.srt.observability.scheduler_stage_metrics import (
SCHEDULER_STAGE_GET_NEXT_BATCH,
SCHEDULER_STAGE_IDLE,
SCHEDULER_STAGE_PROCESS_BATCH_RESULT,
SCHEDULER_STAGE_PROCESS_REQUESTS,
SCHEDULER_STAGE_RUN_BATCH,
SCHEDULER_STAGE_SANITY_CHECK_CACHE,
scheduler_stage_method,
)
from sglang.srt.observability.startup_time import build_scheduler_startup_time
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
from sglang.srt.parser.reasoning_parser import ReasoningParser
@@ -339,7 +348,6 @@ from sglang.srt.utils.hf_transformers_utils import (
)
from sglang.srt.utils.msgspec_utils import msgspec_to_builtins
from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node
from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.srt.utils.weight_versions import (
compute_weight_version_spans,
@@ -642,6 +650,7 @@ class Scheduler(
self.init_diffusion_llm()
self.init_metrics_reporter(tp_rank, pp_rank, dp_rank)
self.scheduler_stage_metrics = self.metrics_reporter.scheduler_stage_metrics
# Init schedule policy and new token estimation
self.init_schedule_policy()
@@ -1542,6 +1551,7 @@ class Scheduler(
gloo_group=self.attn_tp_cpu_group,
max_total_num_tokens=self.max_total_num_tokens,
scheduler=self,
scheduler_stage_metrics=self.scheduler_stage_metrics,
pp_rank=self.ps.pp_rank,
pp_size=self.ps.pp_size,
transfer_backend=self.transfer_backend,
@@ -2022,7 +2032,7 @@ class Scheduler(
for prev_batch, prev_result in self.result_queue:
self.batch_result_processor.advance_grammar_fsm(prev_result, prev_batch)
@scheduler_nvtx_method("scheduler.process_input_requests")
@scheduler_stage_method(SCHEDULER_STAGE_PROCESS_REQUESTS)
def process_input_requests(self, recv_reqs: List):
now = time.monotonic()
self.session_controller.maybe_reap(now)
@@ -2248,6 +2258,7 @@ class Scheduler(
stream_output=lambda *a, **kw: self.output_streamer.stream_output(*a, **kw),
get_last_batch=lambda: self.last_batch,
scripted_scheduler_hook=self.scripted_scheduler_hook,
scheduler_stage_metrics=self.scheduler_stage_metrics,
)
def init_dp_attn_adapter(self) -> None:
@@ -2306,6 +2317,7 @@ class Scheduler(
get_last_batch=lambda: self.last_batch,
get_running_batch=lambda: self.running_batch,
get_chunked_req=lambda: self.chunked_req,
scheduler_stage_metrics=self.scheduler_stage_metrics,
)
def init_rank_consensus_checker(self) -> None:
@@ -3424,7 +3436,7 @@ class Scheduler(
# todo hisparse, maybe other info to contain for the new batch
return batch
@scheduler_nvtx_method("scheduler.get_next_batch_to_run")
@scheduler_stage_method(SCHEDULER_STAGE_GET_NEXT_BATCH)
def get_next_batch_to_run(
self, running_batch: ScheduleBatch, last_batch: Optional[ScheduleBatch]
) -> NextBatchPlan:
@@ -4122,7 +4134,7 @@ class Scheduler(
else:
batch.sampling_info = sched_sampling_info
@scheduler_nvtx_method("scheduler.run_batch")
@scheduler_stage_method(SCHEDULER_STAGE_RUN_BATCH)
def run_batch(
self,
batch: ScheduleBatch,
@@ -4467,7 +4479,7 @@ class Scheduler(
if batch_result.logits_output is not None:
batch_result.logits_output.next_token_logits = None
@scheduler_nvtx_method("scheduler.process_batch_result")
@scheduler_stage_method(SCHEDULER_STAGE_PROCESS_BATCH_RESULT)
def process_batch_result(
self,
batch: ScheduleBatch,
@@ -4594,6 +4606,7 @@ class Scheduler(
if_success = False
return ClearHiCacheReqOutput(success=if_success)
@scheduler_stage_method(SCHEDULER_STAGE_IDLE)
def on_idle(self):
"""Idle housekeeping: guard, check, metrics, reset, sleep."""
# Flush any health-check signal deferred while the engine was busy.
@@ -4634,23 +4647,26 @@ class Scheduler(
self.disaggregation_mode == DisaggregationMode.DECODE
and self.disagg_decode_transfer_queue.has_pending_deferred_releases()
)
if not self.enable_hisparse and not deferred_pending:
has_leak, messages = self.invariant_checker._check_all_pools(
self.pool_stats_observer.get_pool_stats(),
)
if has_leak:
self.invariant_checker._report_leak("pool", "\n".join(messages))
self.invariant_checker._check_req_pool()
# Byte-conservation diagnostic (allocator-owned; static pools
# return [] — the token identity above can't see byte leaks).
byte_violations = self.token_to_kv_pool_allocator.verify_byte_accounting()
if byte_violations:
self.invariant_checker._report_leak(
"pool-bytes", "\n".join(byte_violations)
with self.scheduler_stage_metrics.record(SCHEDULER_STAGE_SANITY_CHECK_CACHE):
if not self.enable_hisparse and not deferred_pending:
has_leak, messages = self.invariant_checker._check_all_pools(
self.pool_stats_observer.get_pool_stats(),
)
if has_leak:
self.invariant_checker._report_leak("pool", "\n".join(messages))
self.invariant_checker._check_req_pool()
# Byte-conservation diagnostic (allocator-owned; static pools
# return [] — the token identity above can't see byte leaks).
byte_violations = (
self.token_to_kv_pool_allocator.verify_byte_accounting()
)
if byte_violations:
self.invariant_checker._report_leak(
"pool-bytes", "\n".join(byte_violations)
)
# tree cache sanity check
self.invariant_checker._check_tree_cache()
# tree cache sanity check
self.invariant_checker._check_tree_cache()
# metrics every 30s
self.metrics_reporter._maybe_log_idle_metrics()
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+7
View File
@@ -26,6 +26,11 @@ from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.multi_ended_allocator import (
UnifiedMambaSWATokenToKVPoolAllocator,
)
from sglang.srt.observability.scheduler_stage_metrics import (
SCHEDULER_STAGE_SANITY_CHECK_CACHE,
SchedulerStageMetricsRecorder,
scheduler_stage_method,
)
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.common import (
ceil_align,
@@ -58,6 +63,7 @@ class SchedulerInvariantChecker:
pool_stats_observer: SchedulerPoolStatsObserver
get_last_batch: Callable
get_running_batch: Callable
scheduler_stage_metrics: SchedulerStageMetricsRecorder
# The chunked-prefill request parked between chunks is in neither batch;
# its uncached tokens must still be counted.
get_chunked_req: Callable = field(default=lambda: None)
@@ -293,6 +299,7 @@ class SchedulerInvariantChecker:
return full_uncached, swa_uncached
@scheduler_stage_method(SCHEDULER_STAGE_SANITY_CHECK_CACHE)
def self_check_during_busy(self):
if self.get_last_batch() is None:
return
python/sglang/srt/managers/scheduler_components/metrics_reporter.py
+16 -1
View File
@@ -22,6 +22,9 @@ from sglang.srt.observability.metrics_collector import (
SchedulerStats,
compute_routing_key_stats,
)
from sglang.srt.observability.scheduler_stage_metrics import (
SchedulerStageMetricsRecorder,
)
from sglang.srt.runtime_context import (
exports_expert_balancedness_to_prometheus,
get_context,
@@ -259,6 +262,9 @@ class SchedulerMetricsReporter:
self._scheduler_time_accounting: Optional[_SchedulerTimeAccountingSnapshot] = (
None
)
self.scheduler_stage_metrics = SchedulerStageMetricsRecorder(
enabled=self.enable_metrics
)
self.forward_pass_device_timer: Optional[DeviceTimer] = None
@@ -1242,9 +1248,11 @@ class SchedulerMetricsReporter:
def start_scheduler_time_accounting(self) -> None:
if not self.enable_metrics:
return
now_wall_ns = time.monotonic_ns()
self._scheduler_time_accounting = _SchedulerTimeAccountingSnapshot.init(
time.monotonic_ns(), time.process_time_ns(), True
now_wall_ns, time.process_time_ns(), True
)
self.scheduler_stage_metrics.start(now_wall_ns)
def record_scheduler_active(self) -> None:
self._record_scheduler_time(is_idle=False)
@@ -1262,6 +1270,7 @@ class SchedulerMetricsReporter:
self._scheduler_time_accounting = _SchedulerTimeAccountingSnapshot.init(
now_wall_ns, time.process_time_ns(), is_idle
)
self.scheduler_stage_metrics.start(now_wall_ns)
return
accounting.sample(now_wall_ns, is_idle)
@@ -1277,6 +1286,12 @@ class SchedulerMetricsReporter:
self.metrics_collector.increment_scheduler_process_cpu_seconds(
elapsed_process_cpu_ns / 1e9
)
for stage, elapsed_wall_ns in self.scheduler_stage_metrics.drain(
now_wall_ns
).items():
self.metrics_collector.increment_scheduler_stage_seconds(
stage=stage, seconds=elapsed_wall_ns / 1e9
)
accounting.reset(now_wall_ns, now_process_cpu_ns, is_idle)
def _reset_device_timer_window(self):
python/sglang/srt/managers/scheduler_components/request_receiver.py
+7 -2
View File
@@ -32,12 +32,16 @@ from sglang.srt.managers.mm_utils import (
has_shm_features,
unwrap_shm_features,
)
from sglang.srt.observability.scheduler_stage_metrics import (
SCHEDULER_STAGE_RECV_REQUESTS,
SchedulerStageMetricsRecorder,
scheduler_stage_method,
)
from sglang.srt.runtime_context import get_disagg, get_parallel, is_ep_scale_joiner
from sglang.srt.utils import (
broadcast_pyobj,
point_to_point_pyobj,
)
from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
@@ -73,13 +77,14 @@ class SchedulerRequestReceiver:
stream_output: Callable[..., None]
get_last_batch: Callable[[], Any]
scripted_scheduler_hook: Optional[ScriptedSchedulerHook] = None
scheduler_stage_metrics: Optional[SchedulerStageMetricsRecorder] = None
def recv_limit_reached(self, num_recv_reqs: int) -> bool:
if self.max_recv_per_poll < 0:
return False
return num_recv_reqs >= self.max_recv_per_poll
@scheduler_nvtx_method("scheduler.recv_requests")
@scheduler_stage_method(SCHEDULER_STAGE_RECV_REQUESTS)
def recv_requests(
self,
) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]:
python/sglang/srt/observability/metrics_collector.py
+17
View File
@@ -25,6 +25,9 @@ from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Set, Union
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.observability.scheduler_stage_metrics import (
SCHEDULER_STAGE_CATEGORIES,
)
from sglang.srt.observability.utils import exponential_buckets, generate_buckets
from sglang.srt.runtime_context import (
exports_expert_balancedness_to_prometheus,
@@ -930,8 +933,17 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
),
labelnames=labels.keys(),
)
self.scheduler_stage_seconds_total = Counter(
name="sglang:scheduler_stage_seconds_total",
documentation=(
"Total scheduler-loop wall time exclusively attributed to each stage."
),
labelnames=list(labels.keys()) + ["category"],
)
self.scheduler_idle_seconds_total.labels(**labels)
self.scheduler_process_cpu_seconds_total.labels(**labels)
for category in SCHEDULER_STAGE_CATEGORIES:
self.scheduler_stage_seconds_total.labels(**labels, category=category)
self.estimated_flops_per_gpu_total = Counter(
name="sglang:estimated_flops_per_gpu_total",
documentation=(
@@ -1324,6 +1336,11 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
def increment_scheduler_process_cpu_seconds(self, t: float) -> None:
self.scheduler_process_cpu_seconds_total.labels(**self.labels).inc(t)
def increment_scheduler_stage_seconds(self, stage: str, seconds: float) -> None:
self.scheduler_stage_seconds_total.labels(**self.labels, category=stage).inc(
seconds
)
def increment_estimated_perf(
self,
num_flops_per_gpu: float = 0.0,
python/sglang/srt/observability/scheduler_stage_metrics.py
+163
View File
@@ -0,0 +1,163 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from __future__ import annotations
import time
from collections import defaultdict
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from dataclasses import dataclass, field
from functools import wraps
from typing import TypeVar, cast
from sglang.srt.utils.nvtx_utils import (
NVTX_SCHEDULER_ENABLED,
profile_range,
scheduler_nvtx_method,
)
SCHEDULER_STAGE_OTHER = "other"
SCHEDULER_STAGE_RECV_REQUESTS = "recv_requests"
SCHEDULER_STAGE_PROCESS_REQUESTS = "process_input_requests"
SCHEDULER_STAGE_GET_NEXT_BATCH = "get_next_batch_to_run"
SCHEDULER_STAGE_PROCESS_QUEUE = "process_queue"
SCHEDULER_STAGE_RUN_BATCH = "run_batch"
SCHEDULER_STAGE_PROCESS_BATCH_RESULT = "process_batch_result"
SCHEDULER_STAGE_SANITY_CHECK_CACHE = "sanity_check_cache"
SCHEDULER_STAGE_IDLE = "idle"
SCHEDULER_STAGE_CATEGORIES = (
SCHEDULER_STAGE_OTHER,
SCHEDULER_STAGE_RECV_REQUESTS,
SCHEDULER_STAGE_PROCESS_REQUESTS,
SCHEDULER_STAGE_GET_NEXT_BATCH,
SCHEDULER_STAGE_PROCESS_QUEUE,
SCHEDULER_STAGE_RUN_BATCH,
SCHEDULER_STAGE_PROCESS_BATCH_RESULT,
SCHEDULER_STAGE_SANITY_CHECK_CACHE,
SCHEDULER_STAGE_IDLE,
)
_SCHEDULER_STAGE_CATEGORY_SET = frozenset(SCHEDULER_STAGE_CATEGORIES)
@dataclass(slots=True)
class SchedulerStageMetricsRecorder:
"""Accumulate mutually exclusive scheduler wall time by stage.
Nested stages temporarily replace their parent, and uncategorized time is
assigned to ``other``. Summing all categories therefore recovers elapsed
scheduler-loop wall time. Active torch profilers receive matching
``scheduler.<stage>`` ranges without requiring Python stacks.
"""
enabled: bool
_current_stage: str = SCHEDULER_STAGE_OTHER
_trace_stage: str | None = None
_last_wall_ns: int | None = None
_wall_ns: dict[str, int] = field(default_factory=lambda: defaultdict(int))
def start(self, wall_ns: int) -> None:
if not self.enabled:
return
self._current_stage = SCHEDULER_STAGE_OTHER
self._last_wall_ns = wall_ns
self._wall_ns.clear()
def enter(self, stage: str) -> str | None:
if (
not self.enabled
or self._last_wall_ns is None
or self._current_stage == stage
):
return None
self._sample(time.monotonic_ns())
previous_stage = self._current_stage
self._current_stage = stage
return previous_stage
def exit(self, previous_stage: str | None) -> None:
if previous_stage is None:
return
self._sample(time.monotonic_ns())
self._current_stage = previous_stage
@contextmanager
def record(self, stage: str) -> Iterator[None]:
if stage not in _SCHEDULER_STAGE_CATEGORY_SET:
raise ValueError(f"Unknown scheduler stage: {stage}")
previous_stage = self.enter(stage)
previous_trace_stage = self._trace_stage
trace_stage_changed = previous_trace_stage != stage
if trace_stage_changed:
self._trace_stage = stage
try:
if trace_stage_changed:
with profile_range(
f"scheduler.{stage}", nvtx_enabled=NVTX_SCHEDULER_ENABLED
):
yield
else:
yield
finally:
if trace_stage_changed:
self._trace_stage = previous_trace_stage
self.exit(previous_stage)
def drain(self, wall_ns: int) -> dict[str, int]:
if not self.enabled or self._last_wall_ns is None:
return {}
self._sample(wall_ns)
wall_ns_by_stage = dict(self._wall_ns)
self._wall_ns.clear()
return wall_ns_by_stage
def _sample(self, wall_ns: int) -> None:
assert self._last_wall_ns is not None
self._wall_ns[self._current_stage] += wall_ns - self._last_wall_ns
self._last_wall_ns = wall_ns
_F = TypeVar("_F", bound=Callable)
def scheduler_stage_method(stage: str) -> Callable[[_F], _F]:
if stage not in _SCHEDULER_STAGE_CATEGORY_SET:
raise ValueError(f"Unknown scheduler stage: {stage}")
trace_name = f"scheduler.{stage}"
def decorator(func: _F) -> _F:
profiled_func = scheduler_nvtx_method(trace_name)(func)
@wraps(func)
def wrapper(self, *args, **kwargs):
recorder = self.scheduler_stage_metrics
if recorder is None:
return profiled_func(self, *args, **kwargs)
previous_stage = recorder.enter(stage)
previous_trace_stage = recorder._trace_stage
trace_stage_changed = previous_trace_stage != stage
if trace_stage_changed:
recorder._trace_stage = stage
try:
return profiled_func(self, *args, **kwargs)
finally:
if trace_stage_changed:
recorder._trace_stage = previous_trace_stage
recorder.exit(previous_stage)
return cast(_F, wrapper)
return decorator