Tiny add metrics for prefill delayer (#16603)

python/sglang/srt/managers/prefill_delayer.py

@@ -1,18 +1,36 @@
 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,
@@ -20,8 +38,10 @@ class PrefillDelayer:
         )
         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 @@ class PrefillDelayer:
         )
 
         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 @@ class Scheduler(
                 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 @@ class SchedulerMetricsCollector:
             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 @@ class SchedulerMetricsCollector:
     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,