From 8b1d8c170360c24c48bdb5a001d891b4c3c51f21 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Fri, 4 Sep 2026 08:06:17 -0700 Subject: [PATCH] [Metrics] Add rolling scheduler utilization counters (#37461) Co-authored-by: Lianmin Zheng Co-authored-by: pranjalssh --- python/sglang/srt/disaggregation/decode.py | 2 + python/sglang/srt/disaggregation/prefill.py | 2 + .../hardware_backend/mlx/scheduler_mixin.py | 1 + python/sglang/srt/managers/scheduler.py | 21 ++++- .../scheduler_components/metrics_reporter.py | 83 ++++++++++++++++ .../srt/observability/metrics_collector.py | 22 +++++ test/registered/observability/test_metrics.py | 3 + .../test_generation_auxiliary_output.py | 1 + .../managers/test_scheduler_on_idle_load.py | 1 + .../test_scheduler_pause_generation.py | 18 ++++ .../test_forward_pass_metrics.py | 94 +++++++++++++++++++ 11 files changed, 247 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 9e0ea879c..88e592721 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -2482,6 +2482,7 @@ class SchedulerDisaggregationDecodeMixin: recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if self._engine_paused: + self._record_scheduler_state_for_paused_engine() continue self.process_decode_queue() @@ -2525,6 +2526,7 @@ class SchedulerDisaggregationDecodeMixin: recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if self._engine_paused: + self._record_scheduler_state_for_paused_engine() continue self.process_decode_queue() diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 4c8072b4e..615874460 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -599,6 +599,7 @@ class SchedulerDisaggregationPrefillMixin: recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if self._engine_paused: + self._record_scheduler_state_for_paused_engine() continue self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() @@ -638,6 +639,7 @@ class SchedulerDisaggregationPrefillMixin: recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if self._engine_paused: + self._record_scheduler_state_for_paused_engine() continue self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index ee2877ee2..850b2da9f 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py @@ -213,6 +213,7 @@ class SchedulerMlxOverlapMixin: recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if self._engine_paused: + self._record_scheduler_state_for_paused_engine() continue # 1. If pending_curr is a pure decode AND no new prefill is waiting, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 7a5342d28..cd2d5d728 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1819,6 +1819,7 @@ class Scheduler( if use_mlx(): # MLX overlap uses mx.async_eval for CPU/GPU overlap, # not PyTorch MPS streams. + self.metrics_reporter.start_scheduler_time_accounting() dispatch_event_loop(self) return @@ -1841,6 +1842,7 @@ class Scheduler( # on the previous forward's read of the unified memory pool. self._war_barrier_enabled = is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get() with self.device_module.StreamContext(self.schedule_stream): + self.metrics_reporter.start_scheduler_time_accounting() dispatch_event_loop(self) def _apply_war_barrier(self): @@ -1865,8 +1867,11 @@ class Scheduler( # Receive requests recv_reqs = self.request_receiver.recv_requests() + if recv_reqs: + self.metrics_reporter.record_scheduler_active() self.process_input_requests(recv_reqs) if self._engine_paused: + self._record_scheduler_state_for_paused_engine() continue # Get the next batch to run @@ -1909,8 +1914,11 @@ class Scheduler( # Receive requests recv_reqs = self.request_receiver.recv_requests() + if recv_reqs: + self.metrics_reporter.record_scheduler_active() self.process_input_requests(recv_reqs) if self._engine_paused: + self._record_scheduler_state_for_paused_engine() continue # Get the next batch to run @@ -4121,6 +4129,7 @@ class Scheduler( pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[GenerationBatchResult, EmbeddingBatchResult]: """Run a batch.""" + self.metrics_reporter.record_scheduler_active() self.forward_ct += 1 batch.forward_iter = self.forward_ct batch.launch_ts = time.monotonic() @@ -4597,7 +4606,9 @@ class Scheduler( # path spins without sleeping, so a wall-clock floor bounds the # O(queue) get_loads for both sinks; the fully-idle publish runs # post-flush below. - if not self.is_fully_idle(): + fully_idle = self.is_fully_idle() + if not fully_idle: + self.metrics_reporter.record_scheduler_active() now = time.monotonic() if now - self._last_stall_publish_ts >= LOAD_STALL_REFRESH_S: self._last_stall_publish_ts = now @@ -4606,6 +4617,7 @@ class Scheduler( self.load_inquirer.get_loads, force=True, snapshot=snapshot ) return + self.metrics_reporter.record_scheduler_idle() if self.enable_unified_memory: try: @@ -4658,6 +4670,13 @@ class Scheduler( # sleep until next event self.maybe_sleep_on_idle() + self.metrics_reporter.record_scheduler_idle() + + def _record_scheduler_state_for_paused_engine(self) -> None: + if self.is_fully_idle(): + self.metrics_reporter.record_scheduler_idle() + else: + self.metrics_reporter.record_scheduler_active() def is_fully_idle(self, for_health_check=False) -> bool: # Health check piggybacks on running requests in process_output. diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 9a7017256..153e1fe13 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -53,6 +53,8 @@ CACHE_HIT_RATE_WINDOW_SECONDS = envs.SGLANG_CACHE_HIT_RATE_WINDOW_SECONDS.get() # be re-exported indefinitely. 30s is far above any healthy decode-stats gap, # so past it the true recent decode throughput is ~0. GEN_THROUGHPUT_STALENESS_SECONDS = 30.0 +# Update scheduler time counters every 1 second. +_SCHEDULER_TIME_ACCOUNTING_INTERVAL_NS = 1_000_000_000 class _CacheHitRateWindow: @@ -124,6 +126,44 @@ class PrefillStats: ) +@dataclass(slots=True) +class _SchedulerTimeAccountingSnapshot: + start_ns: int + last_sample_ns: int + start_process_cpu_ns: int + accumulate_idle_ns: int + # Whether the last sample was idle (for example, no work on the engine). + is_idle: bool + + @classmethod + def init( + cls, now_ns: int, now_process_cpu_ns: int, is_idle: bool + ) -> _SchedulerTimeAccountingSnapshot: + return cls( + start_ns=now_ns, + last_sample_ns=now_ns, + start_process_cpu_ns=now_process_cpu_ns, + accumulate_idle_ns=0, + is_idle=is_idle, + ) + + def sample(self, now_ns: int, is_idle: bool) -> None: + if self.is_idle: + self.accumulate_idle_ns += now_ns - self.last_sample_ns + self.last_sample_ns = now_ns + self.is_idle = is_idle + + def should_record(self, now_ns: int) -> bool: + return now_ns - self.start_ns >= _SCHEDULER_TIME_ACCOUNTING_INTERVAL_NS + + def reset(self, now_ns: int, now_process_cpu_ns: int, is_idle: bool) -> None: + self.start_ns = now_ns + self.last_sample_ns = now_ns + self.start_process_cpu_ns = now_process_cpu_ns + self.accumulate_idle_ns = 0 + self.is_idle = is_idle + + @dataclass(kw_only=True) class SchedulerMetricsReporter: scheduler: Scheduler @@ -216,6 +256,9 @@ class SchedulerMetricsReporter: self._mfu_log_write_bytes = 0.0 self.fwd_occupancy = float("nan") + self._scheduler_time_accounting: Optional[_SchedulerTimeAccountingSnapshot] = ( + None + ) self.forward_pass_device_timer: Optional[DeviceTimer] = None @@ -1196,6 +1239,46 @@ class SchedulerMetricsReporter: if self._device_timer_window_batch_count >= self.decode_log_interval: self._device_timer_window_batch_count = 0 + def start_scheduler_time_accounting(self) -> None: + if not self.enable_metrics: + return + self._scheduler_time_accounting = _SchedulerTimeAccountingSnapshot.init( + time.monotonic_ns(), time.process_time_ns(), True + ) + + def record_scheduler_active(self) -> None: + self._record_scheduler_time(is_idle=False) + + def record_scheduler_idle(self) -> None: + self._record_scheduler_time(is_idle=True) + + def _record_scheduler_time(self, is_idle: bool) -> None: + if not self.enable_metrics: + return + + now_wall_ns = time.monotonic_ns() + accounting = self._scheduler_time_accounting + if accounting is None: + self._scheduler_time_accounting = _SchedulerTimeAccountingSnapshot.init( + now_wall_ns, time.process_time_ns(), is_idle + ) + return + + accounting.sample(now_wall_ns, is_idle) + if not accounting.should_record(now_wall_ns): + return + + now_process_cpu_ns = time.process_time_ns() + elapsed_process_cpu_ns = now_process_cpu_ns - accounting.start_process_cpu_ns + if accounting.accumulate_idle_ns > 0: + self.metrics_collector.increment_scheduler_idle_seconds( + accounting.accumulate_idle_ns / 1e9 + ) + self.metrics_collector.increment_scheduler_process_cpu_seconds( + elapsed_process_cpu_ns / 1e9 + ) + accounting.reset(now_wall_ns, now_process_cpu_ns, is_idle) + def _reset_device_timer_window(self): """Exclude idle time and invalidate the last forward-occupancy sample.""" if ENABLE_METRICS_DEVICE_TIMER: diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 0a9a6f08d..79ab4232c 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -916,6 +916,22 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): ), labelnames=list(labels.keys()) + ["category"], ) + self.scheduler_idle_seconds_total = Counter( + name="sglang:scheduler_idle_seconds_total", + documentation=( + "Total wall time while the scheduler has no runnable or queued work." + ), + labelnames=labels.keys(), + ) + self.scheduler_process_cpu_seconds_total = Counter( + name="sglang:scheduler_process_cpu_seconds_total", + documentation=( + "Total CPU time consumed by the scheduler process across all threads." + ), + labelnames=labels.keys(), + ) + self.scheduler_idle_seconds_total.labels(**labels) + self.scheduler_process_cpu_seconds_total.labels(**labels) self.estimated_flops_per_gpu_total = Counter( name="sglang:estimated_flops_per_gpu_total", documentation=( @@ -1302,6 +1318,12 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): **dp_cooperation_info.to_labels(), ).inc(t) + def increment_scheduler_idle_seconds(self, t: float) -> None: + self.scheduler_idle_seconds_total.labels(**self.labels).inc(t) + + def increment_scheduler_process_cpu_seconds(self, t: float) -> None: + self.scheduler_process_cpu_seconds_total.labels(**self.labels).inc(t) + def increment_estimated_perf( self, num_flops_per_gpu: float = 0.0, diff --git a/test/registered/observability/test_metrics.py b/test/registered/observability/test_metrics.py index 7a079acb0..0c6c66704 100644 --- a/test/registered/observability/test_metrics.py +++ b/test/registered/observability/test_metrics.py @@ -193,6 +193,8 @@ class TestEnableMetrics(CustomTestCase): "sglang:startup_available_gpu_memory_gb", "sglang:startup_time_seconds", "sglang:startup_cuda_graph_time_seconds", + "sglang:scheduler_idle_seconds_total", + "sglang:scheduler_process_cpu_seconds_total", ] mfu_metrics = [ "sglang:estimated_flops_per_gpu_total", @@ -229,6 +231,7 @@ class TestEnableMetrics(CustomTestCase): ("sglang:realtime_tokens_total", {"mode": "decode"}), ("sglang:forward_execution_seconds_total", {"category": "extend"}), ("sglang:forward_execution_seconds_total", {"category": "decode"}), + ("sglang:scheduler_process_cpu_seconds_total", {}), ("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/managers/test_generation_auxiliary_output.py b/test/registered/unit/managers/test_generation_auxiliary_output.py index ae371fd12..7e9998f0c 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.metrics_reporter = Mock() scheduler.forward_ct = 0 scheduler._sched_idled = False scheduler.scripted_scheduler_hook = 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 205ca24da..61e2b49b2 100644 --- a/test/registered/unit/managers/test_scheduler_on_idle_load.py +++ b/test/registered/unit/managers/test_scheduler_on_idle_load.py @@ -27,6 +27,7 @@ class TestOnIdleStallPublish(CustomTestCase): s.publish_load_snapshot = MagicMock(return_value=None) s.load_publisher = MagicMock() s.load_inquirer = MagicMock() + s.metrics_reporter = MagicMock() s._last_stall_publish_ts = float("-inf") return s diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index 31e8650eb..3550eba96 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -143,6 +143,24 @@ class TestSchedulerPauseGeneration(unittest.TestCase): self.assertIs(scheduler.cur_batch_for_debug, original_cur_batch) self.assertIs(scheduler.chunked_req, original_chunked_req) + def test_paused_engine_accounting_uses_current_scheduler_state(self): + scheduler = self._new_scheduler() + scheduler.is_fully_idle = MagicMock() + + for is_idle in (True, False): + with self.subTest(is_idle=is_idle): + scheduler.is_fully_idle.return_value = is_idle + scheduler.metrics_reporter.reset_mock() + + scheduler._record_scheduler_state_for_paused_engine() + + if is_idle: + scheduler.metrics_reporter.record_scheduler_idle.assert_called_once_with() + scheduler.metrics_reporter.record_scheduler_active.assert_not_called() + else: + scheduler.metrics_reporter.record_scheduler_active.assert_called_once_with() + scheduler.metrics_reporter.record_scheduler_idle.assert_not_called() + def test_inplace_does_not_drain_overlap_queue(self): """in_place should not process the overlap result_queue.""" scheduler = self._new_scheduler() diff --git a/test/registered/unit/observability/test_forward_pass_metrics.py b/test/registered/unit/observability/test_forward_pass_metrics.py index 57dd4018e..0d287f147 100644 --- a/test/registered/unit/observability/test_forward_pass_metrics.py +++ b/test/registered/unit/observability/test_forward_pass_metrics.py @@ -407,6 +407,100 @@ class TestIdleMetrics(unittest.TestCase): self.assertEqual(self.published_occupancies, []) +class TestSchedulerTimeAccounting(CustomTestCase): + def setUp(self): + self.reporter = _make_reporter(self, types.SimpleNamespace()) + self.idle_seconds = [] + self.process_cpu_seconds = [] + self.reporter.enable_metrics = 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, + ) + + def test_counts_idle_wall_time_and_process_cpu_time(self): + wall_timestamps = [ + 0, + 1_200_000_000, + 1_500_000_000, + 2_700_000_000, + 3_000_000_000, + 4_100_000_000, + ] + process_cpu_timestamps = [ + 0, + 400_000_000, + 1_200_000_000, + 1_900_000_000, + ] + with ( + patch( + "sglang.srt.managers.scheduler_components.metrics_reporter.time.monotonic_ns", + side_effect=wall_timestamps, + ), + patch( + "sglang.srt.managers.scheduler_components.metrics_reporter.time.process_time_ns", + side_effect=process_cpu_timestamps, + ), + ): + self.reporter.start_scheduler_time_accounting() + self.reporter.record_scheduler_idle() + self.reporter.record_scheduler_active() + self.reporter.record_scheduler_active() + self.reporter.record_scheduler_idle() + self.reporter.record_scheduler_idle() + + self.assertAlmostEqual(sum(self.idle_seconds), 2.6) + self.assertAlmostEqual(sum(self.process_cpu_seconds), 1.9) + + def test_state_transitions_accumulate_until_periodic_update(self): + with ( + patch( + "sglang.srt.managers.scheduler_components.metrics_reporter.time.monotonic_ns", + side_effect=[0, 200_000_000, 400_000_000, 700_000_000, 1_100_000_000], + ), + patch( + "sglang.srt.managers.scheduler_components.metrics_reporter.time.process_time_ns", + side_effect=[0, 300_000_000], + ) as process_time, + ): + self.reporter.start_scheduler_time_accounting() + accounting = self.reporter._scheduler_time_accounting + self.reporter.record_scheduler_active() + self.reporter.record_scheduler_idle() + self.reporter.record_scheduler_active() + self.assertEqual(self.idle_seconds, []) + self.assertEqual(self.process_cpu_seconds, []) + self.assertEqual( + self.reporter._scheduler_time_accounting.accumulate_idle_ns, + 500_000_000, + ) + self.reporter.record_scheduler_active() + + self.assertIs(self.reporter._scheduler_time_accounting, accounting) + self.assertEqual(process_time.call_count, 2) + self.assertEqual(self.idle_seconds, [0.5]) + self.assertAlmostEqual(self.process_cpu_seconds[0], 0.3) + + def test_periodic_update_skips_zero_idle_but_records_cpu_sample(self): + with ( + patch( + "sglang.srt.managers.scheduler_components.metrics_reporter.time.monotonic_ns", + side_effect=[0, 0, 1_000_000_000], + ), + patch( + "sglang.srt.managers.scheduler_components.metrics_reporter.time.process_time_ns", + side_effect=[0, 0], + ), + ): + self.reporter.start_scheduler_time_accounting() + self.reporter.record_scheduler_active() + self.reporter.record_scheduler_active() + + self.assertEqual(self.idle_seconds, []) + self.assertEqual(self.process_cpu_seconds, [0.0]) + + class TestEstimatedPrefillPerf(CustomTestCase): """Causal pair count behind ``est. prefill TFLOPS/s`` and ``estimated_flops``."""