[Metric] Fix the computation of e2e_total_tokens (#519)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/entrypoints/test_omni_llm.py

@@ -304,7 +304,7 @@ def _setup_log_mocks(monkeypatch):
         def on_stage_metrics(self, stage_id, req_id, metrics):
             pass
 
-        def on_finalize_request(self, stage_id, req_id, engine_outputs, start_ts):
+        def on_finalize_request(self, stage_id, req_id, start_ts):
             self.e2e_done.add(req_id)
 
         def on_forward(self, from_stage, to_stage, req_id, size_bytes, tx_ms, use_shm):

vllm_omni/entrypoints/log_utils.py

@@ -170,10 +170,14 @@ def record_stage_metrics(
         rid_key = str(req_id)
         pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
         pr_stages = pr["stages"]  # type: ignore[index]
-        pr_stages[stage_id] = {
+        stage_data: dict[str, Any] = {
             "stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)),
             "num_tokens_out": int(metrics.get("num_tokens_out", 0)),
         }
+        # Only record num_tokens_in for stage 0 (initial prompt)
+        if stage_id == 0:
+            stage_data["num_tokens_in"] = int(metrics.get("num_tokens_in", 0))
+        pr_stages[stage_id] = stage_data
     except Exception:
         pass
 
@@ -352,6 +356,7 @@ class StageStats:
 
 @dataclass
 class StageRequestMetrics:
+    num_tokens_in: int
     num_tokens_out: int
     stage_gen_time_ms: float
     batch_id: int
@@ -500,7 +505,12 @@ class OrchestratorMetrics:
             float(tx_ms),
         )
 
-    def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: list[Any], req_start_ts: float) -> None:
+    def on_finalize_request(
+        self,
+        stage_id: int,
+        req_id: Any,
+        req_start_ts: float,
+    ) -> None:
         rid_key = str(req_id)
         _t0 = float(req_start_ts)
         _t1 = time.time()
@@ -509,21 +519,31 @@ class OrchestratorMetrics:
         self.stage_last_ts[stage_id] = _t1 if prev_last is None else max(prev_last, _t1)
         self.last_finish_ts = max(self.last_finish_ts, _t1)
         e2e_ms = (_t1 - _t0) * 1000.0
-        num_tokens = count_tokens_from_outputs(engine_outputs)
+
+        # Sum tokens from all stages for this request
+        # Include input tokens from stage 0 + output tokens from all stages
+        pr = self.per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+        total_tokens = 0
+        stages_info = pr.get("stages", {})
+        for sid, stage_data in stages_info.items():
+            # Add input tokens only from stage 0 (initial prompt)
+            if sid == 0:
+                total_tokens += int(stage_data.get("num_tokens_in", 0))
+            total_tokens += int(stage_data.get("num_tokens_out", 0))
+
         self.e2e_total_ms += e2e_ms
-        self.e2e_total_tokens += int(num_tokens)
+        self.e2e_total_tokens += total_tokens
         self.e2e_count += 1
         self.e2e_done.add(rid_key)
-        pr = self.per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
         per_req_record = {
             "type": "request_level_metrics",
             "request_id": rid_key,
             "e2e_time_ms": e2e_ms,
-            "e2e_tpt": (e2e_ms / num_tokens) if num_tokens > 0 else 0.0,
-            "num_tokens_out": int(num_tokens),
+            "e2e_tpt": (e2e_ms / total_tokens) if total_tokens > 0 else 0.0,
+            "e2e_total_tokens": total_tokens,
             "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),
             "transfers_total_bytes": int(pr.get("transfers_bytes", 0)),
-            "stages": pr.get("stages", {}),
+            "stages": stages_info,
         }
         self.sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0))
         logger.info(pformat(per_req_record, sort_dicts=False))

vllm_omni/entrypoints/omni.py

@@ -566,8 +566,10 @@ class Omni(OmniBase):
                 # Mark last output time for this stage whenever we receive outputs
                 metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
                 try:
-                    _m = asdict(result.get("metrics"))
+                    _m = result.get("metrics")
                     if _m is not None:
+                        if not isinstance(_m, dict):
+                            _m = asdict(_m)
                         metrics.on_stage_metrics(stage_id, req_id, _m)
                         if pbar:
                             elapsed = pbar.format_dict["elapsed"] or 1e-6
@@ -610,7 +612,6 @@ class Omni(OmniBase):
                             metrics.on_finalize_request(
                                 stage_id,
                                 req_id,
-                                engine_outputs,
                                 _req_start_ts.get(req_id, _wall_start_ts),
                             )
                     except Exception as e:

vllm_omni/entrypoints/omni_stage.py

@@ -1269,6 +1269,19 @@ async def _stage_worker_async(
     logger.info("Stage worker exiting")
 
 
+def count_prompt_tokens_from_outputs(engine_outputs: list[Any]) -> int:
+    """Count prompt tokens from engine outputs."""
+    total = 0
+    for _ro in engine_outputs:
+        try:
+            prompt_token_ids = getattr(_ro, "prompt_token_ids", None)
+            if prompt_token_ids is not None:
+                total += len(prompt_token_ids)
+        except Exception:
+            pass
+    return total
+
+
 def make_request_stats(
     req_output: list[Any],
     stage_gen_time_ms: float,
@@ -1282,8 +1295,10 @@ def make_request_stats(
         StageRequestMetrics,
     )
 
+    num_tokens_in = count_prompt_tokens_from_outputs(req_output)
     num_tokens_out = count_tokens_from_outputs(req_output)
     return StageRequestMetrics(
+        num_tokens_in=num_tokens_in,
         num_tokens_out=num_tokens_out,
         stage_gen_time_ms=stage_gen_time_ms,
         batch_id=batch_id,