diff --git a/tests/metrics/test_prometheus.py b/tests/metrics/test_prometheus.py index f01933dbc..e45365cf1 100644 --- a/tests/metrics/test_prometheus.py +++ b/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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 0bc6a93ec..27b564529 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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, diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index e1b078084..3fee36477 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index af710cba2..a00554209 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index d7011242b..af86b43d5 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/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]] = {}