[Bugfix] Report engine-queued requests as waiting in vllm_omni:num_requests gauges (#6549)

Signed-off-by: Bob Tian <bobzetian@google.com>

tests/metrics/test_prometheus.py

@@ -198,3 +198,32 @@ class TestRequestLifecycleGauges:
         out = generate_latest(registry).decode()
         assert _sample_value(out, 'vllm_omni:num_requests_running{model_name="lifecycle-test-2"}') == 1.0
         assert _sample_value(out, 'vllm_omni:num_requests_waiting{model_name="lifecycle-test-2"}') == 0.0
+
+    def test_engine_queued_requests_count_as_waiting(self, registry: CollectorRegistry) -> None:
+        from types import SimpleNamespace
+
+        from vllm_omni.entrypoints.omni_base import OmniBase
+        from vllm_omni.metrics.prometheus import OmniPrometheusMetrics, OmniRequestCounter
+
+        obj = object.__new__(OmniBase)
+        obj.engine = SimpleNamespace(
+            _running_counter=OmniRequestCounter(),
+            _engines_waiting_counter=OmniRequestCounter(),
+        )
+        obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-3")
+        obj.request_states = {}
+        obj.log_stats = False
+
+        # Three requests dispatched to a stage engine whose scheduler admits
+        # one and queues two (e.g. a diffusion engine with
+        # max_num_running_reqs=1). The orchestrator reports the queued count
+        # via _engines_waiting_counter; they must show as waiting, not running.
+        for _ in range(3):
+            obj.engine._running_counter.increment()
+        obj.engine._engines_waiting_counter.value = 2
+
+        obj._publish_request_gauges(3)
+
+        out = generate_latest(registry).decode()
+        assert _sample_value(out, 'vllm_omni:num_requests_running{model_name="lifecycle-test-3"}') == 1.0
+        assert _sample_value(out, 'vllm_omni:num_requests_waiting{model_name="lifecycle-test-3"}') == 2.0

vllm_omni/engine/async_omni_engine.py

@@ -253,6 +253,7 @@ class AsyncOmniEngine:
         self._correlated_rpc_client: CorrelatedRpcClient | None = None
         self._duplex_control_client: DuplexControlClient | None = None
         self._running_counter = OmniRequestCounter()
+        self._engines_waiting_counter = OmniRequestCounter()
 
         logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages")
 
@@ -417,6 +418,7 @@ class AsyncOmniEngine:
                 pd_config=pd_config,
                 membership_controller=membership_controller,
                 running_counter=self._running_counter,
+                engines_waiting_counter=self._engines_waiting_counter,
                 transfer_emitter=self._transfer_emitter,
                 prom_metrics=self._prom_metrics,
                 log_stats=self._log_stats,

vllm_omni/engine/orchestrator.py

@@ -387,6 +387,7 @@ class Orchestrator:
     # Class-level defaults so tests that bypass __init__ via object.__new__
     # don't AttributeError when transfer / counter emit paths access them.
     _running_counter: OmniRequestCounter | None = None
+    _engines_waiting_counter: OmniRequestCounter | None = None
     _transfer_emitter: Any = None
     _prom_metrics: Any = None
     _stat_logger: OmniPrometheusStatLogger | None = None
@@ -403,6 +404,7 @@ class Orchestrator:
         pd_config: dict[str, Any] | None = None,
         membership_controller: MembershipController | None = None,
         running_counter: OmniRequestCounter | None = None,
+        engines_waiting_counter: OmniRequestCounter | None = None,
         transfer_emitter: Any = None,
         prom_metrics: Any = None,
         log_stats: bool = False,
@@ -439,7 +441,13 @@ class Orchestrator:
             self._pd_bootstrap_addr = pd_config.get("bootstrap_addr")
             self._pd_prefill_engine_id = pd_config.get("prefill_engine_id")
         self.request_states: dict[str, OrchestratorRequestState] = {}
-        self._init_metrics_state(stage_pools, running_counter, transfer_emitter, log_stats=log_stats)
+        self._init_metrics_state(
+            stage_pools,
+            running_counter,
+            transfer_emitter,
+            engines_waiting_counter=engines_waiting_counter,
+            log_stats=log_stats,
+        )
 
         self._cfg_tracker = CfgCompanionTracker()
         self._stage_input_processors: dict[int, Any] = {}
@@ -485,6 +493,7 @@ class Orchestrator:
         stage_pools: list[StagePool],
         running_counter: OmniRequestCounter | None,
         transfer_emitter: Any,
+        engines_waiting_counter: OmniRequestCounter | None = None,
         log_stats: bool = False,
     ) -> None:
         """Wire up all metric-related orchestrator state.
@@ -510,6 +519,7 @@ class Orchestrator:
         from iteration_stats independently of scheduler gauges.
         """
         self._running_counter = running_counter
+        self._engines_waiting_counter = engines_waiting_counter
         self._transfer_emitter = transfer_emitter
 
         # Flat engine_idx ↔ (stage, replica) maps. The reverse map is
@@ -1458,11 +1468,20 @@ class Orchestrator:
         """Update one replica snapshot and expose the stage-wide total."""
         self._stage_replica_waiting[(stage_id, replica_id)] = max(int(n_waiting), 0)
         self._set_stage_waiting_total(stage_id)
+        self._sync_engines_waiting_counter()
 
     def _remove_stage_replica_waiting(self, stage_id: int, replica_id: int) -> None:
         """Drop a dead replica's snapshot and refresh the stage total."""
         self._stage_replica_waiting.pop((stage_id, replica_id), None)
         self._set_stage_waiting_total(stage_id)
+        self._sync_engines_waiting_counter()
+
+    def _sync_engines_waiting_counter(self) -> None:
+        """Aggregate per-replica waiting into the counter shared with the
+        frontend so engine-queued requests show as waiting, not running."""
+        counter = self._engines_waiting_counter
+        if counter is not None:
+            counter.value = sum(self._stage_replica_waiting.values())
 
     def _set_stage_waiting_total(self, stage_id: int) -> None:
         if self._prom_metrics is None:

vllm_omni/entrypoints/async_omni.py

@@ -622,6 +622,8 @@ class AsyncOmni(EngineClient, OmniBase):
             submit_ts = time.time()
             req_state.metrics.stage_first_ts[0] = submit_ts
             req_start_ts[request_id] = submit_ts
+            # Refresh gauges on arrival.
+            self._publish_request_gauges(len(self.request_states))
 
             # Process results based on mode
             # Both sequential and async_chunk modes read the same message stream

vllm_omni/entrypoints/omni_base.py

@@ -406,6 +406,26 @@ class OmniBase(PDDisaggregationMixin):
         prom.request_failed()
         prom.inc_requests_failed(normalize_failure_reason(reason))
 
+    def _publish_request_gauges(self, total: int) -> None:
+        """Publish num_requests_running / num_requests_waiting for ``total`` in-flight requests.
+
+        ``_running_counter`` counts requests from the moment the orchestrator
+        dispatches them to a stage engine, which includes requests still queued
+        inside a stage scheduler. The orchestrator reports those via
+        ``_engines_waiting_counter``.
+        """
+        prom = getattr(self, "prom_metrics", None)
+        if prom is None:
+            return
+        engine = getattr(self, "engine", None)
+        counter = getattr(engine, "_running_counter", None)
+        waiting_counter = getattr(engine, "_engines_waiting_counter", None)
+        dispatched = counter.value if counter is not None else total
+        in_engine_waiting = waiting_counter.value if waiting_counter is not None else 0
+        in_engine_waiting = min(max(0, in_engine_waiting), dispatched)
+        prom.set_running(max(0, dispatched - in_engine_waiting))
+        prom.set_waiting(max(0, total - dispatched) + in_engine_waiting)
+
     def _log_summary_and_cleanup(self, request_id: str, reason: str = "stage_error") -> None:
         req_state = self.request_states.get(request_id)
         try:
@@ -431,13 +451,7 @@ class OmniBase(PDDisaggregationMixin):
             # Republish gauges so any stale value left by the per-stage
             # publish in _process_single_result (which runs while the request
             # is still in self.request_states) is corrected after the pop.
-            prom = getattr(self, "prom_metrics", None)
-            counter = getattr(getattr(self, "engine", None), "_running_counter", None)
-            if prom is not None:
-                total = len(self.request_states)
-                running = counter.value if counter is not None else total
-                prom.set_running(running)
-                prom.set_waiting(max(0, total - running))
+            self._publish_request_gauges(len(self.request_states))
 
     def _compute_final_stage_id(self, output_modalities: list[str] | None) -> int:
         return get_final_stage_id_for_e2e(
@@ -708,12 +722,8 @@ class OmniBase(PDDisaggregationMixin):
         # hasn't popped self.request_states yet — exclude the finalizing
         # request from `total` so waiting doesn't read 1 and stay stuck
         # there until the next request arrives.
-        counter = getattr(self.engine, "_running_counter", None)
         is_finalizing = finished and stage_id == final_stage_id_for_e2e
-        total = max(0, len(self.request_states) - (1 if is_finalizing else 0))
-        running = counter.value if counter is not None else total
-        self.prom_metrics.set_running(running)
-        self.prom_metrics.set_waiting(max(0, total - running))
+        self._publish_request_gauges(max(0, len(self.request_states) - (1 if is_finalizing else 0)))
 
         response_metrics: dict[str, Any] = {}
         stage_metrics: dict[str, dict[str, Any]] = {}