Tiny add metrics for prefill delayer (sgl-project#16603)

Author: fzyzcjy

python/sglang/srt/managers/prefill_delayer.py

@@ -1,27 +1,47 @@
 import logging
-from typing import Optional
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
 from sglang.srt.environ import envs
 from sglang.srt.utils import get_bool_env_var
 
+if TYPE_CHECKING:
+    from sglang.srt.metrics.collector import SchedulerMetricsCollector
+
 _DEBUG_LOG = get_bool_env_var("SGLANG_PREFILL_DELAYER_DEBUG_LOG")
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class _DelayInfo:
+    delayed_count: int = 0
+    start_time: float = field(default_factory=time.perf_counter)
+
+
 class PrefillDelayer:
-    def __init__(self, dp_size, attn_tp_size, tp_worker, server_args):
+    def __init__(
+        self,
+        dp_size,
+        attn_tp_size,
+        tp_worker,
+        server_args,
+        metrics_collector: Optional["SchedulerMetricsCollector"] = None,
+    ):
         self.global_info = torch.empty(
             (dp_size, attn_tp_size, 1),
             dtype=torch.int64,
             device="cpu",
         )
         self.cpu_group = tp_worker.get_tp_group().cpu_group
 
-        self.curr_delayed_count = 0
         self.max_delay_passes = envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.get()
+        self._metrics_collector = metrics_collector
+
+        self._curr_delay_info: Optional[_DelayInfo] = None
 
         assert (
             server_args.enable_dp_attention
@@ -43,18 +63,31 @@ def _negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool:
         )
 
         if global_mixed_prefillable:
-            self.curr_delayed_count += 1
-            if self.curr_delayed_count < self.max_delay_passes:
+            if self._curr_delay_info is None:
+                self._curr_delay_info = _DelayInfo()
+            self._curr_delay_info.delayed_count += 1
+            if self._curr_delay_info.delayed_count < self.max_delay_passes:
                 return False
 
-        if _DEBUG_LOG and global_mixed_prefillable:
+        is_timeout = global_mixed_prefillable
+        if _DEBUG_LOG and is_timeout:
             logger.info(
                 f"PrefillDelayer timeout thus not forbid prefill (prefillable: {global_prefillable.sum()})"
             )
 
-        self.curr_delayed_count = 0
+        self._record_metrics(is_timeout=is_timeout)
+        self._curr_delay_info = None
         return True
 
+    def _record_metrics(self, is_timeout: bool) -> None:
+        if self._curr_delay_info is not None and self._metrics_collector is not None:
+            wait_seconds = time.perf_counter() - self._curr_delay_info.start_time
+            self._metrics_collector.observe_prefill_delayer_wait(
+                forward_passes=self._curr_delay_info.delayed_count,
+                wait_seconds=wait_seconds,
+                is_timeout=is_timeout,
+            )
+
     def _gather_info(self, local_prefillable: bool):
         local_info = torch.tensor(
             [int(local_prefillable)],

python/sglang/srt/managers/scheduler.py

@@ -806,6 +806,9 @@ def init_schedule_policy(self):
                 attn_tp_size=self.attn_tp_size,
                 tp_worker=self.tp_worker,
                 server_args=self.server_args,
+                metrics_collector=(
+                    self.metrics_collector if self.enable_metrics else None
+                ),
             )
         # Enable preemption for priority scheduling.
         self.try_preemption = self.enable_priority_scheduling

python/sglang/srt/metrics/collector.py

@@ -761,6 +761,25 @@ def __init__(
             labelnames=list(labels.keys()) + ["category", "num_prefill_ranks"],
         )
 
+        max_delay_passes = envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.get()
+        self.prefill_delayer_wait_forward_passes = Histogram(
+            name="sglang:prefill_delayer_wait_forward_passes",
+            documentation="Histogram of forward passes waited by prefill delayer.",
+            labelnames=labels.keys(),
+            buckets=[5, 20, max_delay_passes - 1],
+        )
+        self.prefill_delayer_wait_seconds = Histogram(
+            name="sglang:prefill_delayer_wait_seconds",
+            documentation="Histogram of wait time in seconds by prefill delayer.",
+            labelnames=labels.keys(),
+            buckets=[5, 20, 100, 500],
+        )
+        self.prefill_delayer_timeouts_total = Counter(
+            name="sglang:prefill_delayer_timeouts_total",
+            documentation="Total number of prefill delayer timeouts.",
+            labelnames=labels.keys(),
+        )
+
     def _log_gauge(self, gauge, data: Union[int, float]) -> None:
         # Convenience function for logging to gauge.
         gauge.labels(**self.labels).set(data)
@@ -781,6 +800,14 @@ def observe_per_stage_req_latency(self, stage: str, latency: float) -> None:
     def observe_queue_time(self, latency: float) -> None:
         self._log_histogram(self.queue_time, latency)
 
+    def observe_prefill_delayer_wait(
+        self, forward_passes: int, wait_seconds: float, is_timeout: bool
+    ) -> None:
+        self._log_histogram(self.prefill_delayer_wait_forward_passes, forward_passes)
+        self._log_histogram(self.prefill_delayer_wait_seconds, wait_seconds)
+        if is_timeout:
+            self.prefill_delayer_timeouts_total.labels(**self.labels).inc(1)
+
     def increment_retracted_reqs(
         self,
         num_retracted_reqs: int,

test/srt/test_prefill_delayer.py

@@ -2,6 +2,8 @@
 import unittest
 from types import SimpleNamespace
 
+import requests
+
 from sglang.bench_serving import run_benchmark
 from sglang.srt.environ import envs
 from sglang.srt.utils import kill_process_tree
@@ -88,6 +90,7 @@ def _run_throughput_test(
             **other_benchmark_args,
         )
         res = run_benchmark(args)
+        _print_prefill_delayer_metrics(base_url, expect_metrics=prefill_delayer)
     finally:
         kill_process_tree(process.pid)
 
@@ -137,6 +140,7 @@ def _run_accuracy_test(self, prefill_delayer: bool):
 
 def _launch_server(*, model, base_url, prefill_delayer: bool, other_args):
     os.environ["SGLANG_PREFILL_DELAYER_DEBUG_LOG"] = "1"
+    world_size = os.environ.get("SGLANG_TEST_WORLD_SIZE", "8")
 
     with envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.override(
         prefill_delayer
@@ -148,18 +152,35 @@ def _launch_server(*, model, base_url, prefill_delayer: bool, other_args):
             other_args=[
                 "--trust-remote-code",
                 "--tp",
-                "8",
+                world_size,
                 "--enable-dp-attention",
                 "--dp",
-                "8",
+                world_size,
                 "--chunked-prefill-size",
                 "131072",
                 "--mem-fraction-static",
                 "0.6",
+                "--enable-metrics",
                 *(other_args or []),
             ],
         )
 
 
+def _print_prefill_delayer_metrics(base_url: str, expect_metrics: bool):
+    metrics_response = requests.get(f"{base_url}/metrics")
+    assert metrics_response.status_code == 200
+    metrics_text = metrics_response.text
+    prefill_delayer_metrics = [
+        line for line in metrics_text.split("\n") if "prefill_delayer" in line
+    ]
+    print("=== PrefillDelayer Metrics ===")
+    for line in prefill_delayer_metrics:
+        print(line)
+    if expect_metrics:
+        assert "sglang:prefill_delayer_wait_forward_passes" in metrics_text
+        assert "sglang:prefill_delayer_wait_seconds" in metrics_text
+        assert "sglang:prefill_delayer_timeouts_total" in metrics_text
+
+
 if __name__ == "__main__":
     unittest.main()