[Core][Metrics] Leave unobserved diffusion timings out of per-request metrics

docs/design/metrics.md

@@ -312,3 +312,21 @@ The per-request `[StageRequestStats]` table is logged at `DEBUG` (for example `V
 Missing and zero are different. A field is `None` when its interval was not observed: a diffusion stage, `--log-stats` off, or an engine-core event that never arrived. A measured `0` stays `0`. In the `[StageRequestStats]` table a missing value prints as `None`, and, as for every other field, a row whose values are all zero or missing is hidden.
 
 For a streaming-input request, only the output that finishes the request carries the split. On an AR stage, vLLM's output processor finishes the request only after the input stream has ended; at each earlier segment stop it applies the next streaming update instead, which restarts the request stats, so the split covers the last input segment. The last segment's `QUEUED` event can arrive with the previous segment's final output, before that restart, so its `vllm_queued_ms` may be missing. A generation stage finishes on the first output that is not a segment stop and does not restart its stats between segments, so its split spans the whole request: `vllm_prefill_ms` is the first segment's, `vllm_decode_ms` runs from the first segment's first output to the last output, gaps between segments included, and `vllm_queued_ms` is missing when a later segment was queued after the first scheduling. The fields are not added to the `stage_metrics` snapshot returned to clients and benchmarks. The duplex per-response table leaves them out (`DUPLEX_STAGE_TABLE_EXCLUDE`): they describe one engine-core request, not a response turn.
+
+### Diffusion-stage timings in `StageRequestStats`
+
+`DiffusionEngine.step_streaming` attaches its timings to each output it formats, in milliseconds. When a final-output diffusion stage finishes, `OrchestratorAggregator.accumulate_diffusion_metrics` sums them per request, converts them to seconds, and stores them in the stage's `StageRequestStats.diffusion_metrics` under the same name with `_ms` replaced by `_s` (`preprocess_time_ms` becomes `preprocess_time_s`, `scheduler_queue_wait_ms` becomes `scheduler_queue_wait_s`). The `[StageRequestStats]` table shows one row per key; as above, a row whose values are all zero or missing is hidden.
+
+| Engine key | Interval | Left out when |
+| ---------- | -------- | ------------- |
+| `scheduler_queue_wait_ms` | Diffusion scheduler: `add_request` to the first scheduling. Later re-queues are not added. With the step scheduler it also covers waiting for running requests to reach a compatible step phase, for requests that do not allow mixed phases | The request was never scheduled |
+| `preprocess_time_ms` | Preparing the request for admission: the model's pre-processor, plus prefix-cache hashing when enabled | The model has no pre-processor |
+| `diffusion_engine_exec_time_ms` | Adding the request to the scheduler until the engine receives this output. Contains the scheduler wait, the stage-payload and KV receive, the forward pass and the VAE decode | Never |
+| `output_ready_wait_time_ms` | Waiting for an asynchronous device-to-host or shared-memory output to become ready | The output was returned without an asynchronous wait |
+| `postprocess_time_ms` | Engine-side output post-processing: the model's post-processor (`post_process_func`, when set), moving the output to CPU under model-level CPU offload, and formatting the request outputs. This is not the top-level `StageRequestStats.postprocess_time_ms`, which the table hides when a stage of the request has diffusion timings | Never |
+| `forward_time_ms`, `vae_decode_time_ms` | Pipeline-profiler time of the `.diffuse` targets and of `decode` / `decode_latent` on a VAE module. The profiler synchronizes the device around each target | The profiler is off (the default; `--enable-diffusion-pipeline-profiler` turns it on) or no target matched |
+| `kv_recv_time_ms` | Diffusion worker: start of the AR KV receive until the KV is ready, after the stage payload is received. It includes waiting for the sender. With KV prefetch it times consuming the prefetched KV, or the synchronous receive on a prefetch miss | No KV transfer manager, or `0` |
+
+Missing and zero are different here too: an interval that was not observed is left out, not written as `0`. A key that only some stages report leaves the other cells blank in the `[StageRequestStats]` table. One exception keeps an existing Prometheus population unchanged: for a request without a pre-processor, `vllm_omni:diffusion_preprocess_s` still records `0`.
+
+These timings exist only for final-output diffusion stages. A non-final stage reports through `StageMetricsMessage`, which carries no engine output, so its `diffusion_metrics` is `None`. Only the output that finishes the stage is counted, so earlier outputs of a streaming stage are not added: `diffusion_engine_exec_time_ms` already runs from adding the request, while `output_ready_wait_time_ms` and `postprocess_time_ms` cover that output only. The rows overlap: `diffusion_engine_exec_time_ms` contains the scheduler wait, the KV receive and the profiler timings, so do not add the rows together. Like the LLM-stage split, these rows are not added to the `stage_metrics` snapshot returned to clients and benchmarks.

tests/diffusion/test_diffusion_engine_metrics.py

@@ -152,7 +152,8 @@ def test_timing_metric_key_naming_consistency(self) -> None:
                 raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()")
 
         assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source
-        assert '"output_ready_wait_time_ms": output_ready_wait_time * 1000' in step_streaming_source
+        # Written only when observed; see test_diffusion_step_streaming_metrics.py.
+        assert 'metrics_update["output_ready_wait_time_ms"] = output_ready_wait_time * 1000' in step_streaming_source
         # step_total_ms is no longer emitted as a metric (no family consumes
         # it); it lives only in the debug log breakdown below.
         assert '"diffusion_engine_total_time_ms"' not in step_streaming_source

tests/diffusion/test_diffusion_step_streaming_metrics.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Per-request timing keys that ``DiffusionEngine.step_streaming`` attaches.
+
+An interval that was not observed is left out of the output metrics instead of
+being written as ``0.0``; a measured ``0`` stays. The tests drive
+``step_streaming`` on CPU with a stubbed scheduler stream and executor, then
+follow the keys through ``OrchestratorAggregator`` into the per-request
+``StageRequestStats`` record and the ``stage_table`` rows that
+``build_and_log_summary`` returns.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import logging
+import threading
+
+import pytest
+
+from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.metrics import OrchestratorAggregator
+from vllm_omni.metrics.stats import StageRequestStats, StageStats
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_ALWAYS_SET = {"diffusion_engine_exec_time_ms", "postprocess_time_ms"}
+_OPTIONAL = {
+    "preprocess_time_ms",
+    "output_ready_wait_time_ms",
+    "scheduler_queue_wait_ms",
+    "forward_time_ms",
+    "vae_decode_time_ms",
+    "kv_recv_time_ms",
+}
+
+
+def _make_request(*, scheduler_queue_wait_ms: float | None = None) -> OmniDiffusionRequest:
+    request = OmniDiffusionRequest(
+        prompt="prompt",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=4),
+        request_id="req-diffusion",
+    )
+    # Set by the diffusion scheduler at the first admission.
+    request.scheduler_queue_wait_ms = scheduler_queue_wait_ms
+    return request
+
+
+class _ReadyOutputExecutor:
+    """Resolves async outputs the way ``DiffusionExecutor.wait_output_ready``
+    does, from outputs that are already available."""
+
+    def __init__(self, ready_outputs: dict[str, DiffusionOutput]) -> None:
+        self._ready_outputs = ready_outputs
+
+    def wait_output_ready(self, async_output_id: str) -> concurrent.futures.Future[DiffusionOutput]:
+        future: concurrent.futures.Future[DiffusionOutput] = concurrent.futures.Future()
+        future.set_result(self._ready_outputs[async_output_id])
+        return future
+
+
+def _make_engine(
+    outputs: list[DiffusionOutput],
+    *,
+    pre_process_func=None,
+    ready_outputs: dict[str, DiffusionOutput] | None = None,
+) -> DiffusionEngine:
+    """An engine whose scheduler delivers ``outputs`` to the real output
+    stream at admission and whose executor resolves async outputs from
+    ``ready_outputs``. The stream materializes async outputs and records
+    ``output_ready_wait_time`` before ``step_streaming`` sees them."""
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.pre_process_func = pre_process_func
+    engine._scheduler_num_waiting_reqs = 0
+    engine.main_loop = None
+    engine._cv = threading.Condition()
+    engine._out_streams = {}
+    engine._unclaimed_async_outputs = {}
+    engine._shutdown_output_futures = {}
+    engine._streaming_playback = None
+
+    async def _no_background_loop() -> None:
+        return None
+
+    def _admit(request: OmniDiffusionRequest) -> str:
+        engine._out_streams[request.request_id] = asyncio.Queue()
+        for output in outputs:
+            engine._put_output(request.request_id, output)
+        return request.request_id
+
+    engine._check_and_start_background_loop = _no_background_loop
+    engine._add_prepared_request = _admit
+    engine.postprocess_output = lambda request, output: [OmniRequestOutput(request_id=request.request_id)]
+    engine.executor = _ReadyOutputExecutor(ready_outputs or {})
+    return engine
+
+
+async def _collect_metrics(engine: DiffusionEngine, request: OmniDiffusionRequest) -> list[dict]:
+    return [out.metrics async for batch in engine.step_streaming(request) for out in batch]
+
+
+@pytest.mark.asyncio
+async def test_unobserved_intervals_are_left_out() -> None:
+    """No pre-processor, a synchronous output, profiler off, no KV, and no
+    scheduler stamp: only the always-measured keys are written."""
+    engine = _make_engine([DiffusionOutput()])
+
+    (metrics,) = await _collect_metrics(engine, _make_request())
+
+    assert _ALWAYS_SET <= metrics.keys()
+    assert not (_OPTIONAL & metrics.keys())
+    assert metrics["diffusion_engine_exec_time_ms"] >= 0.0
+    assert metrics["postprocess_time_ms"] >= 0.0
+
+
+@pytest.mark.asyncio
+async def test_observed_intervals_are_reported() -> None:
+    def _pre_process(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        return request
+
+    ready = DiffusionOutput(
+        stage_durations={"DemoPipeline.diffuse": 0.5, "DemoPipeline.vae.decode": 0.25},
+        kv_recv_ms=12.0,
+    )
+    engine = _make_engine(
+        [DiffusionOutput(async_output_id="out-1")],
+        pre_process_func=_pre_process,
+        ready_outputs={"out-1": ready},
+    )
+
+    (metrics,) = await _collect_metrics(engine, _make_request(scheduler_queue_wait_ms=3.0))
+
+    assert (_ALWAYS_SET | _OPTIONAL) <= metrics.keys()
+    assert metrics["preprocess_time_ms"] >= 0.0
+    assert metrics["output_ready_wait_time_ms"] >= 0.0
+    assert ready.output_ready_wait_time is not None
+    assert metrics["scheduler_queue_wait_ms"] == pytest.approx(3.0)
+    # Profiler and KV values come from the output the async wait returned.
+    assert metrics["forward_time_ms"] == pytest.approx(500.0)
+    assert metrics["vae_decode_time_ms"] == pytest.approx(250.0)
+    assert metrics["kv_recv_time_ms"] == pytest.approx(12.0)
+
+
+@pytest.mark.asyncio
+async def test_output_without_async_wait_leaves_wait_unset() -> None:
+    """The stream sets ``output_ready_wait_time`` only when it waited for an
+    async output, so a synchronous output does not report a 0 ms wait."""
+    output = DiffusionOutput()
+    engine = _make_engine([output])
+
+    (metrics,) = await _collect_metrics(engine, _make_request())
+
+    assert output.output_ready_wait_time is None
+    assert "output_ready_wait_time_ms" not in metrics
+
+
+@pytest.mark.asyncio
+async def test_measured_zero_queue_wait_is_kept() -> None:
+    engine = _make_engine([DiffusionOutput()])
+
+    (metrics,) = await _collect_metrics(engine, _make_request(scheduler_queue_wait_ms=0.0))
+
+    assert metrics["scheduler_queue_wait_ms"] == 0.0
+
+
+@pytest.mark.asyncio
+async def test_output_ready_wait_is_per_output() -> None:
+    """A streamed request can mix synchronous and async outputs; each output
+    reports only the wait it had."""
+    engine = _make_engine(
+        [DiffusionOutput(finished=False), DiffusionOutput(async_output_id="out-2")],
+        ready_outputs={"out-2": DiffusionOutput()},
+    )
+
+    first, second = await _collect_metrics(engine, _make_request())
+
+    assert "output_ready_wait_time_ms" not in first
+    assert second["output_ready_wait_time_ms"] >= 0.0
+
+
+@pytest.mark.asyncio
+async def test_output_ready_wait_does_not_carry_to_the_next_output() -> None:
+    """A synchronous output after an async one does not inherit its wait."""
+    engine = _make_engine(
+        [DiffusionOutput(async_output_id="out-1", finished=False), DiffusionOutput()],
+        ready_outputs={"out-1": DiffusionOutput(finished=False)},
+    )
+
+    first, second = await _collect_metrics(engine, _make_request())
+
+    assert first["output_ready_wait_time_ms"] >= 0.0
+    assert "output_ready_wait_time_ms" not in second
+
+
+@pytest.mark.asyncio
+async def test_debug_breakdown_formats_unobserved_intervals() -> None:
+    records: list[str] = []
+
+    class _Collect(logging.Handler):
+        def emit(self, record: logging.LogRecord) -> None:
+            records.append(record.getMessage())
+
+    logger = diffusion_engine_module.logger
+    handler = _Collect(level=logging.DEBUG)
+    previous_level = logger.level
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)
+    try:
+        await _collect_metrics(_make_engine([DiffusionOutput()]), _make_request())
+    finally:
+        logger.removeHandler(handler)
+        logger.setLevel(previous_level)
+
+    breakdown = [msg for msg in records if "step_streaming breakdown" in msg]
+    assert len(breakdown) == 1
+    assert "preprocess=0.00 ms" in breakdown[0]
+    assert "output_ready_wait=0.00 ms" in breakdown[0]
+
+
+def _stage_stats() -> StageRequestStats:
+    return StageRequestStats(
+        batch_id=1,
+        batch_size=1,
+        num_tokens_in=0,
+        num_tokens_out=0,
+        stage_gen_time_ms=40.0,
+        rx_transfer_bytes=0,
+        rx_decode_time_ms=0.0,
+        rx_in_flight_time_ms=0.0,
+        stage_stats=StageStats(),
+    )
+
+
+@pytest.mark.asyncio
+async def test_unobserved_intervals_stay_out_of_the_stage_record() -> None:
+    """Follow the engine keys through the aggregator into the per-request
+    record and the ``stage_table`` rows that ``build_and_log_summary``
+    returns."""
+    engine = _make_engine([DiffusionOutput()])
+    request = _make_request(scheduler_queue_wait_ms=2.0)
+    outputs = [out async for batch in engine.step_streaming(request) for out in batch]
+
+    aggregator = OrchestratorAggregator(num_stages=1, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=0)
+    aggregator.accumulate_diffusion_metrics("diffusion", "req-diffusion", outputs)
+    aggregator.on_stage_metrics(0, "req-diffusion", _stage_stats(), "image")
+
+    (event,) = aggregator.stage_events["req-diffusion"]
+    assert event.diffusion_metrics is not None
+    assert "preprocess_time_s" not in event.diffusion_metrics
+    assert "output_ready_wait_time_s" not in event.diffusion_metrics
+    assert event.diffusion_metrics["scheduler_queue_wait_s"] == pytest.approx(0.002)
+    assert "diffusion_engine_exec_time_s" in event.diffusion_metrics
+
+    summary = aggregator.build_and_log_summary()
+    (entry,) = summary["stage_table"]
+    (row,) = entry["stages"]
+    assert "preprocess_time_s" not in row
+    assert "output_ready_wait_time_s" not in row
+    assert row["scheduler_queue_wait_s"] == pytest.approx(0.002)

tests/metrics/test_modality.py

@@ -16,6 +16,7 @@
     observe_modality_at_finalize,
 )
 from vllm_omni.metrics.stats import OrchestratorAggregator
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -578,6 +579,27 @@ def test_accumulator_converts_all_engine_milliseconds_to_seconds(self):
             "kv_recv_time_s": pytest.approx(0.04),
         }
 
+    def test_accumulator_leaves_unobserved_keys_absent(self):
+        """Keys the engine left out stay out of the per-request record; they
+        are not created as 0.0, and finished outputs are summed per key."""
+        aggregator = OrchestratorAggregator(
+            num_stages=1,
+            log_stats=True,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=0,
+        )
+        for _ in range(2):
+            engine_output = OmniRequestOutput(
+                request_id="req-sparse",
+                metrics={"diffusion_engine_exec_time_ms": 1000.0, "postprocess_time_ms": 10.0},
+            )
+            aggregator.accumulate_diffusion_metrics("diffusion", "req-sparse", engine_output)
+
+        assert dict(aggregator.diffusion_metrics["req-sparse"]) == {
+            "diffusion_engine_exec_time_s": pytest.approx(2.0),
+            "postprocess_time_s": pytest.approx(0.02),
+        }
+
 
 class TestObserveDiffusionFinalize:
     @pytest.mark.parametrize(
@@ -600,6 +622,43 @@ def test_emits_optional_breakdown_metric(self, metric_key, call_name, value):
         )
 
         assert (call_name, "2", "0", value) in stub.calls
+        # Without engine timings there is no preprocess sample, as before.
+        assert not any(c[0] == "observe_diffusion_preprocess" for c in stub.calls)
+
+    def test_preprocess_histogram_keeps_zero_sample_without_preprocessor(self):
+        """The engine leaves preprocess out when no pre-processor ran; the
+        histogram still records the 0 it recorded before (RFC #8671 leaves
+        that sample to the metrics owners)."""
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="image",
+            stage_id=2,
+            replica_id=0,
+            stage_metrics=_Bag(
+                diffusion_metrics={"diffusion_engine_exec_time_s": 1.2, "postprocess_time_s": 0.3},
+            ),
+            engine_outputs=_Bag(),
+        )
+
+        preprocess_calls = [c for c in stub.calls if c[0] == "observe_diffusion_preprocess"]
+        assert preprocess_calls == [("observe_diffusion_preprocess", "2", "0", 0.0)]
+
+    def test_measured_preprocess_is_observed_once(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="image",
+            stage_id=2,
+            replica_id=0,
+            stage_metrics=_Bag(
+                diffusion_metrics={"diffusion_engine_exec_time_s": 1.2, "preprocess_time_s": 0.05},
+            ),
+            engine_outputs=_Bag(),
+        )
+
+        preprocess_calls = [c for c in stub.calls if c[0] == "observe_diffusion_preprocess"]
+        assert preprocess_calls == [("observe_diffusion_preprocess", "2", "0", 0.05)]
 
     def test_audio_diffusion_stage_fires_both_paths(self):
         stub = _StubModMetrics()

vllm_omni/diffusion/data.py

@@ -1898,7 +1898,8 @@ class DiffusionOutput:
     stage_durations: dict[str, float] = field(default_factory=dict)
 
     # Engine-side async materialization wait (seconds), separate from model stages.
-    output_ready_wait_time: float = 0.0
+    # None when the output needed no async wait.
+    output_ready_wait_time: float | None = None
 
     # memory usage info
     peak_memory_mb: float = 0.0