From 7ae15c74c8db75861478c989447639979e1fe617 Mon Sep 17 00:00:00 2001 From: Honghan Zhu Date: Wed, 26 Aug 2026 10:54:45 +0800 Subject: [PATCH] text and image metric (#6150) Signed-off-by: zhuhh97 Signed-off-by: Gao Han Co-authored-by: Claude Opus 4.7 Co-authored-by: Gao Han --- docs/design/metrics.md | 106 ++-- .../test_diffusion_engine_metrics.py | 136 ++++- tests/diffusion/test_diffusion_scheduler.py | 13 + .../test_multiproc_engine_concurrency.py | 38 +- tests/engine/test_membership_controller.py | 24 +- tests/engine/test_orchestrator.py | 25 + tests/entrypoints/test_omni_entrypoints.py | 14 + tests/metrics/test_emit_calls.py | 499 ++++++++++++++++++ tests/metrics/test_modality.py | 208 ++++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 78 ++- vllm_omni/core/sched/omni_scheduler_mixin.py | 4 +- vllm_omni/diffusion/data.py | 3 + vllm_omni/diffusion/diffusion_engine.py | 46 +- .../inline_stage_diffusion_client.py | 19 +- vllm_omni/diffusion/request.py | 5 + vllm_omni/diffusion/sched/base_scheduler.py | 3 + vllm_omni/diffusion/sched/interface.py | 1 + vllm_omni/diffusion/stage_diffusion_client.py | 26 +- vllm_omni/diffusion/stage_diffusion_proc.py | 20 +- .../worker/diffusion_model_runner.py | 5 + vllm_omni/engine/async_omni_engine.py | 4 + vllm_omni/engine/membership_controller.py | 9 +- vllm_omni/engine/orchestrator.py | 61 +++ vllm_omni/engine/stage_pool.py | 16 +- vllm_omni/entrypoints/async_omni.py | 5 +- vllm_omni/entrypoints/client_request_state.py | 2 + vllm_omni/entrypoints/omni.py | 3 + vllm_omni/entrypoints/omni_base.py | 64 ++- vllm_omni/metrics/definitions.py | 44 +- vllm_omni/metrics/modality.py | 124 +++-- vllm_omni/metrics/prometheus.py | 135 ++++- vllm_omni/metrics/stats.py | 64 ++- vllm_omni/metrics/utils.py | 160 +++++- 33 files changed, 1802 insertions(+), 162 deletions(-) create mode 100644 tests/metrics/test_emit_calls.py diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 7ee35813f..85eb7ac38 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -6,7 +6,7 @@ This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage - Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope). - Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically. -- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss. +- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss, plus image and diffusion workload, queueing, execution, and memory signals. - Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable. - Keep the metrics collection overhead low enough that it does not regress TTFA or throughput. @@ -20,7 +20,7 @@ vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` col ### The Problem -vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`. +vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm_omni:` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm_omni:` / `vllm_omni`. Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below). @@ -38,7 +38,7 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by | +--------+--------+--------+--------+--------+ | | - vllm:omni_* vllm:* + vllm_omni:* vllm:* collectors collectors | | +----+--------+ +-----------+ +----------+ +-----------+ @@ -58,17 +58,30 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by ### Data Flow -There are four independent paths for metric collection. +There are five independent paths for metric collection. -**Path 1: Pipeline-level metrics (`vllm:omni_*`)** +**Path 1: Pipeline-level metrics (`vllm_omni:*`)** `OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress: - `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`. -- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`. -- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion. +- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm_omni:requests_success_total{finished_reason}`. +- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` on the existing completion Counter. +- `inc_requests_failed(reason)` — increments the dedicated failure-attribution Counter once per request. Reasons are normalized to the bounded set `client_abort`, `client_disconnect`, `stage_error`, and `unknown`. -**Path 2: Audio modality metrics (`vllm:omni_audio_*`)** +**Path 2: Image and diffusion metrics** + +Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetrics` owns pipeline/stage workload families, while `OmniModalityMetrics` owns per-replica diffusion breakdowns. Finished stage messages are consumed once in `OmniBase._process_single_result`, which prevents replayed terminal messages from incrementing Counters or Histograms twice. + +- `stage_gen_time_s`, `stage_in_queue_s`, image workload, inference-step, and per-stage peak-memory values are emitted from `StageRequestStats` when a stage finishes. `image_pixels` and `image_count_total` are restricted to image outputs; `num_inference_steps` is restricted to diffusion stages. +- Diffusion engine timings are accumulated per request by `OrchestratorAggregator.accumulate_diffusion_metrics`. Engine-side millisecond values are converted to seconds before `observe_modality_at_finalize(...)` dispatches the per-replica Histogram observations. +- `request_queue_wait_s` is emitted once at final request completion from `pipeline_timings["queue_wait_ms"]`. A present value of `0` is a valid observation; a missing key produces no sample. +- `stage_waiting_requests` is the sum of the latest scheduler snapshots for all live replicas in a stage. AR snapshots arrive through `SchedulerStats`; diffusion snapshots ride normal diffusion output metrics. A dead replica's snapshot is removed. Because this deliberately reuses existing message paths, the gauge can retain its last value while a replica produces no output. +- `kv_wait_s` measures the AR scheduler interval from entering the KV-transfer-free wait state until extraction acknowledgement. The sample is carried to the Orchestrator through the existing engine output and labeled by connector type. Abort/error/removal paths discard incomplete timers rather than emitting partial waits. + +Zero and missing values have different meanings. Valid zero-duration queue observations are recorded, while unavailable optional measurements are omitted instead of being exported as synthetic zeroes. `peak_memory_mb`, `diffusion_forward_s`, `diffusion_kv_load_s`, `vae_decode_s`, and `kv_wait_s` therefore appear only when their data source is active. + +**Path 3: Audio modality metrics (`vllm_omni:audio_*`)** `OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites: @@ -76,7 +89,7 @@ There are four independent paths for metric collection. - `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry. - `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`. -**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)** +**Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`)** `OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution. @@ -87,7 +100,7 @@ The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transf Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected). -**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)** +**Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap)** The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics. @@ -97,9 +110,9 @@ The Orchestrator runs in a background thread. The API server (OmniBase) runs in ### Metric Registration and Lifecycle -All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data. +All `vllm_omni:*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data. -The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime. +The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm_omni:*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime. ## OmniPrometheusStatLogger Wrap @@ -130,11 +143,11 @@ A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestr ## Gating: `--log-stats` -All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`. +All metrics — both the omni-specific `vllm_omni:*` families and the upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`. Behavior with `--log-stats=off` (default): -- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them. +- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm_omni:*` stay materialized but never have data written to them. - `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all. - The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate. @@ -162,10 +175,43 @@ second. | Metric | Type | Labels | Description | |--------|------|--------|-------------| -| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages | -| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled | -| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | -| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds | +| `vllm_omni:num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages | +| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled | +| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | +| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds | + +### Image and diffusion service-level metrics + +| Metric | Type | Labels | Description | +|--------|------|--------|-------------| +| `vllm_omni:stage_gen_time_s` | Histogram | `model_name`, `stage`, `stage_type` | Stage submit to finished output; includes in-stage queueing | +| `vllm_omni:request_queue_wait_s` | Histogram | `model_name` | Orchestration-layer queue wait; a present zero is recorded | +| `vllm_omni:stage_waiting_requests` | Gauge | `model_name`, `stage` | Sum of the latest waiting snapshots across live replicas in the stage | +| `vllm_omni:stage_in_queue_s` | Histogram | `model_name`, `stage` | Diffusion scheduler wait from enqueue to initial execution admission | +| `vllm_omni:num_inference_steps` | Histogram | `model_name` | Denoising-step distribution for diffusion stages | +| `vllm_omni:image_count_total` | Counter | `model_name` | Cumulative number of generated image outputs | +| `vllm_omni:image_pixels` | Histogram | `model_name` | Pixel-count distribution for image outputs | +| `vllm_omni:peak_memory_mb` | Gauge | `model_name`, `stage` | Positive per-stage peak-memory samples reported by engine outputs | +| `vllm_omni:requests_failed_total` | Counter | `model_name`, `reason` | Request failures with bounded reason values (`client_abort`, `client_disconnect`, `stage_error`, `unknown`) | +| `vllm_omni:kv_wait_s` | Histogram | `model_name`, `connector_type` | Completed AR KV-transfer wait intervals, grouped by connector backend | + +### Diffusion execution breakdown + +Labels: `{model_name, stage, replica}`. + +| Metric | Type | Description | +|--------|------|-------------| +| `vllm_omni:diffusion_exec_s` | Histogram | Core diffusion-step execution time per request | +| `vllm_omni:diffusion_exec_per_step_s` | Histogram | Core diffusion execution divided by inference-step count | +| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time | +| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing time | +| `vllm_omni:vae_decode_s` | Histogram | VAE decode time when the engine reports it | +| `vllm_omni:diffusion_forward_s` | Histogram | Forward-only denoise-loop time when pipeline profiling is enabled | +| `vllm_omni:diffusion_kv_load_s` | Histogram | Diffusion-side KV receive/load time when upstream KV is used | +| `vllm_omni:image_ttfp_s` | Histogram | Image-stage submit to first materialized image output | +| `vllm_omni:denoise_step_latency_s` | Histogram | Mean diffusion forward time per denoise step for image output stages | + +The optional breakdown and memory families are sparse by design: when an engine does not produce the source measurement, no sample is emitted. This differs from queue timings, where an explicitly measured zero is meaningful and is retained. ### Audio (7) @@ -173,13 +219,13 @@ Labels: `{model_name, stage, replica}` plus the listed extra label. | Metric | Type | Extra label | Description | |--------|------|-------------|-------------| -| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | -| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | -| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` | -| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated | -| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps | -| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | -| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | +| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | +| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | +| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` | +| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frames generated | +| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps | +| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | +| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | ### Cross-stage transfer (4) @@ -187,10 +233,10 @@ Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`. | Metric | Type | Description | |--------|------|-------------| -| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | -| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | -| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | -| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) | +| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | +| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | +| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | +| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) | ### LLM stage-level (wrapped `vllm:*`) @@ -203,7 +249,7 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution. - Counters use the `_total` suffix (auto-appended by `prometheus_client`). - Sizes use the `_bytes` suffix. -- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`. +- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm_omni:` / `vllm_omni`. ## Logging vs. Prometheus diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py index 8f434c0e2..8caf4a8dc 100644 --- a/tests/diffusion/test_diffusion_engine_metrics.py +++ b/tests/diffusion/test_diffusion_engine_metrics.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Source-level regression tests for diffusion output/engine helpers. These tests verify naming conventions and patterns by inspecting source code @@ -37,10 +37,50 @@ _FORMATTER_PATH = os.path.normpath( "output_formatter.py", ) ) +_METRICS_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "metrics", + "utils.py", + ) +) +_INLINE_CLIENT_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "diffusion", + "inline_stage_diffusion_client.py", + ) +) +_STAGE_PROC_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "diffusion", + "stage_diffusion_proc.py", + ) +) +_ORCHESTRATOR_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "engine", + "orchestrator.py", + ) +) def _read_source(path: str) -> str: - with open(path) as f: + with open(path, encoding="utf-8") as f: return f.read() @@ -100,9 +140,99 @@ class TestMetricKeys: 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 '"diffusion_engine_total_time_ms": step_total_ms' 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 assert '"postprocess_time_ms": postprocess_time * 1000' in step_streaming_source + def test_step_streaming_piggybacks_scheduler_waiting_snapshot(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "diffusion_scheduler_waiting_metrics(" in step_streaming_source + assert "_scheduler_num_waiting_reqs" in step_streaming_source + + def test_step_streaming_preserves_snapshot_on_postprocess_error(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "scheduler_metrics = diffusion_scheduler_waiting_metrics(" in step_streaming_source + assert 'setattr(exc, "diffusion_metrics", scheduler_metrics)' in step_streaming_source + assert step_streaming_source.index("scheduler_metrics =") < step_streaming_source.index( + "self.postprocess_output(" + ) + + def test_abort_keeps_output_consumers_alive_for_terminal_snapshot(self) -> None: + inline_source = _read_source(_INLINE_CLIENT_PATH) + abort_source = _get_function_source(inline_source, "InlineStageDiffusionClient", "abort_requests_async") + assert ".cancel(" not in abort_source + + proc_source = _read_source(_STAGE_PROC_PATH) + run_loop_source = _get_function_source(proc_source, "StageDiffusionProc", "run_loop") + abort_branch = run_loop_source[run_loop_source.index('elif msg_type == "abort":') :] + abort_branch = abort_branch[: abort_branch.index('elif msg_type == "collective_rpc":')] + assert ".cancel(" not in abort_branch + + def test_orchestrator_consumes_metrics_only_output_without_routing(self) -> None: + source = _read_source(_ORCHESTRATOR_PATH) + loop_source = _get_function_source(source, "Orchestrator", "_orchestration_loop") + snapshot_pos = loop_source.index("_update_stage_replica_waiting(") + sentinel_pos = loop_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID") + route_pos = loop_source.index("pool.record_output_timestamps([diffusion_output])") + assert snapshot_pos < sentinel_pos < route_pos + + +class TestVaeDecodeEmit: + """VAE decode timing is sourced from ``DiffusionOutput.stage_durations`` + (populated by the diffusion pipeline profiler) and emitted as + ``vae_decode_time_ms`` so the ``_MS_TO_S`` accumulator picks it up. + """ + + def test_extract_vae_decode_ms_helper_exists(self) -> None: + source = _read_source(_METRICS_PATH) + helper_src = _get_function_source(source, None, "extract_diffusion_vae_decode_ms") + assert ".vae.decode" in helper_src, "helper must key on the '.vae.decode' suffix" + + def test_step_streaming_emits_vae_decode_time_ms(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "extract_diffusion_vae_decode_ms(" in step_streaming_source, ( + "step_streaming must call extract_diffusion_vae_decode_ms to source VAE decode timing" + ) + assert '"vae_decode_time_ms"' in step_streaming_source, ( + "step_streaming must emit vae_decode_time_ms key for the _MS_TO_S accumulator" + ) + + +class TestDiffuseForwardEmit: + """Forward-only timing → forward_time_ms, mirroring TestVaeDecodeEmit.""" + + def test_extract_diffuse_ms_helper_exists(self) -> None: + source = _read_source(_METRICS_PATH) + helper_src = _get_function_source(source, None, "extract_diffusion_denoise_ms") + assert ".diffuse" in helper_src, "helper must key on the '.diffuse' suffix" + + def test_step_streaming_emits_forward_time_ms(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "extract_diffusion_denoise_ms(" in step_streaming_source, ( + "step_streaming must call extract_diffusion_denoise_ms to source forward-only timing" + ) + assert '"forward_time_ms"' in step_streaming_source, ( + "step_streaming must emit forward_time_ms key for the _MS_TO_S accumulator" + ) + + +class TestKvRecvEmit: + """KV-recv timing → kv_recv_time_ms, sourced from DiffusionOutput.kv_recv_ms + (set by the runner's _prepare_request_for_forward, not the profiler).""" + + def test_step_streaming_emits_kv_recv_time_ms(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "kv_recv_ms" in step_streaming_source, "step_streaming must read output.kv_recv_ms" + assert '"kv_recv_time_ms"' in step_streaming_source, ( + "step_streaming must emit kv_recv_time_ms key for the _MS_TO_S accumulator" + ) + class TestDummyRunAllocation: """Verify _dummy_run generates exact-sized audio arrays.""" diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index 3641ef01b..ae9b86f80 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -773,6 +773,19 @@ class TestRequestScheduler: assert third.num_running_reqs == 1 assert third.num_waiting_reqs == 0 + def test_records_initial_scheduler_queue_wait(self, mocker: MockerFixture) -> None: + perf_counter = mocker.patch( + "vllm_omni.diffusion.sched.base_scheduler.time.perf_counter", + side_effect=[10.0, 10.125], + ) + request = _make_step_request("queue-wait", num_inference_steps=2) + + self.scheduler.add_request(request) + self.scheduler.schedule() + + assert request.scheduler_queue_wait_ms == pytest.approx(125.0) + assert perf_counter.call_count == 2 + def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None: scheduler = RequestScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=2)) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index d3298844f..67e6f4751 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import multiprocessing as mp @@ -1079,6 +1079,42 @@ class TestStageDiffusionClientErrorPropagation: assert client.get_diffusion_output_nowait() is None + def test_error_response_preserves_scheduler_metrics(self): + from vllm_omni.metrics import definitions as metric_defs + + client = self._make_client() + client._response_socket.recv.side_effect = [b"message", zmq.Again()] + client._decoder.decode.return_value = { + "type": "error", + "request_id": "req-error", + "error": "gpu fault", + "metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0}, + } + + output = client.get_diffusion_output_nowait() + + assert output is not None + assert output.error == "gpu fault" + assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0 + + def test_metrics_response_creates_metrics_only_output(self): + from vllm_omni.metrics import definitions as metric_defs + from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID + + client = self._make_client() + client._response_socket.recv.side_effect = [b"message", zmq.Again()] + client._decoder.decode.return_value = { + "type": "metrics", + "metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0}, + } + + output = client.get_diffusion_output_nowait() + + assert output is not None + assert output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID + assert output.error is None + assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0 + def test_check_health_raises_when_proc_dead(self): """``check_health`` detects a dead subprocess via the manager's proc and raises ``EngineDeadError``, setting ``_engine_dead`` as a diff --git a/tests/engine/test_membership_controller.py b/tests/engine/test_membership_controller.py index 410bb3030..5f0416377 100644 --- a/tests/engine/test_membership_controller.py +++ b/tests/engine/test_membership_controller.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -47,6 +47,7 @@ class FakePool: self.invalidated = [] self.hub = None self.lb = None + self.replica_ids = {} def attach_hub(self, hub): self.hub = hub @@ -57,7 +58,9 @@ class FakePool: def add_client(self, input_addr, client, *, replica_id=None): self.added.append((input_addr, client, replica_id)) self.clients.append(client) - return len(self.clients) - 1 if replica_id is None else replica_id + resolved_replica_id = len(self.clients) - 1 if replica_id is None else replica_id + self.replica_ids[input_addr] = resolved_replica_id + return resolved_replica_id def invalidate_addr(self, input_addr): self.invalidated.append(input_addr) @@ -67,6 +70,12 @@ class FakePool: self.removed.append(input_addr) return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown")) + def get_client_by_addr(self, input_addr): + return SimpleNamespace() if input_addr in self.replica_ids else None + + def get_replica_id_by_addr(self, input_addr): + return self.replica_ids.get(input_addr) + def _snapshot(*replicas): return SimpleNamespace(replicas=list(replicas)) @@ -209,6 +218,16 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey return replacement controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=factory) + removed_replicas = [] + + async def cleanup(_request_ids): + return None + + controller.install_unregister_handlers( + output_queue=asyncio.Queue(), + cleanup_callback=cleanup, + replica_removed_callback=lambda stage_id, replica_id: removed_replicas.append((stage_id, replica_id)), + ) await controller.handle_unregister(0, "tcp://old-0") await controller.handle_unregister(0, "tcp://old-1") @@ -219,6 +238,7 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey assert pool.get_replica_id_by_addr("tcp://new-1") == 1 assert pool.get_replica_id_by_addr("tcp://old-1") is None assert pool.available_replica_ids() == [1] + assert removed_replicas == [(0, 0), (0, 1)] @pytest.mark.asyncio diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 981331d45..7103a4948 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import asyncio @@ -2063,6 +2066,28 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None: assert metrics.output_unit_count == 3 +def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None: + stage = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image") + pool = StagePool( + 1, + [stage], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + output = SimpleNamespace(request_id="req-image", _custom_output={"image": "non-empty"}) + pool.record_output_timestamps([output], output_ts=132.0) + + metrics = pool.build_stage_metrics( + [output], + submit_ts=130.0, + request_timestamp=120.0, + replica_id=0, + ) + + assert metrics.serving_time_to_first_output_ms == pytest.approx(12000.0) + assert metrics.image_time_to_first_output_ms == pytest.approx(2000.0) + + @pytest.mark.asyncio async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None: class FailingStageClient(FakeStageClient): diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index f6649e328..4f9fc592b 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -611,7 +611,13 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat # external ID to avoid collisions, so this also tests mapping external_req_id = "req-1" req_id = "req-1-12345678" + recorded_failures = [] try: + monkeypatch.setattr( + app, + "_record_request_failure_once", + lambda request_id, reason: recorded_failures.append((request_id, reason)), + ) app.request_states[req_id] = ClientRequestState( request_id=req_id, external_request_id=external_req_id, @@ -622,6 +628,7 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat assert engine.aborted == [[req_id]] assert external_req_id not in app.request_states + assert recorded_failures == [(req_id, "client_abort")] @pytest.mark.asyncio @@ -814,7 +821,13 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch): _patch_engine(monkeypatch, engine) app = Omni("dummy-model") + recorded_failures = [] try: + monkeypatch.setattr( + app, + "_record_request_failure_once", + lambda request_id, reason: recorded_failures.append((request_id, reason)), + ) app.request_states["req-1"] = object() app.abort("req-1") finally: @@ -822,6 +835,7 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch): assert engine.aborted == [["req-1"]] assert "req-1" not in app.request_states + assert recorded_failures == [("req-1", "client_abort")] def test_omni_forces_final_only_on_llm_stages(monkeypatch: pytest.MonkeyPatch): diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py new file mode 100644 index 000000000..b46f4080f --- /dev/null +++ b/tests/metrics/test_emit_calls.py @@ -0,0 +1,499 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Verify emit-call wiring in production code paths. + +Uses behavior and Prometheus exposition to verify production call semantics. +""" + +from __future__ import annotations + +import time +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest + +from vllm_omni.metrics import OmniPrometheusMetrics +from vllm_omni.metrics import definitions as defs + +if TYPE_CHECKING: + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_MODEL = "emit-calls-test" + + +# --------------------------------------------------------------------------- +# Behavioral pins — verify call semantics with mock OmniPrometheusMetrics. +# --------------------------------------------------------------------------- + + +def _make_omni_base_with_mock_prom(mocker): + """Build a minimal OmniBase shell wired with a mock prom_metrics. + + Mirrors the pattern in ``test_prometheus.py::TestRequestLifecycleGauges`` + — ``object.__new__`` skips ``__init__`` so we don't need a real engine. + """ + from vllm_omni.entrypoints.omni_base import OmniBase + + obj = object.__new__(OmniBase) + obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) + obj.request_states = {} + obj._consumed_metric_messages = {} + obj.log_stats = True + return obj, obj.prom_metrics + + +class TestFailureCounterWiring: + """All failure paths share one request-idempotent counter helper.""" + + def test_failure_reason_taxonomy_is_bounded(self) -> None: + from vllm_omni.metrics.utils import normalize_failure_reason + + assert normalize_failure_reason("client_abort") == "client_abort" + assert normalize_failure_reason("client_disconnect") == "client_disconnect" + assert normalize_failure_reason("stage_error") == "stage_error" + assert normalize_failure_reason("unknown") == "unknown" + assert normalize_failure_reason("timeout for request req-1") == "unknown" + assert normalize_failure_reason("") == "unknown" + assert normalize_failure_reason(None) == "unknown" + + def test_fire_failure_counter_passes_reason_through(self, mocker) -> None: + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace(e2e_done=set()), + failure_recorded=False, + ) + + obj._record_request_failure_once("req-1", reason="client_disconnect") + + prom.request_failed.assert_called_once() + prom.inc_requests_failed.assert_called_once_with("client_disconnect") + + def test_fire_failure_counter_normalizes_unrecognized_reason(self, mocker) -> None: + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace(e2e_done=set()), + failure_recorded=False, + ) + + obj._record_request_failure_once("req-1", reason="CUDA error for req-1") + + prom.inc_requests_failed.assert_called_once_with("unknown") + + def test_fire_failure_counter_skips_when_request_already_succeeded(self, mocker) -> None: + # When the request is in e2e_done (i.e. finalize already fired + # request_succeeded), the failure path must NOT double-count. + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace(e2e_done={"req-1"}), + failure_recorded=False, + ) + + obj._record_request_failure_once("req-1", reason="client_disconnect") + + prom.request_failed.assert_not_called() + prom.inc_requests_failed.assert_not_called() + + def test_fire_failure_counter_skips_when_request_state_missing(self, mocker) -> None: + # No request_states entry — already popped by abort path. Fail-safe + # must NOT raise. + obj, prom = _make_omni_base_with_mock_prom(mocker) + + obj._record_request_failure_once("missing-req", reason="oom") + + prom.request_failed.assert_not_called() + prom.inc_requests_failed.assert_not_called() + + def test_failure_counter_is_idempotent_across_abort_and_cleanup(self, mocker) -> None: + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace( + e2e_done=set(), + build_and_log_summary=lambda: None, + ), + failure_recorded=False, + ) + + obj._record_request_failure_once("req-1", reason="client_disconnect") + obj._log_summary_and_cleanup("req-1", reason="stage_error") + + prom.request_failed.assert_called_once() + prom.inc_requests_failed.assert_called_once_with("client_disconnect") + + def test_log_summary_and_cleanup_default_reason(self, mocker) -> None: + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace( + e2e_done=set(), + build_and_log_summary=lambda: None, + ), + failure_recorded=False, + ) + + obj._log_summary_and_cleanup("req-1") + + prom.inc_requests_failed.assert_called_once_with("stage_error") + + +class TestEarlyReturnOnLogStatsOff: + """When --log-stats is off (default), observe methods must be silent no-ops. + + This is the gating contract — emit sites can call unconditionally and the + helper short-circuits. + """ + + def test_observe_methods_silent_when_log_stats_false(self) -> None: + prom = OmniPrometheusMetrics(model_name=_MODEL, log_stats=False) + # None of these should raise or write to the registry. + prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=1.5) + prom.observe_stage_in_queue(stage=0, in_queue_s=0.2) + prom.observe_queue_wait(queue_wait_s=0.5) + prom.set_stage_waiting_requests(stage=0, n_waiting=3) + prom.observe_num_inference_steps(n_steps=20) + prom.inc_image_count(n_images=1) + prom.observe_image_pixels(n_pixels=512 * 512) + prom.set_peak_memory(stage=0, peak_memory_mb=1024.0) + prom.inc_requests_failed(reason="oom") + prom.observe_kv_wait(connector_type="shm", kv_wait_s=0.01) + + +class TestStageGenerationLabels: + def test_stage_gen_time_exposes_bounded_stage_type(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + model = _MODEL + "-stage-type" + prom = OmniPrometheusMetrics(model_name=model, log_stats=True) + prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=0.25) + prom.observe_stage_gen_time(stage=1, stage_type="diffusion", gen_time_s=1.5) + + out = generate_latest(REGISTRY).decode() + assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="0",stage_type="llm"}} 1.0' in out + assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="1",stage_type="diffusion"}} 1.0' in out + + +class TestCounterExpositionNames: + def test_counter_families_have_exactly_one_total_suffix(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + model = _MODEL + "-counter-suffix" + prom = OmniPrometheusMetrics(model_name=model, log_stats=True) + prom.inc_image_count(2) + prom.inc_requests_failed("stage_error") + + out = generate_latest(REGISTRY).decode() + assert f'{defs.IMAGE_COUNT_METRIC}_total{{model_name="{model}"}} 2.0' in out + assert f'{defs.REQUESTS_FAILED}_total{{model_name="{model}",reason="stage_error"}} 1.0' in out + assert f"{defs.IMAGE_COUNT_METRIC}_total_total" not in out + assert f"{defs.REQUESTS_FAILED}_total_total" not in out + + +class TestQueueWaitExtraction: + def test_present_zero_is_a_valid_observation(self) -> None: + from vllm_omni.metrics.utils import extract_queue_wait_s + + assert extract_queue_wait_s({"queue_wait_ms": 0.0}) == 0.0 + + def test_missing_queue_wait_is_not_synthetic_zero(self) -> None: + from vllm_omni.metrics.utils import extract_queue_wait_s + + assert extract_queue_wait_s({"preprocess_ms": 2.0}) is None + assert extract_queue_wait_s(None) is None + + +class TestStageInQueueObservation: + def test_prometheus_records_valid_zero_wait(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + model = _MODEL + "-zero-stage-wait" + prom = OmniPrometheusMetrics(model_name=model, log_stats=True) + + prom.observe_stage_in_queue(stage=1, in_queue_s=0.0) + + out = generate_latest(REGISTRY).decode() + sample = f'{defs.STAGE_IN_QUEUE_S}_count{{model_name="{model}",stage="1"}} 1.0' + assert sample in out + + +class TestStageWorkloadMetricScope: + @staticmethod + def _observe(mocker, *, stage_type: str, output_unit_type: str): + from vllm_omni.metrics.utils import observe_stage_workload_metrics + + prom = mocker.Mock(spec=OmniPrometheusMetrics) + stage_metrics = SimpleNamespace( + num_inference_steps=50, + output_unit_type=output_unit_type, + image_pixels=1024 * 1024, + output_unit_count=2, + ) + observe_stage_workload_metrics( + prom, + stage_type=stage_type, + stage_metrics=stage_metrics, + ) + return prom + + def test_diffusion_image_observes_steps_and_image_workload(self, mocker) -> None: + prom = self._observe(mocker, stage_type="diffusion", output_unit_type="image") + + prom.observe_num_inference_steps.assert_called_once_with(50) + prom.observe_image_pixels.assert_called_once_with(1024 * 1024) + prom.inc_image_count.assert_called_once_with(2) + + def test_diffusion_video_observes_steps_only(self, mocker) -> None: + prom = self._observe(mocker, stage_type="diffusion", output_unit_type="video") + + prom.observe_num_inference_steps.assert_called_once_with(50) + prom.observe_image_pixels.assert_not_called() + prom.inc_image_count.assert_not_called() + + def test_llm_text_observes_no_workload_metrics(self, mocker) -> None: + prom = self._observe(mocker, stage_type="llm", output_unit_type="text") + + prom.observe_num_inference_steps.assert_not_called() + prom.observe_image_pixels.assert_not_called() + prom.inc_image_count.assert_not_called() + + def test_non_diffusion_image_observes_image_workload_only(self, mocker) -> None: + prom = self._observe(mocker, stage_type="llm", output_unit_type="image") + + prom.observe_num_inference_steps.assert_not_called() + prom.observe_image_pixels.assert_called_once_with(1024 * 1024) + prom.inc_image_count.assert_called_once_with(2) + + @pytest.mark.parametrize("peak_memory_mb", [2048.0, 0.0]) + def test_same_finished_image_message_is_observed_exactly_once(self, mocker, peak_memory_mb: float) -> None: + from vllm_omni.entrypoints.omni_base import OmniBase + + obj = object.__new__(OmniBase) + obj._enable_ar_profiler = False + obj._consumed_metric_messages = {} + obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) + obj.mod_metrics = mocker.Mock() + obj.engine = SimpleNamespace( + get_stage_metadata=lambda _stage_id: SimpleNamespace( + stage_type="diffusion", + final_output=False, + final_output_type="image", + ) + ) + stage_metrics = SimpleNamespace( + stage_gen_time_ms=1000.0, + diffusion_metrics={"diffusion_engine_exec_time_s": 0.8, "scheduler_queue_wait_s": 0.1}, + num_inference_steps=20, + output_unit_type="image", + image_pixels=1024 * 1024, + output_unit_count=1, + serving_time_to_first_output_ms=250.0, + image_time_to_first_output_ms=250.0, + pipeline_timings={}, + ) + result = SimpleNamespace( + request_id="req-replay", + stage_id=1, + replica_id=0, + stage_submit_ts=10.0, + metrics=stage_metrics, + engine_outputs=SimpleNamespace( + stage_durations={}, + peak_memory_mb=peak_memory_mb, + finished=True, + final_output_type="image", + ), + ) + aggregator = mocker.Mock() + aggregator.stage_events = {} + aggregator.stage_first_ts = [None, None] + aggregator.stage_last_ts = [None, None] + + for _ in range(2): + assert obj._process_single_result(result, 1, aggregator, {}, 0.0, 1) is None + + aggregator.on_stage_metrics.assert_called_once_with(1, "req-replay", stage_metrics, "image") + obj.prom_metrics.inc_image_count.assert_called_once_with(1) + obj.prom_metrics.observe_image_pixels.assert_called_once_with(1024 * 1024) + obj.prom_metrics.observe_num_inference_steps.assert_called_once_with(20) + obj.prom_metrics.observe_stage_in_queue.assert_called_once_with(1, 0.1) + if peak_memory_mb > 0: + obj.prom_metrics.set_peak_memory.assert_called_once_with(1, peak_memory_mb) + else: + obj.prom_metrics.set_peak_memory.assert_not_called() + obj.mod_metrics.observe_image_ttfp.assert_called_once_with("1", "0", 0.25) + + +class TestStageWaitingAggregation: + @staticmethod + def _make_orchestrator(mocker): + from vllm_omni.engine.orchestrator import Orchestrator + + orchestrator = object.__new__(Orchestrator) + orchestrator._prom_metrics = mocker.Mock() + orchestrator._stage_replica_waiting = {} + return orchestrator + + def test_multiple_replicas_are_summed(self, mocker) -> None: + orchestrator = self._make_orchestrator(mocker) + + orchestrator._update_stage_replica_waiting(1, 0, 3) + orchestrator._update_stage_replica_waiting(1, 1, 5) + + orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 8) + + def test_zero_update_and_stage_isolation(self, mocker) -> None: + orchestrator = self._make_orchestrator(mocker) + + orchestrator._update_stage_replica_waiting(1, 0, 3) + orchestrator._update_stage_replica_waiting(2, 0, 7) + orchestrator._update_stage_replica_waiting(1, 0, 0) + + orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 0) + assert orchestrator._stage_replica_waiting[(2, 0)] == 7 + + def test_dead_replica_snapshot_is_removed(self, mocker) -> None: + orchestrator = self._make_orchestrator(mocker) + orchestrator._update_stage_replica_waiting(1, 0, 3) + orchestrator._update_stage_replica_waiting(1, 1, 5) + + orchestrator._remove_stage_replica_waiting(1, 1) + + orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 3) + + +# --------------------------------------------------------------------------- +# Behavioral pin: call sites now scope workload observations by stage/output +# type. Keep the Prometheus wrapper's <= 0 guard as a second line of defense. +# --------------------------------------------------------------------------- + + +class TestZeroGuardContract: + """Zero-valued observations must not bump histogram ``_count``.""" + + def test_zero_image_pixels_not_observed(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-guard") + prom.observe_image_pixels(n_pixels=0) + out = generate_latest(REGISTRY).decode() + # ``.labels()`` in OmniPrometheusMetrics.__init__ already creates the + # child sample, so the ``_count`` line exists with value 0.0 even + # without observations. The guard prevents ``_count`` from bumping + # to 1.0 — assert the value stays at 0. + needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-zero-guard"}}' + count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)] + assert count_lines, "expected image_pixels_count line to exist after OmniPrometheusMetrics construction" + assert float(count_lines[0].split()[-1]) == 0.0, ( + "zero-pixel observation leaked to registry — guard should early-return on <= 0" + ) + + def test_zero_num_inference_steps_not_observed(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-steps") + prom.observe_num_inference_steps(n_steps=0) + out = generate_latest(REGISTRY).decode() + needle = f'vllm_omni:num_inference_steps_count{{model_name="{_MODEL}-zero-steps"}}' + count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)] + assert count_lines, "expected num_inference_steps_count line to exist after construction" + assert float(count_lines[0].split()[-1]) == 0.0, ( + "zero-step observation leaked to registry — guard should early-return on <= 0" + ) + + def test_positive_image_pixels_observed(self) -> None: + # Sanity check: positive values DO get observed. + from prometheus_client import REGISTRY, generate_latest + + prom = OmniPrometheusMetrics(model_name=_MODEL + "-pos-guard") + prom.observe_image_pixels(n_pixels=512) + out = generate_latest(REGISTRY).decode() + needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-pos-guard"}}' + count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)] + assert count_lines, "expected image_pixels_count line after positive observation" + assert float(count_lines[0].split()[-1]) == 1.0, "positive-pixel observation did not increment count to 1" + + +# --------------------------------------------------------------------------- +# KV-wait emit wiring — scheduler ENTER/EXIT lifecycle + orchestrator dispatch. +# --------------------------------------------------------------------------- + + +def _make_scheduler_shell() -> OmniARScheduler: + """Minimal OmniARScheduler shell — skips upstream __init__ via object.__new__.""" + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + + obj = object.__new__(OmniARScheduler) + obj.requests = {} + obj.running = [] + obj.waiting = [] + obj._kv_wait_start_ts = {} + obj._omni_kv_config = None + return obj + + +class TestKvWaitSchedulerEmit: + """Scheduler-side half of kv_wait_s: pops start ts, carries the wait across.""" + + def test_skips_when_no_start_ts_recorded(self) -> None: + sched = _make_scheduler_shell() + outputs: dict[int, list] = {} + + sched._emit_kv_wait_output(outputs, "req-no-wait", req=SimpleNamespace(client_index=0)) + + assert outputs == {}, "emit must not append output when no start ts recorded" + assert "req-no-wait" not in sched._kv_wait_start_ts + + def test_emits_wait_duration_and_connector_type(self) -> None: + sched = _make_scheduler_shell() + + sched._kv_wait_start_ts["req-1"] = time.monotonic() - 0.25 + outputs: dict[int, list] = {} + live_req = SimpleNamespace(client_index=3) + + sched._emit_kv_wait_output(outputs, "req-1", req=live_req) + + assert "req-1" not in sched._kv_wait_start_ts, "start ts must be popped after emit" + assert 3 in outputs and len(outputs[3]) == 1 + params = outputs[3][0].kv_transfer_params + assert "kv_wait_s" in params and "connector_type" in params + assert 0.0 < params["kv_wait_s"] < 1.0 + assert params["connector_type"] == "unknown" + + def test_connector_type_resolved_from_omni_kv_config(self) -> None: + sched = _make_scheduler_shell() + sched._omni_kv_config = {"connector_config": {"type": "SharedMemoryConnector"}} + + sched._kv_wait_start_ts["req-3"] = time.monotonic() - 0.01 + outputs: dict[int, list] = {} + + sched._emit_kv_wait_output(outputs, "req-3", req=SimpleNamespace(client_index=0)) + + eco = outputs[0][0] + assert eco.kv_transfer_params["connector_type"] == "SharedMemoryConnector" + + +class TestKvWaitTerminalCleanup: + def test_finish_requests_clears_requested_wait_only(self) -> None: + from vllm.v1.request import RequestStatus + + scheduler = _make_scheduler_shell() + scheduler._kv_wait_start_ts = {"req-1": 10.0, "req-2": 20.0} + + finished = scheduler.finish_requests( + (request_id for request_id in ["req-1"]), + RequestStatus.FINISHED_ABORTED, + ) + + assert finished == [] + assert "req-1" not in scheduler._kv_wait_start_ts + assert scheduler._kv_wait_start_ts["req-2"] == 20.0 + + def test_cleanup_is_idempotent_when_timestamp_is_already_gone(self) -> None: + from vllm.v1.request import RequestStatus + + scheduler = _make_scheduler_shell() + + scheduler.finish_requests("req-already-emitted", RequestStatus.FINISHED_ABORTED) + + assert scheduler._kv_wait_start_ts == {} diff --git a/tests/metrics/test_modality.py b/tests/metrics/test_modality.py index 315594233..694808a31 100644 --- a/tests/metrics/test_modality.py +++ b/tests/metrics/test_modality.py @@ -1,5 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations +from types import SimpleNamespace + import pytest from prometheus_client import REGISTRY, generate_latest @@ -10,6 +15,7 @@ from vllm_omni.metrics.modality import ( observe_audio_streaming_finalize, observe_modality_at_finalize, ) +from vllm_omni.metrics.stats import OrchestratorAggregator pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -46,6 +52,11 @@ _EXPECTED_FAMILIES = [ defs.DIFFUSION_EXEC_PER_STEP_S, defs.DIFFUSION_PREPROCESS_S, defs.DIFFUSION_POSTPROCESS_S, + defs.VAE_DECODE_S, + defs.DIFFUSION_FORWARD_S, + defs.DIFFUSION_KV_LOAD_S, + defs.IMAGE_TTFP_S, + defs.DENOISE_STEP_LATENCY_S, ] @@ -63,6 +74,11 @@ class TestRegistration: mod.observe_diffusion_exec_per_step("s", "r", 0.01) mod.observe_diffusion_preprocess("s", "r", 0.01) mod.observe_diffusion_postprocess("s", "r", 0.2) + mod.observe_vae_decode("s", "r", 0.3) + mod.observe_diffusion_forward("s", "r", 0.8) + mod.observe_diffusion_kv_load("s", "r", 0.04) + mod.observe_image_ttfp("s", "r", 2.5) + mod.observe_denoise_step_latency("s", "r", 0.05) out = generate_latest(REGISTRY).decode() for name in _EXPECTED_FAMILIES: @@ -247,6 +263,21 @@ class _StubModMetrics: def observe_diffusion_postprocess(self, s, r, v): self.calls.append(("observe_diffusion_postprocess", s, r, v)) + def observe_vae_decode(self, s, r, seconds): + self.calls.append(("observe_vae_decode", s, r, seconds)) + + def observe_diffusion_forward(self, s, r, seconds): + self.calls.append(("observe_diffusion_forward", s, r, seconds)) + + def observe_diffusion_kv_load(self, s, r, seconds): + self.calls.append(("observe_diffusion_kv_load", s, r, seconds)) + + def observe_image_ttfp(self, s, r, seconds): + self.calls.append(("observe_image_ttfp", s, r, seconds)) + + def observe_denoise_step_latency(self, s, r, seconds): + self.calls.append(("observe_denoise_step_latency", s, r, seconds)) + class _Bag: def __init__(self, **kw): @@ -306,11 +337,13 @@ class TestObserveModalityAtFinalize: def test_diffusion_path_full(self): stub = _StubModMetrics() stage_metrics = _Bag( + output_unit_type="image", stage_gen_time_ms=1000.0, diffusion_metrics={ "preprocess_time_s": 0.05, "diffusion_engine_exec_time_s": 1.2, "postprocess_time_s": 0.3, + "forward_time_s": 1.0, "image_num": 1, "resolution": 1024, "num_inference_steps": 50, @@ -328,6 +361,71 @@ class TestObserveModalityAtFinalize: assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls + assert ("observe_denoise_step_latency", "2", "0", pytest.approx(0.02)) in stub.calls + + def test_denoise_step_latency_requires_forward_time(self): + stub = _StubModMetrics() + stage_metrics = _Bag( + output_unit_type="image", + num_inference_steps=20, + diffusion_metrics={"diffusion_engine_exec_time_s": 1.2}, + ) + observe_modality_at_finalize( + stub, + output_type="image", + stage_id=2, + replica_id=0, + stage_metrics=stage_metrics, + engine_outputs=_Bag(), + ) + assert not any(c[0] == "observe_denoise_step_latency" for c in stub.calls) + + def test_stage_snapshot_keeps_forward_based_denoise_latency(self): + aggregator = OrchestratorAggregator( + num_stages=1, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=0, + ) + aggregator.diffusion_metrics["req-image"]["forward_time_s"] = 0.8 + stage_metrics = _Bag( + stage_id=None, + request_id=None, + final_output_type=None, + num_tokens_in=0, + num_tokens_out=0, + stage_gen_time_ms=1200.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=_Bag(), + diffusion_metrics=None, + num_inference_steps=20, + output_unit_type="image", + denoise_step_latency_ms=0.0, + ) + + stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image") + + assert stats.denoise_step_latency_ms == pytest.approx(40.0) + + def test_stage_snapshot_falls_back_to_stage_time_without_profiler(self): + aggregator = object.__new__(OrchestratorAggregator) + stage_metrics = _Bag( + stage_id=None, + request_id=None, + final_output_type=None, + stage_gen_time_ms=1200.0, + diffusion_metrics=None, + num_inference_steps=20, + output_unit_type="image", + denoise_step_latency_ms=0.0, + ) + + stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image") + + assert stats.diffusion_metrics is None + assert stats.denoise_step_latency_ms == pytest.approx(60.0) def test_diffusion_per_step_skipped_without_num_inference_steps(self): stub = _StubModMetrics() @@ -422,6 +520,80 @@ class TestObserveModalityAtFinalize: assert stub.calls == [] +class TestDiffusionUnitConversion: + def test_accumulator_converts_all_engine_milliseconds_to_seconds(self): + aggregator = OrchestratorAggregator( + num_stages=1, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=0, + ) + engine_output = SimpleNamespace( + metrics={ + "preprocess_time_ms": 50.0, + "diffusion_engine_exec_time_ms": 1500.0, + "postprocess_time_ms": 20.0, + "vae_decode_time_ms": 300.0, + "forward_time_ms": 800.0, + "scheduler_queue_wait_ms": 125.0, + "kv_recv_time_ms": 40.0, + } + ) + + aggregator.accumulate_diffusion_metrics("diffusion", "req-units", engine_output) + + assert dict(aggregator.diffusion_metrics["req-units"]) == { + "preprocess_time_s": pytest.approx(0.05), + "diffusion_engine_exec_time_s": pytest.approx(1.5), + "postprocess_time_s": pytest.approx(0.02), + "vae_decode_time_s": pytest.approx(0.3), + "forward_time_s": pytest.approx(0.8), + "scheduler_queue_wait_s": pytest.approx(0.125), + "kv_recv_time_s": pytest.approx(0.04), + } + + +class TestObserveDiffusionFinalize: + @pytest.mark.parametrize( + ("metric_key", "call_name", "value"), + [ + ("vae_decode_time_s", "observe_vae_decode", 0.3), + ("forward_time_s", "observe_diffusion_forward", 0.8), + ("kv_recv_time_s", "observe_diffusion_kv_load", 0.04), + ], + ) + def test_emits_optional_breakdown_metric(self, metric_key, call_name, value): + stub = _StubModMetrics() + observe_modality_at_finalize( + stub, + output_type="image", + stage_id=2, + replica_id=0, + stage_metrics=_Bag(diffusion_metrics={metric_key: value}), + engine_outputs=_Bag(), + ) + + assert (call_name, "2", "0", value) in stub.calls + + def test_audio_diffusion_stage_fires_both_paths(self): + stub = _StubModMetrics() + observe_modality_at_finalize( + stub, + output_type="audio", + stage_id=1, + replica_id=0, + stage_metrics=_Bag( + stage_gen_time_ms=1500.0, + audio_generated_frames=24000, + diffusion_metrics={"diffusion_engine_exec_time_s": 1.2}, + ), + engine_outputs=_Bag(multimodal_output={"audio_sample_rate": 24000}), + ) + + assert ("observe_diffusion_exec", "1", "0", 1.2) in stub.calls + assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls + + class TestObserveAudioFirstPacket: def test_observes_with_valid_inputs(self): stub = _StubModMetrics() @@ -541,3 +713,39 @@ class TestBucketSelection: # SECONDS_FAST_BUCKETS includes le=0.001 which is absent from SECONDS_BUCKETS. fast_marker = f'{defs.AUDIO_UNDERRUN_S}_bucket{{le="0.001"' assert fast_marker in out, "audio_underrun_s should use SECONDS_FAST_BUCKETS containing le=0.001" + + def test_vae_decode_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None: + # VAE decode can stretch into seconds for high-resolution latents; + # stays on the wide bucket so the 300 s ceiling covers video VAE. + stage, replica = "vae_decode_sec", "0" + mod.observe_vae_decode(stage, replica, 1.5) + out = generate_latest(REGISTRY).decode() + sec_marker = f'{defs.VAE_DECODE_S}_bucket{{le="300.0"' + assert sec_marker in out, "vae_decode_s should use SECONDS_BUCKETS containing le=300.0" + + def test_diffusion_forward_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None: + # Denoise loop can exceed the 60s fast-bucket ceiling for video gen. + stage, replica = "diff_forward_sec", "0" + mod.observe_diffusion_forward(stage, replica, 2.0) + out = generate_latest(REGISTRY).decode() + sec_marker = f'{defs.DIFFUSION_FORWARD_S}_bucket{{le="300.0"' + assert sec_marker in out, "diffusion_forward_s should use SECONDS_BUCKETS containing le=300.0" + + def test_denoise_step_latency_uses_seconds_fast_buckets(self, mod: OmniModalityMetrics) -> None: + # Per-step denoise is sub-second; stays on the fast bucket for ms-level + # resolution. + stage, replica = "denoise_fast", "0" + mod.observe_denoise_step_latency(stage, replica, 0.075) + out = generate_latest(REGISTRY).decode() + fast_marker = f'{defs.DENOISE_STEP_LATENCY_S}_bucket{{le="0.001"' + assert fast_marker in out, "denoise_step_latency_s should use SECONDS_FAST_BUCKETS containing le=0.001" + + +class TestDenoiseStepLatencyGuard: + @pytest.mark.parametrize(("value", "expected_count"), [(0.0, None), (-0.1, None), (0.05, 1.0)]) + def test_nonpositive_values_are_skipped(self, mod: OmniModalityMetrics, value, expected_count) -> None: + stage, replica = f"denoise_{value}", "0" + mod.observe_denoise_step_latency(stage, replica, value) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.DENOISE_STEP_LATENCY_S}_count{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == expected_count diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index dd7a75821..4a76c8bc5 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -3,8 +3,9 @@ from __future__ import annotations +import time from collections import defaultdict -from collections.abc import Iterable +from collections.abc import Iterable, Iterator from typing import Any import numpy as np @@ -107,6 +108,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Cache per-request flag to avoid repeated deserialization of additional_information self._omits_kv_transfer_cache: dict[str, bool] = {} + + # KV-wait start ts for the vllm_omni:kv_wait_s metric; see + # _emit_kv_wait_output for the engine-core → orchestrator carry. + self._kv_wait_start_ts: dict[str, float] = {} self._init_omni_io_scheduling_state() # Snapshot prompt length for each streaming input update self._new_prompt_len_snapshot: dict[str, int] = {} @@ -116,6 +121,70 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Output placeholders are zero when async scheduling isn't used return request.num_computed_tokens - request.num_output_placeholders + def _resolve_kv_connector_type(self) -> str: + """Connector backend name for the ``kv_wait_s`` label, or ``unknown``.""" + connector_cfg = self._get_omni_kv_config_value("connector_config") + if isinstance(connector_cfg, dict): + c_type = connector_cfg.get("type") + if isinstance(c_type, str) and c_type: + return c_type + return "unknown" + + def _emit_kv_wait_output( + self, + outputs: dict[int, list[OmniEngineCoreOutput]], + req_id: str, + req: Request, + ) -> None: + """Carry the KV-wait duration for ``req_id`` across the process boundary. + + Pops the start ts recorded at ENTER (``_free_request``); skips silently + when none was recorded. The wait rides ``kv_transfer_params`` to the + orchestrator, which calls ``observe_kv_wait``. + """ + start_ts = self._kv_wait_start_ts.pop(req_id, None) + if start_ts is None: + return + kv_wait_s = max(time.monotonic() - start_ts, 0.0) + outputs.setdefault(req.client_index, []).append( + OmniEngineCoreOutput( + request_id=req_id, + new_token_ids=[], + kv_transfer_params={ + "kv_wait_s": kv_wait_s, + "connector_type": self._resolve_kv_connector_type(), + }, + ) + ) + + def _clear_kv_wait_starts(self, request_ids: Iterable[str]) -> None: + """Drop incomplete KV-wait measurements for terminal requests.""" + wait_starts = getattr(self, "_kv_wait_start_ts", None) + if wait_starts is None: + return + for request_id in request_ids: + wait_starts.pop(request_id, None) + + def finish_requests( + self, + request_ids: str | Iterable[str] | None, + finished_status: RequestStatus, + ) -> list[Request]: + """Finish requests and discard any incomplete KV-wait timing.""" + if isinstance(request_ids, str): + cleanup_ids = (request_ids,) + finish_request_ids: str | tuple[str, ...] | None = request_ids + elif request_ids is None: + cleanup_ids = () + finish_request_ids = None + else: + finish_request_ids = list(request_ids) if isinstance(request_ids, Iterator) else request_ids + cleanup_ids = tuple(finish_request_ids) + + finished = super().finish_requests(finish_request_ids, finished_status) + self._clear_kv_wait_starts(cleanup_ids) + return finished + def _get_kv_transfer_criteria(self) -> dict | None: return self._get_omni_kv_config_value("kv_transfer_criteria") @@ -309,6 +378,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): kv_transfer_params={"kv_ready": True}, ) ) + # EXIT: extraction ack ends the wait — carry the duration. + # req is still alive here (the del happens in the cleanup + # loop below); skip if it was already freed upstream. + if req is not None: + self._emit_kv_wait_output(outputs, req_id, req) except Exception: init_logger(__name__).exception("Failed to pre-process KV extraction for %s", req_id) @@ -733,6 +807,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # It triggered but hasn't finished yet. We MUST wait. logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.") self.waiting_for_transfer_free.add(request_id) + self._kv_wait_start_ts[request_id] = time.monotonic() kv_xfer_params = None return kv_xfer_params, None elif request_id in self.waiting_for_transfer_free: @@ -745,6 +820,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): ) else: self.waiting_for_transfer_free.add(request_id) + self._kv_wait_start_ts[request_id] = time.monotonic() confirmed_computed = self._get_confirmed_num_computed_tokens(request) self._mark_request_for_kv_transfer(request_id, confirmed_computed) # Return KV transfer metadata so it propagates to RequestOutput diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 9920e3a27..aef8332fe 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -636,8 +636,8 @@ class OmniSchedulerMixin: ) } - if self.chunk_transfer_adapter: - self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + if chunk_transfer_adapter := getattr(self, "chunk_transfer_adapter", None): + chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) self._realign_request_status_to_queues( request_ids, diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index dece344b2..120f74000 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1528,6 +1528,9 @@ class DiffusionOutput: # memory usage info peak_memory_mb: float = 0.0 + # KV-recv wall-clock (ms) from the runner; feeds vllm_omni:diffusion_kv_load_s. + kv_recv_ms: float = 0.0 + # When True, move tensor fields to CPU at construction time. Useful when # the output is shipped across process boundaries (e.g. step-execution # mode) and the receiving side must not initialise a stray CUDA context. diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index f3294c5fa..800b57e34 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -56,6 +56,11 @@ from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput from vllm_omni.errors import client_error_from_metadata, is_client_error_status from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.metrics.utils import ( + diffusion_scheduler_waiting_metrics, + extract_diffusion_denoise_ms, + extract_diffusion_vae_decode_ms, +) if TYPE_CHECKING: from vllm_omni.outputs import OmniRequestOutput @@ -375,6 +380,9 @@ class DiffusionEngine: self._shutdown_complete = False self.abort_queue: queue.Queue[str] = queue.Queue() self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue() + # Copied onto the existing output metrics payload so queue monitoring + # reuses the normal diffusion result path without additional IPC. + self._scheduler_num_waiting_reqs = 0 def _init_execute_fn(self) -> None: if self.execution_mode == DiffusionExecutionMode.STEP_BATCH: @@ -444,7 +452,14 @@ class DiffusionEngine: ) raise postprocess_start_time = time.perf_counter() - formatted_outputs = self.postprocess_output(request, output) + scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0)) + try: + formatted_outputs = self.postprocess_output(request, output) + except Exception as exc: + # Preserve the latest scheduler snapshot across terminal + # abort/error paths, which do not produce formatted outputs. + setattr(exc, "diffusion_metrics", scheduler_metrics) + raise postprocess_time = time.perf_counter() - postprocess_start_time step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000 logger.debug( @@ -456,14 +471,24 @@ class DiffusionEngine: step_total_ms, ) for request_output in formatted_outputs: - request_output.metrics.update( - { - "preprocess_time_ms": preprocess_time * 1000, - "diffusion_engine_exec_time_ms": exec_total_time * 1000, - "diffusion_engine_total_time_ms": step_total_ms, - "postprocess_time_ms": postprocess_time * 1000, - } - ) + metrics_update = { + "preprocess_time_ms": preprocess_time * 1000, + "diffusion_engine_exec_time_ms": exec_total_time * 1000, + "postprocess_time_ms": postprocess_time * 1000, + **scheduler_metrics, + } + if request.scheduler_queue_wait_ms is not None: + metrics_update["scheduler_queue_wait_ms"] = request.scheduler_queue_wait_ms + vae_decode_ms = extract_diffusion_vae_decode_ms(output) + if vae_decode_ms is not None: + metrics_update["vae_decode_time_ms"] = vae_decode_ms + forward_ms = extract_diffusion_denoise_ms(output) + if forward_ms is not None: + metrics_update["forward_time_ms"] = forward_ms + kv_recv_ms = getattr(output, "kv_recv_ms", 0.0) + if kv_recv_ms > 0: + metrics_update["kv_recv_time_ms"] = kv_recv_ms + request_output.metrics.update(metrics_update) yield formatted_outputs async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]: @@ -554,6 +579,7 @@ class DiffusionEngine: self._wait_for_admission_if_needed_locked() sched_output = self.scheduler.schedule() + self._scheduler_num_waiting_reqs = max(int(sched_output.num_waiting_reqs), 0) if sched_output.is_empty: self._emit_finished_outputs(sched_output.finished_req_ids, None) diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 5082796a2..e5456f0cc 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Inline Stage Diffusion Client for vLLM-Omni multi-stage runtime. Runs DiffusionEngine in a ThreadPoolExecutor inside the Orchestrator process @@ -23,6 +26,7 @@ from vllm_omni.engine.stage_client import StageClientBase from vllm_omni.engine.stage_init_utils import StageMetadata from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt +from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID, diffusion_exception_metrics from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -158,6 +162,15 @@ class InlineStageDiffusionClient(StageClientBase): self._output_queue.put_nowait(result) except DiffusionRequestAbortedError as e: logger.info("request_id: %s aborted: %s", request_id, str(e)) + metrics = diffusion_exception_metrics(e) + if metrics: + self._output_queue.put_nowait( + OmniRequestOutput( + request_id=DIFFUSION_METRICS_ONLY_REQUEST_ID, + finished=True, + metrics=metrics, + ) + ) except Exception as e: logger.exception("Diffusion request %s failed: %s", request_id, e) status_code, error_type = client_error_metadata(e) @@ -167,6 +180,7 @@ class InlineStageDiffusionClient(StageClientBase): status_code=status_code, error_type=error_type, ) + error_output.metrics.update(diffusion_exception_metrics(e)) self._output_queue.put_nowait(error_output) finally: self._tasks.pop(request_id, None) @@ -181,9 +195,8 @@ class InlineStageDiffusionClient(StageClientBase): async def abort_requests_async(self, request_ids: list[str]) -> None: for rid in request_ids: - task = self._tasks.pop(rid, None) - if task: - task.cancel() + # Keep the consumer alive until the engine emits its terminal + # abort output. That path carries the refreshed scheduler gauge. self._engine.abort(rid) async def submit_interaction_async( diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 6173895c3..b2eae3fca 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -67,6 +67,11 @@ class OmniDiffusionRequest: # This is populated by a pipeline preprocessor before the request reaches # the scheduler; ``None`` keeps the default behavior for other pipelines. batch_compatibility_key: tuple[Any, ...] | None = None + # KV-recv wall-clock (ms), set by the runner's _prepare_request_for_forward + # and carried to DiffusionOutput for the vllm_omni:diffusion_kv_load_s metric. + kv_recv_ms: float = 0.0 + # Time spent waiting for initial admission by the diffusion scheduler. + scheduler_queue_wait_ms: float | None = None def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index 17291db37..8fe79a11c 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -3,6 +3,7 @@ from __future__ import annotations +import time from abc import ABC, abstractmethod from collections import deque from dataclasses import fields @@ -117,6 +118,7 @@ class BaseScheduler(ABC): if request_id in self._request_states: raise ValueError(f"request_id {request_id!r} is already active.") state = self._make_request_state(request_id, request) + state.queued_at = time.perf_counter() self._request_states[request_id] = state self._waiting.append(request_id) logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, request_id, len(self._waiting)) @@ -178,6 +180,7 @@ class BaseScheduler(ABC): state.status = DiffusionRequestStatus.RUNNING self._running.append(request_id) if was_new_request: + state.req.scheduler_queue_wait_ms = max((time.perf_counter() - state.queued_at) * 1000.0, 0.0) scheduled_new_reqs.append( NewRequestData.from_state( state, diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index e543ed251..62e5eddef 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -157,6 +157,7 @@ class SchedulerRequestState: diffusion_kv_requests: tuple[DiffusionKVRequest, ...] = () status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING error: str | None = None + queued_at: float = 0.0 def is_finished(self) -> bool: return DiffusionRequestStatus.is_finished(self.status) diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index a8733b8d0..96140ef91 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Stage Diffusion Client for vLLM-Omni multi-stage runtime. Owns the frontend-side ZMQ sockets for StageDiffusionProc and exposes the @@ -30,6 +33,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import ( ) from vllm_omni.engine.stage_client import StageClientBase from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -240,6 +244,14 @@ class StageDiffusionClient(StageClientBase): if msg_type == "result": self._output_queue.put_nowait(msg["output"]) + elif msg_type == "metrics": + self._output_queue.put_nowait( + OmniRequestOutput( + request_id=DIFFUSION_METRICS_ONLY_REQUEST_ID, + finished=True, + metrics=msg.get("metrics") or {}, + ) + ) elif msg_type == "rpc_result": self._rpc_results[msg["rpc_id"]] = msg["result"] elif msg_type == "error": @@ -264,14 +276,14 @@ class StageDiffusionClient(StageClientBase): # Route request errors as error outputs so the Orchestrator # sees the request complete (instead of hanging forever). if req_id is not None: - self._output_queue.put_nowait( - OmniRequestOutput.from_error( - req_id, - error_msg, - status_code=status_code, - error_type=error_type, - ) + error_output = OmniRequestOutput.from_error( + req_id, + error_msg, + status_code=status_code, + error_type=error_type, ) + error_output.metrics.update(msg.get("metrics") or {}) + self._output_queue.put_nowait(error_output) # Fields that are subprocess-local and cannot be serialized across # process boundaries. They are recreated in the subprocess with diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index f01348c74..224a88199 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Subprocess entry point for the diffusion engine. StageDiffusionProc runs DiffusionEngine in a child process, @@ -36,6 +39,7 @@ from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage from vllm_omni.engine.stage_init_utils import set_death_signal from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.metrics.utils import diffusion_exception_metrics from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -361,6 +365,16 @@ class StageDiffusionProc: request_id, str(e), ) + metrics = diffusion_exception_metrics(e) + if metrics: + await response_socket.send( + encoder.encode( + { + "type": "metrics", + "metrics": metrics, + } + ) + ) except Exception as e: logger.exception("Diffusion request %s failed: %s", request_id, e) status_code, error_type = client_error_metadata(e) @@ -372,6 +386,7 @@ class StageDiffusionProc: "error": str(e), "status_code": status_code, "error_type": error_type, + "metrics": diffusion_exception_metrics(e), } ) ) @@ -426,9 +441,8 @@ class StageDiffusionProc: elif msg_type == "abort": for rid in msg.get("request_ids", []): - task = tasks.pop(rid, None) - if task: - task.cancel() + # Let the request task consume the terminal abort + # output so it can publish the scheduler snapshot. self._engine.abort(rid) elif msg_type == "collective_rpc": diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 15680d60e..80c3a1bc4 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -539,6 +539,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): target_device=self._target_device if use_prefetch else getattr(self.pipeline, "device", None), ) kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000 + req.kv_recv_ms = kv_recv_ms logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms) # Kick off the next request's prefetch (+ H2D) to overlap this forward. @@ -596,6 +597,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): reqs: list[OmniDiffusionRequest], outputs: list[DiffusionOutput], ) -> BatchRunnerOutput: + for i in range(len(reqs)): + # Carry the runner-measured KV-recv timing onto the output so the + # engine's step_streaming can surface it as diffusion_kv_load_s. + outputs[i].kv_recv_ms = reqs[i].kv_recv_ms return BatchRunnerOutput.from_list( [ RunnerOutput( diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 4b44e7067..12c3413f1 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -107,6 +107,7 @@ class AsyncOmniEngine: _log_stats: bool = False _coordinator_runtime: Any = None _transfer_emitter: Any = None + _prom_metrics: Any = None _enable_orch_monitor: bool = False def __init__( @@ -117,6 +118,7 @@ class AsyncOmniEngine: diffusion_batch_size: int = 1, single_stage_mode: bool = False, transfer_emitter: Any = None, + prom_metrics: Any = None, log_stats: bool = False, tokenizer: str | None = None, trust_remote_code: bool | None = None, @@ -133,6 +135,7 @@ class AsyncOmniEngine: # Optional: when None, Orchestrator silently skips TX emit (existing # RX path still works via OrchestratorAggregator). self._transfer_emitter = transfer_emitter + self._prom_metrics = prom_metrics # Drives upstream EngineCore + scheduler stats production. When False # the engine skips SchedulerStats / IterationStats; the per-(stage, # replica) vllm:* wrap stays registered but reads zero. Respects the @@ -398,6 +401,7 @@ class AsyncOmniEngine: membership_controller=membership_controller, running_counter=self._running_counter, transfer_emitter=self._transfer_emitter, + prom_metrics=self._prom_metrics, log_stats=self._log_stats, enable_orch_monitor=self._enable_orch_monitor, duplex_runtime_extension=duplex_runtime_extension, diff --git a/vllm_omni/engine/membership_controller.py b/vllm_omni/engine/membership_controller.py index 2053f6160..8f008e286 100644 --- a/vllm_omni/engine/membership_controller.py +++ b/vllm_omni/engine/membership_controller.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MembershipController: distributed replica lifecycle management. @@ -60,6 +60,7 @@ class MembershipController: self._watcher_task: asyncio.Task[None] | None = None self._output_queue: asyncio.Queue[EngineQueueMessage] | None = None self._cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None + self._replica_removed_callback: Callable[[int, int], None] | None = None self._hub = OmniCoordClientForHub(coordinator_pub_address) factory = load_balancer_factory @@ -118,8 +119,12 @@ class MembershipController: return effective_output_queue = output_queue if output_queue is not None else self._output_queue effective_cleanup_callback = cleanup_callback if cleanup_callback is not None else self._cleanup_callback + replica_id = pool.get_replica_id_by_addr(input_addr) + client = pool.get_client_by_addr(input_addr) affected = pool.invalidate_addr(input_addr) self._detach_replica(stage_id, input_addr) + if client is not None and replica_id is not None and self._replica_removed_callback is not None: + self._replica_removed_callback(stage_id, replica_id) if affected and effective_cleanup_callback is not None: await effective_cleanup_callback(affected) if affected and effective_output_queue is not None: @@ -283,10 +288,12 @@ class MembershipController: *, output_queue: asyncio.Queue[EngineQueueMessage], cleanup_callback: Callable[[list[str]], Awaitable[None]], + replica_removed_callback: Callable[[int, int], None] | None = None, ) -> None: """Install shared cleanup sinks for watcher-driven unregister events.""" self._output_queue = output_queue self._cleanup_callback = cleanup_callback + self._replica_removed_callback = replica_removed_callback def _pool_for_stage_id(self, stage_id: int) -> StagePool | None: if not (0 <= stage_id < len(self._stage_pools)): diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 8717f7a1a..a9bc33118 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -57,8 +57,10 @@ from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_k from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError from vllm_omni.errors import DEFAULT_CLIENT_ERROR_TYPE +from vllm_omni.metrics import definitions as metric_defs from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger +from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -360,6 +362,7 @@ class Orchestrator: # don't AttributeError when transfer / counter emit paths access them. _running_counter: OmniRequestCounter | None = None _transfer_emitter: Any = None + _prom_metrics: Any = None _stat_logger: OmniPrometheusStatLogger | None = None duplex_control_plane: DuplexControlPlanePort | None = None @@ -375,6 +378,7 @@ class Orchestrator: membership_controller: MembershipController | None = None, running_counter: OmniRequestCounter | None = None, transfer_emitter: Any = None, + prom_metrics: Any = None, log_stats: bool = False, enable_orch_monitor: bool = False, duplex_runtime_extension: DuplexRuntimeExtension | None = None, @@ -389,6 +393,8 @@ class Orchestrator: self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools self.log_stats = log_stats + self._prom_metrics = prom_metrics + self._stage_replica_waiting: dict[tuple[int, int], int] = {} self._orch_monitor = create_orch_monitor( enabled=enable_orch_monitor, replica_sampler=self._sample_replica_metrics, @@ -542,6 +548,7 @@ class Orchestrator: self._membership.install_unregister_handlers( output_queue=self.output_async_queue, cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True), + replica_removed_callback=self._remove_stage_replica_waiting, ) membership_watcher = self._membership.start() @@ -980,6 +987,16 @@ class Orchestrator: if diffusion_output is None: continue + output_metrics = getattr(diffusion_output, "metrics", None) + if isinstance(output_metrics, dict): + n_waiting = output_metrics.pop(metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY, None) + if n_waiting is not None: + self._update_stage_replica_waiting(stage_id, replica_id, int(n_waiting)) + + if diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID: + idle = False + continue + pool.record_output_timestamps([diffusion_output]) processed = [diffusion_output] else: @@ -989,6 +1006,18 @@ class Orchestrator: await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs) for eco in raw_outputs.outputs: + # Emit kv_wait_s before _handle_kv_ready_raw_outputs' + # async_chunk early-return so it lands in all modes. + kv_params = getattr(eco, "kv_transfer_params", None) + if ( + self._prom_metrics is not None + and isinstance(kv_params, dict) + and (kv_wait_s := kv_params.get("kv_wait_s")) is not None + ): + self._prom_metrics.observe_kv_wait( + kv_params.get("connector_type") or "unknown", + float(kv_wait_s), + ) req_state = self.request_states.get(getattr(eco, "request_id", None)) if req_state is None or not req_state.streaming.enabled: continue @@ -1027,6 +1056,17 @@ class Orchestrator: iteration_stats, engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)], ) + _sched_stats = raw_outputs.scheduler_stats + if ( + self._prom_metrics is not None + and _sched_stats is not None + and getattr(_sched_stats, "num_waiting_reqs", None) is not None + ): + self._update_stage_replica_waiting( + stage_id, + replica_id, + int(_sched_stats.num_waiting_reqs), + ) except asyncio.CancelledError: raise except EngineDeadError as e: @@ -1104,6 +1144,26 @@ class Orchestrator: close_duplex_sessions=True, ) + def _update_stage_replica_waiting(self, stage_id: int, replica_id: int, n_waiting: int) -> None: + """Update one replica snapshot and expose the stage-wide total.""" + self._stage_replica_waiting[(stage_id, replica_id)] = max(int(n_waiting), 0) + self._set_stage_waiting_total(stage_id) + + def _remove_stage_replica_waiting(self, stage_id: int, replica_id: int) -> None: + """Drop a dead replica's snapshot and refresh the stage total.""" + self._stage_replica_waiting.pop((stage_id, replica_id), None) + self._set_stage_waiting_total(stage_id) + + def _set_stage_waiting_total(self, stage_id: int) -> None: + if self._prom_metrics is None: + return + total = sum( + n_waiting + for (snapshot_stage_id, _), n_waiting in self._stage_replica_waiting.items() + if snapshot_stage_id == stage_id + ) + self._prom_metrics.set_stage_waiting_requests(stage_id, total) + async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: EngineDeadError) -> None: """Evict a dead stage replica and fail the requests stranded on it (#4285). @@ -1122,6 +1182,7 @@ class Orchestrator: error, ) pool.evict_replica(replica_id) + self._remove_stage_replica_waiting(stage_id, replica_id) stage_has_live = bool(pool.live_replica_ids()) failed_ids: list[str] = [] for req_id, req_state in list(self.request_states.items()): diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index d30f88a23..5a8305803 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unified stage-local runtime abstraction for vLLM-Omni.""" from __future__ import annotations @@ -653,11 +656,6 @@ class StagePool: ) image_pixels = self._count_image_pixels(request_outputs) if output_unit_type == "image" else 0 num_inference_steps = coerce_positive_int_scalar(getattr(sampling_params, "num_inference_steps", None)) or 0 - denoise_step_latency_ms = ( - defs.compute_denoise_step_latency(stage_gen_time_ms, num_inference_steps) - if output_unit_type == "image" - else 0.0 - ) has_output_timestamps = bool(output_timestamps) first_ts = output_timestamps[0] if has_output_timestamps else now serving_time_to_first_output_ms = ( @@ -665,6 +663,11 @@ class StagePool: if non_empty_first_output_ts is not None else 0.0 ) + image_time_to_first_output_ms = ( + max((non_empty_first_output_ts - submit_ts) * 1000.0, 0.0) + if non_empty_first_output_ts is not None and output_unit_type == "image" + else 0.0 + ) remaining_ms = max((now - first_ts) * 1000.0, 0.0) if output_unit_count > 1 and self._is_streaming_output_unit_type(output_unit_type): time_per_output_unit_ms = remaining_ms / float(output_unit_count - 1) @@ -709,10 +712,11 @@ class StagePool: audio_sample_rate=audio_sample_rate, audio_duration_s=audio_duration_s, image_pixels=image_pixels, - denoise_step_latency_ms=denoise_step_latency_ms, + num_inference_steps=num_inference_steps, output_unit_type=output_unit_type, output_unit_count=output_unit_count, serving_time_to_first_output_ms=serving_time_to_first_output_ms, + image_time_to_first_output_ms=image_time_to_first_output_ms, time_per_output_unit_ms=time_per_output_unit_ms, inter_output_latency_ms=inter_output_latency_ms, inter_output_latencies_ms=inter_output_latencies_ms, diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 5e5b96146..2ab7b2a25 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -645,12 +645,12 @@ class AsyncOmni(EngineClient, OmniBase): except (asyncio.CancelledError, GeneratorExit): if input_stream_task is not None and not input_stream_task.done(): input_stream_task.cancel() - self._fire_failure_counter_if_alive(request_id) + self._record_request_failure_once(request_id, reason="client_disconnect") await self._abort_internal_requests(request_id) logger.info(f"[AsyncOmni] Request {request_id} aborted.") raise except Exception as e: - self._fire_failure_counter_if_alive(request_id) + self._record_request_failure_once(request_id, reason="stage_error") await self._abort_internal_requests(request_id) logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}") raise @@ -1115,6 +1115,7 @@ class AsyncOmni(EngineClient, OmniBase): """ await self.engine.abort_async(request_ids) for rid in request_ids: + self._record_request_failure_once(rid, reason="client_abort") state = self.request_states.pop(rid, None) input_stream_task = getattr(state, "input_stream_task", None) if input_stream_task is not None and not input_stream_task.done(): diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 9526df36a..0aa0c7c5a 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -17,6 +17,8 @@ class ClientRequestState: self.stage_id: int | None = None self.queue = queue if queue is not None else asyncio.Queue() self.metrics: OrchestratorAggregator | None = None + # Request-scoped idempotency guard for Prometheus failure counters. + self.failure_recorded = False # Wall-clock time at which the user's request arrived in the engine # entrypoint. Set in async_omni.generate() before the orchestrator # accepts the request. Used as the t0 anchor for audio_ttfp. diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 14bdc10be..28bbe9219 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -200,6 +200,8 @@ class Omni(OmniBase): self._log_summary_and_cleanup(req_id) except Exception: if "active_reqs" in locals() and active_reqs: + for req_id in active_reqs: + self._record_request_failure_once(req_id, reason="stage_error") self.abort(list(active_reqs)) raise finally: @@ -210,6 +212,7 @@ class Omni(OmniBase): request_ids = [request_id] if isinstance(request_id, str) else list(request_id) self.engine.abort(request_ids) for req_id in request_ids: + self._record_request_failure_once(req_id, reason="client_abort") self.request_states.pop(req_id, None) if self.log_stats: logger.info("[Omni] Aborted request(s) %s", ",".join(request_ids)) diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index ff21f55a2..d7011242b 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import os @@ -30,6 +33,11 @@ from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_ from vllm_omni.metrics.prometheus import OmniPrometheusMetrics from vllm_omni.metrics.stats import OrchestratorAggregator from vllm_omni.metrics.transfer import OmniTransferMetrics +from vllm_omni.metrics.utils import ( + extract_queue_wait_s, + normalize_failure_reason, + observe_stage_workload_metrics, +) from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.tracking_parser import TrackingNamespace @@ -199,6 +207,7 @@ class OmniBase(PDDisaggregationMixin): # (which forwards it to the Orchestrator background thread for # TX-side emit; see Orchestrator._forward_to_next_stage). self.transfer_metrics = OmniTransferMetrics(model_name=model, log_stats=log_stats) + self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats) st = time.time() self.engine = AsyncOmniEngine( model=model, @@ -206,6 +215,7 @@ class OmniBase(PDDisaggregationMixin): stage_init_timeout=stage_init_timeout, diffusion_batch_size=diffusion_batch_size, transfer_emitter=self.transfer_metrics, + prom_metrics=self.prom_metrics, log_stats=log_stats, **kwargs, ) @@ -221,7 +231,6 @@ class OmniBase(PDDisaggregationMixin): self.request_states: dict[str, ClientRequestState] = {} self._consumed_metric_messages: dict[str, set[int]] = {} - self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats) self.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats) self.default_sampling_params_list = self.engine.default_sampling_params_list @@ -384,29 +393,26 @@ class OmniBase(PDDisaggregationMixin): raise TypeError(f"Expected a mapping, dataclass, or msgspec struct, got {type(params).__name__}") return type(params)(**{**values, **constraints}) - def _fire_failure_counter_if_alive(self, request_id: str) -> None: - """Fire the abort/exception bucket of requests_success_total. - - Called from cancel / exception paths in async_omni.generate() BEFORE - _abort_internal_requests pops request_states — that method resolves - the internal id by dict lookup, so popping first would no-op it. We - keep this counter fire separate from _log_summary_and_cleanup (which - pops) so the abort path can still find the state to clean up. - """ + def _record_request_failure_once(self, request_id: str, reason: str) -> None: req_state = self.request_states.get(request_id) prom = getattr(self, "prom_metrics", None) - if req_state is None or req_state.metrics is None or prom is None: + metrics = getattr(req_state, "metrics", None) + if metrics is None or prom is None: + return + if str(request_id) in metrics.e2e_done or getattr(req_state, "failure_recorded", False): return - if str(request_id) not in req_state.metrics.e2e_done: - prom.request_failed() - def _log_summary_and_cleanup(self, request_id: str) -> None: + req_state.failure_recorded = True + prom.request_failed() + prom.inc_requests_failed(normalize_failure_reason(reason)) + + def _log_summary_and_cleanup(self, request_id: str, reason: str = "stage_error") -> None: req_state = self.request_states.get(request_id) try: if req_state is None or req_state.metrics is None: return if str(request_id) not in req_state.metrics.e2e_done: - self.prom_metrics.request_failed() + self._record_request_failure_once(request_id, reason) if self.log_stats: # Emit per-request orchestrator timing (including e2e_total_ms) # before dropping request state. @@ -623,6 +629,30 @@ class OmniBase(PDDisaggregationMixin): metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs) metrics.on_stage_metrics(stage_id, req_id, _m, output_type) consumed.add(msg_id) + if peak_memory_mb > 0: + self.prom_metrics.set_peak_memory(stage_id, peak_memory_mb) + self.prom_metrics.observe_stage_gen_time( + stage_id, + stage_meta.stage_type, + _m.stage_gen_time_ms / 1000.0, + ) + observe_stage_workload_metrics( + self.prom_metrics, + stage_type=stage_meta.stage_type, + stage_metrics=_m, + ) + # Only meaningful for diffusion stages (text stages have no exec key). + _diff_m = _m.diffusion_metrics or {} + stage_in_queue_s = _diff_m.get("scheduler_queue_wait_s") + if stage_in_queue_s is not None: + self.prom_metrics.observe_stage_in_queue(stage_id, float(stage_in_queue_s)) + if _m.output_unit_type == "image": + if result.replica_id is not None: + self.mod_metrics.observe_image_ttfp( + str(stage_id), + str(result.replica_id), + _m.image_time_to_first_output_ms / 1000.0, + ) if not stage_meta.final_output: return None @@ -656,6 +686,10 @@ class OmniBase(PDDisaggregationMixin): _gen_tok += int(evt.num_tokens_out) self.prom_metrics.observe_tokens(_prompt_tok, _gen_tok) + queue_wait_s = extract_queue_wait_s(_m.pipeline_timings if _m is not None else None) + if queue_wait_s is not None: + self.prom_metrics.observe_queue_wait(queue_wait_s) + # Modality observe inside the same finalize guard so it fires # once per request and inherits the try/except isolation. observe_modality_at_finalize( diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index b5d5217f5..68adb44ca 100644 --- a/vllm_omni/metrics/definitions.py +++ b/vllm_omni/metrics/definitions.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Single source of truth for vLLM-Omni Prometheus + bench CLI metric naming. Consumed by: @@ -61,6 +64,7 @@ IMAGE_COUNT = "image_count" IMAGE_GENERATION = "image_generation" IMAGE_GENERATION_TIME_MS = f"{IMAGE_GENERATION}_time_ms" IMAGE_PIXELS = "image_pixels" +DIFFUSION_SCHEDULER_WAITING_KEY = "scheduler_num_waiting_reqs" TOTAL_IMAGES = "total_images" IMAGE_THROUGHPUT = "image_throughput" AVERAGE_PIXELS_PER_IMAGE = "average_pixels_per_image" @@ -91,6 +95,7 @@ POSTPROCESS_TIME_MS = f"{POSTPROCESS_TIME}_ms" POSTPROCESS_TIMES_MS = f"{POSTPROCESS_TIME}s_ms" OUTPUT_UNIT_COUNT = "output_unit_count" SERVING_TIME_TO_FIRST_OUTPUT_MS = "serving_time_to_first_output_ms" +IMAGE_TIME_TO_FIRST_OUTPUT_MS = "image_time_to_first_output_ms" SERVING_TIME_TO_FIRST_OUTPUTS_MS = "serving_time_to_first_outputs_ms" TIME_PER_OUTPUT_UNIT_MS = "time_per_output_unit_ms" TIME_PER_OUTPUT_UNITS_MS = "time_per_output_units_ms" @@ -152,6 +157,30 @@ TRANSFER_RX_S = METRIC_PREFIX + "transfer_rx_s" TRANSFER_IN_FLIGHT_S = METRIC_PREFIX + "transfer_in_flight_s" +# ============================================================================ +# Image / diffusion service-level families +# ============================================================================ +STAGE_GEN_TIME_S = METRIC_PREFIX + "stage_gen_time_s" +REQUEST_QUEUE_WAIT_S = METRIC_PREFIX + "request_queue_wait_s" +STAGE_WAITING_REQUESTS = METRIC_PREFIX + "stage_waiting_requests" +NUM_INFERENCE_STEPS = METRIC_PREFIX + "num_inference_steps" +IMAGE_COUNT_METRIC = METRIC_PREFIX + IMAGE_COUNT +IMAGE_PIXELS_METRIC = METRIC_PREFIX + IMAGE_PIXELS +PEAK_MEMORY_MB = METRIC_PREFIX + "peak_memory_mb" +REQUESTS_FAILED = METRIC_PREFIX + "requests_failed" +KV_WAIT_S = METRIC_PREFIX + "kv_wait_s" +DIFFUSION_FORWARD_S = METRIC_PREFIX + "diffusion_forward_s" +DIFFUSION_EXEC_S = METRIC_PREFIX + "diffusion_exec_s" +DIFFUSION_EXEC_PER_STEP_S = METRIC_PREFIX + "diffusion_exec_per_step_s" +DIFFUSION_PREPROCESS_S = METRIC_PREFIX + "diffusion_preprocess_s" +DIFFUSION_POSTPROCESS_S = METRIC_PREFIX + "diffusion_postprocess_s" +VAE_DECODE_S = METRIC_PREFIX + "vae_decode_s" +DENOISE_STEP_LATENCY_S = METRIC_PREFIX + DENOISE_STEP_LATENCY + "_s" +DIFFUSION_KV_LOAD_S = METRIC_PREFIX + "diffusion_kv_load_s" +IMAGE_TTFP_S = METRIC_PREFIX + "image_ttfp_s" +STAGE_IN_QUEUE_S = METRIC_PREFIX + "stage_in_queue_s" + + # ============================================================================ # Label sets # ============================================================================ @@ -163,6 +192,11 @@ SUCCESS_LABELS = ("model_name", "finished_reason") # ``stage`` + ``replica``. STAGE_LABELS = ("model_name", "stage", "replica") +STAGE_GEN_TIME_LABELS = ("model_name", "stage", "stage_type") +DIFFUSION_LABELS = ("model_name", "stage") +FAILED_LABELS = ("model_name", "reason") +KV_WAIT_LABELS = ("model_name", "connector_type") + # Audio continuity Counter carries an extra ``threshold_ms`` label so multiple # threshold buckets can be tracked simultaneously. The ``_ms`` suffix names a # numeric threshold *value* in ms, not a time-bearing metric. @@ -311,16 +345,6 @@ def compute_audio_frames( return pcm_nbytes // frame_width -def compute_denoise_step_latency(stage_gen_time: float, num_inference_steps: int) -> float: - """Mean denoise step latency = image stage generation time / step count. - - The returned value uses the same time unit as ``stage_gen_time``. - """ - if num_inference_steps <= 0: - return 0.0 - return stage_gen_time / float(num_inference_steps) - - # ============================================================================ # Audio sample-rate resolution # ============================================================================ diff --git a/vllm_omni/metrics/modality.py b/vllm_omni/metrics/modality.py index b83912ed3..e0d51e2b0 100644 --- a/vllm_omni/metrics/modality.py +++ b/vllm_omni/metrics/modality.py @@ -27,7 +27,7 @@ from typing import Any from prometheus_client import Counter, Histogram from vllm_omni.metrics import definitions as defs -from vllm_omni.metrics.utils import count_audio_frames, extract_mm_output +from vllm_omni.metrics.utils import observe_audio_finalize, observe_diffusion_finalize _stage_labels = list(defs.STAGE_LABELS) @@ -105,6 +105,40 @@ _diffusion_postprocess_family = Histogram( labelnames=_stage_labels, buckets=defs.SECONDS_FAST_BUCKETS, ) +_vae_decode_family = Histogram( + defs.VAE_DECODE_S, + "VAE decode latency in seconds (latents -> pixels/audio/video).", + labelnames=_stage_labels, + buckets=defs.SECONDS_BUCKETS, +) +_diffusion_forward_family = Histogram( + defs.DIFFUSION_FORWARD_S, + "Diffusion forward-only latency in seconds (denoise loop; excludes " + "preprocess / postprocess / VAE decode / KV load). Absent when the " + "pipeline profiler is off.", + labelnames=_stage_labels, + buckets=defs.SECONDS_BUCKETS, +) +_diffusion_kv_load_family = Histogram( + defs.DIFFUSION_KV_LOAD_S, + "Diffusion KV-recv latency in seconds (AR→diffusion KV fetch; absent when " + "the stage has no upstream KV to receive).", + labelnames=_stage_labels, + buckets=defs.SECONDS_FAST_BUCKETS, +) +_image_ttfp_family = Histogram( + defs.IMAGE_TTFP_S, + "Image time-to-first-output in seconds (stage submit → first image materialized; non-streaming single-image). " + "Stage-level, not e2e — excludes API queue and inter-stage transfer before the image stage.", + labelnames=_stage_labels, + buckets=defs.SECONDS_BUCKETS, +) +_denoise_step_latency_family = Histogram( + defs.DENOISE_STEP_LATENCY_S, + "Mean per-step denoise forward latency in seconds (forward_time / num_inference_steps).", + labelnames=_stage_labels, + buckets=defs.SECONDS_FAST_BUCKETS, +) class OmniModalityMetrics: @@ -206,6 +240,31 @@ class OmniModalityMetrics: replica=replica, ).observe(seconds) + def observe_vae_decode(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds < 0: + return + _vae_decode_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + + def observe_diffusion_forward(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds < 0: + return + _diffusion_forward_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + + def observe_diffusion_kv_load(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds < 0: + return + _diffusion_kv_load_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + + def observe_image_ttfp(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds < 0: + return + _image_ttfp_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + + def observe_denoise_step_latency(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds <= 0: + return + _denoise_step_latency_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + def observe_modality_at_finalize( mod_metrics: OmniModalityMetrics, @@ -216,61 +275,34 @@ def observe_modality_at_finalize( stage_metrics: Any, engine_outputs: Any, ) -> None: - """Route audio-path observations for a finalized request. + """Route per-modality observations for a finalized request. Used by ``omni_base._process_single_result`` inside the e2e_done finalize - guard so it fires once per request. Skips text path (covered by upstream - ``vllm:*{stage="thinker", ...}``) and any case where required inputs are - missing — caller should not need to pre-validate. + guard so it fires once per request. Text path falls through — covered by + upstream ``vllm:*{stage="thinker", ...}``. Caller should not need to + pre-validate; missing inputs are silently skipped. audio_ttfp is intentionally NOT observed here; it's emitted by the streaming hook at first-packet time, not at finalize. """ - if replica_id is None or stage_metrics is None or output_type is None: + if replica_id is None or stage_metrics is None: return - stage_label = str(stage_id) - replica_label = str(replica_id) - gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0 - mm_out = extract_mm_output(engine_outputs) + observe_diffusion_finalize( + mod_metrics, + stage_id=stage_id, + replica_id=replica_id, + stage_metrics=stage_metrics, + ) if output_type == "audio": - sample_rate = defs.resolve_audio_sample_rate(mm_out) - n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0) - if n_frames == 0: - n_frames = count_audio_frames(mm_out) - mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames) - duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0 - if duration_s > 0: - mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s) - mod_metrics.observe_audio_rtf( - stage_label, - replica_label, - defs.compute_audio_rtf(gen_time_s, duration_s), - ) - else: - mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data") - - dm = getattr(stage_metrics, "diffusion_metrics", None) - if dm: - _key_map = { - "diffusion_engine_exec_time_s": mod_metrics.observe_diffusion_exec, - "preprocess_time_s": mod_metrics.observe_diffusion_preprocess, - "postprocess_time_s": mod_metrics.observe_diffusion_postprocess, - } - for key, observe_fn in _key_map.items(): - val = dm.get(key) - if val is not None: - observe_fn(stage_label, replica_label, float(val)) - - exec_time = dm.get("diffusion_engine_exec_time_s") - num_steps = dm.get("num_inference_steps") - if exec_time is not None and num_steps and num_steps > 0: - mod_metrics.observe_diffusion_exec_per_step( - stage_label, - replica_label, - float(exec_time) / int(num_steps), - ) + observe_audio_finalize( + mod_metrics, + stage_id=stage_id, + replica_id=replica_id, + stage_metrics=stage_metrics, + engine_outputs=engine_outputs, + ) def observe_audio_first_packet( diff --git a/vllm_omni/metrics/prometheus.py b/vllm_omni/metrics/prometheus.py index d372c3d6b..911d21e64 100644 --- a/vllm_omni/metrics/prometheus.py +++ b/vllm_omni/metrics/prometheus.py @@ -38,6 +38,66 @@ _generation_tokens_family = Counter( labelnames=_labelnames, ) +_stage_gen_time_family = Histogram( + defs.STAGE_GEN_TIME_S, + "Per-stage generation time in seconds for diffusion / image stages.", + labelnames=list(defs.STAGE_GEN_TIME_LABELS), + buckets=defs.SECONDS_BUCKETS, +) +_request_queue_wait_family = Histogram( + defs.REQUEST_QUEUE_WAIT_S, + "Per-request queue wait between submit arrival and stage submit, in seconds.", + labelnames=_labelnames, + buckets=defs.SECONDS_BUCKETS, +) +_stage_waiting_requests_family = Gauge( + defs.STAGE_WAITING_REQUESTS, + "Sum of the latest scheduler waiting snapshots across a stage's replicas; " + "diffusion snapshots update when outputs arrive.", + labelnames=list(defs.DIFFUSION_LABELS), +) +_num_inference_steps_family = Histogram( + defs.NUM_INFERENCE_STEPS, + "Diffusion step count distribution per request.", + labelnames=_labelnames, +) +_image_count_family = Counter( + defs.IMAGE_COUNT_METRIC, + "Cumulative image count; per-model rate. Throughput recovered via rate().", + labelnames=_labelnames, +) +_image_pixels_family = Histogram( + defs.IMAGE_PIXELS_METRIC, + "Image pixel count distribution per request.", + labelnames=_labelnames, +) +_peak_memory_family = Gauge( + defs.PEAK_MEMORY_MB, + "Peak GPU memory in MB observed during a stage's generation.", + labelnames=list(defs.DIFFUSION_LABELS), +) +_stage_in_queue_family = Histogram( + defs.STAGE_IN_QUEUE_S, + "Time in seconds from diffusion scheduler enqueue to initial execution admission. " + "Zero is the healthy low-load baseline and is counted.", + labelnames=list(defs.DIFFUSION_LABELS), + buckets=defs.SECONDS_FAST_BUCKETS, +) + +_requests_failed_family = Counter( + defs.REQUESTS_FAILED, + "Total requests by failure reason. Pairs with requests_success_total " + "(abort bucket) to give the full failure surface. Reasons are bounded to " + "client_abort, client_disconnect, stage_error, and unknown.", + labelnames=list(defs.FAILED_LABELS), +) +_kv_wait_s_family = Histogram( + defs.KV_WAIT_S, + "KV cache transfer wait time per request, sliced by connector backend.", + labelnames=list(defs.KV_WAIT_LABELS), + buckets=defs.SECONDS_BUCKETS, +) + class OmniPrometheusMetrics: """Label-bound wrapper around the raw Prometheus metrics. @@ -54,6 +114,10 @@ class OmniPrometheusMetrics: self._e2e_latency = _e2e_latency_family.labels(model_name=model_name) self._prompt_tokens = _prompt_tokens_family.labels(model_name=model_name) self._generation_tokens = _generation_tokens_family.labels(model_name=model_name) + self._request_queue_wait = _request_queue_wait_family.labels(model_name=model_name) + self._num_inference_steps = _num_inference_steps_family.labels(model_name=model_name) + self._image_count = _image_count_family.labels(model_name=model_name) + self._image_pixels = _image_pixels_family.labels(model_name=model_name) def set_running(self, n: int) -> None: if not self._log_stats: @@ -89,13 +153,80 @@ class OmniPrometheusMetrics: def request_failed(self) -> None: if not self._log_stats: return - # Pipeline-level "fail" maps to the upstream FinishReason.ABORT bucket; - # a single counter family now covers both normal stops and aborts. _completion_family.labels( model_name=self._model_name, finished_reason="abort", ).inc() + def observe_stage_gen_time(self, stage: int, stage_type: str, gen_time_s: float) -> None: + if not self._log_stats: + return + _stage_gen_time_family.labels( + model_name=self._model_name, + stage=str(stage), + stage_type=stage_type, + ).observe(max(gen_time_s, 0.0)) + + def observe_stage_in_queue(self, stage: int, in_queue_s: float) -> None: + if not self._log_stats: + return + _stage_in_queue_family.labels( + model_name=self._model_name, + stage=str(stage), + ).observe(max(in_queue_s, 0.0)) + + def observe_queue_wait(self, queue_wait_s: float) -> None: + if not self._log_stats: + return + self._request_queue_wait.observe(max(queue_wait_s, 0.0)) + + def set_stage_waiting_requests(self, stage: int, n_waiting: int) -> None: + if not self._log_stats: + return + _stage_waiting_requests_family.labels( + model_name=self._model_name, + stage=str(stage), + ).set(max(n_waiting, 0)) + + def observe_num_inference_steps(self, n_steps: int) -> None: + if not self._log_stats or n_steps <= 0: + return + self._num_inference_steps.observe(n_steps) + + def inc_image_count(self, n_images: int = 1) -> None: + if not self._log_stats or n_images <= 0: + return + self._image_count.inc(n_images) + + def observe_image_pixels(self, n_pixels: int) -> None: + if not self._log_stats or n_pixels <= 0: + return + self._image_pixels.observe(n_pixels) + + def set_peak_memory(self, stage: int, peak_memory_mb: float) -> None: + if not self._log_stats: + return + _peak_memory_family.labels( + model_name=self._model_name, + stage=str(stage), + ).set(max(peak_memory_mb, 0.0)) + + def inc_requests_failed(self, reason: str) -> None: + if not self._log_stats: + return + _requests_failed_family.labels( + model_name=self._model_name, + reason=reason or "unknown", + ).inc() + + def observe_kv_wait(self, connector_type: str, kv_wait_s: float) -> None: + if not self._log_stats or kv_wait_s < 0: + return + _kv_wait_s_family.labels( + model_name=self._model_name, + connector_type=connector_type or "unknown", + ).observe(kv_wait_s) + class OmniRequestCounter: """Running-request counter written by the orchestrator thread, read by the client thread.""" diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 95198ef6a..51d424508 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import time @@ -49,11 +52,13 @@ class StageRequestStats: audio_sample_rate: int = 0 audio_duration_s: float = 0.0 image_pixels: int = 0 + num_inference_steps: int = 0 denoise_step_latency_ms: float = 0.0 pipeline_timings: dict[str, float] | None = None output_unit_type: str | None = None output_unit_count: int = 0 serving_time_to_first_output_ms: float = 0.0 + image_time_to_first_output_ms: float = 0.0 time_per_output_unit_ms: float = 0.0 inter_output_latency_ms: float = 0.0 inter_output_latencies_ms: list[float] | None = None @@ -480,6 +485,7 @@ class OrchestratorAggregator: "output_unit_type": evt.output_unit_type, defs.OUTPUT_UNIT_COUNT: int(evt.output_unit_count), defs.SERVING_TIME_TO_FIRST_OUTPUT_MS: float(evt.serving_time_to_first_output_ms), + defs.IMAGE_TIME_TO_FIRST_OUTPUT_MS: float(evt.image_time_to_first_output_ms), defs.TIME_PER_OUTPUT_UNIT_MS: float(evt.time_per_output_unit_ms), defs.INTER_OUTPUT_LATENCY_MS: float(evt.inter_output_latency_ms), defs.INTER_OUTPUT_LATENCIES_MS: list(evt.inter_output_latencies_ms or []), @@ -513,6 +519,11 @@ class OrchestratorAggregator: if current_first_output_ms <= 0 < first_output_ms: current[defs.SERVING_TIME_TO_FIRST_OUTPUT_MS] = first_output_ms + image_first_output_ms = float(evt.image_time_to_first_output_ms) + current_image_first_output_ms = float(current.get(defs.IMAGE_TIME_TO_FIRST_OUTPUT_MS, 0.0)) + if current_image_first_output_ms <= 0 < image_first_output_ms: + current[defs.IMAGE_TIME_TO_FIRST_OUTPUT_MS] = image_first_output_ms + if evt.output_unit_type: current["output_unit_type"] = evt.output_unit_type if evt.final_output_type: @@ -578,11 +589,20 @@ class OrchestratorAggregator: stats.request_id = req_id if final_output_type is not None: stats.final_output_type = final_output_type + diffusion_metrics_by_request = getattr(self, "diffusion_metrics", None) stats.diffusion_metrics = ( - {k: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()} - if req_id in self.diffusion_metrics + {k: float(v) for k, v in diffusion_metrics_by_request.pop(req_id, {}).items()} + if diffusion_metrics_by_request is not None and req_id in diffusion_metrics_by_request else None ) + forward_time_s = (stats.diffusion_metrics or {}).get("forward_time_s") + num_inference_steps = int(getattr(stats, "num_inference_steps", 0) or 0) + if num_inference_steps > 0 and getattr(stats, "output_unit_type", None) == "image": + if forward_time_s is not None: + denoise_time_ms = float(forward_time_s) * 1000.0 + else: + denoise_time_ms = float(getattr(stats, "stage_gen_time_ms", 0.0) or 0.0) + stats.denoise_step_latency_ms = denoise_time_ms / num_inference_steps return stats def on_stage_metrics( @@ -628,10 +648,24 @@ class OrchestratorAggregator: _postproc_ms = (time.perf_counter() - _t0) * 1000.0 self.record_stage_postprocess_time(stage_id, req_id, _postproc_ms) + _MS_TO_S: dict[str, str] = { + "preprocess_time_ms": "preprocess_time_s", + "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s", + "postprocess_time_ms": "postprocess_time_s", + "vae_decode_time_ms": "vae_decode_time_s", + "forward_time_ms": "forward_time_s", + "scheduler_queue_wait_ms": "scheduler_queue_wait_s", + "kv_recv_time_ms": "kv_recv_time_s", + } + 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. + Engine emits ``*_ms`` timings; the accumulator converts them to ``_s`` + keys via ``_MS_TO_S`` so downstream observers read a uniform + seconds-bearing dict. Per-chunk timing keys are summed; non-timing + keys (e.g. ``image_num`` / ``resolution`` from ``format_diffusion_outputs``) + preserve the existing ``+=`` semantics. Args: req_id: Request ID @@ -643,20 +677,16 @@ class OrchestratorAggregator: diffusion_metrics: dict = getattr(engine_output, "metrics", {}) if isinstance(diffusion_metrics, list): diffusion_metrics = diffusion_metrics[0] - if diffusion_metrics: - _MS_TO_S = { - "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s", - "preprocess_time_ms": "preprocess_time_s", - "postprocess_time_ms": "postprocess_time_s", - "diffusion_engine_total_time_ms": "diffusion_engine_total_time_s", - } - for key, value in diffusion_metrics.items(): - if value is None: - continue - if key in _MS_TO_S: - self.diffusion_metrics[req_id][_MS_TO_S[key]] = float(value) / 1000.0 - else: - self.diffusion_metrics[req_id][key] = value + if not diffusion_metrics: + return + bucket = self.diffusion_metrics[req_id] + for key, value in diffusion_metrics.items(): + if value is None: + continue + if key in self._MS_TO_S: + bucket[self._MS_TO_S[key]] += float(value) / 1000.0 + else: + bucket[key] += float(value) def on_forward( self, diff --git a/vllm_omni/metrics/utils.py b/vllm_omni/metrics/utils.py index ef8f69ed1..10cdab94a 100644 --- a/vllm_omni/metrics/utils.py +++ b/vllm_omni/metrics/utils.py @@ -1,9 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + from collections.abc import Callable, Mapping, Sequence from dataclasses import fields -from typing import Any +from typing import TYPE_CHECKING, Any from prettytable import PrettyTable +from vllm_omni.metrics import definitions as metric_defs + +if TYPE_CHECKING: + from vllm_omni.metrics.modality import OmniModalityMetrics + from vllm_omni.metrics.prometheus import OmniPrometheusMetrics + from vllm_omni.metrics.stats import StageRequestStats + +FAILURE_REASONS = frozenset({"client_abort", "client_disconnect", "stage_error", "unknown"}) +DIFFUSION_METRICS_ONLY_REQUEST_ID = "__vllm_omni_diffusion_metrics__" + + +def extract_queue_wait_s(pipeline_timings: Mapping[str, float] | None) -> float | None: + if pipeline_timings is None or "queue_wait_ms" not in pipeline_timings: + return None + return float(pipeline_timings["queue_wait_ms"] or 0.0) / 1000.0 + + +def observe_stage_workload_metrics( + prom_metrics: OmniPrometheusMetrics, + *, + stage_type: str, + stage_metrics: StageRequestStats, +) -> None: + if stage_type == "diffusion": + prom_metrics.observe_num_inference_steps(stage_metrics.num_inference_steps) + + if stage_metrics.output_unit_type == "image": + prom_metrics.observe_image_pixels(stage_metrics.image_pixels) + prom_metrics.inc_image_count(stage_metrics.output_unit_count) + + +def observe_audio_finalize( + mod_metrics: OmniModalityMetrics, + *, + stage_id: int, + replica_id: int, + stage_metrics: Any, + engine_outputs: Any, +) -> None: + stage_label = str(stage_id) + replica_label = str(replica_id) + gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0 + mm_out = extract_mm_output(engine_outputs) + + sample_rate = metric_defs.resolve_audio_sample_rate(mm_out) + n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0) + if n_frames == 0: + n_frames = count_audio_frames(mm_out) + mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames) + duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0 + if duration_s > 0: + mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s) + mod_metrics.observe_audio_rtf( + stage_label, + replica_label, + metric_defs.compute_audio_rtf(gen_time_s, duration_s), + ) + else: + mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data") + + +def observe_diffusion_finalize( + mod_metrics: OmniModalityMetrics, + *, + stage_id: int, + replica_id: int, + stage_metrics: Any, +) -> None: + diff_metrics = getattr(stage_metrics, "diffusion_metrics", None) + if not diff_metrics: + return + + stage_label = str(stage_id) + replica_label = str(replica_id) + + exec_s = diff_metrics.get("diffusion_engine_exec_time_s") + if exec_s is not None: + mod_metrics.observe_diffusion_exec(stage_label, replica_label, float(exec_s)) + num_steps = int( + getattr(stage_metrics, "num_inference_steps", 0) or diff_metrics.get("num_inference_steps") or 0 + ) + if num_steps > 0: + mod_metrics.observe_diffusion_exec_per_step(stage_label, replica_label, float(exec_s) / num_steps) + + pre_s = diff_metrics.get("preprocess_time_s") + if pre_s is not None: + mod_metrics.observe_diffusion_preprocess(stage_label, replica_label, float(pre_s)) + + post_s = diff_metrics.get("postprocess_time_s") + if post_s is not None: + mod_metrics.observe_diffusion_postprocess(stage_label, replica_label, float(post_s)) + + vae_s = diff_metrics.get("vae_decode_time_s") + if vae_s is not None: + mod_metrics.observe_vae_decode(stage_label, replica_label, float(vae_s)) + + forward_s = diff_metrics.get("forward_time_s") + if forward_s is not None: + mod_metrics.observe_diffusion_forward(stage_label, replica_label, float(forward_s)) + num_steps = int( + getattr(stage_metrics, "num_inference_steps", 0) or diff_metrics.get("num_inference_steps") or 0 + ) + if num_steps > 0 and getattr(stage_metrics, "output_unit_type", None) == "image": + mod_metrics.observe_denoise_step_latency(stage_label, replica_label, float(forward_s) / num_steps) + + kv_load_s = diff_metrics.get("kv_recv_time_s") + if kv_load_s is not None: + mod_metrics.observe_diffusion_kv_load(stage_label, replica_label, float(kv_load_s)) + + +def normalize_failure_reason(reason: str | None) -> str: + """Map request failure details to the bounded metrics taxonomy.""" + return reason if reason in FAILURE_REASONS else "unknown" + + +def diffusion_scheduler_waiting_metrics(n_waiting: int) -> dict[str, int]: + """Build the diffusion scheduler snapshot consumed by the orchestrator.""" + return {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: max(int(n_waiting), 0)} + + +def diffusion_exception_metrics(exc: BaseException) -> dict[str, Any]: + """Return metrics attached to a terminal diffusion error.""" + metrics = getattr(exc, "diffusion_metrics", None) + return dict(metrics) if isinstance(metrics, dict) else {} + + +def sum_diffusion_stage_durations_ms(output: Any, suffix: str) -> float | None: + """Sum matching diffusion stage durations in milliseconds.""" + durations = getattr(output, "stage_durations", None) + if not isinstance(durations, dict): + return None + + total = 0.0 + found = False + for key, value in durations.items(): + if not key.endswith(suffix): + continue + try: + total += float(value) + found = True + except (TypeError, ValueError): + continue + return total * 1000.0 if found else None + + +def extract_diffusion_vae_decode_ms(output: Any) -> float | None: + return sum_diffusion_stage_durations_ms(output, ".vae.decode") + + +def extract_diffusion_denoise_ms(output: Any) -> float | None: + """Extract denoise-loop time (transformer and per-step scheduler).""" + return sum_diffusion_stage_durations_ms(output, ".diffuse") + def coerce_positive_int_scalar(value: object) -> int | None: """Coerce a value to a positive int without importing tensor libs.