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

Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com>

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}),

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(

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):

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

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

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)

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())

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 (

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()