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
test/registered/observability/test_metrics.py
+2
View File
@@ -195,6 +195,7 @@ class TestEnableMetrics(CustomTestCase):
"sglang:startup_cuda_graph_time_seconds",
"sglang:scheduler_idle_seconds_total",
"sglang:scheduler_process_cpu_seconds_total",
"sglang:scheduler_stage_seconds_total",
]
mfu_metrics = [
"sglang:estimated_flops_per_gpu_total",
@@ -232,6 +233,7 @@ class TestEnableMetrics(CustomTestCase):
("sglang:forward_execution_seconds_total", {"category": "extend"}),
("sglang:forward_execution_seconds_total", {"category": "decode"}),
("sglang:scheduler_process_cpu_seconds_total", {}),
("sglang:scheduler_stage_seconds_total", {"category": "other"}),
("sglang:process_cpu_seconds_total", {"component": "tokenizer"}),
("sglang:weight_memory_usage_gb", {"model_name": _MODEL_NAME}),
("sglang:kv_cache_memory_usage_gb", {"model_name": _MODEL_NAME}),
test/registered/unit/layers/test_minicpm_sparse_cache.py
+2
View File
@@ -182,6 +182,7 @@ def test_reserved_slots_are_excluded_from_full_pool_invariant():
pool_stats_observer=SimpleNamespace(session_held_tokens=lambda: 0),
get_last_batch=lambda: None,
get_running_batch=lambda: None,
scheduler_stage_metrics=None,
)
leak, message = checker._check_full_pool(
@@ -272,6 +273,7 @@ def test_mamba_leak_diagnostic_does_not_report_reserved_slots():
),
get_last_batch=lambda: None,
get_running_batch=lambda: None,
scheduler_stage_metrics=None,
)
leak, message = checker._check_mamba_pool(
test/registered/unit/managers/scheduler_components/test_invariant_checker.py
+1
View File
@@ -42,6 +42,7 @@ class TestCheckTreeCacheGate(CustomTestCase):
pool_stats_observer=MagicMock(),
get_last_batch=lambda: None,
get_running_batch=lambda: None,
scheduler_stage_metrics=None,
)
def test_disabled_by_default(self):
test/registered/unit/managers/test_generation_auxiliary_output.py
+1
View File
@@ -388,6 +388,7 @@ def test_pdmux_split_prefill_schedules_auxiliary_output_copy():
)
copy_done = CopyDone()
scheduler = object.__new__(Scheduler)
scheduler.scheduler_stage_metrics = None
scheduler.metrics_reporter = Mock()
scheduler.forward_ct = 0
scheduler._sched_idled = False
test/registered/unit/managers/test_scheduler_chunked_req_gate.py
+1
View File
@@ -75,6 +75,7 @@ def _make_chunk_cache(req_to_token_pool) -> ChunkCache:
def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
s = Scheduler.__new__(Scheduler)
s.scheduler_stage_metrics = None
s._abort_on_waiting_timeout = MagicMock()
s._abort_on_running_timeout = MagicMock()
s.dllm_config = None
test/registered/unit/managers/test_scheduler_on_idle_load.py
+1
View File
@@ -22,6 +22,7 @@ register_cpu_ci(est_time=2, suite="base-a-test-cpu")
class TestOnIdleStallPublish(CustomTestCase):
def _stalled_scheduler(self) -> Scheduler:
s = Scheduler.__new__(Scheduler)
s.scheduler_stage_metrics = None
s.maybe_send_health_check_signal = MagicMock()
s.is_fully_idle = MagicMock(return_value=False) # stalled, not idle
s.publish_load_snapshot = MagicMock(return_value=None)
test/registered/unit/multimodal/test_gpu_feature_transport.py
+1
View File
@@ -767,6 +767,7 @@ class TestSchedulerMmTransportBoundary(unittest.TestCase):
@staticmethod
def _prepare_scheduler(scheduler):
scheduler.scheduler_stage_metrics = None
scheduler.session_controller = SimpleNamespace(maybe_reap=MagicMock())
scheduler._request_dispatcher = MagicMock(return_value=None)
scheduler.flush_wrapper = SimpleNamespace(check_pending=MagicMock())
test/registered/unit/observability/test_forward_pass_metrics.py
+9
View File
@@ -412,10 +412,15 @@ class TestSchedulerTimeAccounting(CustomTestCase):
self.reporter = _make_reporter(self, types.SimpleNamespace())
self.idle_seconds = []
self.process_cpu_seconds = []
self.stage_seconds = []
self.reporter.enable_metrics = True
self.reporter.scheduler_stage_metrics.enabled = True
self.reporter.metrics_collector = types.SimpleNamespace(
increment_scheduler_idle_seconds=self.idle_seconds.append,
increment_scheduler_process_cpu_seconds=self.process_cpu_seconds.append,
increment_scheduler_stage_seconds=lambda **kwargs: (
self.stage_seconds.append(kwargs)
),
)
def test_counts_idle_wall_time_and_process_cpu_time(self):
@@ -452,6 +457,10 @@ class TestSchedulerTimeAccounting(CustomTestCase):
self.assertAlmostEqual(sum(self.idle_seconds), 2.6)
self.assertAlmostEqual(sum(self.process_cpu_seconds), 1.9)
self.assertAlmostEqual(
sum(sample["seconds"] for sample in self.stage_seconds), 4.1
)
self.assertEqual({sample["stage"] for sample in self.stage_seconds}, {"other"})
def test_state_transitions_accumulate_until_periodic_update(self):
with (
test/registered/unit/observability/test_scheduler_stage_metrics.py
+164
View File
@@ -0,0 +1,164 @@
import unittest
from unittest.mock import patch
import torch
from sglang.srt.observability.scheduler_stage_metrics import (
SCHEDULER_STAGE_CATEGORIES,
SCHEDULER_STAGE_GET_NEXT_BATCH,
SCHEDULER_STAGE_PROCESS_QUEUE,
SCHEDULER_STAGE_PROCESS_REQUESTS,
SCHEDULER_STAGE_RUN_BATCH,
SchedulerStageMetricsRecorder,
scheduler_stage_method,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
class TestSchedulerStageMetricsRecorder(CustomTestCase):
def test_category_names(self):
self.assertEqual(
set(SCHEDULER_STAGE_CATEGORIES),
{
"other",
"recv_requests",
"process_input_requests",
"process_batch_result",
"process_queue",
"get_next_batch_to_run",
"run_batch",
"sanity_check_cache",
"idle",
},
)
def test_nested_stages_are_exclusive(self):
recorder = SchedulerStageMetricsRecorder(enabled=True)
recorder.start(wall_ns=0)
with patch(
"sglang.srt.observability.scheduler_stage_metrics.time.monotonic_ns",
side_effect=[10, 30, 50, 80],
):
outer = recorder.enter(SCHEDULER_STAGE_GET_NEXT_BATCH)
inner = recorder.enter(SCHEDULER_STAGE_PROCESS_QUEUE)
recorder.exit(inner)
recorder.exit(outer)
wall_ns = recorder.drain(wall_ns=100)
self.assertEqual(
wall_ns,
{
"other": 30,
"get_next_batch_to_run": 50,
"process_queue": 20,
},
)
self.assertEqual(sum(wall_ns.values()), 100)
def test_decorator_restores_stage_after_exception(self):
recorder = SchedulerStageMetricsRecorder(enabled=True)
recorder.start(wall_ns=0)
class SchedulerLike:
scheduler_stage_metrics = recorder
@scheduler_stage_method(SCHEDULER_STAGE_RUN_BATCH)
def fail(self):
raise RuntimeError("boom")
with (
patch(
"sglang.srt.observability.scheduler_stage_metrics.time.monotonic_ns",
side_effect=[10, 40],
),
self.assertRaisesRegex(RuntimeError, "boom"),
):
SchedulerLike().fail()
wall_ns = recorder.drain(wall_ns=50)
self.assertEqual(wall_ns, {"other": 20, "run_batch": 30})
def test_nested_same_stage_does_not_double_count(self):
recorder = SchedulerStageMetricsRecorder(enabled=True)
recorder.start(wall_ns=0)
with patch(
"sglang.srt.observability.scheduler_stage_metrics.time.monotonic_ns",
side_effect=[10, 40],
):
with recorder.record(SCHEDULER_STAGE_RUN_BATCH):
with recorder.record(SCHEDULER_STAGE_RUN_BATCH):
pass
wall_ns = recorder.drain(wall_ns=50)
self.assertEqual(wall_ns, {"other": 20, "run_batch": 30})
def test_trace_spans_do_not_require_python_stacks(self):
recorder = SchedulerStageMetricsRecorder(enabled=False)
class SchedulerLike:
scheduler_stage_metrics = recorder
@scheduler_stage_method(SCHEDULER_STAGE_RUN_BATCH)
def run(self):
with self.scheduler_stage_metrics.record(SCHEDULER_STAGE_PROCESS_QUEUE):
pass
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU],
with_stack=False,
acc_events=True,
) as profiler:
SchedulerLike().run()
stage_events = [
event for event in profiler.events() if event.key.startswith("scheduler.")
]
self.assertEqual(
[event.key for event in stage_events],
["scheduler.run_batch", "scheduler.process_queue"],
)
self.assertTrue(all(not event.stack for event in stage_events))
def test_decorator_preserves_existing_trace_names(self):
recorder = SchedulerStageMetricsRecorder(enabled=False)
class SchedulerLike:
scheduler_stage_metrics = recorder
@scheduler_stage_method(SCHEDULER_STAGE_PROCESS_REQUESTS)
def process_input_requests(self):
pass
@scheduler_stage_method(SCHEDULER_STAGE_GET_NEXT_BATCH)
def get_next_batch_to_run(self):
pass
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU],
with_stack=False,
acc_events=True,
) as profiler:
scheduler = SchedulerLike()
scheduler.process_input_requests()
scheduler.get_next_batch_to_run()
stage_events = [
event for event in profiler.events() if event.key.startswith("scheduler.")
]
self.assertEqual(
[event.key for event in stage_events],
[
"scheduler.process_input_requests",
"scheduler.get_next_batch_to_run",
],
)
if __name__ == "__main__":
unittest.main()