From 8fa0394d29758ecb0bf07e97ef486b8b01158dd4 Mon Sep 17 00:00:00 2001 From: blonde <164845202+blondeCS@users.noreply.github.com> Date: Wed, 17 Jun 2026 21:31:44 +0800 Subject: [PATCH] [Bugfix] Wire log_stats to AsyncOmni and add missing token metrics for non-text requests (#4482) Signed-off-by: blonde <164845202+blondeCS@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 (1M context) --- vllm_omni/entrypoints/omni_base.py | 14 +++++++++++--- vllm_omni/entrypoints/openai/api_server.py | 1 + vllm_omni/metrics/definitions.py | 4 ++++ vllm_omni/metrics/prometheus.py | 20 ++++++++++++++++++++ 4 files changed, 36 insertions(+), 3 deletions(-) diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 029665f59..fc22493fc 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -530,6 +530,15 @@ class OmniBase(PDDisaggregationMixin): finished_reason=fr, ) + # Token counters — aggregate across all stages for this request. + _prompt_tok = 0 + _gen_tok = 0 + for evt in metrics.stage_events.get(rid_key, []): + if evt.stage_id == 0: + _prompt_tok += int(evt.num_tokens_in) + _gen_tok += int(evt.num_tokens_out) + self.prom_metrics.observe_tokens(_prompt_tok, _gen_tok) + # Modality observe inside the same finalize guard so it fires # once per request and inherits the try/except isolation. observe_modality_at_finalize( @@ -578,9 +587,8 @@ class OmniBase(PDDisaggregationMixin): if current_stage_metrics is not None: response_metrics["stage_id"] = current_stage_metrics["stage_id"] response_metrics["final_output_type"] = current_stage_metrics["final_output_type"] - if current_stage_metrics["final_output_type"] == "text": - response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"] - response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"] + response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"] + response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"] return OmniRequestOutput( request_id=req_id or "", stage_id=stage_id, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f70a439f5..e396c47d6 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -646,6 +646,7 @@ async def build_async_omni_from_stage_config( try: kwargs = args.get_explicit_kwargs_dict() model = kwargs.pop("model", None) or args.model + kwargs.setdefault("log_stats", not args.disable_log_stats) async_omni = AsyncOmni(model=model, **kwargs) # # Don't keep the dummy data in memory diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index 6c1288edf..00b6ff1cb 100644 --- a/vllm_omni/metrics/definitions.py +++ b/vllm_omni/metrics/definitions.py @@ -114,6 +114,10 @@ E2E_REQUEST_LATENCY_S = METRIC_PREFIX + "e2e_request_latency_s" # ``_total`` at exposition time. REQUESTS_SUCCESS = METRIC_PREFIX + "requests_success" +# Token counters — aggregated across all pipeline stages per request. +PROMPT_TOKENS = METRIC_PREFIX + "prompt_tokens" +GENERATION_TOKENS = METRIC_PREFIX + "generation_tokens" + # ============================================================================ # Audio family (per-stage + per-replica audio path metrics) diff --git a/vllm_omni/metrics/prometheus.py b/vllm_omni/metrics/prometheus.py index fc971df03..225a7c8db 100644 --- a/vllm_omni/metrics/prometheus.py +++ b/vllm_omni/metrics/prometheus.py @@ -27,6 +27,16 @@ _e2e_latency_family = Histogram( labelnames=_labelnames, buckets=defs.SECONDS_BUCKETS, ) +_prompt_tokens_family = Counter( + defs.PROMPT_TOKENS, + "Total prompt (input) tokens processed across all pipeline stages.", + labelnames=_labelnames, +) +_generation_tokens_family = Counter( + defs.GENERATION_TOKENS, + "Total generation (output) tokens produced across all pipeline stages.", + labelnames=_labelnames, +) class OmniPrometheusMetrics: @@ -43,6 +53,8 @@ class OmniPrometheusMetrics: self._running = _running_family.labels(model_name=model_name) self._waiting = _waiting_family.labels(model_name=model_name) self._e2e_latency = _e2e_latency_family.labels(model_name=model_name) + self._prompt_tokens = _prompt_tokens_family.labels(model_name=model_name) + self._generation_tokens = _generation_tokens_family.labels(model_name=model_name) def set_running(self, n: int) -> None: if not self._log_stats: @@ -54,6 +66,14 @@ class OmniPrometheusMetrics: return self._waiting.set(n) + def observe_tokens(self, prompt_tokens: int, generation_tokens: int) -> None: + if not self._log_stats: + return + if prompt_tokens > 0: + self._prompt_tokens.inc(prompt_tokens) + if generation_tokens > 0: + self._generation_tokens.inc(generation_tokens) + def request_succeeded( self, e2e_seconds: float,