[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][]
@@ -105,3 +102,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][]

docs/usage/metrics.md

@@ -0,0 +1,59 @@
+# Production Metrics
+
+## Usage
+
+Users can utilize these metrics in production environments to monitor the health and performance of the vLLM-omni system. Key scenarios include:
+- **Performance Monitoring**: Track throughput (e.g., `e2e_avg_tokens_per_s`), latency (e.g., `e2e_total_ms`), and resource utilization to verify that the system meets expected performance standards.
+- **Debugging and Troubleshooting**: Use detailed per-request metrics to diagnose issues with specific requests, such as high transfer times or unexpected token counts.
+-  **Enable Logging**: Start vLLM-omni with the `--log-stats` flag. This exposes metrics through structured logs.
+
+## Overall Summary
+
+| Field                     | Meaning                                                                                       |
+|---------------------------|----------------------------------------------------------------------------------------------|
+| `e2e_requests`            | Number of completed requests.                                                                |
+| `e2e_wall_time_ms`        | Wall-clock time span from run start to last completion, in ms.                               |
+| `e2e_total_tokens`        | Total tokens counted across all completed requests (stage0 input + all stage outputs).       |
+| `e2e_avg_time_per_request_ms` | Average wall time per request: `e2e_wall_time_ms / e2e_requests`.                        |
+| `e2e_avg_tokens_per_s`    | Average token throughput over wall time: `e2e_total_tokens * 1000 / e2e_wall_time_ms`.      |
+| `stage_wall_time_ms`      | Wall-clock time span for each stage, in ms (list format).                                   |
+
+---
+
+## E2E Table (per request)
+
+| Field                     | Meaning                                                               |
+|---------------------------|-----------------------------------------------------------------------|
+| `e2e_total_ms`            | End-to-end latency in ms.                                             |
+| `e2e_total_tokens`        | Total tokens for the request (stage0 input + all stage outputs).      |
+| `transfers_total_time_ms` | Sum of transfer edge `total_time_ms` for this request.                |
+| `transfers_total_kbytes`  | Sum of transfer kbytes for this request.                              |
+
+
+---
+
+## Stage Table (per stage event / request)
+
+| Field               | Meaning                                                                             |
+|---------------------|-------------------------------------------------------------------------------------|
+| `batch_id`          | Batch index.                                                                        |
+| `batch_size`        | Batch size.                                                                         |
+| `num_tokens_in`     | Input tokens to the stage.                                                          |
+| `num_tokens_out`    | Output tokens from the stage.                                                       |
+| `preprocess_time_ms` | Preprocessing time in ms.                                                           |
+| `stage_gen_time_ms` | Stage compute time in ms, excluding preprocessing time (reported separately as `preprocess_time_ms`). |
+
+---
+
+## Transfer Table (per edge / request)
+
+| Field                | Meaning                                                                   |
+|----------------------|---------------------------------------------------------------------------|
+| `size_kbytes`        | Total kbytes transferred.                                                 |
+| `tx_time_ms`         | Sender transfer time in ms.                                               |
+| `rx_decode_time_ms`  | Receiver decode time in ms.                                               |
+| `in_flight_time_ms`  | In-flight time in ms.                                                     |
+
+## Expectation of the numbers:
+e2e_total_tokens = Stage0 's num_tokens_in + other stage's num_tokens_out
+transfers_total_time_ms = sum(tx_time_ms + rx_decode_time_ms + in_flight_time_ms) in every edge

examples/offline_inference/bagel/end2end.py

@@ -36,7 +36,7 @@ def parse_args():
     )
 
     # OmniLLM init args
-    parser.add_argument("--enable-stats", action="store_true", default=False)
+    parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-sleep-seconds", type=int, default=20)
     parser.add_argument("--batch-timeout", type=int, default=5)
     parser.add_argument("--init-timeout", type=int, default=300)
@@ -120,7 +120,7 @@ def main():
 
         omni_kwargs.update(
             {
-                "log_stats": args.enable_stats,
+                "log_stats": args.log_stats,
                 "init_sleep_seconds": args.init_sleep_seconds,
                 "batch_timeout": args.batch_timeout,
                 "init_timeout": args.init_timeout,

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -322,7 +322,7 @@ def main(args):
         query_result = query_func()
     omni_llm = Omni(
         model=model_name,
-        log_stats=args.enable_stats,
+        log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
         batch_timeout=args.batch_timeout,
         init_timeout=args.init_timeout,
@@ -439,7 +439,7 @@ def parse_args():
         help="Query type.",
     )
     parser.add_argument(
-        "--enable-stats",
+        "--log-stats",
         action="store_true",
         default=False,
         help="Enable writing detailed statistics (default: disabled)",

examples/offline_inference/qwen3_omni/end2end.py

@@ -328,7 +328,7 @@ def main(args):
     omni_llm = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
-        log_stats=args.enable_stats,
+        log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
 
@@ -455,7 +455,7 @@ def parse_args():
         help="Query type.",
     )
     parser.add_argument(
-        "--enable-stats",
+        "--log-stats",
         action="store_true",
         default=False,
         help="Enable writing detailed statistics (default: disabled)",

examples/offline_inference/qwen3_tts/end2end.py

@@ -219,7 +219,7 @@ def main(args):
     omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
-        log_stats=args.enable_stats,
+        log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
 
@@ -275,7 +275,7 @@ def parse_args():
         help="Query type.",
     )
     parser.add_argument(
-        "--enable-stats",
+        "--log-stats",
         action="store_true",
         default=False,
         help="Enable writing detailed statistics (default: disabled)",

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
     "fa3-fwd", # flash attention 3, maintained by @ZJY0516
     "openai-whisper>=20250625",

tests/e2e/online_serving/test_async_omni.py

@@ -154,8 +154,14 @@ 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)
-
-            # 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"])
+            overall = summary["overall_summary"]
+            assert overall["e2e_wall_time_ms"] >= 0.0
+
+            # 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/entrypoints/test_omni_llm.py

@@ -296,9 +296,9 @@ def _setup_log_mocks(monkeypatch):
     # Mock OrchestratorMetrics to be a simple class that doesn't require file operations
 
     class _FakeOrchestratorMetrics:
-        def __init__(self, num_stages, enable_stats, wall_start_ts):
+        def __init__(self, num_stages, log_stats, wall_start_ts):
             self.num_stages = num_stages
-            self.enable_stats = enable_stats
+            self.log_stats = log_stats
             self.stage_first_ts = [None] * num_stages
             self.stage_last_ts = [None] * num_stages
             self.e2e_done = set()
@@ -316,7 +316,7 @@ def build_and_log_summary(self, final_stage_id):
             return "Fake summary"
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.omni.OrchestratorMetrics",
+        "vllm_omni.entrypoints.omni.OrchestratorAggregator",
         _FakeOrchestratorMetrics,
         raising=False,
     )

tests/metrics/test_stats.py

@@ -0,0 +1,122 @@
+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, log_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, log_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, log_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/diffusion/diffusion_engine.py

@@ -104,6 +104,8 @@ def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
             request_id = request.request_ids[0] if request.request_ids else ""
 
             metrics = {}
+            metrics["image_num"] = int(request.sampling_params.num_outputs_per_prompt)
+            metrics["resolution"] = int(request.sampling_params.resolution)
             if output.trajectory_timesteps is not None:
                 metrics["trajectory_timesteps"] = output.trajectory_timesteps
 
@@ -145,6 +147,8 @@ def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
                 output_idx += num_outputs
 
                 metrics = {}
+                metrics["image_num"] = int(request.num_outputs_per_prompt)
+                metrics["resolution"] = int(request.resolution)
                 if output.trajectory_timesteps is not None:
                     metrics["trajectory_timesteps"] = output.trajectory_timesteps
 

vllm_omni/distributed/omni_connectors/adapter.py

@@ -15,6 +15,7 @@
     from .connectors.base import OmniConnectorBase
 
 from vllm_omni.entrypoints.stage_utils import OmniStageTaskType
+from vllm_omni.metrics import OrchestratorAggregator
 
 from .utils.logging import get_connector_logger
 
@@ -30,7 +31,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.