diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index e9c20f3c6..7c30f933c 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/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): diff --git a/vllm_omni/entrypoints/log_utils.py b/vllm_omni/entrypoints/log_utils.py index ac5bfc277..ba147f12c 100644 --- a/vllm_omni/entrypoints/log_utils.py +++ b/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)) diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index b5100ee5b..f994ea494 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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: diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 698c7a87e..8f3fac1d9 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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,