Compare commits

..
Author SHA1 Message Date
hsliuustc0106 9a7ea9b003 [Bugfix] Reconcile env-var inventory with post-#6217 main drift
Merge-window drift broke the new environment-variable ratchet on main
(build 14119): #6255 added the VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT read in
the diffusion engine and #6362 removed the last package reference to
ROSVOT_SOURCE_DIR, both after #6217's CI run.

- Classify VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT as PUBLIC_OMNI and document
  it in the environment-variables reference page.
- Drop the stale ROSVOT_SOURCE_DIR model-specific row now that the
  SoulX-Singer implementation is gone.
- Update the reviewed snapshot counts and disposition tables to match.

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-08-26 06:06:04 +08:00
Hongsheng Liu 780e8d40aa [Doc] Add environment variables configuration reference (#6217)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-26 04:58:54 +08:00
Satellite 6639c82a17 [CI/Build] Pin GGUF plugin for diffusion nightly tests (#6303)
Signed-off-by: Xin Wei <14196318+wxwxwwxxx@users.noreply.github.com>
2026-08-25 09:16:46 -07:00
Alex Brooks 0eb47b76d0 [Bugfix] Handle DoS Overflow Cases (#6598)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-08-25 16:09:03 +00:00
NATURE 497c537c6f [Benchmark] Add local OmniInteract realtime benchmark (#6522)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-08-25 15:32:58 +00:00
Hongsheng Liuandhsliuustc0106 c608910d20 [Diffusion] Register HWR mmap for direct DLO H2D (#6591)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-25 23:30:11 +08:00
Hongsheng Liu 5a8b81ebeb [Model] Optimize MiniMax-H3 DLO component lifecycle (#6526)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-08-25 23:13:41 +08:00
NumberWan dd527b4245 [CI] Reduce Qwen-Image Function and share step-execution perf server (#6613)
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
2026-08-25 22:55:15 +08:00
jl9876andChendi.Xue 0d5d1115d6 [AutoRound] Add offline quantized MXFP4 model support (#5544)
Signed-off-by: jl9876 <jia.li@intel.com>
Co-authored-by: Chendi.Xue <chendi.xue@intel.com>
2026-08-25 14:44:38 +00:00
psv666 683f3c92a4 [CI/Build] Clean up Qwen3-Omni nightly tests (#6570)
Signed-off-by: psv666 <2693925048@qq.com>
2026-08-25 22:14:43 +08:00
Jae 72beeaee86 [Misc] move text-to-audio online examples to a unified folder (#4807)
Signed-off-by: zzehli <jaeli_ottawa@outlook.com>
2026-08-25 22:02:26 +08:00
Bhavyansh Sharma d6b66faba7 [CI] Fix response_format json_schema error expectation (#6290)
Signed-off-by: Bhavyansh Sharma <66562679+clumsylad21@users.noreply.github.com>
2026-08-25 19:47:09 +08:00
NumberWan bc8ea489ef [BugFix] Qwen-image performance regressed - Avoid mapping diffusion_batch_size onto scheduler max_num_seqs (#6525)
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
2026-08-25 03:36:11 -07:00
80 changed files with 6490 additions and 1531 deletions

No files matched your search

+2 -2
View File
@@ -83,7 +83,7 @@ steps:
- export BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and full_model"
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
@@ -375,7 +375,7 @@ steps:
- label: ":full_moon: Diffusion · Quantization Test with L4"
timeout_in_minutes: 60
commands:
- pip install "vllm-gguf-plugin>=0.0.3"
- pip install "vllm-gguf-plugin==0.0.4"
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
mirror_hardwares: l4_1
+13
View File
@@ -74,6 +74,19 @@ steps:
exit $$EXIT
mirror_hardwares: h100_2
- label: ":full_moon: Omni · Perf Test · Async Chunk · Random"
key: weekly-omni-performance-async-chunk-random
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and slow"
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
mirror_hardwares: h100_2
- label: ":full_moon: Omni · Perf Test · Multi-Replica"
key: weekly-omni-performance-multi-replicas
timeout_in_minutes: 180
+21 -6
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# ruff: noqa
# code borrowed from https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py
@@ -494,6 +494,16 @@ def is_xnnpack_available():
def get_env_vars():
try:
from vllm_omni.config.environment_variable_inventory import ENVIRONMENT_VARIABLE_INVENTORY
except Exception:
# Keep environment collection usable when an incomplete/broken Omni
# installation is the reason the user is running this diagnostic.
# Importing an Omni submodule executes package initialization, which
# may fail with RuntimeError, AssertionError, or OSError as well as
# ImportError when vLLM or the accelerator stack is broken.
ENVIRONMENT_VARIABLE_INVENTORY = {}
env_vars = ""
secret_terms = ("secret", "token", "api", "access", "password")
report_prefix = (
@@ -507,12 +517,17 @@ def get_env_vars():
"MKL_",
"NVIDIA",
)
for k, v in os.environ.items():
if any(term in k.lower() for term in secret_terms):
for k, v in sorted(os.environ.items()):
omni_classification = ENVIRONMENT_VARIABLE_INVENTORY.get(k)
if (omni_classification is not None and omni_classification.redact_value) or any(
term in k.lower() for term in secret_terms
):
continue
if k in environment_variables:
env_vars = env_vars + "{}={}".format(k, v) + "\n"
if k.startswith(report_prefix):
if (
k in environment_variables
or (omni_classification is not None and omni_classification.is_public_omni)
or k.startswith(report_prefix)
):
env_vars = env_vars + "{}={}".format(k, v) + "\n"
return env_vars
+7 -1
View File
@@ -2,6 +2,11 @@
The CLI for vLLM-Omni inherits from vllm with some additional arguments.
Environment variables are process configuration rather than CLI-only options.
See the [Environment Variables](../configuration/environment_variables.md)
reference for ownership, precedence, evaluation time, and per-stage `env`
scoping.
## serve
Starts the vLLM-Omni OpenAI Compatible API server.
@@ -39,4 +44,5 @@ vllm bench serve --omni \
--num-prompts 5
```
See [vllm bench serve](./bench/serve.md) for the full reference of all available arguments.
See [vllm bench serve](./bench/serve.md) for serving benchmark arguments and
dataset-specific examples, including OmniInteract native-duplex sessions.
+41 -3
View File
@@ -2,7 +2,7 @@
The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
## Notes
vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` serving benchmark backends.
vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, `daily-omni`, and `openai-realtime-duplex` serving benchmark backends. It also adds the `omniinteract` dataset.
## Basic Parameter Description
You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
@@ -10,7 +10,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
- `--backend`
Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` to the upstream vLLM backend choices.
Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, `daily-omni`, and `openai-realtime-duplex` to the upstream vLLM backend choices.
- `--model`
The model identifier to load, filled according to the models supported by vLLM-Omni.
@@ -19,7 +19,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
The API endpoint exposed externally, to which clients send their requests.
- `--dataset-name`
The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio), while `omniinteract` replays official OmniInteract videos.
- `--num-prompts`
The total number of requests to send, an integer.
@@ -269,6 +269,44 @@ We use audio generation time / audio duration to calculate RTF.
</details>
### OmniInteract Realtime Benchmark
OmniInteract runs each video as one native-duplex WebSocket sample in the standard serving-benchmark lifecycle:
```bash
vllm bench serve --omni \
--backend openai-realtime-duplex \
--dataset-name omniinteract \
--dataset-path /path/to/OmniInteract \
--model openbmb/MiniCPM-o-4_5 \
--base-url http://127.0.0.1:8000 \
--endpoint /v1/realtime \
--omniinteract-ref-audio /path/to/reference.wav \
--omniinteract-output-dir ./omniinteract-artifacts \
--num-prompts 3 \
--num-warmups 0
```
`--dataset-path` accepts an extracted directory, `data.tar[.gz]`, or a Hugging Face dataset ID; omitting it uses
`lucky-lance/OmniInteract`. `--num-prompts` is the total across subsets and defaults to 3 for OmniInteract; explicit `0`
selects all and oversize values use all available cases. Reference audio is required, and OmniInteract uses the
`/v1/realtime` endpoint.
Audio is replayed as 16 kHz PCM16 in 200 ms chunks and video at 1 FPS with real-time pacing. All selected media is decoded
before timing and remains in client memory for the run, so `--max-concurrency` does not limit media preparation memory; use
explicit `--num-prompts 0` only when the client has enough RAM for the full dataset. Media commands are bounded by
`--omniinteract-media-timeout-s`, and concurrency defaults to 1. Standard request-rate, warmup, result-saving, and summary
options apply. Use `--omniinteract-require-response` only for functional E2E cases; LISTEN is a valid benchmark result.
Each completed case writes `output.wav`, `wav_transcript.json`, `events.json`, `result.json`, and a final `.done` marker under
`--omniinteract-output-dir`. The root also contains `batch_summary.json` and `official_eval_manifest.jsonl`; failed cases write
`.failed.json`. Runs sharing one output root are serialized. Completion validates transport, response lifecycle, and artifacts,
not answer accuracy. Transcript timestamps are serialized playback-queue times. Clipped or cancelled outputs are ineligible and
omitted from the official manifest; `audio_clipped_bytes` records output beyond the rounded video horizon.
TTFT, TTFP, and RTF start at client receipt of `response.created`. TPOT/ITL use engine stage-0 timing; ITL is emitted only when
every token interval is present.
### Multi-Modal Benchmark
<details class="admonition abstract" markdown="1">
+3
View File
@@ -7,6 +7,9 @@ For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration g
Each model defines fixed topology in a registered `PipelineConfig` and runtime
overrides in a deploy YAML.
For process-level settings shared by the CLI, deploy YAML, workers, and model
integrations, see [Environment Variables](./environment_variables.md).
For a specific example, see the [Qwen2.5-Omni deploy config](gh-file:vllm_omni/deploy/qwen2_5_omni.yaml). The matching frozen pipeline topology lives at [vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py](gh-file:vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py).
For an introduction, see [Pipeline and deploy configurations](./stage_configs.md).
+237
View File
@@ -0,0 +1,237 @@
# Environment Variables
Environment variables are process-level configuration shared by the CLI,
deploy YAML, Python APIs, worker processes, and model integrations. They are
therefore documented under **User Guide → Configuration**, rather than under
CLI Reference.
This page describes the public vLLM-Omni environment-variable contract on the
`latest` documentation branch. Use the documentation for your installed
vLLM-Omni release when behavior differs.
!!! important "Public contract"
Only names in [Public vLLM-Omni variables](#public-vllm-omni-variables) are
public Omni configuration. A name found in source code is not automatically
supported: model-specific, benchmark, platform, and internal variables have
separate ownership and lifecycle rules.
## Naming and implementation boundary
New public variables owned by vLLM-Omni must use the `VLLM_OMNI_` prefix.
Existing public names with older prefixes remain supported for compatibility,
but they do not establish naming precedent for new settings. Inherited vLLM
and third-party names retain the prefix chosen by their owning project.
The source [environment-variable inventory](gh-file:vllm_omni/config/environment_variable_inventory.py)
records ownership, migration disposition, and redaction metadata. Unlike
`vllm.envs`, it is not an executable value resolver: existing consumers still
own parsing, defaults, precedence, and evaluation time. The public tables below
describe those live consumer contracts. Stable model or stage behavior should
move to typed configuration, and request-varying behavior should move to a
request schema instead of adding another environment variable.
## Precedence and evaluation time
There is no repository-wide rule that makes environment variables override (or
lose to) every CLI, YAML, or request option. The tables below state precedence
for each public setting. In general:
- request fields should control behavior that can vary per request;
- typed deploy or stage configuration should control stable model and stage
behavior;
- environment variables are appropriate for process launch, platform
integration, emergency diagnostics, and compatibility fallbacks.
An environment variable can be read while a module is imported, when server or
model configuration is constructed, when a worker is spawned, or whenever a
property is accessed. Set import- and startup-time variables before invoking
`vllm serve` or importing vLLM-Omni. Changing the parent shell after workers
start does not update existing worker environments.
The lifecycle labels used below are:
- **Stable** — reviewed public configuration; changes require compatibility
consideration.
- **Experimental** — public for evaluation, but its contract may still evolve.
- **Diagnostic** — supported for investigation or temporary mitigation, not as
a normal tuning interface.
- **Deprecated** — compatibility alias; migrate to the named replacement.
## Public vLLM-Omni variables
### Diffusion backend and cache selection
| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
| --- | --- | --- | --- | --- |
| `DIFFUSION_ATTENTION_BACKEND` | Backend name or `auto`; default `auto` (platform selection) | Diffusion stages; read when `OmniDiffusionConfig` is constructed | `diffusion_attention_config.default` wins. An unknown backend fails during backend resolution. | Stable fallback |
| `DIFFUSION_CACHE_BACKEND` | `none`, `cache_dit`, `tea_cache`, `mag_cache`, `step_cache`, `stepcache`, or `step_cache_dit`; default `none` | Diffusion runner startup | Explicit `cache_backend` in config wins. Otherwise this name wins over the deprecated alias. Unsupported values raise `ValueError` during runner setup. | Stable fallback |
| `DIFFUSION_CACHE_ADAPTER` | Same values as `DIFFUSION_CACHE_BACKEND`; default `none` | Diffusion runner startup | Used only when neither explicit `cache_backend` nor `DIFFUSION_CACHE_BACKEND` is set. Unsupported values raise `ValueError`. | Deprecated; use `DIFFUSION_CACHE_BACKEND` |
| `OMNI_DIFFUSION_PROMPT_EMBED_CACHE` | Boolean: `1`, `true`, `yes`, `on`, `0`, `false`, `no`, or `off`; default disabled | Each diffusion runner; resolved during model setup | A recognized environment value overrides the explicit enable setting. An unrecognized value is ignored. | Experimental |
| `OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE` | Positive integer; default `32` entries | Each diffusion runner; resolved during model setup | A valid environment value overrides the explicit cache size. A non-integer logs a warning; a non-positive value is ignored. | Experimental |
| `OMNI_DIFFUSION_SESSION_STATE_MANAGER` | Boolean with the same accepted spellings as the prompt cache; default disabled | Experimental diffusion session manager; model setup | A recognized environment value overrides the explicit enable setting. An unrecognized value is ignored. | Experimental |
| `OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS` | Positive integer; default `64` | Experimental diffusion session manager; model setup | A valid environment value overrides the explicit maximum. A non-integer logs a warning; a non-positive value is ignored. | Experimental |
Backend names for `DIFFUSION_ATTENTION_BACKEND` are the members of
`DiffusionAttentionBackendEnum`, such as `FLASH_ATTN`, `TORCH_SDPA`,
`SAGE_ATTN`, `FLASHINFER_ATTN`, and `TRTLLM_ATTN`. Platform support still
depends on the installed kernels and model path.
### Serving and runtime
| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
| --- | --- | --- | --- | --- |
| `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable |
| `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
| `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
| `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic |
| `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
| `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental |
| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental |
### Server storage
Storage names use Pydantic's nested-settings delimiter (`__`). They are read
when `ServerSettings` is initialized, normally during server import/startup.
Invalid integer values raise a Pydantic validation error.
| Name | Type and default | Meaning | Lifecycle |
| --- | --- | --- | --- |
| `VLLM_OMNI_SERVER_STORAGE__PATH` | Path; default `/tmp/storage` | Directory for completed server files. | Stable |
| `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY` | Integer; default `4` | Maximum concurrent file operations. No additional range validation is currently applied. | Stable |
| `VLLM_OMNI_SERVER_STORAGE__FILE_TTL` | Integer seconds or unset; default unset | Optional lifetime for locally stored files. | Stable |
| `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL` | Integer seconds or unset | Sweep frequency. If file TTL is set and this is unset, the effective default is `300`. | Stable |
| `VLLM_OMNI_STORAGE_PATH` | Path | Deprecated alias for `VLLM_OMNI_SERVER_STORAGE__PATH`. The new name wins when both are set. | Deprecated |
| `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` | Integer | Deprecated alias for `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`. The new name wins when both are set. | Deprecated |
The deprecated aliases emit `DeprecationWarning`. They do not affect file TTL
settings.
### Quantization diagnostics and performance
| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
| --- | --- | --- | --- | --- |
| `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary |
| `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control |
`QUACK_CACHE_DIR` is owned by the external quack library and is not an
Omni-owned variable, even though vLLM-Omni supplies a persistent default when
it is unset.
## Per-stage environment
Deploy configurations can set arbitrary environment keys for one stage:
```yaml
stages:
- stage_id: 0
env:
HCCL_IF_BASE_PORT: 23000
MY_MODEL_CACHE: /mnt/model-cache
```
The legacy nested form under `runtime.env` is also accepted. vLLM-Omni converts
keys and values to strings, applies them while the stage engine or child process
is launched, logs keys but not values, and restores the parent process's prior
environment afterward. A child process inherits the scoped values. This keeps
sibling stages from permanently changing one another's parent environment.
This scoping has two consequences:
1. A setting read inside stage/model initialization can see the stage value.
2. A setting already cached by an import that occurred before stage launch
cannot be changed retroactively with stage `env`.
Stage `env` is an escape hatch for process and third-party settings, not a
replacement for typed stage configuration. Do not put request-varying behavior
there. Secret values may reach child processes; logs and diagnostics must emit
their keys only.
## Inherited vLLM variables
vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
the [vLLM 0.26 environment-variable reference](https://docs.vllm.ai/en/v0.26.0/configuration/env_vars.html)
for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
XPU, ModelScope, and FlashInfer workspace settings.
An inherited name does not become Omni-owned because Omni reads or forwards it.
Where behavior differs, the Omni-specific page should document only that
difference. For example, device-visibility values can be further narrowed by a
stage's `devices` setting.
## Platform, external, and secret variables
Hardware runtimes, distributed launchers, and third-party libraries retain
ownership of their variables. Common groups include:
- device visibility (`ASCEND_RT_VISIBLE_DEVICES`, `HIP_VISIBLE_DEVICES`, and
`MUSA_VISIBLE_DEVICES`);
- distributed launch (`MASTER_ADDR`, `MASTER_PORT`, `RANK`, and `WORLD_SIZE`);
- Hugging Face cache and authentication (`HF_HOME`, `HF_TOKEN`, and related
names);
- accelerator libraries such as NCCL, FlashInfer, MindIE, MORI, and quack.
Consult the owning project's documentation for accepted values and support.
`HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, and `OPENAI_API_KEY` are explicitly marked
for redaction in the reviewed inventory. Never print their values in environment
collection, examples, bug reports, or logs.
## Model-specific variables
The audit found 56 variables read by a single model or pipeline family. They are
not listed as public usage options here because doing so would turn implementation
escape hatches into an accidental compatibility contract.
Every audited model-specific name has a migration disposition in the
[environment-variable inventory](gh-file:vllm_omni/config/environment_variable_inventory.py):
| Disposition | Count | Required outcome |
| --- | ---: | --- |
| Promote | 33 | Move a stable setting into typed stage or model configuration. |
| Request scope | 6 | Move request-varying behavior into a declared request-option schema. |
| External | 0 | Retain only when a supported third-party library owns the contract. |
| Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. |
| Deprecate/remove | 6 | Remove a compatibility escape hatch that has no continuing contract. |
The disposition is a migration target, not a statement that the existing
environment switch is stable. Promote or request-scope work should land in
model-owner-reviewed follow-up changes with validation and model-page updates.
The family-sized migration packages and their ownership gate are tracked in
[#6232](https://github.com/vllm-project/vllm-omni/issues/6232).
## Environment collection and troubleshooting
Run the repository collector when preparing a bug report:
```bash
python collect_env.py
```
The environment section includes the safe public Omni allowlist, registered
vLLM variables, and selected platform prefixes. It omits names marked for
redaction and names containing common secret terms. Review the output before
sharing it. A key supplied through stage `env` is not collected merely because
it appears in deploy configuration; the same allowlist and prefix rules apply.
If a setting appears ineffective, first check its read time in the tables
above. Restart processes after changing import- or startup-time values, verify
that the variable reaches the intended worker, and check whether an explicit
typed configuration field takes precedence.
## Benchmark and internal variables
The reviewed inventory marks 20 benchmark variables as transitional while
benchmark configuration moves to CLI/request schemas in
[issue #5376](https://github.com/vllm-project/vllm-omni/issues/5376). They are not
part of the serving/runtime contract.
Internal process-coordination and debug variables are deliberately excluded
from this page. The CI scanner resolves string literals and module-level string
constants used by direct `os` access, local wrapper helpers, aliased `os`
imports, and environment membership checks. Server-storage names generated by
Pydantic settings are checked separately against their model fields.
Names assembled from cross-module metadata and arbitrary keys supplied through
stage `env` cannot be inferred reliably. Contributors must classify those names
explicitly. CI fails when a statically resolved name has no classification or
when a reviewed public name is absent from this page.
@@ -4,8 +4,8 @@ When you want to add L4-level ***performance test*** cases, add entries to JSON
| Model type | Runner | Example JSON files |
| ---------- | ------ | ------------------ |
| Omni (nightly) | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json` (`full_model` in `mark`) |
| Omni (weekly) | `run_benchmark.py` | `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` (`slow` in `mark`; **Perf Test** in `test-weekly.yml`) |
| Omni (nightly) | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json` (`full_model` without `slow` in `mark`) |
| Omni (weekly) | `run_benchmark.py` | `test_qwen3_omni_async_chunk.json` (CUDA only), `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` (`slow` in `mark`; **Perf Test** in `test-weekly.yml`) |
| TTS | `run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` |
| Diffusion | `run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`, `test_wan22_i2v_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, … |
@@ -20,7 +20,7 @@ Diffusion cases are detected when the JSON has `server_type` (typically `"vllm-o
#### Running perf cases
Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bulk scan above and filter with pytest **`-m`** on each case's JSON `mark` (for example `pytest … run_benchmark.py -m "full_model and tts and H100"` or `-m "slow and omni and H100"` for weekly Omni configs).
Pass **`--test-config-file`** to load one JSON file, or omit it for the bulk scan above. In either mode, pytest **`-m`** filters each case using its JSON `mark`. A single file can contain nightly, weekly, and platform-specific case objects. For example, the nightly object in `test_qwen3_omni_async_chunk.json` declares both H100 and A3, while its `slow` weekly object is CUDA-only; there is no NPU weekly counterpart. Use `-m "H100 and full_model and not slow"`, `-m "H100 and full_model and slow"`, or `-m "npu"` to select the intended schedule and platform.
```JSON
{
@@ -75,7 +75,7 @@ Omit `mark` only for configs not meant to be filtered by `-m`. `server_type` is
Optional top-level field on each perf JSON **case object** (one per `test_name`). `run_benchmark.py` and `run_diffusion_benchmark.py` read it via `tests.dfx.conftest.resolve_pytest_marks` and attach the marks to the corresponding `pytest.param`, so you can filter locally with `-m`.
When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `omni`, `tts`, `diffusion`, or `local_model`.
When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `slow`, `omni`, `tts`, `diffusion`, or `local_model`.
Supported form:
@@ -142,11 +142,18 @@ Examples:
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and H100 and diffusion"
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100"
# Single file (same as nightly CI Perf steps)
# Single file (same selectors as the CI Perf steps)
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
--test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \
-m "H100 and full_model and not slow"
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \
-m "H100 and full_model and slow"
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \
-m "npu"
```
<!-- markdownlint-disable-next-line MD051 -->
@@ -196,16 +196,15 @@ adopt it.
### Final-layout Host Weight Runtime consumer
> **Status:** PR2 implementation. The representation-independent identity,
> producer, and restoration contracts landed in
> [PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445); this consumer
> wires the final-layout BF16 path into eligible no-AllGather DLO.
> **Status:** PR2 landed in
> [PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). PR3 adds
> registered direct H2D plus generic HWR publication/source-digest
> optimizations without changing AllGather behavior.
Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
only to no-AllGather DLO. Enable it with
`--host-weight-runtime-mode preferred` (or `required`) and
`--host-weight-runtime-root <node-local-root>`. Registration of shared mappings
and direct asynchronous H2D remain a separate PR3 transport change.
`--host-weight-runtime-root <node-local-root>`.
The modes express operator fallback policy, not different artifact formats:
@@ -254,6 +253,44 @@ outcomes are authoritative:
Checkpoint mmap remains an unchanged control path whenever final-layout HWR is
not selected.
Local source identity and publication avoid repeated full-payload passes while
preserving the same semantic contract:
- Canonical files without a trusted immutable Hub blob identity are hashed in
parallel. Their content IDs are cached per storage domain under the HWR root
and reused only when the path, inode, size, timestamps, symlink target, and
cache-record checksum still match. Corrupt entries are rebuilt; cache-lock or
cache-I/O failure falls back to hashing the canonical file directly.
- Producers that emit payload bytes in canonical storage-key order may declare
`ordered=True`. The filesystem writer then computes file and tensor SHA256
values during the write and overlaps payload `fsync` with later producer
work. Producers without that guarantee retain unordered writes and use a
parallel readback checksum fallback.
- Manifest and `READY` publication still wait for every payload `fsync` and
checksum to complete. These optimizations do not weaken artifact validation
or change lease ownership.
#### Registered mmap transport
After DLO takes an HWR lease, the no-AllGather backend may register the lease's
complete immutable mapped ranges under the existing `pin_cpu_memory` policy.
Registration is transport state: it does not change artifact identity, the
store, tensor ownership, or H2D payload. On success, each tensor view copies
directly into the existing rotating HBM block buffers and the two private host
staging slots are not allocated.
`--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling
over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned
memory policy, unsupported platform/capability, over-budget mapping, or fully
rolled-back registration error selects the existing two-slot staging path. A
partial registration that cannot be rolled back aborts startup because closing
the lease would unmap memory still owned by the platform.
Direct checkpoint mmap remains unchanged and continues to use staging. It may
require loader-owned per-block transforms, while the HWR artifact already
contains final runtime bytes. DLO AllGather never receives an HWR final-layout
lease and therefore never enters this registration path.
#### Pre-service transaction boundary
The startup transaction does not end at restore commit. A lease-backed model is
@@ -317,6 +354,7 @@ loader owns
-> backend takes once before asynchronous work
-> backend drains pending H2D work
-> backend releases hooks, staging, and model references
-> backend unregisters every mapped range
-> backend closes lease
```
@@ -326,13 +364,15 @@ it. Once the backend takes ownership, only backend abort or teardown may drain
work and close it.
The implementation keeps the final-layout lease through backend setup and
initial prefetch. The backend uses the existing two bounded rank-local staging
slots for the immutable HWR tensors, and its transactional `enable()` cleanup
removes partial hooks and closes the lease before reporting a setup failure.
Preferred mode then uses the runner's fresh canonical retry; required mode
propagates the failure.
initial prefetch. Successful registration copies directly from immutable HWR
views; otherwise the backend uses the existing two bounded rank-local staging
slots. Transactional cleanup drains device work, removes source references,
unregisters mappings, and only then closes the lease. An unregistration failure
retains both registration and lease for retry/process teardown. Preferred mode
then uses the runner's fresh canonical retry; required mode propagates the
failure.
#### PR2 promotion gates
#### PR2 and PR3 promotion gates
- Warm hit performs zero ordinary DiT materialization and zero producer calls.
- Shared warm finalization changes neither restored bytes nor backing pointers.
@@ -344,6 +384,10 @@ propagates the failure.
- Disabled mode and AllGather emit zero final-layout HWR interaction.
- The lease carrier rejects duplicate take and serialization.
- Backend setup or prefetch failure drains asynchronous work before lease close.
- Registration is all-or-nothing; rollback/unregistration failure never closes
a mapping still owned by the platform.
- Registered HWR transport bypasses host staging while preserving output and
H2D payload; unsupported registration retains the bounded staging path.
- Compatible checkpoint mmap remains an unchanged benchmark and control path.
- Prewarm uses one matching producer cohort per TP/SP topology and storage
domain.
@@ -394,11 +438,12 @@ reconstructs those tensors with their recorded layouts. Other online methods
must use `--dlo-no-use-allgather` or disable online quantization until their
runtime layouts are validated.
The Host Weight Runtime representation and publication contracts are merged;
see [RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414) and
[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445). The planned
no-AllGather consumer above is the PR2 implementation; registration and direct
asynchronous H2D remain a separate transport follow-up.
The Host Weight Runtime representation, publication, and no-AllGather consumer
contracts are merged; see
[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414),
[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445), and
[PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). Registered
direct H2D is an optional transport layer over that merged lease contract.
## Validation coverage
@@ -411,8 +456,12 @@ Current source-level validation includes:
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
sharding of finalized weights and scales;
- exact loader-to-backend plan transfer and ordinary-loader fallback;
- ordered publication hashing, overlapped durability, unordered parallel
checksum fallback, and node-local source-digest reuse/invalidation;
- rank-local mmap source retention, bounded two-slot staging, and adapter
transforms without parameter-side flags;
- read-only registration preflight, direct source-to-device copies, safe
fallback, partial rollback, retryable unregistration, and lease ordering;
- resident-layer requests requiring no-AllGather;
- DP request-wave validation for denoising-step compatibility;
- bounded component allocator-cache retention, conservative/forced release,
+14 -1
View File
@@ -16,7 +16,7 @@ validation_paths:
- tests/host_weight_runtime/**
upstream_refs:
- https://github.com/vllm-project/vllm-omni/issues/6414
last_reviewed: 2026-08-21
last_reviewed: 2026-08-25
---
# Host Weight Runtime
@@ -84,6 +84,13 @@ another. `WeightArtifactIdentity` includes:
- producer version, implementation fingerprint, manifest schema, and restorer
schema.
Loader integrations may persist content digests for canonical source shards in
the selected storage domain. Reuse requires the same path, inode, size,
timestamps, symlink target, and cache-record checksum. Corrupt records are
rebuilt, and coordination or cache-I/O failure falls back to hashing the
canonical source directly. This is an identity-computation optimization, not a
second source of truth or an artifact substitution rule.
The identity excludes transfer policy such as registered mmap versus private
pinned staging. Those paths move the same host representation and must not
produce duplicate artifacts.
@@ -237,6 +244,12 @@ Every publication computes full file and tensor hashes. Local lookup supports:
and
- `full_checksum`, which additionally reads and hashes every payload file.
An ordered producer may compute file and tensor digests while payload bytes are
written and overlap each closed payload's `fsync` with later producer work. An
unordered producer retains parallel readback hashing. In both cases manifest
and `READY` publication wait for every payload digest and `fsync`; the atomic
publication and validation contract is unchanged.
Filesystem verity is modeled but unsupported. Invalid artifacts receive an
external deny marker before their shared lookup lock is released; a build owner
then quarantines the immutable directory under the exclusive lock.
@@ -65,6 +65,7 @@ omni = Omni(
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
| `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
| `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
## Host-weight loading
@@ -178,12 +179,48 @@ returning to `required`.
On a cold start, the canonical loader remains authoritative and publishes a
validated final-layout artifact for later workers. A warm start restores the
DiT final tensors without ordinary DiT materialization, then DLO streams them
through the same two bounded host staging slots. The artifact identity includes
the TP rank/size and SP layout, so TP1, TP2 rank-local shards, and distinct SP
layouts do not alias one another. Publication failure remains separate from a
`preferred` serving startup; the next `required` startup provides the explicit
artifact-availability check.
DiT final tensors without ordinary DiT materialization. DLO then attempts to
register the immutable mapping for direct H2D; if registration is unavailable,
it streams through the same two bounded host staging slots. The artifact
identity includes the TP rank/size and SP layout, so TP1, TP2 rank-local shards,
and distinct SP layouts do not alias one another. Publication failure remains
separate from a `preferred` serving startup; the next `required` startup
provides the explicit artifact-availability check.
For local canonical checkpoints, the first eligible worker may hash source
shards to establish immutable identity. HWR caches those digests in the same
node-local storage domain and validates file metadata before reuse, so later
workers normally avoid repeating that read. Cold BF16 publication also hashes
ordered payloads as they are written and overlaps payload durability work with
later shards; this changes startup work only, not artifact contents or runtime
H2D behavior.
#### Registered direct H2D
Registration is attempted automatically only for an eligible warm HWR hit when
`pin_cpu_memory` is enabled. A successful path is:
```text
shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel
```
It removes the recurrent mmap-to-private-staging CPU copy and does not allocate
the two host staging slots. It does not reduce the H2D payload or make GPU
kernels access host memory directly.
CUDA requires read-only host-registration capability for the immutable HWR
mapping. Unsupported capability, a positive registration limit smaller than
the complete page-aligned mapping, or a safely rolled-back registration error
falls back to bounded staging. Programmatic configurations can set
`pin_cpu_memory=False` to disable registration explicitly. A successful
registration locks the complete mapped range in host memory for that worker's
lifetime, so use
`--dlo-host-registration-limit-gib` when an operator-enforced ceiling is
required.
Shutdown drains H2D work, releases hook/source references, unregisters every
range, and only then closes the HWR lease. AllGather and direct checkpoint mmap
retain their existing transfer paths.
When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
not resolve HWR sources or construct its store.
@@ -1,193 +0,0 @@
# Text-To-Audio
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/stable_audio>.
This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
## Supported Models
| Model | Description |
|-------|-------------|
| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
## Start Server
### Basic Start
```bash
vllm-omni serve stabilityai/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8091 \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni
```
## API Calls
### Method 1: Using curl
```bash
# Run all curl examples
bash curl_examples.sh
# Or execute directly
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of a cat purring",
"audio_length": 10.0
}' --output cat.wav
```
### Method 2: Using Python Client
```bash
cd examples/online_serving/stable_audio
# Simple generation
python stable_audio_client.py \
--text "The sound of a cat purring"
# With custom duration
python stable_audio_client.py \
--text "A dog barking" \
--audio_length 5.0
# With all parameters
python stable_audio_client.py \
--text "Thunder and rain" \
--audio_length 15.0 \
--negative_prompt "Low quality" \
--guidance_scale 7.0 \
--num_inference_steps 100 \
--seed 42 \
--output thunder.wav
```
The Python client supports the following command-line arguments:
- `--api_url`: API endpoint URL (default: `http://localhost:8091/v1/audio/generate`)
- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
- `--audio_start`: Audio start time in seconds (default: `0.0`)
- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
- `--num_inference_steps`: Number of inference steps (default: `100`)
- `--seed`: Random seed for reproducibility (default: `None`)
- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
- `--output`: Output file path (default: `stable_audio_output.wav`)
### Method 3: Using Python httpx
```python
import httpx
response = httpx.post(
"http://localhost:8091/v1/audio/generate",
json={
"input": "The sound of ocean waves crashing on a beach",
"audio_length": 10.0,
"negative_prompt": "Low quality, distorted",
"guidance_scale": 7.0,
"num_inference_steps": 100,
},
timeout=300.0,
)
with open("ocean.wav", "wb") as f:
f.write(response.content)
```
## Request Format
### Simple Generation
```json
{
"input": "The sound of ocean waves"
}
```
### Generation with Parameters
```json
{
"input": "A piano playing a gentle melody",
"audio_length": 10.0,
"negative_prompt": "Low quality, distorted, noisy",
"guidance_scale": 8.0,
"num_inference_steps": 150,
"seed": 42,
"response_format": "wav"
}
```
## API Reference
### Endpoint
```
POST /v1/audio/generate
Content-Type: application/json
```
### Generation Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text prompt describing the audio to generate |
| `model` | string | server's model | Model to use (optional, should match server if specified) |
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
| `audio_start` | float | 0.0 | Audio start time in seconds |
| `negative_prompt` | string | null | Text describing what to avoid in generation |
| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
| `seed` | int | null | Random seed for reproducible generation |
### Response Format
Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
## Tuning Tips
1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
2. **Quality vs Speed**:
- 50 steps: Fast, decent quality (quick previews)
- 100 steps: Good balance (general purpose)
- 150+ steps: High quality, slower (final / critical audio)
3. **Guidance Scale**:
- Lower (3 - 5): More creative / varied output
- Default (7): Good balance
- Higher (10+): Strict adherence to the prompt
4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
## File Description
| File | Description |
|------|-------------|
| `curl_examples.sh` | Curl examples covering common use cases |
| `stable_audio_client.py` | Python client with full CLI argument support |
## Troubleshooting
1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
2. **Connection refused**: Make sure the server is running on the correct port.
3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
## Example materials
??? abstract "stable_audio_client.py"
``````py
--8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
``````
??? abstract "curl_examples.sh"
``````sh
--8<-- "examples/online_serving/stable_audio/curl_examples.sh"
``````
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Minimal single-input MiniCPM-o 4.5 Realtime duplex demo.
Run this after starting the duplex server. Strict lifecycle, overlap, and
@@ -19,8 +22,8 @@ if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
PCM16_BYTES_PER_SAMPLE,
PCM16_SAMPLE_RATE,
PCM16_BYTES_PER_SAMPLE, # noqa: F401 - compatibility export
PCM16_SAMPLE_RATE, # noqa: F401 - compatibility export
RealtimeDuplexClient,
RealtimeEventCollector,
build_realtime_url,
@@ -28,6 +31,13 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
wait_for,
write_pcm16_wav,
)
from vllm_omni.experimental.fullduplex.client import chunk_period_ms as _chunk_period_ms # noqa: E402
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
has_residual_model_unit as _has_residual_model_unit,
)
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
reference_audio_data_url as _ref_audio_data_url,
)
class _StreamingOutputWriter:
@@ -133,25 +143,6 @@ def _latest_model_decision(
return decision
def _chunk_period_ms(events: list[dict[str, object]]) -> int:
for event in reversed(events):
session = event.get("session")
if not isinstance(session, dict):
continue
capabilities = session.get("capabilities")
if not isinstance(capabilities, dict):
continue
chunk_period_ms = capabilities.get("chunk_period_ms")
if isinstance(chunk_period_ms, int) and chunk_period_ms > 0:
return chunk_period_ms
return 1000
def _has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool:
unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000
return bool(unit_bytes > 0 and len(pcm16) % unit_bytes)
def _response_in_progress(events: list[dict[str, object]]) -> bool:
return sum(event.get("type") == "response.created" for event in events) > sum(
event.get("type") == "response.done" for event in events
@@ -168,13 +159,6 @@ def _event_count_after(
return sum(event.get("type") == event_type for event in events[index + 1 :])
def _ref_audio_data_url(path: str | None) -> str | None:
if path is None:
return None
ref_path = Path(path).expanduser()
return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii")
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
input_pcm16 = read_pcm16_wav(Path(args.input_wav))
if not input_pcm16:
@@ -1,234 +0,0 @@
# Stable Audio Online Serving
Generate audio from text prompts using Stable Audio models via an OpenAI-compatible API endpoint.
## Features
- **OpenAI-compatible API**: Use `/v1/audio/generate` endpoint
- **Flexible control**: Adjust audio length, guidance scale, inference steps
- **Quality control**: Use negative prompts to avoid unwanted characteristics
- **Reproducible**: Set random seed for deterministic generation
## Quick Start
### 1. Start the Server
```bash
vllm-omni serve stabilityai/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8091 \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni
```
### 2. Generate Audio
#### Using curl
```bash
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of a cat purring",
"audio_length": 10.0
}' --output cat.wav
```
#### Using Python Client
```bash
python stable_audio_client.py \
--text "The sound of a cat purring" \
--audio_length 10.0 \
--output cat.wav
```
#### Using Bash Script
```bash
bash curl_examples.sh
```
## API Reference
### Endpoint
```
POST /v1/audio/generate
```
### Request Body
```json
{
"input": "Text description of the audio",
"audio_length": 10.0,
"audio_start": 0.0,
"negative_prompt": "Low quality",
"guidance_scale": 7.0,
"num_inference_steps": 100,
"seed": 42,
"response_format": "wav"
}
```
### Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text prompt describing the audio to generate |
| `audio_length` | float | ~47s | Audio duration in seconds (max ~47s for stable-audio-open-1.0) |
| `audio_start` | float | 0.0 | Audio start time in seconds |
| `negative_prompt` | string | null | Text describing what to avoid in generation |
| `guidance_scale` | float | 7.0 | Classifier-free guidance scale (higher = more adherence to prompt) |
| `num_inference_steps` | int | 50 | Number of denoising steps (higher = better quality, slower) |
| `seed` | int | null | Random seed for reproducibility |
| `response_format` | string | "wav" | Output format: wav, mp3, flac, pcm |
### Response
Returns audio data in the requested format (default: WAV).
## Usage Examples
### Basic Generation
```bash
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of ocean waves"
}' --output ocean.wav
```
### Custom Duration
```bash
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "A dog barking",
"audio_length": 5.0
}' --output dog_5s.wav
```
### High Quality with Negative Prompt
```bash
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "A piano playing a gentle melody",
"audio_length": 10.0,
"negative_prompt": "Low quality, distorted, noisy",
"guidance_scale": 8.0,
"num_inference_steps": 150
}' --output piano_hq.wav
```
### Reproducible Generation
```bash
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Thunder and rain sounds",
"audio_length": 15.0,
"seed": 42
}' --output thunder.wav
```
### Quick Generation (Fewer Steps)
For faster generation with slightly lower quality:
```bash
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Birds chirping in a forest",
"audio_length": 8.0,
"num_inference_steps": 50
}' --output birds_quick.wav
```
## Python Client Examples
### Simple Generation
```bash
python stable_audio_client.py \
--text "The sound of a cat purring"
```
### Custom Parameters
```bash
python stable_audio_client.py \
--text "Thunder and rain" \
--audio_length 15.0 \
--negative_prompt "Low quality" \
--guidance_scale 7.0 \
--num_inference_steps 100 \
--seed 42 \
--output thunder.wav
```
### Different Output Format
```bash
python stable_audio_client.py \
--text "Guitar playing" \
--response_format mp3 \
--output guitar.mp3
```
## Tips
1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`
2. **Quality vs Speed**:
- 50 steps: Fast, decent quality
- 100 steps: Good balance (default)
- 150+ steps: High quality, slower
3. **Guidance Scale**:
- Lower (3-5): More creative/varied
- Default (7): Good balance
- Higher (10+): More literal to prompt
4. **Negative Prompts**: Use to avoid "Low quality", "distorted", "noisy", etc.
5. **Seeds**: Use same seed for reproducible results
## Performance
| Inference Steps | Quality | Speed | Use Case |
|----------------|---------|-------|----------|
| 50 | Good | Fast | Quick previews |
| 100 (default) | Very Good | Medium | Production |
| 150+ | Excellent | Slow | Final/critical audio |
## Troubleshooting
### Server not responding
- Check if server is running: `curl http://localhost:8091/health`
- Check server logs for errors
### Audio quality issues
- Increase `num_inference_steps` (e.g., 150)
- Add negative prompts: `"Low quality, distorted, noisy"`
- Increase `guidance_scale` for more prompt adherence
### Generation timeout
- Reduce `num_inference_steps`
- Reduce `audio_length`
- Check GPU memory with `nvidia-smi`
### Wrong audio length
- Ensure `audio_length` is within model limits (~47s max)
- Adjust `audio_start` if trimming is needed
## See Also
- [Offline Inference Example](../../offline_inference/text_to_audio/README.md)
- [Stable Audio Model Card](https://huggingface.co/stabilityai/stable-audio-open-1.0)
- [vLLM-Omni Documentation](https://github.com/vllm-project/vllm-omni)
@@ -1,54 +0,0 @@
#!/bin/bash
# Examples for using Stable Audio with curl via /v1/audio/generate endpoint
# Example 1: Simple request with default parameters
echo "Example 1: Simple request with default parameters"
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound audience clapping and cheering in a stadium"
}' --output stadium.wav
# Example 2: Request with custom audio_length
echo "Example 2: Custom audio length (5 seconds)"
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "The sound of a dog barking",
"audio_length": 5.0
}' --output dog_5s.wav
# Example 3: Request with negative prompt for quality control
echo "Example 3: With negative prompt"
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "A piano playing a gentle melody",
"audio_length": 10.0,
"negative_prompt": "Low quality, distorted, noisy"
}' --output piano.wav
# Example 4: Full control with all parameters
echo "Example 4: Full control (custom length, guidance, steps, seed)"
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Thunder and rain sounds",
"audio_length": 15.0,
"negative_prompt": "Low quality",
"guidance_scale": 7.0,
"num_inference_steps": 100,
"seed": 42
}' --output thunder_rain.wav
# Example 5: Quick generation with fewer steps (faster but lower quality)
echo "Example 5: Quick generation (fewer steps)"
curl -X POST http://localhost:8091/v1/audio/generate \
-H "Content-Type: application/json" \
-d '{
"input": "Ocean waves crashing on a beach",
"audio_length": 8.0,
"num_inference_steps": 50
}' --output ocean.wav
echo "All examples completed!"
@@ -1,170 +0,0 @@
#!/usr/bin/env python3
"""
OpenAI-compatible client for Stable Audio via /v1/audio/generate endpoint.
This script demonstrates how to use the OpenAI-compatible speech API
to generate audio from text using Stable Audio models.
Examples:
# Simple generation
python stable_audio_client.py --text "The sound of a cat purring"
# With custom duration
python stable_audio_client.py --text "A dog barking" --audio_length 5.0
# With all parameters
python stable_audio_client.py --text "Thunder and rain" \
--audio_length 15.0 \
--negative_prompt "Low quality" \
--guidance_scale 7.0 \
--num_inference_steps 100 \
--seed 42 \
--output thunder.wav
"""
import argparse
import sys
import requests
def parse_args():
parser = argparse.ArgumentParser(description="Generate audio with Stable Audio via OpenAI-compatible API")
parser.add_argument(
"--api_url",
default="http://localhost:8091/v1/audio/generate",
help="API endpoint URL",
)
parser.add_argument(
"--text",
default="The sound of a cat purring",
help="Text prompt for audio generation",
)
parser.add_argument(
"--audio_length",
type=float,
default=10.0,
help="Audio length in seconds (max ~47s for stable-audio-open-1.0)",
)
parser.add_argument(
"--audio_start",
type=float,
default=0.0,
help="Audio start time in seconds",
)
parser.add_argument(
"--negative_prompt",
default="Low quality",
help="Negative prompt for classifier-free guidance",
)
parser.add_argument(
"--guidance_scale",
type=float,
default=7.0,
help="Guidance scale for diffusion (higher = more adherence to prompt)",
)
parser.add_argument(
"--num_inference_steps",
type=int,
default=100,
help="Number of inference steps (higher = better quality, slower)",
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="Random seed for reproducibility",
)
parser.add_argument(
"--output",
default="stable_audio_output.wav",
help="Output file path",
)
parser.add_argument(
"--response_format",
default="wav",
choices=["wav", "mp3", "flac", "pcm"],
help="Audio output format",
)
return parser.parse_args()
def generate_audio(args):
"""Generate audio using the API."""
# Build request payload
payload = {
"input": args.text,
"audio_length": args.audio_length,
"audio_start": args.audio_start,
"response_format": args.response_format,
}
# Add optional parameters
if args.negative_prompt:
payload["negative_prompt"] = args.negative_prompt
if args.guidance_scale:
payload["guidance_scale"] = args.guidance_scale
if args.num_inference_steps:
payload["num_inference_steps"] = args.num_inference_steps
if args.seed is not None:
payload["seed"] = args.seed
print(f"\n{'=' * 60}")
print("Stable Audio - Text-to-Audio Generation")
print(f"{'=' * 60}")
print(f"API URL: {args.api_url}")
print(f"Prompt: {args.text}")
print(f"Audio length: {args.audio_length}s")
print(f"Negative prompt: {args.negative_prompt}")
print(f"Guidance scale: {args.guidance_scale}")
print(f"Inference steps: {args.num_inference_steps}")
if args.seed is not None:
print(f"Seed: {args.seed}")
print(f"Output: {args.output}")
print(f"{'=' * 60}\n")
try:
# Make the API request
print("Generating audio...")
response = requests.post(
args.api_url,
json=payload,
headers={"Content-Type": "application/json"},
timeout=300, # 5 minute timeout for long generations
)
# Check for errors
if response.status_code != 200:
print(f"Error: API returned status code {response.status_code}")
print(f"Response: {response.text}")
return False
# Save the audio
with open(args.output, "wb") as f:
f.write(response.content)
print(f"✓ Audio saved to {args.output}")
print(f" File size: {len(response.content) / 1024:.1f} KB")
return True
except requests.exceptions.Timeout:
print("Error: Request timed out. Try reducing inference steps or audio length.")
return False
except requests.exceptions.ConnectionError:
print(f"Error: Could not connect to {args.api_url}")
print("Make sure the server is running.")
return False
except Exception as e:
print(f"Error: {e}")
return False
def main():
args = parse_args()
success = generate_audio(args)
sys.exit(0 if success else 1)
if __name__ == "__main__":
main()
+3 -5
View File
@@ -23,8 +23,6 @@ offline validation sample with TeaCache and online serving through the
- Model: <https://huggingface.co/stabilityai/stable-audio-open-1.0>
- Offline example:
[`examples/offline_inference/text_to_audio`](../../examples/offline_inference/text_to_audio)
- Online example:
[`examples/online_serving/stable_audio`](../../examples/online_serving/stable_audio)
- Audio generation API:
[`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md)
- Related issue:
@@ -103,7 +101,7 @@ curl -X POST http://localhost:8091/v1/audio/generate \
"seed": 42,
"response_format": "wav"
}' \
--output examples/online_serving/stable_audio/piano_10s.wav
--output piano_10s.wav
```
##### Verification
@@ -137,12 +135,12 @@ PY
Validate the online output:
```bash
ls -lh examples/online_serving/stable_audio/piano_10s.wav
ls -lh piano_10s.wav
python - <<'PY'
import soundfile as sf
path = "examples/online_serving/stable_audio/piano_10s.wav"
path = "piano_10s.wav"
audio, sample_rate = sf.read(path)
print("sample_rate:", sample_rate)
print("shape:", audio.shape)
+179 -1
View File
@@ -1,10 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for metrics.py
"""
import math
import pytest
from vllm.benchmarks.serve import TaskType
@@ -66,6 +68,22 @@ def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutp
return output
def _calculate_test_metrics(outputs, goodput=None):
return calculate_metrics(
input_requests=[],
outputs=outputs,
dur_s=1.0,
tokenizer=None,
selected_percentiles=[50.0],
goodput_config_dict=goodput or {},
task_type=TaskType.GENERATION,
selected_percentile_metrics=[],
max_concurrency=None,
request_rate=float("inf"),
benchmark_duration=1.0,
)[0]
# ============================================================================
# total_input Tests
# ============================================================================
@@ -126,6 +144,166 @@ def test_audio_continuity_aggregation():
assert p99 is not None and p99 > 0.4
def test_unmeasured_duplex_latency_does_not_add_zero_samples():
measured = _make_output(100, output_tokens=2)
measured.ttft = 0.1
measured.audio_ttfp = 0.2
measured.audio_rtf = 0.5
measured.duplex_session_metrics = {
"mean_ttft_ms": 100.0,
"mean_ttfp_ms": 200.0,
"mean_rtf": 0.5,
}
listen_only = _make_output(100, output_tokens=0)
listen_only.ttft = listen_only.audio_ttfp = listen_only.audio_rtf = 0.0
listen_only.duplex_session_metrics = {
"mean_ttft_ms": None,
"mean_ttfp_ms": None,
"mean_rtf": None,
}
metrics = _calculate_test_metrics([measured, listen_only], {"ttft": 150.0})
assert metrics.mean_ttft_ms == 100.0
assert metrics.mean_audio_ttfp_ms == 200.0
assert metrics.mean_audio_rtf == 0.5
assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (1, 1, 1)
assert metrics.request_goodput == 1.0
def test_all_unmeasured_duplex_latency_is_not_reported_as_zero():
listen_only = _make_output(100, output_tokens=0)
listen_only.duplex_session_metrics = {
"mean_ttft_ms": None,
"mean_ttfp_ms": None,
"mean_rtf": None,
}
metrics = _calculate_test_metrics(
[listen_only],
{"ttft": float("inf"), "audio_ttft": float("inf")},
)
assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (0, 0, 0)
assert math.isnan(metrics.mean_ttft_ms)
assert math.isnan(metrics.mean_audio_ttfp_ms)
assert math.isnan(metrics.mean_audio_rtf)
assert metrics.request_goodput == 0.0
def test_unmeasured_duplex_tpot_does_not_add_zero_or_misalign_goodput():
missing_tpot = _make_output(100, output_tokens=5)
missing_tpot.itl = []
missing_tpot.text_latency = missing_tpot.ttft = 0.1
missing_tpot.tpot_measured = False
missing_tpot.duplex_session_metrics = {"mean_ttft_ms": 100.0}
slow_ttft = _make_output(100, output_tokens=5)
slow_ttft.ttft = 1.0
slow_ttft.text_latency = 1.4
slow_ttft.itl = [0.1] * 4
slow_ttft.duplex_session_metrics = {"mean_ttft_ms": 1000.0}
metrics = _calculate_test_metrics(
[missing_tpot, slow_ttft],
{"ttft": 500.0, "tpot": 200.0},
)
assert metrics.num_tpot_samples == 1
assert metrics.mean_tpot_ms == pytest.approx(100.0)
assert metrics.request_goodput == 0.0
def test_all_unmeasured_duplex_token_timing_is_not_reported_as_zero():
output = _make_output(100, output_tokens=5)
output.itl = []
output.text_latency = output.ttft
output.tpot_measured = False
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
metrics = _calculate_test_metrics([output], {"tpot": float("inf")})
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 0)
assert math.isnan(metrics.mean_tpot_ms)
assert math.isnan(metrics.mean_itl_ms)
assert metrics.request_goodput == 0.0
def test_duplex_response_timings_do_not_build_a_session_token_timeline():
output = _make_output(100, output_tokens=5)
output.latency = 101.0
output.duplex_request_metrics = [
{"response_id": "r1", "stage0_tokens": {"itls_ms": [100.0, 100.0]}},
{"response_id": "r2", "stage0_tokens": {"itls_ms": [100.0, 100.0]}},
]
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
metrics = _calculate_test_metrics([output])
assert math.isnan(metrics.max_output_tokens_per_s)
assert metrics.max_concurrent_requests == 1
assert metrics.mean_tpot_ms == pytest.approx(100.0)
assert metrics.mean_itl_ms == pytest.approx(100.0)
def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback():
output = _make_output(100, output_tokens=0)
output.generated_text = "timing metadata missing"
output.itl = []
output.text_latency = output.ttft
output.tpot_measured = False
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
def tokenizer(text, *, add_special_tokens):
assert text == output.generated_text
assert add_special_tokens is False
return type("Tokenized", (), {"input_ids": [1, 2, 3]})()
metrics, _ = calculate_metrics(
input_requests=[],
outputs=[output],
dur_s=1.0,
tokenizer=tokenizer,
selected_percentiles=[50.0],
goodput_config_dict={"tpot": float("inf")},
task_type=TaskType.GENERATION,
selected_percentile_metrics=[],
max_concurrency=None,
request_rate=float("inf"),
benchmark_duration=1.0,
)
assert metrics.total_output == 3
assert metrics.num_tpot_samples == 0
assert math.isnan(metrics.mean_tpot_ms)
assert metrics.request_goodput == 0.0
def test_measured_zero_itl_is_not_treated_as_missing():
output = _make_output(100, output_tokens=3)
output.itl = [0.0, 0.0]
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
metrics = _calculate_test_metrics([output], {"tpot": 1.0})
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (1, 2)
assert metrics.mean_tpot_ms == 0.0
assert metrics.mean_itl_ms == 0.0
assert metrics.request_goodput == 1.0
def test_duplex_goodput_does_not_pair_measurements_from_different_requests():
text_only, audio_only = _make_output(100), _make_output(100)
text_only.ttft, text_only.audio_ttfp = 0.1, 0.0
text_only.duplex_session_metrics = {"mean_ttft_ms": 100.0, "mean_ttfp_ms": None, "mean_rtf": None}
audio_only.ttft, audio_only.audio_ttfp = 0.0, 0.2
audio_only.duplex_session_metrics = {"mean_ttft_ms": None, "mean_ttfp_ms": 200.0, "mean_rtf": None}
metrics = _calculate_test_metrics([text_only, audio_only], {"ttft": 500.0, "audio_ttft": 500.0})
assert metrics.request_goodput == 0.0
# ============================================================================
# TTFT suppression for pure-audio (TTS) benchmarks
# ============================================================================
+778
View File
@@ -0,0 +1,778 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import argparse
import asyncio
import base64
import io
import json
import subprocess
import tarfile
import threading
import time
import wave
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import pytest
from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
from vllm.benchmarks.serve import TaskType
from vllm_omni.benchmarks import omniinteract as oi
from vllm_omni.benchmarks import serve as benchmark_serve
from vllm_omni.benchmarks.data_modules import omniinteract_dataset as data
from vllm_omni.benchmarks.metrics.metrics import calculate_metrics
from vllm_omni.benchmarks.patch import patch as benchmark_patch
from vllm_omni.entrypoints.cli.benchmark.cli_args import preprocess_serve_args
from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand
from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
MODEL = "openbmb/MiniCPM-o-4_5"
def _case(tmp_path: Path, *, subset: str = "1q1a", name: str = "video.mp4") -> data.OmniInteractCase:
video, annotation = tmp_path / name, tmp_path / f"{Path(name).stem}.json"
video.parent.mkdir(parents=True, exist_ok=True)
video.touch()
annotation.write_text("{}")
return data.OmniInteractCase(subset, name, video, annotation, "multi_turn")
def _collector(*events: tuple[dict[str, object], float]) -> RealtimeEventCollector:
collector = RealtimeEventCollector()
for event, received_at in events:
collector.add(event, received_at_s=received_at)
return collector
def _created(response_id: str = "r1") -> dict[str, object]:
return {"type": "response.created", "response": {"id": response_id}}
def _done(response_id: str = "r1", status: str = "completed") -> dict[str, object]:
return {"type": "response.done", "response": {"id": response_id, "status": status}}
def _audio(
response_id: str = "r1",
*,
samples: int = 2400,
rate: int | None = 24_000,
value: int = 1,
) -> dict[str, object]:
event: dict[str, object] = {
"type": "response.audio.delta",
"response_id": response_id,
"format": "pcm16",
"delta": base64.b64encode(bytes((value, 0)) * samples).decode(),
}
if rate is not None:
event["sample_rate_hz"] = rate
return event
def _text(response_id: str = "r1", value: str = "hello") -> dict[str, object]:
return {"type": "response.audio_transcript.delta", "response_id": response_id, "delta": value}
def _listen(*, buffering: bool = False) -> dict[str, object]:
return {
"type": "response.listen",
"response": {"metadata": {"model_listen": True, "buffering": buffering}},
}
def _session_options(tmp_path: Path) -> data.OmniInteractSessionOptions:
return data.OmniInteractSessionOptions(tmp_path / "artifacts", 10.0, 5.0, "reference.wav", True)
def _sample(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions):
return data.OmniInteractSampleRequest(
prompt="",
prompt_len=0,
expected_output_len=0,
multi_modal_data=None,
request_id="measured-0",
omniinteract_case=case,
omniinteract_options=options,
)
def _write_dataset(root: Path) -> Path:
data_root = root / "data"
for subset in ("1q1a", "1q1a_math"):
subset_root = data_root / subset
(subset_root / "videos").mkdir(parents=True)
(subset_root / "annotations").mkdir()
(subset_root / "videos" / f"{subset}.mp4").touch()
(subset_root / "annotations" / f"{subset}.json").write_text("{}")
(subset_root / "video_json_map.json").write_text(
json.dumps(
{
"entries": [
{
"video": f"videos/{subset}.mp4",
"annotation": f"annotations/{subset}.json",
"scene_type": "multi_turn",
}
]
}
)
)
one_to_many = data_root / "1qna"
(one_to_many / "videos_bench" / "nested").mkdir(parents=True)
(one_to_many / "annotations" / "nested").mkdir(parents=True)
(one_to_many / "videos_bench" / "nested" / "guide.mp4").touch()
(one_to_many / "annotations" / "nested" / "guide.json").write_text("{}")
return data_root
def _successful_output(
tmp_path: Path,
*,
audio_time: float = 10.0,
status: str = "completed",
) -> tuple[data.OmniInteractCase, RealtimeEventCollector, oi.OmniInteractCaseResult]:
case = _case(tmp_path)
collector = _collector(
(_created(), 9.9),
(_audio(), audio_time),
(_text(), audio_time + 0.01),
(_done(status=status), audio_time + 0.1),
)
result = oi.OmniInteractCaseResult(case.subset, str(case.video_path), "")
return case, collector, result
def _write_success(
root: Path,
case: data.OmniInteractCase,
collector: RealtimeEventCollector,
result: oi.OmniInteractCaseResult,
**kwargs,
):
require_response = kwargs.pop("require_response", True)
return oi.prepare_success_artifacts(
root,
case,
collector,
stream_start=10.0,
video_duration_s=1.0,
require_response=require_response,
result=result,
**kwargs,
)
def test_dataset_discovers_official_layouts_and_total_selection(tmp_path: Path):
root = _write_dataset(tmp_path)
all_cases = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=0, disable_shuffle=True)
selected = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=2, disable_shuffle=True)
assert [case.subset for case in all_cases] == ["1q1a", "1q1a_math", "1qna"]
assert all_cases[-1].video_rel == "videos_bench/nested/guide.mp4"
assert [case.subset for case in selected] == ["1q1a", "1q1a_math"]
def _archive(path: Path, member: str) -> None:
with tarfile.open(path, "w") as archive:
content = b"{}"
item = tarfile.TarInfo(member)
item.size = len(content)
archive.addfile(item, io.BytesIO(content))
def test_archive_is_safe_and_atomically_shared(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
bad = tmp_path / "bad.tar"
_archive(bad, "../escape")
with pytest.raises(ValueError, match="Unsafe path"):
data._extract_archive(bad, tmp_path / "bad-cache")
good = tmp_path / "good.tar"
_archive(good, "data/1q1a/video_json_map.json")
target = tmp_path / "cache"
barrier, safe_extract = threading.Barrier(2), data._safe_extract
monkeypatch.setattr(data, "_safe_extract", lambda *args: (barrier.wait(), safe_extract(*args)))
with ThreadPoolExecutor(2) as executor:
roots = list(executor.map(lambda _: data._extract_archive(good, target), range(2)))
assert roots[0] == roots[1]
assert (roots[0] / "1q1a" / "video_json_map.json").is_file()
assert not list(target.glob(".tmp-*"))
def test_hub_archive_uses_vllm_filesystem(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
source = tmp_path / "source.tar"
_archive(source, "data/1q1a/video_json_map.json")
class FS:
def get_file(self, remote: str, local: str) -> None:
assert remote == "datasets/org/repo/data.tar.gz"
Path(local).write_bytes(source.read_bytes())
monkeypatch.setenv("HF_HOME", str(tmp_path / "hf"))
monkeypatch.setattr(data, "hf_fs", lambda: FS())
assert (data.resolve_omniinteract_root(None, "org/repo") / "1q1a").is_dir()
def test_media_commands_are_bounded(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
calls: list[float] = []
commands: list[list[str]] = []
bounded_run = oi._run_media_command
def run(command, *, timeout_s, text=False):
calls.append(timeout_s)
commands.append(command)
if command[0] == "ffprobe":
return subprocess.CompletedProcess(command, 0, "1.0", "")
return subprocess.CompletedProcess(command, 0, b"", b"")
monkeypatch.setattr(oi, "_run_media_command", run)
duration, pcm, frames = oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=1.0)
assert (duration, len(pcm), frames, calls) == (1.0, 32_000, [None], [3.0, 3.0, 3.0])
assert commands[1][commands[1].index("-t") + 1] == "1.0"
assert commands[1][commands[1].index("-fs") + 1] == "32000"
monkeypatch.setattr(subprocess, "run", lambda *a, **k: (_ for _ in ()).throw(subprocess.TimeoutExpired("x", 3)))
with pytest.raises(TimeoutError, match="timed out"):
bounded_run(["ffmpeg"], timeout_s=3.0)
def test_media_duration_is_bounded_before_decode(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
calls = 0
def run(command, *, timeout_s, text=False):
nonlocal calls
calls += 1
return subprocess.CompletedProcess(command, 0, "601.0", "")
monkeypatch.setattr(oi, "_run_media_command", run)
with pytest.raises(ValueError, match="Invalid video duration"):
oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=600.0)
assert calls == 1
def test_final_commit_keeps_a_partial_model_unit():
"""Exact units retain a sample for the server's correlated final flush."""
events = [{"session": {"capabilities": {"chunk_period_ms": 1000}}}]
exact_unit = b"\0\0" * 16_000
assert oi._ensure_final_commit_tail(exact_unit, events) == exact_unit[:-2]
assert oi._ensure_final_commit_tail(exact_unit + b"\0\0", events) == exact_unit + b"\0\0"
@pytest.mark.parametrize(
("events", "match"),
[
(((_done(), 0.0),), "without response.created"),
(((_created(), 0.0), (_created(), 0.1)), "duplicate response.created"),
(((_created(), 0.0), (_done(), 0.1), (_done(), 0.2)), "duplicate response.done"),
(
(
(_created(), 0.0),
({"type": "response.done", "response": {"id": "r1", "status_details": {"type": "failed"}}}, 0.1),
),
"reports failure",
),
],
)
def test_response_ledger_rejects_identity_errors(events, match: str):
with pytest.raises(ValueError, match=match):
oi.response_ledger(_collector(*events))
class _CompletionClient:
def __init__(self, collector: RealtimeEventCollector):
self.events, self.acks = collector, []
def raise_if_reader_stopped(self) -> None:
return None
async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
self.acks.append((response_id, played_ms))
@pytest.mark.asyncio
@pytest.mark.parametrize(
("tail", "succeeds"),
[
(((_created(), 1.1), (_done(), 1.2)), True),
(((_listen(), 1.1),), True),
(((_listen(buffering=True), 1.1),), False),
((({"type": "session.closed", "reason": "timeout"}, 1.1),), False),
],
)
async def test_completion_requires_a_final_model_decision(tail, succeeds: bool):
events = _collector(({"type": "input_audio_buffer.committed"}, 1.0), *tail)
coroutine = oi.wait_for_session_completion(
_CompletionClient(events), oi._Playback(), commit_from=0, timeout_s=0.08, settle_s=0.0
)
if succeeds:
assert await coroutine == 0
else:
with pytest.raises((TimeoutError, RuntimeError)):
await coroutine
@pytest.mark.asyncio
async def test_deferred_commit_ignores_the_prior_response_terminal():
events = _collector(
(_created("old"), 0.0),
({"type": "input_audio_buffer.committed", "event": {"overlap_deferred": True}}, 1.0),
(_done("old"), 1.1),
(_listen(), 1.2),
)
assert (
await oi.wait_for_session_completion(
_CompletionClient(events), oi._Playback(), commit_from=1, timeout_s=0.1, settle_s=0.0
)
== 1
)
@pytest.mark.asyncio
async def test_playback_acks_only_after_serial_audio_drain():
collector = _collector((_audio(samples=2400), 10.0), (_done(), 10.01))
client, playback = _CompletionClient(collector), oi._Playback()
await playback.acknowledge(client, now=10.05)
assert client.acks == []
await playback.acknowledge(client, now=10.1)
assert client.acks == [("r1", 100)]
@pytest.mark.parametrize(
("event", "match"),
[
(_audio(rate=16_000), "24000 Hz"),
({**_audio(), "format": "opus"}, "pcm16"),
({**_audio(), "delta": "!"}, "base64"),
],
)
def test_playback_rejects_non_official_audio(event: dict[str, object], match: str):
playback = oi._Playback()
with pytest.raises(ValueError, match=match):
playback.ingest(_collector((event, 1.0)))
def test_artifacts_publish_official_bundle_and_sparse_deferred_state(tmp_path: Path):
case, collector, result = _successful_output(tmp_path)
summary = _write_success(
tmp_path / "out",
case,
collector,
result,
)
context = result._artifact_context
assert result.success and result.eligible_for_official_eval
assert summary["scene_type"] == "multi_turn"
assert context and context.spans
assert all("delta" not in event for event in context.events if event["type"] == "response.audio.delta")
oi.publish_deferred_case_artifacts(tmp_path / "out", case, result)
directory = oi._output_dir(tmp_path / "out", case)
assert all((directory / name).is_file() for name in oi.SUCCESS_ARTIFACTS)
with wave.open(str(directory / "output.wav")) as output:
assert (output.getframerate(), output.getnframes()) == (24_000, 24_000)
transcript = json.loads((directory / "wav_transcript.json").read_text())
assert transcript["text"] == "hello" and transcript["chunks"][0]["timestamp"] == [0.0, 0.1]
assert "serialized playback queue" in transcript["timestamp_semantics"]
assert result._artifact_context is None
@pytest.mark.parametrize(
("audio_time", "status", "reason"),
[(10.95, "completed", "audio_clipped"), (10.0, "cancelled", "cancelled_response")],
)
def test_ineligible_outputs_are_excluded_from_manifest(tmp_path: Path, audio_time: float, status: str, reason: str):
case, collector, result = _successful_output(tmp_path, audio_time=audio_time, status=status)
_write_success(tmp_path / "out", case, collector, result, require_response=False)
assert reason in result.official_eval_ineligible_reasons
oi.write_batch_artifacts(tmp_path / "out", [case], [result])
assert (tmp_path / "out" / "official_eval_manifest.jsonl").read_text() == ""
def test_artifacts_require_complete_audio_and_transcript_response(tmp_path: Path):
case, _, result = _successful_output(tmp_path)
collector = _collector((_created(), 10.0), (_audio(), 10.0), (_done(), 10.1))
with pytest.raises(ValueError, match="audio and transcript"):
_write_success(tmp_path, case, collector, result)
def test_publication_failure_removes_partial_success_bundle(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
case, collector, result = _successful_output(tmp_path)
_write_success(tmp_path, case, collector, result)
original = oi._atomic_write_json
def fail(path: Path, value: object) -> None:
if path.name == "events.json":
raise OSError("disk full")
original(path, value)
monkeypatch.setattr(oi, "_atomic_write_json", fail)
with pytest.raises(OSError, match="disk full"):
oi.publish_deferred_case_artifacts(tmp_path, case, result)
directory = oi._output_dir(tmp_path, case)
assert not any((directory / name).exists() for name in oi.SUCCESS_ARTIFACTS)
assert not result.success and "artifact_write_failed" in result.official_eval_ineligible_reasons
def test_atomic_write_preserves_destination_on_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
destination = tmp_path / "output.wav"
destination.write_bytes(b"old")
monkeypatch.setattr(oi, "write_pcm16_wav", lambda *a, **k: (_ for _ in ()).throw(OSError("write")))
with pytest.raises(OSError):
oi._atomic_write_wav(destination, b"\0\0", 24_000)
assert destination.read_bytes() == b"old"
assert not list(tmp_path.glob(".*.tmp"))
class _RealtimeClient:
instances: list[_RealtimeClient] = []
def __init__(self, url: str, **kwargs):
self.url, self.events, self.acks, self.configure_kwargs = url, RealtimeEventCollector(), [], {}
self.instances.append(self)
async def __aenter__(self):
return self
async def __aexit__(self, *args):
return None
async def configure(self, model: str, **kwargs) -> None:
self.configure_kwargs = kwargs
self.events.add({"type": "session.created", "session": {"capabilities": {"chunk_period_ms": 1000}}})
async def send(self, event: dict[str, object]) -> None:
return None
async def commit(self) -> None:
now = time.monotonic()
for event, offset in (
({"type": "input_audio_buffer.committed"}, 0.0),
(_created(), 0.001),
(_audio(samples=24), 0.002),
(_text(value="answer"), 0.003),
(_done(), 0.004),
):
self.events.add(event, received_at_s=now + offset)
async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
self.acks.append((response_id, played_ms))
async def close_session(self, **kwargs) -> None:
self.events.add({"type": "session.closed"})
def raise_if_reader_stopped(self) -> None:
return None
@pytest.mark.asyncio
async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
case = _case(tmp_path)
config = oi.OmniInteractBenchmarkConfig(
model=MODEL,
output_root=tmp_path / "out",
ref_audio="ref.wav",
require_response=True,
extra_body={"custom": "value"},
)
prepared = data.OmniInteractPreparedInput(1.0, b"\0\0\0\0", ("frame",), "data:audio/wav;base64,ref")
monkeypatch.setattr(oi, "RealtimeDuplexClient", _RealtimeClient)
monkeypatch.setattr(oi, "_COMPLETION_SETTLE_S", 0.0)
result = await oi.run_omniinteract_case(
case,
config,
request_index=0,
prepared_input=prepared,
)
assert result.success and result.transcript == "answer" and result._artifact_context is not None
assert not config.output_root.exists()
assert "autostart=0" in _RealtimeClient.instances[-1].url
assert _RealtimeClient.instances[-1].configure_kwargs["extra_body"] == {"custom": "value"}
assert _RealtimeClient.instances[-1].acks == [("r1", 1)]
def test_standard_sample_loading_prepares_media_before_timing(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
_write_dataset(tmp_path)
ref = tmp_path / "ref.wav"
ref.touch()
parser = TrackingArgumentParser()
OmniBenchmarkServingSubcommand.add_cli_args(parser)
args = parser.parse_args(
[
"--backend",
"openai-realtime-duplex",
"--dataset-name",
"omniinteract",
"--dataset-path",
str(tmp_path),
"--model",
MODEL,
"--endpoint",
"/v1/realtime",
"--num-prompts",
"0",
"--omniinteract-ref-audio",
str(ref),
]
)
preprocess_serve_args(args)
assert args.endpoint == "/v1/realtime"
monkeypatch.setattr(benchmark_patch, "reference_audio_data_url", lambda _: "data:audio/wav;base64,ref")
monkeypatch.setattr(benchmark_patch, "prepare_media", lambda *a, **k: (1.0, b"pcm", ["frame"]))
samples = benchmark_patch.get_samples(args, None)
assert len(samples) == args.num_prompts == 3
assert all(sample.omniinteract_prepared_input.video_frames == ("frame",) for sample in samples)
def _request(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions) -> RequestFuncInput:
request = RequestFuncInput(
model=MODEL,
model_name=MODEL,
prompt="",
api_url="http://server:8000/v1/realtime",
prompt_len=0,
output_len=0,
logprobs=None,
multi_modal_content=None,
ignore_eos=False,
request_id="measured-0",
)
request.omniinteract_case, request.omniinteract_options = case, options
return request
@pytest.mark.asyncio
async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
case, options = _case(tmp_path), _session_options(tmp_path)
request = _request(case, options)
output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
assert not output.success and "prepared before benchmark request timing" in output.error
prepared = data.OmniInteractPreparedInput(1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref")
request.omniinteract_prepared_input = prepared
request.extra_headers = {"X-Proxy": "test"}
request.extra_body = {"custom": "value"}
monkeypatch.setenv("OPENAI_API_KEY", "secret")
seen = {}
async def run(_case_arg, config, **kwargs):
seen["headers"] = config.extra_headers
seen["extra_body"] = config.extra_body
seen.update(kwargs)
return oi.OmniInteractCaseResult(
case.subset,
str(case.video_path),
str(options.output_root),
success=True,
transcript="timing metadata missing",
output_tokens=0,
duplex_request_metrics=[{"request_metrics": {"ttft_ms": 1.0}}],
duplex_session_metrics={"mean_ttft_ms": 10.0},
)
monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run)
output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
assert output.success and seen["prepared_input"] is prepared
assert seen["headers"] == {"Authorization": "Bearer secret", "X-Proxy": "test", "x-request-id": "measured-0"}
assert seen["extra_body"] == {"custom": "value"}
assert output.itl == []
assert output.tpot_measured is False
assert seen["capture_artifacts"] is True
@pytest.mark.asyncio
@pytest.mark.parametrize(
("request_metrics", "expected_itl", "expected_text_latency"),
[
(
[
{"stage0_tokens": {"output_token_count": 3, "itls_ms": [10.0, 20.0]}},
{"stage0_tokens": {"output_token_count": 2, "itls_ms": [30.0]}},
],
[0.01, 0.02, 0.03],
0.16,
),
(
[
{"stage0_tokens": {"output_token_count": 3, "tpot_ms": 10.0}},
{"stage0_tokens": {"output_token_count": 2, "tpot_ms": 40.0}},
],
[],
0.18,
),
(
[
{"stage0_tokens": {"output_token_count": 3, "tpot_ms": 0.0}},
{"stage0_tokens": {"output_token_count": 2, "tpot_ms": 0.0}},
],
[],
0.1,
),
],
)
async def test_adapter_reports_exact_or_weighted_token_timing(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
request_metrics: list[dict[str, object]],
expected_itl: list[float],
expected_text_latency: float,
):
case, options = _case(tmp_path), _session_options(tmp_path)
request = _request(case, options)
request.omniinteract_prepared_input = data.OmniInteractPreparedInput(
1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref"
)
async def run(*args, **kwargs):
return oi.OmniInteractCaseResult(
case.subset,
str(case.video_path),
str(options.output_root),
success=True,
output_tokens=5,
duplex_request_metrics=request_metrics,
duplex_session_metrics={"mean_ttft_ms": 100.0},
)
monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run)
output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
assert output.success
assert output.itl == expected_itl
assert output.text_latency == pytest.approx(expected_text_latency)
assert output.tpot_measured is True
def test_batch_finalization_publishes_only_measured_results(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
case, options = _case(tmp_path), _session_options(tmp_path)
options.output_root.mkdir()
sample = _sample(case, options)
result = oi.OmniInteractCaseResult(
case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
)
output = benchmark_patch.MixRequestFuncOutput(success=True)
output.omniinteract_case_result = result
published = []
monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: published.append(args))
summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
assert summary and summary["total"] == summary["success"] == summary["eligible_for_official_eval"] == 1
assert summary["failed"] == summary["successful_but_ineligible"] == summary["audio_clipped_bytes"] == 0
assert published == [(options.output_root, case, result)]
def test_case_artifact_failures_do_not_discard_serving_metrics(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
case, options = _case(tmp_path), _session_options(tmp_path)
sample = _sample(case, options)
result = oi.OmniInteractCaseResult(
case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
)
output = benchmark_patch.MixRequestFuncOutput(success=True, error="inference warning")
output.omniinteract_case_result = result
monkeypatch.setattr(
benchmark_patch,
"publish_deferred_case_artifacts",
lambda *args: (_ for _ in ()).throw(OSError("case disk full")),
)
monkeypatch.setattr(
benchmark_patch,
"write_failure_artifacts",
lambda *args: (_ for _ in ()).throw(OSError("failure disk full")),
)
monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", lambda *args: None)
summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
metrics, _ = calculate_metrics(
[],
[output],
1.0,
None,
[50.0],
{},
TaskType.GENERATION,
[],
None,
float("inf"),
1.0,
)
assert summary and summary["failed"] == 1
assert summary["artifacts_complete"] is False
assert len(summary["artifact_errors"]) == 2
assert output.success and metrics.completed == 1
assert output.error.startswith("inference warning\n")
assert "Artifact publication failed: case disk full" in output.error
def test_batch_artifact_failure_removes_partial_batch_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
case, options = _case(tmp_path), _session_options(tmp_path)
sample = _sample(case, options)
result = oi.OmniInteractCaseResult(
case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
)
output = benchmark_patch.MixRequestFuncOutput(success=True)
output.omniinteract_case_result = result
monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: None)
def fail_batch(root, *args):
root.mkdir(parents=True, exist_ok=True)
(root / "batch_summary.json").write_text("partial")
raise OSError("batch disk full")
monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", fail_batch)
summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
assert summary and summary["success"] == 1 and summary["artifacts_complete"] is False
assert summary["artifact_errors"] == ["Batch artifact publication failed: batch disk full"]
assert output.success
assert not (options.output_root / "batch_summary.json").exists()
@pytest.mark.parametrize(
("field", "value", "match"),
[
("omniinteract_ref_audio", None, "ref-audio"),
("endpoint", "/v1/completions", "v1/realtime"),
("max_concurrency", 0, "positive"),
("skip_tokenizer_init", True, "skip-tokenizer-init"),
],
)
def test_cli_rejects_incompatible_omniinteract_options(field: str, value, match: str):
args = argparse.Namespace(
dataset_name="omniinteract",
backend="openai-realtime-duplex",
endpoint="/v1/realtime",
omniinteract_ref_audio="ref.wav",
max_concurrency=1,
skip_tokenizer_init=False,
)
setattr(args, field, value)
with pytest.raises(ValueError, match=match):
preprocess_serve_args(args)
def test_whole_benchmark_lock_serializes_a_shared_output_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
entered, release = threading.Event(), threading.Event()
calls: list[int] = []
async def fake_main_async(args):
calls.append(len(calls))
if len(calls) == 1:
entered.set()
await asyncio.to_thread(release.wait)
return {"total": 1}
monkeypatch.setattr(benchmark_serve, "main_async", fake_main_async)
args = argparse.Namespace(dataset_name="omniinteract", omniinteract_output_dir=tmp_path, extra_body=None)
with ThreadPoolExecutor(2) as executor:
first = executor.submit(benchmark_serve.main, args)
assert entered.wait(1)
second = executor.submit(benchmark_serve.main, args)
time.sleep(0.05)
assert calls == [0]
release.set()
assert first.result() == second.result() == {"total": 1}
assert calls == [0, 1]
+33
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import argparse
import json
import os
@@ -22,6 +25,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
"dataset_name",
[
"daily-omni",
"omniinteract",
"seed-tts",
"seed-tts-text",
"seed-tts-design",
@@ -55,6 +59,7 @@ def test_extend_omni_choices_updates_tracking_parser_shadow(dataset_name: str) -
[
(["--print-stage"], "print_stage", True),
(["--daily-omni-input-mode", "audio"], "daily_omni_input_mode", "audio"),
(["--omniinteract-subsets", "1qna"], "omniinteract_subsets", ["1qna"]),
(["--seed-tts-locale", "zh"], "seed_tts_locale", "zh"),
],
)
@@ -80,12 +85,14 @@ def test_add_omni_args_preserves_implicit_defaults() -> None:
args = parser.parse_args([])
assert args.print_stage is False
assert args.daily_omni_input_mode == "all"
assert args.omniinteract_subsets == ["1q1a", "1q1a_math", "1qna"]
assert args.seed_tts_locale == "en"
assert args.explicit_keys == set()
def test_update_omni_help_updates_upstream_actions() -> None:
parser = TrackingArgumentParser()
parser.add_argument("--num-prompts", default=1000, help="Number of prompts.")
parser.add_argument("--percentile-metrics", help="Upstream percentile help.")
parser.add_argument("--random-mm-limit-mm-per-prompt", help="Upstream limit help.")
parser.add_argument("--random-mm-bucket-config", help="Upstream bucket help.")
@@ -93,6 +100,7 @@ def test_update_omni_help_updates_upstream_actions() -> None:
update_omni_help(parser)
actions = {action.dest: action for action in parser._actions}
assert "OmniInteract uses 3" in actions["num_prompts"].help
assert all(metric in actions["percentile_metrics"].help for metric in ("ttfc", "tpop", "audio_rtf"))
assert "probabilities are renormalized" in actions["random_mm_limit_mm_per_prompt"].help
assert "Currently allows for 3 modalities" in actions["random_mm_bucket_config"].help
@@ -120,6 +128,31 @@ def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body(
assert args.extra_body == expected
@pytest.mark.parametrize(
("argv", "expected"),
[
([], 3),
(["--num-prompts", "0"], 0),
(["--num-prompts", "12"], 12),
],
)
def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default(
argv: list[str],
expected: int,
) -> None:
parser = TrackingArgumentParser()
parser.add_argument("--dataset-name", default="omniinteract")
parser.add_argument("--backend", default="openai-realtime-duplex")
parser.add_argument("--endpoint", default="/v1/realtime")
parser.add_argument("--omniinteract-ref-audio", default="reference.wav")
parser.add_argument("--num-prompts", type=int, default=1000)
args = parser.parse_args(argv)
preprocess_serve_args(args)
assert args.num_prompts == expected
@pytest.mark.parametrize(
("argv", "expected_extra_body", "expected_explicit"),
[
+317
View File
@@ -0,0 +1,317 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Drift checks for the reviewed environment-variable inventory."""
import ast
import builtins
from collections import Counter
from pathlib import Path
import pytest
from vllm_omni.config.environment_variable_inventory import (
ENVIRONMENT_VARIABLE_INVENTORY,
EnvironmentVariableCategory,
ModelEnvironmentVariableDisposition,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_REPO_ROOT = Path(__file__).parents[2]
_PACKAGE_ROOT = _REPO_ROOT / "vllm_omni"
_INVENTORY_MODULE = _PACKAGE_ROOT / "config" / "environment_variable_inventory.py"
_REFERENCE_PAGE = _REPO_ROOT / "docs" / "configuration" / "environment_variables.md"
def _module_string_constants(tree: ast.Module) -> dict[str, str]:
constants: dict[str, str] = {}
for statement in tree.body:
targets: list[ast.expr] = []
value: ast.expr | None = None
if isinstance(statement, ast.Assign):
targets = statement.targets
value = statement.value
elif isinstance(statement, ast.AnnAssign):
targets = [statement.target]
value = statement.value
if isinstance(value, ast.Constant) and isinstance(value.value, str):
for target in targets:
if isinstance(target, ast.Name):
constants[target.id] = value.value
return constants
def _os_import_aliases(tree: ast.Module) -> tuple[set[str], set[str], set[str]]:
os_names: set[str] = set()
getenv_names: set[str] = set()
environ_names: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for imported in node.names:
if imported.name == "os":
os_names.add(imported.asname or "os")
elif isinstance(node, ast.ImportFrom) and node.module == "os":
for imported in node.names:
if imported.name == "getenv":
getenv_names.add(imported.asname or imported.name)
elif imported.name == "environ":
environ_names.add(imported.asname or imported.name)
return os_names, getenv_names, environ_names
def _is_environ(node: ast.expr, os_names: set[str], environ_names: set[str]) -> bool:
return (
isinstance(node, ast.Attribute)
and isinstance(node.value, ast.Name)
and node.value.id in os_names
and node.attr == "environ"
) or (isinstance(node, ast.Name) and node.id in environ_names)
def _environment_key_expressions(
node: ast.AST,
os_names: set[str],
getenv_names: set[str],
environ_names: set[str],
) -> list[ast.expr]:
"""Return key expressions used by one direct environment access."""
if isinstance(node, ast.Call) and node.args:
func = node.func
if (
isinstance(func, ast.Attribute)
and isinstance(func.value, ast.Name)
and func.value.id in os_names
and func.attr == "getenv"
) or (isinstance(func, ast.Name) and func.id in getenv_names):
return [node.args[0]]
if (
isinstance(func, ast.Attribute)
and _is_environ(func.value, os_names, environ_names)
and func.attr in {"get", "pop", "setdefault"}
):
return [node.args[0]]
elif isinstance(node, ast.Subscript) and _is_environ(node.value, os_names, environ_names):
return [node.slice]
elif isinstance(node, ast.Compare):
expressions: list[ast.expr] = []
left_values = [node.left, *node.comparators[:-1]]
for left, operator, right in zip(left_values, node.ops, node.comparators):
if isinstance(operator, (ast.In, ast.NotIn)) and _is_environ(right, os_names, environ_names):
expressions.append(left)
return expressions
return []
def _resolve_environment_name(expression: ast.expr, constants: dict[str, str]) -> str | None:
value: str | None = None
if isinstance(expression, ast.Constant) and isinstance(expression.value, str):
value = expression.value
elif isinstance(expression, ast.Name):
value = constants.get(expression.id)
return value or None
def _environment_wrapper_parameters(
tree: ast.Module,
os_names: set[str],
getenv_names: set[str],
environ_names: set[str],
) -> dict[str, set[tuple[str, int | None]]]:
"""Find local helpers that forward one parameter as an environment key."""
wrappers: dict[str, set[tuple[str, int | None]]] = {}
for function in (node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))):
positional = [*function.args.posonlyargs, *function.args.args]
parameter_indexes: dict[str, int | None] = {parameter.arg: index for index, parameter in enumerate(positional)}
parameter_indexes.update({parameter.arg: None for parameter in function.args.kwonlyargs})
forwarded: set[tuple[str, int | None]] = set()
for node in ast.walk(function):
for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
if isinstance(expression, ast.Name) and expression.id in parameter_indexes:
forwarded.add((expression.id, parameter_indexes[expression.id]))
if forwarded:
wrappers[function.name] = forwarded
return wrappers
def _environment_accesses(path: Path) -> set[str]:
"""Return statically resolvable names accessed through ``os``."""
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
constants = _module_string_constants(tree)
os_names, getenv_names, environ_names = _os_import_aliases(tree)
wrappers = _environment_wrapper_parameters(tree, os_names, getenv_names, environ_names)
names: set[str] = set()
for node in ast.walk(tree):
for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
if name := _resolve_environment_name(expression, constants):
names.add(name)
if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
for parameter_name, position in wrappers.get(node.func.id, set()):
argument = node.args[position] if position is not None and position < len(node.args) else None
if argument is None:
argument = next(
(keyword.value for keyword in node.keywords if keyword.arg == parameter_name),
None,
)
if argument is not None and (name := _resolve_environment_name(argument, constants)):
names.add(name)
return names
def test_inventory_matches_reviewed_snapshot_counts():
"""Make an inventory expansion an explicit review decision."""
category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
assert category_counts == {
EnvironmentVariableCategory.PUBLIC_OMNI: 23,
EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
EnvironmentVariableCategory.INTERNAL: 2,
}
disposition_counts = Counter(
item.model_disposition
for item in ENVIRONMENT_VARIABLE_INVENTORY.values()
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
)
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
ModelEnvironmentVariableDisposition.PROMOTE: 33,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6,
}
def test_new_public_omni_names_use_project_prefix():
"""Grandfather legacy names without allowing more prefix exceptions."""
legacy_public_names = {
"DIFFUSION_ATTENTION_BACKEND",
"DIFFUSION_CACHE_ADAPTER",
"DIFFUSION_CACHE_BACKEND",
"OMNI_DIFFUSION_PROMPT_EMBED_CACHE",
"OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE",
"OMNI_DIFFUSION_SESSION_STATE_MANAGER",
"OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS",
"SPEAKER_MAX_UPLOADED",
"SPEAKER_SAMPLES_DIR",
"VLLM_VIDEO_ASYNC_CHUNK",
"VLLM_VIDEO_AUDIO_DELTA_MODE",
}
public_names = {
name
for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
if item.category is EnvironmentVariableCategory.PUBLIC_OMNI
}
assert {name for name in public_names if not name.startswith("VLLM_OMNI_")} == legacy_public_names
def test_statically_resolvable_environment_accesses_are_classified():
discovered: set[str] = set()
for path in _PACKAGE_ROOT.rglob("*.py"):
discovered.update(_environment_accesses(path))
assert discovered - ENVIRONMENT_VARIABLE_INVENTORY.keys() == set()
def test_model_and_benchmark_inventory_entries_are_still_referenced():
"""Reject transitional rows after their implementation has disappeared."""
package_source = "\n".join(
path.read_text(encoding="utf-8") for path in _PACKAGE_ROOT.rglob("*.py") if path != _INVENTORY_MODULE
)
transitional_categories = {
EnvironmentVariableCategory.MODEL_SPECIFIC,
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL,
}
stale_names = {
name
for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
if item.category in transitional_categories and name not in package_source
}
assert stale_names == set()
def test_environment_scanner_covers_indirection_aliases_membership_and_casing():
expected_by_path = {
"metrics/definitions.py": {
"VLLM_OMNI_BENCH_AUDIO_CHANNELS",
"VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE",
},
"entrypoints/openai/video_stream_envs.py": {
"VLLM_VIDEO_ASYNC_CHUNK",
"VLLM_VIDEO_AUDIO_DELTA_MODE",
},
"model_executor/models/mimo_audio/mimo_audio.py": {"model_stage"},
"model_executor/models/moss_tts/modeling_moss_tts_local_depth.py": {"MOSS_TTS_DEBUG_STOP"},
"distributed/ray_utils/utils.py": {"RAY_RAYLET_PID"},
}
for relative_path, expected in expected_by_path.items():
assert expected <= _environment_accesses(_PACKAGE_ROOT / relative_path)
def test_generated_server_storage_environment_names_are_classified():
from vllm_omni.config.server_settings import FileBackend
generated_names = {
f"VLLM_OMNI_SERVER_STORAGE__{field_name.upper()}"
for field_name in FileBackend.model_fields
if field_name != "type"
}
assert generated_names <= ENVIRONMENT_VARIABLE_INVENTORY.keys()
def test_public_omni_variables_are_in_the_reference_page():
reference = _REFERENCE_PAGE.read_text(encoding="utf-8")
missing = {
name
for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
if item.is_public_omni and f"`{name}`" not in reference
}
assert missing == set()
def test_secret_values_are_marked_for_redaction():
assert {name for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items() if item.redact_value} == {
"HF_TOKEN",
"HUGGINGFACE_HUB_TOKEN",
"OPENAI_API_KEY",
}
def test_collect_env_reports_only_safe_public_omni_values(monkeypatch):
from collect_env import get_env_vars
monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "tea_cache")
monkeypatch.setenv("HF_TOKEN", "hf-secret")
monkeypatch.setenv("OPENAI_API_KEY", "openai-secret")
monkeypatch.setenv("VLLM_OMNI_REPLICA_ID", "internal-replica")
report = get_env_vars()
assert "DIFFUSION_CACHE_BACKEND=tea_cache" in report.splitlines()
assert "hf-secret" not in report
assert "openai-secret" not in report
assert "VLLM_OMNI_REPLICA_ID" not in report
@pytest.mark.parametrize("error_type", [ImportError, OSError, AssertionError, RuntimeError])
def test_collect_env_survives_inventory_import_error(monkeypatch, error_type):
from collect_env import get_env_vars
real_import = builtins.__import__
def fail_inventory_import(name, *args, **kwargs):
if name == "vllm_omni.config.environment_variable_inventory":
raise error_type("broken vllm-omni installation")
return real_import(name, *args, **kwargs)
monkeypatch.setattr(builtins, "__import__", fail_inventory_import)
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
assert "CUDA_VISIBLE_DEVICES=0" in get_env_vars().splitlines()
@@ -21,47 +21,6 @@
]
},
"benchmark_params": [
{
"dataset_name": "random",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [
4,
16,
32,
64,
128
],
"max_concurrency": [
1,
4,
8,
16,
32
],
"random_input_len": 2500,
"random_output_len": 900,
"ignore_eos": true,
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"mean_e2el_ms": [
11701.3818,
27510.3956,
31907.9613,
42173.9832,
63970.0751
],
"mean_audio_rtf": [
0.1175,
0.1647,
0.2003,
0.2726,
0.4195
]
}
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
@@ -170,45 +129,6 @@
}
}
},
{
"dataset_name": "random",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [
4,
16,
32,
64,
128
],
"max_concurrency": [
1,
4,
8,
16,
32
],
"random_input_len": 2500,
"random_output_len": 900,
"ignore_eos": true,
"extra_body": {
"modalities": [
"text"
]
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"H100": {
"mean_e2el_ms": [
4354.8945,
7672.4248,
9861.1687,
13086.8865,
18716.709
]
}
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
@@ -324,5 +244,61 @@
}
}
]
},
{
"test_name": "test_qwen3_omni_chunk_cuda",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 2
}
},
"omni",
"slow"
],
"server_params": {
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"extra_cli_args": [
"--async-chunk"
]
},
"benchmark_params": [
{
"dataset_name": "random",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [
4,
32,
128
],
"max_concurrency": [
1,
8,
32
],
"random_input_len": 100,
"random_output_len": 100,
"ignore_eos": true,
"extra_body": {
"modalities": [
"text"
]
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"H100": {
"mean_e2el_ms": [
4354.8945,
9861.1687,
18716.709
]
}
}
}
]
}
]
@@ -29,16 +29,12 @@
"endpoint": "/v1/completions",
"num_prompts": [
4,
16,
32,
64,
128
],
"max_concurrency": [
1,
4,
8,
16,
32
],
"random_input_len": 2500,
@@ -54,9 +50,7 @@
"H100": {
"mean_e2el_ms": [
4619.0068,
7868.3523,
9760.6175,
13422.3491,
18007.6918
]
}
@@ -99,45 +93,6 @@
}
}
},
{
"dataset_name": "random-mm",
"backend": "vllm",
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"endpoint": "/v1/completions",
"num_prompts": [
40
],
"request_rate": [
0.5
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
"ignore_eos": true,
"extra_body": {
"modalities": [
"text"
]
},
"random_mm_base_items_per_request": 2,
"random_mm_num_mm_items_range_ratio": 0.5,
"random_mm_limit_mm_per_prompt": {
"image": 1,
"video": 1
},
"random_mm_bucket_config": {
"(256, 256, 1)": 0.5,
"(720, 1280, 2)": 0.5
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"H100": {
"mean_e2el_ms": [
524.5321
]
}
}
},
{
"dataset_name": "random-mm",
"backend": "vllm",
@@ -74,13 +74,14 @@
"full_model",
"diffusion"
],
"description": "Single-device baseline (no parallelism) with step execution",
"description": "Single-device step execution: sequential 512/1536 plus 512 concurrency sweep. One server (max-num-seqs 8).",
"server_type": "vllm-omni",
"server_params": {
"model": "Qwen/Qwen-Image",
"serve_args": {
"enable-diffusion-pipeline-profiler": true,
"step-execution": true
"step-execution": true,
"max-num-seqs": 8
}
},
"benchmark_params": [
@@ -119,6 +120,27 @@
"peak_memory_mb_mean": 64500.0
}
}
},
{
"name": "512x512_steps20_high_concurrency",
"dataset": "random",
"task": "t2i",
"width": 512,
"height": 512,
"num-inference-steps": 20,
"num-prompts": [20, 40, 80, 160],
"max-concurrency": [1, 2, 4, 8],
"warmup-requests": 8,
"warmup-concurrency": 8,
"warmup-num-inference-steps": 20,
"enable-negative-prompt": true,
"baseline": {
"H100": {
"throughput_qps": [0.3937, 0.6606, 0.6188, 0.7752],
"latency_mean": [2.5466, 3.0235, 6.3948, 10.1944],
"peak_memory_mb_mean": [55149.0, 55150.0, 55153.1458, 55156.325]
}
}
}
]
},
@@ -243,53 +265,5 @@
}
}
]
},
{
"test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 1
}
},
"full_model",
"diffusion"
],
"description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
"server_type": "vllm-omni",
"server_params": {
"model": "Qwen/Qwen-Image",
"serve_args": {
"enable-diffusion-pipeline-profiler": true,
"step-execution": true,
"max-num-seqs": 8
}
},
"benchmark_params": [
{
"name": "512x512_steps20_high_concurrency",
"dataset": "random",
"task": "t2i",
"width": 512,
"height": 512,
"num-inference-steps": 20,
"num-prompts": [20, 40, 80, 160],
"max-concurrency": [1, 2, 4, 8],
"warmup-requests": 8,
"warmup-concurrency": 8,
"warmup-num-inference-steps": 20,
"enable-negative-prompt": true,
"baseline": {
"H100": {
"throughput_qps": [0.3937, 0.6606, 0.6188, 0.7752],
"latency_mean": [2.5466, 3.0235, 6.3948, 10.1944],
"peak_memory_mb_mean": [55149.0, 55150.0, 55153.1458, 55156.325]
}
}
}
]
}
]
@@ -157,8 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
"ref_audio": None,
"ref_text": None,
},
("CustomVoice", "does not support"),
id="customvoice_task_not_supported",
("Invalid voice", "vivian"),
id="customvoice_invalid_voice",
),
pytest.param(
{"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
@@ -181,13 +181,17 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
("mutually exclusive", "speaker_embedding", "ref_audio"),
id="speaker_embedding_with_ref_audio_mutually_exclusive",
),
pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "greater_than_equal"), id="max_new_tokens_below_min"),
pytest.param(
{"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
("max_new_tokens", "exceed 4096"),
id="max_new_tokens_above_max",
),
pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
pytest.param(
{"seed": -(2**63) - 1},
("seed", "greater_than_equal", "-9223372036854775808"),
id="seed_below_min",
),
pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
pytest.param(
{"initial_codec_chunk_frames": -1},
@@ -275,7 +279,9 @@ def test_speech_invalid_field_values(
pytest.param({"input": "", "response_format": "wav"}, ("input", "empty"), id="input_empty"),
pytest.param({"input": " ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"),
pytest.param(
{"input": "Hello world.", "max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"
{"input": "Hello world.", "max_new_tokens": 0},
("max_new_tokens", "greater_than_equal"),
id="max_new_tokens_below_min",
),
pytest.param(
{
@@ -554,14 +560,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, online_client: Onli
pytest.param(
"batch",
{"max_new_tokens": 0},
("max_new_tokens", "least 1"),
("max_new_tokens", "greater_than_equal"),
id="batch_max_new_tokens_below_min",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{"max_new_tokens": 0},
("max_new_tokens", "least 1"),
("max_new_tokens", "greater_than_equal"),
id="item_max_new_tokens_below_min",
marks=_SKIP_ISSUE_3649,
),
@@ -72,6 +72,9 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
body["top_logprobs"] = 5
elif case_id == "speaker_unknown":
body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz"
elif case_id == "audio_format_unsupported":
body["modalities"] = ["text", "audio"]
body["audio"] = {"voice": "alloy", "format": "aac"}
else:
raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
return {"json": body, "timeout": 120}
@@ -119,7 +122,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
pytest.param(
"response_format_json_schema_incomplete",
400,
("response_format", "value_error", "json_schema"),
("response_format", "BadRequestError", "json_schema"),
id="invalid_response_format_json_schema",
),
pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
@@ -135,6 +138,12 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
("Invalid speaker", "Supported"),
id="speaker_unknown_preset",
),
pytest.param(
"audio_format_unsupported",
400,
("Invalid audio format", "aac", "Supported formats"),
id="unsupported_audio_format",
),
],
)
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
) -> None:
"""Run the same prompts sequentially then in parallel and print a comparison."""
omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
try:
await warmup(omni, WARMUP_PROMPTS)
single_time = await run_single(omni, prompts)
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
return
omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
try:
await warmup(omni, WARMUP_PROMPTS)
@@ -524,7 +524,12 @@ if __name__ == "__main__":
parser = argparse.ArgumentParser(description="E2E diffusion concurrent benchmark / validation")
parser.add_argument("--model", type=str, required=True, help="Model name or path")
parser.add_argument("--num-prompts", type=int, default=8, help="Number of prompts to run")
parser.add_argument("--batch-size", type=int, default=1, help="Diffusion batch size (1 = no batching)")
parser.add_argument(
"--batch-size",
type=int,
default=1,
help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size",
)
parser.add_argument(
"--mode",
choices=["batch", "single", "compare", "validate"],
@@ -24,6 +24,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlanResult,
TensorBinding,
)
from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
from vllm_omni.diffusion.models.helios import HeliosPipeline
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
from vllm_omni.diffusion.registry import initialize_model
@@ -173,6 +174,15 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
str(canonical_root / "model.safetensors"),
)
store_root = tmp_path / "hwr-store"
hash_calls = 0
original_sha256 = source_identity_module._sha256_file
def counted_sha256(path: Path, state: object) -> str:
nonlocal hash_calls
hash_calls += 1
return original_sha256(path, state) # type: ignore[arg-type]
monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
@@ -205,6 +215,8 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
assert warm_plan is not None
assert warm_plan.lease_carrier is not None
warm_plan.lease_carrier.close()
assert hash_calls == 1
assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1
def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch):
@@ -6,6 +6,8 @@ from __future__ import annotations
import dataclasses
import gc
import json
import threading
from collections.abc import Sequence
from pathlib import Path
from types import SimpleNamespace
@@ -24,10 +26,12 @@ from vllm_omni.diffusion.model_loader.host_weights import (
FinalLayoutRequest,
FinalLayoutTensorRestorer,
ImplementationIdentity,
NodeSourceDigestCache,
PreparedWeightSource,
WeightSourceKind,
build_final_layout_identity,
)
from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
from vllm_omni.diffusion.model_loader.host_weights.contracts import (
FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
FinalLayoutArtifactSpec,
@@ -380,6 +384,7 @@ def _identity(
*,
request: FinalLayoutRequest | None = None,
policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY,
source_digest_cache: NodeSourceDigestCache | None = None,
) -> FinalLayoutIdentityContext:
transformer = model.get_submodule("transformer")
return build_final_layout_identity(
@@ -388,6 +393,7 @@ def _identity(
prepared_sources=(source,),
request=request or _request(),
policy=policy,
source_digest_cache=source_digest_cache,
)
@@ -564,6 +570,131 @@ def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path
assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID
def test_node_source_digest_cache_reuses_recovers_and_invalidates(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
model = _TinyPipeline()
source = _prepared_source(tmp_path)
cache_root = tmp_path / "store"
calls = 0
original_sha256 = source_identity_module._sha256_file
def counted_sha256(path: Path, state: object) -> str:
nonlocal calls
calls += 1
return original_sha256(path, state) # type: ignore[arg-type]
monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
first = _identity(
model,
source,
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
)
second = _identity(
model,
source,
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
)
assert second.identity == first.identity
assert calls == 1
cache_entry = next((cache_root / "source-digests-v1" / "entries").glob("*.json"))
cache_entry.write_text("not-json", encoding="utf-8")
recovered = _identity(
model,
source,
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
)
assert recovered.identity == first.identity
assert calls == 2
cache_document = json.loads(cache_entry.read_text(encoding="utf-8"))
cache_document["content_id"] = f"sha256:{'0' * 64}"
cache_entry.write_text(json.dumps(cache_document), encoding="utf-8")
checksum_recovered = _identity(
model,
source,
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
)
assert checksum_recovered.identity == first.identity
assert calls == 3
source.weight_files[0].write_bytes(b"changed-canonical-source-with-a-new-size")
changed = _identity(
model,
source,
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
)
assert changed.identity != first.identity
assert calls == 4
def test_source_identity_hashes_checkpoint_shards_in_parallel(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
model = _TinyPipeline()
source = _prepared_source(tmp_path)
second_file = source.resolved_root / "model-00002.safetensors"
second_file.write_bytes(b"second-checkpoint-shard")
source = dataclasses.replace(source, weight_files=(*source.weight_files, second_file))
barrier = threading.Barrier(2)
thread_ids: set[int] = set()
thread_ids_lock = threading.Lock()
original_sha256 = source_identity_module._sha256_file
def synchronized_sha256(path: Path, state: object) -> str:
with thread_ids_lock:
thread_ids.add(threading.get_ident())
barrier.wait(timeout=5)
return original_sha256(path, state) # type: ignore[arg-type]
monkeypatch.setattr(source_identity_module, "_sha256_file", synchronized_sha256)
_identity(model, source)
assert len(thread_ids) == 2
def test_source_digest_cache_lock_timeout_falls_back_to_direct_hashing(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
model = _TinyPipeline()
source = _prepared_source(tmp_path)
cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
def unavailable_lock(*_args: object, **_kwargs: object) -> object:
raise source_identity_module.FileLockTimeoutError("busy")
monkeypatch.setattr(source_identity_module, "FileLock", unavailable_lock)
context = _identity(model, source, source_digest_cache=cache)
assert context.identity.source.revision.startswith("content-")
assert not tuple(cache.entries.glob("*.json"))
def test_source_digest_cache_directory_failure_falls_back_to_direct_hashing(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
model = _TinyPipeline()
source = _prepared_source(tmp_path)
original_mkdir = Path.mkdir
def failing_cache_mkdir(path: Path, *args: object, **kwargs: object) -> None:
if "source-digests-v1" in path.parts:
raise OSError("injected cache directory failure")
original_mkdir(path, *args, **kwargs) # type: ignore[arg-type]
monkeypatch.setattr(Path, "mkdir", failing_cache_mkdir)
cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
context = _identity(model, source, source_digest_cache=cache)
assert context.identity.source.revision.startswith("content-")
def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None:
model = _TinyPipeline()
source_root = tmp_path / "local-symlink-source"
@@ -0,0 +1,250 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import ctypes
from types import SimpleNamespace
import pytest
import torch
import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
from vllm_omni.diffusion.offloader.cuda_host_registration import (
CudaHostRegistration,
_coalesce_ranges,
_mapped_regions,
)
from vllm_omni.diffusion.offloader.host_registration import (
HostRegistrationBudgetError,
HostRegistrationCleanupError,
HostRegistrationError,
register_host_mappings,
)
from vllm_omni.host_weight_runtime import MappedHostRegion
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
class _FakeRuntime:
def __init__(
self,
register_results: list[int | Exception],
unregister_results: list[int | Exception] | None = None,
*,
read_only_supported: bool = True,
attribute_result: int = 0,
) -> None:
self._register_results = iter(register_results)
self._unregister_results = iter(unregister_results or [])
self._read_only_supported = read_only_supported
self._attribute_result = attribute_result
self.attribute_queries: list[tuple[int, int]] = []
self.registered: list[tuple[int, int, int]] = []
self.unregistered: list[int] = []
def cudaDeviceGetAttribute(self, value, attribute: int, device: int) -> int:
self.attribute_queries.append((attribute, device))
ctypes.cast(value, ctypes.POINTER(ctypes.c_int)).contents.value = int(self._read_only_supported)
return self._attribute_result
def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
self.registered.append((pointer, size, flags))
result = next(self._register_results)
if isinstance(result, Exception):
raise result
return result
def cudaHostUnregister(self, pointer: int) -> int:
self.unregistered.append(pointer)
result = next(self._unregister_results, 0)
if isinstance(result, Exception):
raise result
return result
@staticmethod
def cudaGetErrorString(error: int) -> str:
return f"error-{error}"
def _region(
file_name: str,
address: int,
length: int,
*,
read_only: bool = True,
) -> MappedHostRegion:
return MappedHostRegion(file_name=file_name, address=address, length=length, read_only=read_only)
@pytest.fixture(autouse=True)
def _cuda_runtime_prerequisites(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(registration_module.torch.accelerator, "current_device_index", lambda: 0)
def test_platform_dispatch_and_unsupported_device(monkeypatch: pytest.MonkeyPatch) -> None:
sentinel = object()
regions = (_region("weights", 0x1000, 4096),)
def create(actual_regions, *, max_bytes):
assert actual_regions is regions
assert max_bytes == 4096
return sentinel
monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
assert register_host_mappings(regions, device=torch.device("cuda"), max_bytes=4096) is sentinel
with pytest.raises(HostRegistrationError, match="not supported on cpu"):
register_host_mappings(regions, device=torch.device("cpu"), max_bytes=None)
def test_ranges_align_per_file_without_merging_adjacent_files() -> None:
assert _coalesce_ranges([(0x1003, 4096), (0x2800, 1024)], page_size=4096) == (
registration_module._AddressRange(0x1000, 0x3000),
)
assert _mapped_regions(
(
_region("first.safetensors", 0x1000, 4096),
_region("second.safetensors", 0x2000, 4096),
)
) == (
registration_module._AddressRange(0x1000, 0x2000),
registration_module._AddressRange(0x2000, 0x3000),
)
def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
runtime = _FakeRuntime([0])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
with pytest.raises(HostRegistrationError, match="is writable"):
CudaHostRegistration.create((_region("weights", 0x1000, 1, read_only=False),), max_bytes=None)
with pytest.raises(HostRegistrationBudgetError, match="exceeding"):
CudaHostRegistration.create((_region("weights", 0x1003, 4096),), max_bytes=4096)
assert runtime.registered == []
def test_registration_requires_read_only_capability(monkeypatch: pytest.MonkeyPatch) -> None:
runtime = _FakeRuntime([0], read_only_supported=False)
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
with pytest.raises(HostRegistrationError, match="does not support read-only"):
CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None)
assert runtime.attribute_queries == [
(registration_module._CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, 0)
]
assert runtime.registered == []
def test_capability_error_is_consumed_before_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
runtime = _FakeRuntime([0], attribute_result=7)
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
consumed: list[tuple[object, int]] = []
monkeypatch.setattr(
registration_module,
"_consume_last_cuda_error",
lambda actual_runtime, error: consumed.append((actual_runtime, error)),
)
with pytest.raises(HostRegistrationError, match="error-7"):
CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None)
assert consumed == [(runtime, 7)]
def test_partial_registration_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
runtime = _FakeRuntime([0, 7])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
with pytest.raises(HostRegistrationError, match="error-7"):
CudaHostRegistration.create(
(
_region("first", 0x1000, 4096),
_region("second", 0x9000, 4096),
),
max_bytes=None,
)
assert runtime.unregistered == [0x1000]
def test_failed_rollback_exposes_active_registration_for_retry(monkeypatch: pytest.MonkeyPatch) -> None:
runtime = _FakeRuntime([0, 7], unregister_results=[9, 0])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
with pytest.raises(HostRegistrationCleanupError, match="rollback errors") as error:
CudaHostRegistration.create(
(
_region("first", 0x1000, 4096),
_region("second", 0x9000, 4096),
),
max_bytes=None,
)
registration = error.value.active_registration
assert registration is not None
assert registration.close() == ()
assert runtime.unregistered == [0x1000, 0x1000]
def test_successful_registration_retries_failed_close(monkeypatch: pytest.MonkeyPatch) -> None:
runtime = _FakeRuntime([0], unregister_results=[9, 0])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
registration = CudaHostRegistration.create((_region("weights", 0x1003, 1),), max_bytes=None)
assert registration.total_bytes == 4096
assert registration.region_count == 1
assert registration.close() == ("cudaHostUnregister(0x1000) failed: error-9",)
assert registration.close() == ()
assert runtime.unregistered == [0x1000, 0x1000]
@pytest.mark.parametrize("pending_error", [0, 801])
def test_consume_last_cuda_error_accepts_cleared_or_matching_state(
monkeypatch: pytest.MonkeyPatch,
pending_error: int,
) -> None:
runtime = SimpleNamespace(cudaGetLastError=lambda: pending_error)
monkeypatch.setattr(
registration_module.ctypes,
"CDLL",
lambda _name: pytest.fail("the cudart handle should provide cudaGetLastError"),
)
registration_module._consume_last_cuda_error(runtime, 801)
def test_consume_last_cuda_error_falls_back_to_global_symbol(monkeypatch: pytest.MonkeyPatch) -> None:
global_runtime = SimpleNamespace(cudaGetLastError=lambda: 801)
monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: global_runtime)
registration_module._consume_last_cuda_error(SimpleNamespace(), 801)
def test_clean_rollback_remains_recoverable_when_error_state_cannot_be_cleared(
monkeypatch: pytest.MonkeyPatch,
) -> None:
runtime = _FakeRuntime([0, 7])
monkeypatch.setattr(runtime, "cudaGetLastError", None, raising=False)
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
def missing_global_symbol(_name: object) -> None:
raise OSError("cudart symbols are local")
monkeypatch.setattr(
registration_module.ctypes,
"CDLL",
missing_global_symbol,
)
with pytest.raises(HostRegistrationError, match="cannot clear") as error:
CudaHostRegistration.create(
(
_region("first", 0x1000, 4096),
_region("second", 0x9000, 4096),
),
max_bytes=None,
)
assert not isinstance(error.value, HostRegistrationCleanupError)
assert runtime.unregistered == [0x1000]
@@ -5,6 +5,7 @@
import gc
import json
import weakref
from contextlib import contextmanager
from types import SimpleNamespace
@@ -17,6 +18,7 @@ from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
from tests.helpers.runtime import get_distributed_init_method
from vllm_omni.diffusion.data import validate_dlo_host_registration_options
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
TensorBinding,
@@ -33,12 +35,14 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
DistributedLayerwiseOffloadHook,
PinnedResidentLayerGroup,
)
from vllm_omni.diffusion.offloader.host_registration import HostRegistrationCleanupError
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
from vllm_omni.diffusion.offloader.offload_plan import (
OffloadPlan,
get_offload_plan,
)
from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
from vllm_omni.host_weight_runtime import MappedHostRegion
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -703,11 +707,15 @@ class _HWRPipeline(nn.Module):
class _FakeHostWeightLease:
def __init__(self, resolution_id: str):
def __init__(self, resolution_id: str, events: list[str] | None = None):
self.closed = False
self.provenance = SimpleNamespace(resolution_id=resolution_id)
self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),)
self._events = events
def close(self):
if self._events is not None:
self._events.append("lease")
self.closed = True
@@ -757,6 +765,203 @@ def _hwr_backend():
return _HWRPipeline(), backend, carrier, lease
def test_registered_mmap_hook_bypasses_host_staging(patched_offload_runtime):
current_block = nn.Linear(2, 2, bias=False)
next_block = nn.Linear(2, 2, bias=False)
expected = torch.arange(4, dtype=torch.float32).view(2, 2)
next_block.weight.data.copy_(expected)
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
dp_group=None,
dp_size=1,
rank=0,
pin_memory=False,
rank_local_mmap=True,
)
hook.initialize_hook(current_block)
hook.registered_mmap = True
hook._stage_mmap_sources = lambda _slot: pytest.fail("registered mmap must bypass host staging") # type: ignore[method-assign]
hook.prefetch_layer(slot=0, non_blocking=True)
assert torch.equal(next_block.weight, expected)
def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_runtime):
block = nn.Linear(2, 2, bias=False)
expected = torch.arange(4, dtype=torch.float32).view(2, 2)
block.weight.data.copy_(expected)
group = PinnedResidentLayerGroup(
[block],
device=torch.device("cpu"),
copy_stream=DummyStream(),
pin_memory=False,
rank_local_mmap=True,
)
group.registered_mmap = True
group._cpu_staging_buffers.clear()
group.load()
assert torch.equal(block.weight, expected)
def test_hwr_registration_uses_transport_budget(monkeypatch: pytest.MonkeyPatch):
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
plan, _, lease = _fake_hwr_plan()
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=True,
dp_size=1,
dlo_use_allgather=False,
dlo_host_registration_limit_gib=1.5,
),
torch.device("cuda"),
host_weight_plan=plan,
)
backend._host_weight_lease = lease # type: ignore[assignment]
calls: list[tuple[object, int | None]] = []
class Registration:
total_bytes = 4096
region_count = 1
@staticmethod
def close() -> tuple[str, ...]:
return ()
def register(regions, *, device, max_bytes):
assert device == torch.device("cuda")
calls.append((regions, max_bytes))
return Registration()
monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
pinned_source = SimpleNamespace(numel=lambda: 1, is_pinned=lambda: True)
assert backend._try_register_hwr_mmap((pinned_source,)) # type: ignore[arg-type]
assert calls == [(lease.mapped_regions, int(1.5 * 1024**3))]
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
backend._release_registered_mmap()
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
def test_hwr_registration_budget_validation_is_transport_scoped():
assert (
validate_dlo_host_registration_options(
limit_gib=1.5,
enable_dlo=True,
use_allgather=False,
hwr_mode="preferred",
)
== 1.5
)
invalid = (
dict(limit_gib=-1, enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
dict(limit_gib=float("inf"), enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
dict(limit_gib=1, enable_dlo=False, use_allgather=False, hwr_mode="preferred"),
dict(limit_gib=1, enable_dlo=True, use_allgather=True, hwr_mode="preferred"),
dict(limit_gib=1, enable_dlo=True, use_allgather=False, hwr_mode="disabled"),
)
for options in invalid:
with pytest.raises(ValueError):
validate_dlo_host_registration_options(**options)
def test_unregistration_precedes_lease_close_and_retries_failures(
monkeypatch: pytest.MonkeyPatch,
patched_offload_runtime,
):
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
events: list[str] = []
responses = iter([("busy",), ()])
class Registration:
@staticmethod
def close() -> tuple[str, ...]:
events.append("unregister")
return next(responses)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=True,
dp_size=1,
dlo_use_allgather=False,
),
torch.device("cpu"),
)
lease = _FakeHostWeightLease("retry", events)
backend._host_weight_lease = lease # type: ignore[assignment]
backend._host_registration = Registration()
backend._using_rank_local_mmap = True
backend._using_registered_mmap = True
backend.enabled = True
monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
backend.disable()
assert not lease.closed
assert backend._host_registration is not None
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
backend.disable()
assert lease.closed
assert backend._host_registration is None
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
assert events == ["unregister", "unregister", "lease"]
def test_failed_unregistration_retains_lease_after_backend_is_dropped(
monkeypatch: pytest.MonkeyPatch,
patched_offload_runtime,
):
events: list[str] = []
class Registration:
@staticmethod
def close() -> tuple[str, ...]:
events.append("unregister")
return ("busy",)
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=True,
dp_size=1,
dlo_use_allgather=False,
),
torch.device("cpu"),
)
lease = _FakeHostWeightLease("retained", events)
registration = Registration()
backend._host_weight_lease = lease # type: ignore[assignment]
backend._host_registration = registration
backend._using_rank_local_mmap = True
backend.enabled = True
monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
lease_ref = weakref.ref(lease)
with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
backend.disable()
del backend
del lease
gc.collect()
retained_lease = dist_backend_module._ACTIVE_HWR_REGISTRATIONS[0][1]
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(registration, retained_lease)]
assert lease_ref() is retained_lease
assert not retained_lease.closed
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
retained_lease.close()
class _MultiBlockModel(nn.Module):
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
@@ -0,0 +1,90 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Wan AutoRound-format MXFP4 configuration and layer-mapping tests."""
from unittest.mock import patch
import pytest
from torch import nn
from vllm.model_executor.layers.linear import LinearBase
from vllm.model_executor.model_loader.utils import configure_quant_config
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import (
WanTransformer3DModel,
)
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.inc_config import OmniINCConfig
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _autoround_mxfp4_config() -> dict:
return {
"bits": 4,
"act_bits": 4,
"data_type": "mx_fp",
"act_data_type": "mx_fp",
"group_size": 32,
"act_group_size": 32,
"sym": True,
"act_sym": True,
"act_dynamic": True,
"enable_quanted_input": False,
"block_name_to_quantize": "blocks",
"quant_method": "auto-round",
"packing_format": "auto_round:llm_compressor",
"autoround_version": "0.14.2",
}
def test_build_wan_autoround_mxfp4_config():
config = build_quant_config(_autoround_mxfp4_config())
assert isinstance(config, OmniINCConfig)
assert config.get_name() == "inc"
assert config.weight_bits == 4
assert config.data_type == "mx_fp"
assert config.group_size == 32
assert config.packing_format == "auto_round:llm_compressor"
assert config.block_name_to_quantize == ["blocks"]
def test_wan_autoround_config_uses_runtime_layer_names():
config = build_quant_config(_autoround_mxfp4_config())
configure_quant_config(config, WanTransformer3DModel)
assert config.packed_modules_mapping is WanTransformer3DModel.packed_modules_mapping
assert config.block_name_to_quantize == ["blocks"]
layer = object.__new__(LinearBase)
nn.Module.__init__(layer)
for prefix in (
"blocks.0.ffn.net_0.proj",
"blocks.0.ffn.net_2",
"blocks.0.attn1.to_out",
):
layer_config = config.config_parser.resolve(layer, prefix)
assert layer_config.quantized
assert layer_config.is_mxfp4
def test_wan_autoround_scope_dispatches_only_blocks_to_mxfp4():
config = build_quant_config(_autoround_mxfp4_config())
configure_quant_config(config, WanTransformer3DModel)
layer = object.__new__(LinearBase)
nn.Module.__init__(layer)
layer_config = config.config_parser.resolve(layer, "blocks.0.ffn.net_0.proj")
outside_config = config.config_parser.resolve(layer, "condition_embedder.time_proj")
assert layer_config.quantized
assert layer_config.is_mxfp4
assert not outside_config.quantized
fake_method = object()
with patch(
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_scheme.INCMxfp4Scheme.get_linear_method",
return_value=fake_method,
):
method = config.get_quant_method(layer, "blocks.0.ffn.net_0.proj")
assert method is fake_method
+23 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for WorkerWrapperBase class.
@@ -12,6 +12,7 @@ This module tests the WorkerWrapperBase implementation:
- Dynamic worker class extension
"""
from types import SimpleNamespace
from typing import Any
import pytest
@@ -204,6 +205,27 @@ class TestWorkerWrapperBaseDelegation:
wrapper.worker.shutdown.assert_called_once()
assert result is None
def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture):
events: list[str] = []
offload_backend = mocker.Mock()
offload_backend.disable.side_effect = lambda: events.append("offload")
kv_manager = mocker.Mock()
kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv")
destroy = mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
side_effect=lambda: events.append("distributed"),
)
worker = DiffusionWorker.__new__(DiffusionWorker)
worker.model_runner = SimpleNamespace(
offload_backend=offload_backend,
kv_transfer_manager=kv_manager,
)
worker.shutdown()
assert events == ["offload", "kv", "distributed"]
destroy.assert_called_once_with()
# -------------------------------------------------------------------------
# Tests: execute_method
+81 -3
View File
@@ -1,3 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import base64
import wave
from urllib.parse import parse_qs, urlsplit
@@ -20,7 +24,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_realtime_client_builds_explicit_native_duplex_url():
url = build_realtime_url(
"ws://localhost:8099/v1/realtime?custom=1",
"ws://localhost:8099/v1/realtime?custom=1&duplex=0&model=stale&minicpmo45_native_duplex=0&session_id=stale",
"openbmb/MiniCPM-o-4_5",
session_id="session-a",
)
@@ -50,7 +54,7 @@ def test_seed_tts_initial_text_is_part_of_native_duplex_context():
def test_realtime_client_builds_resume_only_url_when_autostart_disabled():
url = build_realtime_url(
"ws://localhost:8099/v1/realtime?duplex=1",
"ws://localhost:8099/v1/realtime?duplex=1&autostart=1",
"openbmb/MiniCPM-o-4_5",
autostart=False,
)
@@ -73,10 +77,15 @@ async def test_realtime_client_configure_omits_ref_audio_by_default():
client = Client()
await client.configure("openbmb/MiniCPM-o-4_5", timeout_s=1)
await client.configure(
"openbmb/MiniCPM-o-4_5",
idle_timeout_s=900,
timeout_s=1,
)
session = client.sent[0]["session"]
assert "ref_audio" not in session
assert session["idle_timeout_s"] == 900
@pytest.mark.asyncio
@@ -100,6 +109,74 @@ async def test_realtime_client_configure_sends_explicit_ref_audio():
session = client.sent[0]["session"]
assert session["ref_audio"] == "data:audio/wav;base64,AAAA"
assert "idle_timeout_s" not in session
@pytest.mark.asyncio
async def test_realtime_client_reports_reader_exit_to_event_waiters():
client = RealtimeDuplexClient("ws://unused")
async def stopped():
return None
client._reader_task = asyncio.create_task(stopped())
await client._reader_task
with pytest.raises(ConnectionError, match="closed before"):
client.raise_if_reader_stopped()
@pytest.mark.asyncio
async def test_realtime_client_configure_reports_reader_failure_immediately():
class Client(RealtimeDuplexClient):
async def send(self, event):
async def fail():
raise RuntimeError("reader failed")
self._reader_task = asyncio.create_task(fail())
await asyncio.sleep(0)
client = Client("ws://unused")
with pytest.raises(ConnectionError, match="reader failed"):
await client.configure("model", timeout_s=10.0)
@pytest.mark.asyncio
async def test_realtime_client_close_reports_reader_cancellation_immediately():
class Client(RealtimeDuplexClient):
async def send(self, event):
async def wait_forever():
await asyncio.Future()
self._reader_task = asyncio.create_task(wait_forever())
self._reader_task.cancel()
with pytest.raises(asyncio.CancelledError):
await self._reader_task
client = Client("ws://unused")
with pytest.raises(ConnectionError, match="reader was cancelled"):
await client.close_session(timeout_s=10.0)
@pytest.mark.asyncio
async def test_realtime_client_close_waits_for_a_new_session_closed_event():
class Client(RealtimeDuplexClient):
async def send(self, event):
assert event["type"] == "session.close"
async def acknowledge_close():
await asyncio.sleep(0.02)
self.events.add({"type": "session.closed"})
asyncio.create_task(acknowledge_close())
client = Client("ws://unused")
client.events.add({"type": "session.closed", "reason": "old"})
started_at = asyncio.get_running_loop().time()
await client.close_session(timeout_s=0.2)
assert asyncio.get_running_loop().time() - started_at >= 0.01
@pytest.mark.asyncio
@@ -235,6 +312,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals():
"output_token_count": 4,
"ttft_ms": 120.0,
"tpot_ms": 15.0,
"itls_ms": [10.0, 14.0, 18.0],
"inter_token_interval_ms": {
"count": 3,
"mean": 14.0,
@@ -0,0 +1,51 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""E2E smoke test for Wan2.1 AutoRound MXFP4 inference."""
import os
import numpy as np
import pytest
import torch
from tests.helpers.mark import hardware_test
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
QUANTIZED_MODEL = os.environ.get(
"WAN21_AUTOROUND_MXFP4_MODEL",
"INCModel/Wan2.1-T2V-1.3B-Diffusers-MXFP4-AutoRound",
)
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
@hardware_test(res={"xpu": "B60"})
@pytest.mark.parametrize(
"omni_runner",
[(QUANTIZED_MODEL, None)],
indirect=True,
)
def test_wan21_autoround_mxfp4_generates_video(omni_runner, omni_runner_handler):
sampling_params = OmniDiffusionSamplingParams(
height=256,
width=256,
num_frames=5,
num_inference_steps=1,
guidance_scale=1.0,
generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
)
response = omni_runner_handler.send_diffusion_request(
{
"prompt": "A red toy car driving on a table",
"sampling_params": sampling_params,
}
)
assert response.success, response.error_message
assert response.images is not None and len(response.images) == 1
frames = np.asarray(response.images[0])
assert frames.shape[1:4] == (5, 256, 256)
assert frames.std() > 0.01
@@ -58,18 +58,14 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
read_pcm16_wav,
summarize_session_request_metrics,
)
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
reference_audio_data_url as _ref_audio_data_url,
)
_url_with_model = build_realtime_url
_read_wav_pcm16 = read_pcm16_wav
def _ref_audio_data_url(path: str | None) -> str | None:
if path is None:
return None
ref_path = Path(path).expanduser()
return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii")
@dataclass
class DemoState:
events: list[dict[str, object]] = field(default_factory=list)
@@ -7,7 +7,6 @@ E2E Online tests for Qwen3-Omni model.
import os
import pytest
from openai import BadRequestError
from tests.helpers.mark import hardware_test
from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -89,7 +88,7 @@ test_params = [
),
]
test_token_params = [
reduced_token_params = [
pytest.param(
OmniServerParams(
model=model,
@@ -97,7 +96,7 @@ test_token_params = [
use_stage_cli=True,
server_args=["--async-chunk"],
),
id="batch_token_64",
id="batch_token_2048",
)
]
@@ -138,29 +137,7 @@ def get_max_batch_size(size_type="few"):
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_text_to_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text
Output Modal: audio
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
request_config = {
"model": omni_server.model,
"messages": messages,
"modalities": ["audio"],
"stream": True,
"key_words": {"text": ["beijing"]},
}
online_client.send_omni_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
def test_text_to_text_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text
@@ -204,36 +181,13 @@ def test_text_video_to_text_001(omni_server, online_client) -> None:
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
def test_text_audio_to_text_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text, audio
Output Modal: text, audio
Input Setting: stream=False
Datasets: single request
"""
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
messages = dummy_messages_from_mix_data(
audio_data_url=audio_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_audio")
)
request_config = {
"model": omni_server.model,
"messages": messages,
"key_words": {"audio": AUDIO_KEY},
}
online_client.send_omni_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
"""
Input Modal: text, long-duration audio (~LONG_AUDIO_DURATION_SEC s WAV)
Output Modal: text, audio
Input Setting: stream=False
Datasets: single request
Datasets: few requests
"""
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(LONG_AUDIO_DURATION_SEC, 1)['base64']}"
messages = dummy_messages_from_mix_data(
@@ -252,37 +206,13 @@ def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
def test_text_image_to_text_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text, image
Output Modal: text, audio
Input Setting: stream=False
Datasets: single request
"""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
messages = dummy_messages_from_mix_data(
image_data_url=image_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_image")
)
request_config = {
"model": omni_server.model,
"messages": messages,
"key_words": {"image": IMAGE_KEY},
}
online_client.send_omni_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text, high-resolution image (1080p-class JPEG)
Output Modal: text, audio
Input Setting: stream=False
Datasets: single request
Datasets: few requests
"""
image_data_url = (
f"data:image/jpeg;base64,{generate_synthetic_image(LARGE_IMAGE_WIDTH, LARGE_IMAGE_HEIGHT)['base64']}"
@@ -304,7 +234,7 @@ def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text, video
@@ -328,36 +258,6 @@ def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
online_client.send_omni_request(request_config)
@pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
def test_mix_to_text_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text, audio, image, video
Output Modal: text, audio
Input Setting: stream=True
Datasets: few requests
"""
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
video_data_url=video_data_url,
image_data_url=image_data_url,
audio_data_url=audio_data_url,
content_text=get_prompt("mix"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
}
online_client.send_omni_request(request_config, request_num=get_max_batch_size())
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_audio_in_video_001(omni_server, online_client) -> None:
@@ -507,22 +407,22 @@ def test_one_word_prompt_001(omni_server, online_client) -> None:
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_speaker_001(omni_server, online_client) -> None:
"""
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
Input Modal: Chinese text.
Output Modal: Chinese text and audio using an uppercase preset speaker name.
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text=get_prompt("text"),
content_text=get_prompt("text_chinese"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"speaker": "Chelsie",
"key_words": {"text": ["beijing"]},
"speaker": "CHELSIE",
"key_words": {"text": ["北京"]},
}
online_client.send_omni_request(request_config)
@@ -563,55 +463,6 @@ def test_speaker_002(omni_server, online_client) -> None:
print(f"Gender assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_speaker_003(omni_server, online_client) -> None:
"""
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text=get_prompt("text"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"speaker": "CHELSIE",
"key_words": {"text": ["beijing"]},
}
online_client.send_omni_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_language_001(omni_server, online_client) -> None:
"""
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text=get_prompt("text_chinese"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"key_words": {"text": ["北京"]},
}
online_client.send_omni_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
@@ -619,7 +470,7 @@ def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
Input Modal: text only (long-form generation prompt).
Output Modal: text, audio (default ``modalities``);
Input Setting: stream=True
Datasets: single request
Datasets: few requests
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
@@ -631,21 +482,3 @@ def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
text = responses[0].text_content if responses else ""
word_count = len(text.split())
assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True)
def test_invalid_audio_format_rejected(omni_server, openai_client) -> None:
"""Ensure invalid audio format is rejected with 400 before streaming starts."""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text=get_prompt(),
)
with pytest.raises(BadRequestError, match="audio format"):
openai_client.client.chat.completions.create(
model=omni_server.model,
messages=messages,
modalities=["text", "audio"],
audio={"voice": "alloy", "format": "aac"},
stream=True,
)
@@ -8,9 +8,8 @@ and are supported by the following text-to-image models:
- Qwen-Image-2512
One feature per test case, matching the Test Plan in PR #1682 (Qwen-Image-Edit).
Supported features for Qwen-Image series: TeaCache, Cache-DiT, Ulysses-SP, Ring-Attention,
CFG-Parallel, Tensor-Parallel, VAE-Patch-Parallel, CPU offload (layerwise).
See docs/user_guide/diffusion_acceleration.md.
Nightly covers each feature once, alternating models (5× Qwen-Image, 4× Qwen-Image-2512).
Ulysses stays on Qwen-Image. See docs/user_guide/diffusion_acceleration.md.
"""
import pytest
@@ -25,141 +24,116 @@ NEGATIVE_PROMPT = "blurry, low quality"
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
MODEL_IMAGE = "Qwen/Qwen-Image"
MODEL_2512 = "Qwen/Qwen-Image-2512"
def _get_diffusion_feature_cases(model: str):
return [
pytest.param(
OmniServerParams(
model=model,
server_args=["--enable-cpu-offload"],
),
id="cpu_offload",
marks=SINGLE_CARD_FEATURE_MARKS,
# One server per feature. Alternate models; keep feature pytest ids unchanged.
FEATURE_CASES = [
pytest.param(
OmniServerParams(
model=MODEL_IMAGE,
server_args=["--enable-cpu-offload"],
),
pytest.param(
OmniServerParams(model=model, server_args=["--step-execution"]),
id="step_execution",
marks=SINGLE_CARD_FEATURE_MARKS,
id="cpu_offload",
marks=SINGLE_CARD_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(model=MODEL_2512, server_args=["--step-execution"]),
id="step_execution",
marks=SINGLE_CARD_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(model=MODEL_IMAGE, server_args=["--cache-backend", "tea_cache"]),
id="cache_tea_cache",
marks=SINGLE_CARD_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL_2512,
server_args=[
"--cache-backend",
"cache_dit",
"--enable-layerwise-offload",
],
),
pytest.param(
OmniServerParams(model=model, server_args=["--cache-backend", "tea_cache"]),
id="cache_tea_cache",
marks=SINGLE_CARD_FEATURE_MARKS,
id="layerwise_offload",
marks=SINGLE_CARD_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL_IMAGE,
server_args=[
"--cache-backend",
"cache_dit",
"--ulysses-degree",
"2",
],
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--cache-backend",
"cache_dit",
"--enable-layerwise-offload",
],
),
id="layerwise_offload",
marks=SINGLE_CARD_FEATURE_MARKS,
id="ulysses_2",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL_2512,
server_args=[
"--cache-backend",
"cache_dit",
"--ring",
"2",
],
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--cache-backend",
"cache_dit",
"--ulysses-degree",
"2",
],
),
id="ulysses_2",
marks=PARALLEL_FEATURE_MARKS,
id="ring_2",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL_IMAGE,
server_args=[
"--cache-backend",
"tea_cache",
"--cfg-parallel-size",
"2",
],
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--cache-backend",
"cache_dit",
"--ring",
"2",
],
),
id="ring_2",
marks=PARALLEL_FEATURE_MARKS,
id="cfg_parallel_2",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL_2512,
server_args=[
"--cache-backend",
"cache_dit",
"--tensor-parallel-size",
"2",
"--vae-patch-parallel-size",
"2",
"--vae-use-tiling",
"--diffusion-quantization-config",
'{"method":"fp8"}',
],
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--cache-backend",
"tea_cache",
"--cfg-parallel-size",
"2",
],
),
id="cfg_parallel_2",
marks=PARALLEL_FEATURE_MARKS,
id="vae_patch_parallel_2",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL_IMAGE,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--cache-backend",
"cache_dit",
"--tensor-parallel-size",
"2",
"--vae-patch-parallel-size",
"2",
"--vae-use-tiling",
"--diffusion-quantization-config",
'{"method":"fp8"}',
],
),
id="vae_patch_parallel_2",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
id="parallel_hsdp",
marks=PARALLEL_FEATURE_MARKS,
),
]
id="parallel_hsdp",
marks=PARALLEL_FEATURE_MARKS,
),
]
@pytest.mark.parametrize(
"omni_server",
_get_diffusion_feature_cases("Qwen/Qwen-Image"),
indirect=True,
)
@pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True)
def test_qwen_image(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test each diffusion feature with Qwen-Image (text-to-image), one feature per case."""
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
request_config = {
"model": omni_server.model,
"messages": messages,
"extra_body": {
"height": 512,
"width": 512,
"num_inference_steps": 2,
"negative_prompt": NEGATIVE_PROMPT,
"true_cfg_scale": 4.0,
"seed": 42,
},
}
online_client.send_diffusion_request(request_config)
@pytest.mark.parametrize(
"omni_server",
_get_diffusion_feature_cases("Qwen/Qwen-Image-2512"),
indirect=True,
)
def test_qwen_image_2512(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test each diffusion feature with Qwen-Image-2512 (text-to-image), one feature per case."""
"""One diffusion feature per case; model is chosen in FEATURE_CASES."""
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
request_config = {
"model": omni_server.model,
@@ -0,0 +1,71 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E online serving test for Stable Audio Open text-to-audio diffusion.
Stable Audio Open is served through the OpenAI-compatible
`POST /v1/audio/generate` endpoint (JSON in, binary WAV out).
"""
import os
import pytest
from tests.helpers import skip_if_gated_repo_inaccessible
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
# Stable Audio Open is gated; override to a local bundle locally if needed.
STABLE_AUDIO_TEST_MODEL = os.environ.get("STABLE_AUDIO_TEST_MODEL", "stabilityai/stable-audio-open-1.0")
T2A_PROMPT = "A piano playing a gentle melody with soft room ambience."
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
def _stable_audio_server_cases(model: str):
return [
pytest.param(
OmniServerParams(model=model),
id="t2a",
marks=SINGLE_CARD_FEATURE_MARKS,
),
]
@pytest.fixture
def _require_stable_audio_access() -> None:
"""Skip cleanly (before the server boots) if the gated checkpoint is inaccessible."""
skip_if_gated_repo_inaccessible(STABLE_AUDIO_TEST_MODEL)
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _stable_audio_server_cases(STABLE_AUDIO_TEST_MODEL), indirect=True)
def test_stable_audio_t2a_online(
_require_stable_audio_access: None,
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
) -> None:
"""Stable Audio Open text-to-audio: `/v1/audio/generate` returns a non-empty WAV.
Uses tiny steps / short duration to keep CI light, matching the offline smoke
test in `tests/e2e/offline_inference/test_stable_audio_expansion.py`.
"""
request_config = {
"json": {
"model": omni_server.model,
"input": T2A_PROMPT,
"audio_length": 2.0,
"num_inference_steps": 4,
"guidance_scale": 7.0,
"negative_prompt": "Low quality.",
"seed": 42,
"response_format": "wav",
},
"timeout": 300,
}
responses = openai_client.send_audio_generate_http_request(request_config)
assert responses, "no response from /v1/audio/generate"
resp = responses[0]
assert resp.success, f"audio generate failed: {resp.status_code} {resp.error_message}"
assert resp.status_code == 200
@@ -474,7 +474,7 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
assert os.environ[runtime_env_var] == "parent-value"
def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypatch):
def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_seqs(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.engine.stage_init_utils as init_mod
@@ -506,7 +506,7 @@ def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypa
use_inline=True,
)
assert od_config.max_num_seqs == 4
assert od_config.max_num_seqs == 1
assert captured == {
"model": "dummy-model",
"config": od_config,
@@ -517,7 +517,7 @@ def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypa
}
def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch):
def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_size(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
import vllm_omni.engine.stage_engine_startup as startup_mod
@@ -565,10 +565,84 @@ def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch
)
assert result is sentinel_client
assert od_config.max_num_seqs == 4
assert od_config.max_num_seqs == 1
assert resources.manager is proc_manager
def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.engine.stage_init_utils as init_mod
od_config = types.SimpleNamespace(max_num_seqs=8, step_execution=True)
monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
monkeypatch.setattr(client_mod, "create_diffusion_client", lambda *args: object())
metadata = _make_diffusion_metadata(0)
init_mod.initialize_diffusion_stage(
0,
"dummy-model",
types.SimpleNamespace(),
metadata,
stage_init_timeout=12,
batch_size=4,
use_inline=True,
)
assert od_config.max_num_seqs == 8
def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
import vllm_omni.engine.stage_engine_startup as startup_mod
od_config = types.SimpleNamespace(
max_num_seqs=8,
step_execution=True,
parallel_config=types.SimpleNamespace(world_size=1),
)
monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
monkeypatch.setattr(
startup_mod,
"register_stage_with_omni_master",
lambda **kwargs: types.SimpleNamespace(
handshake_address="tcp://127.0.0.1:26001",
input_address="tcp://127.0.0.1:26002",
output_address="tcp://127.0.0.1:26003",
),
)
omni_master_server = types.SimpleNamespace(
address="127.0.0.1",
port=25000,
release_route_port_reservations=lambda *args, **kwargs: None,
)
proc_manager = types.SimpleNamespace(
addresses=types.SimpleNamespace(
inputs=["tcp://127.0.0.1:26002"],
outputs=["tcp://127.0.0.1:26003"],
)
)
monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager)
monkeypatch.setattr(
client_mod.StageDiffusionClient,
"from_addresses",
lambda metadata, **kwargs: object(),
)
startup_mod.launch_diffusion_stage_replica(
model="dummy-model",
stage_config=types.SimpleNamespace(),
metadata=types.SimpleNamespace(stage_id=0),
stage_init_timeout=12,
batch_size=4,
use_inline=False,
omni_master_server=omni_master_server,
)
assert od_config.max_num_seqs == 8
def test_stage_runtime_initializes_stage_pools(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for error propagation paths within the Orchestrator.
Covers:
@@ -16,7 +16,9 @@ from types import SimpleNamespace
import janus
import pytest
from vllm.sampling_params import SamplingParams
from vllm.v1.engine.exceptions import EngineDeadError
from vllm.v1.serial_utils import MsgpackEncoder
from vllm_omni.engine.messages import (
AddCompanionRequestMessage,
@@ -46,8 +48,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _sampling_params(max_tokens: int = 4):
from vllm.sampling_params import SamplingParams
return SamplingParams(max_tokens=max_tokens)
@@ -188,6 +188,45 @@ class _UnavailableDiffusionStageClient(FakeStageClient):
raise StageUnavailableError(f"stage {self.stage_id} has no live replica")
class _OverflowOnEncodeStageClient(FakeStageClient):
"""Msgpack-encodes the request's sampling params on dispatch, as the real
StageEngineCoreClient does. We use this to make sure overflow sampling params
do not break the orchestrator."""
async def add_request_async(self, request, *args, **kwargs) -> None:
MsgpackEncoder().encode(request.sampling_params)
await super().add_request_async(request, *args, **kwargs) # unreachable
@pytest.mark.asyncio
async def test_encode_overflow_should_fail_request_not_loop(orchestrator_factory) -> None:
"""Regression test for over-range seed handling."""
stage0 = _OverflowOnEncodeStageClient(stage_type="llm", final_output=True)
orchestrator_fixture = orchestrator_factory([stage0])
# seed > 2**63-1 cannot be msgpack-serialized -> the malicious input under test.
overflow_params = SamplingParams(seed=2**70)
await _enqueue_add_request(
orchestrator_fixture,
request_id="req-overflow",
prompt=SimpleNamespace(
request_id="req-overflow",
prompt_token_ids=[1, 2],
sampling_params=overflow_params,
),
original_prompt={"prompt": "hi"},
sampling_params_list=[overflow_params],
final_stage_id=0,
)
error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-overflow")
assert error_msg.fatal is False
assert error_msg.status_code == 400
assert orchestrator_fixture.thread.is_alive()
assert "req-overflow" not in orchestrator_fixture.orchestrator.request_states
@pytest.mark.asyncio
async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None:
"""Multi-replica stage: one replica dying must fail only the requests bound
+1 -1
View File
@@ -1175,7 +1175,7 @@ class TestSingleStageReplicaInitialization:
os.environ[device_env_var] = prev_device_env
assert result is sentinel_client
assert od_config.max_num_seqs == 4
assert od_config.max_num_seqs is None
mock_register.assert_called_once_with(
omni_master_address="127.0.0.1",
omni_master_port=25000,
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# tests/entrypoints/openai/test_serving_speech.py
import asyncio
import base64
@@ -4147,16 +4150,15 @@ class TestCosyVoice3Serving:
error = cosyvoice3_server._validate_tts_request(request)
assert error is None
def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server):
request = OpenAICreateSpeechRequest(
input="Hello",
ref_audio="data:audio/wav;base64,abc",
ref_text="hello",
max_new_tokens=0,
)
error = cosyvoice3_server._validate_tts_request(request)
assert error is not None
assert "max_new_tokens" in error
def test_validate_cosyvoice3_max_new_tokens_range(self):
"""Ensure max_new_tokens below the minimum is rejected during request validation."""
with pytest.raises(ValidationError, match="max_new_tokens"):
OpenAICreateSpeechRequest(
input="Hello",
ref_audio="data:audio/wav;base64,abc",
ref_text="hello",
max_new_tokens=0,
)
@pytest.mark.parametrize(
("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),
@@ -378,6 +378,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
"preferred",
"--host-weight-runtime-root",
"/var/cache/vllm-omni/hwr",
"--dlo-host-registration-limit-gib",
"80",
]
)
@@ -387,8 +389,10 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80
assert engine_args["host_weight_runtime_mode"] == "preferred"
assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
assert engine_args["dlo_host_registration_limit_gib"] == 80
def test_serve_cli_accepts_diffusion_compile_controls():
@@ -57,7 +57,9 @@ from vllm_omni.host_weight_runtime.filesystem import (
detect_storage_class,
)
from vllm_omni.host_weight_runtime.filesystem import store as filesystem_store_module
from vllm_omni.host_weight_runtime.filesystem import writer as filesystem_writer_module
from vllm_omni.host_weight_runtime.filesystem.locks import FileLock
from vllm_omni.host_weight_runtime.filesystem.writer import ArtifactWriterError, FilesystemArtifactWriter
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -110,6 +112,150 @@ def _open_fd_targets() -> tuple[str, ...]:
return tuple(targets)
def _direct_writer(tmp_path: Path) -> FilesystemArtifactWriter:
staging = tmp_path / "staging"
staging.mkdir()
return FilesystemArtifactWriter(staging, capacity_check=lambda _projected, _additional: None)
def test_ordered_writer_hashes_payload_during_writes(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
writer = _direct_writer(tmp_path)
weight = _weight()
scale = torch.tensor([1.0, 2.0])
def fail_rescan(_path: Path, _specs: tuple[TensorWriteSpec, ...]) -> object:
raise AssertionError("ordered payload must not be rescanned")
monkeypatch.setattr(filesystem_writer_module, "_scan_payload", fail_rescan)
with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
# The writer's ordered contract follows canonical storage-key order.
output.write_tensor("layer.scale", scale)
output.write_rows("layer.weight", 0, weight[:2])
output.write_rows("layer.weight", 2, weight[2:])
manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
payload = writer.staging_dir / "weights.safetensors"
assert manifest.files[0].sha256 == hashlib.sha256(payload.read_bytes()).hexdigest()
expected_tensor_digests = {
"layer.scale": hashlib.sha256(memoryview(scale.view(torch.uint8).numpy())).hexdigest(),
"layer.weight": hashlib.sha256(memoryview(weight.contiguous().view(torch.uint8).numpy())).hexdigest(),
}
assert {entry.name: entry.sha256 for entry in manifest.tensors} == expected_tensor_digests
def test_ordered_writer_rejects_out_of_order_payload(tmp_path: Path) -> None:
writer = _direct_writer(tmp_path)
output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
with pytest.raises(ArtifactWriterError, match="ordered tensor payload expected offset"):
output.write_tensor("layer.weight", _weight())
output.abort()
assert not list(writer.staging_dir.iterdir())
def test_payload_fsync_overlaps_producer_work(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
writer = _direct_writer(tmp_path)
started = threading.Event()
release = threading.Event()
closed = threading.Event()
original_fsync = filesystem_writer_module._fsync_payload
def blocking_fsync(fd: int) -> None:
started.set()
if not release.wait(timeout=5):
raise TimeoutError("test did not release payload fsync")
original_fsync(fd)
monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", blocking_fsync)
output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
output.write_tensor("layer.weight", _weight())
close_errors: list[BaseException] = []
def close_output() -> None:
try:
output.close()
except BaseException as exc:
close_errors.append(exc)
finally:
closed.set()
close_thread = threading.Thread(target=close_output)
close_thread.start()
try:
assert started.wait(timeout=2)
assert closed.wait(timeout=1), "payload close waited for fsync instead of overlapping it"
finally:
release.set()
close_thread.join(timeout=5)
assert not close_errors
writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
def test_payload_fsync_failure_prevents_publication(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
writer = _direct_writer(tmp_path)
def failing_fsync(fd: int) -> None:
try:
raise OSError(errno.EIO, "injected fsync failure")
finally:
os.close(fd)
monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", failing_fsync)
with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
output.write_tensor("layer.weight", _weight())
with pytest.raises(OSError, match="injected fsync failure"):
writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
assert not (writer.staging_dir / "manifest.json").exists()
assert not (writer.staging_dir / "READY.json").exists()
def test_unordered_payload_fallback_hashes_files_in_parallel(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
writer = _direct_writer(tmp_path)
specs = (
TensorWriteSpec("a", (2,), torch.float32),
TensorWriteSpec("b", (2,), torch.float32),
)
for name, spec in zip(("a.safetensors", "b.safetensors"), specs, strict=True):
with writer.open_tensor_file(name, (spec,)) as output:
output.write_tensor(spec.name, torch.tensor([1.0, 2.0]))
original_scan = filesystem_writer_module._scan_payload
barrier = threading.Barrier(2)
thread_ids: set[int] = set()
thread_ids_lock = threading.Lock()
def synchronized_scan(path: Path, file_specs: tuple[TensorWriteSpec, ...]) -> object:
with thread_ids_lock:
thread_ids.add(threading.get_ident())
barrier.wait(timeout=5)
return original_scan(path, file_specs)
monkeypatch.setattr(filesystem_writer_module, "_scan_payload", synchronized_scan)
manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
assert len(manifest.files) == 2
assert len(thread_ids) == 2
class FakeProducer:
def __init__(
self,
@@ -1026,7 +1172,9 @@ def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path:
free_root = tmp_path / "free-limit"
free_root.mkdir(mode=0o700)
unavailable_reserve = shutil.disk_usage(free_root).free + 1
# Leave enough margin for filesystem accounting to move between the
# preflight measurement and the capacity check.
unavailable_reserve = shutil.disk_usage(free_root).free + (1 << 30)
free_limited = _make_store(free_root, capacity=CapacityPolicy(min_free_bytes=unavailable_reserve))
free_result = free_limited.get_or_build(
BuildRequest(identity),
@@ -0,0 +1,306 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""OmniInteract dataset discovery for the standard serving benchmark."""
from __future__ import annotations
import json
import logging
import os
import random
import shutil
import tarfile
import tempfile
from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from vllm.benchmarks.datasets import BenchmarkDataset, SampleRequest
from vllm.tokenizers import TokenizerLike
from vllm.transformers_utils.repo_utils import hf_fs
OMNIINTERACT_SUBSETS = ("1q1a", "1q1a_math", "1qna")
DEFAULT_OMNIINTERACT_REPO = "lucky-lance/OmniInteract"
DEFAULT_MAX_VIDEO_DURATION_S = 3600.0
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class OmniInteractCase:
subset: str
video_rel: str
video_path: Path
annotation_path: Path
scene_type: str
@dataclass(frozen=True)
class OmniInteractSessionOptions:
output_root: Path
timeout_s: float
media_timeout_s: float
ref_audio: str
require_response: bool = False
max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S
@dataclass(frozen=True)
class OmniInteractPreparedInput:
duration_s: float
pcm16: bytes = field(repr=False)
video_frames: tuple[str | None, ...] = field(repr=False)
ref_audio_data_url: str = field(repr=False)
@dataclass
class OmniInteractSampleRequest(SampleRequest):
omniinteract_case: OmniInteractCase | None = None
omniinteract_options: OmniInteractSessionOptions | None = None
omniinteract_prepared_input: OmniInteractPreparedInput | None = field(default=None, repr=False)
class OmniInteractDataset(BenchmarkDataset):
def __init__(
self,
*,
data_root: str | None,
dataset_repo: str = DEFAULT_OMNIINTERACT_REPO,
subsets: Sequence[str] = OMNIINTERACT_SUBSETS,
random_seed: int = 0,
disable_shuffle: bool = False,
) -> None:
super().__init__(
dataset_path=data_root or dataset_repo,
random_seed=random_seed,
disable_shuffle=disable_shuffle,
)
self.root = resolve_omniinteract_root(data_root, dataset_repo)
self.subsets = tuple(subsets)
def sample(
self,
tokenizer: TokenizerLike | None,
num_requests: int,
*,
request_id_prefix: str = "",
options: OmniInteractSessionOptions,
**_: Any,
) -> list[SampleRequest]:
del tokenizer
cases = discover_omniinteract_cases(
self.root,
self.subsets,
num_prompts=num_requests,
seed=self.random_seed,
disable_shuffle=self.disable_shuffle,
)
return [
OmniInteractSampleRequest(
prompt="",
prompt_len=0,
expected_output_len=0,
multi_modal_data=None,
request_id=f"{request_id_prefix}{index}",
omniinteract_case=case,
omniinteract_options=options,
)
for index, case in enumerate(cases)
]
def _data_dir(root: Path) -> Path:
for candidate in (root, root / "data"):
if any((candidate / subset).is_dir() for subset in OMNIINTERACT_SUBSETS):
return candidate
raise FileNotFoundError(f"OmniInteract data not found under {root}")
def _confined_path(root: Path, relative: object, *, field: str) -> Path:
text, resolved_root = str(relative or ""), root.resolve()
path = Path(text)
destination = (root / path).resolve()
if not text or path.is_absolute() or ".." in path.parts or not destination.is_relative_to(resolved_root):
raise ValueError(f"Unsafe OmniInteract {field} path: {text!r}")
return destination
def _safe_extract(archive: tarfile.TarFile, target: Path) -> None:
root = target.resolve()
members = archive.getmembers()
for member in members:
path = Path(member.name)
if (
path.is_absolute()
or ".." in path.parts
or not (root / path).resolve().is_relative_to(root)
or not (member.isdir() or member.isfile())
):
raise ValueError(f"Unsafe path in OmniInteract archive: {member.name!r}")
archive.extractall(target, members=members, filter="data")
def _extract_archive(archive: Path, target: Path) -> Path:
if target.is_symlink():
raise ValueError(f"Refusing to extract through symlink: {target}")
stat = archive.stat()
published = target / f"{stat.st_size}-{stat.st_mtime_ns}"
if published.is_symlink():
raise ValueError(f"Refusing to use symlinked extraction: {published}")
if published.is_dir():
try:
return _data_dir(published)
except FileNotFoundError:
pass
target.mkdir(parents=True, exist_ok=True)
staging = Path(tempfile.mkdtemp(prefix=".tmp-", dir=target))
try:
with tarfile.open(archive, "r:*") as handle:
_safe_extract(handle, staging)
root = _data_dir(staging)
try:
staging.replace(published)
return published / root.relative_to(staging)
except OSError:
if not published.is_dir():
raise
return _data_dir(published)
finally:
shutil.rmtree(staging, ignore_errors=True)
def resolve_omniinteract_root(data_root: str | None, dataset_repo: str = DEFAULT_OMNIINTERACT_REPO) -> Path:
"""Resolve a local tree/archive or download the official archive."""
if data_root:
local = Path(data_root).expanduser().resolve()
if local.is_file():
return _extract_archive(local, local.parent / f".{local.stem}.vllm_omni_extracted")
if not local.is_dir():
raise FileNotFoundError(f"--dataset-path does not exist: {local}")
try:
return _data_dir(local)
except FileNotFoundError:
for name in ("data.tar.gz", "data.tar"):
if (archive := local / name).is_file():
return _extract_archive(archive, local / ".vllm_omni_extracted")
raise
cache = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface"))
target = cache / "vllm_omni" / "omniinteract" / dataset_repo.replace("/", "__")
downloads = target / ".downloads"
downloads.mkdir(parents=True, exist_ok=True)
filesystem, errors = hf_fs(), []
for name in ("data.tar.gz", "data.tar"):
archive = downloads / name
try:
if not archive.is_file():
descriptor, temporary_name = tempfile.mkstemp(prefix=f".{name}.", dir=downloads)
os.close(descriptor)
temporary = Path(temporary_name)
try:
filesystem.get_file(f"datasets/{dataset_repo}/{name}", str(temporary))
if not temporary.stat().st_size:
raise ValueError("downloaded archive is empty")
temporary.replace(archive)
finally:
temporary.unlink(missing_ok=True)
return _extract_archive(archive, target)
except Exception as exc: # noqa: BLE001, PERF203 - Hub errors vary by version
errors.append(f"{name}: {exc}")
raise FileNotFoundError(f"Could not download OmniInteract from {dataset_repo!r}: {'; '.join(errors)}")
def _mapping_cases(root: Path, subset: str) -> list[OmniInteractCase]:
mapping_path = root / "video_json_map.json"
if not mapping_path.is_file():
raise FileNotFoundError(f"Missing OmniInteract mapping: {mapping_path}")
mapping = json.loads(mapping_path.read_text())
entries = mapping.get("entries") if isinstance(mapping, dict) else None
if not isinstance(entries, list):
raise ValueError(f"Invalid OmniInteract mapping: {mapping_path}")
cases = []
for row in entries:
if not isinstance(row, dict):
raise ValueError(f"Invalid OmniInteract mapping row in {mapping_path}")
video_rel, annotation_rel = str(row.get("video") or ""), str(row.get("annotation") or "")
video = _confined_path(root, video_rel, field="video")
annotation = _confined_path(root, annotation_rel, field="annotation")
if not video.is_file() or not annotation.is_file():
raise FileNotFoundError(f"OmniInteract mapping references missing files: {video_rel!r}, {annotation_rel!r}")
cases.append(
OmniInteractCase(
subset,
video_rel,
video,
annotation,
str(row.get("scene_type") or "multi_turn").lower(),
)
)
return cases
def _one_to_many_cases(root: Path) -> list[OmniInteractCase]:
videos, annotations = root / "videos_bench", root / "annotations"
if not videos.is_dir() or not annotations.is_dir():
raise FileNotFoundError(f"Invalid OmniInteract 1qna layout under {root}")
cases = []
for video in sorted(videos.rglob("*.mp4")):
if not video.resolve().is_relative_to(videos.resolve()):
raise ValueError(f"Unsafe OmniInteract video path: {video}")
relative = video.relative_to(videos)
annotation = (annotations / relative).with_suffix(".json").resolve()
if not annotation.is_relative_to(annotations.resolve()):
raise ValueError(f"Unsafe OmniInteract annotation path: {relative}")
if not annotation.is_file():
raise FileNotFoundError(f"Missing OmniInteract annotation: {annotation}")
cases.append(OmniInteractCase("1qna", f"videos_bench/{relative.as_posix()}", video, annotation, "1qna"))
return cases
def discover_omniinteract_cases(
root: Path,
subsets: Sequence[str],
*,
num_prompts: int,
seed: int = 0,
disable_shuffle: bool = False,
) -> list[OmniInteractCase]:
invalid = set(subsets) - set(OMNIINTERACT_SUBSETS)
if invalid:
raise ValueError(f"Unsupported OmniInteract subsets: {sorted(invalid)}")
if not subsets:
raise ValueError("At least one OmniInteract subset is required")
if len(set(subsets)) != len(subsets):
raise ValueError("OmniInteract subsets must not contain duplicates")
if num_prompts < 0:
raise ValueError("num_prompts must be non-negative")
data_root, cases = _data_dir(root.resolve()), []
for subset in subsets:
subset_root = data_root / subset
selected = _one_to_many_cases(subset_root) if subset == "1qna" else _mapping_cases(subset_root, subset)
if not selected:
raise ValueError(f"No OmniInteract sessions found for requested subset {subset!r}")
cases.extend(selected)
if len({case.video_path for case in cases}) != len(cases):
raise ValueError("OmniInteract dataset contains duplicate video paths")
if not disable_shuffle:
random.Random(seed).shuffle(cases)
if num_prompts:
if num_prompts > len(cases):
logger.warning(
"Requested %d OmniInteract prompts but only %d are available; using all cases", num_prompts, len(cases)
)
cases = cases[:num_prompts]
return cases
def case_manifest(case: OmniInteractCase, output_dir: Path) -> dict[str, Any]:
return {
"sample_id": f"{case.subset}__{output_dir.name}",
"video": str(case.video_path),
"gt_json": str(case.annotation_path.resolve()),
"model_json": str((output_dir / "wav_transcript.json").resolve()),
"scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type,
}
+111 -49
View File
@@ -1,5 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import warnings
from collections import defaultdict
from collections.abc import Sequence
from dataclasses import field, make_dataclass
import numpy as np
@@ -15,6 +19,11 @@ _FLOAT_LIST_TYPE = list[float]
_INT_LIST_TYPE = list[int]
_MULTIMODAL_BENCHMARK_FIELDS = [
("num_ttft_samples", int, field(default=0)),
("num_tpot_samples", int, field(default=0)),
("num_itl_samples", int, field(default=0)),
("num_audio_ttfp_samples", int, field(default=0)),
("num_audio_rtf_samples", int, field(default=0)),
(defs.MEAN_AUDIO_TTFP_MS, float, field(default=0.0)),
(defs.MEDIAN_AUDIO_TTFP_MS, float, field(default=0.0)),
(defs.STD_AUDIO_TTFP_MS, float, field(default=0.0)),
@@ -176,6 +185,17 @@ def _stage_modality_flags(
)
def has_metric_samples(metrics: object, metric_name: str) -> bool:
sample_count_attr = {
"ttft": "num_ttft_samples",
"tpot": "num_tpot_samples",
"itl": "num_itl_samples",
defs.AUDIO_TTFP: "num_audio_ttfp_samples",
defs.AUDIO_RTF: "num_audio_rtf_samples",
}.get(metric_name)
return sample_count_attr is None or getattr(metrics, sample_count_attr, 0) > 0
def print_metrics(
task_type,
selected_percentile_metrics,
@@ -308,6 +328,9 @@ def process_one_metric(
metric_attribute_name: str,
metrics: MultiModalsBenchmarkMetrics,
):
if not has_metric_samples(metrics, metric_attribute_name):
return
metric_header_map = {
"ttft": "Time to First Token",
"tpot": "Time per Output Token (excl. 1st token)",
@@ -718,10 +741,12 @@ def calculate_metrics(
good_completed = 0
itls: list[float] = []
tpots: list[float] = []
all_tpots: list[float] = []
all_tpots: list[float | None] = []
ttfts: list[float] = []
goodput_ttfts: list[float | None] = []
e2els: list[float] = []
audio_ttfps: list[float] = []
goodput_audio_ttfps: list[float | None] = []
audio_rtfs: list[float] = []
audio_duration: list[float] = []
audio_frames: list[int] = []
@@ -748,7 +773,7 @@ def calculate_metrics(
output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
actual_output_lens.append(output_len)
total_input += outputs[i].prompt_len
tpot = 0
tpot: float | None = 0.0
if output_len > 1:
if outputs[i].itl:
# Use mean(ITL) directly so per-request TPOT == mean(ITL).
@@ -756,19 +781,33 @@ def calculate_metrics(
# bundle multiple tokens per chunk, so len(itl)+1 != output_len.
# Using mean(itl) keeps TPOT and ITL on the same footing.
tpot = sum(outputs[i].itl) / len(outputs[i].itl)
else:
elif getattr(outputs[i], "tpot_measured", True):
try:
latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft
except Exception:
latency_minus_ttft = outputs[i].latency - outputs[i].ttft
tpot = latency_minus_ttft / (output_len - 1)
tpots.append(tpot)
else:
tpot = None
if tpot is not None:
tpots.append(tpot)
# Note: if output_len <= 1, we regard tpot as 0 for goodput
all_tpots.append(tpot)
itls += outputs[i].itl
ttfts.append(outputs[i].ttft)
audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
session_metrics = getattr(outputs[i], "duplex_session_metrics", None)
ttft_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_ttft_ms") is not None
audio_ttfp_measured = (
not isinstance(session_metrics, dict) or session_metrics.get("mean_ttfp_ms") is not None
)
audio_rtf_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_rtf") is not None
if ttft_measured:
ttfts.append(outputs[i].ttft)
if audio_ttfp_measured:
audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
if audio_rtf_measured:
audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
goodput_ttfts.append(outputs[i].ttft if ttft_measured else None)
goodput_audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0) if audio_ttfp_measured else None)
audio_duration.append(getattr(outputs[i], defs.AUDIO_DURATION, 0.0))
audio_frames.append(getattr(outputs[i], defs.AUDIO_FRAMES, 0.0))
image_count = int(getattr(outputs[i], defs.IMAGE_COUNT, 0) or 0)
@@ -789,14 +828,14 @@ def calculate_metrics(
actual_output_lens.append(0)
if goodput_config_dict:
valid_metrics = []
valid_metrics: list[Sequence[float | None]] = []
slo_values = []
if "ttft" in goodput_config_dict:
valid_metrics.append(ttfts)
valid_metrics.append(goodput_ttfts)
slo_values.append(goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
if "audio_ttft" in goodput_config_dict:
valid_metrics.append(audio_ttfps)
valid_metrics.append(goodput_audio_ttfps)
slo_values.append(goodput_config_dict["audio_ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
if "tpot" in goodput_config_dict:
valid_metrics.append(all_tpots)
@@ -806,7 +845,7 @@ def calculate_metrics(
slo_values.append(goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION)
for req_metric in zip(*valid_metrics):
is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)])
is_good_req = all(r is not None and s >= r for s, r in zip(slo_values, req_metric))
if is_good_req:
good_completed += 1
@@ -822,6 +861,7 @@ def calculate_metrics(
# Calculate max output tokens per second metric
max_output_tokens_per_s = 0.0
max_concurrent_requests = 0
token_timeline_available = True
# Find the time range across all successful requests
successful_outputs = [output for output in outputs if output.success]
@@ -836,19 +876,26 @@ def calculate_metrics(
concurrent_requests_per_second = np.zeros(duration_seconds)
for i, output in enumerate(successful_outputs):
# Calculate token generation timestamp using
# start_time, ttft, and itl
token_times = [output.start_time + output.ttft]
current_time = token_times[0]
for itl_value in output.itl:
current_time += itl_value
token_times.append(current_time)
session_metrics = getattr(output, "duplex_session_metrics", None)
if isinstance(session_metrics, dict):
# Duplex TTFT/ITL samples are response-local. Without each
# response.created offset they cannot form a session-global
# token timeline, so do not publish a misleading peak rate.
token_timeline_available = False
else:
# Calculate token generation timestamp using
# start_time, ttft, and itl
token_times = [output.start_time + output.ttft]
current_time = token_times[0]
for itl_value in output.itl:
current_time += itl_value
token_times.append(current_time)
# Add tokens to second buckets
for token_time in token_times:
second_bucket = int(token_time - min_start_time)
if 0 <= second_bucket < duration_seconds:
tokens_per_second[second_bucket] += 1
# Add tokens to second buckets
for token_time in token_times:
second_bucket = int(token_time - min_start_time)
if 0 <= second_bucket < duration_seconds:
tokens_per_second[second_bucket] += 1
# Track concurrent requests for each second this request was active
request_start_second = int(output.start_time - min_start_time)
@@ -860,18 +907,22 @@ def calculate_metrics(
# Find the maximum tokens per second and corresponding
# concurrent requests
if len(tokens_per_second) > 0:
max_output_tokens_per_s = float(np.max(tokens_per_second))
if token_timeline_available:
max_output_tokens_per_s = float(np.max(tokens_per_second))
else:
max_output_tokens_per_s = float("nan")
max_concurrent_requests = int(np.max(concurrent_requests_per_second))
if TERM_PLOTLIB_AVAILABLE:
import termplotlib as tpl
fig = tpl.figure()
fig.plot(
np.arange(len(tokens_per_second)),
tokens_per_second,
title="Output tokens per second",
)
if token_timeline_available:
fig.plot(
np.arange(len(tokens_per_second)),
tokens_per_second,
title="Output tokens per second",
)
fig.plot(
np.arange(len(concurrent_requests_per_second)),
concurrent_requests_per_second,
@@ -881,6 +932,10 @@ def calculate_metrics(
else:
print("tip: install termplotlib and gnuplot to plot the metrics")
duplex_metrics_present = any(
isinstance(getattr(output, "duplex_session_metrics", None), dict) for output in outputs
)
missing_duplex_value = float("nan") if duplex_metrics_present else 0
metrics = MultiModalsBenchmarkMetrics(
completed=completed,
failed=len(failed_outputs),
@@ -890,18 +945,18 @@ def calculate_metrics(
request_goodput=good_completed / dur_s,
output_throughput=sum(actual_output_lens) / dur_s,
total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,
mean_ttft_ms=np.mean(ttfts or 0) * 1000, # ttfts is empty if streaming is not supported by the endpoint
std_ttft_ms=np.std(ttfts or 0) * 1000,
median_ttft_ms=np.median(ttfts or 0) * 1000,
percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles],
mean_tpot_ms=np.mean(tpots or 0) * 1000,
std_tpot_ms=np.std(tpots or 0) * 1000,
median_tpot_ms=np.median(tpots or 0) * 1000,
percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles],
mean_itl_ms=np.mean(itls or 0) * 1000,
std_itl_ms=np.std(itls or 0) * 1000,
median_itl_ms=np.median(itls or 0) * 1000,
percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles],
mean_ttft_ms=np.mean(ttfts or missing_duplex_value) * 1000,
std_ttft_ms=np.std(ttfts or missing_duplex_value) * 1000,
median_ttft_ms=np.median(ttfts or missing_duplex_value) * 1000,
percentiles_ttft_ms=[(p, np.percentile(ttfts or missing_duplex_value, p) * 1000) for p in selected_percentiles],
mean_tpot_ms=np.mean(tpots or missing_duplex_value) * 1000,
std_tpot_ms=np.std(tpots or missing_duplex_value) * 1000,
median_tpot_ms=np.median(tpots or missing_duplex_value) * 1000,
percentiles_tpot_ms=[(p, np.percentile(tpots or missing_duplex_value, p) * 1000) for p in selected_percentiles],
mean_itl_ms=np.mean(itls or missing_duplex_value) * 1000,
std_itl_ms=np.std(itls or missing_duplex_value) * 1000,
median_itl_ms=np.median(itls or missing_duplex_value) * 1000,
percentiles_itl_ms=[(p, np.percentile(itls or missing_duplex_value, p) * 1000) for p in selected_percentiles],
mean_e2el_ms=np.mean(e2els or 0) * 1000,
std_e2el_ms=np.std(e2els or 0) * 1000,
median_e2el_ms=np.median(e2els or 0) * 1000,
@@ -910,11 +965,16 @@ def calculate_metrics(
max_concurrent_requests=max_concurrent_requests,
rtfx=input_audio_duration / dur_s,
**{
defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or 0) * 1000,
defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or 0) * 1000,
defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or 0) * 1000,
"num_ttft_samples": len(ttfts),
"num_tpot_samples": len(tpots),
"num_itl_samples": len(itls),
"num_audio_ttfp_samples": len(audio_ttfps),
"num_audio_rtf_samples": len(audio_rtfs),
defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or missing_duplex_value) * 1000,
defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or missing_duplex_value) * 1000,
defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or missing_duplex_value) * 1000,
defs.PERCENTILES_AUDIO_TTFP_MS: [
(p, np.percentile(audio_ttfps or 0, p) * 1000) for p in selected_percentiles
(p, np.percentile(audio_ttfps or missing_duplex_value, p) * 1000) for p in selected_percentiles
],
defs.MEAN_AUDIO_DURATION_S: np.mean(audio_duration or 0),
defs.STD_AUDIO_DURATION_S: np.std(audio_duration or 0),
@@ -925,10 +985,12 @@ def calculate_metrics(
defs.TOTAL_AUDIO_DURATION_S: sum(audio_duration),
defs.TOTAL_AUDIO_FRAMES: sum(audio_frames),
defs.AUDIO_THROUGHPUT: sum(audio_duration) / dur_s,
defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or 0),
defs.STD_AUDIO_RTF: np.std(audio_rtfs or 0),
defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or 0),
defs.PERCENTILES_AUDIO_RTF: [(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles],
defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or missing_duplex_value),
defs.STD_AUDIO_RTF: np.std(audio_rtfs or missing_duplex_value),
defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or missing_duplex_value),
defs.PERCENTILES_AUDIO_RTF: [
(p, np.percentile(audio_rtfs or missing_duplex_value, p)) for p in selected_percentiles
],
defs.TOTAL_IMAGES: total_images,
defs.IMAGE_THROUGHPUT: total_images / dur_s,
defs.AVERAGE_PIXELS_PER_IMAGE: (total_image_pixels / total_images) if total_images > 0 else 0.0,
+933
View File
@@ -0,0 +1,933 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""OmniInteract session execution and artifacts for serving benchmarks."""
from __future__ import annotations
import asyncio
import binascii
import contextlib
import copy
import fcntl
import hashlib
import json
import logging
import math
import os
import subprocess
import tempfile
import time
from collections import Counter
from collections.abc import Callable, Iterator, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, cast
from urllib.parse import urljoin, urlsplit
import pybase64 as base64
from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
DEFAULT_MAX_VIDEO_DURATION_S,
OmniInteractCase,
OmniInteractPreparedInput,
case_manifest,
)
from vllm_omni.experimental.fullduplex.client import (
PCM16_BYTES_PER_SAMPLE,
PCM16_SAMPLE_RATE,
RealtimeDuplexClient,
RealtimeEventCollector,
build_realtime_url,
chunk_period_ms,
has_residual_model_unit,
reference_audio_data_url,
summarize_session_request_metrics,
write_pcm16_wav,
)
OUTPUT_SAMPLE_RATE = 24_000
SUCCESS_ARTIFACTS = (".done", "output.wav", "wav_transcript.json", "events.json", "result.json")
BATCH_ARTIFACTS = ("batch_summary.json", "official_eval_manifest.jsonl")
ARTIFACT_LOCK_FILE = ".omniinteract.lock"
_INPUT_CHUNK_MS = 200
VIDEO_FPS = 1.0
_COMPLETION_SETTLE_S = 2.0
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class OmniInteractBenchmarkConfig:
model: str
base_url: str = "http://127.0.0.1:8000"
endpoint: str = "/v1/realtime"
output_root: Path = Path("omniinteract-output")
timeout_s: float = 900.0
media_timeout_s: float = 600.0
max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S
ref_audio: str | None = None
require_response: bool = False
extra_headers: dict[str, str] | None = None
extra_body: dict[str, object] | None = None
@dataclass
class OmniInteractCaseResult:
subset: str
video: str
output_dir: str
success: bool = False
error: str = ""
session_id: str = ""
latency_s: float = 0.0
input_audio_chunks: int = 0
input_video_frames: int = 0
pacing_mean_lag_s: float = 0.0
pacing_max_lag_s: float = 0.0
responses: int = 0
audio_bytes: int = 0
audio_clipped_bytes: int = 0
transcript: str = ""
eligible_for_official_eval: bool = False
official_eval_ineligible_reasons: list[str] = field(default_factory=list)
artifact_warnings: list[str] = field(default_factory=list)
output_tokens: int = 0
duplex_request_metrics: list[dict[str, object]] = field(default_factory=list)
duplex_session_metrics: dict[str, object] = field(default_factory=dict)
_artifact_context: _DeferredArtifactContext | None = field(default=None, repr=False, compare=False)
def as_dict(self) -> dict[str, Any]:
return copy.deepcopy({key: value for key, value in vars(self).items() if not key.startswith("_")})
@dataclass(frozen=True)
class _DeferredArtifactContext:
horizon_bytes: int
spans: tuple[tuple[int, bytes], ...]
chunks: list[dict[str, object]]
events: list[dict[str, object]]
def benchmark_summary(results: list[OmniInteractCaseResult]) -> dict[str, Any]:
succeeded = sum(result.success for result in results)
eligible = sum(result.success and result.eligible_for_official_eval for result in results)
return {
"total": len(results),
"success": succeeded,
"failed": len(results) - succeeded,
"eligible_for_official_eval": eligible,
"successful_but_ineligible": succeeded - eligible,
"audio_clipped_bytes": sum(result.audio_clipped_bytes for result in results),
"results": [result.as_dict() for result in results],
}
def _run_media_command(
command: list[str],
*,
timeout_s: float,
text: bool = False,
) -> subprocess.CompletedProcess[Any]:
try:
result = subprocess.run(
command,
capture_output=True,
check=False,
text=text,
timeout=timeout_s,
)
except subprocess.TimeoutExpired as exc:
raise TimeoutError(f"Media command timed out after {timeout_s:g}s: {command[0]}") from exc
except FileNotFoundError as exc:
raise RuntimeError(f"Required media command is unavailable: {command[0]}") from exc
if result.returncode:
error = result.stderr if text else result.stderr.decode("utf-8", "ignore")
raise RuntimeError(f"{command[0]} failed: {error.strip()}")
return result
def prepare_media(
video: Path,
fps: float,
*,
timeout_s: float,
max_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S,
) -> tuple[float, bytes, list[str | None]]:
probe = [
"ffprobe",
*"-v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1".split(),
str(video),
]
duration = float(_run_media_command(probe, text=True, timeout_s=timeout_s).stdout.strip())
if not math.isfinite(duration) or duration <= 0 or duration > max_duration_s:
raise ValueError(f"Invalid video duration for {video}: {duration!r}")
target_bytes = math.ceil(duration) * PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
audio = [
"ffmpeg",
*"-loglevel error -i".split(),
str(video),
"-t",
str(duration),
*f"-vn -f s16le -ac 1 -ar {PCM16_SAMPLE_RATE} -fs".split(),
str(target_bytes),
"pipe:1",
]
pcm = _run_media_command(audio, timeout_s=timeout_s).stdout
pcm = (pcm + bytes(max(0, target_bytes - len(pcm))))[:target_bytes]
frame_count = math.ceil(duration * fps)
frames: list[str | None] = [None] * frame_count
with tempfile.TemporaryDirectory(prefix="vllm-omni-frames-") as temp_dir:
output_pattern = Path(temp_dir) / "frame-%06d.jpg"
frame_filter = f"select=gte(t\\,(selected_n+0.5)/{fps}),scale=640:640:force_original_aspect_ratio=decrease"
frame_command = [
"ffmpeg",
*"-loglevel error -i".split(),
str(video),
"-an",
"-vf",
frame_filter,
*f"-frames:v {frame_count} -vsync vfr -q:v 5".split(),
str(output_pattern),
]
_run_media_command(frame_command, timeout_s=timeout_s)
for index, frame_path in enumerate(sorted(Path(temp_dir).glob("frame-*.jpg"))[:frame_count]):
frames[index] = base64.b64encode(frame_path.read_bytes()).decode("ascii")
return duration, pcm, frames
@dataclass(frozen=True)
class _AudioSegment:
event_index: int
response_id: str
start_s: float
end_s: float
pcm16: bytes
class _Playback:
def __init__(self) -> None:
self.cursor = 0
self.end_s = 0.0
self.segments: list[_AudioSegment] = []
self.completed: set[str] = set()
self.completion_acked: set[str] = set()
self.warnings: list[str] = []
self._total_samples: dict[str, int] = {}
self._response_end_s: dict[str, float] = {}
def _warn_once(self, warning: str) -> None:
if warning not in self.warnings:
self.warnings.append(warning)
def ingest(self, events: RealtimeEventCollector) -> None:
while self.cursor < len(events.events):
index, self.cursor = self.cursor, self.cursor + 1
event = events.events[index]
response_id = events.response_id(event)
if event.get("type") == "response.done":
if response_id:
self.completed.add(response_id)
continue
if event.get("type") != "response.audio.delta":
continue
encoded = event.get("delta") or event.get("audio")
if not response_id:
raise ValueError("response audio has no response_id")
if not isinstance(encoded, str):
raise ValueError("response audio payload is missing")
if event.get("format") is None:
self._warn_once("response.audio.delta omitted format; assumed pcm16")
elif event.get("format") != "pcm16":
raise ValueError("OmniInteract output must be pcm16")
rate = event.get("sample_rate_hz")
if rate is None:
rate = events.output_sample_rate_hz or OUTPUT_SAMPLE_RATE
self._warn_once(f"response.audio.delta omitted sample_rate_hz; assumed {rate}")
if rate != OUTPUT_SAMPLE_RATE:
raise ValueError(f"OmniInteract output must use {OUTPUT_SAMPLE_RATE} Hz audio")
try:
raw = base64.b64decode(encoded, validate=True)
except (ValueError, binascii.Error) as exc:
raise ValueError("response audio is not valid base64") from exc
if not raw or len(raw) % PCM16_BYTES_PER_SAMPLE:
raise ValueError("response audio is empty or not PCM16 aligned")
start = max(events.event_received_at_s[index], self.end_s)
samples = len(raw) // PCM16_BYTES_PER_SAMPLE
segment = _AudioSegment(index, response_id, start, start + samples / rate, raw)
self.segments.append(segment)
self.end_s = segment.end_s
self._total_samples[response_id] = self._total_samples.get(response_id, 0) + samples
self._response_end_s[response_id] = segment.end_s
async def acknowledge(self, client: RealtimeDuplexClient, now: float | None = None) -> None:
events = client.events
self.ingest(events)
now = time.monotonic() if now is None else now
for response_id in self.completed - self.completion_acked:
total_samples = self._total_samples.get(response_id, 0)
if not total_samples or now < self._response_end_s[response_id]:
continue
played_ms = total_samples * 1000 // OUTPUT_SAMPLE_RATE
await client.send_playback_ack(response_id, played_ms)
self.completion_acked.add(response_id)
async def stream_inputs(
client: RealtimeDuplexClient,
pcm: bytes,
frames: Sequence[str | None],
playback: _Playback,
) -> tuple[int, int, float, float]:
bytes_per_second = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
chunk_bytes = bytes_per_second * _INPUT_CHUNK_MS // 1000
started_at = time.monotonic()
frame_cursor = sent_frames = 0
lags: list[float] = []
for chunks, offset in enumerate(range(0, len(pcm), chunk_bytes), start=1):
end = min(offset + chunk_bytes, len(pcm))
end_ms = end * 1000 // bytes_per_second
ready: list[str] = []
while frame_cursor < len(frames) and end_ms >= (frame_cursor + 0.5) * 1000 / VIDEO_FPS:
if frames[frame_cursor]:
ready.append(frames[frame_cursor] or "")
frame_cursor += 1
sent_frames += len(ready)
event: dict[str, object] = {
"type": "input_audio_buffer.append",
"audio": base64.b64encode(pcm[offset:end]).decode("ascii"),
"input_audio_format": "pcm16",
"sample_rate_hz": PCM16_SAMPLE_RATE,
"duration_ms": (end - offset) * 1000 // bytes_per_second,
"audio_end_ms": end_ms,
}
if ready:
event["video_frames"] = ready
lags.append(max(0.0, time.monotonic() - started_at - offset / bytes_per_second))
await client.send(event)
await playback.acknowledge(client)
await asyncio.sleep(max(0.0, started_at + end_ms / 1000 - time.monotonic()))
return chunks if pcm else 0, sent_frames, sum(lags) / len(lags) if lags else 0.0, max(lags, default=0.0)
def _response_status(event: dict[str, object]) -> str | None:
response = event.get("response")
if isinstance(response, dict) and response.get("status"):
return str(response["status"])
status = event.get("status")
return str(status) if status else None
def response_ledger(
collector: RealtimeEventCollector,
*,
end: int | None = None,
) -> tuple[set[str], set[str]]:
created: set[str] = set()
done: set[str] = set()
for event in collector.events[:end]:
event_type = event.get("type")
if event_type not in {"response.created", "response.done"}:
continue
response_id = collector.response_id(event)
if not response_id:
raise ValueError(f"{event_type} has no response_id")
if event_type == "response.created":
if response_id in created:
raise ValueError(f"duplicate response.created for {response_id}")
created.add(response_id)
continue
if response_id not in created:
raise ValueError(f"response.done without response.created for {response_id}")
if response_id in done:
raise ValueError(f"duplicate response.done for {response_id}")
response = event.get("response") if isinstance(event.get("response"), dict) else {}
details = event.get("status_details") if isinstance(event.get("status_details"), dict) else {}
nested = response.get("status_details") if isinstance(response.get("status_details"), dict) else {}
if "failed" in {event.get("status"), response.get("status"), details.get("type"), nested.get("type")}:
raise ValueError(f"response.done reports failure for {response_id}")
done.add(response_id)
return created, done
def _is_model_listen_decision(event: dict[str, object]) -> bool:
response = event.get("response")
metadata = response.get("metadata") if isinstance(response, dict) else event.get("metadata")
return isinstance(metadata, dict) and metadata.get("buffering") is not True and metadata.get("model_listen") is True
def _has_post_commit_decision(
collector: RealtimeEventCollector,
events: list[dict[str, object]],
*,
prior_response_ids: set[str],
) -> bool:
pending_prior = set(prior_response_ids)
for event in events:
event_type = event.get("type")
if event_type == "response.done":
response_id = collector.response_id(event)
if response_id in pending_prior:
pending_prior.remove(response_id)
continue
if not pending_prior and _response_status(event) != "cancelled":
return True
elif event_type == "response.listen" and not pending_prior and _is_model_listen_decision(event):
return True
return False
def _raise_if_session_terminated(
collector: RealtimeEventCollector,
from_index: int,
*,
explicit_close_from: int | None = None,
) -> None:
errors = collector.errors()
if errors:
raise RuntimeError(str(errors[-1]))
for index, event in enumerate(collector.events[from_index:], start=from_index):
event_type = event.get("type")
if event_type not in {"session.expired", "session.closed"}:
continue
reason = event.get("reason")
nested = event.get("event")
if reason is None and isinstance(nested, dict):
reason = nested.get("reason")
expected = (
event_type == "session.closed"
and explicit_close_from is not None
and index >= explicit_close_from
and reason is None
)
if expected:
continue
detail = f": {reason}" if reason else ""
prefix = "Unexpected " if explicit_close_from is not None else ""
raise RuntimeError(f"{prefix}{event_type}{detail}")
def _ensure_final_commit_tail(pcm: bytes, events: list[dict[str, object]]) -> bytes:
"""Reserve one PCM sample so an exact model unit is flushed by commit.
A complete unit is otherwise emitted before commit, so its asynchronous
decision cannot be correlated with the accepted final input. Keeping an
almost-full unit buffered makes the server's final residual flush carry
that correlation; the server pads the missing sample back to a full unit.
"""
period_ms = chunk_period_ms(events)
if len(pcm) >= PCM16_BYTES_PER_SAMPLE and not has_residual_model_unit(pcm, chunk_period_ms=period_ms):
return pcm[:-PCM16_BYTES_PER_SAMPLE]
return pcm
async def wait_for_session_completion(
client: RealtimeDuplexClient,
playback: _Playback,
*,
commit_from: int,
session_from: int | None = None,
timeout_s: float,
settle_s: float,
) -> int:
deadline = time.monotonic() + timeout_s
committed_index: int | None = None
prior_response_ids: set[str] = set()
last_event_count = len(client.events.events)
stable_since = time.monotonic()
while time.monotonic() < deadline:
client.raise_if_reader_stopped()
_raise_if_session_terminated(client.events, commit_from if session_from is None else session_from)
await playback.acknowledge(client)
if len(client.events.events) != last_event_count:
last_event_count = len(client.events.events)
stable_since = time.monotonic()
if committed_index is None:
committed_index = next(
(
index
for index in range(commit_from, len(client.events.events))
if client.events.events[index].get("type") == "input_audio_buffer.committed"
),
None,
)
if committed_index is not None:
committed_event = client.events.events[committed_index]
committed = committed_event.get("event")
if isinstance(committed, dict) and committed.get("overlap_deferred") is True:
prior_created, prior_done = response_ledger(client.events, end=committed_index)
prior_response_ids = prior_created - prior_done
stable_since = time.monotonic()
if committed_index is not None:
created, done = response_ledger(client.events)
post_commit = client.events.events[committed_index + 1 :]
decision = _has_post_commit_decision(
client.events,
post_commit,
prior_response_ids=prior_response_ids,
)
if created == done and decision and time.monotonic() - stable_since >= settle_s:
return committed_index
await asyncio.sleep(0.05)
missing: set[str] = set()
with contextlib.suppress(ValueError):
created, done = response_ledger(client.events)
missing = created - done
raise TimeoutError(
"Timed out waiting for committed input and stable responses"
+ (f"; unfinished response_ids={sorted(missing)}" if missing else "")
)
def _output_dir(root: Path, case: OmniInteractCase) -> Path:
relative = case.video_rel.replace("\\", "/")
stem = Path(relative).with_suffix("").as_posix().replace("/", "__")
digest = hashlib.sha256(f"{case.subset}/{relative}".encode()).hexdigest()[:8]
return root / case.subset / f"{stem}--{digest}"
@contextlib.contextmanager
def omniinteract_output_lock(root: Path) -> Iterator[None]:
root.mkdir(parents=True, exist_ok=True)
with (root / ARTIFACT_LOCK_FILE).open("a+b") as lock_file:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
def _atomic_replace(path: Path, writer: Callable[[Path], None]) -> None:
descriptor, temporary_name = tempfile.mkstemp(
dir=path.parent,
prefix=f".{path.name}.",
suffix=".tmp",
)
os.close(descriptor)
temporary = Path(temporary_name)
try:
writer(temporary)
temporary.replace(path)
finally:
temporary.unlink(missing_ok=True)
def _atomic_write_text(path: Path, value: str) -> None:
_atomic_replace(path, lambda temporary: temporary.write_text(value, encoding="utf-8"))
def _atomic_write_json(path: Path, value: object) -> None:
_atomic_write_text(path, json.dumps(value, ensure_ascii=False, indent=2) + "\n")
def _atomic_write_wav(path: Path, pcm: bytes | bytearray, rate: int) -> None:
_atomic_replace(
path,
lambda temporary: write_pcm16_wav(
temporary,
cast(bytes, pcm), # wave.writeframes accepts bytearray.
sample_rate_hz=rate,
),
)
def _clear_artifacts(directory: Path, names: Sequence[str]) -> None:
directory.mkdir(parents=True, exist_ok=True)
for name in names:
(directory / name).unlink(missing_ok=True)
def _artifact_summary(case: OmniInteractCase, result: OmniInteractCaseResult) -> dict[str, Any]:
return {
**result.as_dict(),
"annotation": str(case.annotation_path),
"scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type,
}
def _publish_success_artifacts(
directory: Path,
result: OmniInteractCaseResult,
context: _DeferredArtifactContext,
summary: dict[str, Any],
) -> None:
pcm = bytearray(context.horizon_bytes)
for offset, chunk in context.spans:
pcm[offset : offset + len(chunk)] = chunk
try:
_clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json"))
_atomic_write_wav(directory / "output.wav", pcm, OUTPUT_SAMPLE_RATE)
_atomic_write_json(
directory / "wav_transcript.json",
{
"text": result.transcript,
"chunks": context.chunks,
"timestamp_semantics": "serialized playback queue time relative to input streaming start",
},
)
_atomic_write_json(directory / "events.json", context.events)
_atomic_write_json(directory / "result.json", result.as_dict())
_atomic_write_json(directory / ".done", summary)
except Exception:
_clear_artifacts(directory, SUCCESS_ARTIFACTS)
result.success = result.eligible_for_official_eval = False
if "artifact_write_failed" not in result.official_eval_ineligible_reasons:
result.official_eval_ineligible_reasons.append("artifact_write_failed")
raise
def _collect_output(
collector: RealtimeEventCollector,
playback: _Playback,
*,
stream_start: float,
video_duration_s: float,
require_response: bool,
retain_events: bool,
result: OmniInteractCaseResult,
) -> _DeferredArtifactContext:
playback.ingest(collector)
created, done = response_ledger(collector)
if created != done:
raise ValueError(f"unfinished response_ids: {sorted(created - done)}")
done_events = [event for event in collector.events if event.get("type") == "response.done"]
cancelled = any(_response_status(event) == "cancelled" for event in done_events)
terminal_responses = {
response_id
for event in done_events
if _response_status(event) != "cancelled" and (response_id := collector.response_id(event))
}
horizon_s = math.ceil(video_duration_s)
horizon_bytes = horizon_s * OUTPUT_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
spans: list[tuple[int, bytes]] = []
clipped_bytes = 0
response_times: dict[str, list[float]] = {}
responses_with_audio: set[str] = set()
for segment in playback.segments:
response_id = segment.response_id
if response_id not in created:
raise ValueError("response audio has no matching response.created")
start_s = max(0.0, segment.start_s - stream_start)
end_s = max(start_s, segment.end_s - stream_start)
offset = round(start_s * OUTPUT_SAMPLE_RATE) * PCM16_BYTES_PER_SAMPLE
writable = max(0, min(len(segment.pcm16), horizon_bytes - offset))
clipped_bytes += len(segment.pcm16) - writable
if writable:
clipped = segment.pcm16[:writable]
spans.append((offset, clipped))
if any(clipped):
responses_with_audio.add(response_id)
timing = response_times.setdefault(response_id, [start_s, end_s])
timing[0], timing[1] = min(timing[0], start_s), max(timing[1], end_s)
text_event_types = {"response.audio_transcript.delta", "response.output_text.delta", "response.text.delta"}
text_ids = {collector.response_id(event) for event in collector.events if event.get("type") in text_event_types}
if None in text_ids or not text_ids <= created:
raise ValueError("response transcript has no matching response.created")
chunks: list[dict[str, object]] = []
texts: list[str] = []
for response_id in collector.response_ids:
text = collector.response_text(response_id).strip()
if not text:
continue
response_timing = response_times.get(response_id)
if response_timing is None or response_timing[0] >= horizon_s:
continue
texts.append(text)
chunks.append(
{
"response_id": response_id,
"text": text,
"timestamp": [round(response_timing[0], 6), round(min(response_timing[1], horizon_s), 6)],
}
)
complete_outputs = terminal_responses & responses_with_audio & {str(chunk["response_id"]) for chunk in chunks}
if require_response and not complete_outputs:
raise ValueError("OmniInteract E2E requires a response with audio and transcript")
result.audio_bytes, result.audio_clipped_bytes = (
sum(len(segment.pcm16) for segment in playback.segments),
clipped_bytes,
)
result.transcript, result.responses, result.success = " ".join(texts).strip(), len(created), True
result.artifact_warnings = list(playback.warnings)
result.official_eval_ineligible_reasons = []
if clipped_bytes:
result.official_eval_ineligible_reasons.append("audio_clipped")
if cancelled:
result.official_eval_ineligible_reasons.append("cancelled_response")
result.eligible_for_official_eval = not result.official_eval_ineligible_reasons
audio_sizes = {segment.event_index: len(segment.pcm16) for segment in playback.segments}
events = (
[
{
**{key: value for key, value in event.items() if key not in {"delta", "audio"}},
"audio_bytes": audio_sizes[index],
}
if event.get("type") == "response.audio.delta"
else dict(event)
for index, event in enumerate(collector.events)
]
if retain_events
else []
)
return _DeferredArtifactContext(horizon_bytes, tuple(spans), chunks, events)
def prepare_success_artifacts(
root: Path,
case: OmniInteractCase,
collector: RealtimeEventCollector,
*,
playback: _Playback | None = None,
stream_start: float,
video_duration_s: float,
require_response: bool,
result: OmniInteractCaseResult,
capture_artifacts: bool = True,
) -> dict[str, Any]:
directory = _output_dir(root, case)
context = _collect_output(
collector,
playback or _Playback(),
stream_start=stream_start,
video_duration_s=video_duration_s,
require_response=require_response,
retain_events=capture_artifacts,
result=result,
)
result.output_dir = str(directory.resolve())
summary = _artifact_summary(case, result)
if capture_artifacts:
result._artifact_context = context
return summary
def write_failure_artifacts(root: Path, case: OmniInteractCase, result: OmniInteractCaseResult) -> None:
directory = _output_dir(root, case)
result.output_dir = str(directory.resolve())
result.success = result.eligible_for_official_eval = False
_clear_artifacts(directory, SUCCESS_ARTIFACTS)
_atomic_write_json(directory / ".failed.json", _artifact_summary(case, result))
def publish_deferred_case_artifacts(
root: Path,
case: OmniInteractCase,
result: OmniInteractCaseResult,
) -> None:
context = result._artifact_context
result._artifact_context = None
if not result.success:
write_failure_artifacts(root, case, result)
return
if context is None:
raise RuntimeError("OmniInteract benchmark output lost its deferred artifact context")
directory = _output_dir(root, case)
_publish_success_artifacts(directory, result, context, _artifact_summary(case, result))
def clear_case_artifacts(root: Path, case: OmniInteractCase) -> None:
directory = _output_dir(root, case)
_clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json"))
def clear_batch_artifacts(root: Path) -> None:
_clear_artifacts(root, BATCH_ARTIFACTS)
def write_batch_artifacts(
root: Path,
cases: list[OmniInteractCase],
results: list[OmniInteractCaseResult],
) -> None:
root.mkdir(parents=True, exist_ok=True)
rows = [
case_manifest(case, _output_dir(root, case))
for case, result in zip(cases, results, strict=True)
if result.success and result.eligible_for_official_eval
]
ineligible = [result for result in results if result.success and not result.eligible_for_official_eval]
if ineligible:
reasons = Counter(
reason for result in ineligible for reason in (result.official_eval_ineligible_reasons or ["unspecified"])
)
logger.warning(
"%d successful OmniInteract cases were excluded from official evaluation: %s",
len(ineligible),
", ".join(f"{reason}={count}" for reason, count in sorted(reasons.items())),
)
_atomic_write_json(root / "batch_summary.json", benchmark_summary(results))
_atomic_write_text(
root / "official_eval_manifest.jsonl",
"".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
)
def _websocket_url(config: OmniInteractBenchmarkConfig, session_id: str) -> str:
endpoint = (
config.endpoint
if urlsplit(config.endpoint).scheme
else urljoin(config.base_url.rstrip("/") + "/", config.endpoint.lstrip("/"))
)
return build_realtime_url(
endpoint,
config.model,
autostart=False,
native_duplex=True,
session_id=session_id,
)
def _populate_response_metrics(
result: OmniInteractCaseResult,
collector: RealtimeEventCollector,
*,
stream_start: float,
) -> None:
measurement_origin = {
"ttft": "response.created client receive to first non-empty text delta",
"ttfp": "response.created client receive to first audio packet",
"rtf": "response.created client receive to last audio packet divided by emitted audio duration",
}
request_metrics: list[dict[str, object]] = []
output_tokens = 0
for request_index, response_id in enumerate(collector.response_ids):
timing = collector.timing_summary(
after_s=stream_start,
input_committed_at_s=None,
response_id=response_id,
measurement_origin=measurement_origin,
)
raw_metric = timing.get("request_metrics")
stage0 = timing.get("stage0_tokens")
metric = {
"session_id": result.session_id,
"request_index": request_index,
"response_id": response_id,
**(raw_metric if isinstance(raw_metric, dict) else {}),
}
if isinstance(stage0, dict):
metric["stage0_tokens"] = dict(stage0)
output_tokens += int(stage0.get("output_token_count") or 0)
if isinstance(raw_metric, dict) or isinstance(stage0, dict):
request_metrics.append(metric)
result.output_tokens = output_tokens
result.duplex_request_metrics = request_metrics
result.duplex_session_metrics = summarize_session_request_metrics(
request_metrics,
session_id=result.session_id,
)
async def run_omniinteract_case(
case: OmniInteractCase,
config: OmniInteractBenchmarkConfig,
*,
request_index: int | str,
capture_artifacts: bool = True,
prepared_input: OmniInteractPreparedInput | None = None,
) -> OmniInteractCaseResult:
if not config.ref_audio:
raise ValueError("ref_audio is required for MiniCPM-o native-duplex audio output")
for name, value in (
("timeout_s", config.timeout_s),
("media_timeout_s", config.media_timeout_s),
("max_video_duration_s", config.max_video_duration_s),
):
if not math.isfinite(value) or value <= 0:
raise ValueError(f"{name} must be finite and positive")
session_id = f"omniinteract:{case.subset}:{request_index}:{time.monotonic_ns()}"
result = OmniInteractCaseResult(
case.subset,
str(case.video_path),
str(_output_dir(config.output_root, case).resolve()),
session_id=session_id,
)
started_at = time.monotonic()
try:
if prepared_input is None:
reference_audio = reference_audio_data_url(config.ref_audio)
duration, pcm, frames = await asyncio.to_thread(
prepare_media,
case.video_path,
VIDEO_FPS,
timeout_s=config.media_timeout_s,
max_duration_s=config.max_video_duration_s,
)
else:
reference_audio = prepared_input.ref_audio_data_url
duration = prepared_input.duration_s
pcm = prepared_input.pcm16
frames = prepared_input.video_frames
if not any(frames):
raise ValueError(f"No video frames were decoded from {case.video_path}")
async with RealtimeDuplexClient(
_websocket_url(config, session_id), additional_headers=config.extra_headers
) as client:
session_from = len(client.events.events)
await client.configure(
config.model,
ref_audio=reference_audio,
session_id=session_id,
extra_body=config.extra_body,
idle_timeout_s=config.timeout_s,
timeout_s=min(config.timeout_s, 20.0),
)
pcm = _ensure_final_commit_tail(pcm, client.events.events)
playback = _Playback()
stream_start = time.monotonic()
try:
input_duration_s = len(pcm) / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
upload_timeout_s = config.timeout_s + input_duration_s
try:
chunks, frame_count, mean_lag, max_lag = await asyncio.wait_for(
stream_inputs(client, pcm, frames, playback), timeout=upload_timeout_s
)
except asyncio.TimeoutError as exc:
raise TimeoutError(f"Realtime upload timed out after {upload_timeout_s:g}s") from exc
commit_from = len(client.events.events)
await client.commit()
await wait_for_session_completion(
client,
playback,
commit_from=commit_from,
session_from=session_from,
timeout_s=config.timeout_s,
settle_s=_COMPLETION_SETTLE_S,
)
await playback.acknowledge(client, playback.end_s)
_raise_if_session_terminated(client.events, session_from)
close_from = len(client.events.events)
await client.close_session(timeout_s=min(config.timeout_s, 20.0))
except Exception:
with contextlib.suppress(Exception):
await client.close_session(timeout_s=min(config.timeout_s, 20.0))
raise
errors = client.events.errors()
if errors:
raise RuntimeError(str(errors[-1]))
_raise_if_session_terminated(client.events, session_from, explicit_close_from=close_from)
result.latency_s = time.monotonic() - started_at
result.input_audio_chunks, result.input_video_frames = chunks, frame_count
result.pacing_mean_lag_s, result.pacing_max_lag_s = mean_lag, max_lag
_populate_response_metrics(result, client.events, stream_start=stream_start)
await asyncio.to_thread(
prepare_success_artifacts,
config.output_root,
case,
client.events,
playback=playback,
stream_start=stream_start,
video_duration_s=duration,
require_response=config.require_response,
result=result,
capture_artifacts=capture_artifacts,
)
except Exception as exc:
result.success = result.eligible_for_official_eval = False
if "case_failed" not in result.official_eval_ineligible_reasons:
result.official_eval_ineligible_reasons.append("case_failed")
result.error = str(exc)
result.latency_s = max(0.0, time.monotonic() - started_at)
return result
+308 -24
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import contextlib
import io
@@ -16,7 +19,6 @@ from dataclasses import dataclass, replace
from datetime import datetime
from pathlib import Path
from typing import Any, Literal
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import aiohttp
import numpy as np
@@ -31,6 +33,7 @@ from vllm.benchmarks.lib.endpoint_request_func import (
RequestFuncOutput,
StreamedResponseHandler,
_get_chat_content,
_get_headers,
_update_headers_common,
_update_payload_common,
_validate_api_url,
@@ -46,6 +49,13 @@ from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
daily_omni_local_videos_dir,
resolve_daily_omni_local_root,
)
from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
DEFAULT_OMNIINTERACT_REPO,
OmniInteractDataset,
OmniInteractPreparedInput,
OmniInteractSampleRequest,
OmniInteractSessionOptions,
)
from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT,
@@ -56,8 +66,27 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
)
from vllm_omni.benchmarks.data_modules.sound_effect_dataset import SoundEffectDataset
from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
from vllm_omni.benchmarks.omniinteract import (
VIDEO_FPS,
OmniInteractBenchmarkConfig,
OmniInteractCaseResult,
clear_batch_artifacts,
clear_case_artifacts,
prepare_media,
publish_deferred_case_artifacts,
run_omniinteract_case,
write_failure_artifacts,
)
from vllm_omni.benchmarks.omniinteract import (
benchmark_summary as omniinteract_benchmark_summary,
)
from vllm_omni.benchmarks.omniinteract import (
write_batch_artifacts as write_omniinteract_batch_artifacts,
)
from vllm_omni.experimental.fullduplex.client import (
RealtimeDuplexClient,
build_realtime_url,
reference_audio_data_url,
summarize_session_request_metrics,
wait_for,
)
@@ -221,6 +250,91 @@ def _attach_seed_tts_to_request_func_input(sample: SampleRequest, rfi: RequestFu
rfi.extra_body = base
def _attach_omniinteract_to_request_func_input(sample: SampleRequest, rfi: RequestFuncInput) -> None:
if not isinstance(sample, OmniInteractSampleRequest):
return
setattr(rfi, "omniinteract_case", sample.omniinteract_case)
setattr(rfi, "omniinteract_options", sample.omniinteract_options)
setattr(rfi, "omniinteract_prepared_input", sample.omniinteract_prepared_input)
def _append_error(existing: str, message: str) -> str:
return f"{existing}\n{message}" if existing else message
def _finalize_omniinteract_batch(
input_requests: list[SampleRequest],
outputs: list[RequestFuncOutput],
) -> dict[str, object] | None:
rows = [
(sample, output)
for sample, output in zip(input_requests, outputs, strict=True)
if isinstance(sample, OmniInteractSampleRequest)
]
if not rows:
return None
options = rows[0][0].omniinteract_options
if not isinstance(options, OmniInteractSessionOptions):
raise RuntimeError("OmniInteract benchmark output lost its artifact options")
cases, results = [], []
artifact_errors: list[str] = []
for sample, output in rows:
case = sample.omniinteract_case
result = getattr(output, "omniinteract_case_result", None)
if case is None or not isinstance(result, OmniInteractCaseResult):
raise RuntimeError("OmniInteract benchmark output lost its dataset identity")
cases.append(case)
results.append(result)
try:
publish_deferred_case_artifacts(options.output_root, case, result)
except Exception as exc: # noqa: BLE001 - artifact backends raise heterogeneous errors
result.success = result.eligible_for_official_eval = False
if "artifact_write_failed" not in result.official_eval_ineligible_reasons:
result.official_eval_ineligible_reasons.append("artifact_write_failed")
artifact_error = f"Artifact publication failed: {exc}"
result.error = _append_error(result.error, artifact_error)
output.error = _append_error(output.error, artifact_error)
message = f"{case.subset}/{case.video_rel}: {artifact_error}"
artifact_errors.append(message)
logger.exception(message)
try:
write_failure_artifacts(options.output_root, case, result)
except Exception as recovery_exc: # noqa: BLE001 - preserve benchmark metrics after I/O failure
message = f"{case.subset}/{case.video_rel}: failure artifact publication failed: {recovery_exc}"
artifact_errors.append(message)
logger.exception(message)
try:
write_omniinteract_batch_artifacts(options.output_root, cases, results)
except Exception as exc: # noqa: BLE001 - preserve benchmark metrics after I/O failure
message = f"Batch artifact publication failed: {exc}"
artifact_errors.append(message)
logger.exception(message)
try:
clear_batch_artifacts(options.output_root)
except Exception as cleanup_exc: # noqa: BLE001 - best-effort removal of partial batch artifacts
message = f"Partial batch artifact cleanup failed: {cleanup_exc}"
artifact_errors.append(message)
logger.exception(message)
summary = omniinteract_benchmark_summary(results)
compact_summary = {key: value for key, value in summary.items() if key != "results"}
compact_summary["artifacts_complete"] = not artifact_errors
if artifact_errors:
compact_summary["artifact_errors"] = artifact_errors
return compact_summary
def _prepare_omniinteract_batch(input_requests: list[SampleRequest]) -> None:
roots: set[Path] = set()
for sample in input_requests:
if not isinstance(sample, OmniInteractSampleRequest):
continue
root = sample.omniinteract_options.output_root.resolve()
roots.add(root)
clear_case_artifacts(sample.omniinteract_options.output_root, sample.omniinteract_case)
for root in roots:
clear_batch_artifacts(root)
def _daily_omni_repo_from_args(args) -> str | None:
"""Resolve HuggingFace repo id for Daily-Omni from CLI args.
@@ -250,6 +364,7 @@ def get_samples(args, tokenizer):
"ttsd",
"sound-effect",
)
is_omniinteract = args.dataset_name == "omniinteract"
# Check if we need to handle omni-related backends/datasets
is_omni_backend = args.backend in [
@@ -258,12 +373,68 @@ def get_samples(args, tokenizer):
"openai-realtime-duplex",
"daily-omni",
]
is_omni_dataset = is_daily_omni or is_seed_tts or args.dataset_name == "random-mm"
is_omni_dataset = is_daily_omni or is_seed_tts or is_omniinteract or args.dataset_name == "random-mm"
if not is_omni_backend and not is_omni_dataset:
# Not an omni-related request, delegate to original implementation
return get_samples_old(args, tokenizer)
if is_omniinteract:
dataset_path = getattr(args, "dataset_path", None)
data_root = None
dataset_repo = DEFAULT_OMNIINTERACT_REPO
if dataset_path:
candidate = Path(dataset_path).expanduser()
if candidate.exists() or candidate.is_absolute():
data_root = str(candidate)
else:
dataset_repo = str(dataset_path)
dataset = OmniInteractDataset(
data_root=data_root,
dataset_repo=dataset_repo,
subsets=tuple(getattr(args, "omniinteract_subsets")),
random_seed=args.seed,
disable_shuffle=getattr(args, "disable_shuffle", False),
)
options = OmniInteractSessionOptions(
output_root=Path(getattr(args, "omniinteract_output_dir")),
timeout_s=float(getattr(args, "omniinteract_timeout_s")),
media_timeout_s=float(getattr(args, "omniinteract_media_timeout_s")),
ref_audio=str(getattr(args, "omniinteract_ref_audio")),
require_response=bool(getattr(args, "omniinteract_require_response")),
max_video_duration_s=float(getattr(args, "omniinteract_max_video_duration_s")),
)
requests = dataset.sample(
tokenizer,
args.num_prompts,
request_id_prefix=args.request_id_prefix,
options=options,
)
if not requests:
raise ValueError("No OmniInteract sessions were selected")
encoded_ref_audio = reference_audio_data_url(options.ref_audio)
assert encoded_ref_audio is not None
for request in requests:
case = request.omniinteract_case
assert isinstance(request, OmniInteractSampleRequest) and case is not None
duration, pcm, frames = prepare_media(
case.video_path,
VIDEO_FPS,
timeout_s=options.media_timeout_s,
max_duration_s=options.max_video_duration_s,
)
if not any(frames):
raise ValueError(f"No video frames were decoded from {case.video_path}")
request.omniinteract_prepared_input = OmniInteractPreparedInput(
duration_s=duration,
pcm16=pcm,
video_frames=tuple(frames),
ref_audio_data_url=encoded_ref_audio,
)
# Replace OmniInteract's ``0`` (all) with the measured request count.
args.num_prompts = len(requests)
return requests
# Handle Daily-Omni dataset
if is_daily_omni:
# Support:
@@ -476,6 +647,7 @@ class MixRequestFuncOutput(RequestFuncOutput):
image_pixels: int = 0
denoise_step_latency_ms: float = 0.0
text_latency: float = 0.0
tpot_measured: bool = True
#: Worst-case streaming-audio underrun (wall-clock seconds the player
#: would have been starved). Populated by the audio-speech backend; ``0.0``
#: for backends that do not run continuity analysis.
@@ -1313,11 +1485,115 @@ async def async_request_openai_audio_speech(
def _realtime_websocket_url(api_url: str) -> str:
parts = urlsplit(api_url)
scheme = "wss" if parts.scheme == "https" else "ws"
query = dict(parse_qsl(parts.query, keep_blank_values=True))
query.setdefault("duplex", "1")
return urlunsplit((scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
return build_realtime_url(api_url, None, native_duplex=None)
def _nonnegative_number(value: object) -> bool:
return isinstance(value, int | float) and not isinstance(value, bool) and np.isfinite(value) and value >= 0
async def _async_request_omniinteract(
request_func_input: RequestFuncInput,
*,
pbar: tqdm | None,
) -> MixRequestFuncOutput:
case = getattr(request_func_input, "omniinteract_case", None)
options = getattr(request_func_input, "omniinteract_options", None)
output = MixRequestFuncOutput()
output.prompt_len = request_func_input.prompt_len
output.start_time = time.perf_counter()
try:
if case is None or not isinstance(options, OmniInteractSessionOptions):
raise ValueError("OmniInteract RequestFuncInput is missing its dataset session layout")
prepared_input = getattr(request_func_input, "omniinteract_prepared_input", None)
if not isinstance(prepared_input, OmniInteractPreparedInput):
raise ValueError("OmniInteract media must be prepared before benchmark request timing")
headers = _get_headers()
_update_headers_common(headers, request_func_input)
config = OmniInteractBenchmarkConfig(
endpoint=request_func_input.api_url,
model=request_func_input.model_name or request_func_input.model,
output_root=options.output_root,
timeout_s=options.timeout_s,
media_timeout_s=options.media_timeout_s,
max_video_duration_s=options.max_video_duration_s,
ref_audio=options.ref_audio,
require_response=options.require_response,
extra_headers=headers or None,
extra_body=dict(request_func_input.extra_body or {}) or None,
)
case_result = await run_omniinteract_case(
case,
config,
request_index=request_func_input.request_id or uuid.uuid4().hex,
capture_artifacts=request_func_input.request_id is not None,
prepared_input=prepared_input,
)
output.latency = case_result.latency_s
output.generated_text = case_result.transcript
output.output_tokens = case_result.output_tokens
output.audio_duration = case_result.audio_bytes / (24_000 * 2)
output.audio_frames = case_result.audio_bytes // 2
session_metrics = case_result.duplex_session_metrics
output.ttft = float(session_metrics.get("mean_ttft_ms") or 0.0) / 1000.0
output.audio_ttfp = float(session_metrics.get("mean_ttfp_ms") or 0.0) / 1000.0
output.audio_rtf = float(session_metrics.get("mean_rtf") or 0.0)
stages: list[tuple[int, dict[str, object]]] = []
for request_metric in case_result.duplex_request_metrics:
stage0 = request_metric.get("stage0_tokens")
if not isinstance(stage0, dict):
continue
token_count = stage0.get("output_token_count")
if isinstance(token_count, int) and not isinstance(token_count, bool) and token_count >= 0:
stages.append((token_count, stage0))
timed = [(count - 1, stage) for count, stage in stages if count > 1]
complete = (
len(stages) == len(case_result.duplex_request_metrics)
and sum(count for count, _ in stages) == output.output_tokens
)
exact = complete and all(
isinstance(values := stage.get("itls_ms"), list)
and len(values) == intervals
and all(_nonnegative_number(value) for value in values)
for intervals, stage in timed
)
output.itl = [float(value) / 1000.0 for _, stage in timed for value in stage["itls_ms"]] if exact else []
if output.itl:
output.text_latency = output.ttft + sum(output.itl)
elif complete and timed and all(_nonnegative_number(stage.get("tpot_ms")) for _, stage in timed):
intervals = sum(count for count, _ in timed)
weighted_tpot = sum(float(stage["tpot_ms"]) * count for count, stage in timed) / intervals / 1000.0
output.text_latency = output.ttft + weighted_tpot * (output.output_tokens - 1)
else:
output.text_latency = output.ttft
output.tpot_measured = False
if output.output_tokens > 1 or (output.output_tokens == 0 and output.generated_text):
logger.warning(
"OmniInteract session %s omitted complete engine token timing; standard TPOT/ITL are unavailable",
case_result.session_id,
)
output.duplex_request_metrics = case_result.duplex_request_metrics
output.duplex_session_metrics = session_metrics
output.success = case_result.success
output.error = case_result.error
setattr(output, "omniinteract_case", case)
setattr(output, "omniinteract_case_result", case_result)
except Exception:
output.success = False
output.error = traceback.format_exc()
logger.error("OmniInteract Realtime request failed: %s", output.error)
if case is not None:
case_result = OmniInteractCaseResult(
subset=case.subset,
video=str(case.video_path),
output_dir="",
error=output.error,
)
setattr(output, "omniinteract_case", case)
setattr(output, "omniinteract_case_result", case_result)
if pbar:
pbar.update(1)
return output
async def async_request_openai_realtime_duplex(
@@ -1325,8 +1601,9 @@ async def async_request_openai_realtime_duplex(
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
) -> MixRequestFuncOutput:
"""Run one or more Seed-TTS turns through one explicit Realtime TTS session."""
del session
if getattr(request_func_input, "omniinteract_case", None) is not None:
return await _async_request_omniinteract(request_func_input, pbar=pbar)
output = MixRequestFuncOutput()
output.prompt_len = request_func_input.prompt_len
output.start_time = time.perf_counter()
@@ -1387,8 +1664,10 @@ async def async_request_openai_realtime_duplex(
)
await client.send({"type": "response.create"})
await wait_for(
lambda: client.events.count("response.done") > done_before
or len(client.events.errors()) > errors_before,
lambda: (
client.events.count("response.done") > done_before
or len(client.events.errors()) > errors_before
),
timeout_s=180.0,
label=f"Seed-TTS Realtime TTS turn {request_index} response.done",
)
@@ -1433,19 +1712,7 @@ async def async_request_openai_realtime_duplex(
channels=1,
)
)
turn_transcripts.append(
"".join(
str(event.get("delta") or "")
for event in client.events.events
if client.events.response_id(event) == response_id
and event.get("type")
in {
"response.audio_transcript.delta",
"response.output_text.delta",
"response.text.delta",
}
)
)
turn_transcripts.append(client.events.response_text(response_id))
await client.acknowledge_playback()
request_finished_at = time.perf_counter()
session_metrics = summarize_session_request_metrics(
@@ -1520,6 +1787,7 @@ from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, ge
from vllm_omni.benchmarks.metrics.metrics import (
MultiModalsBenchmarkMetrics,
calculate_metrics,
has_metric_samples,
)
# ruff: noqa: E402
@@ -1626,6 +1894,7 @@ async def benchmark(
)
_attach_daily_omni_to_request_func_input(input_requests[0], test_input)
_attach_seed_tts_to_request_func_input(input_requests[0], test_input)
_attach_omniinteract_to_request_func_input(input_requests[0], test_input)
if ready_check_timeout_sec > 0:
test_output = await wait_for_endpoint(
@@ -1741,6 +2010,7 @@ async def benchmark(
print(f"Probe request rate: {probe_request_rate} req/s")
probe_task = asyncio.create_task(probe_loop())
_prepare_omniinteract_batch(input_requests)
benchmark_start_time = time.perf_counter()
tasks: list[asyncio.Task] = []
@@ -1801,6 +2071,7 @@ async def benchmark(
)
_attach_daily_omni_to_request_func_input(request, request_func_input)
_attach_seed_tts_to_request_func_input(request, request_func_input)
_attach_omniinteract_to_request_func_input(request, request_func_input)
tasks.append(
asyncio.create_task(limited_request_func(request_func_input=request_func_input, session=session, pbar=pbar))
)
@@ -1815,6 +2086,8 @@ async def benchmark(
benchmark_duration = time.perf_counter() - benchmark_start_time
omniinteract_summary = _finalize_omniinteract_batch(input_requests, outputs)
if task_type == TaskType.GENERATION:
metrics, actual_output_lens = calculate_metrics(
input_requests=input_requests,
@@ -1839,6 +2112,13 @@ async def benchmark(
actual_output_lens = 0
if isinstance(metrics, MultiModalsBenchmarkMetrics):
def measured_ttft(output: RequestFuncOutput) -> float | None:
session_metrics = getattr(output, "duplex_session_metrics", None)
if isinstance(session_metrics, dict) and session_metrics.get("mean_ttft_ms") is None:
return None
return output.ttft
result = {
"duration": benchmark_duration,
"completed": metrics.completed,
@@ -1859,7 +2139,7 @@ async def benchmark(
"input_lens": [output.prompt_len for output in outputs],
"start_times": [output.start_time for output in outputs],
"output_lens": actual_output_lens,
"ttfts": [output.ttft for output in outputs],
"ttfts": [measured_ttft(output) for output in outputs],
"itls": [output.itl for output in outputs],
"generated_texts": [output.generated_text for output in outputs],
"errors": [output.error for output in outputs],
@@ -1893,6 +2173,8 @@ async def benchmark(
]
if duplex_session_metrics:
result["duplex_session_metrics"] = duplex_session_metrics
if omniinteract_summary is not None:
result["omniinteract"] = omniinteract_summary
from vllm_omni.benchmarks.data_modules.daily_omni_eval import (
compute_daily_omni_accuracy_metrics,
@@ -1949,6 +2231,8 @@ async def benchmark(
# metric.
if metric_attribute_name not in result_percentile_metrics:
return
if not has_metric_samples(metrics, metric_attribute_name):
return
# No text tokens generated (e.g. pure TTS speech endpoint): per-token
# latency metrics (ttft/tpot/itl) are undefined, so skip them.
is_text_token_metric = not (metric_attribute_name == "e2el" or metric_attribute_name.startswith("audio"))
+15 -1
View File
@@ -1,6 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import argparse
import asyncio
import contextlib
import os
from pathlib import Path
from typing import Any
from vllm.benchmarks.serve import main_async
@@ -9,6 +16,7 @@ from vllm.benchmarks.serve import main_async
# This monkey-patches vllm.benchmarks.datasets.get_samples before it's used
# Must be imported before any vllm.benchmarks module usage
import vllm_omni.benchmarks.patch.patch # noqa: F401
from vllm_omni.benchmarks.omniinteract import omniinteract_output_lock
from vllm_omni.benchmarks.patch.patch import (
maybe_enable_stage_metrics,
set_print_stage,
@@ -28,4 +36,10 @@ def main(args: argparse.Namespace) -> dict[str, Any]:
getattr(args, "extra_body", None),
enabled=should_request_stage_metrics(args),
)
return asyncio.run(main_async(args))
lock = (
omniinteract_output_lock(Path(args.omniinteract_output_dir))
if getattr(args, "dataset_name", None) == "omniinteract"
else contextlib.nullcontext()
)
with lock:
return asyncio.run(main_async(args))
@@ -0,0 +1,291 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Reviewed inventory of environment variables used by vLLM-Omni.
This module classifies environment-variable names; it does not read them or
make every classified name part of the public configuration contract. In
particular, model-specific variables remain transitional until their recorded
disposition is implemented.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from enum import Enum
from types import MappingProxyType
class EnvironmentVariableCategory(str, Enum):
"""Ownership and documentation boundary for an environment variable."""
PUBLIC_OMNI = "public_omni"
INHERITED_VLLM = "inherited_vllm"
PLATFORM_EXTERNAL = "platform_external"
MODEL_SPECIFIC = "model_specific"
BENCHMARK_TRANSITIONAL = "benchmark_transitional"
INTERNAL = "internal"
class ModelEnvironmentVariableDisposition(str, Enum):
"""Required destination for a model-specific environment variable."""
PROMOTE = "promote"
REQUEST_SCOPE = "request_scope"
EXTERNAL = "external"
INTERNALIZE = "internalize"
DEPRECATE_REMOVE = "deprecate_remove"
@dataclass(frozen=True)
class EnvironmentVariableClassification:
"""Classification metadata for one statically identifiable name."""
name: str
category: EnvironmentVariableCategory
model_disposition: ModelEnvironmentVariableDisposition | None = None
redact_value: bool = False
@property
def is_public_omni(self) -> bool:
return self.category is EnvironmentVariableCategory.PUBLIC_OMNI
# New entries must use the VLLM_OMNI_ prefix. The regression test explicitly
# grandfathers the older public names below.
_PUBLIC_OMNI = (
"DIFFUSION_ATTENTION_BACKEND",
"DIFFUSION_CACHE_ADAPTER",
"DIFFUSION_CACHE_BACKEND",
"OMNI_DIFFUSION_PROMPT_EMBED_CACHE",
"OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE",
"OMNI_DIFFUSION_SESSION_STATE_MANAGER",
"OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS",
"SPEAKER_MAX_UPLOADED",
"SPEAKER_SAMPLES_DIR",
"VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT",
"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S",
"VLLM_OMNI_ORCH_MONITOR_PATH",
"VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY",
"VLLM_OMNI_SERVER_STORAGE__FILE_TTL",
"VLLM_OMNI_SERVER_STORAGE__PATH",
"VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL",
"VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP",
"VLLM_OMNI_STORAGE_MAX_CONCURRENCY",
"VLLM_OMNI_STORAGE_PATH",
"VLLM_OMNI_USE_QUACK_FP8",
"VLLM_OMNI_VIDEO_SYNC_TIMEOUT",
"VLLM_VIDEO_ASYNC_CHUNK",
"VLLM_VIDEO_AUDIO_DELTA_MODE",
)
_INHERITED_VLLM = (
"CUDA_HOME",
"CUDA_VISIBLE_DEVICES",
"LOCAL_RANK",
"NO_COLOR",
"VLLM_ALLOW_LONG_MAX_MODEL_LEN",
"VLLM_BATCH_INVARIANT",
"VLLM_CACHE_ROOT",
"VLLM_DISABLE_LOG_LOGO",
"VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE",
"VLLM_HTTP_TIMEOUT_KEEP_ALIVE",
"VLLM_LOGGING_COLOR",
"VLLM_MOONCAKE_BOOTSTRAP_PORT",
"VLLM_PLUGINS",
"VLLM_ROCM_USE_AITER",
"VLLM_ROCM_USE_AITER_RMSNORM",
"VLLM_TUNED_CONFIG_FOLDER",
"VLLM_USE_MODELSCOPE",
"VLLM_USE_OINK_OPS",
"VLLM_WORKER_MULTIPROC_METHOD",
"VLLM_XPU_USE_SAMPLER_KERNEL",
)
_PLATFORM_EXTERNAL = (
"ASCEND_LAUNCH_BLOCKING",
"ASCEND_RT_VISIBLE_DEVICES",
"FLASHINFER_DISABLE_VERSION_CHECK",
"HF_HOME",
"HF_HUB_DISABLE_XET",
"HF_TOKEN",
"HIP_VISIBLE_DEVICES",
"HUGGINGFACE_HUB_TOKEN",
"MASTER_ADDR",
"MASTER_PORT",
"MINDIE_SD_FA_TYPE",
"MORI_RDMA_DEVICES",
"MUSA_VISIBLE_DEVICES",
"MV2_COMM_WORLD_LOCAL_RANK",
"NCCL_ASYNC_ERROR_HANDLING",
"OMPI_COMM_WORLD_LOCAL_RANK",
"OPENAI_API_KEY",
"PYTHONPATH",
"QUACK_CACHE_DIR",
"RANK",
"RANK_ID",
"RAY_RAYLET_PID",
"RDMA_DEVICE_NAME",
"SAGE_ATTN_TENSOR_LAYOUT",
"VLLM_LOCAL_RANK",
"WORLD_SIZE",
"XDG_CACHE_HOME",
)
_MODEL_PROMOTE = (
"AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH",
"COSMOS3_SOUND_TOKENIZER_CONFIG_PATH",
"COSMOS3_SOUND_TOKENIZER_PATH",
"COSYVOICE3_ESTIMATOR_ONNX",
"COSYVOICE3_TRT",
"COSYVOICE3_TRT_CACHE",
"DYNIN_CONFIG_PATH",
"DYNIN_MAGVIT_REMOTE_CODE_LOCAL_FILES_ONLY",
"DYNIN_MAGVIT_REMOTE_CODE_REPO_ID",
"DYNIN_MAGVIT_REMOTE_CODE_REVISION",
"DYNIN_REMOTE_CODE_LOCAL_FILES_ONLY",
"DYNIN_REMOTE_CODE_REPO_ID",
"DYNIN_REMOTE_CODE_REVISION",
"HIGGS_AUDIO_TOKENIZER_PATH",
"INTERNVLA_A1_COSMOS_DECODER_PATH",
"INTERNVLA_A1_COSMOS_DIR",
"INTERNVLA_A1_COSMOS_ENCODER_PATH",
"INTERNVLA_A1_PROCESSOR_DIR",
"MIMO_AUDIO_TOKENIZER_CUDA_GRAPH",
"MIMO_AUDIO_TOKENIZER_DEVICE",
"MIMO_AUDIO_TOKENIZER_PATH",
"MING_CFM_CUDAGRAPH",
"MINICPMO_TOKEN2WAV_TRT",
"MINICPMO_TOKEN2WAV_TRT_CACHE",
"MINICPMO_TOKEN2WAV_TRT_DTYPE",
"NEMOTRON_VOICECHAT_LLM_PATH",
"OMNIVOICE_CUDA_GRAPH",
"OMNIVOICE_TF32",
"VLLM_GEPARD_CHUNK_FRAMES",
"VLLM_GEPARD_FIRST_CHUNK_FRAMES",
"VLLM_GEPARD_LOOKBACK_FRAMES",
"VLLM_OMNI_FISH_KVCACHE_ATTN",
"VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER",
)
_MODEL_REQUEST_SCOPE = (
"DOTS_TTS_DIT_NUM_STEPS",
"GR00T_NOISE_SEED",
"STEP_AUDIO2_DEFAULT_PROMPT_WAV",
"VLLM_GEPARD_END_FADE_MS",
"VLLM_GEPARD_END_SILENCE_MS",
"VLLM_GEPARD_GREEDY",
)
# No audited model variable is currently owned by a supported third-party
# contract. Keep the disposition explicit so future reviews cannot silently
# fold an externally owned name into "promote".
_MODEL_EXTERNAL: tuple[str, ...] = ()
_MODEL_INTERNALIZE = (
"DOTS_TTS_BETA_TRACE",
"DYNIN_S2U_VENDOR_ROOT",
"DZ_PHASE_TIMING",
"MING_TTS_STAGE1_FINAL_LOG",
"MOSS_TTS_DEBUG_STOP",
"NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE",
"VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS",
"VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS_LOG_EVERY",
"VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS",
"VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_FILE",
"VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_LOG_EVERY",
)
_MODEL_DEPRECATE_REMOVE = (
"HIGGS_AUDIO_V2_TOKENIZER_PATH",
"VLLM_OMNI_LINGBOT_ACTION_ROOT",
"VLLM_OMNI_VOXCPM_CODE_PATH",
"VLLM_USE_FLASHINFER_MOE_FP16",
"XCODEC1_PATH",
"model_stage",
)
_BENCHMARK_TRANSITIONAL = (
"DAILY_OMNI_EXTRACT_MODE",
"DAILY_OMNI_SAVE_EVAL_ITEMS",
"MAX_NUM_FRAMES",
"SEED_TTS_EVAL_DEVICE",
"SEED_TTS_HF_WHISPER_MODEL",
"SEED_TTS_SIM_DEVICE",
"SEED_TTS_SIM_EVAL",
"SEED_TTS_UTMOS_EVAL",
"SEED_TTS_UTMOS_HF_REPO",
"SEED_TTS_UTMOS_JIT_FILE",
"SEED_TTS_WAVLM_MAX_SECONDS",
"SEED_TTS_WAVLM_MIN_SAMPLES",
"SEED_TTS_WAVLM_MODEL",
"SEED_TTS_WER_EVAL",
"SEED_TTS_WER_SAVE_AUDIO_DIR",
"SEED_TTS_WER_SAVE_ITEMS",
"VLLM_DAILY_OMNI_MEDIA_REPO",
"VLLM_OMNI_BENCH_AUDIO_CHANNELS",
"VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S",
"VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE",
)
_INTERNAL = (
"VLLM_OMNI_DLO_DP_WAVE_TIMEOUT",
"VLLM_OMNI_REPLICA_ID",
)
_REDACTED = frozenset(
{
"HF_TOKEN",
"HUGGINGFACE_HUB_TOKEN",
"OPENAI_API_KEY",
}
)
def _build_inventory() -> Mapping[str, EnvironmentVariableClassification]:
inventory: dict[str, EnvironmentVariableClassification] = {}
def add(
names: tuple[str, ...],
category: EnvironmentVariableCategory,
disposition: ModelEnvironmentVariableDisposition | None = None,
) -> None:
for name in names:
if name in inventory:
raise ValueError(f"Environment variable {name!r} is classified more than once")
inventory[name] = EnvironmentVariableClassification(
name=name,
category=category,
model_disposition=disposition,
redact_value=name in _REDACTED,
)
add(_PUBLIC_OMNI, EnvironmentVariableCategory.PUBLIC_OMNI)
add(_INHERITED_VLLM, EnvironmentVariableCategory.INHERITED_VLLM)
add(_PLATFORM_EXTERNAL, EnvironmentVariableCategory.PLATFORM_EXTERNAL)
add(_MODEL_PROMOTE, EnvironmentVariableCategory.MODEL_SPECIFIC, ModelEnvironmentVariableDisposition.PROMOTE)
add(
_MODEL_REQUEST_SCOPE,
EnvironmentVariableCategory.MODEL_SPECIFIC,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE,
)
add(_MODEL_EXTERNAL, EnvironmentVariableCategory.MODEL_SPECIFIC, ModelEnvironmentVariableDisposition.EXTERNAL)
add(
_MODEL_INTERNALIZE,
EnvironmentVariableCategory.MODEL_SPECIFIC,
ModelEnvironmentVariableDisposition.INTERNALIZE,
)
add(
_MODEL_DEPRECATE_REMOVE,
EnvironmentVariableCategory.MODEL_SPECIFIC,
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE,
)
add(_BENCHMARK_TRANSITIONAL, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL)
add(_INTERNAL, EnvironmentVariableCategory.INTERNAL)
return MappingProxyType(dict(sorted(inventory.items())))
# All reviewed environment-variable names, indexed by exact name. This is
# classification metadata, not the executable value registry used by
# ``vllm.envs``.
ENVIRONMENT_VARIABLE_INVENTORY: Mapping[str, EnvironmentVariableClassification] = _build_inventory()
+8
View File
@@ -699,6 +699,7 @@ class _DiffusionConfigProjection:
dlo_resident_layers: int = Field(default=0, ge=0)
host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
host_weight_runtime_root: str | None = None
dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)
pin_cpu_memory: bool = True
diffusion_compile_granularity: Literal["regional", "full"] = "regional"
diffusion_compile_dynamic: bool = Field(default=True, strict=True)
@@ -756,6 +757,7 @@ class _DiffusionConfigProjection:
TransformerConfig,
build_attention_config,
parse_kv_cache_skip_selector,
validate_dlo_host_registration_options,
validate_host_weight_runtime_options,
)
from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
@@ -838,6 +840,12 @@ class _DiffusionConfigProjection:
mode=self.host_weight_runtime_mode,
root=self.host_weight_runtime_root,
)
self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
limit_gib=self.dlo_host_registration_limit_gib,
enable_dlo=self.enable_distributed_layerwise_offload,
use_allgather=self.dlo_use_allgather,
hwr_mode=self.host_weight_runtime_mode,
)
if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
raise ValueError(
+1
View File
@@ -391,6 +391,7 @@ class StageDeployConfig:
dlo_resident_layers: int | None = None
host_weight_runtime_mode: str | None = None
host_weight_runtime_root: str | None = None
dlo_host_registration_limit_gib: float | None = None
# Diffusion-specific debug and observability knobs.
enable_diffusion_pipeline_profiler: bool | None = None
+2 -1
View File
@@ -1,6 +1,7 @@
# Wan2.2 TI2V-5B deploy.
#
# Single GPU by default. Increase tensor_parallel_size and devices for A14B MoE variants.
# Request-level batching uses scheduler max_num_seqs, not diffusion_batch_size.
async_chunk: false
trust_remote_code: true
distributed_executor_backend: mp
@@ -8,7 +9,7 @@ distributed_executor_backend: mp
stages:
- stage_id: 0
devices: "0"
max_num_seqs: 1
max_num_seqs: 4 # in-flight request batch width (CLI: --max-num-seqs)
tensor_parallel_size: 1
enforce_eager: true
# Keep VAE tiling on by default for higher resolutions / longer clips.
+25
View File
@@ -113,6 +113,22 @@ def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
def validate_dlo_host_registration_options(
*,
limit_gib: object,
enable_dlo: bool,
use_allgather: bool,
hwr_mode: object,
) -> float:
"""Validate the optional transport budget without probing CUDA or HWR."""
value = float(limit_gib)
if not math.isfinite(value) or value < 0:
raise ValueError("dlo_host_registration_limit_gib must be finite and >= 0")
if value and (not enable_dlo or use_allgather or hwr_mode == "disabled"):
raise ValueError("dlo_host_registration_limit_gib requires enabled no-AllGather DLO and Host Weight Runtime")
return value
def parse_kv_cache_skip_selector(
selector: str | list[int] | tuple[int, ...] | set[int] | None,
) -> set[int] | None:
@@ -784,6 +800,9 @@ class OmniDiffusionConfig:
# existing loader/storage path.
host_weight_runtime_mode: str = "disabled"
host_weight_runtime_root: str | None = None
# Optional per-worker ceiling for registering final-layout HWR mappings.
# Zero adds no ceiling; pin_cpu_memory controls whether registration is tried.
dlo_host_registration_limit_gib: float = 0.0
pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading
@@ -1164,6 +1183,12 @@ class OmniDiffusionConfig:
mode=self.host_weight_runtime_mode,
root=self.host_weight_runtime_root,
)
self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
limit_gib=self.dlo_host_registration_limit_gib,
enable_dlo=self.enable_distributed_layerwise_offload,
use_allgather=self.dlo_use_allgather,
hwr_mode=self.host_weight_runtime_mode,
)
if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
raise ValueError(
@@ -27,7 +27,10 @@ logger = init_logger(__name__)
if TYPE_CHECKING:
from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext
from vllm_omni.diffusion.model_loader.host_weights.source_identity import PreparedWeightSource
from vllm_omni.diffusion.model_loader.host_weights.source_identity import (
NodeSourceDigestCache,
PreparedWeightSource,
)
class _HWRCommitError(RuntimeError):
@@ -215,6 +218,7 @@ class HWRLoaderMixin:
*,
load_format: str,
sources: Sequence[object],
source_digest_cache: NodeSourceDigestCache | None = None,
) -> FinalLayoutIdentityContext:
from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
@@ -293,6 +297,7 @@ class HWRLoaderMixin:
prepared_sources=prepared_sources,
request=request,
policy=FINAL_LAYOUT_BF16_POLICY,
source_digest_cache=source_digest_cache,
)
def _resolve_hwr(
@@ -306,7 +311,10 @@ class HWRLoaderMixin:
sources: Sequence[object],
) -> dict[str, object] | None:
"""Resolve an eligible no-AllGather final-layout HWR transaction."""
from vllm_omni.diffusion.model_loader.host_weights import FinalLayoutTensorRestorer
from vllm_omni.diffusion.model_loader.host_weights import (
FinalLayoutTensorRestorer,
NodeSourceDigestCache,
)
from vllm_omni.host_weight_runtime import (
HostWeightLeaseCarrier,
HostWeightRuntime,
@@ -347,7 +355,6 @@ class HWRLoaderMixin:
raise ValueError(f"required Host Weight Runtime path is ineligible: {message}")
logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message)
return None
context = self._build_hwr_context(model, modules, load_format=load_format, sources=sources)
root_value = getattr(self.od_config, "host_weight_runtime_root", None)
if not isinstance(root_value, str) or not root_value.strip():
raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
@@ -362,6 +369,17 @@ class HWRLoaderMixin:
),
)
)
source_digest_cache = NodeSourceDigestCache(
root,
timeout_seconds=runtime.config.wait.coordination_timeout_seconds,
)
context = self._build_hwr_context(
model,
modules,
load_format=load_format,
sources=sources,
source_digest_cache=source_digest_cache,
)
resolution = runtime.resolve(context.identity)
state: dict[str, object] = {
"mode": mode,
@@ -19,7 +19,7 @@ from .producers import (
FinalLayoutBF16Producer,
)
from .restorer import FinalLayoutTensorRestorer
from .source_identity import PreparedWeightSource, WeightSourceKind
from .source_identity import NodeSourceDigestCache, PreparedWeightSource, WeightSourceKind
__all__ = [
"FINAL_LAYOUT_BF16_POLICY",
@@ -32,6 +32,7 @@ __all__ = [
"FinalLayoutRequest",
"FinalLayoutTensorRestorer",
"ImplementationIdentity",
"NodeSourceDigestCache",
"PreparedWeightSource",
"WeightSourceKind",
"build_final_layout_identity",
@@ -27,6 +27,7 @@ from .contracts import (
)
from .source_identity import (
FinalLayoutSourceContext,
NodeSourceDigestCache,
PreparedWeightSource,
resolve_final_layout_source_identity,
)
@@ -78,6 +79,7 @@ def build_final_layout_identity(
prepared_sources: Sequence[PreparedWeightSource],
request: FinalLayoutRequest,
policy: FinalLayoutTensorPolicy,
source_digest_cache: NodeSourceDigestCache | None = None,
) -> FinalLayoutIdentityContext:
"""Build exact source, representation, layout, and implementation identity."""
if not isinstance(request, FinalLayoutRequest):
@@ -99,6 +101,7 @@ def build_final_layout_identity(
prepared_sources,
model_id=request.model_id,
target_names=target_names,
digest_cache=source_digest_cache,
)
dit_names = tuple(name for name, _ in dit_modules)
model_contracts = tuple(binding.contract for binding in bindings)
@@ -286,7 +286,7 @@ class FinalLayoutBF16Producer:
for record in shard
)
file_name = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
with writer.open_tensor_file(file_name, specs) as output:
with writer.open_tensor_file(file_name, specs, ordered=True) as output:
for record in shard:
output.write_tensor(record.name, record.tensor.detach())
@@ -5,9 +5,13 @@
from __future__ import annotations
import hashlib
import json
import os
import stat
import time
import uuid
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
@@ -15,6 +19,7 @@ from pathlib import Path
import regex as re
from vllm_omni.host_weight_runtime import CanonicalJson, WeightSourceIdentity
from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, FileLockTimeoutError
from vllm_omni.host_weight_runtime.identity import canonical_json
from .contracts import (
@@ -101,6 +106,140 @@ class _FileSnapshot:
)
@dataclass(frozen=True)
class _FileState:
size: int
device: int
inode: int
mtime_ns: int
ctime_ns: int
symlink_target: str | None
def to_dict(self) -> dict[str, object]:
return {
"size": self.size,
"device": self.device,
"inode": self.inode,
"mtime_ns": self.mtime_ns,
"ctime_ns": self.ctime_ns,
"symlink_target": self.symlink_target,
}
def _observe(path: Path) -> _FileState:
try:
current = path.stat()
except OSError as exc:
raise ValueError(f"cannot stat canonical weight file {path}") from exc
if not stat.S_ISREG(current.st_mode):
raise ValueError(f"canonical weight source is not a regular file: {path}")
return _FileState(
size=current.st_size,
device=current.st_dev,
inode=current.st_ino,
mtime_ns=current.st_mtime_ns,
ctime_ns=current.st_ctime_ns,
symlink_target=os.readlink(path) if path.is_symlink() else None,
)
def _sha256_file(path: Path, expected: _FileState) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
before = os.fstat(handle.fileno())
while chunk := handle.read(_HASH_CHUNK_BYTES):
digest.update(chunk)
after = os.fstat(handle.fileno())
observed = (before.st_size, before.st_dev, before.st_ino, before.st_mtime_ns, before.st_ctime_ns)
completed = (after.st_size, after.st_dev, after.st_ino, after.st_mtime_ns, after.st_ctime_ns)
expected_stat = (expected.size, expected.device, expected.inode, expected.mtime_ns, expected.ctime_ns)
if observed != expected_stat or completed != expected_stat or _observe(path) != expected:
raise FinalLayoutContractError(
FinalLayoutContractCode.SOURCE_CHANGED,
f"canonical weight file changed while hashing: {path}",
)
return digest.hexdigest()
class NodeSourceDigestCache:
"""Share content-backed local-source fingerprints across node workers."""
def __init__(self, root: Path, *, timeout_seconds: float) -> None:
self.entries = root / "source-digests-v1" / "entries"
self.locks = root / "source-digests-v1" / "locks"
self.timeout_seconds = timeout_seconds
self._available = True
try:
self.entries.mkdir(parents=True, exist_ok=True, mode=0o700)
self.locks.mkdir(parents=True, exist_ok=True, mode=0o700)
except OSError:
self._available = False
@staticmethod
def _cached_content_id(entry: Path, path: Path, state: _FileState) -> str | None:
try:
cached = json.loads(entry.read_bytes())
if not isinstance(cached, dict):
return None
cached_state = cached["state"]
content_id = cached["content_id"]
record_sha256 = cached["record_sha256"]
except (FileNotFoundError, OSError, json.JSONDecodeError, KeyError, TypeError):
return None
record = {
"content_id": content_id,
"schema_version": cached.get("schema_version"),
"state": cached_state,
}
if (
record["schema_version"] != 1
or cached_state != state.to_dict()
or not isinstance(content_id, str)
or re.fullmatch(r"sha256:[0-9a-f]{64}", content_id) is None
or not isinstance(record_sha256, str)
or hashlib.sha256(canonical_json(record)).hexdigest() != record_sha256
or _observe(path) != state
):
return None
return content_id
def content_id(self, path: Path, state: _FileState) -> str:
if not self._available:
return f"sha256:{_sha256_file(path, state)}"
key = hashlib.sha256(os.fsencode(path.absolute())).hexdigest()
entry = self.entries / f"{key}.json"
content_id: str | None = None
try:
with FileLock(
self.locks / f"{key}.lock",
exclusive=True,
deadline=time.monotonic() + self.timeout_seconds,
):
if (cached := self._cached_content_id(entry, path, state)) is not None:
return cached
content_id = f"sha256:{_sha256_file(path, state)}"
record: dict[str, object] = {
"content_id": content_id,
"schema_version": 1,
"state": state.to_dict(),
}
record["record_sha256"] = hashlib.sha256(canonical_json(record)).hexdigest()
temporary = entry.with_name(f".{entry.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp")
try:
temporary.write_bytes(canonical_json(record))
os.chmod(temporary, 0o600)
os.replace(temporary, entry)
finally:
temporary.unlink(missing_ok=True)
return content_id
except (FileLockTimeoutError, OSError):
# Digest sharing is an optimization. Coordination or cache I/O
# failures must not make a readable canonical source unavailable.
if content_id is not None and _observe(path) == state:
return content_id
return f"sha256:{_sha256_file(path, state)}"
@dataclass(frozen=True)
class _SourceSnapshot:
semantic: CanonicalJson
@@ -129,14 +268,6 @@ class FinalLayoutSourceContext:
)
def _sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
while chunk := handle.read(_HASH_CHUNK_BYTES):
digest.update(chunk)
return digest.hexdigest()
def _hf_blob_content_id(
path: Path,
*,
@@ -200,7 +331,10 @@ def _logical_model_id(value: str) -> str:
return value
def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
def _snapshot_source(
source: PreparedWeightSource,
digest_cache: NodeSourceDigestCache | None,
) -> _SourceSnapshot:
root = source.resolved_root.expanduser().resolve()
if not root.is_dir():
raise ValueError(f"canonical source root is not a directory: {root}")
@@ -210,8 +344,7 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None:
immutable_revision = requested.lower()
files: list[_FileSnapshot] = []
for path in sorted(source.weight_files, key=lambda item: str(item)):
def snapshot_file(path: Path) -> _FileSnapshot:
candidate = path.expanduser()
if not candidate.is_absolute():
candidate = root / candidate
@@ -220,34 +353,34 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
relative_name = candidate.relative_to(root).as_posix()
except ValueError as exc:
raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc
try:
current = candidate.stat()
except OSError as exc:
raise ValueError(f"cannot stat canonical weight file {candidate}") from exc
if not stat.S_ISREG(current.st_mode):
raise ValueError(f"canonical weight source is not a regular file: {candidate}")
symlink_target = os.readlink(candidate) if candidate.is_symlink() else None
current = _observe(candidate)
content_id = _hf_blob_content_id(
candidate,
source=source,
source_root=root,
)
if content_id is None:
content_id = f"sha256:{_sha256_file(candidate)}"
files.append(
_FileSnapshot(
path=candidate,
relative_name=relative_name,
size=current.st_size,
device=current.st_dev,
inode=current.st_ino,
mtime_ns=current.st_mtime_ns,
ctime_ns=current.st_ctime_ns,
symlink_target=symlink_target,
content_id=content_id,
content_id = (
digest_cache.content_id(candidate, current)
if digest_cache is not None
else f"sha256:{_sha256_file(candidate, current)}"
)
return _FileSnapshot(
path=candidate,
relative_name=relative_name,
size=current.size,
device=current.device,
inode=current.inode,
mtime_ns=current.mtime_ns,
ctime_ns=current.ctime_ns,
symlink_target=current.symlink_target,
content_id=content_id,
)
paths = sorted(source.weight_files, key=lambda item: str(item))
with ThreadPoolExecutor(max_workers=min(8, len(paths))) as executor:
files = list(executor.map(snapshot_file, paths))
file_semantics = [file.semantic_dict() for file in files]
file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest()
revision = immutable_revision or f"content-{file_fingerprint}"
@@ -306,11 +439,12 @@ def resolve_final_layout_source_identity(
*,
model_id: str,
target_names: frozenset[str],
digest_cache: NodeSourceDigestCache | None = None,
) -> FinalLayoutSourceContext:
"""Resolve exact source identity for files covering the owned tensors."""
bindings = _resolve_target_sources(prepared_sources, target_names)
selected_sources = frozenset(bindings.values())
snapshot_by_source = {source: _snapshot_source(source) for source in selected_sources}
snapshot_by_source = {source: _snapshot_source(source, digest_cache) for source in selected_sources}
snapshots = tuple(sorted(snapshot_by_source.values(), key=lambda snapshot: snapshot.semantic.encoded))
source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots]
target_bindings = [
@@ -345,6 +479,7 @@ def resolve_final_layout_source_identity(
__all__ = [
"FinalLayoutSourceContext",
"NodeSourceDigestCache",
"PreparedWeightSource",
"WeightSourceKind",
"resolve_final_layout_source_identity",
+12 -2
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from abc import ABC, abstractmethod
from dataclasses import dataclass
@@ -10,7 +10,7 @@ import torch
from torch import nn
from vllm.logger import init_logger
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.data import OmniDiffusionConfig, validate_dlo_host_registration_options
logger = init_logger(__name__)
@@ -52,6 +52,9 @@ class OffloadConfig:
# blocks from the loader-selected host backing with H2D only.
dlo_use_allgather: bool = True
dlo_resident_layers: int = 0 # leading DiT layers kept on device
# Optional per-worker ceiling for registering an HWR mmap. Zero means no
# additional ceiling; pin_cpu_memory controls whether registration is tried.
dlo_host_registration_limit_gib: float = 0.0
@classmethod
def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -122,6 +125,12 @@ class OffloadConfig:
# requirements (concurrent requests, dummy run skip).
dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0),
enable_dlo=enable_distributed_layerwise_offload,
use_allgather=dlo_use_allgather,
hwr_mode=getattr(od_config, "host_weight_runtime_mode", "disabled"),
)
if dlo_resident_layers < 0:
raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
if dlo_resident_layers and dlo_use_allgather:
@@ -157,6 +166,7 @@ class OffloadConfig:
dp_size=dp_size,
dlo_use_allgather=dlo_use_allgather,
dlo_resident_layers=dlo_resident_layers,
dlo_host_registration_limit_gib=dlo_host_registration_limit_gib,
)
@@ -0,0 +1,248 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CUDA implementation of read-only host-mapping registration."""
from __future__ import annotations
import ctypes
import mmap
from collections import defaultdict
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Protocol
import torch
from vllm_omni.host_weight_runtime import MappedHostRegion
from .host_registration import (
HostRegistrationBudgetError,
HostRegistrationCleanupError,
HostRegistrationError,
)
_CUDA_HOST_REGISTER_READ_ONLY = 0x08
_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
class _CudaRuntime(Protocol):
def cudaHostRegister(self, address: int, size: int, flags: int) -> int: ...
def cudaHostUnregister(self, address: int) -> int: ...
def cudaGetErrorString(self, error: int) -> str | bytes: ...
def cudaGetLastError(self) -> int: ...
@dataclass(frozen=True)
class _AddressRange:
start: int
end: int
@property
def size(self) -> int:
return self.end - self.start
def _coalesce_ranges(
ranges: Sequence[tuple[int, int]],
page_size: int = mmap.PAGESIZE,
) -> tuple[_AddressRange, ...]:
"""Page-align and merge overlapping ranges from one backing mapping."""
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
aligned: list[_AddressRange] = []
for start, size in ranges:
if start <= 0 or size < 0:
raise ValueError(f"invalid host range start={start}, size={size}")
if size == 0:
continue
aligned_start = start - start % page_size
end = start + size
aligned_end = ((end + page_size - 1) // page_size) * page_size
aligned.append(_AddressRange(aligned_start, aligned_end))
merged: list[_AddressRange] = []
for region in sorted(aligned, key=lambda item: (item.start, item.end)):
if merged and region.start <= merged[-1].end:
previous = merged[-1]
merged[-1] = _AddressRange(previous.start, max(previous.end, region.end))
else:
merged.append(region)
return tuple(merged)
def _mapped_regions(regions: Sequence[MappedHostRegion]) -> tuple[_AddressRange, ...]:
"""Resolve page spans without merging unrelated artifact files."""
by_file: dict[str, list[tuple[int, int]]] = defaultdict(list)
for region in regions:
if not region.read_only:
raise HostRegistrationError(
f"CUDA direct H2D requires an immutable mapping, but {region.file_name!r} is writable"
)
by_file[region.file_name].append((region.address, region.length))
return tuple(span for ranges in by_file.values() for span in _coalesce_ranges(ranges))
def _error_message(runtime: _CudaRuntime, error: int) -> str:
try:
message = runtime.cudaGetErrorString(error)
except Exception:
return str(error)
if isinstance(message, bytes):
return message.decode(errors="replace")
return str(message)
def _consume_last_cuda_error(runtime: _CudaRuntime, expected_error: int) -> None:
"""Clear one handled CUDA Runtime error before returning to PyTorch."""
try:
get_last_error = getattr(runtime, "cudaGetLastError", None)
if get_last_error is None:
get_last_error = ctypes.CDLL(None).cudaGetLastError
get_last_error.argtypes = []
get_last_error.restype = ctypes.c_int
pending_error = int(get_last_error())
except (AttributeError, OSError, RuntimeError, TypeError, ValueError) as exc:
raise HostRegistrationError("cannot clear CUDA's pending error after host-registration failure") from exc
if pending_error not in (0, expected_error):
raise HostRegistrationError(
f"host registration returned CUDA error {expected_error}, but cudaGetLastError reported {pending_error}"
)
def _handled_error_message(runtime: _CudaRuntime, error: int) -> str:
error_code = int(error)
message = _error_message(runtime, error_code)
try:
_consume_last_cuda_error(runtime, error_code)
except HostRegistrationError as exc:
raise HostRegistrationError(f"{message}; {exc}") from exc
return message
def _supports_read_only_host_registration(runtime: _CudaRuntime) -> bool:
"""Query support required by immutable artifact mappings."""
try:
get_attribute = getattr(runtime, "cudaDeviceGetAttribute", None)
if get_attribute is None:
get_attribute = ctypes.CDLL(None).cudaDeviceGetAttribute
get_attribute.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int]
get_attribute.restype = ctypes.c_int
except (AttributeError, OSError) as exc:
raise HostRegistrationError("cannot query CUDA read-only host-registration support") from exc
supported = ctypes.c_int()
try:
error = get_attribute(
ctypes.byref(supported),
_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED,
torch.accelerator.current_device_index(),
)
except Exception as exc:
raise HostRegistrationError(f"cannot query CUDA read-only host-registration support: {exc}") from exc
if int(error) != 0:
raise HostRegistrationError(
f"cudaDeviceGetAttribute(read-only host registration) failed: {_handled_error_message(runtime, error)}"
)
return bool(supported.value)
class CudaHostRegistration:
"""Own CUDA registrations for immutable file-backed CPU regions."""
def __init__(self, runtime: _CudaRuntime, regions: tuple[_AddressRange, ...]) -> None:
self._runtime = runtime
self._regions = regions
@classmethod
def create(
cls,
regions: Sequence[MappedHostRegion],
*,
max_bytes: int | None,
) -> CudaHostRegistration:
mapped = _mapped_regions(regions)
total_bytes = sum(region.size for region in mapped)
if max_bytes is not None and total_bytes > max_bytes:
raise HostRegistrationBudgetError(
f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
)
if not mapped:
raise HostRegistrationError("no non-empty host ranges were available for registration")
if not torch.cuda.is_available():
raise HostRegistrationError("CUDA is not available")
try:
runtime = torch.cuda.cudart()
except Exception as exc:
raise HostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc
if not _supports_read_only_host_registration(runtime):
raise HostRegistrationError(
"CUDA device does not support read-only host registration required by immutable HWR mappings"
)
registered: list[_AddressRange] = []
try:
for region in mapped:
error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY)
if int(error) != 0:
raise HostRegistrationError(
"cudaHostRegister(read-only) failed for "
f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}"
)
registered.append(region)
except Exception as exc:
rollback_errors: list[str] = []
rollback_failed: list[_AddressRange] = []
for region in reversed(registered):
try:
error = runtime.cudaHostUnregister(region.start)
if int(error) != 0:
rollback_errors.append(
f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(runtime, error)}"
)
rollback_failed.append(region)
except Exception as rollback_exc:
rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
rollback_failed.append(region)
if rollback_errors:
active_registration = cls(runtime, tuple(reversed(rollback_failed)))
raise HostRegistrationCleanupError(
f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}",
active_registration=active_registration,
) from exc
if isinstance(exc, HostRegistrationError):
raise
raise HostRegistrationError(f"CUDA host registration raised: {exc}") from exc
return cls(runtime, mapped)
@property
def total_bytes(self) -> int:
return sum(region.size for region in self._regions)
@property
def region_count(self) -> int:
return len(self._regions)
def close(self) -> tuple[str, ...]:
"""Unregister every range and retain failures for a later retry."""
errors: list[str] = []
failed: list[_AddressRange] = []
for region in reversed(self._regions):
try:
error = self._runtime.cudaHostUnregister(region.start)
if int(error) != 0:
errors.append(
f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(self._runtime, error)}"
)
failed.append(region)
except Exception as exc:
errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}")
failed.append(region)
self._regions = tuple(reversed(failed))
return tuple(errors)
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Distributed Layerwise Offload backend with double-buffered H2D.
This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
@@ -18,6 +18,8 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
from __future__ import annotations
import threading
import time
from itertools import chain
from typing import Any
@@ -35,6 +37,12 @@ from vllm_omni.platforms import current_omni_platform
from .base import OffloadBackend, OffloadConfig
from .block_discovery import get_blocks_from_dit
from .host_registration import (
HostRegistration,
HostRegistrationCleanupError,
HostRegistrationError,
register_host_mappings,
)
from .module_collector import ModuleDiscovery
from .offload_plan import (
OffloadPlan,
@@ -51,6 +59,38 @@ from .tensor_utils import (
logger = init_logger(__name__)
# A backend normally owns both objects below. This process-lifetime safety
# owner prevents HostWeightLease.__del__ from unmapping storage when cleanup
# failure unwinds startup and the backend itself becomes unreachable. A clean
# retry removes the pair before closing the lease.
_ACTIVE_HWR_REGISTRATIONS: list[tuple[HostRegistration, HostWeightLease]] = []
_ACTIVE_HWR_REGISTRATIONS_LOCK = threading.Lock()
def _retain_active_hwr_registration(
registration: HostRegistration,
lease: HostWeightLease,
) -> None:
with _ACTIVE_HWR_REGISTRATIONS_LOCK:
if not any(
candidate is registration and candidate_lease is lease
for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS
):
_ACTIVE_HWR_REGISTRATIONS.append((registration, lease))
def _forget_active_hwr_registration(
registration: HostRegistration,
lease: HostWeightLease,
) -> None:
with _ACTIVE_HWR_REGISTRATIONS_LOCK:
_ACTIVE_HWR_REGISTRATIONS[:] = [
(candidate, candidate_lease)
for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS
if candidate is not registration or candidate_lease is not lease
]
# Threshold (in MB) for deciding whether a non-block DiT submodule should
# use layerwise offload (streaming hooks) or be moved to GPU as a resident
# module. Submodules larger than this are offloaded to save HBM; smaller
@@ -94,6 +134,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
self.rank = rank
self.pin_memory = pin_memory
self.rank_local_mmap = rank_local_mmap
self.registered_mmap = False
self.tensor_transforms = tensor_transforms or {}
self.copy_stream = copy_stream or current_omni_platform.Stream()
@@ -445,6 +486,26 @@ class DistributedLayerwiseOffloadHook(ModelHook):
# Prefetch: H2D + AllGather (overlapped on dedicated streams) #
# ------------------------------------------------------------------ #
@staticmethod
def _resolve_mmap_source(
source_info: dict[str, Any],
meta: dict[str, Any],
dtype: torch.dtype,
) -> torch.Tensor:
source = source_info["tensor"]
transform = source_info["transform"]
if callable(transform):
source = transform(source)
if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
raise ValueError(
"mmap weight transform changed tensor layout for "
f"{source_info['name']!r}: expected dtype={dtype}, "
f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
f"stride={source.stride()}"
)
return source
@staticmethod
def _pack_mmap_sources(
cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
@@ -457,18 +518,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
destination = slot_buffers[dtype][:total_numel]
sources = cpu_sources[dtype]
for source_info, meta in zip(sources, metas, strict=True):
source = source_info["tensor"]
transform = source_info["transform"]
if callable(transform):
source = transform(source)
if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
raise ValueError(
"mmap weight transform changed tensor layout for "
f"{source_info['name']!r}: expected dtype={dtype}, "
f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
f"stride={source.stride()}"
)
source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
start = meta["offset"]
physical_storage = destination[start : start + meta["numel"]]
if source.is_contiguous():
@@ -500,6 +550,32 @@ class DistributedLayerwiseOffloadHook(ModelHook):
raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated")
return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers)
@staticmethod
def _copy_mmap_sources_to_device(
cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
metadata: dict[torch.dtype, list[dict[str, Any]]],
gpu_buffers: dict[torch.dtype, torch.Tensor],
*,
non_blocking: bool,
) -> None:
"""Copy registered source views directly into flattened device buffers."""
for dtype, metas in metadata.items():
destination = gpu_buffers[dtype]
sources = cpu_sources[dtype]
for source_info, meta in zip(sources, metas, strict=True):
source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
start = meta["offset"]
physical_storage = destination[start : start + meta["numel"]]
async_copy = non_blocking and source.is_pinned()
if source.is_contiguous():
physical_storage.copy_(source.flatten(), non_blocking=async_copy)
else:
torch.as_strided(
physical_storage,
size=source.shape,
stride=source.stride(),
).copy_(source, non_blocking=async_copy)
@torch.compiler.disable
def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None:
"""Prepare next block's weights into the shared device buffer for *slot*.
@@ -515,17 +591,27 @@ class DistributedLayerwiseOffloadHook(ModelHook):
assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
if self.dp_size <= 1 or self.dp_group is None:
cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
with current_omni_platform.stream(self.copy_stream):
for dtype, cpu_shard in cpu_weights.items():
gw = gpu_weights[dtype]
async_copy = non_blocking and cpu_shard.is_pinned()
gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
evt.record(self.copy_stream)
if self.rank_local_mmap:
# The CPU slot may be overwritten only after this H2D copy has
# finished. The shared event protects reuse by another hook.
self.cpu_staging_events[slot] = evt
if self.rank_local_mmap and self.registered_mmap:
with current_omni_platform.stream(self.copy_stream):
self._copy_mmap_sources_to_device(
self.cpu_sources,
self.metadata,
gpu_weights,
non_blocking=non_blocking,
)
evt.record(self.copy_stream)
else:
cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
with current_omni_platform.stream(self.copy_stream):
for dtype, cpu_shard in cpu_weights.items():
gw = gpu_weights[dtype]
async_copy = non_blocking and cpu_shard.is_pinned()
gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
evt.record(self.copy_stream)
if self.rank_local_mmap:
# The CPU slot may be overwritten only after this H2D copy has
# finished. The shared event protects reuse by another hook.
self.cpu_staging_events[slot] = evt
else:
gpu_shards: dict[torch.dtype, torch.Tensor] = {}
shard_bufs = self.gpu_shard_buffers[slot]
@@ -735,6 +821,7 @@ class PinnedResidentLayerGroup:
self.copy_stream = copy_stream
self.loaded = False
self.rank_local_mmap = rank_local_mmap
self.registered_mmap = False
self.pin_memory = pin_memory
self._states: list[dict[str, Any]] = []
self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
@@ -805,6 +892,14 @@ class PinnedResidentLayerGroup:
ready = current_omni_platform.Event()
with current_omni_platform.stream(self.copy_stream):
for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
if self.rank_local_mmap and self.registered_mmap:
DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
state["cpu_sources"],
state["metadata"],
block_buffers,
non_blocking=True,
)
continue
if self.rank_local_mmap:
slot = index % 2
previous_copy = self._cpu_staging_events[slot]
@@ -902,8 +997,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
self._resident_layer_group: PinnedResidentLayerGroup | None = None
self._using_mmap = False
self._using_rank_local_mmap = False
self._using_registered_mmap = False
self.host_weight_plan = host_weight_plan
self._host_weight_lease: HostWeightLease | None = None
self._host_registration: HostRegistration | None = None
self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
def load_resident_layers(self) -> None:
@@ -916,6 +1013,136 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
if self._resident_layer_group is not None:
self._resident_layer_group.offload()
def _rank_local_source_tensors(
self,
hooks: list[DistributedLayerwiseOffloadHook],
) -> tuple[torch.Tensor, ...]:
"""Return the exact CPU sources that direct H2D would consume."""
tensors: list[torch.Tensor] = []
seen: set[int] = set()
def collect(cpu_sources: dict[torch.dtype, list[dict[str, Any]]]) -> None:
for sources in cpu_sources.values():
for source in sources:
tensor = source["tensor"]
if id(tensor) not in seen:
seen.add(id(tensor))
tensors.append(tensor)
for hook in hooks:
collect(hook.cpu_sources)
if self._resident_layer_group is not None:
for state in self._resident_layer_group._states:
collect(state["cpu_sources"])
return tuple(tensors)
def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bool:
"""Register the complete final-layout lease under pinned-memory policy."""
lease = self._host_weight_lease
if lease is None:
return False
if not self.config.pin_cpu_memory:
logger.info("HWR mmap registration disabled by pin_cpu_memory=False; using bounded host staging")
return False
if not lease.mapped_regions or not source_tensors:
logger.warning("HWR mmap registration found no mapped sources; using bounded host staging")
return False
limit_gib = self.config.dlo_host_registration_limit_gib
max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None
started = time.perf_counter()
try:
registration = register_host_mappings(
lease.mapped_regions,
device=self.device,
max_bytes=max_bytes,
)
try:
unpinned = [tensor for tensor in source_tensors if tensor.numel() and not tensor.is_pinned()]
except Exception as exc:
errors = registration.close()
if errors:
self._host_registration = registration
_retain_active_hwr_registration(registration, lease)
raise HostRegistrationCleanupError(
"CUDA registration succeeded but pinned-source verification failed, "
f"and rollback failed: {errors[:3]}"
) from exc
raise HostRegistrationError(f"cannot verify registered HWR sources: {exc}") from exc
if unpinned:
errors = registration.close()
if errors:
self._host_registration = registration
_retain_active_hwr_registration(registration, lease)
raise HostRegistrationCleanupError(
"CUDA registration succeeded but PyTorch rejected mapped sources, "
f"and rollback failed: {errors[:3]}"
)
raise HostRegistrationError(
"CUDA registration succeeded but PyTorch did not recognize "
f"{len(unpinned)} mapped source(s) as pinned"
)
except HostRegistrationCleanupError as exc:
# Falling back could close a lease while the platform still owns
# one of its mappings. Fail startup and retain ownership for retry.
active_registration = exc.active_registration
if active_registration is not None:
self._host_registration = active_registration
_retain_active_hwr_registration(active_registration, lease)
logger.exception("HWR mmap registration rollback failed")
raise
except HostRegistrationError as exc:
logger.warning("HWR registered direct H2D unavailable (%s); using bounded host staging", exc)
return False
self._host_registration = registration
_retain_active_hwr_registration(registration, lease)
logger.info(
"Registered %.2f GiB of HWR mmap in %d range(s) for direct H2D in %.3f s",
registration.total_bytes / 1024**3,
registration.region_count,
time.perf_counter() - started,
)
return True
def _configure_hwr_transfer(self, hooks: list[DistributedLayerwiseOffloadHook]) -> None:
"""Select registered direct H2D or bounded staging once per backend."""
plan = self.host_weight_plan
if (
(not hooks and self._resident_layer_group is None)
or not self._using_rank_local_mmap
or plan is None
or plan.backing_kind != "host_weight_runtime"
):
return
source_tensors = self._rank_local_source_tensors(hooks)
self._using_registered_mmap = self._try_register_hwr_mmap(source_tensors)
for hook in hooks:
hook.registered_mmap = self._using_registered_mmap
if self._resident_layer_group is not None:
self._resident_layer_group.registered_mmap = self._using_registered_mmap
if self._using_registered_mmap:
self._resident_layer_group._cpu_staging_buffers.clear()
def _release_registered_mmap(self) -> None:
"""Release every platform registration before closing the HWR lease."""
registration = self._host_registration
if registration is None:
return
errors = registration.close()
if errors:
lease = self._host_weight_lease
if lease is not None and not lease.closed:
_retain_active_hwr_registration(registration, lease)
logger.error("HWR mmap unregistration failed; retaining lease mappings for retry: %s", errors[:3])
raise HostRegistrationCleanupError(f"failed to unregister {len(errors)} HWR mmap range(s)")
lease = self._host_weight_lease
if lease is not None:
_forget_active_hwr_registration(registration, lease)
self._host_registration = None
logger.info("Unregistered HWR mmap ranges")
def _load_weights_via_mmap(
self,
pipeline: nn.Module,
@@ -1465,11 +1692,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
raise RuntimeError("DLO received a closed Host Weight Runtime lease")
# The final-layout restorer has already rebound the model to
# immutable host tensors. Treat those tensors as mmap-like
# sources for the bounded two-slot rank-local staging path;
# do not reinitialize or rematerialize the restored model.
# sources; transport setup below selects registered direct
# H2D or the bounded two-slot staging fallback.
self._using_rank_local_mmap = True
logger.info(
"DLO consuming final-layout Host Weight Runtime lease %s with bounded rank-local staging",
"DLO consuming final-layout Host Weight Runtime lease %s",
self._host_weight_lease.provenance.resolution_id,
)
elif self.host_weight_plan.backing_kind == "checkpoint_mmap":
@@ -1483,7 +1710,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
if self._using_rank_local_mmap:
logger.info(
"DLO rank-local host storage enabled: source pages are "
"node-shared; each worker owns only two bounded host staging slots"
"node-shared; transfer setup will select registered direct H2D or bounded host staging"
)
else:
remaining_meta = [
@@ -1653,6 +1880,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
)
pipeline._dlo_residency_controller = self
all_hooks = [hook for group in self._all_hook_groups for hook in group]
self._configure_hwr_transfer(all_hooks)
if not self._all_hook_groups:
self.enabled = bool(self._resident_blocks)
if self._using_mmap and not self.enabled:
@@ -1662,17 +1892,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
# Unified allocation: 2 shared output buffers + 2 shared shard buffers
# sized to the max block across ALL module groups (gen_layers +
# language_model). Groups execute sequentially, so 2 buffers suffice.
all_hooks: list[DistributedLayerwiseOffloadHook] = []
for group in self._all_hook_groups:
all_hooks.extend(group)
unified_buffers = self._allocate_shared_buffers(all_hooks)
unified_shard_buffers = None
if self.dp_size > 1:
unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
unified_cpu_staging = None
cpu_staging_events = None
if self._using_rank_local_mmap:
if self._using_rank_local_mmap and not self._using_registered_mmap:
unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
all_hooks,
self._resident_layer_group,
@@ -1753,6 +1979,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
def _release_mmap_handles(self) -> None:
"""Release source handles and the transport-owned HWR lease."""
if self._host_registration is not None:
raise HostRegistrationCleanupError("cannot close HWR mappings while host registration is still active")
self._mmap_transforms_by_tensor_id.clear()
if hasattr(self, "_mmap_file_cache"):
self._mmap_file_cache.clear()
@@ -1781,6 +2009,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
def disable(self) -> None:
has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed
has_registration = self._host_registration is not None
has_carrier = (
self.host_weight_plan is not None
and self.host_weight_plan.lease_carrier is not None
@@ -1796,12 +2025,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
not self.enabled
and not hasattr(self, "_mmap_file_cache")
and not has_open_lease
and not has_registration
and not has_carrier
and not has_partial_hooks
):
return
if self._using_rank_local_mmap:
if self._using_rank_local_mmap or has_registration:
current_omni_platform.synchronize()
for blocks in self._blocks:
@@ -1818,9 +2048,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
self._all_hook_groups.clear()
self._resident_blocks.clear()
self._resident_layer_group = None
self._release_registered_mmap()
self._release_mmap_handles()
self._using_mmap = False
self._using_rank_local_mmap = False
self._using_registered_mmap = False
self.enabled = False
logger.info("Distributed layer-wise offloading disabled")
@@ -0,0 +1,63 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Platform-neutral lifecycle for registering existing host mappings."""
from __future__ import annotations
from collections.abc import Sequence
from typing import Protocol
import torch
from vllm_omni.host_weight_runtime import MappedHostRegion
class HostRegistrationError(RuntimeError):
"""Raised when existing host ranges cannot be registered safely."""
class HostRegistrationBudgetError(HostRegistrationError):
"""Raised before mutation when the complete mapping exceeds its budget."""
class HostRegistrationCleanupError(HostRegistrationError):
"""Raised when registered ranges cannot be released safely."""
def __init__(self, message: str, *, active_registration: HostRegistration | None = None) -> None:
super().__init__(message)
self.active_registration = active_registration
class HostRegistration(Protocol):
"""Transport-owned registration for one process-local host mapping."""
@property
def total_bytes(self) -> int: ...
@property
def region_count(self) -> int: ...
def close(self) -> tuple[str, ...]: ...
def register_host_mappings(
regions: Sequence[MappedHostRegion],
*,
device: torch.device,
max_bytes: int | None,
) -> HostRegistration:
"""Register immutable ranges with the active platform implementation."""
if device.type == "cuda":
from .cuda_host_registration import CudaHostRegistration
return CudaHostRegistration.create(regions, max_bytes=max_bytes)
raise HostRegistrationError(f"host-mapping registration is not supported on {device.type}")
__all__ = [
"HostRegistration",
"HostRegistrationBudgetError",
"HostRegistrationCleanupError",
"HostRegistrationError",
"register_host_mappings",
]
+13 -6
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Diffusion Worker for vLLM-Omni.
@@ -897,11 +897,18 @@ class DiffusionWorker:
def shutdown(self) -> None:
"""Shutdown the worker and cleanup distributed environment."""
if self.model_runner is not None:
mgr = getattr(self.model_runner, "kv_transfer_manager", None)
if mgr is not None:
mgr.shutdown_prefetch()
destroy_distributed_env()
try:
if self.model_runner is not None:
mgr = getattr(self.model_runner, "kv_transfer_manager", None)
try:
offload_backend = getattr(self.model_runner, "offload_backend", None)
if offload_backend is not None:
offload_backend.disable()
finally:
if mgr is not None:
mgr.shutdown_prefetch()
finally:
destroy_distributed_env()
class CustomPipelineWorkerExtension:
+1
View File
@@ -571,6 +571,7 @@ class OrchestratorArgs:
dlo_resident_layers: int = 0
host_weight_runtime_mode: str = "disabled"
host_weight_runtime_root: str | None = None
dlo_host_registration_limit_gib: float = 0.0
boundary_ratio: float | None = None
flow_shift: float | None = None
diffusion_kv_cache_dtype: str | None = None
+1
View File
@@ -1061,6 +1061,7 @@ class AsyncOmniEngine:
"dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
"host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
"host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
"dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0),
"enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
"diffusion_compile_granularity": (
"regional"
+60 -18
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Orchestrator for vLLM-Omni multi-stage runtime.
@@ -363,8 +366,8 @@ class Orchestrator:
def __init__(
self,
request_async_queue: janus.AsyncQueue[EngineQueueMessage],
output_async_queue: janus.AsyncQueue[dict[str, Any]],
rpc_async_queue: janus.AsyncQueue[dict[str, Any]],
output_async_queue: janus.AsyncQueue[EngineQueueMessage],
rpc_async_queue: janus.AsyncQueue[EngineQueueMessage],
stage_pools: list[StagePool],
*,
async_chunk: bool = False,
@@ -1173,6 +1176,35 @@ class Orchestrator:
abort=True,
)
async def _fail_request_client_error(
self,
req_id: str,
stage_id: int,
error: str,
*,
close_duplex_sessions: bool = False,
) -> None:
"""Fail one request with a non-fatal 400 (bad input, engine survives).
The non-fatal counterpart of `_fail_request_dead_stage`: emits a
client-error ErrorMessage (default `fatal=False`) so the engine keeps
serving, then releases the request's state across every stage pool.
"""
await self.output_async_queue.put(
ErrorMessage(
error=error,
status_code=HTTPStatus.BAD_REQUEST.value,
error_type=DEFAULT_CLIENT_ERROR_TYPE,
request_id=req_id,
stage_id=stage_id,
)
)
await self._cleanup_request_ids(
[req_id, *self._cfg_tracker.cleanup_parent(req_id)],
abort=True,
close_duplex_sessions=close_duplex_sessions,
)
async def _dispatch_or_fail_request(
self,
dispatch: Callable[[], Awaitable[Any]],
@@ -1194,7 +1226,9 @@ class Orchestrator:
StageUnavailableError (no live replica / evicted slot) or
EngineDeadError (replica died mid-submit before the poll loop evicted
it). Both mean "this request cannot be placed", not "the server is
broken": fail the request and keep serving (#4285). Unrelated errors
broken": fail the request and keep serving (#4285). An OverflowError
from encoding an out-of-range request value is likewise failed as a
client error (400) rather than killing the loop. Other unrelated errors
propagate. Returns False when the request was failed.
``req_id`` is the request the failure is attributed to; ``dispatch_req_id``
@@ -1237,6 +1271,24 @@ class Orchestrator:
if dead_replica is not None and dead_replica in pool.live_replica_ids():
await self._handle_dead_replica(stage_id, dead_replica, e)
return False
except OverflowError as e:
# Catch overflow errors in the request payload to ensure that msgpack
# encoding can't kill the engine; instead, we just fail this request.
# The frontend should generally validate this as well, but this is needed
# in case we get values like seed > 2**64-1 from sampling params.
logger.warning(
"[Orchestrator] %s dispatch for req=%s hit an unserializable value on stage-%s: %s",
operation,
req_id,
stage_id,
e,
)
await self._fail_request_client_error(
req_id,
stage_id,
f"Request contains a value that cannot be serialized: {e}",
)
return False
# ---- Shared helpers ----
@@ -2292,21 +2344,11 @@ class Orchestrator:
"and none were provided; failing the request",
request_id,
)
await self.output_async_queue.put(
ErrorMessage(
error=(
"async_chunk requires prompt_token_ids on the stage-0 request; "
"an embeds-only prompt cannot prewarm downstream stages"
),
status_code=HTTPStatus.BAD_REQUEST.value,
error_type=DEFAULT_CLIENT_ERROR_TYPE,
request_id=request_id,
stage_id=0,
)
)
await self._cleanup_request_ids(
[request_id, *self._cfg_tracker.cleanup_parent(request_id)],
abort=True,
await self._fail_request_client_error(
request_id,
0,
"async_chunk requires prompt_token_ids on the stage-0 request; "
"an embeds-only prompt cannot prewarm downstream stages",
close_duplex_sessions=True,
)
return False
-1
View File
@@ -1558,7 +1558,6 @@ def launch_diffusion_stage_replica(
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
od_config = build_diffusion_config(model, stage_config, metadata)
od_config.max_num_seqs = batch_size
parallel_config = getattr(od_config, "parallel_config", None)
world_size = getattr(parallel_config, "world_size", 1)
try:
+3 -4
View File
@@ -1506,15 +1506,14 @@ def initialize_diffusion_stage(
stage_cfg: Stage configuration.
metadata: Extracted stage metadata.
stage_init_timeout: Timeout in seconds for stage initialization handshake
batch_size: Maximum number of requests to batch together in the
diffusion engine. Passed through to ``StageDiffusionClient``
and ultimately to ``AsyncOmni``.
batch_size: Client-side request batch width. Does not set scheduler
``max_num_seqs``; pass ``--max-num-seqs`` or stage YAML for that.
Forwarded to ``StageDiffusionClient``.
use_inline: If True, uses the inline diffusion client instead of subprocess.
"""
from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
od_config = build_diffusion_config(model, stage_cfg, metadata)
od_config.max_num_seqs = batch_size
return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""vLLM-Omni extensions for the ``vllm bench serve`` CLI.
Core functions:
@@ -16,6 +19,63 @@ by ``add_omni_args``.
"""
import argparse
import math
from pathlib import Path
_DEFAULT_OMNIINTERACT_NUM_PROMPTS = 3
def _positive_finite_float(value: str) -> float:
parsed = float(value)
if not math.isfinite(parsed) or parsed <= 0:
raise argparse.ArgumentTypeError(f"must be a finite positive number, got {value!r}")
return parsed
def _existing_file(value: str) -> str:
path = Path(value).expanduser()
if not path.is_file():
raise argparse.ArgumentTypeError(f"file does not exist: {value!r}")
return str(path)
def add_omniinteract_cli_args(parser: argparse.ArgumentParser) -> None:
from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
DEFAULT_MAX_VIDEO_DURATION_S,
OMNIINTERACT_SUBSETS,
)
group = parser.add_argument_group("OmniInteract Benchmark Options")
group.add_argument(
"--omniinteract-subsets", nargs="+", choices=OMNIINTERACT_SUBSETS, default=list(OMNIINTERACT_SUBSETS)
)
group.add_argument(
"--omniinteract-timeout-s", type=_positive_finite_float, default=900.0, help="Complete session timeout."
)
group.add_argument(
"--omniinteract-media-timeout-s",
type=_positive_finite_float,
default=600.0,
help="Per-command media timeout.",
)
group.add_argument(
"--omniinteract-max-video-duration-s",
type=_positive_finite_float,
default=DEFAULT_MAX_VIDEO_DURATION_S,
help="Reject media longer than this safety limit before decoding.",
)
group.add_argument(
"--omniinteract-ref-audio", type=_existing_file, help="Reference WAV for native-duplex audio output."
)
group.add_argument(
"--omniinteract-require-response", action="store_true", help="Fail LISTEN-only functional E2E cases."
)
group.add_argument(
"--omniinteract-output-dir",
type=Path,
default=Path("omniinteract-output"),
help="Directory for case and evaluator artifacts.",
)
def add_multi_stage_cli_args(parser: argparse.ArgumentParser) -> None:
@@ -188,6 +248,7 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None:
_OMNI_BENCH_DATASET_CHOICES = (
"daily-omni",
"omniinteract",
"seed-tts",
"seed-tts-text",
"seed-tts-design",
@@ -218,6 +279,11 @@ def extend_omni_choices(parser: argparse.ArgumentParser) -> None:
def update_omni_help(parser: argparse.ArgumentParser) -> None:
"""Update upstream argument help text to describe Omni-specific behavior."""
for action in parser._actions:
if action.dest == "num_prompts":
action.help = (
f"{action.help} OmniInteract uses {_DEFAULT_OMNIINTERACT_NUM_PROMPTS} when this option is omitted; "
"0 selects all available cases."
)
if action.dest == "percentile_metrics":
action.help = (
"Comma-separated list of selected metrics to report percentiles. "
@@ -254,6 +320,7 @@ def update_omni_help(parser: argparse.ArgumentParser) -> None:
def add_omni_args(parser: argparse.ArgumentParser) -> None:
"""Register all vLLM-Omni serving benchmark arguments."""
add_daily_omni_cli_args(parser)
add_omniinteract_cli_args(parser)
add_seed_tts_cli_args(parser)
add_multi_stage_cli_args(parser)
add_diffusion_cli_args(parser)
@@ -261,6 +328,28 @@ def add_omni_args(parser: argparse.ArgumentParser) -> None:
def preprocess_serve_args(args: argparse.Namespace) -> None:
"""Apply serving benchmark CLI transformations after parsing."""
if getattr(args, "dataset_name", None) == "omniinteract":
if getattr(args, "backend", None) != "openai-realtime-duplex":
raise ValueError("OmniInteract requires --backend openai-realtime-duplex")
if getattr(args, "endpoint", None) != "/v1/realtime":
raise ValueError("OmniInteract requires --endpoint /v1/realtime")
if not getattr(args, "omniinteract_ref_audio", None):
raise ValueError("OmniInteract requires --omniinteract-ref-audio")
if getattr(args, "ignore_eos", False):
raise ValueError("OmniInteract does not support --ignore-eos")
if getattr(args, "profile", False):
raise ValueError("OmniInteract does not support --profile")
if getattr(args, "skip_tokenizer_init", False):
raise ValueError("OmniInteract does not support --skip-tokenizer-init")
if float(getattr(args, "probe_request_rate", 0.0) or 0.0) > 0:
raise ValueError("OmniInteract does not support --probe-request-rate")
if "num_prompts" not in getattr(args, "explicit_keys", ()):
args.num_prompts = _DEFAULT_OMNIINTERACT_NUM_PROMPTS
max_concurrency = getattr(args, "max_concurrency", None)
if max_concurrency is None:
args.max_concurrency = 1
elif max_concurrency <= 0:
raise ValueError("OmniInteract requires --max-concurrency to be positive")
extra_body = dict(getattr(args, "extra_body", None) or {})
bot_task = getattr(args, "bot_task", None)
if getattr(args, "backend", None) == "openai-image-edits-omni" and bot_task is not None:
+4 -1
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Online serving benchmark subcommand for vLLM-Omni.
``OmniBenchmarkServingSubcommand`` starts from vLLM's serving benchmark
@@ -24,7 +27,7 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
"""The `serve` subcommand for vllm bench."""
name = "serve"
help = "Benchmark the online serving throughput. Supports Daily-Omni and Seed-TTS datasets."
help = "Benchmark online serving. Supports Daily-Omni, OmniInteract, and Seed-TTS datasets."
@classmethod
def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
+10
View File
@@ -760,6 +760,16 @@ class OmniServeCommand(CLISubcommand):
"the same persistent path for population and serving."
),
)
omni_config_group.add_argument(
"--dlo-host-registration-limit-gib",
type=float,
default=0.0,
help=(
"Optional per-worker GiB ceiling for registering an HWR mmap for direct H2D. "
"Zero applies no additional ceiling. Eligible no-AllGather HWR hits attempt registration "
"under the existing pinned-memory policy and fall back to bounded staging when unavailable."
),
)
# Video model parameters (e.g., Wan2.2) - engine-level
omni_config_group.add_argument(
"--boundary-ratio",
+21 -7
View File
@@ -1,9 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import math
from typing import Any, Literal
import numpy as np
from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
# Bound int request fields to avoid overflow issues.
_INT64_MIN = -(2**63)
_INT64_MAX = 2**63 - 1
_MAX_EMBEDDING_DIM = 8192
SUPPORTED_AUDIO_FORMATS: frozenset[str] = frozenset({"wav", "pcm", "flac", "mp3", "opus"})
@@ -142,18 +148,21 @@ class OpenAICreateSpeechRequest(BaseModel):
)
max_new_tokens: int | None = Field(
default=None,
ge=1,
le=_INT64_MAX,
description="Maximum tokens to generate",
)
seed: int | None = Field(
default=None,
ge=0,
le=2**63 - 1,
ge=_INT64_MIN,
le=_INT64_MAX,
description="Random seed for reproducible generation. When set, ensures "
"deterministic output for the same input text and seed value.",
)
initial_codec_chunk_frames: int | None = Field(
default=None,
ge=0,
le=_INT64_MAX,
description="Per-request initial chunk size override. If null, computed dynamically based on server load.",
)
non_streaming_mode: bool | None = Field(
@@ -365,10 +374,14 @@ class OpenAICreateAudioGenerateRequest(BaseModel):
)
num_inference_steps: int | None = Field(
default=None,
ge=1,
le=_INT64_MAX,
description="Number of inference steps",
)
seed: int | None = Field(
default=None,
ge=_INT64_MIN,
le=_INT64_MAX,
description="Random seed for reproducibility",
)
@@ -413,8 +426,8 @@ class SpeechBatchItem(BaseModel):
ref_audio: str | None = None
ref_text: str | None = None
x_vector_only_mode: bool | None = None
max_new_tokens: int | None = None
initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX)
non_streaming_mode: bool | None = None
@@ -433,8 +446,8 @@ class BatchSpeechRequest(BaseModel):
ref_audio: str | None = None
ref_text: str | None = None
x_vector_only_mode: bool | None = None
max_new_tokens: int | None = None
initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX)
non_streaming_mode: bool | None = None
@@ -517,10 +530,11 @@ class StreamingSpeechSessionConfig(BaseModel):
instructions: str | None = None
response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
max_new_tokens: int | None = Field(default=None, ge=1)
max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
initial_codec_chunk_frames: int | None = Field(
default=None,
ge=0,
le=_INT64_MAX,
description="Initial chunk size for reduced TTFA. Overrides the deploy config for this session.",
)
non_streaming_mode: bool | None = Field(
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
OpenAI-compatible protocol definitions for image generation.
@@ -21,6 +21,10 @@ from pydantic import BaseModel, Field, field_validator
from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
# Bound int request fields to avoid overflow issues.
_INT64_MIN = -(2**63)
_INT64_MAX = 2**63 - 1
_IMAGE_FILE_METADATA = {
"jpg": ("jpg", "image/jpeg"),
"jpeg": ("jpeg", "image/jpeg"),
@@ -144,7 +148,7 @@ class ImageGenerationRequest(BaseModel):
default=None,
description="Optional model-specific parameters passed directly to the model's extra_args.",
)
seed: int | None = Field(default=None, description="Random seed for reproducibility")
seed: int | None = Field(default=None, ge=_INT64_MIN, le=_INT64_MAX, description="Random seed for reproducibility")
generator_device: str | None = Field(
default=None,
description="Device for the seeded torch.Generator (e.g. 'cpu', 'cuda'). Defaults to the runner's device.",
+21 -11
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
OpenAI-compatible protocol definitions for video generation.
@@ -20,6 +20,10 @@ from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_vali
from vllm_omni.entrypoints.openai.image_api_utils import parse_size
from vllm_omni.inputs.data import DIFFUSION_QUALITY_LEVELS
# Bound int request fields to avoid overflow issues.
_INT64_MIN = -(2**63)
_INT64_MAX = 2**63 - 1
class VideoGenerationStatus(str, Enum):
QUEUED = "queued"
@@ -48,10 +52,10 @@ def file_extension(media_type: str):
class VideoParams(BaseModel):
"""Optional block for video-specific parameters."""
width: int | None = Field(default=None, ge=1, description="Video width in pixels")
height: int | None = Field(default=None, ge=1, description="Video height in pixels")
num_frames: int | None = Field(default=None, ge=1, description="Number of frames")
fps: int | None = Field(default=None, ge=1, description="Frames per second for output video")
width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels")
num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames")
fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video")
@property
def size(self) -> str | None:
@@ -140,10 +144,10 @@ class VideoGenerationRequest(BaseModel):
user: str | None = Field(default=None, description="User identifier for tracking")
# Video-specific fields (top-level for OpenAI-style compatibility)
width: int | None = Field(default=None, ge=1, description="Video width in pixels")
height: int | None = Field(default=None, ge=1, description="Video height in pixels")
fps: int | None = Field(default=None, ge=1, description="Frames per second for output video")
num_frames: int | None = Field(default=None, ge=1, description="Number of frames to generate")
width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels")
fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video")
num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames to generate")
aspect_ratio: str | None = Field(
default=None,
description=(
@@ -151,7 +155,12 @@ class VideoGenerationRequest(BaseModel):
"FL2VA follows the input image; Ref2VA defaults to 16:9."
),
)
short_edge: int | None = Field(default=None, ge=1, description="MiniMax H3 output short edge in pixels")
short_edge: int | None = Field(
default=None,
ge=1,
le=_INT64_MAX,
description="MiniMax H3 output short edge in pixels",
)
num_outputs_per_prompt: int = Field(
default=1,
ge=1,
@@ -218,7 +227,7 @@ class VideoGenerationRequest(BaseModel):
le=20.0,
description="True CFG scale (model-specific parameter, may be ignored if not supported)",
)
seed: int | None = Field(default=None, description="Random seed for reproducibility")
seed: int | None = Field(default=None, ge=_INT64_MIN, le=_INT64_MAX, description="Random seed for reproducibility")
generate_sound: bool = Field(
default=False,
description="Request model-generated audio for video models that support sound generation.",
@@ -237,6 +246,7 @@ class VideoGenerationRequest(BaseModel):
frame_interpolation_exp: int = Field(
default=1,
ge=1,
le=_INT64_MAX,
description="Interpolation exponent: 1=2x temporal resolution, 2=4x, etc.",
)
frame_interpolation_scale: float = Field(
+130 -31
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Reusable Realtime WebSocket, PCM, and event helpers for MiniCPM-o demos."""
from __future__ import annotations
@@ -30,6 +33,13 @@ def _rounded_ms(value: float) -> float:
return round(float(value), 3)
def _finite_number(value: object, *, nonnegative: bool = False) -> float | None:
if isinstance(value, bool) or not isinstance(value, int | float):
return None
parsed = float(value)
return parsed if math.isfinite(parsed) and (not nonnegative or parsed >= 0) else None
def _interval_summary(values: list[float]) -> dict[str, float | int]:
clean = sorted(_rounded_ms(value) for value in values if math.isfinite(value) and value >= 0)
if not clean:
@@ -56,11 +66,7 @@ def summarize_session_request_metrics(
"""Average client-observed metrics across turns that emitted audio."""
def mean(metric: str, *, digits: int = 3) -> float | None:
values = [
float(request[metric])
for request in request_metrics
if isinstance(request.get(metric), int | float) and math.isfinite(float(request[metric]))
]
values = [value for request in request_metrics if (value := _finite_number(request.get(metric))) is not None]
return round(sum(values) / len(values), digits) if values else None
return {
@@ -93,24 +99,57 @@ def _event_stage_metrics(event: dict[str, object]) -> dict[str, object] | None:
def build_realtime_url(
url: str,
model: str,
model: str | None,
*,
autostart: bool | None = None,
native_duplex: bool | None = True,
session_id: str | None = None,
) -> str:
"""Add the explicit native-duplex query parameters to a Realtime URL."""
parts = urlsplit(url)
if parts.scheme in {"http", "https"}:
parts = parts._replace(scheme="ws" if parts.scheme == "http" else "wss")
if parts.scheme not in {"ws", "wss"} or not parts.netloc:
raise ValueError(f"Unsupported Realtime URL: {url!r}")
query = dict(parse_qsl(parts.query, keep_blank_values=True))
query.setdefault("duplex", "1")
query.setdefault("model", model)
query.setdefault("minicpmo45_native_duplex", "1")
query["duplex"] = "1"
if model:
query["model"] = model
if native_duplex is not None:
query["minicpmo45_native_duplex"] = "1" if native_duplex else "0"
if autostart is not None:
query.setdefault("autostart", "1" if autostart else "0")
query["autostart"] = "1" if autostart else "0"
if session_id:
query.setdefault("session_id", session_id)
query["session_id"] = session_id
return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
def reference_audio_data_url(path: str | None) -> str | None:
"""Encode a local reference WAV for a Realtime session update."""
if path is None:
return None
audio = Path(path).expanduser().resolve()
if not audio.is_file():
raise FileNotFoundError(f"Reference audio does not exist: {audio}")
return "data:audio/wav;base64," + base64.b64encode(audio.read_bytes()).decode("ascii")
def chunk_period_ms(events: list[dict[str, object]], *, default: int = 1000) -> int:
"""Read the negotiated native-duplex model-unit duration."""
for event in reversed(events):
session = event.get("session")
capabilities = session.get("capabilities") if isinstance(session, dict) else None
value = capabilities.get("chunk_period_ms") if isinstance(capabilities, dict) else None
if isinstance(value, int) and not isinstance(value, bool) and value > 0:
return value
return default
def has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool:
unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000
return bool(unit_bytes and len(pcm16) % unit_bytes)
def read_pcm16_wav(path: Path) -> bytes:
"""Read a mono, uncompressed, 16 kHz PCM16 WAV file."""
with wave.open(str(path), "rb") as wav_file:
@@ -202,6 +241,20 @@ class RealtimeEventCollector:
chunk for response_id in self.response_ids for chunk in self.response_audio.get(response_id, ())
)
def response_text(self, response_id: str) -> str:
"""Join all text/transcript deltas for one response identity."""
return "".join(
str(event.get("delta") or "")
for event in self.events
if self.response_id(event) == response_id
and event.get("type")
in {
"response.audio_transcript.delta",
"response.output_text.delta",
"response.text.delta",
}
)
def errors(self) -> list[dict[str, object]]:
return [event for event in self.events if event.get("type") == "error"]
@@ -282,15 +335,16 @@ class RealtimeEventCollector:
if stage0_metrics is not None:
raw_itls = stage0_metrics.get("vllm_itls_ms")
itls = (
[float(value) for value in raw_itls if isinstance(value, int | float)]
[parsed for value in raw_itls if (parsed := _finite_number(value, nonnegative=True)) is not None]
if isinstance(raw_itls, list)
else []
)
result["stage0_tokens"] = {
"source": "engine_stage_metrics",
"output_token_count": int(stage0_metrics.get("num_tokens_out") or 0),
"ttft_ms": float(stage0_metrics.get("vllm_ttft_ms") or 0.0),
"tpot_ms": float(stage0_metrics.get("vllm_tpot_ms") or 0.0),
"output_token_count": int(_finite_number(stage0_metrics.get("num_tokens_out"), nonnegative=True) or 0),
"ttft_ms": _finite_number(stage0_metrics.get("vllm_ttft_ms"), nonnegative=True) or 0.0,
"tpot_ms": _finite_number(stage0_metrics.get("vllm_tpot_ms"), nonnegative=True),
"itls_ms": itls,
"inter_token_interval_ms": _interval_summary(itls),
}
@@ -368,15 +422,26 @@ class RealtimeEventCollector:
class RealtimeDuplexClient:
"""Small async client used by the user demo and reusable smoke probes."""
def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None:
def __init__(
self,
url: str,
*,
max_size: int = 64 * 1024 * 1024,
additional_headers: dict[str, str] | None = None,
) -> None:
self.url = url
self.max_size = max_size
self.additional_headers = additional_headers
self.events = RealtimeEventCollector()
self._ws: Any = None
self._reader_task: asyncio.Task[None] | None = None
async def __aenter__(self) -> RealtimeDuplexClient:
self._ws = await websockets.connect(self.url, max_size=self.max_size)
self._ws = await websockets.connect(
self.url,
max_size=self.max_size,
additional_headers=self.additional_headers,
)
self._reader_task = asyncio.create_task(self._read_events())
return self
@@ -405,6 +470,19 @@ class RealtimeDuplexClient:
async def send(self, event: dict[str, object]) -> None:
await self._ws.send(json.dumps(event))
def raise_if_reader_stopped(self) -> None:
"""Fail a caller waiting on events after the WebSocket reader exits."""
task = self._reader_task
if task is None or not task.done():
return
if task.cancelled():
raise ConnectionError("Realtime WebSocket reader was cancelled")
error = task.exception()
if error is not None:
raise ConnectionError("Realtime WebSocket reader failed") from error
raise ConnectionError("Realtime WebSocket closed before the requested event arrived")
async def configure(
self,
model: str,
@@ -418,6 +496,7 @@ class RealtimeDuplexClient:
temperature: float | None = None,
extra_body: dict[str, object] | None = None,
session_id: str | None = None,
idle_timeout_s: float | None = None,
timeout_s: float = 20.0,
) -> None:
session_extra_body = dict(extra_body or {})
@@ -448,14 +527,23 @@ class RealtimeDuplexClient:
if instructions is not None:
session["instructions"] = instructions
if initial_user_text is not None:
extra_body = session["extra_body"]
assert isinstance(extra_body, dict)
extra_body["duplex_initial_user_text"] = initial_user_text
session_extra = session["extra_body"]
assert isinstance(session_extra, dict)
session_extra["duplex_initial_user_text"] = initial_user_text
if session_id:
session["session_id"] = session_id
if idle_timeout_s is not None:
session["idle_timeout_s"] = idle_timeout_s
await self.send({"type": "session.update", "session": session})
def session_created() -> bool:
if self.events.count("session.created") > 0:
return True
self.raise_if_reader_stopped()
return False
await wait_for(
lambda: self.events.count("session.created") > 0,
session_created,
timeout_s=timeout_s,
label="session.created",
)
@@ -498,20 +586,31 @@ class RealtimeDuplexClient:
if not pcm16:
continue
played_ms = len(pcm16) * 1000 // (self.events.output_sample_rate_hz * PCM16_BYTES_PER_SAMPLE)
await self.send(
{
"type": "playback.ack",
"response_id": response_id,
"item_id": f"item_{response_id}",
"played_ms": played_ms,
"committed_ms": played_ms,
}
)
await self.send_playback_ack(response_id, played_ms)
async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
await self.send(
{
"type": "playback.ack",
"response_id": response_id,
"item_id": f"item_{response_id}",
"played_ms": played_ms,
"committed_ms": played_ms,
}
)
async def close_session(self, *, timeout_s: float = 20.0) -> None:
from_index = len(self.events.events)
await self.send({"type": "session.close"})
def session_closed() -> bool:
if any(event.get("type") == "session.closed" for event in self.events.events[from_index:]):
return True
self.raise_if_reader_stopped()
return False
await wait_for(
lambda: self.events.count("session.closed") > 0,
session_closed,
timeout_s=timeout_s,
label="session.closed",
)
@@ -12,17 +12,19 @@ import os
import struct
import uuid
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path, PurePosixPath
from types import TracebackType
import torch
from safetensors import safe_open
from ..identity import WeightArtifactIdentity, canonical_json
from ..manifest import FileManifestEntry, PublicationMarker, TensorManifestEntry, WeightManifest
from ..protocols import ProductionMetadata, TensorWriteSpec
_HASH_CHUNK_BYTES = 8 * 1024**2
_MAX_HASH_WORKERS = 8
_MANIFEST_FILE = "manifest.json"
_READY_FILE = "READY.json"
@@ -79,21 +81,47 @@ def _fsync_directory(path: Path) -> None:
os.close(fd)
def _file_sha256(path: Path) -> str:
digest = hashlib.sha256()
def _fsync_payload(fd: int) -> None:
try:
os.fsync(fd)
finally:
os.close(fd)
@dataclass(frozen=True)
class _PayloadDigests:
file: str
tensors: dict[str, str]
def _scan_payload(path: Path, specs: tuple[TensorWriteSpec, ...]) -> _PayloadDigests:
"""Calculate file and tensor digests in one ordered read."""
file_digest = hashlib.sha256()
tensor_digests: dict[str, str] = {}
with path.open("rb") as handle:
while chunk := handle.read(_HASH_CHUNK_BYTES):
digest.update(chunk)
return digest.hexdigest()
def _tensor_sha256(tensor: torch.Tensor) -> str:
digest = hashlib.sha256()
byte_view = tensor.detach().reshape(-1).view(torch.uint8)
for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
digest.update(memoryview(chunk.numpy()))
return digest.hexdigest()
prefix = handle.read(8)
if len(prefix) != 8:
raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}")
header_size = struct.unpack("<Q", prefix)[0]
header = handle.read(header_size)
if len(header) != header_size:
raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}")
file_digest.update(prefix)
file_digest.update(header)
for spec in specs:
digest = hashlib.sha256()
remaining = spec.nbytes
while remaining:
chunk = handle.read(min(remaining, _HASH_CHUNK_BYTES))
if not chunk:
raise ArtifactWriterError(f"truncated tensor {spec.name!r} in {path.name!r}")
file_digest.update(chunk)
digest.update(chunk)
remaining -= len(chunk)
tensor_digests[spec.name] = digest.hexdigest()
if handle.read(1):
raise ArtifactWriterError(f"unexpected trailing bytes in {path.name!r}")
return _PayloadDigests(file_digest.hexdigest(), tensor_digests)
class _Coverage:
@@ -130,6 +158,8 @@ class StreamingSafeTensorFile:
owner: FilesystemArtifactWriter,
relative_name: str,
specs: tuple[TensorWriteSpec, ...],
*,
ordered: bool,
) -> None:
self._owner = owner
self.relative_name = _safe_relative_name(relative_name)
@@ -144,6 +174,7 @@ class StreamingSafeTensorFile:
self.specs = tuple(sorted(specs, key=lambda spec: spec.name))
self._specs = {spec.name: spec for spec in self.specs}
self._coverage = {spec.name: _Coverage(spec.nbytes) for spec in self.specs}
self._ordered = ordered
self._closed = False
offsets: dict[str, int] = {}
@@ -167,6 +198,7 @@ class StreamingSafeTensorFile:
total_size = self._payload_offset + cursor
owner._reserve_file(total_size)
self._fd = os.open(self._path, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_CLOEXEC, 0o600)
prefix = struct.pack("<Q", len(encoded_header))
try:
if hasattr(os, "posix_fallocate"):
try:
@@ -177,13 +209,16 @@ class StreamingSafeTensorFile:
os.ftruncate(self._fd, total_size)
else: # pragma: no cover - Linux production path has posix_fallocate
os.ftruncate(self._fd, total_size)
_pwrite_all(self._fd, struct.pack("<Q", len(encoded_header)), 0)
_pwrite_all(self._fd, prefix, 0)
_pwrite_all(self._fd, encoded_header, 8)
except BaseException:
os.close(self._fd)
self._path.unlink(missing_ok=True)
owner._forget_reserved_file(total_size)
raise
self._file_digest = hashlib.sha256(prefix + encoded_header) if ordered else None
self._tensor_digests = {spec.name: hashlib.sha256() for spec in self.specs} if ordered else {}
self._write_cursor = self._payload_offset
owner._opened(self)
@staticmethod
@@ -226,9 +261,19 @@ class StreamingSafeTensorFile:
raise ArtifactWriterError(f"tensor {name!r} was not declared in this payload file") from exc
def _write(self, name: str, byte_offset: int, data: memoryview) -> None:
file_offset = self._payload_offset + self._offsets[name] + byte_offset
if self._ordered and file_offset != self._write_cursor:
raise ArtifactWriterError(
f"ordered tensor payload expected offset {self._write_cursor}, got {file_offset} for {name!r}"
)
coverage = self._coverage[name]
coverage.add(byte_offset, byte_offset + len(data))
_pwrite_all(self._fd, data, self._payload_offset + self._offsets[name] + byte_offset)
_pwrite_all(self._fd, data, file_offset)
if self._ordered:
assert self._file_digest is not None
self._file_digest.update(data)
self._tensor_digests[name].update(data)
self._write_cursor += len(data)
def close(self) -> None:
if self._closed:
@@ -238,11 +283,19 @@ class StreamingSafeTensorFile:
self.abort()
raise ArtifactWriterError(f"tensor payload closed with incomplete writes: {missing}")
self._closed = True
digests = None
if self._ordered:
assert self._file_digest is not None
digests = _PayloadDigests(
self._file_digest.hexdigest(),
{name: digest.hexdigest() for name, digest in self._tensor_digests.items()},
)
try:
os.fsync(self._fd)
finally:
os.fchmod(self._fd, 0o444)
self._owner._closed(self, digests, self._fd)
except BaseException:
os.close(self._fd)
self._owner._closed(self)
raise
def abort(self) -> None:
if self._closed:
@@ -281,20 +334,25 @@ class FilesystemArtifactWriter:
self._capacity_check = capacity_check
self._active: set[StreamingSafeTensorFile] = set()
self._files: dict[str, tuple[TensorWriteSpec, ...]] = {}
self._digests: dict[str, _PayloadDigests] = {}
self._reserved_bytes = 0
self._finalized = False
self._fsync_executor: ThreadPoolExecutor | None = None
self._fsync_futures: list[Future[None]] = []
def open_tensor_file(
self,
relative_name: str,
specs: tuple[TensorWriteSpec, ...],
*,
ordered: bool = False,
) -> StreamingSafeTensorFile:
if self._finalized:
raise ArtifactWriterError("cannot add payloads after artifact finalization")
relative_name = _safe_relative_name(relative_name)
if relative_name in self._files or any(item.relative_name == relative_name for item in self._active):
raise ArtifactWriterError(f"artifact payload {relative_name!r} was declared more than once")
return StreamingSafeTensorFile(self, relative_name, specs)
return StreamingSafeTensorFile(self, relative_name, specs, ordered=ordered)
def _reserve_file(self, size: int) -> None:
self._capacity_check(self._reserved_bytes + size, size)
@@ -306,9 +364,14 @@ class FilesystemArtifactWriter:
def _opened(self, writer: StreamingSafeTensorFile) -> None:
self._active.add(writer)
def _closed(self, writer: StreamingSafeTensorFile) -> None:
def _closed(self, writer: StreamingSafeTensorFile, digests: _PayloadDigests | None, fd: int) -> None:
self._active.remove(writer)
self._files[writer.relative_name] = writer.specs
if digests is not None:
self._digests[writer.relative_name] = digests
if self._fsync_executor is None:
self._fsync_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="hwr-fsync")
self._fsync_futures.append(self._fsync_executor.submit(_fsync_payload, fd))
def _aborted(self, writer: StreamingSafeTensorFile, size: int) -> None:
self._active.discard(writer)
@@ -318,6 +381,19 @@ class FilesystemArtifactWriter:
for writer in tuple(self._active):
with contextlib.suppress(Exception):
writer.abort()
with contextlib.suppress(Exception):
self._wait_for_payload_fsync()
def _wait_for_payload_fsync(self) -> None:
if self._fsync_executor is None:
return
executor = self._fsync_executor
futures = tuple(self._fsync_futures)
self._fsync_executor = None
self._fsync_futures.clear()
executor.shutdown(wait=True)
for future in futures:
future.result()
def finalize(
self,
@@ -335,46 +411,50 @@ class FilesystemArtifactWriter:
if metadata.restorer_schema != identity.producer.restorer_schema:
raise ArtifactWriterError("production metadata does not match the identity's restorer schema")
self._wait_for_payload_fsync()
tensor_entries: list[TensorManifestEntry] = []
file_entries: list[FileManifestEntry] = []
seen_tensors: set[str] = set()
for file_name in sorted(self._files):
file_names = sorted(self._files)
unhashed = [file_name for file_name in file_names if file_name not in self._digests]
if unhashed:
with ThreadPoolExecutor(max_workers=min(_MAX_HASH_WORKERS, len(unhashed))) as executor:
scans = executor.map(
lambda file_name: _scan_payload(self.staging_dir / file_name, self._files[file_name]),
unhashed,
)
self._digests.update(zip(unhashed, scans, strict=True))
for file_name in file_names:
path = self.staging_dir / file_name
os.chmod(path, 0o444)
with path.open("rb") as handle:
os.fsync(handle.fileno())
specs = self._files[file_name]
digests = self._digests[file_name]
file_entries.append(
FileManifestEntry(
file_name=file_name,
size=path.stat().st_size,
sha256=_file_sha256(path),
sha256=digests.file,
)
)
expected = {spec.name: spec for spec in self._files[file_name]}
with safe_open(str(path), framework="pt", device="cpu") as handle:
if set(handle.keys()) != set(expected):
raise ArtifactWriterError(f"published tensor keys differ in {file_name!r}")
for name in sorted(expected):
if name in seen_tensors:
raise ArtifactWriterError(f"artifact tensor {name!r} occurs in multiple payload files")
seen_tensors.add(name)
spec = expected[name]
tensor = handle.get_tensor(name)
if tuple(tensor.shape) != spec.shape or tensor.dtype != spec.dtype:
raise ArtifactWriterError(f"published tensor {name!r} differs from its declared metadata")
entry = TensorManifestEntry(
name=name,
for spec in specs:
if spec.name in seen_tensors:
raise ArtifactWriterError(f"artifact tensor {spec.name!r} occurs in multiple payload files")
seen_tensors.add(spec.name)
tensor_entries.append(
TensorManifestEntry(
name=spec.name,
kind=spec.kind,
role=spec.role,
file_name=file_name,
storage_key=name,
shape=tuple(tensor.shape),
dtype=str(tensor.dtype),
stride=tuple(tensor.stride()),
nbytes=tensor.numel() * tensor.element_size(),
sha256=_tensor_sha256(tensor),
storage_key=spec.name,
shape=spec.shape,
dtype=str(spec.dtype),
stride=tuple(torch.empty(spec.shape, dtype=spec.dtype, device="meta").stride()),
nbytes=spec.nbytes,
sha256=digests.tensors[spec.name],
)
tensor_entries.append(entry)
)
ordered_tensor_entries = tuple(sorted(tensor_entries, key=lambda entry: entry.name))
content_digest = hashlib.sha256()
@@ -88,6 +88,8 @@ class ArtifactWriter(Protocol):
self,
relative_name: str,
specs: tuple[TensorWriteSpec, ...],
*,
ordered: bool = False,
) -> TensorFileWriter: ...
+7 -3
View File
@@ -470,10 +470,14 @@ def resolve_quant_config_from_disk(
)
return build_quant_config(qc_method, **qc_kwargs)
# AutoRound MXFP8 checkpoints use data_type="mx_fp" instead of
# is_checkpoint_*_serialized; rebuild so the offline path is selected.
# AutoRound MXFP checkpoints use data_type="mx_fp" instead of
# is_checkpoint_*_serialized; rebuild so the offline MXFP4/MXFP8 path is
# selected according to the checkpoint's bit width.
if qc_kwargs.get("data_type") == "mx_fp":
logger.info("config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP8.")
logger.info(
"config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP%d.",
qc_kwargs.get("bits", getattr(quant_config, "weight_bits", 0)),
)
return build_quant_config(qc_method, **qc_kwargs)
if (