diff --git a/benchmarks/qwen3-omni/README.md b/benchmarks/qwen3-omni/README.md index fdc385445..c76fea895 100644 --- a/benchmarks/qwen3-omni/README.md +++ b/benchmarks/qwen3-omni/README.md @@ -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. diff --git a/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh b/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh index 61e46f8c3..935753605 100644 --- a/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh +++ b/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh @@ -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" diff --git a/docs/api/README.md b/docs/api/README.md index 6219aee40..eb3a40950 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -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][] diff --git a/docs/contributing/metrics.md b/docs/contributing/metrics.md new file mode 100644 index 000000000..7a46a2861 --- /dev/null +++ b/docs/contributing/metrics.md @@ -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). diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index 397fd333e..8fa412065 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/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, diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 63b08202a..7cd8e1873 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/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)", diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 0b484c507..263377762 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/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)", diff --git a/examples/offline_inference/qwen3_tts/end2end.py b/examples/offline_inference/qwen3_tts/end2end.py index 18bb9dedc..93aeba3ca 100644 --- a/examples/offline_inference/qwen3_tts/end2end.py +++ b/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)", diff --git a/requirements/common.txt b/requirements/common.txt index 19ebca949..2331e0f5d 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -12,3 +12,4 @@ torchsde>=0.2.6 openai-whisper>=20250625 imageio[ffmpeg]>=2.37.2 sox>=1.5.0 +prettytable>=3.8.0 diff --git a/tests/e2e/online_serving/test_async_omni.py b/tests/e2e/online_serving/test_async_omni.py index ba5084ad4..cab3e6e22 100644 --- a/tests/e2e/online_serving/test_async_omni.py +++ b/tests/e2e/online_serving/test_async_omni.py @@ -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"]) diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index f33ca6d59..4f05575ca 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/tests/entrypoints/test_omni_llm.py @@ -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, ) diff --git a/tests/metrics/test_stats.py b/tests/metrics/test_stats.py new file mode 100644 index 000000000..8b18c5dff --- /dev/null +++ b/tests/metrics/test_stats.py @@ -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 diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 8e19f1242..2ab94d5b7 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -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( diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 4cd225078..66aebddc1 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -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. diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 1545e17a1..6a3d7393a 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -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}", diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 3e68abb17..1c9103f79 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/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 @@ 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 diff --git a/vllm_omni/entrypoints/log_utils.py b/vllm_omni/entrypoints/log_utils.py deleted file mode 100644 index 2031038da..000000000 --- a/vllm_omni/entrypoints/log_utils.py +++ /dev/null @@ -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 diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 97357dc3b..9e49f951e 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -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}") diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 0356ea231..118a0ba7f 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -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: diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 4935c43cc..ca5b2f43e 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -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): diff --git a/vllm_omni/metrics/__init__.py b/vllm_omni/metrics/__init__.py new file mode 100644 index 000000000..deceb2333 --- /dev/null +++ b/vllm_omni/metrics/__init__.py @@ -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", +] diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py new file mode 100644 index 000000000..037db2ea8 --- /dev/null +++ b/vllm_omni/metrics/stats.py @@ -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, + } diff --git a/vllm_omni/metrics/utils.py b/vllm_omni/metrics/utils.py new file mode 100644 index 000000000..ea83d959e --- /dev/null +++ b/vllm_omni/metrics/utils.py @@ -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}] " + + 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