[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) <noreply@anthropic.com>

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,

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

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)

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,