diff --git a/docs/design/index.md b/docs/design/index.md index 61aa50483..8789f2e1b 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -13,6 +13,10 @@ This section contains design documents and architecture specifications for vLLM- - [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md) - [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md) +## Infrastructure Design Documents + +- [Prometheus Metrics](metrics.md) + ## Module Design Documents - [AR Module](module/ar_module.md) diff --git a/docs/design/metrics.md b/docs/design/metrics.md new file mode 100644 index 000000000..09f0907df --- /dev/null +++ b/docs/design/metrics.md @@ -0,0 +1,208 @@ +# Prometheus Metrics Design + +This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up. + +## Objectives + +- 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-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. + +## Background + +### Upstream vLLM Metrics + +Upstream vLLM defines ~37 Prometheus metric families under the `vllm:` prefix. These are registered by `PrometheusStatLogger` and cover engine-level state: KV cache usage, running/waiting request counts, token throughput, TTFT, inter-token latency, e2e latency, and so on. They are served via the `/metrics` HTTP endpoint provided by `prometheus_fastapi_instrumentator` and the default `prometheus_client` WSGI handler. + +vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` collector whose `_name` attribute contains the substring `"vllm"`. This runs during engine initialization to clean up stale collectors from prior instantiations within the same process. + +### 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`. + +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). + +## Architecture + +### Component Overview + +```text + +------------------------+ + | API Server (FastAPI) | + | GET /metrics | + +-----------+------------+ + | + prometheus_client default registry + | + +--------+--------+--------+--------+--------+ + | | + vllm:omni_* vllm:* + collectors collectors + | | + +----+--------+ +-----------+ +----------+ +-----------+ + | OmniPromet- | | OmniMod- | | OmniTra- | | OmniProm- | + | heusMetrics | | alityMet- | | nsferMe- | | etheusSt- | + | | | rics | | trics | | atLogger | + +----+--------+ +-----+-----+ +----+-----+ +----+------+ + | | | | + OmniBase OmniBase Orchestrator Orchestrator + (request life- (finalize + (record_trans- (per-(stage, + cycle, success/ streaming fer_tx/rx replica) + fail counter) hooks via hooks via scheduler/ + observe_* emit hook in iteration + APIs) OrchestratorAg- stats) + gregator) +``` + +### Data Flow + +There are four independent paths for metric collection. + +**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. + +**Path 2: 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: + +- `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path). +- `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_*`)** + +`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. + +The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transfer_rx`. After the existing `TransferEdgeStats` accumulation, the aggregator calls `_emit_transfer_tx` / `_emit_transfer_rx`. Those: + +1. Resolve `from_replica` / `to_replica` via a `replica_resolver` callback supplied by `async_omni.py`. The resolver delegates to `stage_pool.get_bound_replica_id(request_id)` — i.e. the orchestrator's existing sticky-routing binding is the source of truth. +2. Convert the underlying `_ms` accumulators to seconds and call the `_s`-suffixed observe methods on `OmniTransferMetrics`. + +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)** + +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. + +### Shared State Between Threads + +The Orchestrator runs in a background thread. The API server (OmniBase) runs in the asyncio event loop thread. `OmniRequestCounter` bridges them — a plain Python object with an `int` field. The Orchestrator increments/decrements it; the entrypoint reads it for gauge updates. No lock is needed because the counter is advisory (a stale read by one Prometheus scrape interval is acceptable). It is created by `AsyncOmniEngine.__init__()` and passed to the Orchestrator at construction time. + +### 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. + +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 + +Upstream `PrometheusStatLogger.__init__` hard-codes `labelnames = ["model_name", "engine"]` as a local variable, references it across ~37 metric-family construction sites, and uses the `engine` label value in five different `.labels()` call shapes (kwarg with int engine, kwarg with str engine, positional with str engine in the middle, plus a `metrics_info["engine"] = str(...)` dict pattern). To reshape `engine` into `stage` + `replica` without forking the entire upstream `__init__`, the wrap uses three coordinated mechanisms: + +1. **Class-level metric class slot overrides.** `OmniPrometheusStatLogger` overrides `_gauge_cls`, `_counter_cls`, `_histogram_cls` (which upstream calls via `self._gauge_cls(...)` etc.) with `_RelabelGauge` / `_RelabelCounter` / `_RelabelHistogram` wrapper classes. These intercept the `labelnames` kwarg at metric family creation time and replace `engine` with `("stage", "replica")`. +2. **Property descriptor for `per_engine_labelvalues`.** Upstream builds `self.per_engine_labelvalues = {idx: [model_name, str(idx)]}` inside `__init__` and then captures it into a local variable for `create_metric_per_engine` calls. By making `per_engine_labelvalues` a Python property on the subclass, the setter intercepts upstream's assignment and rewrites each 2-tuple into a 3-tuple `[model_name, stage, replica]` using the `stage_replica_map` supplied at construction time. The captured local then sees the rewritten dict. +3. **Override of `.labels()` on the wrapper classes.** For the five call sites that pass `engine` directly (kwarg or positional, int or str), `_RelabelMixin.labels()` translates the engine value back to `(stage, replica)` via a process-level `_ENGINE_INDEX_MAP` populated by `OmniPrometheusStatLogger.__init__`. This handles `gauge_engine_sleep_state.labels(engine=idx, ...)`, `counter_request_success_base.labels(model_name, str(idx), str(reason))`, `info_gauge.labels(**metrics_info)`, etc. + +The three sub-helpers that upstream `PrometheusStatLogger.__init__` constructs (`spec_decoding_prom` / `kv_connector_prom` / `perf_metrics_prom`) use their own `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots and would otherwise build families with the raw 2-element labelnames. `_OmniPerfMetricsProm` / `_OmniSpecDecodingProm` / `_OmniKVConnectorProm` subclass each helper to route the same relabel mixin through their internal family construction. + +The `Orchestrator` constructs `stage_replica_map` from the static `stage_pools` configuration at startup: + +```python +stage_replica_map = { + flat_idx: (str(stage_id), str(replica_id)) + for flat_idx, (stage_id, replica_id) in enumerate( + (s, r) + for s, pool in enumerate(stage_pools) + for r in range(pool.num_replicas) + ) +} +``` + +A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestrator so the per-replica `record(engine_idx=...)` call site can look up the right flat index. + +> Dynamic add/remove of replicas at runtime is intentionally out of scope — the upstream `PrometheusStatLogger` materializes per-engine_idx child metrics at init time, and supporting hot-add would require non-trivial intervention into upstream's per-family child dictionaries. + +## 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(...)`. + +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. +- `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. + +Behavior with `--log-stats=on`: all metric paths fire normally; the orchestrator's per-replica recording is bounded only by `OmniSchedulerMixin.make_stats()`'s per-scheduler 1 Hz throttle (see next section). + +The overhead with the flag on is small enough that an A/B benchmark on Qwen3-Omni-30B single replica (30 sequential audio requests) showed a mean latency delta of +0.6% (Welch's t = 0.318, n=30, not statistically significant at α=0.05); the `/metrics` line count drops from 1358 to 124 lines when the flag is off. + +## Throttling: `make_stats()` Override + +Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, returning a SchedulerStats object for the orchestrator. Under vLLM's architecture, this is fine. But since vLLM-Omni requires that the object be serialized and transferred over ZMQ, receiving a SchedulerStats object on every step can introduce unacceptable overhead to the system. + +`OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead. + +The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second. + +## Metric Definitions + +### Pipeline (4) + +| 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 | + +### Audio (7) + +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 | + +### Cross-stage transfer (4) + +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) | + +### LLM stage-level (wrapped `vllm:*`) + +After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — carries `{model_name, stage, replica}` labels. For the full upstream catalog see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md); note that metrics depending on features unsupported in vLLM-Omni (e.g. speculative decoding, LoRA) will not be available. + +## Naming Convention + +- All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used: + - `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values. + - `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`. + +## Logging vs. Prometheus + +`OrchestratorAggregator` (in `vllm_omni/metrics/stats.py`) is the logging-oriented metrics path. It collects detailed per-request, per-stage, and per-transfer statistics and prints formatted tables to the `INFO` log. This is designed for development and debugging — individual request traces, transfer bandwidth, inter-stage timing. + +`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics` form the Prometheus-oriented path. They record aggregate counters, gauges, and histograms suitable for time-series monitoring and alerting. Both paths share the same source data (`StageRequestStats`, `TransferEdgeStats`) — `OrchestratorAggregator.record_transfer_tx/rx` in particular calls both the existing accumulator code and the Prometheus emit hook in the same method body. The two consumption models can run simultaneously without coupling. + +The separation follows upstream vLLM's pattern of `LoggingStatLogger` vs. `PrometheusStatLogger` — same underlying data, different consumption models. diff --git a/docs/usage/metrics.md b/docs/usage/metrics.md new file mode 100644 index 000000000..68a38043e --- /dev/null +++ b/docs/usage/metrics.md @@ -0,0 +1,98 @@ +# Production Metrics + +vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR. + +```bash +vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats +curl http://localhost:8000/metrics +``` + +**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate. + +## Metric Namespaces + +| Prefix | Source | Present when | +|--------|--------|--------------| +| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent | +| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage | +| `http_` / `process_` | Uvicorn / Python runtime | Always | + +## Pipeline-Level Metrics (`vllm:omni_`) + +Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline. + +### Request counts + +| Metric | Type | Labels | Description | +|--------|------|--------|-------------| +| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) | +| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue | +| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | + +### Latency + +| Metric | Type | Labels | Description | +|--------|------|--------|-------------| +| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds | + +## Audio Modality Metrics (`vllm:omni_`) + +Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{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`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` | +| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` | +| `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 | + +The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition. + +## Cross-Stage Transfer Metrics (`vllm:omni_`) + +Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed. + +| 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 Engine Metrics (`vllm:`) + +When the pipeline includes an LLM stage, the upstream vLLM engine exposes its full set of ~37 metric families under the `vllm:` prefix. + +vLLM-Omni wraps the upstream `vllm.v1.metrics.loggers.PrometheusStatLogger` with `OmniPrometheusStatLogger` so that the original `engine` single label is reshaped into `stage` + `replica`. Every `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — therefore gains per-`(stage, replica)` visibility automatically. No omni-side duplicate is needed for the text path. + +```text +# Before wrap: +vllm:num_requests_running{model_name="...", engine="1"} 3.0 + +# After wrap: +vllm:num_requests_running{model_name="...", stage="1", replica="0"} 2.0 +vllm:num_requests_running{model_name="...", stage="1", replica="1"} 1.0 +``` + +For the full list of upstream metrics, see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md). + +## Metric Availability by Pipeline Type + +| Metric group | Multi-stage LLM (Qwen3-Omni) | +|---|---| +| `vllm:omni_` request tracking + latency | With `--log-stats` | +| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage | +| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages | +| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` | +| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` | + +## Naming Convention + +- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values. +- 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 (see `vllm_omni/patch.py`) 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`. +- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label. diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 9d474e81c..51be7d4ce 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -312,7 +312,7 @@ def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_inpu monkeypatch.setattr( engine_mod, "build_llm_stage_output_processor", - lambda plan, _cfg: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor, + lambda plan, _cfg, **_kw: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor, ) monkeypatch.setattr(engine_mod, "build_stage0_input_processor", lambda _cfg: top_level_input_processor) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 55e3edc41..0f6163b73 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -1157,3 +1157,37 @@ async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator assert stage0_r1.add_request_calls[1][0].request_id == "parent-neg" finally: await _shutdown_orchestrator(orchestrator_fixture) + + +def test_orchestrator_does_not_re_introduce_global_stats_throttle() -> None: + """Regression: each (stage, replica) must independently publish its wrapped + vllm:* stats when its scheduler emits non-None scheduler_stats. + + A previous version of Orchestrator carried a global self._last_stats_ts / + _stats_interval_s gate around _stat_logger.record(). Because + OmniSchedulerMixin.make_stats() already throttles at 1 Hz per scheduler + (one per (stage, replica)), the extra global gate starved every replica + other than the first to emit within each second — their {stage, replica} + gauges/counters went stale. + + The fix removed the global gate entirely; the only signal needed is + 'this replica's scheduler emitted non-None scheduler_stats'. This test + fails loudly if someone reintroduces the global throttle. + """ + import inspect + + from vllm_omni.engine.orchestrator import Orchestrator + + source = inspect.getsource(Orchestrator) + assert "_last_stats_ts" not in source, ( + "Orchestrator must not gate stat recording on a global timestamp. " + "OmniSchedulerMixin.make_stats() already throttles per scheduler " + "(per (stage, replica)); an outer global gate starves all but the " + "first replica to emit within each 1s window." + ) + assert "_stats_interval_s" not in source + assert "raw_outputs.scheduler_stats is not None" in source, ( + "Orchestrator must gate stat recording solely on " + "raw_outputs.scheduler_stats being non-None — the per-scheduler 1Hz " + "throttle in OmniSchedulerMixin.make_stats() is the only gate needed." + ) diff --git a/tests/metrics/test_modality.py b/tests/metrics/test_modality.py new file mode 100644 index 000000000..ba1785639 --- /dev/null +++ b/tests/metrics/test_modality.py @@ -0,0 +1,399 @@ +from __future__ import annotations + +import pytest +from prometheus_client import REGISTRY, generate_latest + +from vllm_omni.metrics import definitions as defs +from vllm_omni.metrics.modality import ( + OmniModalityMetrics, + observe_audio_first_packet, + observe_audio_streaming_finalize, + observe_modality_at_finalize, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +_MODEL = "test-modality-model" + + +@pytest.fixture(scope="module") +def mod() -> OmniModalityMetrics: + return OmniModalityMetrics(model_name=_MODEL) + + +def _sample_value(output: str, line_prefix: str) -> float | None: + for line in output.splitlines(): + if line.startswith(line_prefix): + return float(line.split()[-1]) + return None + + +# --------------------------------------------------------------------------- +# Family registration — 7 audio families served from modality.py +# --------------------------------------------------------------------------- + + +_EXPECTED_FAMILIES = [ + defs.AUDIO_TTFP_S, + defs.AUDIO_DURATION_S, + defs.AUDIO_RTF_METRIC, + defs.AUDIO_FRAMES_METRIC, + defs.AUDIO_UNDERRUN_S, + defs.AUDIO_CONTINUITY_OK_METRIC, + defs.AUDIO_SKIPPED_REQUESTS_METRIC, +] + + +class TestRegistration: + def test_all_locked_families_present(self, mod: OmniModalityMetrics) -> None: + # Trigger at least one observation per family so the registry exposes them. + mod.observe_audio_ttfp("s", "r", 0.1) + mod.observe_audio_duration("s", "r", 1.0) + mod.observe_audio_rtf("s", "r", 0.5) + mod.inc_audio_frames("s", "r", 1) + mod.observe_audio_underrun("s", "r", 0.01) + mod.inc_audio_continuity_ok("s", "r", 100) + mod.inc_audio_skipped("s", "r", "malformed_codec") + + out = generate_latest(REGISTRY).decode() + for name in _EXPECTED_FAMILIES: + assert f"# HELP {name}" in out, f"missing family: {name}" + + +# --------------------------------------------------------------------------- +# Audio observe API +# --------------------------------------------------------------------------- + + +class TestAudio: + def test_audio_ttfp_observed(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_ttfp", "0" + mod.observe_audio_ttfp(stage, replica, 0.42) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.AUDIO_TTFP_S}_count{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == 1.0 + + def test_audio_duration_observed(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_dur", "0" + mod.observe_audio_duration(stage, replica, 3.5) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.AUDIO_DURATION_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == 3.5 + + def test_audio_rtf_observed(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_rtf", "0" + mod.observe_audio_rtf(stage, replica, 0.45) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.AUDIO_RTF_METRIC}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == 0.45 + + def test_audio_frames_inc(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_frames", "0" + mod.inc_audio_frames(stage, replica, 240) + mod.inc_audio_frames(stage, replica, 60) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.AUDIO_FRAMES_METRIC}_total{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == 300.0 + + def test_audio_frames_zero_or_negative_skipped(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_zero", "0" + mod.inc_audio_frames(stage, replica, 0) + mod.inc_audio_frames(stage, replica, -5) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.AUDIO_FRAMES_METRIC}_total{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) is None + + def test_audio_underrun_observed(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_under", "0" + mod.observe_audio_underrun(stage, replica, 0.05) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.AUDIO_UNDERRUN_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == 0.05 + + def test_audio_underrun_negative_clamped_to_zero(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_under_neg", "0" + mod.observe_audio_underrun(stage, replica, -0.1) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.AUDIO_UNDERRUN_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == 0.0 + + def test_audio_continuity_ok_carries_threshold_label(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_cont", "0" + mod.inc_audio_continuity_ok(stage, replica, threshold_ms=100) + mod.inc_audio_continuity_ok(stage, replica, threshold_ms=200) + out = generate_latest(REGISTRY).decode() + # Counter family auto-suffixes _total at exposition. + v100 = _sample_value( + out, + f'{defs.AUDIO_CONTINUITY_OK_METRIC}_total{{model_name="{_MODEL}",' + f'replica="{replica}",stage="{stage}",threshold_ms="100"}}', + ) + v200 = _sample_value( + out, + f'{defs.AUDIO_CONTINUITY_OK_METRIC}_total{{model_name="{_MODEL}",' + f'replica="{replica}",stage="{stage}",threshold_ms="200"}}', + ) + assert v100 == 1.0 + assert v200 == 1.0 + + def test_audio_skipped_carries_reason_label(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_skip", "0" + mod.inc_audio_skipped(stage, replica, "malformed_codec") + mod.inc_audio_skipped(stage, replica, "") # → "unknown" + out = generate_latest(REGISTRY).decode() + good = _sample_value( + out, + f'{defs.AUDIO_SKIPPED_REQUESTS_METRIC}_total{{model_name="{_MODEL}",' + f'reason="malformed_codec",replica="{replica}",stage="{stage}"}}', + ) + unknown = _sample_value( + out, + f'{defs.AUDIO_SKIPPED_REQUESTS_METRIC}_total{{model_name="{_MODEL}",' + f'reason="unknown",replica="{replica}",stage="{stage}"}}', + ) + assert good == 1.0 + assert unknown == 1.0 + + +# --------------------------------------------------------------------------- +# Stub-driven routing tests for the finalize / streaming helpers. +# --------------------------------------------------------------------------- + + +class _StubModMetrics: + def __init__(self): + self.calls: list[tuple] = [] + + def inc_audio_frames(self, s, r, n): + self.calls.append(("inc_audio_frames", s, r, n)) + + def observe_audio_duration(self, s, r, d): + self.calls.append(("observe_audio_duration", s, r, d)) + + def observe_audio_rtf(self, s, r, rtf): + self.calls.append(("observe_audio_rtf", s, r, rtf)) + + def observe_audio_underrun(self, s, r, u): + self.calls.append(("observe_audio_underrun", s, r, u)) + + def inc_audio_continuity_ok(self, s, r, threshold_ms): + self.calls.append(("inc_audio_continuity_ok", s, r, threshold_ms)) + + def inc_audio_skipped(self, s, r, reason): + self.calls.append(("inc_audio_skipped", s, r, reason)) + + +class _Bag: + def __init__(self, **kw): + self.__dict__.update(kw) + + +class TestObserveModalityAtFinalize: + def test_audio_path_full(self): + stub = _StubModMetrics() + stage_metrics = _Bag(stage_gen_time_ms=500.0, audio_generated_frames=24000) + engine_outputs = _Bag(multimodal_output={"audio_sample_rate": 24000}) + + observe_modality_at_finalize( + stub, + output_type="audio", + stage_id=1, + replica_id=0, + stage_metrics=stage_metrics, + engine_outputs=engine_outputs, + ) + # 24000 frames / 24000 Hz = 1.0s duration; gen 0.5s → rtf 0.5 + assert ("inc_audio_frames", "1", "0", 24000) in stub.calls + assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls + assert ("observe_audio_rtf", "1", "0", 0.5) in stub.calls + # Continuity/underrun NOT emitted from finalize — they come from the + # streaming hook because they need the per-chunk arrival timeline. + assert not any(c[0] == "observe_audio_underrun" for c in stub.calls) + assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls) + + def test_audio_path_zero_frames_emits_skipped(self): + stub = _StubModMetrics() + observe_modality_at_finalize( + stub, + output_type="audio", + stage_id=1, + replica_id=0, + stage_metrics=_Bag(stage_gen_time_ms=300.0, audio_generated_frames=0), + engine_outputs=_Bag(multimodal_output={}), + ) + assert ("inc_audio_frames", "1", "0", 0) in stub.calls + assert not any(c[0] == "observe_audio_duration" for c in stub.calls) + assert not any(c[0] == "observe_audio_rtf" for c in stub.calls) + assert ("inc_audio_skipped", "1", "0", "no_audio_data") in stub.calls + + def test_audio_uses_resolved_sample_rate_from_multimodal_output(self): + stub = _StubModMetrics() + observe_modality_at_finalize( + stub, + output_type="audio", + stage_id=1, + replica_id=0, + stage_metrics=_Bag(stage_gen_time_ms=1000.0, audio_generated_frames=16000), + engine_outputs=_Bag(multimodal_output={"sample_rate": 16000}), + ) + assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls + + def test_non_audio_output_type_skipped(self): + # Image / video / text output types are out of scope for this module. + stub = _StubModMetrics() + for output_type in ("text", "image", "video"): + observe_modality_at_finalize( + stub, + output_type=output_type, + stage_id=0, + replica_id=0, + stage_metrics=_Bag(stage_gen_time_ms=100.0), + engine_outputs=_Bag(), + ) + assert stub.calls == [] + + def test_replica_id_none_skipped(self): + stub = _StubModMetrics() + observe_modality_at_finalize( + stub, + output_type="audio", + stage_id=1, + replica_id=None, + stage_metrics=_Bag(stage_gen_time_ms=500.0, audio_generated_frames=240), + engine_outputs=_Bag(multimodal_output={}), + ) + assert stub.calls == [] + + def test_stage_metrics_none_skipped(self): + stub = _StubModMetrics() + observe_modality_at_finalize( + stub, + output_type="audio", + stage_id=1, + replica_id=0, + stage_metrics=None, + engine_outputs=_Bag(multimodal_output={}), + ) + assert stub.calls == [] + + +class TestObserveAudioFirstPacket: + def test_observes_with_valid_inputs(self): + stub = _StubModMetrics() + stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) + + observe_audio_first_packet( + stub, + stage_id=1, + replica_id=0, + arrival_ts=100.0, + now_ts=100.42, + ) + assert stub.calls == [("observe_audio_ttfp", "1", "0", pytest.approx(0.42))] + + def test_replica_none_skipped(self): + stub = _StubModMetrics() + stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) + observe_audio_first_packet(stub, stage_id=1, replica_id=None, arrival_ts=100.0, now_ts=100.5) + assert stub.calls == [] + + def test_arrival_ts_zero_skipped(self): + stub = _StubModMetrics() + stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) + observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=0.0, now_ts=100.5) + assert stub.calls == [] + + def test_clock_skew_clamped_to_zero(self): + stub = _StubModMetrics() + stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) + observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=100.5, now_ts=100.0) + assert stub.calls == [("observe_audio_ttfp", "1", "0", 0.0)] + + +class TestObserveAudioStreamingFinalize: + def test_continuous_run_emits_underrun_and_continuity_ok(self): + stub = _StubModMetrics() + # Three back-to-back chunks of 0.5s each at 24 kHz s16le mono: + # 24000 Hz * 2 bytes * 0.5s = 24000 bytes per chunk. + arrivals = [0.5, 1.0, 1.5] + bytes_ = [24000, 24000, 24000] + observe_audio_streaming_finalize( + stub, + stage_id=1, + replica_id=0, + chunk_arrival_times_s=arrivals, + chunk_bytes=bytes_, + sample_rate=24000, + threshold_s=0.1, + ) + # 0 underrun -> observe with 0.0, continuity_ok incremented at threshold 100. + assert ("observe_audio_underrun", "1", "0", 0.0) in stub.calls + assert ("inc_audio_continuity_ok", "1", "0", 100) in stub.calls + + def test_late_chunk_emits_nonzero_underrun_and_no_continuity_inc(self): + stub = _StubModMetrics() + # Two chunks but second arrives 1s late; underrun > 0.1s threshold. + arrivals = [0.0, 2.0] + bytes_ = [24000, 24000] + observe_audio_streaming_finalize( + stub, + stage_id=1, + replica_id=0, + chunk_arrival_times_s=arrivals, + chunk_bytes=bytes_, + sample_rate=24000, + threshold_s=0.1, + ) + underrun_calls = [c for c in stub.calls if c[0] == "observe_audio_underrun"] + assert underrun_calls + assert underrun_calls[0][-1] > 0.1 + assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls) + + def test_empty_arrivals_skipped(self): + stub = _StubModMetrics() + observe_audio_streaming_finalize( + stub, + stage_id=1, + replica_id=0, + chunk_arrival_times_s=[], + chunk_bytes=[], + sample_rate=24000, + ) + assert stub.calls == [] + + def test_replica_none_skipped(self): + stub = _StubModMetrics() + observe_audio_streaming_finalize( + stub, + stage_id=1, + replica_id=None, + chunk_arrival_times_s=[0.5], + chunk_bytes=[24000], + sample_rate=24000, + ) + assert stub.calls == [] + + +class TestBucketSelection: + def test_audio_rtf_uses_rtf_buckets(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_buckets", "0" + mod.observe_audio_rtf(stage, replica, 0.5) + out = generate_latest(REGISTRY).decode() + rtf_marker = f'{defs.AUDIO_RTF_METRIC}_bucket{{le="0.9"' + assert rtf_marker in out, "audio_rtf should use RTF_BUCKETS containing le=0.9" + + def test_audio_ttfp_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_seconds", "0" + mod.observe_audio_ttfp(stage, replica, 0.1) + out = generate_latest(REGISTRY).decode() + sec_marker = f'{defs.AUDIO_TTFP_S}_bucket{{le="0.05"' + assert sec_marker in out, "audio_ttfp should use SECONDS_BUCKETS containing le=0.05" + + def test_audio_underrun_uses_seconds_fast_buckets(self, mod: OmniModalityMetrics) -> None: + stage, replica = "talker_fast", "0" + mod.observe_audio_underrun(stage, replica, 0.005) + out = generate_latest(REGISTRY).decode() + # 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" diff --git a/tests/metrics/test_prometheus.py b/tests/metrics/test_prometheus.py new file mode 100644 index 000000000..63becf036 --- /dev/null +++ b/tests/metrics/test_prometheus.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import re + +import pytest +from prometheus_client import REGISTRY, CollectorRegistry, generate_latest + +from vllm_omni.metrics import OmniPrometheusMetrics + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_MODEL = "test-model" + +_PIPELINE_METRICS = [ + "vllm:omni_num_requests_running", + "vllm:omni_num_requests_waiting", + "vllm:omni_requests_success", + "vllm:omni_e2e_request_latency_s", +] + + +@pytest.fixture(scope="module") +def registry() -> CollectorRegistry: + return REGISTRY + + +@pytest.fixture(scope="module") +def prom() -> OmniPrometheusMetrics: + return OmniPrometheusMetrics(model_name=_MODEL) + + +@pytest.fixture(scope="module") +def scrape_output(prom: OmniPrometheusMetrics, registry: CollectorRegistry) -> str: + # Two natural completions (stop) + one length-cap + one failure (abort) + # exercise three distinct finished_reason buckets in the merged Counter. + prom.request_succeeded(e2e_seconds=1.5, finished_reason="stop") + prom.request_succeeded(e2e_seconds=2.0, finished_reason="stop") + prom.request_succeeded(e2e_seconds=3.0, finished_reason="length") + prom.request_failed() # → finished_reason="abort" + prom.set_running(5) + prom.set_waiting(2) + return generate_latest(registry).decode() + + +def _sample_value(output: str, metric_line: str) -> float | None: + for line in output.splitlines(): + if line.startswith(metric_line): + return float(line.split()[-1]) + return None + + +class TestMetricObservation: + def test_all_metric_families_present(self, scrape_output: str) -> None: + for name in _PIPELINE_METRICS: + assert f"# HELP {name}" in scrape_output, f"missing metric family: {name}" + + def test_counter_values(self, scrape_output: str) -> None: + # Per-reason buckets sourced from the merged completion Counter. + stop = _sample_value( + scrape_output, + f'vllm:omni_requests_success_total{{finished_reason="stop",model_name="{_MODEL}"}}', + ) + assert stop == 2.0 + + length = _sample_value( + scrape_output, + f'vllm:omni_requests_success_total{{finished_reason="length",model_name="{_MODEL}"}}', + ) + assert length == 1.0 + + abort = _sample_value( + scrape_output, + f'vllm:omni_requests_success_total{{finished_reason="abort",model_name="{_MODEL}"}}', + ) + assert abort == 1.0 + + def test_gauge_values(self, scrape_output: str) -> None: + running = _sample_value( + scrape_output, + f'vllm:omni_num_requests_running{{model_name="{_MODEL}"}}', + ) + assert running == 5.0 + + waiting = _sample_value( + scrape_output, + f'vllm:omni_num_requests_waiting{{model_name="{_MODEL}"}}', + ) + assert waiting == 2.0 + + def test_histogram_counts(self, scrape_output: str) -> None: + # 3 successful completions (stop x2 + length x1) all observe e2e; + # the 1 failed completion only increments the Counter without + # observing the latency histogram, so the count stays at 3. + e2e_count = _sample_value( + scrape_output, + f'vllm:omni_e2e_request_latency_s_count{{model_name="{_MODEL}"}}', + ) + assert e2e_count == 3.0 + + +class TestLabelCorrectness: + def test_pipeline_metrics_carry_model_name(self, scrape_output: str) -> None: + for name in _PIPELINE_METRICS: + pattern = rf'^{re.escape(name)}.*model_name="{re.escape(_MODEL)}"' + assert re.search(pattern, scrape_output, re.MULTILINE), f"{name} missing model_name label" + + def test_no_legacy_engine_label(self, scrape_output: str) -> None: + assert 'engine="' not in scrape_output + + def test_no_legacy_seconds_or_ms_families(self, scrape_output: str) -> None: + # Renamed: *_time_ms → *_s; *_time_seconds dropped. + for legacy in ( + "vllm:omni_request_queue_time_seconds", + "vllm:omni_e2e_request_latency_seconds", + ): + assert legacy not in scrape_output, f"legacy family {legacy} still registered" + + +class TestScrapeOutput: + def test_omni_metrics_in_default_registry(self, scrape_output: str) -> None: + for name in _PIPELINE_METRICS: + assert name in scrape_output + + def test_process_metrics_in_default_registry(self, scrape_output: str) -> None: + # vllm:* metrics require a full PrometheusStatLogger with VllmConfig + # and are registered by the Orchestrator at server startup. Verifying + # their presence is covered by integration tests. Here we confirm the + # default registry is being scraped by checking for process_* metrics + # from the Python prometheus_client runtime. + assert "process_" in scrape_output + + +class TestRequestLifecycleGauges: + """Regression tests for the running/waiting gauge race at request finalize. + + Before the fix, the per-stage publish in OmniBase._process_single_result + ran while the completed request was still in self.request_states even + though the orchestrator had already decremented _running_counter — so a + single completed request briefly exposed running=0, total=1 → waiting=1, + which stuck on /metrics until another request triggered a refresh. + + Two-layer fix: _process_single_result excludes the finalizing request + from `total`, and _log_summary_and_cleanup republishes gauges after the + pop as a fallback. These tests pin the post-cleanup state. + """ + + def test_running_and_waiting_zero_after_request_completes(self, registry: CollectorRegistry) -> None: + from types import SimpleNamespace + + from vllm_omni.entrypoints.omni_base import OmniBase + from vllm_omni.metrics.prometheus import OmniPrometheusMetrics, OmniRequestCounter + + obj = object.__new__(OmniBase) + obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter()) + obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test") + obj.request_states = {} + obj.log_stats = False + + # Simulate request lifecycle: start (counter 0→1, dict {} → {req}), + # then finalize (counter 1→0, dict still holds the request because + # _log_summary_and_cleanup hasn't run yet — this is the exact race + # window the fix addresses). + obj.engine._running_counter.increment() + obj.request_states["req-1"] = SimpleNamespace(metrics=SimpleNamespace(e2e_done={"req-1"})) + obj.engine._running_counter.decrement() + + obj._log_summary_and_cleanup("req-1") + + assert obj.engine._running_counter.value == 0 + assert len(obj.request_states) == 0 + out = generate_latest(registry).decode() + assert _sample_value(out, 'vllm:omni_num_requests_running{model_name="lifecycle-test"}') == 0.0 + assert _sample_value(out, 'vllm:omni_num_requests_waiting{model_name="lifecycle-test"}') == 0.0 + + def test_gauges_reflect_remaining_requests_after_one_completes(self, registry: CollectorRegistry) -> None: + from types import SimpleNamespace + + from vllm_omni.entrypoints.omni_base import OmniBase + from vllm_omni.metrics.prometheus import OmniPrometheusMetrics, OmniRequestCounter + + obj = object.__new__(OmniBase) + obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter()) + obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-2") + obj.request_states = {} + obj.log_stats = False + + # Two in flight, one finalizes — running should report 1, waiting 0. + obj.engine._running_counter.increment() + obj.engine._running_counter.increment() + obj.request_states["req-1"] = SimpleNamespace(metrics=SimpleNamespace(e2e_done={"req-1"})) + obj.request_states["req-2"] = SimpleNamespace(metrics=SimpleNamespace(e2e_done=set())) + obj.engine._running_counter.decrement() + + obj._log_summary_and_cleanup("req-1") + + out = generate_latest(registry).decode() + assert _sample_value(out, 'vllm:omni_num_requests_running{model_name="lifecycle-test-2"}') == 1.0 + assert _sample_value(out, 'vllm:omni_num_requests_waiting{model_name="lifecycle-test-2"}') == 0.0 diff --git a/tests/metrics/test_stat_logger.py b/tests/metrics/test_stat_logger.py new file mode 100644 index 000000000..3591579d8 --- /dev/null +++ b/tests/metrics/test_stat_logger.py @@ -0,0 +1,456 @@ +from __future__ import annotations + +import pytest +from prometheus_client import CollectorRegistry, generate_latest + +from vllm_omni.metrics.stat_logger import ( + _ENGINE_INDEX_MAP, + OmniPrometheusStatLogger, + _engine_to_stage_replica, + _OmniKVConnectorProm, + _OmniPerfMetricsProm, + _OmniSpecDecodingProm, + _RelabelCounter, + _RelabelGauge, + _RelabelHistogram, + _rewrite_labelnames, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _isolate_engine_map(): + """Each test gets a clean _ENGINE_INDEX_MAP.""" + _ENGINE_INDEX_MAP.clear() + yield + _ENGINE_INDEX_MAP.clear() + + +@pytest.fixture +def registry() -> CollectorRegistry: + return CollectorRegistry() + + +# --------------------------------------------------------------------------- +# _rewrite_labelnames +# --------------------------------------------------------------------------- + + +class TestRewriteLabelnames: + def test_engine_at_end(self): + assert _rewrite_labelnames(["model_name", "engine"]) == [ + "model_name", + "stage", + "replica", + ] + + def test_engine_in_middle(self): + # Upstream uses `labelnames + ["reason"]` etc., putting engine in middle. + assert _rewrite_labelnames(["model_name", "engine", "reason"]) == [ + "model_name", + "stage", + "replica", + "reason", + ] + + def test_no_engine_label(self): + # Unaffected (e.g. omni's own families that don't use engine). + assert _rewrite_labelnames(["model_name"]) == ["model_name"] + + def test_tuple_input_returns_tuple(self): + out = _rewrite_labelnames(("model_name", "engine")) + assert isinstance(out, tuple) + assert out == ("model_name", "stage", "replica") + + def test_none_passthrough(self): + assert _rewrite_labelnames(None) is None + + +# --------------------------------------------------------------------------- +# _engine_to_stage_replica +# --------------------------------------------------------------------------- + + +class TestEngineToStageReplica: + def test_int_engine_value(self): + # Mirrors upstream `.labels(engine=idx, ...)` with int (loggers.py:510). + _ENGINE_INDEX_MAP[7] = ("talker", "1") + assert _engine_to_stage_replica(7) == ("talker", "1") + + def test_str_engine_value(self): + # Mirrors upstream `metrics_info["engine"] = str(idx)` (loggers.py:1055). + _ENGINE_INDEX_MAP[2] = ("thinker", "0") + assert _engine_to_stage_replica("2") == ("thinker", "0") + + def test_missing_engine_idx_raises(self): + # Empty map → fail-fast rather than emit a wrong (stage, replica). + with pytest.raises(KeyError): + _engine_to_stage_replica(999) + + +# --------------------------------------------------------------------------- +# Wrapper class behavior +# --------------------------------------------------------------------------- + + +class TestRelabelGauge: + def test_labelnames_rewritten_at_creation(self, registry): + g = _RelabelGauge( + name="omni_test_gauge", + documentation="test", + labelnames=["model_name", "engine"], + registry=registry, + ) + assert g._labelnames == ("model_name", "stage", "replica") + + def test_labels_kwarg_translated(self, registry): + _ENGINE_INDEX_MAP[5] = ("diffusion", "0") + g = _RelabelGauge( + name="omni_test_gauge_kwarg", + documentation="test", + labelnames=["model_name", "engine"], + registry=registry, + ) + g.labels(engine=5, model_name="qwen-omni").set(42.0) + + out = generate_latest(registry).decode() + assert 'omni_test_gauge_kwarg{model_name="qwen-omni",replica="0",stage="diffusion"} 42.0' in out + + def test_labels_positional_passthrough(self, registry): + # Double-rewrite guard: the per_engine_labelvalues setter rewrites + # the values to 3-tuple [model_name, stage, replica] BEFORE + # create_metric_per_engine fans them into .labels(*values). The mixin + # must detect that args length already matches the rewritten + # labelnames and pass through, otherwise it would re-interpret + # args[engine_label_index] as an engine_idx and splice (stage, replica) + # again, blowing label count to 4. + g = _RelabelGauge( + name="omni_test_gauge_pos", + documentation="test", + labelnames=["model_name", "engine"], + registry=registry, + ) + g.labels("qwen-omni", "thinker", "0").set(7.0) + + out = generate_latest(registry).decode() + assert 'omni_test_gauge_pos{model_name="qwen-omni",replica="0",stage="thinker"} 7.0' in out + + def test_multiprocess_mode_kwarg_passthrough(self, registry): + # Upstream creates Gauges with multiprocess_mode="mostrecent" — must not + # be eaten by our mixin. + g = _RelabelGauge( + name="omni_test_gauge_mp", + documentation="test", + labelnames=["model_name", "engine"], + multiprocess_mode="mostrecent", + registry=registry, + ) + assert g._multiprocess_mode == "mostrecent" + + +class TestRelabelCounter: + def test_labelnames_rewritten(self, registry): + c = _RelabelCounter( + name="omni_test_counter", + documentation="test", + labelnames=["model_name", "engine", "finished_reason"], + registry=registry, + ) + assert c._labelnames == ( + "model_name", + "stage", + "replica", + "finished_reason", + ) + + def test_labels_kwarg_translated(self, registry): + _ENGINE_INDEX_MAP[2] = ("thinker", "0") + c = _RelabelCounter( + name="omni_test_counter_kwarg", + documentation="test", + labelnames=["model_name", "engine", "finished_reason"], + registry=registry, + ) + c.labels(engine=2, model_name="m", finished_reason="stop").inc(3) + + out = generate_latest(registry).decode() + assert ( + 'omni_test_counter_kwarg_total{finished_reason="stop",model_name="m",replica="0",stage="thinker"} 3.0' + in out + ) + + +class TestRelabelHistogram: + def test_labelnames_rewritten(self, registry): + h = _RelabelHistogram( + name="omni_test_histo", + documentation="test", + labelnames=["model_name", "engine"], + registry=registry, + ) + assert h._labelnames == ("model_name", "stage", "replica") + + def test_labels_kwarg_translated_and_observe(self, registry): + _ENGINE_INDEX_MAP[0] = ("talker", "0") + h = _RelabelHistogram( + name="omni_test_histo_obs", + documentation="test", + labelnames=["model_name", "engine"], + registry=registry, + ) + h.labels(engine=0, model_name="m").observe(0.5) + + out = generate_latest(registry).decode() + assert 'omni_test_histo_obs_count{model_name="m",replica="0",stage="talker"} 1.0' in out + + def test_no_engine_label_unaffected(self, registry): + # Families without engine label (e.g. omni-side own metrics) pass through. + h = _RelabelHistogram( + name="omni_test_no_engine", + documentation="test", + labelnames=["model_name"], + registry=registry, + ) + assert h._labelnames == ("model_name",) + h.labels(model_name="m").observe(1.0) + + +# --------------------------------------------------------------------------- +# Positional .labels() with engine value (loggers.py:646, 679) +# --------------------------------------------------------------------------- + + +class TestPositionalEngine: + def test_positional_engine_at_middle_index(self, registry): + # Mirrors `counter_prompt_tokens_by_source.labels(model_name, str(idx), source)`. + # Family original labelnames = ["model_name", "engine", "source"]. + _ENGINE_INDEX_MAP[5] = ("talker", "0") + c = _RelabelCounter( + name="omni_test_pos_mid", + documentation="test", + labelnames=["model_name", "engine", "source"], + registry=registry, + ) + c.labels("m", "5", "decoder").inc(2) + + out = generate_latest(registry).decode() + assert 'omni_test_pos_mid_total{model_name="m",replica="0",source="decoder",stage="talker"} 2.0' in out + + def test_positional_engine_with_int_value(self, registry): + # Defensive: positional form may also receive an int (we accept both). + _ENGINE_INDEX_MAP[3] = ("thinker", "1") + c = _RelabelCounter( + name="omni_test_pos_int", + documentation="test", + labelnames=["model_name", "engine", "reason"], + registry=registry, + ) + c.labels("m", 3, "stop").inc() + + out = generate_latest(registry).decode() + assert 'omni_test_pos_int_total{model_name="m",reason="stop",replica="1",stage="thinker"} 1.0' in out + + +# --------------------------------------------------------------------------- +# String-form engine kwarg (loggers.py:1056 info_gauge) +# --------------------------------------------------------------------------- + + +class TestStrEngineKwarg: + def test_engine_kwarg_str_form(self, registry): + # Mirrors `info_gauge.labels(**metrics_info)` where metrics_info["engine"]="0". + _ENGINE_INDEX_MAP[0] = ("thinker", "0") + g = _RelabelGauge( + name="omni_test_info", + documentation="test", + labelnames=["cache_size", "engine"], + multiprocess_mode="mostrecent", + registry=registry, + ) + # Upstream pattern: pass everything as kwargs from the metrics_info dict. + g.labels(cache_size="big", engine="0").set(1) + + out = generate_latest(registry).decode() + assert 'omni_test_info{cache_size="big",replica="0",stage="thinker"} 1.0' in out + + +# --------------------------------------------------------------------------- +# Child metric does not re-trigger relabel logic +# --------------------------------------------------------------------------- + + +class TestChildNoRecursion: + def test_child_set_does_not_relookup(self, registry): + # Once .labels() returns a child, subsequent .set()/.inc() must not + # consult _ENGINE_INDEX_MAP again. We verify by clearing the map + # AFTER labels() and proving .set() still works. + _ENGINE_INDEX_MAP[4] = ("diffusion", "0") + g = _RelabelGauge( + name="omni_test_child", + documentation="test", + labelnames=["model_name", "engine"], + registry=registry, + ) + child = g.labels(engine=4, model_name="m") + _ENGINE_INDEX_MAP.clear() # would break a second .labels() lookup + child.set(99.0) # but set() is on the bound child — no map needed + + # Re-populate so generate_latest doesn't trip on anything else. + _ENGINE_INDEX_MAP[4] = ("diffusion", "0") + out = generate_latest(registry).decode() + assert 'omni_test_child{model_name="m",replica="0",stage="diffusion"} 99.0' in out + + +# --------------------------------------------------------------------------- +# OmniPrometheusStatLogger — focused on the wrap mechanics (full +# PrometheusStatLogger init requires a real VllmConfig and is exercised by +# the orchestrator integration tests). +# --------------------------------------------------------------------------- + + +class TestOmniPrometheusStatLogger: + def test_class_slots_point_to_wrappers(self): + # Upstream's __init__ uses self._gauge_cls(...) etc. when constructing + # families; class-level slot override is how we inject the relabel logic. + assert OmniPrometheusStatLogger._gauge_cls is _RelabelGauge + assert OmniPrometheusStatLogger._counter_cls is _RelabelCounter + assert OmniPrometheusStatLogger._histogram_cls is _RelabelHistogram + + def test_per_engine_labelvalues_setter_rewrites_to_3tuple(self): + # Construct via __new__ to skip the upstream PrometheusStatLogger __init__ + # (which needs a real VllmConfig). We only verify the property descriptor. + sl = OmniPrometheusStatLogger.__new__(OmniPrometheusStatLogger) + sl._stage_replica_map = { + 0: ("thinker", "0"), + 1: ("talker", "0"), + 2: ("talker", "1"), + } + + # Mirror upstream's loggers.py:433 assignment shape. + sl.per_engine_labelvalues = { + 0: ["my-model", "0"], + 1: ["my-model", "1"], + 2: ["my-model", "2"], + } + + # Getter should return the 3-tuple form for downstream + # create_metric_per_engine consumers. + assert sl.per_engine_labelvalues == { + 0: ["my-model", "thinker", "0"], + 1: ["my-model", "talker", "0"], + 2: ["my-model", "talker", "1"], + } + + def test_per_engine_labelvalues_getter_returns_internal_dict(self): + sl = OmniPrometheusStatLogger.__new__(OmniPrometheusStatLogger) + sl._stage_replica_map = {0: ("thinker", "0")} + sl._omni_per_engine_labelvalues = {0: ["m", "thinker", "0"]} + assert sl.per_engine_labelvalues == {0: ["m", "thinker", "0"]} + + def test_stage_replica_map_property_exposed(self): + sl = OmniPrometheusStatLogger.__new__(OmniPrometheusStatLogger) + srm = {0: ("thinker", "0"), 1: ("diffusion", "0")} + sl._stage_replica_map = srm + assert sl.stage_replica_map is srm + + def test_init_populates_engine_index_map(self): + # Simulate the bookkeeping portion of __init__ (clear + update) without + # calling super, since super needs a real VllmConfig. + _ENGINE_INDEX_MAP[99] = ("stale", "stale") # leftover from prior + srm = {0: ("thinker", "0"), 1: ("talker", "0")} + + # Manually invoke the bookkeeping the way __init__ does it. + _ENGINE_INDEX_MAP.clear() + _ENGINE_INDEX_MAP.update(srm) + + assert dict(_ENGINE_INDEX_MAP) == srm + assert 99 not in _ENGINE_INDEX_MAP # old entry was cleared + + +# --------------------------------------------------------------------------- +# Helper-class wraps for upstream's spec_decoding / kv_connector / +# perf_metrics sub-collectors. Without these, OmniPrometheusStatLogger +# crashes at startup with `Incorrect label count` because each helper builds +# its internal Counter/Gauge/Histogram families with raw 2-element labelnames +# (passed via constructor arg) while consuming the rewritten 3-element +# per_engine_labelvalues from the property descriptor. +# --------------------------------------------------------------------------- + + +class TestHelperClassWraps: + def test_perf_metrics_wrap_routes_through_relabel_counter(self): + assert _OmniPerfMetricsProm._counter_cls is _RelabelCounter + + def test_spec_decoding_wrap_routes_through_relabel_counter(self): + assert _OmniSpecDecodingProm._counter_cls is _RelabelCounter + + def test_kv_connector_wrap_routes_through_all_three_relabel_classes(self): + # KVConnector lets each connector build any of Gauge/Counter/Histogram, + # so all three slots must be intercepted. + assert _OmniKVConnectorProm._gauge_cls is _RelabelGauge + assert _OmniKVConnectorProm._counter_cls is _RelabelCounter + assert _OmniKVConnectorProm._histogram_cls is _RelabelHistogram + + def test_omni_logger_slots_point_to_helper_subclasses(self): + # Upstream's PrometheusStatLogger.__init__ instantiates each sub-helper + # via `self.__cls(...)`, so the slot overrides on the omni + # subclass are what routes through to the relabel mixin. + assert OmniPrometheusStatLogger._perf_metrics_cls is _OmniPerfMetricsProm + assert OmniPrometheusStatLogger._spec_decoding_cls is _OmniSpecDecodingProm + assert OmniPrometheusStatLogger._kv_connector_cls is _OmniKVConnectorProm + + +# --------------------------------------------------------------------------- +# Double-rewrite guard. The mixin's positional-args path used to +# unconditionally splice (stage, replica) at engine_label_index. Because +# per_engine_labelvalues is rewritten to 3-tuples *before* feeding into +# create_metric_per_engine, that splice would run a second time on the +# already-rewritten values and blow the label count to 4. The guard now +# detects len(args) == len(self._labelnames) and short-circuits to passthrough. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def fresh_registry(): + return CollectorRegistry() + + +class TestDoubleRewriteGuard: + def test_pre_rewritten_3tuple_passes_through(self, fresh_registry): + # 2-label original → 3-label rewritten family. Caller passes 3 values + # (the rewritten shape) and they should land verbatim, not get + # re-spliced. + _ENGINE_INDEX_MAP.clear() + _ENGINE_INDEX_MAP[0] = ("0", "0") + _ENGINE_INDEX_MAP[1] = ("1", "0") + g = _RelabelGauge( + name="dr_pre_rewritten", + documentation="t", + labelnames=["model_name", "engine"], + registry=fresh_registry, + ) + # 3 positional args matching the rewritten 3-label family. + g.labels("m", "1", "0").set(42) + out = generate_latest(fresh_registry).decode() + assert 'dr_pre_rewritten{model_name="m",replica="0",stage="1"} 42.0' in out + + def test_legacy_2tuple_with_extra_label_still_splices(self, fresh_registry): + # 3-label original (engine in middle) → 4-label rewritten family. + # Caller passes 3 values matching the ORIGINAL labelnames (the + # gauge_waiting_by_reason / counter_request_success pattern from + # upstream loggers.py:646, 679). The mixin must splice + # (stage, replica) at engine's position to reach the 4-label family. + _ENGINE_INDEX_MAP.clear() + _ENGINE_INDEX_MAP[1] = ("1", "0") + c = _RelabelCounter( + name="dr_legacy_with_extra", + documentation="t", + labelnames=["model_name", "engine", "reason"], + registry=fresh_registry, + ) + # 3 positional args matching ORIGINAL labelnames (model_name, + # engine_str, reason). + c.labels("m", "1", "stop").inc(3) + out = generate_latest(fresh_registry).decode() + assert 'dr_legacy_with_extra_total{model_name="m",reason="stop",replica="0",stage="1"} 3.0' in out diff --git a/tests/metrics/test_transfer.py b/tests/metrics/test_transfer.py new file mode 100644 index 000000000..18301f80e --- /dev/null +++ b/tests/metrics/test_transfer.py @@ -0,0 +1,346 @@ +from __future__ import annotations + +import pytest +from prometheus_client import REGISTRY, generate_latest + +from vllm_omni.metrics import definitions as defs +from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats, StageStats +from vllm_omni.metrics.transfer import OmniTransferMetrics + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +_MODEL = "test-transfer-model" + + +@pytest.fixture(scope="module") +def tx() -> OmniTransferMetrics: + return OmniTransferMetrics(model_name=_MODEL) + + +def _sample_value(output: str, line_prefix: str) -> float | None: + for line in output.splitlines(): + if line.startswith(line_prefix): + return float(line.split()[-1]) + return None + + +_EXPECTED_FAMILIES = [ + defs.TRANSFER_SIZE_BYTES, + defs.TRANSFER_TX_S, + defs.TRANSFER_RX_S, + defs.TRANSFER_IN_FLIGHT_S, +] + + +# --------------------------------------------------------------------------- +# Family registration +# --------------------------------------------------------------------------- + + +class TestRegistration: + def test_all_four_families_present(self, tx: OmniTransferMetrics) -> None: + tx.observe_size(0, 0, 1, 0, 1024) + tx.observe_tx_time(0, 0, 1, 0, 0.005) + tx.observe_rx_time(0, 0, 1, 0, 0.008) + tx.observe_in_flight_time(0, 0, 1, 0, 0.002) + + out = generate_latest(REGISTRY).decode() + for name in _EXPECTED_FAMILIES: + assert f"# HELP {name}" in out, f"missing family: {name}" + + def test_legacy_ms_family_names_dropped(self, tx: OmniTransferMetrics) -> None: + out = generate_latest(REGISTRY).decode() + for legacy in ( + "vllm:omni_transfer_tx_time_ms", + "vllm:omni_transfer_rx_decode_time_ms", + "vllm:omni_transfer_in_flight_time_ms", + ): + assert legacy not in out, f"legacy family {legacy} still registered" + + +# --------------------------------------------------------------------------- +# Observe APIs +# --------------------------------------------------------------------------- + + +class TestObserveSize: + def test_size_observed_with_correct_labels(self, tx: OmniTransferMetrics) -> None: + tx.observe_size(2, 0, 3, 1, 65536) + out = generate_latest(REGISTRY).decode() + prefix = ( + f"{defs.TRANSFER_SIZE_BYTES}_sum" + f'{{from_replica="0",from_stage="2",model_name="{_MODEL}",' + f'to_replica="1",to_stage="3"}}' + ) + assert _sample_value(out, prefix) == 65536.0 + + +class TestObserveTxTime: + def test_tx_time_observed_in_seconds(self, tx: OmniTransferMetrics) -> None: + tx.observe_tx_time(2, 1, 3, 0, 0.0125) + out = generate_latest(REGISTRY).decode() + prefix = ( + f"{defs.TRANSFER_TX_S}_sum" + f'{{from_replica="1",from_stage="2",model_name="{_MODEL}",' + f'to_replica="0",to_stage="3"}}' + ) + assert _sample_value(out, prefix) == pytest.approx(0.0125) + + +class TestObserveRxTime: + def test_rx_time_observed_in_seconds(self, tx: OmniTransferMetrics) -> None: + # Use a (from, to) combo not shared with TestRegistration's setup + # observe on the rx_s Histogram, otherwise the per-test sum picks + # up accumulated values from earlier observes. + tx.observe_rx_time(3, 0, 4, 0, 0.0042) + out = generate_latest(REGISTRY).decode() + prefix = ( + f"{defs.TRANSFER_RX_S}_sum" + f'{{from_replica="0",from_stage="3",model_name="{_MODEL}",' + f'to_replica="0",to_stage="4"}}' + ) + assert _sample_value(out, prefix) == pytest.approx(0.0042) + + +class TestObserveInFlightTime: + def test_in_flight_time_observed_in_seconds(self, tx: OmniTransferMetrics) -> None: + tx.observe_in_flight_time(3, 0, 4, 0, 0.0017) + out = generate_latest(REGISTRY).decode() + prefix = ( + f"{defs.TRANSFER_IN_FLIGHT_S}_sum" + f'{{from_replica="0",from_stage="3",model_name="{_MODEL}",' + f'to_replica="0",to_stage="4"}}' + ) + assert _sample_value(out, prefix) == pytest.approx(0.0017) + + +# --------------------------------------------------------------------------- +# Multi (from, to) cardinality +# --------------------------------------------------------------------------- + + +class TestCardinality: + def test_multiple_edges_produce_independent_series(self, tx: OmniTransferMetrics) -> None: + tx.observe_size(5, 0, 6, 0, 100) + tx.observe_size(5, 0, 6, 1, 200) + tx.observe_size(5, 1, 6, 0, 300) + + out = generate_latest(REGISTRY).decode() + + prefix_a = ( + f"{defs.TRANSFER_SIZE_BYTES}_sum" + f'{{from_replica="0",from_stage="5",model_name="{_MODEL}",' + f'to_replica="0",to_stage="6"}}' + ) + prefix_b = ( + f"{defs.TRANSFER_SIZE_BYTES}_sum" + f'{{from_replica="0",from_stage="5",model_name="{_MODEL}",' + f'to_replica="1",to_stage="6"}}' + ) + prefix_c = ( + f"{defs.TRANSFER_SIZE_BYTES}_sum" + f'{{from_replica="1",from_stage="5",model_name="{_MODEL}",' + f'to_replica="0",to_stage="6"}}' + ) + assert _sample_value(out, prefix_a) == 100.0 + assert _sample_value(out, prefix_b) == 200.0 + assert _sample_value(out, prefix_c) == 300.0 + + +# --------------------------------------------------------------------------- +# Bucket selection +# --------------------------------------------------------------------------- + + +class TestBucketSelection: + def test_size_uses_bytes_buckets(self, tx: OmniTransferMetrics) -> None: + tx.observe_size(7, 0, 8, 0, 4096) + out = generate_latest(REGISTRY).decode() + # prometheus_client renders bucket boundary as float (e.g. ``1024.0``) + # regardless of the declared int in BYTES_BUCKETS. + marker = f'{defs.TRANSFER_SIZE_BYTES}_bucket{{from_replica="0",from_stage="7",le="1024.0"' + assert marker in out + + def test_time_families_use_seconds_fast_buckets(self, tx: OmniTransferMetrics) -> None: + tx.observe_tx_time(7, 0, 8, 0, 0.001) + out = generate_latest(REGISTRY).decode() + # SECONDS_FAST_BUCKETS has le=0.001 — distinctive vs SECONDS_BUCKETS (no 0.001). + marker = f'{defs.TRANSFER_TX_S}_bucket{{from_replica="0",from_stage="7",le="0.001"' + assert marker in out + + +# --------------------------------------------------------------------------- +# OrchestratorAggregator emit hook +# --------------------------------------------------------------------------- + + +class _StubTransferEmitter: + """Records every observe_* call so the hook routing can be asserted + without standing up a Prometheus registry.""" + + def __init__(self) -> None: + self.calls: list[tuple] = [] + + def observe_size(self, fs, fr, ts, tr, n): + self.calls.append(("observe_size", fs, fr, ts, tr, n)) + + def observe_tx_time(self, fs, fr, ts, tr, t): + self.calls.append(("observe_tx_time", fs, fr, ts, tr, t)) + + def observe_rx_time(self, fs, fr, ts, tr, t): + self.calls.append(("observe_rx_time", fs, fr, ts, tr, t)) + + def observe_in_flight_time(self, fs, fr, ts, tr, t): + self.calls.append(("observe_in_flight_time", fs, fr, ts, tr, t)) + + +def _make_stats(stage_id, request_id, *, rx_decode=0.0, rx_in_flight=0.0, rx_bytes=0): + """Minimal StageRequestStats for record_transfer_rx input.""" + return StageRequestStats( + batch_id=1, + batch_size=1, + num_tokens_in=0, + num_tokens_out=0, + stage_gen_time_ms=0.0, + rx_transfer_bytes=rx_bytes, + rx_decode_time_ms=rx_decode, + rx_in_flight_time_ms=rx_in_flight, + stage_stats=StageStats(), + stage_id=stage_id, + request_id=request_id, + ) + + +class TestEmitHookTx: + def test_record_transfer_tx_emits_size_and_tx_time_in_seconds(self): + emitter = _StubTransferEmitter() + agg = OrchestratorAggregator( + num_stages=3, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=2, + transfer_emitter=emitter, + replica_resolver=lambda s, rid: {0: 1, 1: 0}.get(s), + ) + agg.record_transfer_tx( + from_stage=0, + to_stage=1, + request_id="r-tx-1", + size_bytes=2048, + tx_time_ms=7.5, + used_shm=False, + ) + # tx_time_ms=7.5 -> 0.0075s at the emit boundary. + assert emitter.calls == [ + ("observe_size", 0, 1, 1, 0, 2048), + ("observe_tx_time", 0, 1, 1, 0, pytest.approx(0.0075)), + ] + + def test_record_transfer_tx_no_emit_when_emitter_none(self): + agg = OrchestratorAggregator( + num_stages=3, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=2, + transfer_emitter=None, + replica_resolver=lambda s, rid: 0, + ) + evt = agg.record_transfer_tx( + from_stage=0, + to_stage=1, + request_id="r-tx-2", + size_bytes=128, + tx_time_ms=1.0, + used_shm=True, + ) + # Underlying accumulator keeps ms; only the Prometheus emit converts. + assert evt is not None + assert evt.size_bytes == 128 + assert evt.tx_time_ms == 1.0 + + def test_record_transfer_tx_no_emit_when_resolver_returns_none(self): + emitter = _StubTransferEmitter() + agg = OrchestratorAggregator( + num_stages=3, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=2, + transfer_emitter=emitter, + replica_resolver=lambda s, rid: None, + ) + agg.record_transfer_tx( + from_stage=0, + to_stage=1, + request_id="r-tx-3", + size_bytes=512, + tx_time_ms=2.0, + used_shm=False, + ) + assert emitter.calls == [] + + def test_record_transfer_tx_no_emit_when_one_side_resolves_to_none(self): + emitter = _StubTransferEmitter() + agg = OrchestratorAggregator( + num_stages=3, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=2, + transfer_emitter=emitter, + replica_resolver=lambda s, rid: 0 if s == 0 else None, + ) + agg.record_transfer_tx( + from_stage=0, + to_stage=1, + request_id="r-tx-4", + size_bytes=64, + tx_time_ms=0.5, + used_shm=False, + ) + assert emitter.calls == [] + + +class TestEmitHookRx: + def test_record_transfer_rx_emits_rx_and_in_flight_in_seconds(self): + emitter = _StubTransferEmitter() + agg = OrchestratorAggregator( + num_stages=3, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=2, + transfer_emitter=emitter, + replica_resolver=lambda s, rid: {0: 1, 1: 0}.get(s), + ) + stats = _make_stats(stage_id=1, request_id="r-rx-1", rx_decode=4.2, rx_in_flight=1.7) + agg.record_transfer_rx(stats) + assert emitter.calls == [ + ("observe_rx_time", 0, 1, 1, 0, pytest.approx(0.0042)), + ("observe_in_flight_time", 0, 1, 1, 0, pytest.approx(0.0017)), + ] + + def test_record_transfer_rx_skips_stage_zero(self): + emitter = _StubTransferEmitter() + agg = OrchestratorAggregator( + num_stages=3, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=2, + transfer_emitter=emitter, + replica_resolver=lambda s, rid: 0, + ) + stats = _make_stats(stage_id=0, request_id="r-rx-2", rx_decode=4.2) + agg.record_transfer_rx(stats) + assert emitter.calls == [] + + def test_record_transfer_rx_no_emit_when_emitter_none(self): + agg = OrchestratorAggregator( + num_stages=3, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=2, + ) + stats = _make_stats(stage_id=1, request_id="r-rx-3", rx_decode=1.0, rx_in_flight=0.5) + evt = agg.record_transfer_rx(stats) + assert evt is not None + assert evt.rx_decode_time_ms == 1.0 + assert evt.in_flight_time_ms == 0.5 diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index 56ebb20ea..55f4b8830 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/vllm_omni/benchmarks/metrics/metrics.py @@ -7,6 +7,8 @@ from vllm.benchmarks.lib.endpoint_request_func import RequestFuncOutput from vllm.benchmarks.serve import MILLISECONDS_TO_SECONDS_CONVERSION, TERM_PLOTLIB_AVAILABLE, BenchmarkMetrics, TaskType from vllm.tokenizers import TokenizerLike +from vllm_omni.metrics import definitions as defs + @dataclass class MultiModalsBenchmarkMetrics(BenchmarkMetrics): @@ -109,17 +111,17 @@ def process_one_metric( "tpot": "Time per Output Token (excl. 1st token)", "itl": "Inter-token Latency", "e2el": "End-to-end Latency", - "audio_ttfp": "Time to First Packet", - "audio_rtf": "Real Time Factor", - "audio_duration": "Audio Duration", - "audio_underrun": "Streaming Audio Underrun", + defs.AUDIO_TTFP: "Time to First Packet", + defs.AUDIO_RTF: "Real Time Factor", + defs.AUDIO_DURATION: "Audio Duration", + defs.AUDIO_UNDERRUN: "Streaming Audio Underrun", } header = metric_header_map.get(metric_attribute_name, metric_attribute_name) print("{s:{c}^{n}}".format(s=header, n=50, c="-")) - is_audio_rtf = metric_attribute_name == "audio_rtf" - is_audio_duration_or_underrun = metric_attribute_name in ("audio_duration", "audio_underrun") + is_audio_rtf = metric_attribute_name == defs.AUDIO_RTF + is_audio_duration_or_underrun = metric_attribute_name in (defs.AUDIO_DURATION, defs.AUDIO_UNDERRUN) suffix = "_ms" unit_suffix = " (ms)" @@ -216,12 +218,12 @@ def calculate_metrics( all_tpots.append(tpot) itls += outputs[i].itl ttfts.append(outputs[i].ttft) - audio_ttfps.append(getattr(outputs[i], "audio_ttfp", 0.0)) - audio_rtfs.append(getattr(outputs[i], "audio_rtf", 0.0)) - audio_duration.append(getattr(outputs[i], "audio_duration", 0.0)) - audio_frames.append(getattr(outputs[i], "audio_frames", 0.0)) - audio_underruns.append(getattr(outputs[i], "audio_underrun_s", 0.0)) - audio_continuity_ok.append(bool(getattr(outputs[i], "audio_continuity_ok", True))) + audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0)) + audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0)) + audio_duration.append(getattr(outputs[i], defs.AUDIO_DURATION, 0.0)) + audio_frames.append(getattr(outputs[i], defs.AUDIO_FRAMES, 0.0)) + audio_underruns.append(getattr(outputs[i], f"{defs.AUDIO_UNDERRUN}_s", 0.0)) + audio_continuity_ok.append(bool(getattr(outputs[i], defs.AUDIO_CONTINUITY_OK, True))) e2els.append(outputs[i].latency) input_audio_duration += outputs[i].input_audio_duration completed += 1 diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index c4b1e2429..c33d16dac 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -1,16 +1,20 @@ from __future__ import annotations import os +import time from typing import Any from vllm.logger import init_logger from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.stats import SchedulerStats from vllm.v1.request import RequestStatus from vllm_omni.core.sched.output import OmniChunkRecvHandle, OmniSchedulerOutput logger = init_logger(__name__) +_STATS_INTERVAL_S = 1.0 + # Upper bound on how long a request may sit in full-payload-input wait # (the state ``OmniSchedulingCoordinator`` records via ``_waiting_since``) # before the scheduler force-fails it. Defends against stuck consumer-side @@ -145,3 +149,10 @@ class OmniSchedulerMixin: finished_requests_needing_kv_transfer=finished_requests_needing_kv_transfer or {}, pending_input_registrations=pending_input_registrations, ) + + def make_stats(self, *args, **kwargs) -> SchedulerStats | None: + now = time.monotonic() + if now - getattr(self, "_last_stats_time", 0.0) < _STATS_INTERVAL_S: + return None + self._last_stats_time = now + return super().make_stats(*args, **kwargs) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 939f9b224..5163de016 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -110,6 +110,7 @@ from vllm_omni.entrypoints.utils import ( load_and_resolve_stage_configs, ) from vllm_omni.inputs.data import OmniSamplingParams +from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: @@ -259,6 +260,12 @@ class AsyncOmniEngine: **kwargs: Additional arguments """ + # Class-level defaults so tests that bypass __init__ via object.__new__ + # don't AttributeError when stage-init / forward paths touch these attrs. + _log_stats: bool = False + _coordinator_runtime: Any = None + _transfer_emitter: Any = None + def __init__( self, model: str, @@ -267,11 +274,23 @@ class AsyncOmniEngine: init_timeout: int = 600, diffusion_batch_size: int = 1, single_stage_mode: bool = False, + transfer_emitter: Any = None, + log_stats: bool = False, **kwargs: Any, ) -> None: self.model = model self.diffusion_batch_size = diffusion_batch_size startup_timeout = int(init_timeout) + # Forwarded into Orchestrator so its _forward_to_next_stage path can + # emit per-edge transfer_tx_s / transfer_size_bytes histograms. + # Optional: when None, Orchestrator silently skips TX emit (existing + # RX path still works via OrchestratorAggregator). + self._transfer_emitter = transfer_emitter + # 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 + # --log-stats CLI flag set by the user via OmniBase. + self._log_stats = log_stats logger.info(f"[AsyncOmniEngine] Initializing with model {model}") @@ -357,6 +376,7 @@ class AsyncOmniEngine: self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None self._weak_finalizer: weakref.finalize | None = None self._rpc_lock = threading.Lock() + self._running_counter = OmniRequestCounter() logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages") @@ -905,7 +925,7 @@ class AsyncOmniEngine: launch_omni_core_engines( vllm_config=vllm_config, executor_class=executor_class, - log_stats=False, + log_stats=self._log_stats, omni_master_server=self._omni_master_server, stage_id=plan.metadata.stage_id, stage_config=stage_cfg, @@ -917,7 +937,7 @@ class AsyncOmniEngine: addresses, proc, handshake_address = spawn_stage_core( vllm_config=vllm_config, executor_class=executor_class, - log_stats=False, + log_stats=self._log_stats, ) logger.info( "[AsyncOmniEngine] Stage %s engine launch started", @@ -1229,7 +1249,7 @@ class AsyncOmniEngine: if plan.replicas[0].metadata.stage_type != "diffusion": stage_vllm_config = plan.replicas[0].stage_vllm_config assert stage_vllm_config is not None - output_processor = build_llm_stage_output_processor(plan, stage_vllm_config) + output_processor = build_llm_stage_output_processor(plan, stage_vllm_config, log_stats=self._log_stats) stage_pools.append( StagePool( @@ -1374,9 +1394,12 @@ class AsyncOmniEngine: stage_pools=self.stage_pools, async_chunk=self.async_chunk, pd_config=pd_config, + running_counter=self._running_counter, coordinator_pub_address=coordinator_pub_address, load_balancer_factory=load_balancer_factory, remote_replica_factory=remote_replica_factory, + transfer_emitter=self._transfer_emitter, + log_stats=self._log_stats, ) if not startup_future.done(): startup_future.set_result(asyncio.get_running_loop()) diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 0e55105f1..0ced58391 100644 --- a/vllm_omni/engine/messages.py +++ b/vllm_omni/engine/messages.py @@ -80,6 +80,7 @@ class OutputMessage(EngineQueueMessage, kw_only=True): type: Literal["output"] = "output" request_id: str stage_id: int + replica_id: int | None = None engine_outputs: OmniRequestOutput metrics: StageRequestMetrics | None = None finished: bool @@ -90,6 +91,7 @@ class StageMetricsMessage(EngineQueueMessage, kw_only=True): type: Literal["stage_metrics"] = "stage_metrics" request_id: str stage_id: int + replica_id: int | None = None metrics: StageRequestMetrics stage_submit_ts: float | None = None diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 791b6d5d9..99ee7f93b 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -29,6 +29,7 @@ from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreOutputs from vllm.v1.engine.exceptions import EngineDeadError +from vllm.v1.metrics.stats import IterationStats from vllm_omni.distributed.omni_coordinator import ( LoadBalancer, @@ -54,6 +55,8 @@ from vllm_omni.engine.messages import ( ) from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.metrics.prometheus import OmniRequestCounter +from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger from vllm_omni.outputs import OmniRequestOutput # Factory signature for building a head-side stage client for a @@ -155,6 +158,12 @@ class Orchestrator: _WATCH_REPLICA_INTERVAL_S: float = 0.5 _WATCH_REPLICA_IDLE_INTERVAL_S: float = 1.0 + # Class-level defaults so tests that bypass __init__ via object.__new__ + # don't AttributeError when transfer / counter emit paths access them. + _running_counter: OmniRequestCounter | None = None + _transfer_emitter: Any = None + _stat_logger: OmniPrometheusStatLogger | None = None + def __init__( self, request_async_queue: janus.AsyncQueue[EngineQueueMessage], @@ -164,9 +173,12 @@ class Orchestrator: *, async_chunk: bool = False, pd_config: dict[str, Any] | None = None, + running_counter: OmniRequestCounter | None = None, coordinator_pub_address: str | None = None, load_balancer_factory: Callable[[], LoadBalancer] | None = None, remote_replica_factory: RemoteReplicaFactory | None = None, + transfer_emitter: Any = None, + log_stats: bool = False, ) -> None: self.request_async_queue = request_async_queue self.output_async_queue = output_async_queue @@ -186,6 +198,8 @@ class Orchestrator: self._pd_bootstrap_addr = pd_config.get("bootstrap_addr") self._pd_prefill_engine_id = pd_config.get("prefill_engine_id") self.request_states: dict[str, OrchestratorRequestState] = {} + self._init_metrics_state(stage_pools, running_counter, transfer_emitter, log_stats=log_stats) + self._cfg_tracker = CfgCompanionTracker() self._shutdown_event = asyncio.Event() @@ -215,6 +229,71 @@ class Orchestrator: pool.attach_hub(self._hub) pool.attach_load_balancer(factory()) + def _init_metrics_state( + self, + stage_pools: list[StagePool], + running_counter: OmniRequestCounter | None, + transfer_emitter: Any, + log_stats: bool = False, + ) -> None: + """Wire up all metric-related orchestrator state. + + Sets ``self._running_counter`` and ``self._transfer_emitter`` + (both optional, used by request-add / forward paths), builds the + ``(stage_id, replica_id) ↔ engine_idx`` lookup used at record() time, + and best-effort constructs the ``OmniPrometheusStatLogger`` wrap + that exposes ~37 upstream ``vllm:*`` families with per-(stage, + replica) labels. Failure to build the wrap is logged and metrics + are simply disabled — orchestrator construction continues so unit + tests with a minimal ``vllm_config`` still pass. + + ``log_stats=False`` short-circuits the wrap entirely so the + ~65 upstream ``vllm:*`` families are not registered in the + Prometheus default registry at all. The per-step record() path + already no-ops on ``scheduler_stats is None`` (which is what + the upstream scheduler returns when its own log_stats is False), + so this gate is mainly to keep the ``/metrics`` surface clean + when the user did not request stats. + """ + self._running_counter = running_counter + self._transfer_emitter = transfer_emitter + + # Flat engine_idx ↔ (stage, replica) maps. The reverse map is + # consulted at record() time to translate the orchestrator's + # (stage_id, replica_id) loop variables into an engine_idx the + # underlying PrometheusStatLogger can address. + stage_replica_map: dict[int, tuple[str, str]] = {} + self._stage_replica_to_engine_idx: dict[tuple[int, int], int] = {} + flat_idx = 0 + for stage_id, pool in enumerate(stage_pools): + for replica_id in range(pool.num_replicas): + stage_replica_map[flat_idx] = (str(stage_id), str(replica_id)) + self._stage_replica_to_engine_idx[(stage_id, replica_id)] = flat_idx + flat_idx += 1 + + if not log_stats: + self._stat_logger = None + return + + vllm_config_for_stats = next( + (p.stage_vllm_config for p in stage_pools if p.stage_vllm_config is not None), + None, + ) + if vllm_config_for_stats is None: + self._stat_logger = None + return + try: + self._stat_logger = OmniPrometheusStatLogger( + vllm_config=vllm_config_for_stats, + stage_replica_map=stage_replica_map, + ) + except Exception: + # Minimal vllm_config in unit-test contexts can lack fields the + # upstream PrometheusStatLogger expects. Skip wrap rather than + # break orchestrator construction. + logger.exception("[Orchestrator] OmniPrometheusStatLogger init failed; metrics wrap disabled") + self._stat_logger = None + async def run(self) -> None: """Main entry point for the Orchestrator event loop.""" logger.info("[Orchestrator] Starting event loop") @@ -393,6 +472,8 @@ class Orchestrator: mm_features=getattr(prompt, "mm_features", None), ) self.request_states[request_id] = req_state + if self._running_counter is not None: + self._running_counter.increment() req_state.streaming.enabled = bool(getattr(prompt, "resumable", False)) req_state.stage_submit_ts[stage_id] = _time.time() enqueue_ts = msg.enqueue_ts @@ -603,7 +684,24 @@ class Orchestrator: "new_prompt_len_snapshot", None, ) - raw_output = await pool.process_llm_raw_outputs(replica_id, raw_outputs) + # OmniSchedulerMixin.make_stats() already throttles + # per-scheduler at 1 Hz, so raw_outputs.scheduler_stats + # being non-None means this replica passed its own gate. + # A second global throttle here would drop stats for + # other (stage, replica) pairs in the same 1s window. + record_stats = self._stat_logger is not None and raw_outputs.scheduler_stats is not None + iteration_stats = IterationStats() if record_stats else None + raw_output = await pool.process_llm_raw_outputs( + replica_id, + raw_outputs, + iteration_stats=iteration_stats, + ) + if record_stats: + self._stat_logger.record( + raw_outputs.scheduler_stats, + iteration_stats, + engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)], + ) except asyncio.CancelledError: raise except EngineDeadError as e: @@ -681,7 +779,7 @@ class Orchestrator: ) stage_metrics.pipeline_timings = dict(req_state.pipeline_timings) - await self._route_output(stage_id, output, req_state, stage_metrics) + await self._route_output(stage_id, replica_id, output, req_state, stage_metrics) async def _handle_stage_error(self, stage_id: int, output: Any) -> None: """Emit a frontend-visible error and clean up request state.""" @@ -713,7 +811,8 @@ class Orchestrator: self._release_request_bindings(request_ids) for request_id in request_ids: self._pd_kv_params.pop(request_id, None) - self.request_states.pop(request_id, None) + if self.request_states.pop(request_id, None) is not None and self._running_counter is not None: + self._running_counter.decrement() def _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any: """Attach CFG companion ids to diffusion sampling params when needed.""" @@ -735,6 +834,7 @@ class Orchestrator: async def _route_output( self, stage_id: int, + replica_id: int, output: Any, req_state: OrchestratorRequestState, stage_metrics: Any, @@ -757,6 +857,7 @@ class Orchestrator: OutputMessage( request_id=req_id, stage_id=stage_id, + replica_id=replica_id, engine_outputs=output, metrics=stage_metrics, finished=finished and stage_id == req_state.final_stage_id, @@ -768,6 +869,7 @@ class Orchestrator: StageMetricsMessage( request_id=req_id, stage_id=stage_id, + replica_id=replica_id, metrics=stage_metrics, stage_submit_ts=submit_ts, ) @@ -797,6 +899,7 @@ class Orchestrator: stage_id, output, req_state, + src_replica_id=replica_id, is_streaming_session=req_state.streaming.enabled, is_final_update=False, ) @@ -807,6 +910,7 @@ class Orchestrator: stage_id, output, req_state, + src_replica_id=replica_id, is_streaming_session=True, is_final_update=True, ) @@ -918,6 +1022,41 @@ class Orchestrator: ) return sp + def _emit_tx_edge( + self, + *, + from_stage: int, + from_replica: int, + to_stage: int, + to_pool: StagePool, + request_id: str, + tx_ms: float, + ) -> None: + """Emit per-edge transfer_tx_s + transfer_size_bytes histograms. + + ``tx_ms`` is the orchestrator-side wall-clock spent in ``next_pool. + submit_*`` (serialize + queue submit to the receiving worker). Best- + effort size_bytes left at 0 — orchestrator doesn't have a cheap handle + on the serialized payload size; a follow-up can plumb that from the + connector adapter. + """ + if self._transfer_emitter is None: + return + to_replica = to_pool.get_bound_replica_id(request_id) + if to_replica is None: + return + try: + self._transfer_emitter.observe_size(from_stage, from_replica, to_stage, to_replica, 0) + self._transfer_emitter.observe_tx_time(from_stage, from_replica, to_stage, to_replica, tx_ms / 1000.0) + except Exception: + logger.debug( + "[Orchestrator] transfer_tx emit failed for edge %d->%d req=%s", + from_stage, + to_stage, + request_id, + exc_info=True, + ) + async def _forward_to_next_stage( self, req_id: str, @@ -925,6 +1064,7 @@ class Orchestrator: output: Any, req_state: OrchestratorRequestState, *, + src_replica_id: int | None = None, is_streaming_session: bool = False, is_final_update: bool = False, ) -> None: @@ -937,6 +1077,7 @@ class Orchestrator: next_stage_resumable = is_streaming_session and not is_final_update already_submitted = self._next_stage_already_submitted(src_stage_id, req_state) requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False) + _t_submit_start = _time.perf_counter() if next_pool.stage_type == "diffusion": companion_outputs = self._cfg_tracker.pop_companion_outputs(req_id) @@ -1024,6 +1165,15 @@ class Orchestrator: params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), ) req_state.stage_submit_ts[next_logical] = _time.time() + _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 + self._emit_tx_edge( + from_stage=src_stage_id, + from_replica=src_replica_id if src_replica_id is not None else 0, + to_stage=next_logical, + to_pool=next_pool, + request_id=req_id, + tx_ms=_tx_ms, + ) return # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params @@ -1063,6 +1213,15 @@ class Orchestrator: await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) req_state.stage_submit_ts[next_logical] = _time.time() + _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 + self._emit_tx_edge( + from_stage=src_stage_id, + from_replica=src_replica_id if src_replica_id is not None else 0, + to_stage=next_logical, + to_pool=next_pool, + request_id=req_id, + tx_ms=_tx_ms, + ) return if req_state.pd_prefill_multimodal_output is not None: @@ -1107,6 +1266,15 @@ class Orchestrator: await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) req_state.stage_submit_ts[next_logical] = _time.time() + _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 + self._emit_tx_edge( + from_stage=src_stage_id, + from_replica=src_replica_id if src_replica_id is not None else 0, + to_stage=next_logical, + to_pool=next_pool, + request_id=req_id, + tx_ms=_tx_ms, + ) async def _prewarm_async_chunk_stages( self, @@ -1131,6 +1299,7 @@ class Orchestrator: params = req_state.sampling_params_list[next_stage_id] req_state.stage_submit_ts[next_stage_id] = _time.time() + _t_submit_start = _time.perf_counter() if next_pool.stage_type == "diffusion": await next_pool.submit_initial( @@ -1174,6 +1343,20 @@ class Orchestrator: request.external_req_id = request.request_id await next_pool.submit_initial(request_id, req_state, request, prompt_text=None) + # async_chunk pre-submit fires per stage edge (N-1 -> N). Source + # replica is stage 0's bound replica (single-replica thinker in + # all current configs); fall back to 0 if unknown. + _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 + src_replica = self.stage_pools[next_stage_id - 1].get_bound_replica_id(request_id) + self._emit_tx_edge( + from_stage=next_stage_id - 1, + from_replica=src_replica if src_replica is not None else 0, + to_stage=next_stage_id, + to_pool=next_pool, + request_id=request_id, + tx_ms=_tx_ms, + ) + def _build_kv_sender_info( self, sender_stage_ids: list[int], diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 75024abdb..d870e0950 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -831,8 +831,18 @@ def build_vllm_config( return vllm_config, executor_class -def build_llm_stage_output_processor(plan: LogicalStageInitPlan, stage_vllm_config: Any) -> Any | None: - """Build one output processor per logical LLM stage.""" +def build_llm_stage_output_processor( + plan: LogicalStageInitPlan, + stage_vllm_config: Any, + log_stats: bool = False, +) -> Any | None: + """Build one output processor per logical LLM stage. + + ``log_stats`` controls whether the processor populates per-request + IterationStats (consumed by the Prometheus wrap). Default False matches + the upstream MultimodalOutputProcessor default and respects the + --log-stats CLI flag plumbed through AsyncOmniEngine. + """ metadata = plan.replicas[0].metadata if stage_vllm_config.model_config.skip_tokenizer_init: @@ -843,7 +853,7 @@ def build_llm_stage_output_processor(plan: LogicalStageInitPlan, stage_vllm_conf ) return MultimodalOutputProcessor( tokenizer=tokenizer, - log_stats=False, + log_stats=log_stats, engine_core_output_type=metadata.engine_output_type, ) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 8305c23d4..c6a40f55e 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, cast from vllm.logger import init_logger from vllm.v1.engine import EngineCoreOutputs +from vllm.v1.metrics.stats import IterationStats from vllm_omni.distributed.omni_coordinator import ( LoadBalancer, @@ -510,6 +511,7 @@ class StagePool: stage_gen_time_ms=stage_gen_time_ms, batch_id=batch_id, batch_size=1, + replica_id=replica_id, rx_decode_time_ms=0.0, rx_transfer_bytes=0, rx_in_flight_time_ms=0.0, @@ -648,6 +650,7 @@ class StagePool: self, replica_id: int, raw_outputs: EngineCoreOutputs, + iteration_stats: IterationStats | None = None, ) -> list[Any]: """Run the shared LLM output processor on one raw poll result.""" raw_client = self.clients[replica_id] @@ -658,7 +661,7 @@ class StagePool: processed = processor.process_outputs( raw_outputs.outputs, raw_outputs.timestamp, - None, + iteration_stats, ) if processed.reqs_to_abort: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 8a4a1d48e..ae770346d 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -160,6 +160,20 @@ class AsyncOmni(EngineClient, OmniBase): renderer = renderer_from_config(vllm_config) self.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin) + def _resolve_transfer_replica(self, stage_id: int, request_id: str) -> int | None: + """Look up the sticky-routed replica for (stage_id, request_id). + + Used as the ``replica_resolver`` callback by ``OrchestratorAggregator`` + to label transfer_* metrics without plumbing replica ids through + ``TransferEdgeStats`` / ``StageRequestStats`` / connector adapters. + Returns None when stage_id is out of range or the request hasn't been + bound to a replica yet — the metric emit then defensive-skips. + """ + pools = getattr(self.engine, "stage_pools", None) + if pools is None or not (0 <= stage_id < len(pools)): + return None + return pools[stage_id].get_bound_replica_id(request_id) + def _get_comprehension_stage_index(self) -> int | None: fallback_idx: int | None = None for idx, stage_client in enumerate(self.engine.stage_clients): @@ -331,6 +345,8 @@ class AsyncOmni(EngineClient, OmniBase): self.log_stats, wall_start_ts, final_stage_id_for_e2e, + transfer_emitter=getattr(self, "transfer_metrics", None), + replica_resolver=self._resolve_transfer_replica, ) req_state = ClientRequestState( @@ -338,6 +354,7 @@ class AsyncOmni(EngineClient, OmniBase): external_request_id=external_request_id, ) req_state.metrics = metrics + req_state.request_arrival_ts = wall_start_ts self.request_states[request_id] = req_state # PD disaggregation: modify prefill-stage sampling params per request @@ -387,10 +404,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) 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) await self._abort_internal_requests(request_id) logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}") raise diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 680ab9838..316d4c2aa 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -17,3 +17,23 @@ class ClientRequestState: self.stage_id: int | None = None self.queue = queue if queue is not None else asyncio.Queue() self.metrics: OrchestratorAggregator | None = None + # 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. + self.request_arrival_ts: float = 0.0 + # Wall-clock time at which the first audio packet was observed for + # this request. None means the streaming hook hasn't fired yet. + # Used as the once-per-request guard for audio_ttfp_s emit. + self.first_audio_ts: float | None = None + # Per-chunk timeline (seconds since request_arrival_ts) and PCM byte + # counts for the audio streaming response. Populated by the streaming + # endpoint on every audio.chunk emit; consumed at request finalize to + # compute audio_underrun_s and audio_continuity_ok_total. + self.audio_chunk_arrivals_s: list[float] = [] + self.audio_chunk_bytes: list[int] = [] + self.audio_sample_rate: int | None = None + # Stage / replica that produced the audio packets — captured at the + # first-packet hook so the finalize-time emit can label correctly + # without re-querying stage_pools. + self.audio_emit_stage_id: int | None = None + self.audio_emit_replica_id: int | None = None diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 9ddd32e9d..f538645b4 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -23,7 +23,10 @@ from vllm_omni.engine.messages import ( from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e +from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_finalize +from vllm_omni.metrics.prometheus import OmniPrometheusMetrics from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics +from vllm_omni.metrics.transfer import OmniTransferMetrics from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.outputs import OmniRequestOutput @@ -170,6 +173,10 @@ class OmniBase(PDDisaggregationMixin): self.tts_batch_max_items: int = kwargs.pop("tts_batch_max_items", 32) logger.info("[%s] Initializing with model %s", self.__class__.__name__, model) + # Construct transfer_metrics first so we can hand it to AsyncOmniEngine + # (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) st = time.time() self.engine = AsyncOmniEngine( model=model, @@ -177,6 +184,8 @@ class OmniBase(PDDisaggregationMixin): init_timeout=init_timeout, stage_init_timeout=stage_init_timeout, diffusion_batch_size=diffusion_batch_size, + transfer_emitter=self.transfer_metrics, + log_stats=log_stats, **kwargs, ) self._shutdown_called = False @@ -190,6 +199,8 @@ class OmniBase(PDDisaggregationMixin): self.async_chunk = bool(getattr(self.engine, "async_chunk", False)) self.request_states: dict[str, ClientRequestState] = {} + 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 if not self.output_modalities: @@ -272,11 +283,29 @@ class OmniBase(PDDisaggregationMixin): raise ValueError(f"Expected {self.num_stages} sampling params, got {len(normalized)}") return normalized + 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. + """ + 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: + 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 = 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() if self.log_stats: # Emit per-request orchestrator timing (including e2e_total_ms) # before dropping request state. @@ -289,6 +318,16 @@ class OmniBase(PDDisaggregationMixin): ) finally: self.request_states.pop(request_id, None) + # Republish gauges so any stale value left by the per-stage + # publish in _process_single_result (which runs while the request + # is still in self.request_states) is corrected after the pop. + prom = getattr(self, "prom_metrics", None) + counter = getattr(getattr(self, "engine", None), "_running_counter", None) + if prom is not None: + total = len(self.request_states) + running = counter.value if counter is not None else total + prom.set_running(running) + prom.set_waiting(max(0, total - running)) def _compute_final_stage_id(self, output_modalities: list[str] | None) -> int: return get_final_stage_id_for_e2e( @@ -435,6 +474,8 @@ class OmniBase(PDDisaggregationMixin): if not stage_meta.final_output: return None + output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type) + try: rid_key = str(req_id) if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished: @@ -443,10 +484,41 @@ class OmniBase(PDDisaggregationMixin): req_id, req_start_ts.get(req_id, wall_start_ts), ) + e2e_seconds = now - req_start_ts.get(req_id, wall_start_ts) + # Extract finished_reason from upstream CompletionOutput so + # the per-reason completion Counter is labelled correctly. + completion_outputs = getattr(engine_outputs, "outputs", None) or [] + fr = (getattr(completion_outputs[0], "finish_reason", None) if completion_outputs else None) or "stop" + self.prom_metrics.request_succeeded( + e2e_seconds, + finished_reason=fr, + ) + + # Modality observe inside the same finalize guard so it fires + # once per request and inherits the try/except isolation. + observe_modality_at_finalize( + self.mod_metrics, + output_type=output_type, + stage_id=stage_id, + replica_id=result.replica_id, + stage_metrics=_m, + engine_outputs=engine_outputs, + ) except Exception: logger.exception("[%s] Finalize request handling error", self.__class__.__name__) - output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type) + # When this result finalizes the request, the orchestrator has + # already decremented _running_counter but _log_summary_and_cleanup + # hasn't popped self.request_states yet — exclude the finalizing + # request from `total` so waiting doesn't read 1 and stay stuck + # there until the next request arrives. + counter = getattr(self.engine, "_running_counter", None) + is_finalizing = finished and stage_id == final_stage_id_for_e2e + total = max(0, len(self.request_states) - (1 if is_finalizing else 0)) + running = counter.value if counter is not None else total + self.prom_metrics.set_running(running) + self.prom_metrics.set_waiting(max(0, total - running)) + images = getattr(engine_outputs, "images", []) if output_type == "image" else [] return OmniRequestOutput( request_id=req_id or "", diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 2abfefe82..7e1d3027b 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -20,6 +20,11 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse from vllm_omni.entrypoints.utils import coerce_param_message_types from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.metrics import definitions as _metric_defs +from vllm_omni.metrics.modality import ( + observe_audio_first_packet, + observe_audio_streaming_finalize, +) try: import soundfile @@ -110,6 +115,7 @@ from vllm_omni.entrypoints.openai.utils import ( ) from vllm_omni.lora.request import LoRARequest from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate logger = init_logger(__name__) @@ -1032,6 +1038,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): include_usage, include_continuous_usage = should_include_usage(stream_options, self.enable_force_include_usage) last_metrics: dict[str, Any] | None = None + # Hold a strong reference to the audio request state so the + # streaming-finalize hook below survives the inner generator's + # _log_summary_and_cleanup, which pops request_states[request_id] + # before this outer block runs. + req_state_audio_ref: Any = None try: async for omni_res in result_generator: final_output_type = omni_res.final_output_type @@ -1605,8 +1616,53 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): yield f"data: {data}\n\n" elif final_output_type == "audio": + # Observe audio_ttfp_s on first audio packet for this request_id + # (once-per-request guard via first_audio_ts). The same hook + # also captures (stage, replica) for the streaming-continuity + # emit at request finalize. self.engine_client.request_states + # is keyed by the internal UUID-suffixed id (set by + # AsyncOmni.generate via _get_unique_request_id), but the + # `request_id` we have here is the external (user-visible) id, + # so resolve via the external_request_id field. + req_state = next( + (s for s in self.engine_client.request_states.values() if s.external_request_id == request_id), + None, + ) + if req_state is not None and req_state_audio_ref is None: + req_state_audio_ref = req_state + now_ts = time.time() + if req_state is not None and req_state.first_audio_ts is None: + req_state.first_audio_ts = now_ts + stage_pools = getattr(self.engine_client.engine, "stage_pools", None) + # The orchestrator binds requests by their internal id, + # not the user-visible external id, so look up the + # replica with req_state.request_id (internal). + replica_id = ( + stage_pools[omni_res.stage_id].get_bound_replica_id(req_state.request_id) + if stage_pools is not None and 0 <= omni_res.stage_id < len(stage_pools) + else None + ) + req_state.audio_emit_stage_id = omni_res.stage_id + req_state.audio_emit_replica_id = replica_id + observe_audio_first_packet( + self.engine_client.mod_metrics, + stage_id=omni_res.stage_id, + replica_id=replica_id, + arrival_ts=req_state.request_arrival_ts, + now_ts=now_ts, + ) + role = self.get_chat_request_role(request) choices_data = self._create_audio_choice(omni_res, role, request, stream=True) + # Record per-chunk PCM byte count + arrival timestamp for + # audio_underrun_s / audio_continuity_ok_total at finalize. + if req_state is not None and req_state.request_arrival_ts > 0: + chunk_bytes = audio_chunk_pcm_bytes(omni_res) + if chunk_bytes > 0: + req_state.audio_chunk_arrivals_s.append(max(now_ts - req_state.request_arrival_ts, 0.0)) + req_state.audio_chunk_bytes.append(chunk_bytes) + if req_state.audio_sample_rate is None: + req_state.audio_sample_rate = audio_chunk_sample_rate(omni_res) chunk = OmniChatCompletionStreamResponse( id=request_id, object=chunk_object_type, @@ -1627,6 +1683,24 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): logger.warning(f"Unsupported streaming final output type: {final_output_type}") continue + # Emit audio_underrun_s + audio_continuity_ok_total once per + # request after the audio chunk stream is exhausted. The + # captured reference (req_state_audio_ref) outlives the inner + # _log_summary_and_cleanup that pops request_states[request_id]. + if ( + req_state_audio_ref is not None + and req_state_audio_ref.audio_chunk_arrivals_s + and req_state_audio_ref.audio_emit_replica_id is not None + ): + observe_audio_streaming_finalize( + self.engine_client.mod_metrics, + stage_id=req_state_audio_ref.audio_emit_stage_id or 0, + replica_id=req_state_audio_ref.audio_emit_replica_id, + chunk_arrival_times_s=req_state_audio_ref.audio_chunk_arrivals_s, + chunk_bytes=req_state_audio_ref.audio_chunk_bytes, + sample_rate=(req_state_audio_ref.audio_sample_rate or _metric_defs.DEFAULT_AUDIO_SAMPLE_RATE), + ) + # once the final token is handled, if stream_options.include_usage # is sent, send the usage if include_usage: diff --git a/vllm_omni/metrics/__init__.py b/vllm_omni/metrics/__init__.py index deceb2333..6814a5891 100644 --- a/vllm_omni/metrics/__init__.py +++ b/vllm_omni/metrics/__init__.py @@ -1,7 +1,10 @@ +from .prometheus import OmniPrometheusMetrics, OmniRequestCounter from .stats import OrchestratorAggregator, StageRequestStats, StageStats from .utils import count_tokens_from_outputs __all__ = [ + "OmniPrometheusMetrics", + "OmniRequestCounter", "OrchestratorAggregator", "StageStats", "StageRequestStats", diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py new file mode 100644 index 000000000..b45217abc --- /dev/null +++ b/vllm_omni/metrics/definitions.py @@ -0,0 +1,221 @@ +"""Single source of truth for vLLM-Omni Prometheus + bench CLI metric naming. + +Consumed by: +- vllm_omni.metrics.prometheus (server-side /metrics pipeline families) +- vllm_omni.metrics.modality (audio families) +- vllm_omni.metrics.transfer (cross-stage transfer families) +- vllm_omni.benchmarks.metrics.metrics (bench CLI MultiModalsBenchmarkMetrics) + +Naming conventions for the ``vllm_omni:*`` families exposed here: +time-bearing metrics use the ``_s`` suffix (values in seconds), counters use +``_total`` (auto-suffixed by the prometheus client), sizes use ``_bytes``. +""" + +# vllm:omni_ avoids upstream's unregister_vllm_metrics() stripping, which +# removes every collector whose ``_name`` does not start with ``vllm``. +METRIC_PREFIX = "vllm:omni_" + + +# ============================================================================ +# Bench-side stems (also used as RequestFuncOutput attribute names) +# ============================================================================ +AUDIO_TTFP = "audio_ttfp" +AUDIO_DURATION = "audio_duration" +AUDIO_RTF = "audio_rtf" +AUDIO_FRAMES = "audio_frames" +AUDIO_UNDERRUN = "audio_underrun" +AUDIO_CONTINUITY_OK = "audio_continuity_ok" +AUDIO_SKIPPED_REQUESTS = "audio_skipped_requests" + + +# ============================================================================ +# Pipeline-level metric families (request counts + e2e latency) +# ============================================================================ +NUM_REQUESTS_RUNNING = METRIC_PREFIX + "num_requests_running" +NUM_REQUESTS_WAITING = METRIC_PREFIX + "num_requests_waiting" +E2E_REQUEST_LATENCY_S = METRIC_PREFIX + "e2e_request_latency_s" + +# Per-finished_reason Counter; finished_reason ∈ {stop, length, abort, ...}. +# Aborts include client disconnect / cancellation paths. Counter auto-suffixes +# ``_total`` at exposition time. +REQUESTS_SUCCESS = METRIC_PREFIX + "requests_success" + + +# ============================================================================ +# Audio family (per-stage + per-replica audio path metrics) +# ============================================================================ +AUDIO_TTFP_S = METRIC_PREFIX + AUDIO_TTFP + "_s" +AUDIO_DURATION_S = METRIC_PREFIX + AUDIO_DURATION + "_s" +AUDIO_RTF_METRIC = METRIC_PREFIX + AUDIO_RTF +AUDIO_FRAMES_METRIC = METRIC_PREFIX + AUDIO_FRAMES +AUDIO_UNDERRUN_S = METRIC_PREFIX + AUDIO_UNDERRUN + "_s" +AUDIO_CONTINUITY_OK_METRIC = METRIC_PREFIX + AUDIO_CONTINUITY_OK +AUDIO_SKIPPED_REQUESTS_METRIC = METRIC_PREFIX + AUDIO_SKIPPED_REQUESTS + + +# ============================================================================ +# Cross-stage Transfer family (per-physical-hop TX/RX/in-flight timings) +# ============================================================================ +TRANSFER_SIZE_BYTES = METRIC_PREFIX + "transfer_size_bytes" +TRANSFER_TX_S = METRIC_PREFIX + "transfer_tx_s" +TRANSFER_RX_S = METRIC_PREFIX + "transfer_rx_s" +TRANSFER_IN_FLIGHT_S = METRIC_PREFIX + "transfer_in_flight_s" + + +# ============================================================================ +# Label sets +# ============================================================================ +PIPELINE_LABELS = ("model_name",) +SUCCESS_LABELS = ("model_name", "finished_reason") + +# Per-stage / per-replica label set used by the audio families and by the +# OmniPrometheusStatLogger wrap which relabels upstream ``engine`` into +# ``stage`` + ``replica``. +STAGE_LABELS = ("model_name", "stage", "replica") + +# 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. +AUDIO_CONTINUITY_LABELS = ("model_name", "stage", "replica", "threshold_ms") + +# Audio skipped-requests Counter carries a `reason` label so the silent-loss +# path (e.g. code2wav rejecting malformed codec input) can be distinguished +# from other "200 OK + empty audio" cases. +AUDIO_SKIPPED_LABELS = ("model_name", "stage", "replica", "reason") + +# Cross-stage transfer label set. Each observation is one physical hop from +# (from_stage, from_replica) to (to_stage, to_replica). +TRANSFER_LABELS = ( + "model_name", + "from_stage", + "from_replica", + "to_stage", + "to_replica", +) + + +# ============================================================================ +# Histogram buckets +# ============================================================================ +# Seconds bucket for e2e / generation / TTFP-style metrics that range from +# ~10 ms to several minutes. +SECONDS_BUCKETS = ( + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.0, + 5.0, + 10.0, + 20.0, + 30.0, + 60.0, + 120.0, + 300.0, +) + +# Seconds bucket for fine-grained metrics (cross-stage transfer + audio +# underrun) that need millisecond-level resolution. +SECONDS_FAST_BUCKETS = ( + 0.001, + 0.005, + 0.01, + 0.025, + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.5, + 5.0, + 10.0, + 60.0, +) + +# RTF SLO red line is 1.0 — TTS must generate faster than playback. +RTF_BUCKETS = ( + 0.1, + 0.25, + 0.5, + 0.75, + 0.9, + 1.0, + 1.25, + 1.5, + 2.0, + 5.0, + 10.0, +) + +# Bytes bucket for transfer payload size. +BYTES_BUCKETS = ( + 1024, + 4096, + 16384, + 65536, + 262144, + 1048576, + 4194304, + 16777216, + 67108864, + 268435456, +) + + +# ============================================================================ +# Audio-continuity defaults +# ============================================================================ +# Default underrun threshold — kept aligned with the bench-side default and +# the commonly-cited "audible gap" threshold for streaming TTS. +AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S = 0.1 + + +# ============================================================================ +# Formula helpers (shared by server-side observe and bench-side calculation) +# ============================================================================ +def compute_audio_rtf(stage_gen_time_s: float, audio_duration_s: float) -> float: + """RTF = stage_gen_time / audio_content_duration. + + SLO red line < 1 — must generate faster than content plays back to stream. + Returns 0.0 when audio_duration_s is non-positive (caller decides whether + to observe; we don't want to divide by zero or emit negative samples). + """ + if audio_duration_s <= 0: + return 0.0 + return stage_gen_time_s / audio_duration_s + + +# ============================================================================ +# Audio sample-rate resolution +# ============================================================================ +# Most common across vllm-omni talker variants (cosyvoice3, omnivoice, +# qwen3_tts, mimo_audio). voxcpm2 uses 48000, stable_audio 44100, +# ming_flash 16000 — these models populate multimodal_output["audio_sample_rate"] +# at runtime so this default only kicks in when the field is missing. +DEFAULT_AUDIO_SAMPLE_RATE = 24000 + +_SAMPLE_RATE_KEYS = ("audio_sample_rate", "sample_rate", "sampling_rate", "sr") + + +def resolve_audio_sample_rate(multimodal_output: dict | None) -> int: + """Extract audio sample_rate from a multimodal_output dict, with fallbacks. + + Tries the same key chain as serving_chat.py's audio response path so + /metrics audio_duration_s = audio_frames / sample_rate stays consistent + with what the OpenAI streaming endpoint reports back to clients. + Returns DEFAULT_AUDIO_SAMPLE_RATE when no usable value is present. + """ + if not multimodal_output: + return DEFAULT_AUDIO_SAMPLE_RATE + for key in _SAMPLE_RATE_KEYS: + raw = multimodal_output.get(key) + if raw is None: + continue + try: + value = int(raw) + except (TypeError, ValueError): + continue + if value > 0: + return value + return DEFAULT_AUDIO_SAMPLE_RATE diff --git a/vllm_omni/metrics/modality.py b/vllm_omni/metrics/modality.py new file mode 100644 index 000000000..73d46ee8b --- /dev/null +++ b/vllm_omni/metrics/modality.py @@ -0,0 +1,303 @@ +"""OmniModalityMetrics — per-modality Prometheus families (audio path only). + +7 audio business-semantic metric families. Text-path metrics (TTFT / ITL / +TPOT / e2e) are NOT here — they come from the upstream +``vllm:*{stage="thinker", ...}`` families exposed via the +``OmniPrometheusStatLogger`` wrap. + +Contents: +- Audio family declarations (Histograms + Counters) +- ``OmniModalityMetrics``: label-bound observe API for the audio family +- ``observe_modality_at_finalize``: dispatcher called from omni_base's e2e + finalize hook; currently handles the audio path only. +- ``observe_audio_first_packet``: TTFP emit from the streaming SSE first + audio packet. +- ``observe_audio_streaming_finalize``: emits ``audio_underrun_s`` + + ``audio_continuity_ok_total`` at SSE close using accumulated per-chunk + arrival timestamps. +- ``_extract_mm_output`` / ``_count_audio_frames``: shape-tolerant helpers + for the heterogeneous multimodal_output payloads emitted by different + audio pipelines. +""" + +from __future__ import annotations + +from typing import Any + +from prometheus_client import Counter, Histogram + +from vllm_omni.metrics import definitions as defs + +_stage_labels = list(defs.STAGE_LABELS) + + +# ---------------------------------------------------------------------------- +# Audio family +# ---------------------------------------------------------------------------- +_audio_ttfp_family = Histogram( + defs.AUDIO_TTFP_S, + "Time from request arrival to first audio packet/frame, in seconds.", + labelnames=_stage_labels, + buckets=defs.SECONDS_BUCKETS, +) +_audio_duration_family = Histogram( + defs.AUDIO_DURATION_S, + "Generated audio content duration in seconds (audio_frames / sample_rate).", + labelnames=_stage_labels, + buckets=defs.SECONDS_BUCKETS, +) +_audio_rtf_family = Histogram( + defs.AUDIO_RTF_METRIC, + "Audio real-time factor (stage_gen_time_s / audio_duration_s); streaming TTS requires < 1.", + labelnames=_stage_labels, + buckets=defs.RTF_BUCKETS, +) +_audio_frames_family = Counter( + defs.AUDIO_FRAMES_METRIC, + "Cumulative audio frame count; per-model rate (not summable across models). Throughput recovered via rate().", + labelnames=_stage_labels, +) +_audio_underrun_family = Histogram( + defs.AUDIO_UNDERRUN_S, + "Per-request worst-case player-deficit in seconds (max time the player " + "ran out of buffered audio). > 0 indicates listener experienced silent gaps.", + labelnames=_stage_labels, + buckets=defs.SECONDS_FAST_BUCKETS, +) +_audio_continuity_ok_family = Counter( + defs.AUDIO_CONTINUITY_OK_METRIC, + "Incremented when the request's worst underrun stayed below threshold_ms. " + "Pair with requests_success_total to compute streaming-continuity health rate.", + labelnames=list(defs.AUDIO_CONTINUITY_LABELS), +) +_audio_skipped_family = Counter( + defs.AUDIO_SKIPPED_REQUESTS_METRIC, + "Silent-loss counter — code2wav rejected malformed codec input and returned 200 OK with empty audio.", + labelnames=list(defs.AUDIO_SKIPPED_LABELS), +) + + +class OmniModalityMetrics: + """Per-modality observe API. Stage/replica are passed at observe time + because a single OmniModalityMetrics instance per pipeline serves all + stage+replica combinations. + """ + + def __init__(self, model_name: str, log_stats: bool = True) -> None: + self._model_name = model_name + self._log_stats = log_stats + + # ---- Audio ------------------------------------------------------------ + + def observe_audio_ttfp(self, stage: str, replica: str, ttfp_seconds: float) -> None: + if not self._log_stats: + return + _audio_ttfp_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(ttfp_seconds) + + def observe_audio_duration(self, stage: str, replica: str, duration_seconds: float) -> None: + if not self._log_stats: + return + _audio_duration_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe( + duration_seconds + ) + + def observe_audio_rtf(self, stage: str, replica: str, rtf: float) -> None: + if not self._log_stats: + return + _audio_rtf_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(rtf) + + def inc_audio_frames(self, stage: str, replica: str, n_frames: int) -> None: + if not self._log_stats or n_frames <= 0: + return + _audio_frames_family.labels(model_name=self._model_name, stage=stage, replica=replica).inc(n_frames) + + def observe_audio_underrun(self, stage: str, replica: str, underrun_s: float) -> None: + if not self._log_stats: + return + _audio_underrun_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe( + max(underrun_s, 0.0) + ) + + def inc_audio_continuity_ok(self, stage: str, replica: str, threshold_ms: int) -> None: + if not self._log_stats: + return + _audio_continuity_ok_family.labels( + model_name=self._model_name, + stage=stage, + replica=replica, + threshold_ms=str(int(threshold_ms)), + ).inc() + + def inc_audio_skipped(self, stage: str, replica: str, reason: str) -> None: + if not self._log_stats: + return + _audio_skipped_family.labels( + model_name=self._model_name, + stage=stage, + replica=replica, + reason=reason or "unknown", + ).inc() + + +def _extract_mm_output(engine_outputs: Any) -> dict[str, Any]: + """Return the multimodal_output dict regardless of where it's nested. + + Three shapes seen in the wild: + * ``engine_outputs.multimodal_output`` — synthesized on OmniRequestOutput + for some pipelines (often empty for AR audio) + * ``engine_outputs.outputs[0].multimodal_output`` — vllm CompletionOutput + nesting (where actual qwen3-omni audio data lives) + * neither — returns ``{}`` + """ + mm = getattr(engine_outputs, "multimodal_output", None) + if isinstance(mm, dict) and mm: + return mm + outs = getattr(engine_outputs, "outputs", None) + if outs: + nested = getattr(outs[0], "multimodal_output", None) + if isinstance(nested, dict): + return nested + return {} + + +def _count_audio_frames(mm_out: dict[str, Any]) -> int: + """Sum the per-tensor sample count of audio chunks in mm_out["audio"]. + + Returns the total number of audio frames (samples) across all chunks. + For multi-dim tensors (e.g. shape [channels, samples]) the last axis is + treated as the sample dim; for 1-D tensors the only axis is the sample + dim; scalars count as 1. + """ + audio_chunks = mm_out.get("audio") if isinstance(mm_out, dict) else None + if audio_chunks is None: + return 0 + chunks = audio_chunks if isinstance(audio_chunks, list) else [audio_chunks] + n = 0 + for t in chunks: + try: + ndim = getattr(t, "ndim", 0) + shape = getattr(t, "shape", None) + if ndim == 0 or shape is None or len(shape) == 0: + n += 1 + else: + n += int(shape[-1]) + except Exception: + continue + return n + + +def observe_modality_at_finalize( + mod_metrics: OmniModalityMetrics, + *, + output_type: str | None, + stage_id: int, + replica_id: int | None, + stage_metrics: Any, + engine_outputs: Any, +) -> None: + """Route audio-path 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. + + 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: + return + if output_type != "audio": + 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) + + sample_rate = defs.resolve_audio_sample_rate(mm_out) + # `stage_metrics.audio_generated_frames` is the legacy per-chunk + # accumulator field on StageRequestStats. No production path currently + # fills it, so the fallback below is the live source — but we leave the + # field lookup in place in case the accumulator gets re-wired upstream. + 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: + # Request completed (finish_reason ∈ {stop, length} — error paths + # don't reach finalize) but no audio samples were produced. Covers + # silent `return None` skips in the talker→code2wav stage + # processors and the `parsed.append((0,0))` malformed-length path + # in qwen3-tts code2wav. raise-paths surface via the upstream + # vllm:request_success_total{finished_reason="error"} channel and + # never reach this branch. + mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data") + # audio_underrun / continuity are emitted from the streaming path in + # observe_audio_streaming_finalize; finalize is too late for the + # per-chunk timeline they need. + + +def observe_audio_first_packet( + mod_metrics: OmniModalityMetrics, + *, + stage_id: int, + replica_id: int | None, + arrival_ts: float, + now_ts: float, +) -> None: + """Observe audio_ttfp_s on a request's first audio packet. + + Caller is responsible for the once-per-request guard (e.g. checking + ``ClientRequestState.first_audio_ts is None``) so this function fires at + most once per request_id. Defensive-skips when ``replica_id`` or + ``arrival_ts`` is insufficient — both can legitimately be missing in error + paths and we'd rather drop the sample than emit a wrong (stage, replica). + """ + if replica_id is None or arrival_ts <= 0: + return + ttfp = max(now_ts - arrival_ts, 0.0) + mod_metrics.observe_audio_ttfp(str(stage_id), str(replica_id), ttfp) + + +def observe_audio_streaming_finalize( + mod_metrics: OmniModalityMetrics, + *, + stage_id: int, + replica_id: int | None, + chunk_arrival_times_s: list[float], + chunk_bytes: list[int], + sample_rate: int, + threshold_s: float = defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S, +) -> None: + """Emit audio_underrun_s + audio_continuity_ok_total at request end. + + Reuses the math from ``vllm_omni.benchmarks.audio_continuity`` so the + server-side and bench-side definitions stay aligned. Caller is responsible + for collecting per-chunk arrival timestamps and byte sizes during the + streaming response. + """ + if replica_id is None or not chunk_arrival_times_s: + return + # Local import to keep the bench module optional at import time. + from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats + + stats = compute_continuity_stats( + chunk_arrival_times_s=chunk_arrival_times_s, + chunk_bytes=chunk_bytes, + sample_rate=sample_rate, + threshold_s=threshold_s, + ) + stage_label = str(stage_id) + replica_label = str(replica_id) + mod_metrics.observe_audio_underrun(stage_label, replica_label, stats.max_underrun_s) + if stats.is_continuous: + mod_metrics.inc_audio_continuity_ok(stage_label, replica_label, int(threshold_s * 1000)) diff --git a/vllm_omni/metrics/prometheus.py b/vllm_omni/metrics/prometheus.py new file mode 100644 index 000000000..fc971df03 --- /dev/null +++ b/vllm_omni/metrics/prometheus.py @@ -0,0 +1,91 @@ +from prometheus_client import Counter, Gauge, Histogram + +from vllm_omni.metrics import definitions as defs + +_labelnames = list(defs.PIPELINE_LABELS) + +_running_family = Gauge( + defs.NUM_REQUESTS_RUNNING, + "Number of requests currently running across all pipeline stages.", + labelnames=_labelnames, +) +_waiting_family = Gauge( + defs.NUM_REQUESTS_WAITING, + "Number of requests waiting to be scheduled.", + labelnames=_labelnames, +) +_completion_family = Counter( + defs.REQUESTS_SUCCESS, + "Total requests by completion reason " + "(stop / length / abort / ...). Aborts cover client-disconnect / " + "cancellation paths in addition to upstream FinishReason.ABORT.", + labelnames=list(defs.SUCCESS_LABELS), +) +_e2e_latency_family = Histogram( + defs.E2E_REQUEST_LATENCY_S, + "Pipeline-global end-to-end request latency in seconds (user arrival to complete response).", + labelnames=_labelnames, + buckets=defs.SECONDS_BUCKETS, +) + + +class OmniPrometheusMetrics: + """Label-bound wrapper around the raw Prometheus metrics. + + Metric collectors use the ``vllm:omni_`` prefix to avoid being + removed by upstream vLLM's ``unregister_vllm_metrics()``, which + strips every collector whose ``_name`` contains ``"vllm"``. + """ + + def __init__(self, model_name: str, log_stats: bool = True) -> None: + self._model_name = model_name + self._log_stats = log_stats + self._running = _running_family.labels(model_name=model_name) + self._waiting = _waiting_family.labels(model_name=model_name) + self._e2e_latency = _e2e_latency_family.labels(model_name=model_name) + + def set_running(self, n: int) -> None: + if not self._log_stats: + return + self._running.set(n) + + def set_waiting(self, n: int) -> None: + if not self._log_stats: + return + self._waiting.set(n) + + def request_succeeded( + self, + e2e_seconds: float, + finished_reason: str = "stop", + ) -> None: + if not self._log_stats: + return + _completion_family.labels( + model_name=self._model_name, + finished_reason=finished_reason, + ).inc() + self._e2e_latency.observe(e2e_seconds) + + 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() + + +class OmniRequestCounter: + """Running-request counter written by the orchestrator thread, read by the client thread.""" + + def __init__(self) -> None: + self.value = 0 + + def increment(self) -> None: + self.value += 1 + + def decrement(self) -> None: + self.value = max(0, self.value - 1) diff --git a/vllm_omni/metrics/stat_logger.py b/vllm_omni/metrics/stat_logger.py new file mode 100644 index 000000000..9f8dc14b1 --- /dev/null +++ b/vllm_omni/metrics/stat_logger.py @@ -0,0 +1,243 @@ +"""OmniPrometheusStatLogger — wrap upstream PrometheusStatLogger. + +Rewrites the upstream ``engine`` single-label scheme into a ``stage`` + +``replica`` two-label scheme so the ~37 ``vllm:*`` metric families gain +per-(stage, replica) visibility for multi-replica deployments. + +Contents: +- ``_ENGINE_INDEX_MAP``: process-wide engine_idx → (stage_name, replica_id) + lookup, populated by ``OmniPrometheusStatLogger.__init__``. +- ``_RelabelMixin``: rewrites ``labelnames`` at family creation and translates + ``.labels()`` calls; applied via ``_RelabelGauge`` / ``_RelabelCounter`` / + ``_RelabelHistogram``. +- ``_OmniPerfMetricsProm`` / ``_OmniSpecDecodingProm`` / ``_OmniKVConnectorProm``: + helper-class wraps so the upstream sub-collectors construct their internal + families through the relabel mixin too. +- ``OmniPrometheusStatLogger``: the subclass that wires everything together + and rewrites ``per_engine_labelvalues`` from 2-tuple to 3-tuple at setter + time. +""" + +from __future__ import annotations + +from prometheus_client import Counter, Gauge, Histogram +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorProm +from vllm.v1.metrics.loggers import PrometheusStatLogger +from vllm.v1.metrics.perf import PerfMetricsProm +from vllm.v1.spec_decode.metrics import SpecDecodingProm + +# Process-wide translation table written by OmniPrometheusStatLogger at init. +# Keys are flat engine_idx values (as upstream PrometheusStatLogger sees them); +# values are the (stage_name, replica_id_str) tuple we expose as labels. +# +# Module-level rather than per-instance because the wrapper metric classes are +# constructed by upstream's __init__ and never get a back-reference to the +# StatLogger that owns them. vLLM runs a single Orchestrator/StatLogger per +# process, so a module global is safe; tests isolate by .clear()ing first. +_ENGINE_INDEX_MAP: dict[int, tuple[str, str]] = {} + + +def _rewrite_labelnames(labelnames): + """Replace `engine` in ``labelnames`` with (`stage`, `replica`) in place. + + Preserves ordering (so ``["model_name", "engine", "reason"]`` becomes + ``["model_name", "stage", "replica", "reason"]``) and the original + container type (list vs tuple). + """ + if labelnames is None: + return labelnames + seq = list(labelnames) + if "engine" not in seq: + return labelnames + out: list[str] = [] + for name in seq: + if name == "engine": + out.extend(("stage", "replica")) + else: + out.append(name) + return type(labelnames)(out) if not isinstance(labelnames, list) else out + + +def _engine_to_stage_replica(engine_value) -> tuple[str, str]: + """Look up (stage, replica) for an engine_idx, accepting int or str input. + + Upstream emits engine values in two flavors: + - int form, e.g. ``gauge_engine_sleep_state.labels(engine=idx, ...)`` (loggers.py:510) + - str form, e.g. ``info_gauge.labels(**metrics_info)`` where ``metrics_info["engine"] = str(idx)`` (loggers.py:1055) + + Raises ``KeyError`` when the value is missing from the map — fail-fast is + preferable to silently emitting series under a wrong (stage, replica). + """ + key = int(engine_value) if isinstance(engine_value, str) else engine_value + return _ENGINE_INDEX_MAP[key] + + +class _RelabelMixin: + """Mixin: rewrite ``labelnames`` at family creation and ``.labels()`` calls. + + Handles all four upstream forms encountered in + ``vllm.v1.metrics.loggers.PrometheusStatLogger``: + + 1. ``.labels(engine=idx, ...)`` kwarg with int engine (loggers.py:510) + 2. ``.labels(model_name, str(idx), source)`` positional with str engine + (loggers.py:646, 679) + 3. ``.labels(**metrics_info)`` kwarg with str engine (loggers.py:1056) + 4. Families without an ``engine`` label — passthrough (e.g. lora_info) + + Drops into upstream's ``_gauge_cls`` / ``_counter_cls`` / ``_histogram_cls`` + class slots. + """ + + def __init__(self, *args, **kwargs): + # Remember where `engine` sat in the original labelnames so positional + # `.labels()` calls can splice (stage, replica) at the right offset. + labelnames = kwargs.get("labelnames") + if labelnames is not None: + original = list(labelnames) + self._engine_label_index = original.index("engine") if "engine" in original else -1 + kwargs["labelnames"] = _rewrite_labelnames(labelnames) + else: + self._engine_label_index = -1 + super().__init__(*args, **kwargs) + + def labels(self, *args, **kwargs): + if self._engine_label_index >= 0: + if args: + # Positional form. There are TWO upstream patterns: + # + # (a) Pre-rewritten path: create_metric_per_engine fans + # `per_engine_labelvalues` (already a 3-tuple + # [model_name, stage, replica] thanks to the property- + # descriptor setter on OmniPrometheusStatLogger) into + # `metric.labels(*values)`. len(args) matches the + # rewritten label set already, so just pass through. + # + # (b) Legacy 2-tuple path: upstream sites like + # `counter_request_success.labels(model_name, str(idx), + # str(reason))` pass values shaped to the *original* + # labelnames (engine still present at idx). Here + # len(args) is short by 1 — splice (stage, replica) + # in place of the engine value at engine_label_index. + if len(args) == len(self._labelnames): + return super().labels(*args, **kwargs) + idx = self._engine_label_index + if idx < len(args): + stage, replica = _engine_to_stage_replica(args[idx]) + args = (*args[:idx], stage, replica, *args[idx + 1 :]) + elif "engine" in kwargs: + stage, replica = _engine_to_stage_replica(kwargs.pop("engine")) + kwargs["stage"] = stage + kwargs["replica"] = replica + return super().labels(*args, **kwargs) + + +class _RelabelGauge(_RelabelMixin, Gauge): + pass + + +class _RelabelCounter(_RelabelMixin, Counter): + pass + + +class _RelabelHistogram(_RelabelMixin, Histogram): + pass + + +# ---------------------------------------------------------------------------- +# Helper-class wraps for the three sub-metric collectors that upstream +# PrometheusStatLogger constructs in its __init__ (loggers.py:438-446): +# +# self.spec_decoding_prom = self._spec_decoding_cls(...) +# self.kv_connector_prom = self._kv_connector_cls(...) +# self.perf_metrics_prom = self._perf_metrics_cls(...) +# +# Each helper receives raw `labelnames` as a constructor argument and uses +# its own class-level `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots +# to build internal Counter/Gauge/Histogram families. The slot overrides on +# OmniPrometheusStatLogger only reach families created via *its* slots, so +# the helpers would otherwise still construct 2-label families and then hit +# `Incorrect label count` when create_metric_per_engine feeds the rewritten +# 3-element per_engine_labelvalues. Subclassing each helper and overriding +# its slots routes the relabel mixin through to the helper-internal families +# too. The helper kept seeing the OLD 2-element labelnames param, but that +# is fine because the wrapper rewrites it at family-creation time. +# ---------------------------------------------------------------------------- + + +class _OmniPerfMetricsProm(PerfMetricsProm): + _counter_cls = _RelabelCounter + + +class _OmniSpecDecodingProm(SpecDecodingProm): + _counter_cls = _RelabelCounter + + +class _OmniKVConnectorProm(KVConnectorProm): + _gauge_cls = _RelabelGauge + _counter_cls = _RelabelCounter + _histogram_cls = _RelabelHistogram + + +class OmniPrometheusStatLogger(PrometheusStatLogger): + """Wrap upstream PrometheusStatLogger to expose per-(stage, replica) labels. + + Replaces the upstream single ``engine`` label with two labels ``stage`` and + ``replica`` so the ~37 ``vllm:*`` metric families gain per-replica + visibility for multi-replica deployments. + + The orchestrator builds ``stage_replica_map`` from the static stage_pools + config; flat engine_idx values map 1:1 to (stage_name, replica_id) tuples. + Dynamic add/remove of replicas at runtime is intentionally not supported — + the map is built once at construction and never mutated afterward. + """ + + # Inject our wrapper metric classes into upstream's class-level slots so + # every ~37 family is created with `engine` rewritten to `stage`+`replica`. + _gauge_cls = _RelabelGauge + _counter_cls = _RelabelCounter + _histogram_cls = _RelabelHistogram + # Inject helper-class wraps too so the perf / spec-decoding / kv-connector + # sub-collectors get the same labelname rewrite and don't crash with + # `Incorrect label count` when create_metric_per_engine fans out the + # rewritten 3-element per_engine_labelvalues over their internal families. + _perf_metrics_cls = _OmniPerfMetricsProm + _spec_decoding_cls = _OmniSpecDecodingProm + _kv_connector_cls = _OmniKVConnectorProm + + def __init__( + self, + vllm_config: VllmConfig, + stage_replica_map: dict[int, tuple[str, str]], + ) -> None: + self._stage_replica_map = stage_replica_map + # Populate the process-level translation table that wrapper metric + # classes consult on every `.labels()` call. Cleared first so a + # second OmniPrometheusStatLogger in the same process (e.g. tests, + # orchestrator restart) starts from a clean slate. + _ENGINE_INDEX_MAP.clear() + _ENGINE_INDEX_MAP.update(stage_replica_map) + super().__init__( + vllm_config=vllm_config, + engine_indexes=list(stage_replica_map.keys()), + ) + + @property + def stage_replica_map(self) -> dict[int, tuple[str, str]]: + return self._stage_replica_map + + @property + def per_engine_labelvalues(self) -> dict[int, list[object]]: + return self._omni_per_engine_labelvalues + + @per_engine_labelvalues.setter + def per_engine_labelvalues(self, value: dict[int, list[object]]) -> None: + # Upstream sets {idx: [model_name, str(idx)]} (loggers.py:433); we drop + # the engine str and append (stage, replica) so labelvalues match the + # 3-element labelnames our wrapper classes produce. + rewritten: dict[int, list[object]] = {} + for idx, vals in value.items(): + model_name = vals[0] + stage, replica = self._stage_replica_map[idx] + rewritten[idx] = [model_name, stage, replica] + self._omni_per_engine_labelvalues = rewritten diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 4245deb54..88e7e770a 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -5,12 +5,15 @@ from collections import defaultdict from collections.abc import Callable from contextlib import contextmanager from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any from vllm.logger import init_logger from vllm_omni.metrics.utils import _build_field_defs, _build_row, _format_table +if TYPE_CHECKING: + from vllm_omni.metrics.transfer import OmniTransferMetrics + logger = init_logger(__name__) @@ -36,6 +39,7 @@ class StageRequestStats: rx_in_flight_time_ms: float stage_stats: StageStats stage_id: int | None = None + replica_id: int | None = None final_output_type: str | None = None request_id: str | None = None postprocess_time_ms: float = 0.0 @@ -121,6 +125,9 @@ class OrchestratorAggregator: log_stats: bool, wall_start_ts: float, final_stage_id_for_e2e: dict[str, int] | int, + *, + transfer_emitter: OmniTransferMetrics | None = None, + replica_resolver: Callable[[int, str], int | None] | None = None, ) -> None: self.num_stages = int(num_stages) self.log_stats = bool(log_stats) @@ -131,6 +138,20 @@ class OrchestratorAggregator: tuple[int, int, str], TransferEdgeStats ] = {} # Key: (from_stage, to_stage, request_id) self.e2e_events: list[RequestE2EStats] = [] + # Emit per-physical-transfer Histogram observations to Prometheus + # alongside the existing TransferEdgeStats accumulation. Both deps + # are optional so OrchestratorAggregator stays usable in contexts + # that don't have a Prometheus registry (e.g. unit tests). + self._transfer_emitter = transfer_emitter + self._replica_resolver = replica_resolver + # Snapshot of (stage_id -> replica_id) per request, captured at + # ``record_transfer_rx`` time from the receiving stage's own + # ``StageRequestStats.replica_id``. Survives pool-side + # ``release_binding`` so the transfer emit can resolve label values + # even when the downstream stage's binding has already been released + # by the orchestrator's cleanup path (e.g. final-stage cleanup races + # the omni_base finalize emit). + self._replica_cache: dict[tuple[int, str], int] = {} def init_run_state(self, wall_start_ts: float) -> None: # Per-run aggregates and timing state @@ -191,6 +212,14 @@ class OrchestratorAggregator: evt.size_bytes += int(size_bytes) evt.tx_time_ms += float(tx_time_ms) evt.used_shm = evt.used_shm or bool(used_shm) + # Emit per-physical-transfer Histogram observations to Prometheus. + self._emit_transfer_tx( + from_stage=int(from_stage), + to_stage=int(to_stage), + request_id=str(request_id), + size_bytes=int(size_bytes), + tx_time_ms=float(tx_time_ms), + ) return evt except Exception: return None @@ -205,6 +234,11 @@ class OrchestratorAggregator: from_stage = int(stats.stage_id) - 1 to_stage = int(stats.stage_id) rid_key = str(stats.request_id) + # Snapshot the receiving (stage, replica) binding before the + # pool may release it on cleanup. Used by _resolve_edge_replicas + # as a fallback when get_bound_replica_id has already cleared. + if stats.replica_id is not None: + self._replica_cache[(to_stage, rid_key)] = int(stats.replica_id) evt = self._get_or_create_transfer_event(from_stage, to_stage, rid_key) # Accumulate rx metrics if evt.size_bytes == 0: @@ -212,10 +246,79 @@ class OrchestratorAggregator: evt.size_bytes = int(stats.rx_transfer_bytes) evt.rx_decode_time_ms += float(stats.rx_decode_time_ms) evt.in_flight_time_ms += float(stats.rx_in_flight_time_ms) + # Emit per-physical-receive Histogram observations to Prometheus. + self._emit_transfer_rx( + from_stage=from_stage, + to_stage=to_stage, + request_id=rid_key, + rx_decode_time_ms=float(stats.rx_decode_time_ms), + in_flight_time_ms=float(stats.rx_in_flight_time_ms), + ) return evt except Exception: return None + # ------------------------------------------------------------------ + # Prometheus emit hooks. Both helpers are no-ops when transfer_emitter + # or replica_resolver is None, or when the resolver cannot find a + # (stage_id, request_id) -> replica_id mapping. We deliberately + # fail-safe (skip) rather than emit a series with a wrong or invented + # replica label. + # ------------------------------------------------------------------ + + def _resolve_edge_replicas(self, from_stage: int, to_stage: int, request_id: str) -> tuple[int, int] | None: + # Prefer the per-request replica snapshot captured by + # record_transfer_rx (survives pool binding release); fall back to the + # resolver callback (pool.get_bound_replica_id) when the cache + # misses, e.g. for the TX side that fires before the receiving stage + # has recorded its stats. + from_r = self._replica_cache.get((from_stage, request_id)) + to_r = self._replica_cache.get((to_stage, request_id)) + if (from_r is None or to_r is None) and self._replica_resolver is not None: + if from_r is None: + from_r = self._replica_resolver(from_stage, request_id) + if to_r is None: + to_r = self._replica_resolver(to_stage, request_id) + if from_r is None or to_r is None: + return None + return int(from_r), int(to_r) + + def _emit_transfer_tx( + self, + *, + from_stage: int, + to_stage: int, + request_id: str, + size_bytes: int, + tx_time_ms: float, + ) -> None: + if self._transfer_emitter is None: + return + replicas = self._resolve_edge_replicas(from_stage, to_stage, request_id) + if replicas is None: + return + from_r, to_r = replicas + self._transfer_emitter.observe_size(from_stage, from_r, to_stage, to_r, size_bytes) + self._transfer_emitter.observe_tx_time(from_stage, from_r, to_stage, to_r, tx_time_ms / 1000.0) + + def _emit_transfer_rx( + self, + *, + from_stage: int, + to_stage: int, + request_id: str, + rx_decode_time_ms: float, + in_flight_time_ms: float, + ) -> None: + if self._transfer_emitter is None: + return + replicas = self._resolve_edge_replicas(from_stage, to_stage, request_id) + if replicas is None: + return + from_r, to_r = replicas + self._transfer_emitter.observe_rx_time(from_stage, from_r, to_stage, to_r, rx_decode_time_ms / 1000.0) + self._transfer_emitter.observe_in_flight_time(from_stage, from_r, to_stage, to_r, in_flight_time_ms / 1000.0) + def record_audio_generated_frames( self, output_to_yield: Any, @@ -462,6 +565,10 @@ class OrchestratorAggregator: ), ) self.e2e_events.append(per_req_record) + # Drop per-request replica snapshots once the request is finalized. + if self._replica_cache: + for cached_key in [k for k in self._replica_cache if k[1] == rid_key]: + self._replica_cache.pop(cached_key, None) def build_and_log_summary(self) -> dict[str, Any]: if not self.log_stats: diff --git a/vllm_omni/metrics/transfer.py b/vllm_omni/metrics/transfer.py new file mode 100644 index 000000000..0569c139d --- /dev/null +++ b/vllm_omni/metrics/transfer.py @@ -0,0 +1,145 @@ +"""OmniTransferMetrics — cross-stage transfer Prometheus families. + +Four families with ``{model_name, from_stage, from_replica, to_stage, +to_replica}`` labels. Each ``observe_*()`` call corresponds to one physical +transfer event (one chunk hop from a sender replica to a receiver replica), +so the Histogram tracks the distribution of physical transfers, not +request-aggregated sums. + +Data source: ``vllm_omni.metrics.stats.TransferEdgeStats`` accumulators in +``OrchestratorAggregator.record_transfer_tx`` / ``record_transfer_rx``. The +emit hook lives in stats.py; this module only registers the families and +exposes the typed observe API. +""" + +from __future__ import annotations + +from prometheus_client import Histogram + +from vllm_omni.metrics import definitions as defs + +_labelnames = list(defs.TRANSFER_LABELS) + + +# ---------------------------------------------------------------------------- +# TX-side families (observed when record_transfer_tx fires) +# ---------------------------------------------------------------------------- +_transfer_size_bytes_family = Histogram( + defs.TRANSFER_SIZE_BYTES, + "Per-transfer payload size in bytes (one observation per physical hop).", + labelnames=_labelnames, + buckets=defs.BYTES_BUCKETS, +) +_transfer_tx_family = Histogram( + defs.TRANSFER_TX_S, + "Sender-side time in seconds (serialize + submit to connector).", + labelnames=_labelnames, + buckets=defs.SECONDS_FAST_BUCKETS, +) + + +# ---------------------------------------------------------------------------- +# RX-side families (observed when record_transfer_rx fires) +# ---------------------------------------------------------------------------- +_transfer_rx_family = Histogram( + defs.TRANSFER_RX_S, + "Receiver-side time in seconds (recv + deserialize).", + labelnames=_labelnames, + buckets=defs.SECONDS_FAST_BUCKETS, +) +_transfer_in_flight_family = Histogram( + defs.TRANSFER_IN_FLIGHT_S, + "Network in-flight time in seconds (TX done -> RX recv start).", + labelnames=_labelnames, + buckets=defs.SECONDS_FAST_BUCKETS, +) + + +class OmniTransferMetrics: + """Per-(from, to) replica observe API for cross-stage transfers. + + A single instance per pipeline; ``model_name`` is bound at init and + every observe call carries it in the label set. Stage/replica are + passed at observe time because the same instance serves all + (from_stage, from_replica) -> (to_stage, to_replica) edges. + """ + + def __init__(self, model_name: str, log_stats: bool = True) -> None: + self._model_name = model_name + self._log_stats = log_stats + + # ---- TX side (record_transfer_tx hook) ------------------------------- + + def observe_size( + self, + from_stage: int, + from_replica: int, + to_stage: int, + to_replica: int, + size_bytes: int, + ) -> None: + if not self._log_stats: + return + _transfer_size_bytes_family.labels( + model_name=self._model_name, + from_stage=str(from_stage), + from_replica=str(from_replica), + to_stage=str(to_stage), + to_replica=str(to_replica), + ).observe(size_bytes) + + def observe_tx_time( + self, + from_stage: int, + from_replica: int, + to_stage: int, + to_replica: int, + tx_time_s: float, + ) -> None: + if not self._log_stats: + return + _transfer_tx_family.labels( + model_name=self._model_name, + from_stage=str(from_stage), + from_replica=str(from_replica), + to_stage=str(to_stage), + to_replica=str(to_replica), + ).observe(tx_time_s) + + # ---- RX side (record_transfer_rx hook) ------------------------------- + + def observe_rx_time( + self, + from_stage: int, + from_replica: int, + to_stage: int, + to_replica: int, + rx_time_s: float, + ) -> None: + if not self._log_stats: + return + _transfer_rx_family.labels( + model_name=self._model_name, + from_stage=str(from_stage), + from_replica=str(from_replica), + to_stage=str(to_stage), + to_replica=str(to_replica), + ).observe(rx_time_s) + + def observe_in_flight_time( + self, + from_stage: int, + from_replica: int, + to_stage: int, + to_replica: int, + in_flight_time_s: float, + ) -> None: + if not self._log_stats: + return + _transfer_in_flight_family.labels( + model_name=self._model_name, + from_stage=str(from_stage), + from_replica=str(from_replica), + to_stage=str(to_stage), + to_replica=str(to_replica), + ).observe(in_flight_time_s) diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index af774fa21..e5dc955eb 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,3 +1,4 @@ +import logging import sys from functools import cached_property @@ -129,6 +130,43 @@ for module_name, module in list(sys.modules.items()): if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest: module.EngineCoreRequest = OmniEngineCoreRequest +# ============================================================================= +# Patch unregister_vllm_metrics to skip vllm:omni_* collectors +# ============================================================================= +# WHY: Upstream unregister_vllm_metrics() uses `"vllm" in collector._name` to +# strip collectors before each new PrometheusStatLogger registers, which is +# how it avoids "Duplicated timeseries" when the same process spawns multiple +# engines / orchestrators. But its substring match also wipes our vllm:omni_* +# families, so we must replace it with a scoped version that keeps ours. +# +# REMOVAL: Remove this patch once upstream vLLM adds +# _STAT_LOGGER_METRIC_NAMES to vllm.v1.metrics.prometheus and scopes +# unregister_vllm_metrics() to that set. Track: +# https://github.com/vllm-project/vllm/pull/42331 +import vllm.v1.metrics.prometheus as _vllm_prometheus # noqa: E402 + +_logger = logging.getLogger(__name__) + + +def _scoped_unregister_vllm_metrics(): + """Drop upstream vllm:* collectors but preserve vllm:omni_* (ours).""" + from prometheus_client import REGISTRY + + for collector in list(REGISTRY._collector_to_names): + name = getattr(collector, "_name", "") + if "vllm" not in name: + continue + if name.startswith("vllm:omni_") or name.startswith("vllm_omni"): + continue + REGISTRY.unregister(collector) + + +_vllm_prometheus.unregister_vllm_metrics = _scoped_unregister_vllm_metrics +_logger.warning( + "Monkey-patched unregister_vllm_metrics() to scope drops to non-omni vllm:* collectors. " + "Remove this patch once vLLM adds _STAT_LOGGER_METRIC_NAMES." +) + # Patch: add qwen3_omni_moe to vllm's chat template fallback registry. # Qwen/Qwen3-Omni-30B-A3B-Instruct stores its chat_template in a standalone diff --git a/vllm_omni/utils/audio.py b/vllm_omni/utils/audio.py index cc25c1794..bc2498efd 100644 --- a/vllm_omni/utils/audio.py +++ b/vllm_omni/utils/audio.py @@ -7,6 +7,9 @@ import numpy as np import torch from torchaudio.functional import melscale_fbanks +from vllm_omni.metrics import definitions as _metric_defs +from vllm_omni.outputs import OmniRequestOutput + def mel_filter_bank( sr: int, @@ -66,3 +69,39 @@ def peak_normalize( return audio target = 10.0 ** (db_level / 20.0) return audio * (target / peak) + + +def audio_chunk_pcm_bytes(omni_res: OmniRequestOutput) -> int: + """Best-effort PCM byte count of the last audio chunk in ``omni_res``. + + Used by the audio-streaming continuity tracker to size the player buffer. + Returns 0 when the chunk shape can't be interpreted — caller drops the + sample rather than recording a wrong byte count. + """ + try: + final_res = omni_res.request_output + mm_output = final_res.outputs[0].multimodal_output + audio_data = mm_output.get("audio") + if isinstance(audio_data, list): + if not audio_data: + return 0 + audio_tensor = audio_data[-1] + else: + audio_tensor = audio_data + if audio_tensor is None: + return 0 + n_samples = int(audio_tensor.numel() if hasattr(audio_tensor, "numel") else audio_tensor.size) + # PCM s16le mono → 2 bytes per sample. This matches what + # _create_audio_choice serialises via CreateAudio (response_format="wav"). + return max(n_samples, 0) * 2 + except Exception: + return 0 + + +def audio_chunk_sample_rate(omni_res: OmniRequestOutput) -> int: + """Resolve audio sample rate for the request's audio stream.""" + try: + mm_output = omni_res.request_output.outputs[0].multimodal_output + except Exception: + return _metric_defs.DEFAULT_AUDIO_SAMPLE_RATE + return _metric_defs.resolve_audio_sample_rate(mm_output)