diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 88e592721..03cd52a33 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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() diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 615874460..22c265640 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -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]: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index cd2d5d728..33c08792d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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() diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index d49db49aa..973f47849 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 153e1fe13..047631862 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -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): diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index 0ff41adab..74f039949 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -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]]: diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 79ab4232c..b8129a4a8 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -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, diff --git a/python/sglang/srt/observability/scheduler_stage_metrics.py b/python/sglang/srt/observability/scheduler_stage_metrics.py new file mode 100644 index 000000000..0d3bf13bb --- /dev/null +++ b/python/sglang/srt/observability/scheduler_stage_metrics.py @@ -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.`` 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 diff --git a/test/registered/observability/test_metrics.py b/test/registered/observability/test_metrics.py index 0c6c66704..a0ed860cb 100644 --- a/test/registered/observability/test_metrics.py +++ b/test/registered/observability/test_metrics.py @@ -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}), diff --git a/test/registered/unit/layers/test_minicpm_sparse_cache.py b/test/registered/unit/layers/test_minicpm_sparse_cache.py index c34c39d46..675c7ae00 100644 --- a/test/registered/unit/layers/test_minicpm_sparse_cache.py +++ b/test/registered/unit/layers/test_minicpm_sparse_cache.py @@ -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( diff --git a/test/registered/unit/managers/scheduler_components/test_invariant_checker.py b/test/registered/unit/managers/scheduler_components/test_invariant_checker.py index d4ea9dff3..33a0f9a8e 100644 --- a/test/registered/unit/managers/scheduler_components/test_invariant_checker.py +++ b/test/registered/unit/managers/scheduler_components/test_invariant_checker.py @@ -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): diff --git a/test/registered/unit/managers/test_generation_auxiliary_output.py b/test/registered/unit/managers/test_generation_auxiliary_output.py index 7e9998f0c..f7462155e 100644 --- a/test/registered/unit/managers/test_generation_auxiliary_output.py +++ b/test/registered/unit/managers/test_generation_auxiliary_output.py @@ -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 diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index d60870210..02f77d69a 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -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 diff --git a/test/registered/unit/managers/test_scheduler_on_idle_load.py b/test/registered/unit/managers/test_scheduler_on_idle_load.py index 61e2b49b2..770d6f990 100644 --- a/test/registered/unit/managers/test_scheduler_on_idle_load.py +++ b/test/registered/unit/managers/test_scheduler_on_idle_load.py @@ -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) diff --git a/test/registered/unit/multimodal/test_gpu_feature_transport.py b/test/registered/unit/multimodal/test_gpu_feature_transport.py index bef625a72..590f8b2c4 100644 --- a/test/registered/unit/multimodal/test_gpu_feature_transport.py +++ b/test/registered/unit/multimodal/test_gpu_feature_transport.py @@ -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()) diff --git a/test/registered/unit/observability/test_forward_pass_metrics.py b/test/registered/unit/observability/test_forward_pass_metrics.py index 0d287f147..4fcf909f7 100644 --- a/test/registered/unit/observability/test_forward_pass_metrics.py +++ b/test/registered/unit/observability/test_forward_pass_metrics.py @@ -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 ( diff --git a/test/registered/unit/observability/test_scheduler_stage_metrics.py b/test/registered/unit/observability/test_scheduler_stage_metrics.py new file mode 100644 index 000000000..611920c9e --- /dev/null +++ b/test/registered/unit/observability/test_scheduler_stage_metrics.py @@ -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()