Fix num_running_reqs gauge on disagg prefill servers (#31495)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
2026-07-16 15:43:06 -07:00
parent 640101da13
commit d28e35b1a1
1 changed files with 11 additions and 1 deletions
@@ -635,7 +635,17 @@ class SchedulerMetricsReporter:
            )

            # Basics
-            self.stats.num_running_reqs = prefill_stats.num_running_reqs
+            if (
+                self.scheduler.disaggregation_mode == DisaggregationMode.PREFILL
+                and batch is not None
+            ):
+                # running_batch is never populated in PD prefill, so count the
+                # current forward batch instead.
+                self.stats.num_running_reqs = QueueCount.from_reqs(
+                    batch.reqs, priority_enabled
+                )
+            else:
+                self.stats.num_running_reqs = prefill_stats.num_running_reqs
            self.stats.num_queue_reqs = QueueCount.from_reqs(
                self.scheduler.waiting_queue, priority_enabled
            )