[Feature] Opt metrics structure

docs/api/README.md

@@ -12,9 +12,6 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
-- [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
-- [vllm_omni.entrypoints.log_utils.StageRequestMetrics][]
-- [vllm_omni.entrypoints.log_utils.StageStats][]
 - [vllm_omni.entrypoints.omni.Omni][]
 - [vllm_omni.entrypoints.omni.OmniBase][]
 - [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
@@ -120,3 +117,8 @@ Worker classes and model runners for distributed inference.
 - [vllm_omni.worker.npu.npu_generation_model_runner.NPUGenerationModelRunner][]
 - [vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker][]
 - [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][]
+
+
+## Metrics
+
+- [vllm_omni.metrics.OrchestratorAggregator][]

pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "soundfile>=0.13.1",
     "cache-dit==1.2.0",
     "tqdm>=4.66.0",
+    "prettytable>=3.9.0",
     "torchsde>=0.2.6",  # Required for Stable Audio scheduler
     "openai-whisper>=20250625",
     # "vllm==0.14.0",  # TODO: fix the entrypoints overwrite problem

tests/e2e/online_serving/test_async_omni.py

@@ -153,8 +153,19 @@ def __init__(self, request_id: str):
                 output_modalities, engine.output_modalities, engine.stage_list
             )
             summary = capture_metrics[request_ids[idx]].metrics.build_and_log_summary(final_stage_id_for_e2e)
+            overall = summary["overall_summary"]
+            assert overall["e2e_wall_time_ms"] >= 0.0
 
-            # Check that total tokens matches sum of stage tokens.
-            assert summary["e2e_total_tokens"] == sum(stage["tokens"] for stage in summary["stages"])
-            # Check that total time matches sum of stage times.
-            assert summary["e2e_total_time_ms"] >= sum(stage["total_time_ms"] for stage in summary["stages"])
+            # Check that total tokens matches sum of stage tokens for this request.
+            stage_entry = next(
+                entry for entry in summary["stage_table"] if entry["request_id"] == request_ids[idx]
+            )
+            stage_sum = sum(
+                (stage.get("num_tokens_in", 0) if stage.get("stage_id") == 0 else 0)
+                + stage.get("num_tokens_out", 0)
+                for stage in stage_entry["stages"]
+            )
+            e2e_entry = next(
+                entry for entry in summary["e2e_table"] if entry["request_id"] == request_ids[idx]
+            )
+            assert e2e_entry["e2e_total_tokens"] == stage_sum

tests/metrics/test_stats.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+from vllm_omni.metrics import OrchestratorAggregator
+from vllm_omni.metrics.stats import RequestE2EStats
+
+
+def _get_request_entry(table: list[dict], request_id: str) -> dict:
+    for entry in table:
+        if entry.get("request_id") == request_id:
+            return entry
+    raise AssertionError(f"request_id={request_id} not found")
+
+
+def test_orchestrator_aggregator_builds_summary() -> None:
+    agg = OrchestratorAggregator(num_stages=2, enable_stats=False, wall_start_ts=0.0)
+    agg.stage_first_ts[0] = 0.0
+    agg.stage_last_ts[0] = 0.03
+    agg.stage_first_ts[1] = 0.05
+    agg.stage_last_ts[1] = 0.07
+
+    agg.on_forward(0, 1, "r1", size_bytes=1024, tx_ms=5.0, used_shm=False)
+    agg.on_stage_metrics(
+        0,
+        "r1",
+        {
+            "num_tokens_in": 3,
+            "num_tokens_out": 3,
+            "stage_gen_time_ms": 30.0,
+            "batch_id": 1,
+            "batch_size": 1,
+            "rx_transfer_bytes": 0,
+            "rx_decode_time_ms": 0.0,
+        },
+    )
+    agg.on_stage_metrics(
+        1,
+        "r1",
+        {
+            "num_tokens_out": 4,
+            "stage_gen_time_ms": 20.0,
+            "batch_id": 1,
+            "batch_size": 1,
+            "rx_transfer_bytes": 1024,
+            "rx_decode_time_ms": 5.0,
+            "rx_in_flight_time_ms": 2.0,
+        },
+    )
+    agg.on_finalize_request(1, "r1", req_start_ts=0.0)
+
+    summary = agg.build_and_log_summary(final_stage_id_to_prompt={"r1": 1})
+    overall = summary["overall_summary"]
+    assert overall["e2e_requests"] == 1
+
+    stage_entry = _get_request_entry(summary["stage_table"], "r1")
+    stage_ids = [row["stage_id"] for row in stage_entry["stages"]]
+    assert stage_ids == [0, 1]
+
+    transfer_entry = _get_request_entry(summary["trans_table"], "r1")
+    assert transfer_entry["transfers"][0]["edge"] == "0->1"
+    assert transfer_entry["transfers"][0]["size_kbytes"] == 1.0
+
+    e2e_entry = _get_request_entry(summary["e2e_table"], "r1")
+    assert e2e_entry["e2e_total_tokens"] == 10
+
+
+def test_build_and_log_summary_e2e_only() -> None:
+    agg = OrchestratorAggregator(num_stages=1, enable_stats=False, wall_start_ts=0.0)
+    agg.e2e_events.append(
+        RequestE2EStats(
+            request_id="r",
+            e2e_total_ms=10.0,
+            e2e_total_tokens=5,
+            transfers_total_time_ms=0.0,
+            transfers_total_bytes=0,
+        )
+    )
+
+    summary = agg.build_and_log_summary(final_stage_id_to_prompt=0)
+    e2e_entry = _get_request_entry(summary["e2e_table"], "r")
+    assert e2e_entry["e2e_total_tokens"] == 5
+    stage_entry = _get_request_entry(summary["stage_table"], "r")
+    assert stage_entry["stages"] == []
+
+
+def test_build_and_log_summary_multiple_requests() -> None:
+    agg = OrchestratorAggregator(num_stages=1, enable_stats=False, wall_start_ts=0.0)
+
+    agg.on_stage_metrics(
+        0,
+        "r1",
+        {
+            "num_tokens_in": 2,
+            "num_tokens_out": 4,
+            "batch_id": 1,
+            "batch_size": 1,
+            "stage_gen_time_ms": 10.0,
+            "rx_transfer_bytes": 0,
+            "rx_decode_time_ms": 0.0,
+            "rx_in_flight_time_ms": 0.0,
+        },
+    )
+    agg.on_finalize_request(0, "r1", req_start_ts=0.0)
+
+    agg.on_stage_metrics(
+        0,
+        "r2",
+        {
+            "num_tokens_in": 1,
+            "num_tokens_out": 2,
+            "batch_id": 2,
+            "batch_size": 1,
+            "stage_gen_time_ms": 12.0,
+            "rx_transfer_bytes": 0,
+            "rx_decode_time_ms": 0.0,
+            "rx_in_flight_time_ms": 0.0,
+        },
+    )
+    agg.on_finalize_request(0, "r2", req_start_ts=0.0)
+
+    summary = agg.build_and_log_summary(final_stage_id_to_prompt=0)
+    assert len(summary["stage_table"]) == 2
+    assert {entry["request_id"] for entry in summary["e2e_table"]} == {"r1", "r2"}

vllm_omni/distributed/omni_connectors/adapter.py

@@ -8,6 +8,7 @@
 from typing import Any
 
 from vllm_omni.entrypoints.stage_utils import OmniStageTaskType
+from vllm_omni.metrics import OrchestratorAggregator
 
 from .utils.logging import get_connector_logger
 
@@ -23,7 +24,7 @@ def try_send_via_connector(
     sampling_params: Any,
     original_prompt: Any,
     next_stage_queue_submit_fn: Callable[[dict[str, Any]], None],
-    metrics: Any,
+    metrics: OrchestratorAggregator,
 ) -> bool:
     """
     Attempts to send data via OmniConnector.

vllm_omni/entrypoints/async_omni.py

@@ -16,6 +16,7 @@
 from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import TokenizerLike
 from vllm.v1.engine.exceptions import EngineDeadError
+import vllm.envs as envs
 
 # Internal imports (our code)
 from vllm_omni.config import OmniModelConfig
@@ -24,16 +25,14 @@
 from vllm_omni.distributed.ray_utils.utils import try_close_ray
 from vllm_omni.engine.input_processor import OmniInputProcessor
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
-from vllm_omni.entrypoints.log_utils import (
-    OrchestratorMetrics,
-)
 from vllm_omni.entrypoints.omni import OmniBase
 from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
 )
+from vllm_omni.metrics import OrchestratorAggregator
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -57,7 +56,6 @@ def _weak_close_cleanup_async(stage_list, stage_in_queues, ray_pg, output_handle
     if output_handler is not None:
         output_handler.cancel()
 
-
 class AsyncOmni(OmniBase):
     """Asynchronous unified entry point supporting multi-stage pipelines for LLM and Diffusion models.
 
@@ -320,27 +318,27 @@ async def generate(self, *args: Any, **kwargs: dict[str, Any]) -> AsyncGenerator
             )
 
             # Metrics/aggregation helper
-            metrics = OrchestratorMetrics(
-                num_stages,
-                self._enable_stats,
-                _wall_start_ts,
+            metrics = OrchestratorAggregator(
+                num_stages=num_stages,
+                enable_stats=self._enable_stats,
+                wall_start_ts=_wall_start_ts, # will be reset at generate() time, just a placeholder here
             )
             # Seed stage-0 queue with all requests
             logger.debug(f"[{self._name}] Seeding request into stage-0")
             req_state = ClientRequestState(request_id)
             req_state.metrics = metrics
             self.request_states[request_id] = req_state
+
+            _req_start_ts[request_id] = time.time()
             # Mark first input time for stage-0
             metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
-
             sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
             task = {
                 "request_id": request_id,
                 "engine_inputs": prompt,
                 "sampling_params": sp0,
             }
             self.stage_list[0].submit(task)
-            _req_start_ts[request_id] = time.time()
             logger.debug(f"[{self._name}] Enqueued request {request_id} to stage-0")
 
             logger.debug(f"[{self._name}] Entering scheduling loop: stages={num_stages}")
@@ -366,6 +364,8 @@ async def generate(self, *args: Any, **kwargs: dict[str, Any]) -> AsyncGenerator
                     metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
                     try:
                         _m = asdict(result.get("metrics"))
+                        # stage_gen_time_ms is the time of generating every chunk in this stage
+                        metrics.accumulated_gen_time_ms[req_id] += _m.get("stage_gen_time_ms", 0.0)
                         if _m is not None and finished:
                             metrics.on_stage_metrics(stage_id, req_id, _m)
                     except Exception as e:
@@ -423,7 +423,11 @@ async def generate(self, *args: Any, **kwargs: dict[str, Any]) -> AsyncGenerator
                 next_stage_id = stage_id + 1
                 if next_stage_id <= final_stage_id_for_e2e and finished:
                     next_stage: OmniStage = self.stage_list[next_stage_id]
+                    # Derive inputs for the next stage, record preprocess time
+                    _prep_t0 = time.perf_counter()
                     next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
+                    _prep_ms = (time.perf_counter() - _prep_t0) * 1000.0
+                    metrics.record_stage_preprocess_time(next_stage_id, req_id, _prep_ms)
                     sp_next: SamplingParams = sampling_params_list[next_stage_id]
 
                     # Check if we have a connector for this edge
@@ -460,11 +464,9 @@ async def generate(self, *args: Any, **kwargs: dict[str, Any]) -> AsyncGenerator
                     logger.debug(f"[{self._name}] Request {req_id} fully completed")
 
             logger.debug(f"[{self._name}] All requests completed")
-
             # Summarize and print stats
             try:
-                summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
-                logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
+                metrics.build_and_log_summary(final_stage_id_for_e2e)
             except Exception as e:
                 logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
             finally:

vllm_omni/entrypoints/client_request_state.py

@@ -1,6 +1,6 @@
 import asyncio
 
-from vllm_omni.entrypoints.log_utils import OrchestratorMetrics
+from vllm_omni.metrics import OrchestratorAggregator
 
 
 class ClientRequestState:
@@ -10,4 +10,4 @@ def __init__(self, request_id: str, queue: asyncio.Queue | None = None):
         self.request_id = request_id
         self.stage_id: int | None = None
         self.queue = queue if queue is not None else asyncio.Queue()
-        self.metrics: OrchestratorMetrics | None = None
+        self.metrics: OrchestratorAggregator | None = None