text and image metric (#6150)
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com> Signed-off-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
This commit is contained in:
33 files changed
+1802
-162
No files matched your search
+76
-30
@@ -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
|
||||
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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 == {}
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)):
|
||||
|
||||
@@ -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()):
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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))
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
# ============================================================================
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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."""
|
||||
|
||||
+47
-17
@@ -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,
|
||||
|
||||
+159
-1
@@ -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.
|
||||
|
||||
Reference in new issue
Block a user