anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feature] Opt metrics structure (#891)

Browse Source 
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
Signed-off-by: Junhong Liu <ljh_lbj@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Junhong LiuandHongsheng Liu authored and GitHub committed 2026-02-09 19:17:32 +08:00
1 parent ac397ae2a0
commit 9af6fb9a1f
23 files changed
+1151 -728
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

benchmarks/qwen3-omni/README.md
+1 -1
View File
@@ -53,7 +53,7 @@ bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
```
What it does:
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--enable-stats`.
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
- Uses `benchmarks/build_dataset/top100.txt` and writes to:
- Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
- `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
+1 -1
View File
@@ -23,7 +23,7 @@ else
python $end2end_script_path --output-wav $outputs_dir \
--query-type text \
--txt-prompts $build_dataset_path \
--enable-stats \
--log-stats \
--log-dir $log_dir
echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
echo " - omni_llm_pipeline_text run dir/base name"
docs/api/README.md
+5 -3
View File
@@ -15,9 +15,6 @@ Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
- [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][]
@@ -114,3 +111,8 @@ Worker classes and model runners for distributed inference.
- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
- [vllm_omni.worker.mixins.OmniWorkerMixin][]
## Metrics
- [vllm_omni.metrics.OrchestratorAggregator][]
docs/contributing/metrics.md
+165
View File
@@ -0,0 +1,165 @@
# Metrics vLLM-Omni:
You can use these metrics in production to monitor the health and performance of the vLLM-omni system. Typical 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 standards.
- **Debugging and Troubleshooting**: Use detailed per-request metrics to diagnose issues, such as high transfer times or unexpected token counts.
## How to Enable and View Metrics
### 1. Start the Service with Metrics Logging
```bash
vllm serve /workspace/models/Qwen3-Omni-30B-A3B-Instruct --omni --port 8014 --log-stats
```
### 2. Send a Request
```bash
python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
```
### 3. What You Will See
With `--log-stats` enabled, the server will output detailed metrics logs after each request. Example output:
#### Overall Summary
| Field | Value |
|-----------------------------|--------------|
| e2e_requests | 1 |
| e2e_wall_time_ms | 41,299.190 |
| e2e_total_tokens | 5,202 |
| e2e_avg_time_per_request_ms | 41,299.190 |
| e2e_avg_tokens_per_s | 125.959 |
| e2e_stage_0_wall_time_ms | 10,192.289 |
| e2e_stage_1_wall_time_ms | 30,541.409 |
| e2e_stage_2_wall_time_ms | 207.496 |
#### RequestE2EStats
| Field | Value |
|-------------------------|------------|
| e2e_total_ms | 41,299.133 |
| e2e_total_tokens | 5,202 |
| transfers_total_time_ms | 245.895 |
| transfers_total_kbytes | 138,089.939|
#### StageRequestStats
| Field | 0 | 1 | 2 |
|------------------------|--------|--------|--------|
| audio_generated_frames | 0 | 0 | 525,525|
| batch_id | 38 | 274 | 0 |
| batch_size | 1 | 1 | 1 |
| num_tokens_in | 4,860 | 4,826 | 4,384 |
| num_tokens_out | 67 | 275 | 0 |
| postprocess_time_ms | 256.158| 0.491 | 0.000 |
| stage_gen_time_ms | 9,910.007|30,379.198|160.745|
#### TransferEdgeStats
| Field | 0->1 | 1->2 |
|---------------------|-------------|------------|
| size_kbytes | 109,277.349 | 28,812.591 |
| tx_time_ms | 78.701 | 18.790 |
| rx_decode_time_ms | 111.865 | 31.706 |
| in_flight_time_ms | 2.015 | 2.819 |
These logs include:
- **Overall summary**: total requests, wall time, average tokens/sec, etc.
- **E2E table**: per-request latency and token counts.
- **Stage table**: per-stage batch and timing details.
- **Transfer table**: data transfer and timing for each edge.
You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above.
## Metrics Scope: Offline vs Online Inference
For **offline inference** (batch mode), the summary includes both system-level metrics (aggregated across all requests) and per-request metrics. In this case, `e2e_requests` can be greater than 1, reflecting multiple completed requests in a batch.
For **online inference** (serving mode), the summary is always per-request. `e2e_requests` is always 1, and only request-level metrics are reported for each completion.
---
## Parameter Details
| 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`. |
| `e2e_stage_{i}_wall_time_ms` | Wall-clock time span for stage i, in ms. Each stage's wall time is reported as a separate field, e.g., `e2e_stage_0_wall_time_ms`, `e2e_stage_1_wall_time_ms`, etc. |
---
## 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. |
| `stage_gen_time_ms` | Stage compute time in ms, excluding postprocessing time (reported separately as `postprocess_time_ms`). |
| `image_num` | Number of images generated (for diffusion/image stages). |
| `resolution` | Image resolution (for diffusion/image stages). |
| `postprocess_time_ms` | Diffusion/image: post-processing time in 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 (Verification)
**Formulas:**
- `e2e_total_tokens = Stage0's num_tokens_in + sum(all stages' num_tokens_out)`
- `transfers_total_time_ms = sum(tx_time_ms + rx_decode_time_ms + in_flight_time_ms)` for every edge
**Using the example above:**
### e2e_total_tokens
- Stage0's `num_tokens_in`: **4,860**
- Stage0's `num_tokens_out`: **67**
- Stage1's `num_tokens_out`: **275**
- Stage2's `num_tokens_out`: **0**
So,
```
e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202
```
This matches the table value: `e2e_total_tokens = 5,202`.
### transfers_total_time_ms
For each edge:
- 0->1: tx_time_ms (**78.701**) + rx_decode_time_ms (**111.865**) + in_flight_time_ms (**2.015**) = **192.581**
- 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315**
Sum: 192.581 + 53.315 = **245.896**
The table shows `transfers_total_time_ms = 245.895`, which matches the calculation (difference is due to rounding).
examples/offline_inference/bagel/end2end.py
+2 -2
View File
@@ -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
+2 -2
View File
@@ -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
+2 -2
View File
@@ -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
+2 -2
View File
@@ -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)",
requirements/common.txt
+1
View File
@@ -12,3 +12,4 @@ torchsde>=0.2.6
openai-whisper>=20250625
imageio[ffmpeg]>=2.37.2
sox>=1.5.0
prettytable>=3.8.0
tests/e2e/online_serving/test_async_omni.py
+1 -52
View File
@@ -9,7 +9,7 @@ from vllm import SamplingParams
from vllm.inputs import PromptType
from tests.utils import hardware_test
from vllm_omni.entrypoints.async_omni import AsyncOmni, ClientRequestState
from vllm_omni.entrypoints.async_omni import AsyncOmni
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -115,54 +115,3 @@ async def test_abort():
num_generated_tokens, request_id = await task
assert num_generated_tokens == NUM_EXPECTED_TOKENS
await asyncio.sleep(5)
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.asyncio
async def test_build_and_log_summary(monkeypatch):
from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e
RealCRS = ClientRequestState
capture_metrics = {}
class MockCRS(RealCRS):
def __init__(self, request_id: str):
super().__init__(request_id)
capture_metrics[request_id] = self
monkeypatch.setattr("vllm_omni.entrypoints.async_omni.ClientRequestState", MockCRS)
monkeypatch.setattr("vllm_omni.entrypoints.client_request_state.ClientRequestState", MockCRS)
with ExitStack() as after:
# Avoid SHM IPC in tests to prevent /dev/shm exhaustion and SIGBUS.
engine = AsyncOmni(
model=model,
stage_configs_path=stage_config,
shm_threshold_bytes=sys.maxsize,
)
after.callback(engine.shutdown)
prompt = "Hello my name is Robert and "
NUM_EXPECTED_TOKENS = 64
NUM_REQUESTS = 3
request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
# Create concurrent requests.
tasks: list[asyncio.Task] = []
for idx, request_id in enumerate(request_ids):
tasks.append(asyncio.create_task(generate(engine, request_id, prompt, NUM_EXPECTED_TOKENS)))
# Confirm the requests are okay.
for idx, task in enumerate(tasks):
await task
output_modalities = ["text"]
final_stage_id_for_e2e = get_final_stage_id_for_e2e(
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"])
tests/entrypoints/test_omni_llm.py
+3 -3
View File
@@ -298,9 +298,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()
@@ -318,7 +318,7 @@ def _setup_log_mocks(monkeypatch):
return "Fake summary"
monkeypatch.setattr(
"vllm_omni.entrypoints.omni.OrchestratorMetrics",
"vllm_omni.entrypoints.omni.OrchestratorAggregator",
_FakeOrchestratorMetrics,
raising=False,
)
tests/metrics/test_stats.py
+157
View File
@@ -0,0 +1,157 @@
from __future__ import annotations
import pytest
from vllm_omni.metrics import OrchestratorAggregator
from vllm_omni.metrics.stats import RequestE2EStats, StageRequestStats, StageStats
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
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=True, wall_start_ts=0.0, final_stage_id_for_e2e=1)
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",
StageRequestStats(
batch_id=1,
batch_size=1,
num_tokens_in=3,
num_tokens_out=3,
stage_gen_time_ms=30.0,
rx_transfer_bytes=0,
rx_decode_time_ms=0.0,
rx_in_flight_time_ms=0.0,
stage_stats=StageStats(),
),
)
agg.on_stage_metrics(
1,
"r1",
StageRequestStats(
batch_id=1,
batch_size=1,
num_tokens_in=0,
num_tokens_out=4,
stage_gen_time_ms=20.0,
rx_transfer_bytes=1024,
rx_decode_time_ms=5.0,
rx_in_flight_time_ms=2.0,
stage_stats=StageStats(),
),
)
agg.on_finalize_request(1, "r1", req_start_ts=0.0)
summary = agg.build_and_log_summary()
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=True, wall_start_ts=0.0, final_stage_id_for_e2e=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()
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=2, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e={"r1": 1, "r2": 0}
)
# Request r1 goes through both stages
agg.on_stage_metrics(
0,
"r1",
StageRequestStats(
batch_id=1,
batch_size=1,
num_tokens_in=2,
num_tokens_out=4,
stage_gen_time_ms=10.0,
rx_transfer_bytes=0,
rx_decode_time_ms=0.0,
rx_in_flight_time_ms=0.0,
stage_stats=StageStats(),
),
)
agg.on_stage_metrics(
1,
"r1",
StageRequestStats(
batch_id=1,
batch_size=1,
num_tokens_in=4,
num_tokens_out=5,
stage_gen_time_ms=8.0,
rx_transfer_bytes=0,
rx_decode_time_ms=0.0,
rx_in_flight_time_ms=0.0,
stage_stats=StageStats(),
),
)
agg.on_finalize_request(1, "r1", req_start_ts=0.0)
# Request r2 only goes through stage 0
agg.on_stage_metrics(
0,
"r2",
StageRequestStats(
batch_id=2,
batch_size=1,
num_tokens_in=1,
num_tokens_out=2,
stage_gen_time_ms=12.0,
rx_transfer_bytes=0,
rx_decode_time_ms=0.0,
rx_in_flight_time_ms=0.0,
stage_stats=StageStats(),
),
)
agg.on_finalize_request(0, "r2", req_start_ts=0.0)
summary = agg.build_and_log_summary()
assert len(summary["stage_table"]) == 2
assert {entry["request_id"] for entry in summary["e2e_table"]} == {"r1", "r2"}
# Check that r1 has two stages and r2 has one
r1_stage_entry = next(e for e in summary["stage_table"] if e["request_id"] == "r1")
r2_stage_entry = next(e for e in summary["stage_table"] if e["request_id"] == "r2")
assert len(r1_stage_entry["stages"]) == 2
assert len(r2_stage_entry["stages"]) == 1
vllm_omni/diffusion/diffusion_engine.py
+8 -8
View File
@@ -102,16 +102,20 @@ class DiffusionEngine:
if not isinstance(outputs, list):
outputs = [outputs] if outputs is not None else []
metrics = {
"image_num": int(request.sampling_params.num_outputs_per_prompt),
"resolution": int(request.sampling_params.resolution),
"postprocess_time_ms": postprocess_time * 1000,
}
if self.pre_process_func is not None:
metrics["preprocessing_time_ms"] = preprocess_time * 1000
# Handle single request or multiple requests
if len(request.prompts) == 1:
# Single request: return single OmniRequestOutput
prompt = request.prompts[0]
request_id = request.request_ids[0] if request.request_ids else ""
metrics = {}
if output.trajectory_timesteps is not None:
metrics["trajectory_timesteps"] = output.trajectory_timesteps
if supports_audio_output(self.od_config.model_class_name):
audio_payload = outputs[0] if len(outputs) == 1 else outputs
return [
@@ -149,10 +153,6 @@ class DiffusionEngine:
request_outputs = outputs[output_idx : output_idx + num_outputs] if output_idx < len(outputs) else []
output_idx += num_outputs
metrics = {}
if output.trajectory_timesteps is not None:
metrics["trajectory_timesteps"] = output.trajectory_timesteps
if supports_audio_output(self.od_config.model_class_name):
audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs
results.append(
vllm_omni/distributed/omni_connectors/adapter.py
+2 -1
View File
@@ -15,6 +15,7 @@ if TYPE_CHECKING:
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.
vllm_omni/entrypoints/async_omni.py
+34 -32
View File
@@ -5,8 +5,6 @@ import copy
import time
import weakref
from collections.abc import AsyncGenerator, Iterable, Sequence
from dataclasses import asdict
from pprint import pformat
from typing import Any
from vllm.config import VllmConfig
@@ -23,9 +21,6 @@ from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_
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
@@ -37,6 +32,7 @@ from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams
# Internal imports (our code)
from vllm_omni.lora.request import LoRARequest
from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -298,10 +294,11 @@ class AsyncOmni(OmniBase):
)
# Metrics/aggregation helper
metrics = OrchestratorMetrics(
num_stages,
self._enable_stats,
_wall_start_ts,
metrics = OrchestratorAggregator(
num_stages=num_stages,
log_stats=self.log_stats,
wall_start_ts=_wall_start_ts,
final_stage_id_for_e2e=final_stage_id_for_e2e,
)
req_state = ClientRequestState(request_id)
req_state.metrics = metrics
@@ -328,8 +325,6 @@ class AsyncOmni(OmniBase):
req_state,
metrics,
final_stage_id_for_e2e,
_req_start_ts,
_wall_start_ts,
):
yield output
else:
@@ -338,8 +333,6 @@ class AsyncOmni(OmniBase):
req_state,
metrics,
final_stage_id_for_e2e,
_req_start_ts,
_wall_start_ts,
sampling_params_list,
prompt,
):
@@ -348,17 +341,15 @@ class AsyncOmni(OmniBase):
logger.debug(f"[{self._name}] Request {request_id} finalized at stage-{final_stage_id_for_e2e}")
try:
# Finalize E2E metrics if not already done
if str(request_id) not in metrics.e2e_done:
metrics.on_finalize_request(
final_stage_id_for_e2e,
request_id,
_req_start_ts.get(request_id, _wall_start_ts),
)
metrics.on_finalize_request(
final_stage_id_for_e2e,
request_id,
_req_start_ts.get(request_id, _wall_start_ts),
)
logger.debug(f"[{self._name}] All requests completed")
# Summarize and print stats
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()
except Exception as e:
logger.exception(f"[{self._name}] Request {request_id} Failed to finalized/build/log summary: {e}")
finally:
@@ -374,10 +365,8 @@ class AsyncOmni(OmniBase):
prompt: Any,
sampling_params_list: list[SamplingParams],
req_state: ClientRequestState,
metrics: OrchestratorMetrics,
metrics: OrchestratorAggregator,
final_stage_id_for_e2e: int,
req_start_ts: dict[int, float],
wall_start_ts: float,
) -> AsyncGenerator[OmniRequestOutput, None]:
all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)}
submit_flag = True
@@ -413,16 +402,17 @@ class AsyncOmni(OmniBase):
all_stages_finished[stage_id] = finished
if output_to_yield:
metrics.record_audio_generated_frames(
output_to_yield, engine_outputs.finished, stage_id, request_id
)
yield output_to_yield
async def _process_sequential_results(
self,
request_id: str,
req_state: ClientRequestState,
metrics: OrchestratorMetrics,
metrics: OrchestratorAggregator,
final_stage_id_for_e2e: int,
req_start_ts: dict[int, float],
wall_start_ts: float,
sampling_params_list: list[SamplingParams],
prompt: Any,
) -> AsyncGenerator[OmniRequestOutput, None]:
@@ -438,6 +428,9 @@ class AsyncOmni(OmniBase):
metrics,
)
if output_to_yield:
metrics.record_audio_generated_frames(
output_to_yield, engine_outputs.finished, stage_id, request_id
)
yield output_to_yield
if not isinstance(engine_outputs, list):
engine_outputs = [engine_outputs]
@@ -446,7 +439,9 @@ class AsyncOmni(OmniBase):
next_stage_id = stage_id + 1
if next_stage_id <= final_stage_id_for_e2e:
next_stage: OmniStage = self.stage_list[next_stage_id]
next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
# Derive inputs for the next stage, record postprocess time
with metrics.stage_postprocess_timer(stage_id, request_id):
next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
sp_next: SamplingParams = sampling_params_list[next_stage_id]
# Check if we have a connector for this edge
@@ -487,7 +482,7 @@ class AsyncOmni(OmniBase):
result: dict[str, Any],
stage: OmniStage,
stage_id: int,
metrics: OrchestratorMetrics,
metrics: OrchestratorAggregator,
) -> tuple[Any, bool, OmniRequestOutput | None]:
"""
Process a single result dictionary from a stage.
@@ -513,9 +508,16 @@ class AsyncOmni(OmniBase):
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
try:
_m = asdict(result.get("metrics"))
if _m is not None and finished:
metrics.on_stage_metrics(stage_id, req_id, _m)
_m: StageRequestStats = result.get("metrics")
if _m is not None:
# Accumulate generation time
metrics.accumulated_gen_time_ms[req_id][stage_id] += _m.stage_gen_time_ms
# For diffusion stages, we also accumulate diffusion time
metrics.accumulate_diffusion_metrics(stage.stage_type, req_id, engine_outputs)
if finished:
metrics.on_stage_metrics(stage_id, req_id, _m, stage.final_output_type)
except Exception as e:
logger.exception(
f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
vllm_omni/entrypoints/client_request_state.py
+2 -2
View File
@@ -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 @@ class ClientRequestState:
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
vllm_omni/entrypoints/log_utils.py
-591
View File
@@ -1,591 +0,0 @@
from __future__ import annotations
import time
from dataclasses import dataclass
from pprint import pformat
from typing import Any
from vllm.logger import init_logger
logger = init_logger(__name__)
def log_transfer_tx(
from_stage: int,
to_stage: int,
request_id: Any,
size_bytes: int,
tx_time_ms: float,
used_shm: bool,
) -> None:
logger.info(
pformat(
{
"type": "transfer_stats",
"from_stage": from_stage,
"to_stage": to_stage,
"request_id": request_id,
"size_bytes": int(size_bytes),
"tx_time_ms": float(tx_time_ms),
"tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
"used_shm": bool(used_shm),
},
sort_dicts=False,
)
)
def log_transfer_rx(
from_stage: int,
to_stage: int,
request_id: Any,
rx_bytes: int,
rx_decode_time_ms: float,
in_flight_time_ms: float,
) -> None:
logger.info(
pformat(
{
"type": "transfer_rx_stats",
"from_stage": from_stage,
"to_stage": to_stage,
"request_id": request_id,
"rx_bytes": int(rx_bytes),
"rx_decode_time_ms": float(rx_decode_time_ms),
"in_flight_time_ms": float(in_flight_time_ms),
"rx_time_per_kb_ms": (
(float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0
),
},
sort_dicts=False,
)
)
def log_transfer_total(
from_stage: int,
to_stage: int,
request_id: Any,
size_bytes: int,
tx_time_ms: float,
in_flight_time_ms: float,
rx_decode_time_ms: float,
total_time_ms: float,
) -> None:
logger.info(
pformat(
{
"type": "transfer_total_stats",
"from_stage": from_stage,
"to_stage": to_stage,
"request_id": request_id,
"size_bytes": int(size_bytes),
"tx_time_ms": float(tx_time_ms),
"in_flight_time_ms": float(in_flight_time_ms),
"rx_decode_time_ms": float(rx_decode_time_ms),
"total_time_ms": float(total_time_ms),
"total_time_per_kb_ms": (
float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0
),
},
sort_dicts=False,
)
)
def log_stage_request_stats(
stage_id: int,
request_id: Any,
batch_size: int,
num_tokens_out: int,
stage_gen_time_ms: float,
tokens_per_s: float,
rx_transfer_bytes: int,
rx_decode_time_ms: float,
rx_mbps: float,
) -> None:
logger.info(
pformat(
{
"type": "Request_stage_stats",
"stage_id": stage_id,
"request_id": request_id,
"batch_size": int(batch_size),
"num_tokens_out": int(num_tokens_out),
"stage_gen_time_ms": float(stage_gen_time_ms),
"tokens_per_s": float(tokens_per_s),
"rx_transfer_bytes": int(rx_transfer_bytes),
"rx_decode_time_ms": float(rx_decode_time_ms),
"rx_mbps": float(rx_mbps),
},
sort_dicts=False,
)
)
def compute_and_log_stage_request_stats(
stage_id: int,
request_id: Any,
batch_size: int,
num_engine_outputs: int,
stage_gen_time_ms: float,
rx_transfer_bytes: int,
rx_decode_time_ms: float,
) -> None:
"""Compute per-request metrics and log them in one call."""
tokens_per_s = (num_engine_outputs * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
rx_mbps = (
(float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0)
if rx_transfer_bytes > 0
else 0.0
)
log_stage_request_stats(
stage_id,
request_id,
int(batch_size),
int(num_engine_outputs),
float(stage_gen_time_ms),
float(tokens_per_s),
int(rx_transfer_bytes),
float(rx_decode_time_ms),
float(rx_mbps),
)
# ----------------- Aggregation helpers for orchestrator -----------------
def record_stage_metrics(
per_request: dict[str, dict[str, Any]],
stage_req_counts: list[int],
stage_total_time_ms: list[float],
stage_total_tokens: list[int],
stage_id: int,
req_id: Any,
metrics: dict[str, Any],
) -> None:
try:
stage_req_counts[stage_id] += 1
stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
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]
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))
stage_total_tokens[stage_id] += int(metrics.get("num_tokens_in", 0))
pr_stages[stage_id] = stage_data
except Exception:
pass
def aggregate_rx_and_maybe_total(
transfer_edge_req: dict[tuple[int, int, str], dict[str, float]],
transfer_agg: dict[tuple[int, int], dict[str, float]],
per_request: dict[str, dict[str, Any]],
stage_id: int,
req_id: Any,
rx_bytes: float,
rx_ms: float,
in_flight_ms: float,
) -> tuple[int, float, float] | None:
try:
# Update RX aggregates for (stage_id-1 -> stage_id)
if stage_id > 0:
key = (stage_id - 1, stage_id)
agg = transfer_agg.get(key)
if agg is None:
agg = {
"sum_bytes": 0.0,
"sum_ms": 0.0,
"count": 0.0,
"sum_rx_bytes": 0.0,
"sum_rx_ms": 0.0,
"rx_count": 0.0,
"sum_total_ms": 0.0,
"total_count": 0.0,
}
transfer_agg[key] = agg
agg["sum_rx_bytes"] += float(rx_bytes)
agg["sum_rx_ms"] += float(rx_ms)
agg["rx_count"] += 1.0
# Try combine with sender-side timing if present
rid_key = str(req_id)
s = transfer_edge_req.get((stage_id - 1, stage_id, rid_key))
if s is None:
return None
tx_ms = float(s.get("tx_ms", 0.0))
size_b = float(s.get("size_bytes", rx_bytes))
total_ms = tx_ms + float(in_flight_ms) + float(rx_ms)
agg["sum_total_ms"] += total_ms
agg["total_count"] += 1.0
# accumulate per-request transfer totals
try:
pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index]
pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes) # type: ignore[index]
except Exception:
pass
return int(size_b), float(tx_ms), float(total_ms)
return None
except Exception:
return None
def record_sender_transfer_agg(
transfer_agg: dict[tuple[int, int], dict[str, float]],
transfer_edge_req: dict[tuple[int, int, str], dict[str, float]],
from_stage: int,
to_stage: int,
req_id: Any,
size_bytes: int,
tx_ms: float,
) -> None:
try:
key = (from_stage, to_stage)
agg = transfer_agg.get(key)
if agg is None:
agg = {
"sum_bytes": 0.0,
"sum_ms": 0.0,
"count": 0.0,
"sum_rx_bytes": 0.0,
"sum_rx_ms": 0.0,
"rx_count": 0.0,
"sum_total_ms": 0.0,
"total_count": 0.0,
}
transfer_agg[key] = agg
agg["sum_bytes"] += float(size_bytes)
agg["sum_ms"] += float(tx_ms)
agg["count"] += 1.0
# Store sender-side timing for per-request combination
rid_key = str(req_id)
transfer_edge_req[(from_stage, to_stage, rid_key)] = {
"tx_ms": float(tx_ms),
"size_bytes": float(size_bytes),
}
except Exception:
pass
def count_tokens_from_outputs(engine_outputs: list[Any]) -> int:
total = 0
for _ro in engine_outputs:
try:
outs = getattr(_ro, "outputs", None)
if outs and len(outs) > 0:
tokens = getattr(outs[0], "token_ids", None)
if tokens is not None:
total += len(tokens)
except Exception:
pass
return total
def build_stage_summary(
stage_req_counts: list[int],
stage_total_tokens: list[int],
stage_total_time_ms: list[float],
) -> list[dict[str, Any]]:
summary: list[dict[str, Any]] = []
for sid in range(len(stage_req_counts)):
reqs = stage_req_counts[sid]
tokens = stage_total_tokens[sid]
total_ms = float(stage_total_time_ms[sid])
avg_req = (total_ms / reqs) if reqs > 0 else 0.0
avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
summary.append(
{
"stage_id": sid,
"requests": int(reqs),
"tokens": int(tokens),
"total_time_ms": total_ms,
"avg_time_per_request_ms": avg_req,
"avg_tokens_per_s": avg_tok,
}
)
return summary
def build_transfer_summary(
transfer_agg: dict[tuple[int, int], dict[str, float]],
) -> list[dict[str, Any]]:
summary: list[dict[str, Any]] = []
for (src, dst), agg in transfer_agg.items():
sum_bytes = float(agg.get("sum_bytes", 0.0))
sum_ms = float(agg.get("sum_ms", 0.0))
samples = int(agg.get("count", 0.0))
tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0))
sum_rx_ms = float(agg.get("sum_rx_ms", 0.0))
samples_rx = int(agg.get("rx_count", 0.0))
rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0
sum_total_ms = float(agg.get("sum_total_ms", 0.0))
samples_total = int(agg.get("total_count", 0.0))
total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
summary.append(
{
"from_stage": src,
"to_stage": dst,
"samples": samples,
"total_bytes": int(sum_bytes),
"total_time_ms": sum_ms,
"tx_mbps": tx_mbps,
"rx_samples": samples_rx,
"rx_total_bytes": int(sum_rx_bytes),
"rx_total_time_ms": sum_rx_ms,
"rx_mbps": rx_mbps,
"total_samples": samples_total,
"total_transfer_time_ms": sum_total_ms,
"total_mbps": total_mbps,
}
)
return summary
@dataclass
class StageStats:
total_token: int
total_gen_time: float
@dataclass
class StageRequestMetrics:
num_tokens_in: int
num_tokens_out: int
stage_gen_time_ms: float
batch_id: int
batch_size: int
rx_decode_time_ms: float
rx_transfer_bytes: int
rx_in_flight_time_ms: float
stage_stats: StageStats
class OrchestratorMetrics:
def __init__(
self,
num_stages: int,
enable_stats: bool,
wall_start_ts: float,
) -> None:
self.num_stages = int(num_stages)
self.enable_stats = bool(enable_stats)
self.stage_total_time_ms: list[float] = [0.0 for _ in range(self.num_stages)]
self.stage_total_tokens: list[int] = [0 for _ in range(self.num_stages)]
self.stage_req_counts: list[int] = [0 for _ in range(self.num_stages)]
self.transfer_agg: dict[tuple[int, int], dict[str, float]] = {}
self.transfer_edge_req: dict[tuple[int, int, str], dict[str, float]] = {}
self.e2e_total_ms: float = 0.0
self.e2e_total_tokens: int = 0
self.e2e_count: int = 0
self.e2e_done: set[str] = set()
self.per_request: dict[str, dict[str, Any]] = {}
self.sum_per_request_transfer_ms: float = 0.0
self.wall_start_ts: float = float(wall_start_ts)
self.last_finish_ts: float = float(wall_start_ts)
self.stage_seen_batches: dict[int, set] = {sid: set() for sid in range(self.num_stages)}
self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)]
self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)]
def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: dict[str, Any]) -> None:
record_stage_metrics(
self.per_request,
self.stage_req_counts,
self.stage_total_time_ms,
self.stage_total_tokens,
stage_id,
req_id,
metrics,
)
if self.enable_stats:
compute_and_log_stage_request_stats(
stage_id=stage_id,
request_id=req_id,
batch_size=metrics.get("batch_size"),
num_engine_outputs=metrics.get("num_tokens_out"),
stage_gen_time_ms=metrics.get("stage_gen_time_ms"),
rx_decode_time_ms=metrics.get("rx_decode_time_ms"),
rx_transfer_bytes=metrics.get("rx_transfer_bytes"),
)
if stage_stats := metrics.get("stage_stats", None):
total_token = int(stage_stats.get("total_token"))
total_gen_time = float(stage_stats.get("total_gen_time"))
_avg_tokens_per_s = (total_token * 1000.0 / total_gen_time) if total_gen_time > 0 else 0.0
logger.info(
pformat(
{
"type": "Stage_running_avg",
"stage_id": stage_id,
"total_tokens": total_token,
"total_gen_time_ms": total_gen_time,
"avg_tokens_per_s": _avg_tokens_per_s,
},
sort_dicts=False,
)
)
try:
batch_id_raw = metrics.get("batch_id", None)
if batch_id_raw is not None:
batch_id = int(batch_id_raw)
if batch_id not in self.stage_seen_batches[stage_id]:
self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0))
self.stage_seen_batches[stage_id].add(batch_id)
except Exception:
pass
rx_b = float(metrics.get("rx_transfer_bytes", 0.0))
rx_ms = float(metrics.get("rx_decode_time_ms", 0.0))
in_flight_ms = float(metrics.get("rx_in_flight_time_ms", 0.0))
combined = aggregate_rx_and_maybe_total(
self.transfer_edge_req,
self.transfer_agg,
self.per_request,
stage_id,
req_id,
rx_b,
rx_ms,
in_flight_ms,
)
if self.enable_stats and stage_id > 0:
log_transfer_rx(
stage_id - 1,
stage_id,
req_id,
int(rx_b),
float(rx_ms),
float(in_flight_ms),
)
if combined is not None:
size_b_c, tx_ms_c, total_ms_c = combined
log_transfer_total(
stage_id - 1,
stage_id,
req_id,
int(size_b_c),
float(tx_ms_c),
float(in_flight_ms),
float(rx_ms),
float(total_ms_c),
)
def on_forward(
self,
from_stage: int,
to_stage: int,
req_id: Any,
size_bytes: int,
tx_ms: float,
used_shm: bool,
) -> None:
# Mark first input time for the destination stage if not set
if self.stage_first_ts[to_stage] is None:
self.stage_first_ts[to_stage] = time.time()
if self.enable_stats:
log_transfer_tx(
from_stage,
to_stage,
req_id,
int(size_bytes),
float(tx_ms),
bool(used_shm),
)
record_sender_transfer_agg(
self.transfer_agg,
self.transfer_edge_req,
from_stage,
to_stage,
req_id,
int(size_bytes),
float(tx_ms),
)
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()
# Update last output time for this stage
prev_last = self.stage_last_ts[stage_id]
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
# 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 += total_tokens
self.e2e_count += 1
self.e2e_done.add(rid_key)
per_req_record = {
"type": "request_level_metrics",
"request_id": rid_key,
"e2e_time_ms": e2e_ms,
"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": stages_info,
}
self.sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0))
logger.info(pformat(per_req_record, sort_dicts=False))
def build_and_log_summary(self, final_stage_id_to_prompt: dict[str, int]) -> dict[str, Any]:
# Compute stage summary using wall time between first input and last output per stage
stage_summary: list[dict[str, Any]] = []
for sid in range(self.num_stages):
first_ts = self.stage_first_ts[sid]
last_ts = self.stage_last_ts[sid]
total_ms = (
(max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
)
reqs = self.stage_req_counts[sid]
tokens = self.stage_total_tokens[sid]
avg_req = (total_ms / reqs) if reqs > 0 else 0.0
avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
stage_summary.append(
{
"stage_id": sid,
"requests": int(reqs),
"tokens": int(tokens),
"total_time_ms": float(total_ms),
"avg_time_per_request_ms": float(avg_req),
"avg_tokens_per_s": float(avg_tok),
}
)
transfer_summary = build_transfer_summary(self.transfer_agg)
e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0
wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
summary: dict[str, Any] = {
"e2e_requests": int(self.e2e_count),
"e2e_total_time_ms": float(wall_time_ms),
"e2e_sum_time_ms": float(self.e2e_total_ms),
"e2e_total_tokens": int(self.e2e_total_tokens),
"e2e_avg_time_per_request_ms": e2e_avg_req,
"e2e_avg_tokens_per_s": e2e_avg_tok,
"wall_time_ms": wall_time_ms,
"final_stage_id": final_stage_id_to_prompt,
"stages": stage_summary,
"transfers": transfer_summary,
}
return summary
vllm_omni/entrypoints/omni.py
+42 -16
View File
@@ -8,8 +8,6 @@ import uuid
import weakref
from collections.abc import Callable, Generator, Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import asdict
from pprint import pformat
from typing import Any, Literal, overload
from omegaconf import OmegaConf
@@ -30,7 +28,6 @@ from vllm_omni.distributed.ray_utils.utils import (
get_ray_queue_class,
try_close_ray,
)
from vllm_omni.entrypoints.log_utils import OrchestratorMetrics
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
@@ -42,6 +39,7 @@ from vllm_omni.entrypoints.utils import (
resolve_model_config_path,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams
from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -240,7 +238,7 @@ class OmniBase:
)
# Initialize stats paths
self._enable_stats: bool = bool(log_stats)
self.log_stats: bool = bool(log_stats)
self.worker_backend = worker_backend
self.ray_address = ray_address
@@ -670,10 +668,11 @@ class Omni(OmniBase):
final_stage_id_to_prompt[rid] = final_stage_id_for_e2e
# Metrics/aggregation helper
metrics = OrchestratorMetrics(
metrics = OrchestratorAggregator(
num_stages,
self._enable_stats,
self.log_stats,
_wall_start_ts,
final_stage_id_to_prompt,
)
it = request_id_to_prompt.items()
@@ -740,11 +739,15 @@ 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 = result.get("metrics")
_m: StageRequestStats = 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)
# Accumulate generation time
metrics.accumulated_gen_time_ms[req_id][stage_id] += _m.stage_gen_time_ms
# For diffusion stages, we also accumulate diffusion time
metrics.accumulate_diffusion_metrics(stage.stage_type, req_id, engine_outputs)
metrics.on_stage_metrics(stage_id, req_id, _m, stage.final_output_type)
if pbar:
elapsed = pbar.format_dict["elapsed"] or 1e-6
# Aggregate total tokens/images across all stages
@@ -781,8 +784,7 @@ class Omni(OmniBase):
# End-to-end timing and time-per-token for final output
# (only once per request at the designated final stage)
try:
rid_key = str(req_id)
if stage_id == final_stage_id_to_prompt[req_id] and rid_key not in metrics.e2e_done:
if stage_id == final_stage_id_to_prompt[req_id]:
metrics.on_finalize_request(
stage_id,
req_id,
@@ -792,17 +794,42 @@ class Omni(OmniBase):
logger.exception(
f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
)
yield OmniRequestOutput(
output_to_yield = OmniRequestOutput(
stage_id=stage_id,
final_output_type=stage.final_output_type, # type: ignore[attr-defined]
request_output=engine_outputs,
)
# Record audio generated frames with unified signature
try:
finished = (
engine_outputs.finished
if hasattr(engine_outputs, "finished")
else (
engine_outputs[0].finished
if isinstance(engine_outputs, list)
and engine_outputs
and hasattr(engine_outputs[0], "finished")
else False
)
)
metrics.record_audio_generated_frames(output_to_yield, finished, stage_id, req_id)
except Exception as e:
logger.exception(
f"[{self._name}] Failed to record audio metrics for req {req_id} at stage {stage_id}: {e}",
)
yield output_to_yield
next_stage_id = stage_id + 1
if next_stage_id <= final_stage_id_to_prompt[req_id]:
next_stage: OmniStage = self.stage_list[next_stage_id]
try:
next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
# Derive inputs for the next stage, record preprocess time
with metrics.stage_postprocess_timer(stage_id, req_id):
next_inputs = next_stage.process_engine_inputs(
self.stage_list, [request_id_to_prompt[req_id]]
)
except Exception as e:
logger.exception(
f"[{self._name}] Process engine inputs error for req {req_id}"
@@ -856,8 +883,7 @@ class Omni(OmniBase):
# Summarize and print stats
try:
summary = metrics.build_and_log_summary(final_stage_id_to_prompt)
logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
metrics.build_and_log_summary()
except Exception as e:
logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
vllm_omni/entrypoints/omni_llm.py
+1 -1
View File
@@ -81,7 +81,7 @@ class OmniLLM(LLM):
self.worker_backend = kwargs.get("worker_backend", "multi_process")
self.ray_address = kwargs.get("ray_address", None)
self.batch_timeout = batch_timeout
self._enable_stats: bool = bool(log_stats)
self.log_stats: bool = bool(log_stats)
# Load stage configurations
if stage_configs_path is None:
vllm_omni/entrypoints/omni_stage.py
+7 -9
View File
@@ -39,7 +39,6 @@ from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_acto
from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
from vllm_omni.entrypoints.log_utils import count_tokens_from_outputs
from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
from vllm_omni.entrypoints.omni_llm import OmniLLM
from vllm_omni.entrypoints.stage_utils import (
@@ -51,6 +50,7 @@ from vllm_omni.entrypoints.stage_utils import (
set_stage_devices,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams, OmniTokensPrompt
from vllm_omni.metrics import count_tokens_from_outputs
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -844,7 +844,7 @@ def _stage_worker(
try:
sent_ts = float(t.get("sent_ts", None)) if isinstance(t, dict) else None
if sent_ts is not None:
_in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0
_in_flight_ms_by_rid[rid] = max(0.0, (_recv_dequeue_ts - sent_ts) * 1000.0)
else:
_in_flight_ms_by_rid[rid] = 0.0
except Exception:
@@ -1205,7 +1205,7 @@ async def _stage_worker_async(
try:
sent_ts = float(task.get("sent_ts", None)) if isinstance(task, dict) else None
if sent_ts is not None:
_in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0
_in_flight_ms_by_rid[rid] = max(0.0, (_recv_dequeue_ts - sent_ts) * 1000.0)
else:
_in_flight_ms_by_rid[rid] = 0.0
except Exception:
@@ -1384,13 +1384,11 @@ def make_request_stats(
rx_transfer_bytes: int,
rx_in_flight_time_ms: float,
):
from vllm_omni.entrypoints.log_utils import (
StageRequestMetrics,
)
from vllm_omni.metrics import StageRequestStats
num_tokens_in = count_prompt_tokens_from_outputs(req_output)
num_tokens_out = count_tokens_from_outputs(req_output)
return StageRequestMetrics(
return StageRequestStats(
num_tokens_in=num_tokens_in,
num_tokens_out=num_tokens_out,
stage_gen_time_ms=stage_gen_time_ms,
@@ -1404,9 +1402,9 @@ def make_request_stats(
def make_stage_stats(_agg_total_tokens: int, _agg_total_gen_time_ms: float):
from vllm_omni.entrypoints.log_utils import StageStats
from vllm_omni.metrics import StageStats
return StageStats(total_token=_agg_total_tokens, total_gen_time=_agg_total_gen_time_ms)
return StageStats(total_token=_agg_total_tokens, total_gen_time_ms=_agg_total_gen_time_ms)
def output_strip(r_output: RequestOutput | OmniRequestOutput, omni_stage: OmniStage):
vllm_omni/metrics/__init__.py
+9
View File
@@ -0,0 +1,9 @@
from .stats import OrchestratorAggregator, StageRequestStats, StageStats
from .utils import count_tokens_from_outputs
__all__ = [
"OrchestratorAggregator",
"StageStats",
"StageRequestStats",
"count_tokens_from_outputs",
]
vllm_omni/metrics/stats.py
+561
View File
@@ -0,0 +1,561 @@
from __future__ import annotations
import time
from collections import defaultdict
from collections.abc import Callable
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any
from vllm.logger import init_logger
from vllm_omni.metrics.utils import _build_field_defs, _build_row, _format_table
logger = init_logger(__name__)
@dataclass
class StageStats:
total_token: int = 0
total_gen_time_ms: float = 0.0
@property
def avg_tokens_per_s(self) -> float:
return (self.total_token * 1000.0 / self.total_gen_time_ms) if self.total_gen_time_ms > 0 else 0.0
@dataclass
class StageRequestStats:
batch_id: int
batch_size: int
num_tokens_in: int
num_tokens_out: int
stage_gen_time_ms: float
rx_transfer_bytes: int
rx_decode_time_ms: float
rx_in_flight_time_ms: float
stage_stats: StageStats
stage_id: int | None = None
final_output_type: str | None = None
request_id: str | None = None
postprocess_time_ms: float = 0.0
diffusion_metrics: dict[str, int] = None
audio_generated_frames: int = 0
@property
def rx_mbps(self) -> float:
return (
(float(self.rx_transfer_bytes) * 8.0) / (max(float(self.rx_decode_time_ms), 1e-6) * 1000.0)
if self.rx_transfer_bytes > 0
else 0.0
)
@property
def tokens_per_s(self) -> float:
return (self.num_tokens_out * 1000.0 / self.stage_gen_time_ms) if (self.stage_gen_time_ms > 0) else 0.0
@dataclass
class TransferEdgeStats:
from_stage: int
to_stage: int
request_id: str
size_bytes: int
tx_time_ms: float
used_shm: bool = False
rx_decode_time_ms: float = 0.0
in_flight_time_ms: float = 0.0
@property
def total_time_ms(self) -> float:
return float(self.tx_time_ms) + float(self.rx_decode_time_ms) + float(self.in_flight_time_ms)
@dataclass
class RequestE2EStats:
request_id: str
e2e_total_ms: float
e2e_total_tokens: int
transfers_total_time_ms: float
transfers_total_bytes: int
@property
def e2e_tpt(self) -> float:
return (self.e2e_total_ms / self.e2e_total_tokens) if self.e2e_total_tokens > 0 else 0.0
# === Field Configuration ===
# Fields requiring unit conversion: original_field_name -> (display_name, transform_fn)
FIELD_TRANSFORMS: dict[str, tuple[str, Callable[[Any], Any]]] = {
"rx_transfer_bytes": ("rx_transfer_kbytes", lambda v: v / 1024.0),
"size_bytes": ("size_kbytes", lambda v: v / 1024.0),
"transfers_total_bytes": ("transfers_total_kbytes", lambda v: v / 1024.0),
}
# Fields to exclude from table display for each event type
STAGE_EXCLUDE = {
"stage_stats",
"stage_id",
"request_id",
"rx_transfer_bytes",
"rx_decode_time_ms",
"rx_in_flight_time_ms",
"final_output_type",
}
TRANSFER_EXCLUDE = {"from_stage", "to_stage", "request_id", "used_shm"}
E2E_EXCLUDE = {"request_id"}
# Decide the order of overall summary fields, or None for auto
OVERALL_FIELDS: list[str] | None = None
STAGE_FIELDS = _build_field_defs(StageRequestStats, STAGE_EXCLUDE, FIELD_TRANSFORMS)
TRANSFER_FIELDS = _build_field_defs(TransferEdgeStats, TRANSFER_EXCLUDE, FIELD_TRANSFORMS)
E2E_FIELDS = _build_field_defs(RequestE2EStats, E2E_EXCLUDE, FIELD_TRANSFORMS)
class OrchestratorAggregator:
def __init__(
self,
num_stages: int,
log_stats: bool,
wall_start_ts: float,
final_stage_id_for_e2e: dict[str, int] | int,
) -> None:
self.num_stages = int(num_stages)
self.log_stats = bool(log_stats)
self.final_stage_id_for_e2e = final_stage_id_for_e2e
self.init_run_state(wall_start_ts)
self.stage_events: dict[str, list[StageRequestStats]] = {}
self.transfer_events: dict[
tuple[int, int, str], TransferEdgeStats
] = {} # Key: (from_stage, to_stage, request_id)
self.e2e_events: list[RequestE2EStats] = []
def init_run_state(self, wall_start_ts: float) -> None:
# Per-run aggregates and timing state
self.stage_total_tokens = [0 for _ in range(self.num_stages)]
self.e2e_total_ms = 0.0
self.e2e_total_tokens = 0
self.e2e_count = 0
self.e2e_done = set()
self.wall_start_ts = float(wall_start_ts)
self.last_finish_ts = float(wall_start_ts)
self.stage_first_ts = [None for _ in range(self.num_stages)]
self.stage_last_ts = [None for _ in range(self.num_stages)]
self.accumulated_gen_time_ms: defaultdict[str, defaultdict[int, float]] = defaultdict(
lambda: defaultdict(float)
) # {request_id: {stage_id:accumulated_gen_time_ms}}
self.diffusion_metrics: defaultdict[str, defaultdict[str, float]] = defaultdict(
lambda: defaultdict(float)
) # {request_id: {diffusion_metrics_key: accumulated_metrics_data}}
def _get_or_create_transfer_event(
self,
from_stage: int,
to_stage: int,
request_id: str,
) -> TransferEdgeStats:
key = (from_stage, to_stage, request_id)
evt = self.transfer_events.get(key)
if evt is None:
evt = TransferEdgeStats(
from_stage=from_stage,
to_stage=to_stage,
request_id=request_id,
size_bytes=0,
tx_time_ms=0.0,
used_shm=False,
rx_decode_time_ms=0.0,
in_flight_time_ms=0.0,
)
self.transfer_events[key] = evt
return evt
def record_transfer_tx(
self,
from_stage: int,
to_stage: int,
request_id: Any,
size_bytes: int,
tx_time_ms: float,
used_shm: bool,
) -> TransferEdgeStats | None:
try:
evt = self._get_or_create_transfer_event(
int(from_stage),
int(to_stage),
str(request_id),
)
# Accumulate tx metrics
evt.size_bytes += int(size_bytes)
evt.tx_time_ms += float(tx_time_ms)
evt.used_shm = evt.used_shm or bool(used_shm)
return evt
except Exception:
return None
def record_transfer_rx(
self,
stats: StageRequestStats,
) -> TransferEdgeStats | None:
try:
if stats.stage_id is None or stats.stage_id <= 0:
return None
from_stage = int(stats.stage_id) - 1
to_stage = int(stats.stage_id)
rid_key = str(stats.request_id)
evt = self._get_or_create_transfer_event(from_stage, to_stage, rid_key)
# Accumulate rx metrics
if evt.size_bytes == 0:
# size_bytes has been recorded in tx phase
evt.size_bytes = int(stats.rx_transfer_bytes)
evt.rx_decode_time_ms += float(stats.rx_decode_time_ms)
evt.in_flight_time_ms += float(stats.rx_in_flight_time_ms)
return evt
except Exception:
return None
def record_audio_generated_frames(
self,
output_to_yield: Any,
finished: bool,
stage_id: int,
request_id: str,
) -> None:
if (
output_to_yield.final_output_type == "audio"
and finished
and (multimodal_output := output_to_yield.multimodal_output.get("audio")) is not None
):
nframes = int(multimodal_output[-1].shape[0])
stage_events_for_req = self.stage_events.get(request_id, [])
if stage_events_for_req:
for stage_event in stage_events_for_req:
if stage_event.stage_id == stage_id:
stage_event.audio_generated_frames += nframes
break
else:
logger.warning(
"Failed to record audio generated frames for request %s at stage %s: no stage event found",
request_id,
stage_id,
)
def _as_stage_request_stats(
self,
stage_id: int,
req_id: str,
metrics: StageRequestStats,
final_output_type: str | None = None,
) -> StageRequestStats:
"Convert dict to StageRequestStats if needed."
stats = metrics
stats.stage_id = stage_id
stats.request_id = req_id
stats.final_output_type = final_output_type
stats.diffusion_metrics = (
{k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
if req_id in self.diffusion_metrics
else None
)
return stats
def on_stage_metrics(
self,
stage_id: int,
req_id: Any,
metrics: StageRequestStats,
final_output_type: str | None = None,
) -> None:
stats = self._as_stage_request_stats(stage_id, req_id, metrics, final_output_type)
self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_out)
if stats.stage_id == 0:
self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_in)
self.stage_events.setdefault(str(stats.request_id), []).append(stats)
self.record_transfer_rx(stats)
def record_stage_postprocess_time(self, stage_id: int, req_id: Any, postproc_time_ms: float) -> None:
if req_id in self.stage_events:
for stats in self.stage_events[req_id]:
if stats.stage_id == stage_id:
stats.postprocess_time_ms = float(postproc_time_ms)
break
else:
logger.warning(
"Failed to record postprocess time for request %s at stage %s: no stage event found",
req_id,
stage_id,
)
@contextmanager
def stage_postprocess_timer(self, stage_id: int, req_id: Any):
"""Context manager for measuring and recording stage postprocessing time.
Usage:
with metrics.stage_postprocess_timer(stage_id, request_id):
next_inputs = next_stage.process_engine_inputs(...)
"""
_t0 = time.perf_counter()
try:
yield
finally:
_postproc_ms = (time.perf_counter() - _t0) * 1000.0
self.record_stage_postprocess_time(stage_id, req_id, _postproc_ms)
def accumulate_diffusion_metrics(self, stage_type: str, req_id: Any, engine_outputs: Any) -> None:
"""Accumulate diffusion metrics for a request.
Handles extraction and accumulation of diffusion stage metrics.
Args:
req_id: Request ID
engine_outputs: Engine output object containing metrics
"""
if stage_type != "diffusion":
return
engine_output = engine_outputs[0] if isinstance(engine_outputs, list) and engine_outputs else engine_outputs
diffusion_metrics: dict = getattr(engine_output, "metrics", {})
if isinstance(diffusion_metrics, list):
diffusion_metrics = diffusion_metrics[0]
if diffusion_metrics:
for key, value in diffusion_metrics.items():
self.diffusion_metrics[req_id][key] += value
def on_forward(
self,
from_stage: int,
to_stage: int,
req_id: Any,
size_bytes: int,
tx_ms: float,
used_shm: bool,
) -> None:
# Mark first input time for the destination stage if not set
if self.stage_first_ts[to_stage] is None:
self.stage_first_ts[to_stage] = time.time()
self.record_transfer_tx(
from_stage=from_stage,
to_stage=to_stage,
request_id=req_id,
size_bytes=size_bytes,
tx_time_ms=tx_ms,
used_shm=used_shm,
)
def on_finalize_request(
self,
stage_id: int,
req_id: Any,
req_start_ts: float,
) -> None:
rid_key = str(req_id)
if rid_key in self.e2e_done:
return # Already finalized
_t0 = float(req_start_ts)
_t1 = time.time()
# Update last output time for this stage
prev_last = self.stage_last_ts[stage_id]
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
# Sum tokens from all stages for this request
# Include input tokens from stage 0 + output tokens from all stages
total_tokens = 0
if rid_key in self.stage_events:
for evt in self.stage_events[rid_key]:
if evt.stage_id == 0:
total_tokens += int(evt.num_tokens_in)
total_tokens += int(evt.num_tokens_out)
self.e2e_total_ms += e2e_ms
self.e2e_total_tokens += total_tokens
self.e2e_count += 1
self.e2e_done.add(rid_key)
per_req_record = RequestE2EStats(
request_id=rid_key,
e2e_total_ms=e2e_ms,
e2e_total_tokens=total_tokens,
transfers_total_time_ms=float(
sum(evt.total_time_ms for evt in self.transfer_events.values() if evt.request_id == rid_key)
),
transfers_total_bytes=int(
sum(evt.size_bytes for evt in self.transfer_events.values() if evt.request_id == rid_key)
),
)
self.e2e_events.append(per_req_record)
def build_and_log_summary(self) -> dict[str, Any]:
if not self.log_stats:
return {}
wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
e2e_avg_req = (wall_time_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / wall_time_ms) if wall_time_ms > 0 else 0.0
if isinstance(self.final_stage_id_for_e2e, int):
final_stage_id_map: dict[str, int] = {"*": int(self.final_stage_id_for_e2e)}
else:
final_stage_id_map = self.final_stage_id_for_e2e
stage_wall_time_ms = [
((self.stage_last_ts[i] - self.stage_first_ts[i]) * 1000.0)
if (self.stage_first_ts[i] is not None and self.stage_last_ts[i] is not None)
else 0.0
for i in range(self.num_stages)
]
overall_summary = {
"e2e_requests": int(self.e2e_count),
"e2e_wall_time_ms": float(wall_time_ms),
"e2e_total_tokens": int(self.e2e_total_tokens),
"e2e_avg_time_per_request_ms": float(e2e_avg_req),
"e2e_avg_tokens_per_s": float(e2e_avg_tok),
}
# Add stage_wall_time_ms as separate fields for each stage
for idx, wall_time in enumerate(stage_wall_time_ms):
overall_summary[f"e2e_stage_{idx}_wall_time_ms"] = wall_time
# Print overall summary
# filter out all-zero fields for logging
overall_fields = []
for k in OVERALL_FIELDS or list(overall_summary.keys()):
v = overall_summary.get(k, None)
if v not in (0, 0.0, 0.000, None, ""):
overall_fields.append(k)
if overall_fields:
logger.info(
"\n%s",
_format_table("Overall Summary", overall_summary, overall_fields),
)
all_request_ids = sorted(set(self.stage_events.keys()) | {e.request_id for e in self.e2e_events})
result_stage_table = []
result_trans_table = []
result_e2e_table = []
for rid in all_request_ids:
# === E2E table (single column) ===
e2e_evt = next((e for e in self.e2e_events if e.request_id == rid), None)
if e2e_evt:
e2e_data = _build_row(e2e_evt, E2E_FIELDS)
result_e2e_table.append({"request_id": rid, **e2e_data})
# filter out all-zero fields for logging
nonzero_e2e_fields = set()
for k, v in e2e_data.items():
if v not in (0, 0.000, None, ""):
nonzero_e2e_fields.add(k)
value_fields_e2e = sorted(nonzero_e2e_fields)
if value_fields_e2e:
logger.info(
"\n%s",
_format_table(
f"RequestE2EStats [request_id={rid}]",
e2e_data,
value_fields=value_fields_e2e,
),
)
# === Stage table (columns = stage_id) ===
stage_evts = sorted(
self.stage_events.get(rid, []),
key=lambda e: e.stage_id if e.stage_id is not None else -1,
)
# if any stage has diffusion_metrics, remove postprocess_time_ms field
# because it is already included in diffusion_metrics
local_exclude = STAGE_EXCLUDE.copy()
has_diffusion_metrics = any(getattr(evt, "diffusion_metrics", None) for evt in stage_evts)
if has_diffusion_metrics:
local_exclude.add("postprocess_time_ms")
local_stage_fields = _build_field_defs(StageRequestStats, local_exclude, FIELD_TRANSFORMS)
# if diffusion_metrics is present, expand it into multiple columns
# then remove diffusion_metrics from the table
stage_rows = []
for evt in stage_evts:
row = {"stage_id": evt.stage_id, **_build_row(evt, local_stage_fields)}
if evt.diffusion_metrics:
row.update(evt.diffusion_metrics)
row.pop("diffusion_metrics", None) # Remove the dict itself
stage_rows.append(row)
result_stage_table.append({"request_id": rid, "stages": stage_rows})
if stage_rows:
# filter out all-zero fields for logging
all_value_fields = set()
for row in stage_rows:
for k in row.keys():
if k != "stage_id":
all_value_fields.add(k)
value_fields_list = []
for field in sorted(all_value_fields):
all_zero = True
for row in stage_rows:
v = row.get(field, None)
if v not in (0, 0.0, 0.000, None, ""):
all_zero = False
break
if not all_zero:
value_fields_list.append(field)
if value_fields_list:
logger.info(
"\n%s",
_format_table(
f"StageRequestStats [request_id={rid}]",
stage_rows,
column_key="stage_id",
value_fields=value_fields_list,
),
)
# === Transfer table (columns = edge) ===
transfer_evts = sorted(
[e for e in self.transfer_events.values() if e.request_id == rid],
key=lambda e: (e.from_stage, e.to_stage),
)
transfer_rows = [
{"edge": f"{evt.from_stage}->{evt.to_stage}", **_build_row(evt, TRANSFER_FIELDS)}
for evt in transfer_evts
]
result_trans_table.append({"request_id": rid, "transfers": transfer_rows})
if transfer_rows:
# filter out all-zero fields for logging
all_value_fields = set()
for row in transfer_rows:
for k in row.keys():
if k != "edge":
all_value_fields.add(k)
value_fields_list = []
for field in sorted(all_value_fields):
all_zero = True
for row in transfer_rows:
v = row.get(field, None)
if v not in (0, 0.0, 0.000, None, ""):
all_zero = False
break
if not all_zero:
value_fields_list.append(field)
if value_fields_list:
logger.info(
"\n%s",
_format_table(
f"TransferEdgeStats [request_id={rid}]",
transfer_rows,
column_key="edge",
value_fields=value_fields_list,
),
)
return {
"final_stage_id": final_stage_id_map,
"overall_summary": overall_summary,
"stage_table": result_stage_table,
"trans_table": result_trans_table,
"e2e_table": result_e2e_table,
}
vllm_omni/metrics/utils.py
+143
View File
@@ -0,0 +1,143 @@
from collections.abc import Callable
from dataclasses import fields
from typing import Any
from prettytable import PrettyTable
def _build_field_defs(
cls: type,
exclude: set[str],
transforms: dict[str, tuple[str, Callable]] | None = None,
) -> list[tuple[str, Callable[[Any], Any]]]:
"""Auto-generate field definitions from a dataclass.
Args:
cls: The dataclass type to extract fields from.
exclude: Set of field names to exclude from output.
transforms: Optional mapping of field transformations.
Format: {original_name: (display_name, transform_fn)}
Returns:
List of (display_name, getter_fn) tuples for table generation.
"""
transforms = transforms or {}
result = []
for f in fields(cls):
if f.name in exclude:
continue
if f.name in transforms:
display_name, transform_fn = transforms[f.name]
# Capture variables in closure to avoid late binding issues
result.append((display_name, lambda e, fn=transform_fn, n=f.name: fn(getattr(e, n))))
else:
result.append((f.name, lambda e, n=f.name: getattr(e, n)))
return result
def _build_row(evt: Any, field_defs: list[tuple[str, Callable]]) -> dict[str, Any]:
"""Build a row dict from an event object using field definitions.
Args:
evt: The event object (dataclass instance).
field_defs: List of (field_name, getter_fn) tuples.
Returns:
Dict mapping field names to their values.
"""
return {name: getter(evt) for name, getter in field_defs}
def _get_field_names(field_defs: list[tuple[str, Callable]]) -> list[str]:
"""Extract field names from field definitions.
Args:
field_defs: List of (field_name, getter_fn) tuples.
Returns:
List of field names.
"""
return [name for name, _ in field_defs]
def _format_table(
title: str,
data: dict[str, Any] | list[dict[str, Any]],
value_fields: list[str],
column_key: str | None = None,
column_prefix: str = "",
) -> str:
"""Format a table for display.
Supports two modes:
1. Single-column mode: data is a dict, displays as Field | Value
2. Multi-column mode: data is a list of dicts, displays as Field | col1 | col2 | ...
Args:
title: Table title.
data: Either a single dict (single-column) or list of dicts (multi-column).
value_fields: List of field names to display as rows.
column_key: Key in each dict used as column header (required for multi-column mode).
column_prefix: Optional prefix for column headers (multi-column mode only).
Returns:
Formatted table string.
"""
if not data:
return f"[{title}] <empty>"
def _format_value(value: Any) -> str:
if isinstance(value, bool):
return str(value)
if isinstance(value, int):
return f"{value:,}"
if isinstance(value, float):
return f"{value:,.3f}"
if isinstance(value, list):
return ", ".join(str(f"{v:,.3f}") for v in value)
return str(value)
table = PrettyTable()
# Single-column mode: data is a dict
if isinstance(data, dict):
table.field_names = ["Field", "Value"]
table.align["Field"] = "l"
table.align["Value"] = "r"
for field in value_fields:
if field in data:
if isinstance(data[field], dict):
for sub_key, sub_value in data[field].items():
table.add_row([f"{sub_key}", _format_value(sub_value)])
else:
table.add_row([field, _format_value(data[field])])
# Multi-column mode: data is a list of dicts
else:
if column_key is None:
raise ValueError("column_key is required for multi-column mode")
col_headers = [f"{column_prefix}{row.get(column_key, '?')}" for row in data]
table.field_names = ["Field"] + col_headers
table.align["Field"] = "l"
for col in col_headers:
table.align[col] = "r"
for field in value_fields:
row_values = [_format_value(r.get(field, "")) for r in data]
table.add_row([field] + row_values)
return "\n".join([f"[{title}]", table.get_string()])
def count_tokens_from_outputs(engine_outputs: list[Any]) -> int:
total = 0
for _ro in engine_outputs:
try:
outs = getattr(_ro, "outputs", None)
if outs and len(outs) > 0:
tokens = getattr(outs[0], "token_ids", None)
if tokens is not None:
total += len(tokens)
except Exception:
# Ignore any issues with individual outputs to keep token counting best-effort.
pass
return total