Compare commits
13
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7db3a0cd5e | ||
|
|
780e8d40aa | ||
|
|
6639c82a17 | ||
|
|
0eb47b76d0 | ||
|
|
497c537c6f | ||
|
|
c608910d20 | ||
|
|
5a8b81ebeb | ||
|
|
dd527b4245 | ||
|
|
0d5d1115d6 | ||
|
|
683f3c92a4 | ||
|
|
72beeaee86 | ||
|
|
d6b66faba7 | ||
|
|
bc8ea489ef |
No files matched your search
@@ -83,7 +83,7 @@ steps:
|
|||||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||||
- |
|
- |
|
||||||
set +e
|
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=$$?
|
EXIT=$$?
|
||||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||||
exit $$EXIT
|
exit $$EXIT
|
||||||
@@ -375,7 +375,7 @@ steps:
|
|||||||
- label: ":full_moon: Diffusion · Quantization Test with L4"
|
- label: ":full_moon: Diffusion · Quantization Test with L4"
|
||||||
timeout_in_minutes: 60
|
timeout_in_minutes: 60
|
||||||
commands:
|
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"
|
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
|
||||||
mirror_hardwares: l4_1
|
mirror_hardwares: l4_1
|
||||||
|
|
||||||
|
|||||||
@@ -74,6 +74,19 @@ steps:
|
|||||||
exit $$EXIT
|
exit $$EXIT
|
||||||
mirror_hardwares: h100_2
|
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"
|
- label: ":full_moon: Omni · Perf Test · Multi-Replica"
|
||||||
key: weekly-omni-performance-multi-replicas
|
key: weekly-omni-performance-multi-replicas
|
||||||
timeout_in_minutes: 180
|
timeout_in_minutes: 180
|
||||||
|
|||||||
+21
-6
@@ -1,5 +1,5 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# 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
|
# ruff: noqa
|
||||||
# code borrowed from https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py
|
# 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():
|
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 = ""
|
env_vars = ""
|
||||||
secret_terms = ("secret", "token", "api", "access", "password")
|
secret_terms = ("secret", "token", "api", "access", "password")
|
||||||
report_prefix = (
|
report_prefix = (
|
||||||
@@ -507,12 +517,17 @@ def get_env_vars():
|
|||||||
"MKL_",
|
"MKL_",
|
||||||
"NVIDIA",
|
"NVIDIA",
|
||||||
)
|
)
|
||||||
for k, v in os.environ.items():
|
for k, v in sorted(os.environ.items()):
|
||||||
if any(term in k.lower() for term in secret_terms):
|
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
|
continue
|
||||||
if k in environment_variables:
|
if (
|
||||||
env_vars = env_vars + "{}={}".format(k, v) + "\n"
|
k in environment_variables
|
||||||
if k.startswith(report_prefix):
|
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"
|
env_vars = env_vars + "{}={}".format(k, v) + "\n"
|
||||||
|
|
||||||
return env_vars
|
return env_vars
|
||||||
|
|||||||
+7
-1
@@ -2,6 +2,11 @@
|
|||||||
|
|
||||||
The CLI for vLLM-Omni inherits from vllm with some additional arguments.
|
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
|
## serve
|
||||||
|
|
||||||
Starts the vLLM-Omni OpenAI Compatible API server.
|
Starts the vLLM-Omni OpenAI Compatible API server.
|
||||||
@@ -39,4 +44,5 @@ vllm bench serve --omni \
|
|||||||
--num-prompts 5
|
--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
@@ -2,7 +2,7 @@
|
|||||||
The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
|
The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
|
||||||
|
|
||||||
## Notes
|
## 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
|
## 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:
|
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.
|
Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
|
||||||
|
|
||||||
- `--backend`
|
- `--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`
|
- `--model`
|
||||||
The model identifier to load, filled according to the models supported by vLLM-Omni.
|
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.
|
The API endpoint exposed externally, to which clients send their requests.
|
||||||
|
|
||||||
- `--dataset-name`
|
- `--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`
|
- `--num-prompts`
|
||||||
The total number of requests to send, an integer.
|
The total number of requests to send, an integer.
|
||||||
@@ -269,6 +269,44 @@ We use audio generation time / audio duration to calculate RTF.
|
|||||||
|
|
||||||
</details>
|
</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
|
### Multi-Modal Benchmark
|
||||||
|
|
||||||
<details class="admonition abstract" markdown="1">
|
<details class="admonition abstract" markdown="1">
|
||||||
|
|||||||
@@ -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
|
Each model defines fixed topology in a registered `PipelineConfig` and runtime
|
||||||
overrides in a deploy YAML.
|
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 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).
|
For an introduction, see [Pipeline and deploy configurations](./stage_configs.md).
|
||||||
|
|||||||
@@ -0,0 +1,236 @@
|
|||||||
|
# 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_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 57 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 | 7 | 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 |
|
| 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 (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_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` (`slow` in `mark`; **Perf Test** in `test-weekly.yml`) |
|
| 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` |
|
| 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`, … |
|
| 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
|
#### 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
|
```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`.
|
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:
|
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_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"
|
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 \
|
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
|
||||||
--test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
--test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
|
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 -->
|
<!-- markdownlint-disable-next-line MD051 -->
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsu
|
|||||||
| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
|
| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
|
||||||
| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
|
| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
|
||||||
| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
|
| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
|
||||||
|
| **Auxiliary component policy** | ✅ Encoders/VAEs may remain resident. | ✅ Encoders/VAEs may remain resident. |
|
||||||
|
|
||||||
See [Parallelism compatibility](#parallelism-compatibility) and
|
See [Parallelism compatibility](#parallelism-compatibility) and
|
||||||
[Request and loading constraints](#request-and-loading-constraints) for the
|
[Request and loading constraints](#request-and-loading-constraints) for the
|
||||||
@@ -196,16 +197,15 @@ adopt it.
|
|||||||
|
|
||||||
### Final-layout Host Weight Runtime consumer
|
### Final-layout Host Weight Runtime consumer
|
||||||
|
|
||||||
> **Status:** PR2 implementation. The representation-independent identity,
|
> **Status:** PR2 landed in
|
||||||
> producer, and restoration contracts landed in
|
> [PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). PR3 adds
|
||||||
> [PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445); this consumer
|
> registered direct H2D plus generic HWR publication/source-digest
|
||||||
> wires the final-layout BF16 path into eligible no-AllGather DLO.
|
> optimizations without changing AllGather behavior.
|
||||||
|
|
||||||
Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
|
Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
|
||||||
only to no-AllGather DLO. Enable it with
|
only to no-AllGather DLO. Enable it with
|
||||||
`--host-weight-runtime-mode preferred` (or `required`) and
|
`--host-weight-runtime-mode preferred` (or `required`) and
|
||||||
`--host-weight-runtime-root <node-local-root>`. Registration of shared mappings
|
`--host-weight-runtime-root <node-local-root>`.
|
||||||
and direct asynchronous H2D remain a separate PR3 transport change.
|
|
||||||
|
|
||||||
The modes express operator fallback policy, not different artifact formats:
|
The modes express operator fallback policy, not different artifact formats:
|
||||||
|
|
||||||
@@ -254,6 +254,44 @@ outcomes are authoritative:
|
|||||||
Checkpoint mmap remains an unchanged control path whenever final-layout HWR is
|
Checkpoint mmap remains an unchanged control path whenever final-layout HWR is
|
||||||
not selected.
|
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
|
#### Pre-service transaction boundary
|
||||||
|
|
||||||
The startup transaction does not end at restore commit. A lease-backed model is
|
The startup transaction does not end at restore commit. A lease-backed model is
|
||||||
@@ -317,6 +355,7 @@ loader owns
|
|||||||
-> backend takes once before asynchronous work
|
-> backend takes once before asynchronous work
|
||||||
-> backend drains pending H2D work
|
-> backend drains pending H2D work
|
||||||
-> backend releases hooks, staging, and model references
|
-> backend releases hooks, staging, and model references
|
||||||
|
-> backend unregisters every mapped range
|
||||||
-> backend closes lease
|
-> backend closes lease
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -326,13 +365,15 @@ it. Once the backend takes ownership, only backend abort or teardown may drain
|
|||||||
work and close it.
|
work and close it.
|
||||||
|
|
||||||
The implementation keeps the final-layout lease through backend setup and
|
The implementation keeps the final-layout lease through backend setup and
|
||||||
initial prefetch. The backend uses the existing two bounded rank-local staging
|
initial prefetch. Successful registration copies directly from immutable HWR
|
||||||
slots for the immutable HWR tensors, and its transactional `enable()` cleanup
|
views; otherwise the backend uses the existing two bounded rank-local staging
|
||||||
removes partial hooks and closes the lease before reporting a setup failure.
|
slots. Transactional cleanup drains device work, removes source references,
|
||||||
Preferred mode then uses the runner's fresh canonical retry; required mode
|
unregisters mappings, and only then closes the lease. An unregistration failure
|
||||||
propagates the 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.
|
- Warm hit performs zero ordinary DiT materialization and zero producer calls.
|
||||||
- Shared warm finalization changes neither restored bytes nor backing pointers.
|
- Shared warm finalization changes neither restored bytes nor backing pointers.
|
||||||
@@ -344,6 +385,10 @@ propagates the failure.
|
|||||||
- Disabled mode and AllGather emit zero final-layout HWR interaction.
|
- Disabled mode and AllGather emit zero final-layout HWR interaction.
|
||||||
- The lease carrier rejects duplicate take and serialization.
|
- The lease carrier rejects duplicate take and serialization.
|
||||||
- Backend setup or prefetch failure drains asynchronous work before lease close.
|
- 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.
|
- Compatible checkpoint mmap remains an unchanged benchmark and control path.
|
||||||
- Prewarm uses one matching producer cohort per TP/SP topology and storage
|
- Prewarm uses one matching producer cohort per TP/SP topology and storage
|
||||||
domain.
|
domain.
|
||||||
@@ -375,6 +420,22 @@ propagates the failure.
|
|||||||
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
|
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
|
||||||
configuration.
|
configuration.
|
||||||
|
|
||||||
|
### Auxiliary component policy
|
||||||
|
|
||||||
|
`dlo_offload_components` maps discovered encoder/VAE paths to booleans. A
|
||||||
|
component set to `False` is materialized on the target device once and is not
|
||||||
|
entered through the model's manual offload lifecycle. A `True` value permits
|
||||||
|
the existing `OffloadPlan.on_demand_component_paths` and encoder block-hook
|
||||||
|
behavior; it does not add support where the model declares none. The optional
|
||||||
|
`default` key handles unmatched auxiliary components, and an empty map
|
||||||
|
preserves the current policy.
|
||||||
|
|
||||||
|
The policy is resolved after component discovery and before DLO mutates DiT
|
||||||
|
storage. Unknown names, including DiT paths, are rejected. DLO continues to
|
||||||
|
own DiT streaming regardless of this map, which avoids an invalid state where
|
||||||
|
the loader has skipped ordinary DiT materialization but the backend is asked
|
||||||
|
to keep that DiT resident.
|
||||||
|
|
||||||
## Request and loading constraints
|
## Request and loading constraints
|
||||||
|
|
||||||
AllGather DP multi-concurrency requires:
|
AllGather DP multi-concurrency requires:
|
||||||
@@ -386,6 +447,15 @@ AllGather DP multi-concurrency requires:
|
|||||||
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
||||||
requirements.
|
requirements.
|
||||||
|
|
||||||
|
MiniMax-H3 currently builds one global encoder TP group and broadcasts the
|
||||||
|
resulting conditioning from global rank 0. Therefore an encoder TP group that
|
||||||
|
spans multiple DP replicas is not valid for independent-prompt AllGather
|
||||||
|
waves. Exact-recipe DP2/TP2 and DP2/SP2 experiments with encoder TP4 completed
|
||||||
|
the collective schedule, but two distinct prompts with the same seed produced
|
||||||
|
byte-identical video and audio in every measured wave. Enabling or disabling
|
||||||
|
encoder offload produced paired-identical results, so this is a pre-existing
|
||||||
|
model topology limitation rather than component-policy behavior.
|
||||||
|
|
||||||
Direct checkpoint mmap can back either transfer path. It is currently limited
|
Direct checkpoint mmap can back either transfer path. It is currently limited
|
||||||
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
|
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
|
||||||
ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
|
ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
|
||||||
@@ -394,11 +464,12 @@ reconstructs those tensors with their recorded layouts. Other online methods
|
|||||||
must use `--dlo-no-use-allgather` or disable online quantization until their
|
must use `--dlo-no-use-allgather` or disable online quantization until their
|
||||||
runtime layouts are validated.
|
runtime layouts are validated.
|
||||||
|
|
||||||
The Host Weight Runtime representation and publication contracts are merged;
|
The Host Weight Runtime representation, publication, and no-AllGather consumer
|
||||||
see [RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414) and
|
contracts are merged; see
|
||||||
[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445). The planned
|
[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414),
|
||||||
no-AllGather consumer above is the PR2 implementation; registration and direct
|
[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445), and
|
||||||
asynchronous H2D remain a separate transport follow-up.
|
[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
|
## Validation coverage
|
||||||
|
|
||||||
@@ -411,8 +482,12 @@ Current source-level validation includes:
|
|||||||
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
|
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
|
||||||
sharding of finalized weights and scales;
|
sharding of finalized weights and scales;
|
||||||
- exact loader-to-backend plan transfer and ordinary-loader fallback;
|
- 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
|
- rank-local mmap source retention, bounded two-slot staging, and adapter
|
||||||
transforms without parameter-side flags;
|
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;
|
- resident-layer requests requiring no-AllGather;
|
||||||
- DP request-wave validation for denoising-step compatibility;
|
- DP request-wave validation for denoising-step compatibility;
|
||||||
- bounded component allocator-cache retention, conservative/forced release,
|
- bounded component allocator-cache retention, conservative/forced release,
|
||||||
@@ -445,6 +520,9 @@ AllGather, because DP replicas did not share checkpoint-backed runtime
|
|||||||
weights. This is a functional and memory smoke test, not a production-quality
|
weights. This is a functional and memory smoke test, not a production-quality
|
||||||
performance or output-quality benchmark.
|
performance or output-quality benchmark.
|
||||||
|
|
||||||
|
These B300 DP waves intentionally reused the same prompt and seed, so they do
|
||||||
|
not establish independent-prompt isolation across DP replicas.
|
||||||
|
|
||||||
### Host-memory measurement
|
### Host-memory measurement
|
||||||
|
|
||||||
A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
|
A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ validation_paths:
|
|||||||
- tests/host_weight_runtime/**
|
- tests/host_weight_runtime/**
|
||||||
upstream_refs:
|
upstream_refs:
|
||||||
- https://github.com/vllm-project/vllm-omni/issues/6414
|
- https://github.com/vllm-project/vllm-omni/issues/6414
|
||||||
last_reviewed: 2026-08-21
|
last_reviewed: 2026-08-25
|
||||||
---
|
---
|
||||||
|
|
||||||
# Host Weight Runtime
|
# Host Weight Runtime
|
||||||
@@ -84,6 +84,13 @@ another. `WeightArtifactIdentity` includes:
|
|||||||
- producer version, implementation fingerprint, manifest schema, and restorer
|
- producer version, implementation fingerprint, manifest schema, and restorer
|
||||||
schema.
|
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
|
The identity excludes transfer policy such as registered mmap versus private
|
||||||
pinned staging. Those paths move the same host representation and must not
|
pinned staging. Those paths move the same host representation and must not
|
||||||
produce duplicate artifacts.
|
produce duplicate artifacts.
|
||||||
@@ -237,6 +244,12 @@ Every publication computes full file and tensor hashes. Local lookup supports:
|
|||||||
and
|
and
|
||||||
- `full_checksum`, which additionally reads and hashes every payload file.
|
- `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
|
Filesystem verity is modeled but unsupported. Invalid artifacts receive an
|
||||||
external deny marker before their shared lookup lock is released; a build owner
|
external deny marker before their shared lookup lock is released; a build owner
|
||||||
then quarantines the immutable directory under the exclusive lock.
|
then quarantines the immutable directory under the exclusive lock.
|
||||||
|
|||||||
@@ -36,7 +36,8 @@ vllm serve /path/to/model --omni \
|
|||||||
vllm serve /path/to/model --omni \
|
vllm serve /path/to/model --omni \
|
||||||
--enable-distributed-layerwise-offload \
|
--enable-distributed-layerwise-offload \
|
||||||
--data-parallel-size 4 \
|
--data-parallel-size 4 \
|
||||||
--dlo-no-use-allgather
|
--dlo-no-use-allgather \
|
||||||
|
--dlo-offload-components '{"text_encoder": false}'
|
||||||
|
|
||||||
# Sequence parallel deployment
|
# Sequence parallel deployment
|
||||||
vllm serve /path/to/model --omni \
|
vllm serve /path/to/model --omni \
|
||||||
@@ -51,6 +52,7 @@ omni = Omni(
|
|||||||
model="/path/to/model",
|
model="/path/to/model",
|
||||||
enable_distributed_layerwise_offload=True,
|
enable_distributed_layerwise_offload=True,
|
||||||
dlo_use_allgather=True,
|
dlo_use_allgather=True,
|
||||||
|
dlo_offload_components={"text_encoder": False},
|
||||||
)
|
)
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -63,8 +65,10 @@ omni = Omni(
|
|||||||
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
|
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
|
||||||
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
|
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
|
||||||
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
|
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
|
||||||
|
| `--dlo-offload-components JSON` | Permit or disable offload for discovered encoder/VAE components | `{}` |
|
||||||
| `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
|
| `--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 |
|
| `--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
|
## Host-weight loading
|
||||||
|
|
||||||
@@ -178,12 +182,48 @@ returning to `required`.
|
|||||||
|
|
||||||
On a cold start, the canonical loader remains authoritative and publishes a
|
On a cold start, the canonical loader remains authoritative and publishes a
|
||||||
validated final-layout artifact for later workers. A warm start restores the
|
validated final-layout artifact for later workers. A warm start restores the
|
||||||
DiT final tensors without ordinary DiT materialization, then DLO streams them
|
DiT final tensors without ordinary DiT materialization. DLO then attempts to
|
||||||
through the same two bounded host staging slots. The artifact identity includes
|
register the immutable mapping for direct H2D; if registration is unavailable,
|
||||||
the TP rank/size and SP layout, so TP1, TP2 rank-local shards, and distinct SP
|
it streams through the same two bounded host staging slots. The artifact
|
||||||
layouts do not alias one another. Publication failure remains separate from a
|
identity includes the TP rank/size and SP layout, so TP1, TP2 rank-local shards,
|
||||||
`preferred` serving startup; the next `required` startup provides the explicit
|
and distinct SP layouts do not alias one another. Publication failure remains
|
||||||
artifact-availability check.
|
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
|
When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
|
||||||
not resolve HWR sources or construct its store.
|
not resolve HWR sources or construct its store.
|
||||||
@@ -207,6 +247,25 @@ class MyPipeline(nn.Module):
|
|||||||
When no plan exists, discovery falls back to
|
When no plan exists, discovery falls back to
|
||||||
`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
|
`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
|
||||||
|
|
||||||
|
## Auxiliary component policy
|
||||||
|
|
||||||
|
`dlo_offload_components` follows the same component-map style as diffusion
|
||||||
|
quantization. Keys are discovered encoder or VAE paths, values are booleans,
|
||||||
|
and `default` is the optional fallback:
|
||||||
|
|
||||||
|
```python
|
||||||
|
dlo_offload_components = {
|
||||||
|
"text_encoder": False,
|
||||||
|
"default": True,
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
`True` permits the model's existing DLO behavior for that component. `False`
|
||||||
|
keeps the whole component resident, including its child blocks and non-block
|
||||||
|
modules. An empty map preserves current behavior. Unknown keys are rejected;
|
||||||
|
the map does not control DiT blocks because enabling DLO gives the backend
|
||||||
|
ownership of DiT streaming.
|
||||||
|
|
||||||
## Data-parallel concurrency
|
## Data-parallel concurrency
|
||||||
|
|
||||||
With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
|
With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
|
||||||
@@ -231,6 +290,13 @@ must enter each collective.
|
|||||||
are validated.
|
are validated.
|
||||||
- Resident leading layers require `--dlo-no-use-allgather` and a model
|
- Resident leading layers require `--dlo-no-use-allgather` and a model
|
||||||
`OffloadPlan` that declares eligible `resident_dit_paths`.
|
`OffloadPlan` that declares eligible `resident_dit_paths`.
|
||||||
|
- Component policy applies only to discovered auxiliary encoders and VAEs;
|
||||||
|
it cannot disable DLO for a DiT.
|
||||||
|
- MiniMax-H3's encoder TP group is global rather than DP-replica-local. An
|
||||||
|
encoder TP size that spans multiple DP replicas (for example TP4 with
|
||||||
|
DP2/TP2 or DP2/SP2 on four ranks) is not supported for independent-prompt
|
||||||
|
AllGather waves: the current pipeline broadcasts rank 0's conditioning to
|
||||||
|
both replicas. This limitation is independent of the component policy.
|
||||||
- DP concurrency requires an explicit, identical inference-step count.
|
- DP concurrency requires an explicit, identical inference-step count.
|
||||||
|
|
||||||
Sharing quantized or otherwise unvalidated transformed layouts through a
|
Sharing quantized or otherwise unvalidated transformed layouts through a
|
||||||
|
|||||||
@@ -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,5 +1,5 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||||
"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
|
"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
|
||||||
|
|
||||||
The DLO modes accept any positive data-parallel size. With AllGather enabled,
|
The DLO modes accept any positive data-parallel size. With AllGather enabled,
|
||||||
@@ -85,12 +85,20 @@ def parse_args() -> argparse.Namespace:
|
|||||||
parser.add_argument("--height", type=int, default=768)
|
parser.add_argument("--height", type=int, default=768)
|
||||||
parser.add_argument("--batch-wait-ms", type=float, default=500.0)
|
parser.add_argument("--batch-wait-ms", type=float, default=500.0)
|
||||||
parser.add_argument("--init-timeout", type=float, default=1800.0)
|
parser.add_argument("--init-timeout", type=float, default=1800.0)
|
||||||
|
parser.add_argument(
|
||||||
|
"--dlo-offload-components",
|
||||||
|
type=json.loads,
|
||||||
|
default=None,
|
||||||
|
help="Component policy JSON, e.g. '{\"text_encoder\": false}'",
|
||||||
|
)
|
||||||
parser.add_argument("--output", type=Path)
|
parser.add_argument("--output", type=Path)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
||||||
is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
|
is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
|
||||||
|
if args.dlo_offload_components and not is_dlo:
|
||||||
|
raise ValueError("--dlo-offload-components requires a DLO mode")
|
||||||
common: dict[str, Any] = {
|
common: dict[str, Any] = {
|
||||||
"model": args.model,
|
"model": args.model,
|
||||||
"trust_remote_code": True,
|
"trust_remote_code": True,
|
||||||
@@ -114,6 +122,7 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
|||||||
enable_distributed_layerwise_offload=True,
|
enable_distributed_layerwise_offload=True,
|
||||||
dlo_use_allgather=args.mode == "dlo",
|
dlo_use_allgather=args.mode == "dlo",
|
||||||
dlo_resident_layers=0,
|
dlo_resident_layers=0,
|
||||||
|
dlo_offload_components=({} if args.dlo_offload_components is None else args.dlo_offload_components),
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
common.update(
|
common.update(
|
||||||
|
|||||||
@@ -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.
|
"""Minimal single-input MiniCPM-o 4.5 Realtime duplex demo.
|
||||||
|
|
||||||
Run this after starting the duplex server. Strict lifecycle, overlap, and
|
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))
|
sys.path.insert(0, str(REPO_ROOT))
|
||||||
|
|
||||||
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||||
PCM16_BYTES_PER_SAMPLE,
|
PCM16_BYTES_PER_SAMPLE, # noqa: F401 - compatibility export
|
||||||
PCM16_SAMPLE_RATE,
|
PCM16_SAMPLE_RATE, # noqa: F401 - compatibility export
|
||||||
RealtimeDuplexClient,
|
RealtimeDuplexClient,
|
||||||
RealtimeEventCollector,
|
RealtimeEventCollector,
|
||||||
build_realtime_url,
|
build_realtime_url,
|
||||||
@@ -28,6 +31,13 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
|||||||
wait_for,
|
wait_for,
|
||||||
write_pcm16_wav,
|
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:
|
class _StreamingOutputWriter:
|
||||||
@@ -133,25 +143,6 @@ def _latest_model_decision(
|
|||||||
return 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:
|
def _response_in_progress(events: list[dict[str, object]]) -> bool:
|
||||||
return sum(event.get("type") == "response.created" for event in events) > sum(
|
return sum(event.get("type") == "response.created" for event in events) > sum(
|
||||||
event.get("type") == "response.done" for event in events
|
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 :])
|
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]:
|
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||||
input_pcm16 = read_pcm16_wav(Path(args.input_wav))
|
input_pcm16 = read_pcm16_wav(Path(args.input_wav))
|
||||||
if not input_pcm16:
|
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()
|
|
||||||
@@ -23,8 +23,6 @@ offline validation sample with TeaCache and online serving through the
|
|||||||
- Model: <https://huggingface.co/stabilityai/stable-audio-open-1.0>
|
- Model: <https://huggingface.co/stabilityai/stable-audio-open-1.0>
|
||||||
- Offline example:
|
- Offline example:
|
||||||
[`examples/offline_inference/text_to_audio`](../../examples/offline_inference/text_to_audio)
|
[`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:
|
- Audio generation API:
|
||||||
[`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md)
|
[`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md)
|
||||||
- Related issue:
|
- Related issue:
|
||||||
@@ -103,7 +101,7 @@ curl -X POST http://localhost:8091/v1/audio/generate \
|
|||||||
"seed": 42,
|
"seed": 42,
|
||||||
"response_format": "wav"
|
"response_format": "wav"
|
||||||
}' \
|
}' \
|
||||||
--output examples/online_serving/stable_audio/piano_10s.wav
|
--output piano_10s.wav
|
||||||
```
|
```
|
||||||
|
|
||||||
##### Verification
|
##### Verification
|
||||||
@@ -137,12 +135,12 @@ PY
|
|||||||
Validate the online output:
|
Validate the online output:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
ls -lh examples/online_serving/stable_audio/piano_10s.wav
|
ls -lh piano_10s.wav
|
||||||
|
|
||||||
python - <<'PY'
|
python - <<'PY'
|
||||||
import soundfile as sf
|
import soundfile as sf
|
||||||
|
|
||||||
path = "examples/online_serving/stable_audio/piano_10s.wav"
|
path = "piano_10s.wav"
|
||||||
audio, sample_rate = sf.read(path)
|
audio, sample_rate = sf.read(path)
|
||||||
print("sample_rate:", sample_rate)
|
print("sample_rate:", sample_rate)
|
||||||
print("shape:", audio.shape)
|
print("shape:", audio.shape)
|
||||||
|
|||||||
@@ -1,10 +1,12 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# 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
|
Unit tests for metrics.py
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import math
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from vllm.benchmarks.serve import TaskType
|
from vllm.benchmarks.serve import TaskType
|
||||||
|
|
||||||
@@ -66,6 +68,22 @@ def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutp
|
|||||||
return output
|
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
|
# total_input Tests
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
@@ -126,6 +144,166 @@ def test_audio_continuity_aggregation():
|
|||||||
assert p99 is not None and p99 > 0.4
|
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
|
# TTFT suppression for pure-audio (TTS) benchmarks
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
|
|||||||
@@ -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]
|
||||||
@@ -1,3 +1,6 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
@@ -22,6 +25,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
|||||||
"dataset_name",
|
"dataset_name",
|
||||||
[
|
[
|
||||||
"daily-omni",
|
"daily-omni",
|
||||||
|
"omniinteract",
|
||||||
"seed-tts",
|
"seed-tts",
|
||||||
"seed-tts-text",
|
"seed-tts-text",
|
||||||
"seed-tts-design",
|
"seed-tts-design",
|
||||||
@@ -55,6 +59,7 @@ def test_extend_omni_choices_updates_tracking_parser_shadow(dataset_name: str) -
|
|||||||
[
|
[
|
||||||
(["--print-stage"], "print_stage", True),
|
(["--print-stage"], "print_stage", True),
|
||||||
(["--daily-omni-input-mode", "audio"], "daily_omni_input_mode", "audio"),
|
(["--daily-omni-input-mode", "audio"], "daily_omni_input_mode", "audio"),
|
||||||
|
(["--omniinteract-subsets", "1qna"], "omniinteract_subsets", ["1qna"]),
|
||||||
(["--seed-tts-locale", "zh"], "seed_tts_locale", "zh"),
|
(["--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([])
|
args = parser.parse_args([])
|
||||||
assert args.print_stage is False
|
assert args.print_stage is False
|
||||||
assert args.daily_omni_input_mode == "all"
|
assert args.daily_omni_input_mode == "all"
|
||||||
|
assert args.omniinteract_subsets == ["1q1a", "1q1a_math", "1qna"]
|
||||||
assert args.seed_tts_locale == "en"
|
assert args.seed_tts_locale == "en"
|
||||||
assert args.explicit_keys == set()
|
assert args.explicit_keys == set()
|
||||||
|
|
||||||
|
|
||||||
def test_update_omni_help_updates_upstream_actions() -> None:
|
def test_update_omni_help_updates_upstream_actions() -> None:
|
||||||
parser = TrackingArgumentParser()
|
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("--percentile-metrics", help="Upstream percentile help.")
|
||||||
parser.add_argument("--random-mm-limit-mm-per-prompt", help="Upstream limit 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.")
|
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)
|
update_omni_help(parser)
|
||||||
|
|
||||||
actions = {action.dest: action for action in parser._actions}
|
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 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 "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
|
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
|
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(
|
@pytest.mark.parametrize(
|
||||||
("argv", "expected_extra_body", "expected_explicit"),
|
("argv", "expected_extra_body", "expected_explicit"),
|
||||||
[
|
[
|
||||||
|
|||||||
@@ -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: 22,
|
||||||
|
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||||
|
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||||
|
EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
|
||||||
|
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: 7,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
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": [
|
"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",
|
"dataset_name": "random-mm",
|
||||||
"backend": "openai-chat-omni",
|
"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",
|
"dataset_name": "random-mm",
|
||||||
"backend": "openai-chat-omni",
|
"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",
|
"endpoint": "/v1/completions",
|
||||||
"num_prompts": [
|
"num_prompts": [
|
||||||
4,
|
4,
|
||||||
16,
|
|
||||||
32,
|
32,
|
||||||
64,
|
|
||||||
128
|
128
|
||||||
],
|
],
|
||||||
"max_concurrency": [
|
"max_concurrency": [
|
||||||
1,
|
1,
|
||||||
4,
|
|
||||||
8,
|
8,
|
||||||
16,
|
|
||||||
32
|
32
|
||||||
],
|
],
|
||||||
"random_input_len": 2500,
|
"random_input_len": 2500,
|
||||||
@@ -54,9 +50,7 @@
|
|||||||
"H100": {
|
"H100": {
|
||||||
"mean_e2el_ms": [
|
"mean_e2el_ms": [
|
||||||
4619.0068,
|
4619.0068,
|
||||||
7868.3523,
|
|
||||||
9760.6175,
|
9760.6175,
|
||||||
13422.3491,
|
|
||||||
18007.6918
|
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",
|
"dataset_name": "random-mm",
|
||||||
"backend": "vllm",
|
"backend": "vllm",
|
||||||
|
|||||||
@@ -74,13 +74,14 @@
|
|||||||
"full_model",
|
"full_model",
|
||||||
"diffusion"
|
"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_type": "vllm-omni",
|
||||||
"server_params": {
|
"server_params": {
|
||||||
"model": "Qwen/Qwen-Image",
|
"model": "Qwen/Qwen-Image",
|
||||||
"serve_args": {
|
"serve_args": {
|
||||||
"enable-diffusion-pipeline-profiler": true,
|
"enable-diffusion-pipeline-profiler": true,
|
||||||
"step-execution": true
|
"step-execution": true,
|
||||||
|
"max-num-seqs": 8
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"benchmark_params": [
|
"benchmark_params": [
|
||||||
@@ -119,6 +120,27 @@
|
|||||||
"peak_memory_mb_mean": 64500.0
|
"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_audio": None,
|
||||||
"ref_text": None,
|
"ref_text": None,
|
||||||
},
|
},
|
||||||
("CustomVoice", "does not support"),
|
("Invalid voice", "vivian"),
|
||||||
id="customvoice_task_not_supported",
|
id="customvoice_invalid_voice",
|
||||||
),
|
),
|
||||||
pytest.param(
|
pytest.param(
|
||||||
{"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
|
{"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"),
|
("mutually exclusive", "speaker_embedding", "ref_audio"),
|
||||||
id="speaker_embedding_with_ref_audio_mutually_exclusive",
|
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(
|
pytest.param(
|
||||||
{"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
|
{"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
|
||||||
("max_new_tokens", "exceed 4096"),
|
("max_new_tokens", "exceed 4096"),
|
||||||
id="max_new_tokens_above_max",
|
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({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
|
||||||
pytest.param(
|
pytest.param(
|
||||||
{"initial_codec_chunk_frames": -1},
|
{"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_empty"),
|
||||||
pytest.param({"input": " ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"),
|
pytest.param({"input": " ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"),
|
||||||
pytest.param(
|
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(
|
pytest.param(
|
||||||
{
|
{
|
||||||
@@ -554,14 +560,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, online_client: Onli
|
|||||||
pytest.param(
|
pytest.param(
|
||||||
"batch",
|
"batch",
|
||||||
{"max_new_tokens": 0},
|
{"max_new_tokens": 0},
|
||||||
("max_new_tokens", "least 1"),
|
("max_new_tokens", "greater_than_equal"),
|
||||||
id="batch_max_new_tokens_below_min",
|
id="batch_max_new_tokens_below_min",
|
||||||
marks=_SKIP_ISSUE_3649,
|
marks=_SKIP_ISSUE_3649,
|
||||||
),
|
),
|
||||||
pytest.param(
|
pytest.param(
|
||||||
"item",
|
"item",
|
||||||
{"max_new_tokens": 0},
|
{"max_new_tokens": 0},
|
||||||
("max_new_tokens", "least 1"),
|
("max_new_tokens", "greater_than_equal"),
|
||||||
id="item_max_new_tokens_below_min",
|
id="item_max_new_tokens_below_min",
|
||||||
marks=_SKIP_ISSUE_3649,
|
marks=_SKIP_ISSUE_3649,
|
||||||
),
|
),
|
||||||
|
|||||||
@@ -72,6 +72,9 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
|||||||
body["top_logprobs"] = 5
|
body["top_logprobs"] = 5
|
||||||
elif case_id == "speaker_unknown":
|
elif case_id == "speaker_unknown":
|
||||||
body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz"
|
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:
|
else:
|
||||||
raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
|
raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
|
||||||
return {"json": body, "timeout": 120}
|
return {"json": body, "timeout": 120}
|
||||||
@@ -119,7 +122,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
|||||||
pytest.param(
|
pytest.param(
|
||||||
"response_format_json_schema_incomplete",
|
"response_format_json_schema_incomplete",
|
||||||
400,
|
400,
|
||||||
("response_format", "value_error", "json_schema"),
|
("response_format", "BadRequestError", "json_schema"),
|
||||||
id="invalid_response_format_json_schema",
|
id="invalid_response_format_json_schema",
|
||||||
),
|
),
|
||||||
pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
|
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"),
|
("Invalid speaker", "Supported"),
|
||||||
id="speaker_unknown_preset",
|
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)
|
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
|
||||||
|
|||||||
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""Run the same prompts sequentially then in parallel and print a comparison."""
|
"""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:
|
try:
|
||||||
await warmup(omni, WARMUP_PROMPTS)
|
await warmup(omni, WARMUP_PROMPTS)
|
||||||
single_time = await run_single(omni, 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)
|
await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
|
||||||
return
|
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:
|
try:
|
||||||
await warmup(omni, WARMUP_PROMPTS)
|
await warmup(omni, WARMUP_PROMPTS)
|
||||||
|
|
||||||
@@ -524,7 +524,12 @@ if __name__ == "__main__":
|
|||||||
parser = argparse.ArgumentParser(description="E2E diffusion concurrent benchmark / validation")
|
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("--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("--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(
|
parser.add_argument(
|
||||||
"--mode",
|
"--mode",
|
||||||
choices=["batch", "single", "compare", "validate"],
|
choices=["batch", "single", "compare", "validate"],
|
||||||
|
|||||||
@@ -24,6 +24,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
|
|||||||
HostWeightPlanResult,
|
HostWeightPlanResult,
|
||||||
TensorBinding,
|
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.helios import HeliosPipeline
|
||||||
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
|
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
|
||||||
from vllm_omni.diffusion.registry import initialize_model
|
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"),
|
str(canonical_root / "model.safetensors"),
|
||||||
)
|
)
|
||||||
store_root = tmp_path / "hwr-store"
|
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]:
|
def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
|
||||||
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
|
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 is not None
|
||||||
assert warm_plan.lease_carrier is not None
|
assert warm_plan.lease_carrier is not None
|
||||||
warm_plan.lease_carrier.close()
|
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):
|
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 dataclasses
|
||||||
import gc
|
import gc
|
||||||
|
import json
|
||||||
|
import threading
|
||||||
from collections.abc import Sequence
|
from collections.abc import Sequence
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
@@ -24,10 +26,12 @@ from vllm_omni.diffusion.model_loader.host_weights import (
|
|||||||
FinalLayoutRequest,
|
FinalLayoutRequest,
|
||||||
FinalLayoutTensorRestorer,
|
FinalLayoutTensorRestorer,
|
||||||
ImplementationIdentity,
|
ImplementationIdentity,
|
||||||
|
NodeSourceDigestCache,
|
||||||
PreparedWeightSource,
|
PreparedWeightSource,
|
||||||
WeightSourceKind,
|
WeightSourceKind,
|
||||||
build_final_layout_identity,
|
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 (
|
from vllm_omni.diffusion.model_loader.host_weights.contracts import (
|
||||||
FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
|
FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
|
||||||
FinalLayoutArtifactSpec,
|
FinalLayoutArtifactSpec,
|
||||||
@@ -380,6 +384,7 @@ def _identity(
|
|||||||
*,
|
*,
|
||||||
request: FinalLayoutRequest | None = None,
|
request: FinalLayoutRequest | None = None,
|
||||||
policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY,
|
policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY,
|
||||||
|
source_digest_cache: NodeSourceDigestCache | None = None,
|
||||||
) -> FinalLayoutIdentityContext:
|
) -> FinalLayoutIdentityContext:
|
||||||
transformer = model.get_submodule("transformer")
|
transformer = model.get_submodule("transformer")
|
||||||
return build_final_layout_identity(
|
return build_final_layout_identity(
|
||||||
@@ -388,6 +393,7 @@ def _identity(
|
|||||||
prepared_sources=(source,),
|
prepared_sources=(source,),
|
||||||
request=request or _request(),
|
request=request or _request(),
|
||||||
policy=policy,
|
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
|
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:
|
def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None:
|
||||||
model = _TinyPipeline()
|
model = _TinyPipeline()
|
||||||
source_root = tmp_path / "local-symlink-source"
|
source_root = tmp_path / "local-symlink-source"
|
||||||
|
|||||||
@@ -99,3 +99,22 @@ def test_manual_component_offload_failure_forces_retained_cache_release(mocker):
|
|||||||
assert component.load_to_device.call_count == 2
|
assert component.load_to_device.call_count == 2
|
||||||
assert component.offload_to_cpu.call_count == 2
|
assert component.offload_to_cpu.call_count == 2
|
||||||
pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)
|
pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)
|
||||||
|
|
||||||
|
|
||||||
|
def test_resident_component_skips_manual_dlo_lifecycle(mocker):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||||
|
torch.nn.Module.__init__(pipeline)
|
||||||
|
pipeline.od_config = mocker.Mock()
|
||||||
|
pipeline.od_config.enable_layerwise_offload = False
|
||||||
|
pipeline.od_config.enable_distributed_layerwise_offload = True
|
||||||
|
pipeline._model_cpu_offload_modules = []
|
||||||
|
component = mocker.Mock()
|
||||||
|
component._omni_dlo_offload_enabled = False
|
||||||
|
|
||||||
|
with pipeline._component_on_device(component):
|
||||||
|
pass
|
||||||
|
|
||||||
|
component.load_to_device.assert_not_called()
|
||||||
|
component.offload_to_cpu.assert_not_called()
|
||||||
@@ -0,0 +1,142 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
|
||||||
|
from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
|
||||||
|
DistributedLayerwiseOffloadBackend,
|
||||||
|
)
|
||||||
|
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
|
||||||
|
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||||
|
|
||||||
|
|
||||||
|
def _od_config(**overrides):
|
||||||
|
values = {
|
||||||
|
"enable_cpu_offload": False,
|
||||||
|
"enable_layerwise_offload": False,
|
||||||
|
"enable_distributed_layerwise_offload": True,
|
||||||
|
"dlo_use_allgather": False,
|
||||||
|
"dlo_resident_layers": 0,
|
||||||
|
"dlo_offload_components": {},
|
||||||
|
"pin_cpu_memory": False,
|
||||||
|
"parallel_config": SimpleNamespace(
|
||||||
|
use_hsdp=False,
|
||||||
|
data_parallel_size=1,
|
||||||
|
sequence_parallel_size=1,
|
||||||
|
),
|
||||||
|
}
|
||||||
|
values.update(overrides)
|
||||||
|
return SimpleNamespace(**values)
|
||||||
|
|
||||||
|
|
||||||
|
def _backend(policy: dict[str, bool]) -> DistributedLayerwiseOffloadBackend:
|
||||||
|
backend = object.__new__(DistributedLayerwiseOffloadBackend)
|
||||||
|
backend.config = OffloadConfig(
|
||||||
|
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||||
|
pin_cpu_memory=False,
|
||||||
|
dlo_use_allgather=False,
|
||||||
|
dlo_offload_components=policy,
|
||||||
|
)
|
||||||
|
backend.device = torch.device("cpu")
|
||||||
|
return backend
|
||||||
|
|
||||||
|
|
||||||
|
def _modules() -> PipelineModules:
|
||||||
|
return PipelineModules(
|
||||||
|
dits=[],
|
||||||
|
dit_names=[],
|
||||||
|
encoders=[nn.Module()],
|
||||||
|
encoder_names=["text_encoder"],
|
||||||
|
vaes=[nn.Module(), nn.Module()],
|
||||||
|
vae_names=["video_vae", "audio_vae"],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_component_policy_resolves_exact_then_default():
|
||||||
|
backend = _backend({"text_encoder": False, "default": True})
|
||||||
|
|
||||||
|
assert not backend._component_offload_enabled("text_encoder")
|
||||||
|
assert backend._component_offload_enabled("video_vae")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("value", "error", "match"),
|
||||||
|
[
|
||||||
|
([], TypeError, "must be a mapping"),
|
||||||
|
(["text_encoder"], TypeError, "must be a mapping"),
|
||||||
|
(False, TypeError, "must be a mapping"),
|
||||||
|
({1: False}, TypeError, "keys must be non-empty strings"),
|
||||||
|
({"text_encoder": 0}, TypeError, "must be a boolean"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_component_policy_rejects_invalid_shapes(value, error, match):
|
||||||
|
with pytest.raises(error, match=match):
|
||||||
|
OffloadConfig.from_od_config(_od_config(dlo_offload_components=value))
|
||||||
|
|
||||||
|
|
||||||
|
def test_component_policy_accepts_none_as_default():
|
||||||
|
config = OffloadConfig.from_od_config(_od_config(dlo_offload_components=None))
|
||||||
|
|
||||||
|
assert config.dlo_offload_components == {}
|
||||||
|
|
||||||
|
|
||||||
|
def test_component_policy_requires_dlo():
|
||||||
|
with pytest.raises(ValueError, match="requires distributed layerwise offload"):
|
||||||
|
OffloadConfig.from_od_config(
|
||||||
|
_od_config(
|
||||||
|
enable_distributed_layerwise_offload=False,
|
||||||
|
dlo_offload_components={"text_encoder": False},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_component_policy_rejects_unknown_or_dit_names():
|
||||||
|
backend = _backend({"transformer": False})
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="unknown auxiliary components.*transformer"):
|
||||||
|
backend._validate_component_offload_policy(_modules())
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("policy", "encoder_stage", "video_stage", "audio_stage", "try_encoder"),
|
||||||
|
[
|
||||||
|
({}, True, True, True, True),
|
||||||
|
({"text_encoder": False}, False, True, True, False),
|
||||||
|
({"default": False, "video_vae": True}, False, True, False, False),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_prepare_auxiliary_components_applies_policy(
|
||||||
|
mocker,
|
||||||
|
policy,
|
||||||
|
encoder_stage,
|
||||||
|
video_stage,
|
||||||
|
audio_stage,
|
||||||
|
try_encoder,
|
||||||
|
):
|
||||||
|
backend = _backend(policy)
|
||||||
|
backend._try_layerwise_offload_encoder = mocker.Mock()
|
||||||
|
backend._register_on_demand_hook = mocker.Mock()
|
||||||
|
modules = _modules()
|
||||||
|
plan = OffloadPlan(
|
||||||
|
encoder_block_attrs={"text_encoder": ("layers",)},
|
||||||
|
on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
|
||||||
|
)
|
||||||
|
|
||||||
|
backend._prepare_auxiliary_components(modules, plan)
|
||||||
|
|
||||||
|
assert backend._try_layerwise_offload_encoder.called is try_encoder
|
||||||
|
assert modules.encoders[0]._omni_dlo_offload_enabled is encoder_stage
|
||||||
|
assert modules.vaes[0]._omni_dlo_offload_enabled is video_stage
|
||||||
|
assert modules.vaes[1]._omni_dlo_offload_enabled is audio_stage
|
||||||
|
assert backend._register_on_demand_hook.call_args_list == [
|
||||||
|
mocker.call(modules.encoders[0], "text_encoder", stage_on_demand=encoder_stage),
|
||||||
|
mocker.call(modules.vaes[0], "video_vae", stage_on_demand=video_stage),
|
||||||
|
mocker.call(modules.vaes[1], "audio_vae", stage_on_demand=audio_stage),
|
||||||
|
]
|
||||||
@@ -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 gc
|
||||||
import json
|
import json
|
||||||
|
import weakref
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
from types import SimpleNamespace
|
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
|
import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
|
||||||
from tests.helpers.runtime import get_distributed_init_method
|
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 (
|
from vllm_omni.diffusion.model_loader.host_weight_plan import (
|
||||||
HostWeightPlan,
|
HostWeightPlan,
|
||||||
TensorBinding,
|
TensorBinding,
|
||||||
@@ -33,12 +35,14 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
|
|||||||
DistributedLayerwiseOffloadHook,
|
DistributedLayerwiseOffloadHook,
|
||||||
PinnedResidentLayerGroup,
|
PinnedResidentLayerGroup,
|
||||||
)
|
)
|
||||||
|
from vllm_omni.diffusion.offloader.host_registration import HostRegistrationCleanupError
|
||||||
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
|
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
|
||||||
from vllm_omni.diffusion.offloader.offload_plan import (
|
from vllm_omni.diffusion.offloader.offload_plan import (
|
||||||
OffloadPlan,
|
OffloadPlan,
|
||||||
get_offload_plan,
|
get_offload_plan,
|
||||||
)
|
)
|
||||||
from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
|
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
|
from vllm_omni.platforms import current_omni_platform
|
||||||
|
|
||||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||||
@@ -703,11 +707,15 @@ class _HWRPipeline(nn.Module):
|
|||||||
|
|
||||||
|
|
||||||
class _FakeHostWeightLease:
|
class _FakeHostWeightLease:
|
||||||
def __init__(self, resolution_id: str):
|
def __init__(self, resolution_id: str, events: list[str] | None = None):
|
||||||
self.closed = False
|
self.closed = False
|
||||||
self.provenance = SimpleNamespace(resolution_id=resolution_id)
|
self.provenance = SimpleNamespace(resolution_id=resolution_id)
|
||||||
|
self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),)
|
||||||
|
self._events = events
|
||||||
|
|
||||||
def close(self):
|
def close(self):
|
||||||
|
if self._events is not None:
|
||||||
|
self._events.append("lease")
|
||||||
self.closed = True
|
self.closed = True
|
||||||
|
|
||||||
|
|
||||||
@@ -757,6 +765,203 @@ def _hwr_backend():
|
|||||||
return _HWRPipeline(), backend, carrier, lease
|
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):
|
class _MultiBlockModel(nn.Module):
|
||||||
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
|
_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
|
||||||
@@ -1,5 +1,5 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# 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.
|
Unit tests for WorkerWrapperBase class.
|
||||||
@@ -12,6 +12,7 @@ This module tests the WorkerWrapperBase implementation:
|
|||||||
- Dynamic worker class extension
|
- Dynamic worker class extension
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@@ -204,6 +205,27 @@ class TestWorkerWrapperBaseDelegation:
|
|||||||
wrapper.worker.shutdown.assert_called_once()
|
wrapper.worker.shutdown.assert_called_once()
|
||||||
assert result is None
|
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
|
# Tests: execute_method
|
||||||
|
|||||||
@@ -1,3 +1,7 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import base64
|
import base64
|
||||||
import wave
|
import wave
|
||||||
from urllib.parse import parse_qs, urlsplit
|
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():
|
def test_realtime_client_builds_explicit_native_duplex_url():
|
||||||
url = build_realtime_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",
|
"openbmb/MiniCPM-o-4_5",
|
||||||
session_id="session-a",
|
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():
|
def test_realtime_client_builds_resume_only_url_when_autostart_disabled():
|
||||||
url = build_realtime_url(
|
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",
|
"openbmb/MiniCPM-o-4_5",
|
||||||
autostart=False,
|
autostart=False,
|
||||||
)
|
)
|
||||||
@@ -73,10 +77,15 @@ async def test_realtime_client_configure_omits_ref_audio_by_default():
|
|||||||
|
|
||||||
client = Client()
|
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"]
|
session = client.sent[0]["session"]
|
||||||
assert "ref_audio" not in session
|
assert "ref_audio" not in session
|
||||||
|
assert session["idle_timeout_s"] == 900
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -100,6 +109,74 @@ async def test_realtime_client_configure_sends_explicit_ref_audio():
|
|||||||
|
|
||||||
session = client.sent[0]["session"]
|
session = client.sent[0]["session"]
|
||||||
assert session["ref_audio"] == "data:audio/wav;base64,AAAA"
|
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
|
@pytest.mark.asyncio
|
||||||
@@ -235,6 +312,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals():
|
|||||||
"output_token_count": 4,
|
"output_token_count": 4,
|
||||||
"ttft_ms": 120.0,
|
"ttft_ms": 120.0,
|
||||||
"tpot_ms": 15.0,
|
"tpot_ms": 15.0,
|
||||||
|
"itls_ms": [10.0, 14.0, 18.0],
|
||||||
"inter_token_interval_ms": {
|
"inter_token_interval_ms": {
|
||||||
"count": 3,
|
"count": 3,
|
||||||
"mean": 14.0,
|
"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,
|
read_pcm16_wav,
|
||||||
summarize_session_request_metrics,
|
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
|
_url_with_model = build_realtime_url
|
||||||
_read_wav_pcm16 = read_pcm16_wav
|
_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
|
@dataclass
|
||||||
class DemoState:
|
class DemoState:
|
||||||
events: list[dict[str, object]] = field(default_factory=list)
|
events: list[dict[str, object]] = field(default_factory=list)
|
||||||
|
|||||||
@@ -7,7 +7,6 @@ E2E Online tests for Qwen3-Omni model.
|
|||||||
import os
|
import os
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from openai import BadRequestError
|
|
||||||
|
|
||||||
from tests.helpers.mark import hardware_test
|
from tests.helpers.mark import hardware_test
|
||||||
from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
|
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(
|
pytest.param(
|
||||||
OmniServerParams(
|
OmniServerParams(
|
||||||
model=model,
|
model=model,
|
||||||
@@ -97,7 +96,7 @@ test_token_params = [
|
|||||||
use_stage_cli=True,
|
use_stage_cli=True,
|
||||||
server_args=["--async-chunk"],
|
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)
|
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
@pytest.mark.parametrize("omni_server", test_params + reduced_token_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)
|
|
||||||
def test_text_to_text_audio_001(omni_server, online_client) -> None:
|
def test_text_to_text_audio_001(omni_server, online_client) -> None:
|
||||||
"""
|
"""
|
||||||
Input Modal: text
|
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)
|
@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_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)
|
|
||||||
def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
|
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)
|
Input Modal: text, long-duration audio (~LONG_AUDIO_DURATION_SEC s WAV)
|
||||||
Output Modal: text, audio
|
Output Modal: text, audio
|
||||||
Input Setting: stream=False
|
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']}"
|
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(LONG_AUDIO_DURATION_SEC, 1)['base64']}"
|
||||||
messages = dummy_messages_from_mix_data(
|
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)
|
@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_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)
|
|
||||||
def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
|
def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
|
||||||
"""
|
"""
|
||||||
Input Modal: text, high-resolution image (1080p-class JPEG)
|
Input Modal: text, high-resolution image (1080p-class JPEG)
|
||||||
Output Modal: text, audio
|
Output Modal: text, audio
|
||||||
Input Setting: stream=False
|
Input Setting: stream=False
|
||||||
Datasets: single request
|
Datasets: few requests
|
||||||
"""
|
"""
|
||||||
image_data_url = (
|
image_data_url = (
|
||||||
f"data:image/jpeg;base64,{generate_synthetic_image(LARGE_IMAGE_WIDTH, LARGE_IMAGE_HEIGHT)['base64']}"
|
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)
|
@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:
|
def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
|
||||||
"""
|
"""
|
||||||
Input Modal: text, video
|
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)
|
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)
|
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||||
def test_audio_in_video_001(omni_server, online_client) -> None:
|
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)
|
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||||
def test_speaker_001(omni_server, online_client) -> None:
|
def test_speaker_001(omni_server, online_client) -> None:
|
||||||
"""
|
"""
|
||||||
Input Modal: text only (one-word answer constraint).
|
Input Modal: Chinese text.
|
||||||
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
|
Output Modal: Chinese text and audio using an uppercase preset speaker name.
|
||||||
Input Setting: stream=True
|
Input Setting: stream=True
|
||||||
Datasets: single request
|
Datasets: single request
|
||||||
"""
|
"""
|
||||||
messages = dummy_messages_from_mix_data(
|
messages = dummy_messages_from_mix_data(
|
||||||
system_prompt=get_system_prompt(),
|
system_prompt=get_system_prompt(),
|
||||||
content_text=get_prompt("text"),
|
content_text=get_prompt("text_chinese"),
|
||||||
)
|
)
|
||||||
|
|
||||||
request_config = {
|
request_config = {
|
||||||
"model": omni_server.model,
|
"model": omni_server.model,
|
||||||
"messages": messages,
|
"messages": messages,
|
||||||
"stream": True,
|
"stream": True,
|
||||||
"speaker": "Chelsie",
|
"speaker": "CHELSIE",
|
||||||
"key_words": {"text": ["beijing"]},
|
"key_words": {"text": ["北京"]},
|
||||||
}
|
}
|
||||||
|
|
||||||
online_client.send_omni_request(request_config)
|
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}")
|
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)
|
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||||
def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
|
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).
|
Input Modal: text only (long-form generation prompt).
|
||||||
Output Modal: text, audio (default ``modalities``);
|
Output Modal: text, audio (default ``modalities``);
|
||||||
Input Setting: stream=True
|
Input Setting: stream=True
|
||||||
Datasets: single request
|
Datasets: few requests
|
||||||
"""
|
"""
|
||||||
messages = dummy_messages_from_mix_data(
|
messages = dummy_messages_from_mix_data(
|
||||||
system_prompt=get_system_prompt(),
|
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 ""
|
text = responses[0].text_content if responses else ""
|
||||||
word_count = len(text.split())
|
word_count = len(text.split())
|
||||||
assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
|
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
|
- Qwen-Image-2512
|
||||||
|
|
||||||
One feature per test case, matching the Test Plan in PR #1682 (Qwen-Image-Edit).
|
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,
|
Nightly covers each feature once, alternating models (5× Qwen-Image, 4× Qwen-Image-2512).
|
||||||
CFG-Parallel, Tensor-Parallel, VAE-Patch-Parallel, CPU offload (layerwise).
|
Ulysses stays on Qwen-Image. See docs/user_guide/diffusion_acceleration.md.
|
||||||
See docs/user_guide/diffusion_acceleration.md.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@@ -25,141 +24,116 @@ NEGATIVE_PROMPT = "blurry, low quality"
|
|||||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||||
PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
|
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):
|
# One server per feature. Alternate models; keep feature pytest ids unchanged.
|
||||||
return [
|
FEATURE_CASES = [
|
||||||
pytest.param(
|
pytest.param(
|
||||||
OmniServerParams(
|
OmniServerParams(
|
||||||
model=model,
|
model=MODEL_IMAGE,
|
||||||
server_args=["--enable-cpu-offload"],
|
server_args=["--enable-cpu-offload"],
|
||||||
),
|
|
||||||
id="cpu_offload",
|
|
||||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
|
||||||
),
|
),
|
||||||
pytest.param(
|
id="cpu_offload",
|
||||||
OmniServerParams(model=model, server_args=["--step-execution"]),
|
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||||
id="step_execution",
|
),
|
||||||
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(
|
id="layerwise_offload",
|
||||||
OmniServerParams(model=model, server_args=["--cache-backend", "tea_cache"]),
|
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||||
id="cache_tea_cache",
|
),
|
||||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
pytest.param(
|
||||||
|
OmniServerParams(
|
||||||
|
model=MODEL_IMAGE,
|
||||||
|
server_args=[
|
||||||
|
"--cache-backend",
|
||||||
|
"cache_dit",
|
||||||
|
"--ulysses-degree",
|
||||||
|
"2",
|
||||||
|
],
|
||||||
),
|
),
|
||||||
pytest.param(
|
id="ulysses_2",
|
||||||
OmniServerParams(
|
marks=PARALLEL_FEATURE_MARKS,
|
||||||
model=model,
|
),
|
||||||
server_args=[
|
pytest.param(
|
||||||
"--cache-backend",
|
OmniServerParams(
|
||||||
"cache_dit",
|
model=MODEL_2512,
|
||||||
"--enable-layerwise-offload",
|
server_args=[
|
||||||
],
|
"--cache-backend",
|
||||||
),
|
"cache_dit",
|
||||||
id="layerwise_offload",
|
"--ring",
|
||||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
"2",
|
||||||
|
],
|
||||||
),
|
),
|
||||||
pytest.param(
|
id="ring_2",
|
||||||
OmniServerParams(
|
marks=PARALLEL_FEATURE_MARKS,
|
||||||
model=model,
|
),
|
||||||
server_args=[
|
pytest.param(
|
||||||
"--cache-backend",
|
OmniServerParams(
|
||||||
"cache_dit",
|
model=MODEL_IMAGE,
|
||||||
"--ulysses-degree",
|
server_args=[
|
||||||
"2",
|
"--cache-backend",
|
||||||
],
|
"tea_cache",
|
||||||
),
|
"--cfg-parallel-size",
|
||||||
id="ulysses_2",
|
"2",
|
||||||
marks=PARALLEL_FEATURE_MARKS,
|
],
|
||||||
),
|
),
|
||||||
pytest.param(
|
id="cfg_parallel_2",
|
||||||
OmniServerParams(
|
marks=PARALLEL_FEATURE_MARKS,
|
||||||
model=model,
|
),
|
||||||
server_args=[
|
pytest.param(
|
||||||
"--cache-backend",
|
OmniServerParams(
|
||||||
"cache_dit",
|
model=MODEL_2512,
|
||||||
"--ring",
|
server_args=[
|
||||||
"2",
|
"--cache-backend",
|
||||||
],
|
"cache_dit",
|
||||||
),
|
"--tensor-parallel-size",
|
||||||
id="ring_2",
|
"2",
|
||||||
marks=PARALLEL_FEATURE_MARKS,
|
"--vae-patch-parallel-size",
|
||||||
|
"2",
|
||||||
|
"--vae-use-tiling",
|
||||||
|
"--diffusion-quantization-config",
|
||||||
|
'{"method":"fp8"}',
|
||||||
|
],
|
||||||
),
|
),
|
||||||
pytest.param(
|
id="vae_patch_parallel_2",
|
||||||
OmniServerParams(
|
marks=PARALLEL_FEATURE_MARKS,
|
||||||
model=model,
|
),
|
||||||
server_args=[
|
pytest.param(
|
||||||
"--cache-backend",
|
OmniServerParams(
|
||||||
"tea_cache",
|
model=MODEL_IMAGE,
|
||||||
"--cfg-parallel-size",
|
server_args=[
|
||||||
"2",
|
"--use-hsdp",
|
||||||
],
|
"--hsdp-shard-size",
|
||||||
),
|
"2",
|
||||||
id="cfg_parallel_2",
|
],
|
||||||
marks=PARALLEL_FEATURE_MARKS,
|
|
||||||
),
|
),
|
||||||
pytest.param(
|
id="parallel_hsdp",
|
||||||
OmniServerParams(
|
marks=PARALLEL_FEATURE_MARKS,
|
||||||
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,
|
|
||||||
),
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True)
|
||||||
"omni_server",
|
|
||||||
_get_diffusion_feature_cases("Qwen/Qwen-Image"),
|
|
||||||
indirect=True,
|
|
||||||
)
|
|
||||||
def test_qwen_image(omni_server: OmniServer, online_client: OnlineOmniClient):
|
def test_qwen_image(omni_server: OmniServer, online_client: OnlineOmniClient):
|
||||||
"""Test each diffusion feature with Qwen-Image (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,
|
|
||||||
"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."""
|
|
||||||
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
|
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
|
||||||
request_config = {
|
request_config = {
|
||||||
"model": omni_server.model,
|
"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"
|
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.diffusion.stage_diffusion_client as client_mod
|
||||||
import vllm_omni.engine.stage_init_utils as init_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,
|
use_inline=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
assert od_config.max_num_seqs == 4
|
assert od_config.max_num_seqs == 1
|
||||||
assert captured == {
|
assert captured == {
|
||||||
"model": "dummy-model",
|
"model": "dummy-model",
|
||||||
"config": od_config,
|
"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_client as client_mod
|
||||||
import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
|
import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
|
||||||
import vllm_omni.engine.stage_engine_startup as startup_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 result is sentinel_client
|
||||||
assert od_config.max_num_seqs == 4
|
assert od_config.max_num_seqs == 1
|
||||||
assert resources.manager is proc_manager
|
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):
|
def test_stage_runtime_initializes_stage_pools(monkeypatch):
|
||||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||||
|
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# 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.
|
"""Tests for error propagation paths within the Orchestrator.
|
||||||
|
|
||||||
Covers:
|
Covers:
|
||||||
@@ -16,7 +16,9 @@ from types import SimpleNamespace
|
|||||||
|
|
||||||
import janus
|
import janus
|
||||||
import pytest
|
import pytest
|
||||||
|
from vllm.sampling_params import SamplingParams
|
||||||
from vllm.v1.engine.exceptions import EngineDeadError
|
from vllm.v1.engine.exceptions import EngineDeadError
|
||||||
|
from vllm.v1.serial_utils import MsgpackEncoder
|
||||||
|
|
||||||
from vllm_omni.engine.messages import (
|
from vllm_omni.engine.messages import (
|
||||||
AddCompanionRequestMessage,
|
AddCompanionRequestMessage,
|
||||||
@@ -46,8 +48,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|||||||
|
|
||||||
|
|
||||||
def _sampling_params(max_tokens: int = 4):
|
def _sampling_params(max_tokens: int = 4):
|
||||||
from vllm.sampling_params import SamplingParams
|
|
||||||
|
|
||||||
return SamplingParams(max_tokens=max_tokens)
|
return SamplingParams(max_tokens=max_tokens)
|
||||||
|
|
||||||
|
|
||||||
@@ -188,6 +188,45 @@ class _UnavailableDiffusionStageClient(FakeStageClient):
|
|||||||
raise StageUnavailableError(f"stage {self.stage_id} has no live replica")
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None:
|
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
|
"""Multi-replica stage: one replica dying must fail only the requests bound
|
||||||
|
|||||||
@@ -1175,7 +1175,7 @@ class TestSingleStageReplicaInitialization:
|
|||||||
os.environ[device_env_var] = prev_device_env
|
os.environ[device_env_var] = prev_device_env
|
||||||
|
|
||||||
assert result is sentinel_client
|
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(
|
mock_register.assert_called_once_with(
|
||||||
omni_master_address="127.0.0.1",
|
omni_master_address="127.0.0.1",
|
||||||
omni_master_port=25000,
|
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
|
# tests/entrypoints/openai/test_serving_speech.py
|
||||||
import asyncio
|
import asyncio
|
||||||
import base64
|
import base64
|
||||||
@@ -4147,16 +4150,15 @@ class TestCosyVoice3Serving:
|
|||||||
error = cosyvoice3_server._validate_tts_request(request)
|
error = cosyvoice3_server._validate_tts_request(request)
|
||||||
assert error is None
|
assert error is None
|
||||||
|
|
||||||
def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server):
|
def test_validate_cosyvoice3_max_new_tokens_range(self):
|
||||||
request = OpenAICreateSpeechRequest(
|
"""Ensure max_new_tokens below the minimum is rejected during request validation."""
|
||||||
input="Hello",
|
with pytest.raises(ValidationError, match="max_new_tokens"):
|
||||||
ref_audio="data:audio/wav;base64,abc",
|
OpenAICreateSpeechRequest(
|
||||||
ref_text="hello",
|
input="Hello",
|
||||||
max_new_tokens=0,
|
ref_audio="data:audio/wav;base64,abc",
|
||||||
)
|
ref_text="hello",
|
||||||
error = cosyvoice3_server._validate_tts_request(request)
|
max_new_tokens=0,
|
||||||
assert error is not None
|
)
|
||||||
assert "max_new_tokens" in error
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),
|
("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),
|
||||||
|
|||||||
@@ -332,7 +332,7 @@ def test_serve_cli_forwards_distilled_lora_to_diffusion_stage():
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
def test_serve_cli_forwards_distributed_offload_residency():
|
def test_serve_cli_forwards_distributed_offload_options():
|
||||||
"""Ensure the two-GPU DLO placement controls reach the diffusion stage."""
|
"""Ensure the two-GPU DLO placement controls reach the diffusion stage."""
|
||||||
parser = TrackingArgumentParser()
|
parser = TrackingArgumentParser()
|
||||||
subparsers = parser.add_subparsers(dest="command")
|
subparsers = parser.add_subparsers(dest="command")
|
||||||
@@ -347,6 +347,8 @@ def test_serve_cli_forwards_distributed_offload_residency():
|
|||||||
"--dlo-no-use-allgather",
|
"--dlo-no-use-allgather",
|
||||||
"--dlo-resident-layers",
|
"--dlo-resident-layers",
|
||||||
"20",
|
"20",
|
||||||
|
"--dlo-offload-components",
|
||||||
|
'{"text_encoder": false, "default": true}',
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -357,9 +359,14 @@ def test_serve_cli_forwards_distributed_offload_residency():
|
|||||||
assert args.enable_distributed_layerwise_offload is True
|
assert args.enable_distributed_layerwise_offload is True
|
||||||
assert args.dlo_use_allgather is False
|
assert args.dlo_use_allgather is False
|
||||||
assert args.dlo_resident_layers == 20
|
assert args.dlo_resident_layers == 20
|
||||||
|
assert args.dlo_offload_components == {"text_encoder": False, "default": True}
|
||||||
assert engine_args["enable_distributed_layerwise_offload"] is True
|
assert engine_args["enable_distributed_layerwise_offload"] is True
|
||||||
assert engine_args["dlo_use_allgather"] is False
|
assert engine_args["dlo_use_allgather"] is False
|
||||||
assert engine_args["dlo_resident_layers"] == 20
|
assert engine_args["dlo_resident_layers"] == 20
|
||||||
|
assert engine_args["dlo_offload_components"] == {
|
||||||
|
"text_encoder": False,
|
||||||
|
"default": True,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
|
def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
|
||||||
@@ -378,6 +385,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
|
|||||||
"preferred",
|
"preferred",
|
||||||
"--host-weight-runtime-root",
|
"--host-weight-runtime-root",
|
||||||
"/var/cache/vllm-omni/hwr",
|
"/var/cache/vllm-omni/hwr",
|
||||||
|
"--dlo-host-registration-limit-gib",
|
||||||
|
"80",
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -387,8 +396,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_mode"] == "preferred"
|
||||||
assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
|
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_mode"] == "preferred"
|
||||||
assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
|
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():
|
def test_serve_cli_accepts_diffusion_compile_controls():
|
||||||
|
|||||||
@@ -57,7 +57,9 @@ from vllm_omni.host_weight_runtime.filesystem import (
|
|||||||
detect_storage_class,
|
detect_storage_class,
|
||||||
)
|
)
|
||||||
from vllm_omni.host_weight_runtime.filesystem import store as filesystem_store_module
|
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.locks import FileLock
|
||||||
|
from vllm_omni.host_weight_runtime.filesystem.writer import ArtifactWriterError, FilesystemArtifactWriter
|
||||||
|
|
||||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
@@ -110,6 +112,150 @@ def _open_fd_targets() -> tuple[str, ...]:
|
|||||||
return tuple(targets)
|
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:
|
class FakeProducer:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
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 = tmp_path / "free-limit"
|
||||||
free_root.mkdir(mode=0o700)
|
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_limited = _make_store(free_root, capacity=CapacityPolicy(min_free_bytes=unavailable_reserve))
|
||||||
free_result = free_limited.get_or_build(
|
free_result = free_limited.get_or_build(
|
||||||
BuildRequest(identity),
|
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,
|
||||||
|
}
|
||||||
@@ -1,5 +1,9 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||||
|
|
||||||
import warnings
|
import warnings
|
||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
|
from collections.abc import Sequence
|
||||||
from dataclasses import field, make_dataclass
|
from dataclasses import field, make_dataclass
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -15,6 +19,11 @@ _FLOAT_LIST_TYPE = list[float]
|
|||||||
_INT_LIST_TYPE = list[int]
|
_INT_LIST_TYPE = list[int]
|
||||||
|
|
||||||
_MULTIMODAL_BENCHMARK_FIELDS = [
|
_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.MEAN_AUDIO_TTFP_MS, float, field(default=0.0)),
|
||||||
(defs.MEDIAN_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)),
|
(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(
|
def print_metrics(
|
||||||
task_type,
|
task_type,
|
||||||
selected_percentile_metrics,
|
selected_percentile_metrics,
|
||||||
@@ -308,6 +328,9 @@ def process_one_metric(
|
|||||||
metric_attribute_name: str,
|
metric_attribute_name: str,
|
||||||
metrics: MultiModalsBenchmarkMetrics,
|
metrics: MultiModalsBenchmarkMetrics,
|
||||||
):
|
):
|
||||||
|
if not has_metric_samples(metrics, metric_attribute_name):
|
||||||
|
return
|
||||||
|
|
||||||
metric_header_map = {
|
metric_header_map = {
|
||||||
"ttft": "Time to First Token",
|
"ttft": "Time to First Token",
|
||||||
"tpot": "Time per Output Token (excl. 1st token)",
|
"tpot": "Time per Output Token (excl. 1st token)",
|
||||||
@@ -718,10 +741,12 @@ def calculate_metrics(
|
|||||||
good_completed = 0
|
good_completed = 0
|
||||||
itls: list[float] = []
|
itls: list[float] = []
|
||||||
tpots: list[float] = []
|
tpots: list[float] = []
|
||||||
all_tpots: list[float] = []
|
all_tpots: list[float | None] = []
|
||||||
ttfts: list[float] = []
|
ttfts: list[float] = []
|
||||||
|
goodput_ttfts: list[float | None] = []
|
||||||
e2els: list[float] = []
|
e2els: list[float] = []
|
||||||
audio_ttfps: list[float] = []
|
audio_ttfps: list[float] = []
|
||||||
|
goodput_audio_ttfps: list[float | None] = []
|
||||||
audio_rtfs: list[float] = []
|
audio_rtfs: list[float] = []
|
||||||
audio_duration: list[float] = []
|
audio_duration: list[float] = []
|
||||||
audio_frames: list[int] = []
|
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)
|
output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
|
||||||
actual_output_lens.append(output_len)
|
actual_output_lens.append(output_len)
|
||||||
total_input += outputs[i].prompt_len
|
total_input += outputs[i].prompt_len
|
||||||
tpot = 0
|
tpot: float | None = 0.0
|
||||||
if output_len > 1:
|
if output_len > 1:
|
||||||
if outputs[i].itl:
|
if outputs[i].itl:
|
||||||
# Use mean(ITL) directly so per-request TPOT == mean(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.
|
# bundle multiple tokens per chunk, so len(itl)+1 != output_len.
|
||||||
# Using mean(itl) keeps TPOT and ITL on the same footing.
|
# Using mean(itl) keeps TPOT and ITL on the same footing.
|
||||||
tpot = sum(outputs[i].itl) / len(outputs[i].itl)
|
tpot = sum(outputs[i].itl) / len(outputs[i].itl)
|
||||||
else:
|
elif getattr(outputs[i], "tpot_measured", True):
|
||||||
try:
|
try:
|
||||||
latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft
|
latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft
|
||||||
except Exception:
|
except Exception:
|
||||||
latency_minus_ttft = outputs[i].latency - outputs[i].ttft
|
latency_minus_ttft = outputs[i].latency - outputs[i].ttft
|
||||||
tpot = latency_minus_ttft / (output_len - 1)
|
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
|
# Note: if output_len <= 1, we regard tpot as 0 for goodput
|
||||||
all_tpots.append(tpot)
|
all_tpots.append(tpot)
|
||||||
itls += outputs[i].itl
|
itls += outputs[i].itl
|
||||||
ttfts.append(outputs[i].ttft)
|
session_metrics = getattr(outputs[i], "duplex_session_metrics", None)
|
||||||
audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
|
ttft_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_ttft_ms") is not None
|
||||||
audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
|
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_duration.append(getattr(outputs[i], defs.AUDIO_DURATION, 0.0))
|
||||||
audio_frames.append(getattr(outputs[i], defs.AUDIO_FRAMES, 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)
|
image_count = int(getattr(outputs[i], defs.IMAGE_COUNT, 0) or 0)
|
||||||
@@ -789,14 +828,14 @@ def calculate_metrics(
|
|||||||
actual_output_lens.append(0)
|
actual_output_lens.append(0)
|
||||||
|
|
||||||
if goodput_config_dict:
|
if goodput_config_dict:
|
||||||
valid_metrics = []
|
valid_metrics: list[Sequence[float | None]] = []
|
||||||
slo_values = []
|
slo_values = []
|
||||||
|
|
||||||
if "ttft" in goodput_config_dict:
|
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)
|
slo_values.append(goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
|
||||||
if "audio_ttft" in goodput_config_dict:
|
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)
|
slo_values.append(goodput_config_dict["audio_ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
|
||||||
if "tpot" in goodput_config_dict:
|
if "tpot" in goodput_config_dict:
|
||||||
valid_metrics.append(all_tpots)
|
valid_metrics.append(all_tpots)
|
||||||
@@ -806,7 +845,7 @@ def calculate_metrics(
|
|||||||
slo_values.append(goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION)
|
slo_values.append(goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION)
|
||||||
|
|
||||||
for req_metric in zip(*valid_metrics):
|
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:
|
if is_good_req:
|
||||||
good_completed += 1
|
good_completed += 1
|
||||||
|
|
||||||
@@ -822,6 +861,7 @@ def calculate_metrics(
|
|||||||
# Calculate max output tokens per second metric
|
# Calculate max output tokens per second metric
|
||||||
max_output_tokens_per_s = 0.0
|
max_output_tokens_per_s = 0.0
|
||||||
max_concurrent_requests = 0
|
max_concurrent_requests = 0
|
||||||
|
token_timeline_available = True
|
||||||
|
|
||||||
# Find the time range across all successful requests
|
# Find the time range across all successful requests
|
||||||
successful_outputs = [output for output in outputs if output.success]
|
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)
|
concurrent_requests_per_second = np.zeros(duration_seconds)
|
||||||
|
|
||||||
for i, output in enumerate(successful_outputs):
|
for i, output in enumerate(successful_outputs):
|
||||||
# Calculate token generation timestamp using
|
session_metrics = getattr(output, "duplex_session_metrics", None)
|
||||||
# start_time, ttft, and itl
|
if isinstance(session_metrics, dict):
|
||||||
token_times = [output.start_time + output.ttft]
|
# Duplex TTFT/ITL samples are response-local. Without each
|
||||||
current_time = token_times[0]
|
# response.created offset they cannot form a session-global
|
||||||
for itl_value in output.itl:
|
# token timeline, so do not publish a misleading peak rate.
|
||||||
current_time += itl_value
|
token_timeline_available = False
|
||||||
token_times.append(current_time)
|
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
|
# Add tokens to second buckets
|
||||||
for token_time in token_times:
|
for token_time in token_times:
|
||||||
second_bucket = int(token_time - min_start_time)
|
second_bucket = int(token_time - min_start_time)
|
||||||
if 0 <= second_bucket < duration_seconds:
|
if 0 <= second_bucket < duration_seconds:
|
||||||
tokens_per_second[second_bucket] += 1
|
tokens_per_second[second_bucket] += 1
|
||||||
|
|
||||||
# Track concurrent requests for each second this request was active
|
# Track concurrent requests for each second this request was active
|
||||||
request_start_second = int(output.start_time - min_start_time)
|
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
|
# Find the maximum tokens per second and corresponding
|
||||||
# concurrent requests
|
# concurrent requests
|
||||||
if len(tokens_per_second) > 0:
|
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))
|
max_concurrent_requests = int(np.max(concurrent_requests_per_second))
|
||||||
|
|
||||||
if TERM_PLOTLIB_AVAILABLE:
|
if TERM_PLOTLIB_AVAILABLE:
|
||||||
import termplotlib as tpl
|
import termplotlib as tpl
|
||||||
|
|
||||||
fig = tpl.figure()
|
fig = tpl.figure()
|
||||||
fig.plot(
|
if token_timeline_available:
|
||||||
np.arange(len(tokens_per_second)),
|
fig.plot(
|
||||||
tokens_per_second,
|
np.arange(len(tokens_per_second)),
|
||||||
title="Output tokens per second",
|
tokens_per_second,
|
||||||
)
|
title="Output tokens per second",
|
||||||
|
)
|
||||||
fig.plot(
|
fig.plot(
|
||||||
np.arange(len(concurrent_requests_per_second)),
|
np.arange(len(concurrent_requests_per_second)),
|
||||||
concurrent_requests_per_second,
|
concurrent_requests_per_second,
|
||||||
@@ -881,6 +932,10 @@ def calculate_metrics(
|
|||||||
else:
|
else:
|
||||||
print("tip: install termplotlib and gnuplot to plot the metrics")
|
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(
|
metrics = MultiModalsBenchmarkMetrics(
|
||||||
completed=completed,
|
completed=completed,
|
||||||
failed=len(failed_outputs),
|
failed=len(failed_outputs),
|
||||||
@@ -890,18 +945,18 @@ def calculate_metrics(
|
|||||||
request_goodput=good_completed / dur_s,
|
request_goodput=good_completed / dur_s,
|
||||||
output_throughput=sum(actual_output_lens) / dur_s,
|
output_throughput=sum(actual_output_lens) / dur_s,
|
||||||
total_token_throughput=(total_input + 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
|
mean_ttft_ms=np.mean(ttfts or missing_duplex_value) * 1000,
|
||||||
std_ttft_ms=np.std(ttfts or 0) * 1000,
|
std_ttft_ms=np.std(ttfts or missing_duplex_value) * 1000,
|
||||||
median_ttft_ms=np.median(ttfts or 0) * 1000,
|
median_ttft_ms=np.median(ttfts or missing_duplex_value) * 1000,
|
||||||
percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles],
|
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 0) * 1000,
|
mean_tpot_ms=np.mean(tpots or missing_duplex_value) * 1000,
|
||||||
std_tpot_ms=np.std(tpots or 0) * 1000,
|
std_tpot_ms=np.std(tpots or missing_duplex_value) * 1000,
|
||||||
median_tpot_ms=np.median(tpots or 0) * 1000,
|
median_tpot_ms=np.median(tpots or missing_duplex_value) * 1000,
|
||||||
percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles],
|
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 0) * 1000,
|
mean_itl_ms=np.mean(itls or missing_duplex_value) * 1000,
|
||||||
std_itl_ms=np.std(itls or 0) * 1000,
|
std_itl_ms=np.std(itls or missing_duplex_value) * 1000,
|
||||||
median_itl_ms=np.median(itls or 0) * 1000,
|
median_itl_ms=np.median(itls or missing_duplex_value) * 1000,
|
||||||
percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles],
|
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,
|
mean_e2el_ms=np.mean(e2els or 0) * 1000,
|
||||||
std_e2el_ms=np.std(e2els or 0) * 1000,
|
std_e2el_ms=np.std(e2els or 0) * 1000,
|
||||||
median_e2el_ms=np.median(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,
|
max_concurrent_requests=max_concurrent_requests,
|
||||||
rtfx=input_audio_duration / dur_s,
|
rtfx=input_audio_duration / dur_s,
|
||||||
**{
|
**{
|
||||||
defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or 0) * 1000,
|
"num_ttft_samples": len(ttfts),
|
||||||
defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or 0) * 1000,
|
"num_tpot_samples": len(tpots),
|
||||||
defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or 0) * 1000,
|
"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: [
|
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.MEAN_AUDIO_DURATION_S: np.mean(audio_duration or 0),
|
||||||
defs.STD_AUDIO_DURATION_S: np.std(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_DURATION_S: sum(audio_duration),
|
||||||
defs.TOTAL_AUDIO_FRAMES: sum(audio_frames),
|
defs.TOTAL_AUDIO_FRAMES: sum(audio_frames),
|
||||||
defs.AUDIO_THROUGHPUT: sum(audio_duration) / dur_s,
|
defs.AUDIO_THROUGHPUT: sum(audio_duration) / dur_s,
|
||||||
defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or 0),
|
defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or missing_duplex_value),
|
||||||
defs.STD_AUDIO_RTF: np.std(audio_rtfs or 0),
|
defs.STD_AUDIO_RTF: np.std(audio_rtfs or missing_duplex_value),
|
||||||
defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or 0),
|
defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or missing_duplex_value),
|
||||||
defs.PERCENTILES_AUDIO_RTF: [(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles],
|
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.TOTAL_IMAGES: total_images,
|
||||||
defs.IMAGE_THROUGHPUT: total_images / dur_s,
|
defs.IMAGE_THROUGHPUT: total_images / dur_s,
|
||||||
defs.AVERAGE_PIXELS_PER_IMAGE: (total_image_pixels / total_images) if total_images > 0 else 0.0,
|
defs.AVERAGE_PIXELS_PER_IMAGE: (total_image_pixels / total_images) if total_images > 0 else 0.0,
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -1,3 +1,6 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import contextlib
|
import contextlib
|
||||||
import io
|
import io
|
||||||
@@ -16,7 +19,6 @@ from dataclasses import dataclass, replace
|
|||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Literal
|
from typing import Any, Literal
|
||||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
|
||||||
|
|
||||||
import aiohttp
|
import aiohttp
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -31,6 +33,7 @@ from vllm.benchmarks.lib.endpoint_request_func import (
|
|||||||
RequestFuncOutput,
|
RequestFuncOutput,
|
||||||
StreamedResponseHandler,
|
StreamedResponseHandler,
|
||||||
_get_chat_content,
|
_get_chat_content,
|
||||||
|
_get_headers,
|
||||||
_update_headers_common,
|
_update_headers_common,
|
||||||
_update_payload_common,
|
_update_payload_common,
|
||||||
_validate_api_url,
|
_validate_api_url,
|
||||||
@@ -46,6 +49,13 @@ from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
|
|||||||
daily_omni_local_videos_dir,
|
daily_omni_local_videos_dir,
|
||||||
resolve_daily_omni_local_root,
|
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.random_multi_modal_dataset import OmniRandomMultiModalDataset
|
||||||
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
|
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
|
||||||
SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT,
|
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.sound_effect_dataset import SoundEffectDataset
|
||||||
from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
|
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 (
|
from vllm_omni.experimental.fullduplex.client import (
|
||||||
RealtimeDuplexClient,
|
RealtimeDuplexClient,
|
||||||
|
build_realtime_url,
|
||||||
|
reference_audio_data_url,
|
||||||
summarize_session_request_metrics,
|
summarize_session_request_metrics,
|
||||||
wait_for,
|
wait_for,
|
||||||
)
|
)
|
||||||
@@ -221,6 +250,91 @@ def _attach_seed_tts_to_request_func_input(sample: SampleRequest, rfi: RequestFu
|
|||||||
rfi.extra_body = base
|
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:
|
def _daily_omni_repo_from_args(args) -> str | None:
|
||||||
"""Resolve HuggingFace repo id for Daily-Omni from CLI args.
|
"""Resolve HuggingFace repo id for Daily-Omni from CLI args.
|
||||||
|
|
||||||
@@ -250,6 +364,7 @@ def get_samples(args, tokenizer):
|
|||||||
"ttsd",
|
"ttsd",
|
||||||
"sound-effect",
|
"sound-effect",
|
||||||
)
|
)
|
||||||
|
is_omniinteract = args.dataset_name == "omniinteract"
|
||||||
|
|
||||||
# Check if we need to handle omni-related backends/datasets
|
# Check if we need to handle omni-related backends/datasets
|
||||||
is_omni_backend = args.backend in [
|
is_omni_backend = args.backend in [
|
||||||
@@ -258,12 +373,68 @@ def get_samples(args, tokenizer):
|
|||||||
"openai-realtime-duplex",
|
"openai-realtime-duplex",
|
||||||
"daily-omni",
|
"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:
|
if not is_omni_backend and not is_omni_dataset:
|
||||||
# Not an omni-related request, delegate to original implementation
|
# Not an omni-related request, delegate to original implementation
|
||||||
return get_samples_old(args, tokenizer)
|
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
|
# Handle Daily-Omni dataset
|
||||||
if is_daily_omni:
|
if is_daily_omni:
|
||||||
# Support:
|
# Support:
|
||||||
@@ -476,6 +647,7 @@ class MixRequestFuncOutput(RequestFuncOutput):
|
|||||||
image_pixels: int = 0
|
image_pixels: int = 0
|
||||||
denoise_step_latency_ms: float = 0.0
|
denoise_step_latency_ms: float = 0.0
|
||||||
text_latency: float = 0.0
|
text_latency: float = 0.0
|
||||||
|
tpot_measured: bool = True
|
||||||
#: Worst-case streaming-audio underrun (wall-clock seconds the player
|
#: Worst-case streaming-audio underrun (wall-clock seconds the player
|
||||||
#: would have been starved). Populated by the audio-speech backend; ``0.0``
|
#: would have been starved). Populated by the audio-speech backend; ``0.0``
|
||||||
#: for backends that do not run continuity analysis.
|
#: 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:
|
def _realtime_websocket_url(api_url: str) -> str:
|
||||||
parts = urlsplit(api_url)
|
return build_realtime_url(api_url, None, native_duplex=None)
|
||||||
scheme = "wss" if parts.scheme == "https" else "ws"
|
|
||||||
query = dict(parse_qsl(parts.query, keep_blank_values=True))
|
|
||||||
query.setdefault("duplex", "1")
|
def _nonnegative_number(value: object) -> bool:
|
||||||
return urlunsplit((scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
|
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(
|
async def async_request_openai_realtime_duplex(
|
||||||
@@ -1325,8 +1601,9 @@ async def async_request_openai_realtime_duplex(
|
|||||||
session: aiohttp.ClientSession,
|
session: aiohttp.ClientSession,
|
||||||
pbar: tqdm | None = None,
|
pbar: tqdm | None = None,
|
||||||
) -> MixRequestFuncOutput:
|
) -> MixRequestFuncOutput:
|
||||||
"""Run one or more Seed-TTS turns through one explicit Realtime TTS session."""
|
|
||||||
del 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 = MixRequestFuncOutput()
|
||||||
output.prompt_len = request_func_input.prompt_len
|
output.prompt_len = request_func_input.prompt_len
|
||||||
output.start_time = time.perf_counter()
|
output.start_time = time.perf_counter()
|
||||||
@@ -1387,8 +1664,10 @@ async def async_request_openai_realtime_duplex(
|
|||||||
)
|
)
|
||||||
await client.send({"type": "response.create"})
|
await client.send({"type": "response.create"})
|
||||||
await wait_for(
|
await wait_for(
|
||||||
lambda: client.events.count("response.done") > done_before
|
lambda: (
|
||||||
or len(client.events.errors()) > errors_before,
|
client.events.count("response.done") > done_before
|
||||||
|
or len(client.events.errors()) > errors_before
|
||||||
|
),
|
||||||
timeout_s=180.0,
|
timeout_s=180.0,
|
||||||
label=f"Seed-TTS Realtime TTS turn {request_index} response.done",
|
label=f"Seed-TTS Realtime TTS turn {request_index} response.done",
|
||||||
)
|
)
|
||||||
@@ -1433,19 +1712,7 @@ async def async_request_openai_realtime_duplex(
|
|||||||
channels=1,
|
channels=1,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
turn_transcripts.append(
|
turn_transcripts.append(client.events.response_text(response_id))
|
||||||
"".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",
|
|
||||||
}
|
|
||||||
)
|
|
||||||
)
|
|
||||||
await client.acknowledge_playback()
|
await client.acknowledge_playback()
|
||||||
request_finished_at = time.perf_counter()
|
request_finished_at = time.perf_counter()
|
||||||
session_metrics = summarize_session_request_metrics(
|
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 (
|
from vllm_omni.benchmarks.metrics.metrics import (
|
||||||
MultiModalsBenchmarkMetrics,
|
MultiModalsBenchmarkMetrics,
|
||||||
calculate_metrics,
|
calculate_metrics,
|
||||||
|
has_metric_samples,
|
||||||
)
|
)
|
||||||
|
|
||||||
# ruff: noqa: E402
|
# ruff: noqa: E402
|
||||||
@@ -1626,6 +1894,7 @@ async def benchmark(
|
|||||||
)
|
)
|
||||||
_attach_daily_omni_to_request_func_input(input_requests[0], test_input)
|
_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_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:
|
if ready_check_timeout_sec > 0:
|
||||||
test_output = await wait_for_endpoint(
|
test_output = await wait_for_endpoint(
|
||||||
@@ -1741,6 +2010,7 @@ async def benchmark(
|
|||||||
print(f"Probe request rate: {probe_request_rate} req/s")
|
print(f"Probe request rate: {probe_request_rate} req/s")
|
||||||
probe_task = asyncio.create_task(probe_loop())
|
probe_task = asyncio.create_task(probe_loop())
|
||||||
|
|
||||||
|
_prepare_omniinteract_batch(input_requests)
|
||||||
benchmark_start_time = time.perf_counter()
|
benchmark_start_time = time.perf_counter()
|
||||||
tasks: list[asyncio.Task] = []
|
tasks: list[asyncio.Task] = []
|
||||||
|
|
||||||
@@ -1801,6 +2071,7 @@ async def benchmark(
|
|||||||
)
|
)
|
||||||
_attach_daily_omni_to_request_func_input(request, request_func_input)
|
_attach_daily_omni_to_request_func_input(request, request_func_input)
|
||||||
_attach_seed_tts_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(
|
tasks.append(
|
||||||
asyncio.create_task(limited_request_func(request_func_input=request_func_input, session=session, pbar=pbar))
|
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
|
benchmark_duration = time.perf_counter() - benchmark_start_time
|
||||||
|
|
||||||
|
omniinteract_summary = _finalize_omniinteract_batch(input_requests, outputs)
|
||||||
|
|
||||||
if task_type == TaskType.GENERATION:
|
if task_type == TaskType.GENERATION:
|
||||||
metrics, actual_output_lens = calculate_metrics(
|
metrics, actual_output_lens = calculate_metrics(
|
||||||
input_requests=input_requests,
|
input_requests=input_requests,
|
||||||
@@ -1839,6 +2112,13 @@ async def benchmark(
|
|||||||
actual_output_lens = 0
|
actual_output_lens = 0
|
||||||
|
|
||||||
if isinstance(metrics, MultiModalsBenchmarkMetrics):
|
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 = {
|
result = {
|
||||||
"duration": benchmark_duration,
|
"duration": benchmark_duration,
|
||||||
"completed": metrics.completed,
|
"completed": metrics.completed,
|
||||||
@@ -1859,7 +2139,7 @@ async def benchmark(
|
|||||||
"input_lens": [output.prompt_len for output in outputs],
|
"input_lens": [output.prompt_len for output in outputs],
|
||||||
"start_times": [output.start_time for output in outputs],
|
"start_times": [output.start_time for output in outputs],
|
||||||
"output_lens": actual_output_lens,
|
"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],
|
"itls": [output.itl for output in outputs],
|
||||||
"generated_texts": [output.generated_text for output in outputs],
|
"generated_texts": [output.generated_text for output in outputs],
|
||||||
"errors": [output.error for output in outputs],
|
"errors": [output.error for output in outputs],
|
||||||
@@ -1893,6 +2173,8 @@ async def benchmark(
|
|||||||
]
|
]
|
||||||
if duplex_session_metrics:
|
if duplex_session_metrics:
|
||||||
result["duplex_session_metrics"] = 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 (
|
from vllm_omni.benchmarks.data_modules.daily_omni_eval import (
|
||||||
compute_daily_omni_accuracy_metrics,
|
compute_daily_omni_accuracy_metrics,
|
||||||
@@ -1949,6 +2231,8 @@ async def benchmark(
|
|||||||
# metric.
|
# metric.
|
||||||
if metric_attribute_name not in result_percentile_metrics:
|
if metric_attribute_name not in result_percentile_metrics:
|
||||||
return
|
return
|
||||||
|
if not has_metric_samples(metrics, metric_attribute_name):
|
||||||
|
return
|
||||||
# No text tokens generated (e.g. pure TTS speech endpoint): per-token
|
# No text tokens generated (e.g. pure TTS speech endpoint): per-token
|
||||||
# latency metrics (ttft/tpot/itl) are undefined, so skip them.
|
# 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"))
|
is_text_token_metric = not (metric_attribute_name == "e2el" or metric_attribute_name.startswith("audio"))
|
||||||
|
|||||||
@@ -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 argparse
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import contextlib
|
||||||
import os
|
import os
|
||||||
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from vllm.benchmarks.serve import main_async
|
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
|
# This monkey-patches vllm.benchmarks.datasets.get_samples before it's used
|
||||||
# Must be imported before any vllm.benchmarks module usage
|
# Must be imported before any vllm.benchmarks module usage
|
||||||
import vllm_omni.benchmarks.patch.patch # noqa: F401
|
import vllm_omni.benchmarks.patch.patch # noqa: F401
|
||||||
|
from vllm_omni.benchmarks.omniinteract import omniinteract_output_lock
|
||||||
from vllm_omni.benchmarks.patch.patch import (
|
from vllm_omni.benchmarks.patch.patch import (
|
||||||
maybe_enable_stage_metrics,
|
maybe_enable_stage_metrics,
|
||||||
set_print_stage,
|
set_print_stage,
|
||||||
@@ -28,4 +36,10 @@ def main(args: argparse.Namespace) -> dict[str, Any]:
|
|||||||
getattr(args, "extra_body", None),
|
getattr(args, "extra_body", None),
|
||||||
enabled=should_request_stage_metrics(args),
|
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_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",
|
||||||
|
"ROSVOT_SOURCE_DIR",
|
||||||
|
"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()
|
||||||
@@ -697,8 +697,10 @@ class _DiffusionConfigProjection:
|
|||||||
enable_distributed_layerwise_offload: bool = False
|
enable_distributed_layerwise_offload: bool = False
|
||||||
dlo_use_allgather: bool = True
|
dlo_use_allgather: bool = True
|
||||||
dlo_resident_layers: int = Field(default=0, ge=0)
|
dlo_resident_layers: int = Field(default=0, ge=0)
|
||||||
|
dlo_offload_components: dict[str, bool] = Field(default_factory=dict)
|
||||||
host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
|
host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
|
||||||
host_weight_runtime_root: str | None = None
|
host_weight_runtime_root: str | None = None
|
||||||
|
dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)
|
||||||
pin_cpu_memory: bool = True
|
pin_cpu_memory: bool = True
|
||||||
diffusion_compile_granularity: Literal["regional", "full"] = "regional"
|
diffusion_compile_granularity: Literal["regional", "full"] = "regional"
|
||||||
diffusion_compile_dynamic: bool = Field(default=True, strict=True)
|
diffusion_compile_dynamic: bool = Field(default=True, strict=True)
|
||||||
@@ -756,6 +758,7 @@ class _DiffusionConfigProjection:
|
|||||||
TransformerConfig,
|
TransformerConfig,
|
||||||
build_attention_config,
|
build_attention_config,
|
||||||
parse_kv_cache_skip_selector,
|
parse_kv_cache_skip_selector,
|
||||||
|
validate_dlo_host_registration_options,
|
||||||
validate_host_weight_runtime_options,
|
validate_host_weight_runtime_options,
|
||||||
)
|
)
|
||||||
from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
|
from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
|
||||||
@@ -838,6 +841,12 @@ class _DiffusionConfigProjection:
|
|||||||
mode=self.host_weight_runtime_mode,
|
mode=self.host_weight_runtime_mode,
|
||||||
root=self.host_weight_runtime_root,
|
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):
|
if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
|
|||||||
@@ -389,8 +389,10 @@ class StageDeployConfig:
|
|||||||
enable_distributed_layerwise_offload: bool | None = None
|
enable_distributed_layerwise_offload: bool | None = None
|
||||||
dlo_use_allgather: bool | None = None
|
dlo_use_allgather: bool | None = None
|
||||||
dlo_resident_layers: int | None = None
|
dlo_resident_layers: int | None = None
|
||||||
|
dlo_offload_components: dict[str, bool] | None = None
|
||||||
host_weight_runtime_mode: str | None = None
|
host_weight_runtime_mode: str | None = None
|
||||||
host_weight_runtime_root: str | None = None
|
host_weight_runtime_root: str | None = None
|
||||||
|
dlo_host_registration_limit_gib: float | None = None
|
||||||
# Diffusion-specific debug and observability knobs.
|
# Diffusion-specific debug and observability knobs.
|
||||||
enable_diffusion_pipeline_profiler: bool | None = None
|
enable_diffusion_pipeline_profiler: bool | None = None
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
# Wan2.2 TI2V-5B deploy.
|
# Wan2.2 TI2V-5B deploy.
|
||||||
#
|
#
|
||||||
# Single GPU by default. Increase tensor_parallel_size and devices for A14B MoE variants.
|
# 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
|
async_chunk: false
|
||||||
trust_remote_code: true
|
trust_remote_code: true
|
||||||
distributed_executor_backend: mp
|
distributed_executor_backend: mp
|
||||||
@@ -8,7 +9,7 @@ distributed_executor_backend: mp
|
|||||||
stages:
|
stages:
|
||||||
- stage_id: 0
|
- stage_id: 0
|
||||||
devices: "0"
|
devices: "0"
|
||||||
max_num_seqs: 1
|
max_num_seqs: 4 # in-flight request batch width (CLI: --max-num-seqs)
|
||||||
tensor_parallel_size: 1
|
tensor_parallel_size: 1
|
||||||
enforce_eager: true
|
enforce_eager: true
|
||||||
# Keep VAE tiling on by default for higher resolutions / longer clips.
|
# Keep VAE tiling on by default for higher resolutions / longer clips.
|
||||||
|
|||||||
@@ -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")
|
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(
|
def parse_kv_cache_skip_selector(
|
||||||
selector: str | list[int] | tuple[int, ...] | set[int] | None,
|
selector: str | list[int] | tuple[int, ...] | set[int] | None,
|
||||||
) -> set[int] | None:
|
) -> set[int] | None:
|
||||||
@@ -779,11 +795,17 @@ class OmniDiffusionConfig:
|
|||||||
dlo_use_allgather: bool = True
|
dlo_use_allgather: bool = True
|
||||||
# Leading main-DiT blocks kept resident by distributed layerwise offload.
|
# Leading main-DiT blocks kept resident by distributed layerwise offload.
|
||||||
dlo_resident_layers: int = 0
|
dlo_resident_layers: int = 0
|
||||||
|
# Auxiliary encoder/VAE policy. False keeps the whole component resident;
|
||||||
|
# DiT offload remains owned by DLO.
|
||||||
|
dlo_offload_components: dict[str, bool] = field(default_factory=dict)
|
||||||
# Final-layout Host Weight Runtime policy. The loader only activates this
|
# Final-layout Host Weight Runtime policy. The loader only activates this
|
||||||
# for eligible no-AllGather DLO; all other configurations preserve their
|
# for eligible no-AllGather DLO; all other configurations preserve their
|
||||||
# existing loader/storage path.
|
# existing loader/storage path.
|
||||||
host_weight_runtime_mode: str = "disabled"
|
host_weight_runtime_mode: str = "disabled"
|
||||||
host_weight_runtime_root: str | None = None
|
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
|
pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading
|
||||||
|
|
||||||
@@ -1164,6 +1186,12 @@ class OmniDiffusionConfig:
|
|||||||
mode=self.host_weight_runtime_mode,
|
mode=self.host_weight_runtime_mode,
|
||||||
root=self.host_weight_runtime_root,
|
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):
|
if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
|
|||||||
@@ -27,7 +27,10 @@ logger = init_logger(__name__)
|
|||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext
|
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):
|
class _HWRCommitError(RuntimeError):
|
||||||
@@ -215,6 +218,7 @@ class HWRLoaderMixin:
|
|||||||
*,
|
*,
|
||||||
load_format: str,
|
load_format: str,
|
||||||
sources: Sequence[object],
|
sources: Sequence[object],
|
||||||
|
source_digest_cache: NodeSourceDigestCache | None = None,
|
||||||
) -> FinalLayoutIdentityContext:
|
) -> FinalLayoutIdentityContext:
|
||||||
from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
|
from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
|
||||||
|
|
||||||
@@ -293,6 +297,7 @@ class HWRLoaderMixin:
|
|||||||
prepared_sources=prepared_sources,
|
prepared_sources=prepared_sources,
|
||||||
request=request,
|
request=request,
|
||||||
policy=FINAL_LAYOUT_BF16_POLICY,
|
policy=FINAL_LAYOUT_BF16_POLICY,
|
||||||
|
source_digest_cache=source_digest_cache,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _resolve_hwr(
|
def _resolve_hwr(
|
||||||
@@ -306,7 +311,10 @@ class HWRLoaderMixin:
|
|||||||
sources: Sequence[object],
|
sources: Sequence[object],
|
||||||
) -> dict[str, object] | None:
|
) -> dict[str, object] | None:
|
||||||
"""Resolve an eligible no-AllGather final-layout HWR transaction."""
|
"""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 (
|
from vllm_omni.host_weight_runtime import (
|
||||||
HostWeightLeaseCarrier,
|
HostWeightLeaseCarrier,
|
||||||
HostWeightRuntime,
|
HostWeightRuntime,
|
||||||
@@ -347,7 +355,6 @@ class HWRLoaderMixin:
|
|||||||
raise ValueError(f"required Host Weight Runtime path is ineligible: {message}")
|
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)
|
logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message)
|
||||||
return None
|
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)
|
root_value = getattr(self.od_config, "host_weight_runtime_root", None)
|
||||||
if not isinstance(root_value, str) or not root_value.strip():
|
if not isinstance(root_value, str) or not root_value.strip():
|
||||||
raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
|
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)
|
resolution = runtime.resolve(context.identity)
|
||||||
state: dict[str, object] = {
|
state: dict[str, object] = {
|
||||||
"mode": mode,
|
"mode": mode,
|
||||||
|
|||||||
@@ -19,7 +19,7 @@ from .producers import (
|
|||||||
FinalLayoutBF16Producer,
|
FinalLayoutBF16Producer,
|
||||||
)
|
)
|
||||||
from .restorer import FinalLayoutTensorRestorer
|
from .restorer import FinalLayoutTensorRestorer
|
||||||
from .source_identity import PreparedWeightSource, WeightSourceKind
|
from .source_identity import NodeSourceDigestCache, PreparedWeightSource, WeightSourceKind
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"FINAL_LAYOUT_BF16_POLICY",
|
"FINAL_LAYOUT_BF16_POLICY",
|
||||||
@@ -32,6 +32,7 @@ __all__ = [
|
|||||||
"FinalLayoutRequest",
|
"FinalLayoutRequest",
|
||||||
"FinalLayoutTensorRestorer",
|
"FinalLayoutTensorRestorer",
|
||||||
"ImplementationIdentity",
|
"ImplementationIdentity",
|
||||||
|
"NodeSourceDigestCache",
|
||||||
"PreparedWeightSource",
|
"PreparedWeightSource",
|
||||||
"WeightSourceKind",
|
"WeightSourceKind",
|
||||||
"build_final_layout_identity",
|
"build_final_layout_identity",
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ from .contracts import (
|
|||||||
)
|
)
|
||||||
from .source_identity import (
|
from .source_identity import (
|
||||||
FinalLayoutSourceContext,
|
FinalLayoutSourceContext,
|
||||||
|
NodeSourceDigestCache,
|
||||||
PreparedWeightSource,
|
PreparedWeightSource,
|
||||||
resolve_final_layout_source_identity,
|
resolve_final_layout_source_identity,
|
||||||
)
|
)
|
||||||
@@ -78,6 +79,7 @@ def build_final_layout_identity(
|
|||||||
prepared_sources: Sequence[PreparedWeightSource],
|
prepared_sources: Sequence[PreparedWeightSource],
|
||||||
request: FinalLayoutRequest,
|
request: FinalLayoutRequest,
|
||||||
policy: FinalLayoutTensorPolicy,
|
policy: FinalLayoutTensorPolicy,
|
||||||
|
source_digest_cache: NodeSourceDigestCache | None = None,
|
||||||
) -> FinalLayoutIdentityContext:
|
) -> FinalLayoutIdentityContext:
|
||||||
"""Build exact source, representation, layout, and implementation identity."""
|
"""Build exact source, representation, layout, and implementation identity."""
|
||||||
if not isinstance(request, FinalLayoutRequest):
|
if not isinstance(request, FinalLayoutRequest):
|
||||||
@@ -99,6 +101,7 @@ def build_final_layout_identity(
|
|||||||
prepared_sources,
|
prepared_sources,
|
||||||
model_id=request.model_id,
|
model_id=request.model_id,
|
||||||
target_names=target_names,
|
target_names=target_names,
|
||||||
|
digest_cache=source_digest_cache,
|
||||||
)
|
)
|
||||||
dit_names = tuple(name for name, _ in dit_modules)
|
dit_names = tuple(name for name, _ in dit_modules)
|
||||||
model_contracts = tuple(binding.contract for binding in bindings)
|
model_contracts = tuple(binding.contract for binding in bindings)
|
||||||
|
|||||||
@@ -286,7 +286,7 @@ class FinalLayoutBF16Producer:
|
|||||||
for record in shard
|
for record in shard
|
||||||
)
|
)
|
||||||
file_name = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
|
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:
|
for record in shard:
|
||||||
output.write_tensor(record.name, record.tensor.detach())
|
output.write_tensor(record.name, record.tensor.detach())
|
||||||
|
|
||||||
|
|||||||
@@ -5,9 +5,13 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import hashlib
|
import hashlib
|
||||||
|
import json
|
||||||
import os
|
import os
|
||||||
import stat
|
import stat
|
||||||
|
import time
|
||||||
|
import uuid
|
||||||
from collections.abc import Sequence
|
from collections.abc import Sequence
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -15,6 +19,7 @@ from pathlib import Path
|
|||||||
import regex as re
|
import regex as re
|
||||||
|
|
||||||
from vllm_omni.host_weight_runtime import CanonicalJson, WeightSourceIdentity
|
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 vllm_omni.host_weight_runtime.identity import canonical_json
|
||||||
|
|
||||||
from .contracts import (
|
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)
|
@dataclass(frozen=True)
|
||||||
class _SourceSnapshot:
|
class _SourceSnapshot:
|
||||||
semantic: CanonicalJson
|
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(
|
def _hf_blob_content_id(
|
||||||
path: Path,
|
path: Path,
|
||||||
*,
|
*,
|
||||||
@@ -200,7 +331,10 @@ def _logical_model_id(value: str) -> str:
|
|||||||
return value
|
return value
|
||||||
|
|
||||||
|
|
||||||
def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
|
def _snapshot_source(
|
||||||
|
source: PreparedWeightSource,
|
||||||
|
digest_cache: NodeSourceDigestCache | None,
|
||||||
|
) -> _SourceSnapshot:
|
||||||
root = source.resolved_root.expanduser().resolve()
|
root = source.resolved_root.expanduser().resolve()
|
||||||
if not root.is_dir():
|
if not root.is_dir():
|
||||||
raise ValueError(f"canonical source root is not a directory: {root}")
|
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:
|
if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None:
|
||||||
immutable_revision = requested.lower()
|
immutable_revision = requested.lower()
|
||||||
|
|
||||||
files: list[_FileSnapshot] = []
|
def snapshot_file(path: Path) -> _FileSnapshot:
|
||||||
for path in sorted(source.weight_files, key=lambda item: str(item)):
|
|
||||||
candidate = path.expanduser()
|
candidate = path.expanduser()
|
||||||
if not candidate.is_absolute():
|
if not candidate.is_absolute():
|
||||||
candidate = root / candidate
|
candidate = root / candidate
|
||||||
@@ -220,34 +353,34 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
|
|||||||
relative_name = candidate.relative_to(root).as_posix()
|
relative_name = candidate.relative_to(root).as_posix()
|
||||||
except ValueError as exc:
|
except ValueError as exc:
|
||||||
raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc
|
raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc
|
||||||
try:
|
current = _observe(candidate)
|
||||||
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
|
|
||||||
content_id = _hf_blob_content_id(
|
content_id = _hf_blob_content_id(
|
||||||
candidate,
|
candidate,
|
||||||
source=source,
|
source=source,
|
||||||
source_root=root,
|
source_root=root,
|
||||||
)
|
)
|
||||||
if content_id is None:
|
if content_id is None:
|
||||||
content_id = f"sha256:{_sha256_file(candidate)}"
|
content_id = (
|
||||||
files.append(
|
digest_cache.content_id(candidate, current)
|
||||||
_FileSnapshot(
|
if digest_cache is not None
|
||||||
path=candidate,
|
else f"sha256:{_sha256_file(candidate, current)}"
|
||||||
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,
|
|
||||||
)
|
)
|
||||||
|
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_semantics = [file.semantic_dict() for file in files]
|
||||||
file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest()
|
file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest()
|
||||||
revision = immutable_revision or f"content-{file_fingerprint}"
|
revision = immutable_revision or f"content-{file_fingerprint}"
|
||||||
@@ -306,11 +439,12 @@ def resolve_final_layout_source_identity(
|
|||||||
*,
|
*,
|
||||||
model_id: str,
|
model_id: str,
|
||||||
target_names: frozenset[str],
|
target_names: frozenset[str],
|
||||||
|
digest_cache: NodeSourceDigestCache | None = None,
|
||||||
) -> FinalLayoutSourceContext:
|
) -> FinalLayoutSourceContext:
|
||||||
"""Resolve exact source identity for files covering the owned tensors."""
|
"""Resolve exact source identity for files covering the owned tensors."""
|
||||||
bindings = _resolve_target_sources(prepared_sources, target_names)
|
bindings = _resolve_target_sources(prepared_sources, target_names)
|
||||||
selected_sources = frozenset(bindings.values())
|
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))
|
snapshots = tuple(sorted(snapshot_by_source.values(), key=lambda snapshot: snapshot.semantic.encoded))
|
||||||
source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots]
|
source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots]
|
||||||
target_bindings = [
|
target_bindings = [
|
||||||
@@ -345,6 +479,7 @@ def resolve_final_layout_source_identity(
|
|||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"FinalLayoutSourceContext",
|
"FinalLayoutSourceContext",
|
||||||
|
"NodeSourceDigestCache",
|
||||||
"PreparedWeightSource",
|
"PreparedWeightSource",
|
||||||
"WeightSourceKind",
|
"WeightSourceKind",
|
||||||
"resolve_final_layout_source_identity",
|
"resolve_final_layout_source_identity",
|
||||||
|
|||||||
@@ -1436,7 +1436,7 @@ class MiniMaxH3Pipeline(
|
|||||||
with sequential_offload_component(component):
|
with sequential_offload_component(component):
|
||||||
yield
|
yield
|
||||||
return
|
return
|
||||||
staged = self._uses_manual_component_offload()
|
staged = self._uses_manual_component_offload() and getattr(component, "_omni_dlo_offload_enabled", True)
|
||||||
try:
|
try:
|
||||||
if staged:
|
if staged:
|
||||||
component.load_to_device()
|
component.load_to_device()
|
||||||
|
|||||||
@@ -1,8 +1,9 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# 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 abc import ABC, abstractmethod
|
||||||
from dataclasses import dataclass
|
from collections.abc import Mapping
|
||||||
|
from dataclasses import dataclass, field
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
from typing import Protocol, runtime_checkable
|
from typing import Protocol, runtime_checkable
|
||||||
|
|
||||||
@@ -10,7 +11,7 @@ import torch
|
|||||||
from torch import nn
|
from torch import nn
|
||||||
from vllm.logger import init_logger
|
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__)
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
@@ -52,6 +53,12 @@ class OffloadConfig:
|
|||||||
# blocks from the loader-selected host backing with H2D only.
|
# blocks from the loader-selected host backing with H2D only.
|
||||||
dlo_use_allgather: bool = True
|
dlo_use_allgather: bool = True
|
||||||
dlo_resident_layers: int = 0 # leading DiT layers kept on device
|
dlo_resident_layers: int = 0 # leading DiT layers kept on device
|
||||||
|
# Per auxiliary component: True permits model-declared DLO offload;
|
||||||
|
# False keeps the whole component resident. "default" is the fallback.
|
||||||
|
dlo_offload_components: dict[str, bool] = field(default_factory=dict)
|
||||||
|
# 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
|
@classmethod
|
||||||
def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
|
def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
|
||||||
@@ -122,6 +129,31 @@ class OffloadConfig:
|
|||||||
# requirements (concurrent requests, dummy run skip).
|
# requirements (concurrent requests, dummy run skip).
|
||||||
dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
|
dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
|
||||||
dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
|
dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
|
||||||
|
raw_component_policy = getattr(od_config, "dlo_offload_components", {})
|
||||||
|
if raw_component_policy is None:
|
||||||
|
raw_component_policy = {}
|
||||||
|
if not isinstance(raw_component_policy, Mapping):
|
||||||
|
raise TypeError(
|
||||||
|
"dlo_offload_components must be a mapping of component names to booleans, "
|
||||||
|
f"got {type(raw_component_policy).__name__}"
|
||||||
|
)
|
||||||
|
dlo_offload_components: dict[str, bool] = {}
|
||||||
|
for component, enabled in raw_component_policy.items():
|
||||||
|
if not isinstance(component, str) or not component:
|
||||||
|
raise TypeError("dlo_offload_components keys must be non-empty strings")
|
||||||
|
if not isinstance(enabled, bool):
|
||||||
|
raise TypeError(
|
||||||
|
f"dlo_offload_components[{component!r}] must be a boolean, got {type(enabled).__name__}"
|
||||||
|
)
|
||||||
|
dlo_offload_components[component] = enabled
|
||||||
|
if dlo_offload_components and not enable_distributed_layerwise_offload:
|
||||||
|
raise ValueError("dlo_offload_components requires distributed layerwise offload")
|
||||||
|
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:
|
if dlo_resident_layers < 0:
|
||||||
raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
|
raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
|
||||||
if dlo_resident_layers and dlo_use_allgather:
|
if dlo_resident_layers and dlo_use_allgather:
|
||||||
@@ -157,6 +189,8 @@ class OffloadConfig:
|
|||||||
dp_size=dp_size,
|
dp_size=dp_size,
|
||||||
dlo_use_allgather=dlo_use_allgather,
|
dlo_use_allgather=dlo_use_allgather,
|
||||||
dlo_resident_layers=dlo_resident_layers,
|
dlo_resident_layers=dlo_resident_layers,
|
||||||
|
dlo_offload_components=dlo_offload_components,
|
||||||
|
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-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.
|
"""Distributed Layerwise Offload backend with double-buffered H2D.
|
||||||
|
|
||||||
This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
|
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
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
from itertools import chain
|
from itertools import chain
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -35,7 +37,13 @@ from vllm_omni.platforms import current_omni_platform
|
|||||||
|
|
||||||
from .base import OffloadBackend, OffloadConfig
|
from .base import OffloadBackend, OffloadConfig
|
||||||
from .block_discovery import get_blocks_from_dit
|
from .block_discovery import get_blocks_from_dit
|
||||||
from .module_collector import ModuleDiscovery
|
from .host_registration import (
|
||||||
|
HostRegistration,
|
||||||
|
HostRegistrationCleanupError,
|
||||||
|
HostRegistrationError,
|
||||||
|
register_host_mappings,
|
||||||
|
)
|
||||||
|
from .module_collector import ModuleDiscovery, PipelineModules
|
||||||
from .offload_plan import (
|
from .offload_plan import (
|
||||||
OffloadPlan,
|
OffloadPlan,
|
||||||
get_offload_plan,
|
get_offload_plan,
|
||||||
@@ -51,6 +59,38 @@ from .tensor_utils import (
|
|||||||
|
|
||||||
logger = init_logger(__name__)
|
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
|
# 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
|
# use layerwise offload (streaming hooks) or be moved to GPU as a resident
|
||||||
# module. Submodules larger than this are offloaded to save HBM; smaller
|
# module. Submodules larger than this are offloaded to save HBM; smaller
|
||||||
@@ -94,6 +134,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
|
|||||||
self.rank = rank
|
self.rank = rank
|
||||||
self.pin_memory = pin_memory
|
self.pin_memory = pin_memory
|
||||||
self.rank_local_mmap = rank_local_mmap
|
self.rank_local_mmap = rank_local_mmap
|
||||||
|
self.registered_mmap = False
|
||||||
self.tensor_transforms = tensor_transforms or {}
|
self.tensor_transforms = tensor_transforms or {}
|
||||||
|
|
||||||
self.copy_stream = copy_stream or current_omni_platform.Stream()
|
self.copy_stream = copy_stream or current_omni_platform.Stream()
|
||||||
@@ -445,6 +486,26 @@ class DistributedLayerwiseOffloadHook(ModelHook):
|
|||||||
# Prefetch: H2D + AllGather (overlapped on dedicated streams) #
|
# 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
|
@staticmethod
|
||||||
def _pack_mmap_sources(
|
def _pack_mmap_sources(
|
||||||
cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
|
cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
|
||||||
@@ -457,18 +518,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
|
|||||||
destination = slot_buffers[dtype][:total_numel]
|
destination = slot_buffers[dtype][:total_numel]
|
||||||
sources = cpu_sources[dtype]
|
sources = cpu_sources[dtype]
|
||||||
for source_info, meta in zip(sources, metas, strict=True):
|
for source_info, meta in zip(sources, metas, strict=True):
|
||||||
source = source_info["tensor"]
|
source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
|
||||||
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()}"
|
|
||||||
)
|
|
||||||
start = meta["offset"]
|
start = meta["offset"]
|
||||||
physical_storage = destination[start : start + meta["numel"]]
|
physical_storage = destination[start : start + meta["numel"]]
|
||||||
if source.is_contiguous():
|
if source.is_contiguous():
|
||||||
@@ -500,6 +550,32 @@ class DistributedLayerwiseOffloadHook(ModelHook):
|
|||||||
raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated")
|
raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated")
|
||||||
return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers)
|
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
|
@torch.compiler.disable
|
||||||
def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None:
|
def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None:
|
||||||
"""Prepare next block's weights into the shared device buffer for *slot*.
|
"""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"
|
assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
|
||||||
|
|
||||||
if self.dp_size <= 1 or self.dp_group is None:
|
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
|
if self.rank_local_mmap and self.registered_mmap:
|
||||||
with current_omni_platform.stream(self.copy_stream):
|
with current_omni_platform.stream(self.copy_stream):
|
||||||
for dtype, cpu_shard in cpu_weights.items():
|
self._copy_mmap_sources_to_device(
|
||||||
gw = gpu_weights[dtype]
|
self.cpu_sources,
|
||||||
async_copy = non_blocking and cpu_shard.is_pinned()
|
self.metadata,
|
||||||
gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
|
gpu_weights,
|
||||||
evt.record(self.copy_stream)
|
non_blocking=non_blocking,
|
||||||
if self.rank_local_mmap:
|
)
|
||||||
# The CPU slot may be overwritten only after this H2D copy has
|
evt.record(self.copy_stream)
|
||||||
# finished. The shared event protects reuse by another hook.
|
else:
|
||||||
self.cpu_staging_events[slot] = evt
|
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:
|
else:
|
||||||
gpu_shards: dict[torch.dtype, torch.Tensor] = {}
|
gpu_shards: dict[torch.dtype, torch.Tensor] = {}
|
||||||
shard_bufs = self.gpu_shard_buffers[slot]
|
shard_bufs = self.gpu_shard_buffers[slot]
|
||||||
@@ -735,6 +821,7 @@ class PinnedResidentLayerGroup:
|
|||||||
self.copy_stream = copy_stream
|
self.copy_stream = copy_stream
|
||||||
self.loaded = False
|
self.loaded = False
|
||||||
self.rank_local_mmap = rank_local_mmap
|
self.rank_local_mmap = rank_local_mmap
|
||||||
|
self.registered_mmap = False
|
||||||
self.pin_memory = pin_memory
|
self.pin_memory = pin_memory
|
||||||
self._states: list[dict[str, Any]] = []
|
self._states: list[dict[str, Any]] = []
|
||||||
self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
|
self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
|
||||||
@@ -805,6 +892,14 @@ class PinnedResidentLayerGroup:
|
|||||||
ready = current_omni_platform.Event()
|
ready = current_omni_platform.Event()
|
||||||
with current_omni_platform.stream(self.copy_stream):
|
with current_omni_platform.stream(self.copy_stream):
|
||||||
for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
|
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:
|
if self.rank_local_mmap:
|
||||||
slot = index % 2
|
slot = index % 2
|
||||||
previous_copy = self._cpu_staging_events[slot]
|
previous_copy = self._cpu_staging_events[slot]
|
||||||
@@ -902,8 +997,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
self._resident_layer_group: PinnedResidentLayerGroup | None = None
|
self._resident_layer_group: PinnedResidentLayerGroup | None = None
|
||||||
self._using_mmap = False
|
self._using_mmap = False
|
||||||
self._using_rank_local_mmap = False
|
self._using_rank_local_mmap = False
|
||||||
|
self._using_registered_mmap = False
|
||||||
self.host_weight_plan = host_weight_plan
|
self.host_weight_plan = host_weight_plan
|
||||||
self._host_weight_lease: HostWeightLease | None = None
|
self._host_weight_lease: HostWeightLease | None = None
|
||||||
|
self._host_registration: HostRegistration | None = None
|
||||||
self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
|
self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
|
||||||
|
|
||||||
def load_resident_layers(self) -> None:
|
def load_resident_layers(self) -> None:
|
||||||
@@ -916,6 +1013,136 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
if self._resident_layer_group is not None:
|
if self._resident_layer_group is not None:
|
||||||
self._resident_layer_group.offload()
|
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(
|
def _load_weights_via_mmap(
|
||||||
self,
|
self,
|
||||||
pipeline: nn.Module,
|
pipeline: nn.Module,
|
||||||
@@ -1173,6 +1400,39 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
module.to(self.device)
|
module.to(self.device)
|
||||||
logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__)
|
logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__)
|
||||||
|
|
||||||
|
def _component_offload_enabled(self, name: str) -> bool:
|
||||||
|
policy = self.config.dlo_offload_components
|
||||||
|
return policy.get(name, policy.get("default", True))
|
||||||
|
|
||||||
|
def _validate_component_offload_policy(self, modules: PipelineModules) -> None:
|
||||||
|
known_components = {"default", *modules.encoder_names, *modules.vae_names}
|
||||||
|
unknown = sorted(self.config.dlo_offload_components.keys() - known_components)
|
||||||
|
if unknown:
|
||||||
|
raise ValueError(
|
||||||
|
f"dlo_offload_components contains unknown auxiliary components: {unknown}. DLO always owns DiT offload."
|
||||||
|
)
|
||||||
|
|
||||||
|
def _prepare_auxiliary_components(self, modules: PipelineModules, plan: OffloadPlan | None) -> None:
|
||||||
|
for encoder, name in zip(modules.encoders, modules.encoder_names, strict=True):
|
||||||
|
offload = self._component_offload_enabled(name)
|
||||||
|
encoder._omni_dlo_offload_enabled = offload
|
||||||
|
if offload:
|
||||||
|
self._try_layerwise_offload_encoder(encoder, name, plan)
|
||||||
|
self._register_on_demand_hook(
|
||||||
|
encoder,
|
||||||
|
name,
|
||||||
|
stage_on_demand=(offload and plan is not None and name in plan.on_demand_component_paths),
|
||||||
|
)
|
||||||
|
|
||||||
|
for vae, name in zip(modules.vaes, modules.vae_names, strict=True):
|
||||||
|
offload = self._component_offload_enabled(name)
|
||||||
|
vae._omni_dlo_offload_enabled = offload
|
||||||
|
self._register_on_demand_hook(
|
||||||
|
vae,
|
||||||
|
name,
|
||||||
|
stage_on_demand=(offload and plan is not None and name in plan.on_demand_component_paths),
|
||||||
|
)
|
||||||
|
|
||||||
def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool:
|
def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool:
|
||||||
"""Stream plan-declared encoder blocks on each rank without AllGather."""
|
"""Stream plan-declared encoder blocks on each rank without AllGather."""
|
||||||
if plan is None or name not in plan.encoder_block_attrs:
|
if plan is None or name not in plan.encoder_block_attrs:
|
||||||
@@ -1441,6 +1701,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
# Retrieve optional declarative OffloadPlan from the pipeline.
|
# Retrieve optional declarative OffloadPlan from the pipeline.
|
||||||
# When present, replaces heuristic block discovery.
|
# When present, replaces heuristic block discovery.
|
||||||
plan = get_offload_plan(pipeline)
|
plan = get_offload_plan(pipeline)
|
||||||
|
self._validate_component_offload_policy(modules)
|
||||||
|
|
||||||
if self.config.dlo_resident_layers and (plan is None or not plan.resident_dit_paths):
|
if self.config.dlo_resident_layers and (plan is None or not plan.resident_dit_paths):
|
||||||
logger.warning(
|
logger.warning(
|
||||||
@@ -1465,11 +1726,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
raise RuntimeError("DLO received a closed Host Weight Runtime lease")
|
raise RuntimeError("DLO received a closed Host Weight Runtime lease")
|
||||||
# The final-layout restorer has already rebound the model to
|
# The final-layout restorer has already rebound the model to
|
||||||
# immutable host tensors. Treat those tensors as mmap-like
|
# immutable host tensors. Treat those tensors as mmap-like
|
||||||
# sources for the bounded two-slot rank-local staging path;
|
# sources; transport setup below selects registered direct
|
||||||
# do not reinitialize or rematerialize the restored model.
|
# H2D or the bounded two-slot staging fallback.
|
||||||
self._using_rank_local_mmap = True
|
self._using_rank_local_mmap = True
|
||||||
logger.info(
|
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,
|
self._host_weight_lease.provenance.resolution_id,
|
||||||
)
|
)
|
||||||
elif self.host_weight_plan.backing_kind == "checkpoint_mmap":
|
elif self.host_weight_plan.backing_kind == "checkpoint_mmap":
|
||||||
@@ -1483,7 +1744,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
if self._using_rank_local_mmap:
|
if self._using_rank_local_mmap:
|
||||||
logger.info(
|
logger.info(
|
||||||
"DLO rank-local host storage enabled: source pages are "
|
"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:
|
else:
|
||||||
remaining_meta = [
|
remaining_meta = [
|
||||||
@@ -1501,21 +1762,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
)
|
)
|
||||||
logger.info("DLO is using host tensors materialized by the ordinary loader")
|
logger.info("DLO is using host tensors materialized by the ordinary loader")
|
||||||
|
|
||||||
# Keep VAE/encoders on CPU; move to GPU on-demand via hooks.
|
# Apply the user policy to discovered auxiliary components. The default
|
||||||
# This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9)
|
# keeps the model-declared staged lifecycle; False leaves a component
|
||||||
# during the DiT forward pass. They are only needed briefly for
|
# resident for repeated requests.
|
||||||
# text-encoding (before DiT) and VAE-decode (after DiT).
|
self._prepare_auxiliary_components(modules, plan)
|
||||||
for enc, enc_name in zip(modules.encoders, modules.encoder_names):
|
|
||||||
self._try_layerwise_offload_encoder(enc, enc_name, plan)
|
|
||||||
self._register_on_demand_hook(
|
|
||||||
enc, "encoder", stage_on_demand=plan is not None and enc_name in plan.on_demand_component_paths
|
|
||||||
)
|
|
||||||
for vae, vae_name in zip(modules.vaes, modules.vae_names):
|
|
||||||
self._register_on_demand_hook(
|
|
||||||
vae,
|
|
||||||
"vae",
|
|
||||||
stage_on_demand=(plan is not None and vae_name in plan.on_demand_component_paths),
|
|
||||||
)
|
|
||||||
|
|
||||||
# Move resident modules to GPU (small modules needed every forward)
|
# Move resident modules to GPU (small modules needed every forward)
|
||||||
for name, module in zip(modules.resident_names, modules.resident_modules):
|
for name, module in zip(modules.resident_names, modules.resident_modules):
|
||||||
@@ -1653,6 +1903,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
)
|
)
|
||||||
pipeline._dlo_residency_controller = self
|
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:
|
if not self._all_hook_groups:
|
||||||
self.enabled = bool(self._resident_blocks)
|
self.enabled = bool(self._resident_blocks)
|
||||||
if self._using_mmap and not self.enabled:
|
if self._using_mmap and not self.enabled:
|
||||||
@@ -1662,17 +1915,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
# Unified allocation: 2 shared output buffers + 2 shared shard buffers
|
# Unified allocation: 2 shared output buffers + 2 shared shard buffers
|
||||||
# sized to the max block across ALL module groups (gen_layers +
|
# sized to the max block across ALL module groups (gen_layers +
|
||||||
# language_model). Groups execute sequentially, so 2 buffers suffice.
|
# 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_buffers = self._allocate_shared_buffers(all_hooks)
|
||||||
unified_shard_buffers = None
|
unified_shard_buffers = None
|
||||||
if self.dp_size > 1:
|
if self.dp_size > 1:
|
||||||
unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
|
unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
|
||||||
unified_cpu_staging = None
|
unified_cpu_staging = None
|
||||||
cpu_staging_events = 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(
|
unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
|
||||||
all_hooks,
|
all_hooks,
|
||||||
self._resident_layer_group,
|
self._resident_layer_group,
|
||||||
@@ -1753,6 +2002,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
|
|
||||||
def _release_mmap_handles(self) -> None:
|
def _release_mmap_handles(self) -> None:
|
||||||
"""Release source handles and the transport-owned HWR lease."""
|
"""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()
|
self._mmap_transforms_by_tensor_id.clear()
|
||||||
if hasattr(self, "_mmap_file_cache"):
|
if hasattr(self, "_mmap_file_cache"):
|
||||||
self._mmap_file_cache.clear()
|
self._mmap_file_cache.clear()
|
||||||
@@ -1781,6 +2032,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
|
|
||||||
def disable(self) -> None:
|
def disable(self) -> None:
|
||||||
has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed
|
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 = (
|
has_carrier = (
|
||||||
self.host_weight_plan is not None
|
self.host_weight_plan is not None
|
||||||
and self.host_weight_plan.lease_carrier is not None
|
and self.host_weight_plan.lease_carrier is not None
|
||||||
@@ -1796,12 +2048,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
not self.enabled
|
not self.enabled
|
||||||
and not hasattr(self, "_mmap_file_cache")
|
and not hasattr(self, "_mmap_file_cache")
|
||||||
and not has_open_lease
|
and not has_open_lease
|
||||||
|
and not has_registration
|
||||||
and not has_carrier
|
and not has_carrier
|
||||||
and not has_partial_hooks
|
and not has_partial_hooks
|
||||||
):
|
):
|
||||||
return
|
return
|
||||||
|
|
||||||
if self._using_rank_local_mmap:
|
if self._using_rank_local_mmap or has_registration:
|
||||||
current_omni_platform.synchronize()
|
current_omni_platform.synchronize()
|
||||||
|
|
||||||
for blocks in self._blocks:
|
for blocks in self._blocks:
|
||||||
@@ -1818,9 +2071,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
|
|||||||
self._all_hook_groups.clear()
|
self._all_hook_groups.clear()
|
||||||
self._resident_blocks.clear()
|
self._resident_blocks.clear()
|
||||||
self._resident_layer_group = None
|
self._resident_layer_group = None
|
||||||
|
self._release_registered_mmap()
|
||||||
self._release_mmap_handles()
|
self._release_mmap_handles()
|
||||||
self._using_mmap = False
|
self._using_mmap = False
|
||||||
self._using_rank_local_mmap = False
|
self._using_rank_local_mmap = False
|
||||||
|
self._using_registered_mmap = False
|
||||||
self.enabled = False
|
self.enabled = False
|
||||||
logger.info("Distributed layer-wise offloading disabled")
|
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",
|
||||||
|
]
|
||||||
@@ -1,5 +1,5 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# 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.
|
Diffusion Worker for vLLM-Omni.
|
||||||
@@ -897,11 +897,18 @@ class DiffusionWorker:
|
|||||||
|
|
||||||
def shutdown(self) -> None:
|
def shutdown(self) -> None:
|
||||||
"""Shutdown the worker and cleanup distributed environment."""
|
"""Shutdown the worker and cleanup distributed environment."""
|
||||||
if self.model_runner is not None:
|
try:
|
||||||
mgr = getattr(self.model_runner, "kv_transfer_manager", None)
|
if self.model_runner is not None:
|
||||||
if mgr is not None:
|
mgr = getattr(self.model_runner, "kv_transfer_manager", None)
|
||||||
mgr.shutdown_prefetch()
|
try:
|
||||||
destroy_distributed_env()
|
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:
|
class CustomPipelineWorkerExtension:
|
||||||
|
|||||||
@@ -569,8 +569,10 @@ class OrchestratorArgs:
|
|||||||
enable_distributed_layerwise_offload: bool = False
|
enable_distributed_layerwise_offload: bool = False
|
||||||
dlo_use_allgather: bool = True
|
dlo_use_allgather: bool = True
|
||||||
dlo_resident_layers: int = 0
|
dlo_resident_layers: int = 0
|
||||||
|
dlo_offload_components: dict[str, bool] = field(default_factory=dict)
|
||||||
host_weight_runtime_mode: str = "disabled"
|
host_weight_runtime_mode: str = "disabled"
|
||||||
host_weight_runtime_root: str | None = None
|
host_weight_runtime_root: str | None = None
|
||||||
|
dlo_host_registration_limit_gib: float = 0.0
|
||||||
boundary_ratio: float | None = None
|
boundary_ratio: float | None = None
|
||||||
flow_shift: float | None = None
|
flow_shift: float | None = None
|
||||||
diffusion_kv_cache_dtype: str | None = None
|
diffusion_kv_cache_dtype: str | None = None
|
||||||
|
|||||||
@@ -1059,8 +1059,12 @@ class AsyncOmniEngine:
|
|||||||
"enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
|
"enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
|
||||||
"dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
|
"dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
|
||||||
"dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
|
"dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
|
||||||
|
"dlo_offload_components": (
|
||||||
|
{} if kwargs.get("dlo_offload_components") is None else kwargs["dlo_offload_components"]
|
||||||
|
),
|
||||||
"host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
|
"host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
|
||||||
"host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
|
"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"),
|
"enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
|
||||||
"diffusion_compile_granularity": (
|
"diffusion_compile_granularity": (
|
||||||
"regional"
|
"regional"
|
||||||
|
|||||||
@@ -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.
|
Orchestrator for vLLM-Omni multi-stage runtime.
|
||||||
|
|
||||||
@@ -363,8 +366,8 @@ class Orchestrator:
|
|||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
request_async_queue: janus.AsyncQueue[EngineQueueMessage],
|
request_async_queue: janus.AsyncQueue[EngineQueueMessage],
|
||||||
output_async_queue: janus.AsyncQueue[dict[str, Any]],
|
output_async_queue: janus.AsyncQueue[EngineQueueMessage],
|
||||||
rpc_async_queue: janus.AsyncQueue[dict[str, Any]],
|
rpc_async_queue: janus.AsyncQueue[EngineQueueMessage],
|
||||||
stage_pools: list[StagePool],
|
stage_pools: list[StagePool],
|
||||||
*,
|
*,
|
||||||
async_chunk: bool = False,
|
async_chunk: bool = False,
|
||||||
@@ -1173,6 +1176,35 @@ class Orchestrator:
|
|||||||
abort=True,
|
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(
|
async def _dispatch_or_fail_request(
|
||||||
self,
|
self,
|
||||||
dispatch: Callable[[], Awaitable[Any]],
|
dispatch: Callable[[], Awaitable[Any]],
|
||||||
@@ -1194,7 +1226,9 @@ class Orchestrator:
|
|||||||
StageUnavailableError (no live replica / evicted slot) or
|
StageUnavailableError (no live replica / evicted slot) or
|
||||||
EngineDeadError (replica died mid-submit before the poll loop evicted
|
EngineDeadError (replica died mid-submit before the poll loop evicted
|
||||||
it). Both mean "this request cannot be placed", not "the server is
|
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.
|
propagate. Returns False when the request was failed.
|
||||||
|
|
||||||
``req_id`` is the request the failure is attributed to; ``dispatch_req_id``
|
``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():
|
if dead_replica is not None and dead_replica in pool.live_replica_ids():
|
||||||
await self._handle_dead_replica(stage_id, dead_replica, e)
|
await self._handle_dead_replica(stage_id, dead_replica, e)
|
||||||
return False
|
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 ----
|
# ---- Shared helpers ----
|
||||||
|
|
||||||
@@ -2292,21 +2344,11 @@ class Orchestrator:
|
|||||||
"and none were provided; failing the request",
|
"and none were provided; failing the request",
|
||||||
request_id,
|
request_id,
|
||||||
)
|
)
|
||||||
await self.output_async_queue.put(
|
await self._fail_request_client_error(
|
||||||
ErrorMessage(
|
request_id,
|
||||||
error=(
|
0,
|
||||||
"async_chunk requires prompt_token_ids on the stage-0 request; "
|
"async_chunk requires prompt_token_ids on the stage-0 request; "
|
||||||
"an embeds-only prompt cannot prewarm downstream stages"
|
"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,
|
|
||||||
close_duplex_sessions=True,
|
close_duplex_sessions=True,
|
||||||
)
|
)
|
||||||
return False
|
return False
|
||||||
|
|||||||
@@ -1558,7 +1558,6 @@ def launch_diffusion_stage_replica(
|
|||||||
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
|
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
|
||||||
|
|
||||||
od_config = build_diffusion_config(model, stage_config, metadata)
|
od_config = build_diffusion_config(model, stage_config, metadata)
|
||||||
od_config.max_num_seqs = batch_size
|
|
||||||
parallel_config = getattr(od_config, "parallel_config", None)
|
parallel_config = getattr(od_config, "parallel_config", None)
|
||||||
world_size = getattr(parallel_config, "world_size", 1)
|
world_size = getattr(parallel_config, "world_size", 1)
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -1506,15 +1506,14 @@ def initialize_diffusion_stage(
|
|||||||
stage_cfg: Stage configuration.
|
stage_cfg: Stage configuration.
|
||||||
metadata: Extracted stage metadata.
|
metadata: Extracted stage metadata.
|
||||||
stage_init_timeout: Timeout in seconds for stage initialization handshake
|
stage_init_timeout: Timeout in seconds for stage initialization handshake
|
||||||
batch_size: Maximum number of requests to batch together in the
|
batch_size: Client-side request batch width. Does not set scheduler
|
||||||
diffusion engine. Passed through to ``StageDiffusionClient``
|
``max_num_seqs``; pass ``--max-num-seqs`` or stage YAML for that.
|
||||||
and ultimately to ``AsyncOmni``.
|
Forwarded to ``StageDiffusionClient``.
|
||||||
use_inline: If True, uses the inline diffusion client instead of subprocess.
|
use_inline: If True, uses the inline diffusion client instead of subprocess.
|
||||||
"""
|
"""
|
||||||
from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
|
from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
|
||||||
|
|
||||||
od_config = build_diffusion_config(model, stage_cfg, metadata)
|
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)
|
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.
|
"""vLLM-Omni extensions for the ``vllm bench serve`` CLI.
|
||||||
|
|
||||||
Core functions:
|
Core functions:
|
||||||
@@ -16,6 +19,63 @@ by ``add_omni_args``.
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import argparse
|
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:
|
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 = (
|
_OMNI_BENCH_DATASET_CHOICES = (
|
||||||
"daily-omni",
|
"daily-omni",
|
||||||
|
"omniinteract",
|
||||||
"seed-tts",
|
"seed-tts",
|
||||||
"seed-tts-text",
|
"seed-tts-text",
|
||||||
"seed-tts-design",
|
"seed-tts-design",
|
||||||
@@ -218,6 +279,11 @@ def extend_omni_choices(parser: argparse.ArgumentParser) -> None:
|
|||||||
def update_omni_help(parser: argparse.ArgumentParser) -> None:
|
def update_omni_help(parser: argparse.ArgumentParser) -> None:
|
||||||
"""Update upstream argument help text to describe Omni-specific behavior."""
|
"""Update upstream argument help text to describe Omni-specific behavior."""
|
||||||
for action in parser._actions:
|
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":
|
if action.dest == "percentile_metrics":
|
||||||
action.help = (
|
action.help = (
|
||||||
"Comma-separated list of selected metrics to report percentiles. "
|
"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:
|
def add_omni_args(parser: argparse.ArgumentParser) -> None:
|
||||||
"""Register all vLLM-Omni serving benchmark arguments."""
|
"""Register all vLLM-Omni serving benchmark arguments."""
|
||||||
add_daily_omni_cli_args(parser)
|
add_daily_omni_cli_args(parser)
|
||||||
|
add_omniinteract_cli_args(parser)
|
||||||
add_seed_tts_cli_args(parser)
|
add_seed_tts_cli_args(parser)
|
||||||
add_multi_stage_cli_args(parser)
|
add_multi_stage_cli_args(parser)
|
||||||
add_diffusion_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:
|
def preprocess_serve_args(args: argparse.Namespace) -> None:
|
||||||
"""Apply serving benchmark CLI transformations after parsing."""
|
"""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 {})
|
extra_body = dict(getattr(args, "extra_body", None) or {})
|
||||||
bot_task = getattr(args, "bot_task", None)
|
bot_task = getattr(args, "bot_task", None)
|
||||||
if getattr(args, "backend", None) == "openai-image-edits-omni" and bot_task is not None:
|
if getattr(args, "backend", None) == "openai-image-edits-omni" and bot_task is not None:
|
||||||
|
|||||||
@@ -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.
|
"""Online serving benchmark subcommand for vLLM-Omni.
|
||||||
|
|
||||||
``OmniBenchmarkServingSubcommand`` starts from vLLM's serving benchmark
|
``OmniBenchmarkServingSubcommand`` starts from vLLM's serving benchmark
|
||||||
@@ -24,7 +27,7 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
|
|||||||
"""The `serve` subcommand for vllm bench."""
|
"""The `serve` subcommand for vllm bench."""
|
||||||
|
|
||||||
name = "serve"
|
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
|
@classmethod
|
||||||
def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
|
def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
|
||||||
|
|||||||
@@ -760,6 +760,28 @@ class OmniServeCommand(CLISubcommand):
|
|||||||
"the same persistent path for population and serving."
|
"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."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
omni_config_group.add_argument(
|
||||||
|
"--dlo-offload-components",
|
||||||
|
type=json.loads,
|
||||||
|
default=None,
|
||||||
|
metavar="JSON",
|
||||||
|
help=(
|
||||||
|
"Per auxiliary encoder/VAE DLO policy as a JSON object, for example "
|
||||||
|
"'{\"text_encoder\": false}'. False keeps the whole component "
|
||||||
|
"resident; true permits the model's existing offload behavior. "
|
||||||
|
"DiT offload is not controlled by this map."
|
||||||
|
),
|
||||||
|
)
|
||||||
# Video model parameters (e.g., Wan2.2) - engine-level
|
# Video model parameters (e.g., Wan2.2) - engine-level
|
||||||
omni_config_group.add_argument(
|
omni_config_group.add_argument(
|
||||||
"--boundary-ratio",
|
"--boundary-ratio",
|
||||||
|
|||||||
@@ -1,9 +1,15 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||||
import math
|
import math
|
||||||
from typing import Any, Literal
|
from typing import Any, Literal
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
|
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
|
_MAX_EMBEDDING_DIM = 8192
|
||||||
|
|
||||||
SUPPORTED_AUDIO_FORMATS: frozenset[str] = frozenset({"wav", "pcm", "flac", "mp3", "opus"})
|
SUPPORTED_AUDIO_FORMATS: frozenset[str] = frozenset({"wav", "pcm", "flac", "mp3", "opus"})
|
||||||
@@ -142,18 +148,21 @@ class OpenAICreateSpeechRequest(BaseModel):
|
|||||||
)
|
)
|
||||||
max_new_tokens: int | None = Field(
|
max_new_tokens: int | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
|
ge=1,
|
||||||
|
le=_INT64_MAX,
|
||||||
description="Maximum tokens to generate",
|
description="Maximum tokens to generate",
|
||||||
)
|
)
|
||||||
seed: int | None = Field(
|
seed: int | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
ge=0,
|
ge=_INT64_MIN,
|
||||||
le=2**63 - 1,
|
le=_INT64_MAX,
|
||||||
description="Random seed for reproducible generation. When set, ensures "
|
description="Random seed for reproducible generation. When set, ensures "
|
||||||
"deterministic output for the same input text and seed value.",
|
"deterministic output for the same input text and seed value.",
|
||||||
)
|
)
|
||||||
initial_codec_chunk_frames: int | None = Field(
|
initial_codec_chunk_frames: int | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
ge=0,
|
ge=0,
|
||||||
|
le=_INT64_MAX,
|
||||||
description="Per-request initial chunk size override. If null, computed dynamically based on server load.",
|
description="Per-request initial chunk size override. If null, computed dynamically based on server load.",
|
||||||
)
|
)
|
||||||
non_streaming_mode: bool | None = Field(
|
non_streaming_mode: bool | None = Field(
|
||||||
@@ -365,10 +374,14 @@ class OpenAICreateAudioGenerateRequest(BaseModel):
|
|||||||
)
|
)
|
||||||
num_inference_steps: int | None = Field(
|
num_inference_steps: int | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
|
ge=1,
|
||||||
|
le=_INT64_MAX,
|
||||||
description="Number of inference steps",
|
description="Number of inference steps",
|
||||||
)
|
)
|
||||||
seed: int | None = Field(
|
seed: int | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
|
ge=_INT64_MIN,
|
||||||
|
le=_INT64_MAX,
|
||||||
description="Random seed for reproducibility",
|
description="Random seed for reproducibility",
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -413,8 +426,8 @@ class SpeechBatchItem(BaseModel):
|
|||||||
ref_audio: str | None = None
|
ref_audio: str | None = None
|
||||||
ref_text: str | None = None
|
ref_text: str | None = None
|
||||||
x_vector_only_mode: bool | None = None
|
x_vector_only_mode: bool | None = None
|
||||||
max_new_tokens: int | None = None
|
max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
|
||||||
initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
|
initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX)
|
||||||
non_streaming_mode: bool | None = None
|
non_streaming_mode: bool | None = None
|
||||||
|
|
||||||
|
|
||||||
@@ -433,8 +446,8 @@ class BatchSpeechRequest(BaseModel):
|
|||||||
ref_audio: str | None = None
|
ref_audio: str | None = None
|
||||||
ref_text: str | None = None
|
ref_text: str | None = None
|
||||||
x_vector_only_mode: bool | None = None
|
x_vector_only_mode: bool | None = None
|
||||||
max_new_tokens: int | None = None
|
max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
|
||||||
initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
|
initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX)
|
||||||
non_streaming_mode: bool | None = None
|
non_streaming_mode: bool | None = None
|
||||||
|
|
||||||
|
|
||||||
@@ -517,10 +530,11 @@ class StreamingSpeechSessionConfig(BaseModel):
|
|||||||
instructions: str | None = None
|
instructions: str | None = None
|
||||||
response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
|
response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
|
||||||
speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
|
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(
|
initial_codec_chunk_frames: int | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
ge=0,
|
ge=0,
|
||||||
|
le=_INT64_MAX,
|
||||||
description="Initial chunk size for reduced TTFA. Overrides the deploy config for this session.",
|
description="Initial chunk size for reduced TTFA. Overrides the deploy config for this session.",
|
||||||
)
|
)
|
||||||
non_streaming_mode: bool | None = Field(
|
non_streaming_mode: bool | None = Field(
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# 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.
|
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
|
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 = {
|
_IMAGE_FILE_METADATA = {
|
||||||
"jpg": ("jpg", "image/jpeg"),
|
"jpg": ("jpg", "image/jpeg"),
|
||||||
"jpeg": ("jpeg", "image/jpeg"),
|
"jpeg": ("jpeg", "image/jpeg"),
|
||||||
@@ -144,7 +148,7 @@ class ImageGenerationRequest(BaseModel):
|
|||||||
default=None,
|
default=None,
|
||||||
description="Optional model-specific parameters passed directly to the model's extra_args.",
|
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(
|
generator_device: str | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
description="Device for the seeded torch.Generator (e.g. 'cpu', 'cuda'). Defaults to the runner's device.",
|
description="Device for the seeded torch.Generator (e.g. 'cpu', 'cuda'). Defaults to the runner's device.",
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# 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.
|
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.entrypoints.openai.image_api_utils import parse_size
|
||||||
from vllm_omni.inputs.data import DIFFUSION_QUALITY_LEVELS
|
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):
|
class VideoGenerationStatus(str, Enum):
|
||||||
QUEUED = "queued"
|
QUEUED = "queued"
|
||||||
@@ -48,10 +52,10 @@ def file_extension(media_type: str):
|
|||||||
class VideoParams(BaseModel):
|
class VideoParams(BaseModel):
|
||||||
"""Optional block for video-specific parameters."""
|
"""Optional block for video-specific parameters."""
|
||||||
|
|
||||||
width: int | None = Field(default=None, ge=1, description="Video width in pixels")
|
width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
|
||||||
height: int | None = Field(default=None, ge=1, description="Video height 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, description="Number of frames")
|
num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames")
|
||||||
fps: int | None = Field(default=None, ge=1, description="Frames per second for output video")
|
fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video")
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def size(self) -> str | None:
|
def size(self) -> str | None:
|
||||||
@@ -140,10 +144,10 @@ class VideoGenerationRequest(BaseModel):
|
|||||||
user: str | None = Field(default=None, description="User identifier for tracking")
|
user: str | None = Field(default=None, description="User identifier for tracking")
|
||||||
|
|
||||||
# Video-specific fields (top-level for OpenAI-style compatibility)
|
# Video-specific fields (top-level for OpenAI-style compatibility)
|
||||||
width: int | None = Field(default=None, ge=1, description="Video width in pixels")
|
width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
|
||||||
height: int | None = Field(default=None, ge=1, description="Video height 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, description="Frames per second for output video")
|
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, description="Number of frames to generate")
|
num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames to generate")
|
||||||
aspect_ratio: str | None = Field(
|
aspect_ratio: str | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
description=(
|
description=(
|
||||||
@@ -151,7 +155,12 @@ class VideoGenerationRequest(BaseModel):
|
|||||||
"FL2VA follows the input image; Ref2VA defaults to 16:9."
|
"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(
|
num_outputs_per_prompt: int = Field(
|
||||||
default=1,
|
default=1,
|
||||||
ge=1,
|
ge=1,
|
||||||
@@ -218,7 +227,7 @@ class VideoGenerationRequest(BaseModel):
|
|||||||
le=20.0,
|
le=20.0,
|
||||||
description="True CFG scale (model-specific parameter, may be ignored if not supported)",
|
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(
|
generate_sound: bool = Field(
|
||||||
default=False,
|
default=False,
|
||||||
description="Request model-generated audio for video models that support sound generation.",
|
description="Request model-generated audio for video models that support sound generation.",
|
||||||
@@ -237,6 +246,7 @@ class VideoGenerationRequest(BaseModel):
|
|||||||
frame_interpolation_exp: int = Field(
|
frame_interpolation_exp: int = Field(
|
||||||
default=1,
|
default=1,
|
||||||
ge=1,
|
ge=1,
|
||||||
|
le=_INT64_MAX,
|
||||||
description="Interpolation exponent: 1=2x temporal resolution, 2=4x, etc.",
|
description="Interpolation exponent: 1=2x temporal resolution, 2=4x, etc.",
|
||||||
)
|
)
|
||||||
frame_interpolation_scale: float = Field(
|
frame_interpolation_scale: float = Field(
|
||||||
|
|||||||
@@ -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."""
|
"""Reusable Realtime WebSocket, PCM, and event helpers for MiniCPM-o demos."""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -30,6 +33,13 @@ def _rounded_ms(value: float) -> float:
|
|||||||
return round(float(value), 3)
|
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]:
|
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)
|
clean = sorted(_rounded_ms(value) for value in values if math.isfinite(value) and value >= 0)
|
||||||
if not clean:
|
if not clean:
|
||||||
@@ -56,11 +66,7 @@ def summarize_session_request_metrics(
|
|||||||
"""Average client-observed metrics across turns that emitted audio."""
|
"""Average client-observed metrics across turns that emitted audio."""
|
||||||
|
|
||||||
def mean(metric: str, *, digits: int = 3) -> float | None:
|
def mean(metric: str, *, digits: int = 3) -> float | None:
|
||||||
values = [
|
values = [value for request in request_metrics if (value := _finite_number(request.get(metric))) is not None]
|
||||||
float(request[metric])
|
|
||||||
for request in request_metrics
|
|
||||||
if isinstance(request.get(metric), int | float) and math.isfinite(float(request[metric]))
|
|
||||||
]
|
|
||||||
return round(sum(values) / len(values), digits) if values else None
|
return round(sum(values) / len(values), digits) if values else None
|
||||||
|
|
||||||
return {
|
return {
|
||||||
@@ -93,24 +99,57 @@ def _event_stage_metrics(event: dict[str, object]) -> dict[str, object] | None:
|
|||||||
|
|
||||||
def build_realtime_url(
|
def build_realtime_url(
|
||||||
url: str,
|
url: str,
|
||||||
model: str,
|
model: str | None,
|
||||||
*,
|
*,
|
||||||
autostart: bool | None = None,
|
autostart: bool | None = None,
|
||||||
|
native_duplex: bool | None = True,
|
||||||
session_id: str | None = None,
|
session_id: str | None = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Add the explicit native-duplex query parameters to a Realtime URL."""
|
"""Add the explicit native-duplex query parameters to a Realtime URL."""
|
||||||
parts = urlsplit(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 = dict(parse_qsl(parts.query, keep_blank_values=True))
|
||||||
query.setdefault("duplex", "1")
|
query["duplex"] = "1"
|
||||||
query.setdefault("model", model)
|
if model:
|
||||||
query.setdefault("minicpmo45_native_duplex", "1")
|
query["model"] = model
|
||||||
|
if native_duplex is not None:
|
||||||
|
query["minicpmo45_native_duplex"] = "1" if native_duplex else "0"
|
||||||
if autostart is not None:
|
if autostart is not None:
|
||||||
query.setdefault("autostart", "1" if autostart else "0")
|
query["autostart"] = "1" if autostart else "0"
|
||||||
if session_id:
|
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))
|
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:
|
def read_pcm16_wav(path: Path) -> bytes:
|
||||||
"""Read a mono, uncompressed, 16 kHz PCM16 WAV file."""
|
"""Read a mono, uncompressed, 16 kHz PCM16 WAV file."""
|
||||||
with wave.open(str(path), "rb") as 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, ())
|
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]]:
|
def errors(self) -> list[dict[str, object]]:
|
||||||
return [event for event in self.events if event.get("type") == "error"]
|
return [event for event in self.events if event.get("type") == "error"]
|
||||||
|
|
||||||
@@ -282,15 +335,16 @@ class RealtimeEventCollector:
|
|||||||
if stage0_metrics is not None:
|
if stage0_metrics is not None:
|
||||||
raw_itls = stage0_metrics.get("vllm_itls_ms")
|
raw_itls = stage0_metrics.get("vllm_itls_ms")
|
||||||
itls = (
|
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)
|
if isinstance(raw_itls, list)
|
||||||
else []
|
else []
|
||||||
)
|
)
|
||||||
result["stage0_tokens"] = {
|
result["stage0_tokens"] = {
|
||||||
"source": "engine_stage_metrics",
|
"source": "engine_stage_metrics",
|
||||||
"output_token_count": int(stage0_metrics.get("num_tokens_out") or 0),
|
"output_token_count": int(_finite_number(stage0_metrics.get("num_tokens_out"), nonnegative=True) or 0),
|
||||||
"ttft_ms": float(stage0_metrics.get("vllm_ttft_ms") or 0.0),
|
"ttft_ms": _finite_number(stage0_metrics.get("vllm_ttft_ms"), nonnegative=True) or 0.0,
|
||||||
"tpot_ms": float(stage0_metrics.get("vllm_tpot_ms") 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),
|
"inter_token_interval_ms": _interval_summary(itls),
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -368,15 +422,26 @@ class RealtimeEventCollector:
|
|||||||
class RealtimeDuplexClient:
|
class RealtimeDuplexClient:
|
||||||
"""Small async client used by the user demo and reusable smoke probes."""
|
"""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.url = url
|
||||||
self.max_size = max_size
|
self.max_size = max_size
|
||||||
|
self.additional_headers = additional_headers
|
||||||
self.events = RealtimeEventCollector()
|
self.events = RealtimeEventCollector()
|
||||||
self._ws: Any = None
|
self._ws: Any = None
|
||||||
self._reader_task: asyncio.Task[None] | None = None
|
self._reader_task: asyncio.Task[None] | None = None
|
||||||
|
|
||||||
async def __aenter__(self) -> RealtimeDuplexClient:
|
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())
|
self._reader_task = asyncio.create_task(self._read_events())
|
||||||
return self
|
return self
|
||||||
|
|
||||||
@@ -405,6 +470,19 @@ class RealtimeDuplexClient:
|
|||||||
async def send(self, event: dict[str, object]) -> None:
|
async def send(self, event: dict[str, object]) -> None:
|
||||||
await self._ws.send(json.dumps(event))
|
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(
|
async def configure(
|
||||||
self,
|
self,
|
||||||
model: str,
|
model: str,
|
||||||
@@ -418,6 +496,7 @@ class RealtimeDuplexClient:
|
|||||||
temperature: float | None = None,
|
temperature: float | None = None,
|
||||||
extra_body: dict[str, object] | None = None,
|
extra_body: dict[str, object] | None = None,
|
||||||
session_id: str | None = None,
|
session_id: str | None = None,
|
||||||
|
idle_timeout_s: float | None = None,
|
||||||
timeout_s: float = 20.0,
|
timeout_s: float = 20.0,
|
||||||
) -> None:
|
) -> None:
|
||||||
session_extra_body = dict(extra_body or {})
|
session_extra_body = dict(extra_body or {})
|
||||||
@@ -448,14 +527,23 @@ class RealtimeDuplexClient:
|
|||||||
if instructions is not None:
|
if instructions is not None:
|
||||||
session["instructions"] = instructions
|
session["instructions"] = instructions
|
||||||
if initial_user_text is not None:
|
if initial_user_text is not None:
|
||||||
extra_body = session["extra_body"]
|
session_extra = session["extra_body"]
|
||||||
assert isinstance(extra_body, dict)
|
assert isinstance(session_extra, dict)
|
||||||
extra_body["duplex_initial_user_text"] = initial_user_text
|
session_extra["duplex_initial_user_text"] = initial_user_text
|
||||||
if session_id:
|
if session_id:
|
||||||
session["session_id"] = 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})
|
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(
|
await wait_for(
|
||||||
lambda: self.events.count("session.created") > 0,
|
session_created,
|
||||||
timeout_s=timeout_s,
|
timeout_s=timeout_s,
|
||||||
label="session.created",
|
label="session.created",
|
||||||
)
|
)
|
||||||
@@ -498,20 +586,31 @@ class RealtimeDuplexClient:
|
|||||||
if not pcm16:
|
if not pcm16:
|
||||||
continue
|
continue
|
||||||
played_ms = len(pcm16) * 1000 // (self.events.output_sample_rate_hz * PCM16_BYTES_PER_SAMPLE)
|
played_ms = len(pcm16) * 1000 // (self.events.output_sample_rate_hz * PCM16_BYTES_PER_SAMPLE)
|
||||||
await self.send(
|
await self.send_playback_ack(response_id, played_ms)
|
||||||
{
|
|
||||||
"type": "playback.ack",
|
async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
|
||||||
"response_id": response_id,
|
await self.send(
|
||||||
"item_id": f"item_{response_id}",
|
{
|
||||||
"played_ms": played_ms,
|
"type": "playback.ack",
|
||||||
"committed_ms": played_ms,
|
"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:
|
async def close_session(self, *, timeout_s: float = 20.0) -> None:
|
||||||
|
from_index = len(self.events.events)
|
||||||
await self.send({"type": "session.close"})
|
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(
|
await wait_for(
|
||||||
lambda: self.events.count("session.closed") > 0,
|
session_closed,
|
||||||
timeout_s=timeout_s,
|
timeout_s=timeout_s,
|
||||||
label="session.closed",
|
label="session.closed",
|
||||||
)
|
)
|
||||||
@@ -12,17 +12,19 @@ import os
|
|||||||
import struct
|
import struct
|
||||||
import uuid
|
import uuid
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
|
from concurrent.futures import Future, ThreadPoolExecutor
|
||||||
|
from dataclasses import dataclass
|
||||||
from pathlib import Path, PurePosixPath
|
from pathlib import Path, PurePosixPath
|
||||||
from types import TracebackType
|
from types import TracebackType
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from safetensors import safe_open
|
|
||||||
|
|
||||||
from ..identity import WeightArtifactIdentity, canonical_json
|
from ..identity import WeightArtifactIdentity, canonical_json
|
||||||
from ..manifest import FileManifestEntry, PublicationMarker, TensorManifestEntry, WeightManifest
|
from ..manifest import FileManifestEntry, PublicationMarker, TensorManifestEntry, WeightManifest
|
||||||
from ..protocols import ProductionMetadata, TensorWriteSpec
|
from ..protocols import ProductionMetadata, TensorWriteSpec
|
||||||
|
|
||||||
_HASH_CHUNK_BYTES = 8 * 1024**2
|
_HASH_CHUNK_BYTES = 8 * 1024**2
|
||||||
|
_MAX_HASH_WORKERS = 8
|
||||||
_MANIFEST_FILE = "manifest.json"
|
_MANIFEST_FILE = "manifest.json"
|
||||||
_READY_FILE = "READY.json"
|
_READY_FILE = "READY.json"
|
||||||
|
|
||||||
@@ -79,21 +81,47 @@ def _fsync_directory(path: Path) -> None:
|
|||||||
os.close(fd)
|
os.close(fd)
|
||||||
|
|
||||||
|
|
||||||
def _file_sha256(path: Path) -> str:
|
def _fsync_payload(fd: int) -> None:
|
||||||
digest = hashlib.sha256()
|
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:
|
with path.open("rb") as handle:
|
||||||
while chunk := handle.read(_HASH_CHUNK_BYTES):
|
prefix = handle.read(8)
|
||||||
digest.update(chunk)
|
if len(prefix) != 8:
|
||||||
return digest.hexdigest()
|
raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}")
|
||||||
|
header_size = struct.unpack("<Q", prefix)[0]
|
||||||
|
header = handle.read(header_size)
|
||||||
def _tensor_sha256(tensor: torch.Tensor) -> str:
|
if len(header) != header_size:
|
||||||
digest = hashlib.sha256()
|
raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}")
|
||||||
byte_view = tensor.detach().reshape(-1).view(torch.uint8)
|
file_digest.update(prefix)
|
||||||
for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
|
file_digest.update(header)
|
||||||
chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
|
for spec in specs:
|
||||||
digest.update(memoryview(chunk.numpy()))
|
digest = hashlib.sha256()
|
||||||
return digest.hexdigest()
|
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:
|
class _Coverage:
|
||||||
@@ -130,6 +158,8 @@ class StreamingSafeTensorFile:
|
|||||||
owner: FilesystemArtifactWriter,
|
owner: FilesystemArtifactWriter,
|
||||||
relative_name: str,
|
relative_name: str,
|
||||||
specs: tuple[TensorWriteSpec, ...],
|
specs: tuple[TensorWriteSpec, ...],
|
||||||
|
*,
|
||||||
|
ordered: bool,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._owner = owner
|
self._owner = owner
|
||||||
self.relative_name = _safe_relative_name(relative_name)
|
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 = tuple(sorted(specs, key=lambda spec: spec.name))
|
||||||
self._specs = {spec.name: spec for spec in self.specs}
|
self._specs = {spec.name: spec for spec in self.specs}
|
||||||
self._coverage = {spec.name: _Coverage(spec.nbytes) for spec in self.specs}
|
self._coverage = {spec.name: _Coverage(spec.nbytes) for spec in self.specs}
|
||||||
|
self._ordered = ordered
|
||||||
self._closed = False
|
self._closed = False
|
||||||
|
|
||||||
offsets: dict[str, int] = {}
|
offsets: dict[str, int] = {}
|
||||||
@@ -167,6 +198,7 @@ class StreamingSafeTensorFile:
|
|||||||
total_size = self._payload_offset + cursor
|
total_size = self._payload_offset + cursor
|
||||||
owner._reserve_file(total_size)
|
owner._reserve_file(total_size)
|
||||||
self._fd = os.open(self._path, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_CLOEXEC, 0o600)
|
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:
|
try:
|
||||||
if hasattr(os, "posix_fallocate"):
|
if hasattr(os, "posix_fallocate"):
|
||||||
try:
|
try:
|
||||||
@@ -177,13 +209,16 @@ class StreamingSafeTensorFile:
|
|||||||
os.ftruncate(self._fd, total_size)
|
os.ftruncate(self._fd, total_size)
|
||||||
else: # pragma: no cover - Linux production path has posix_fallocate
|
else: # pragma: no cover - Linux production path has posix_fallocate
|
||||||
os.ftruncate(self._fd, total_size)
|
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)
|
_pwrite_all(self._fd, encoded_header, 8)
|
||||||
except BaseException:
|
except BaseException:
|
||||||
os.close(self._fd)
|
os.close(self._fd)
|
||||||
self._path.unlink(missing_ok=True)
|
self._path.unlink(missing_ok=True)
|
||||||
owner._forget_reserved_file(total_size)
|
owner._forget_reserved_file(total_size)
|
||||||
raise
|
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)
|
owner._opened(self)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
@@ -226,9 +261,19 @@ class StreamingSafeTensorFile:
|
|||||||
raise ArtifactWriterError(f"tensor {name!r} was not declared in this payload file") from exc
|
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:
|
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 = self._coverage[name]
|
||||||
coverage.add(byte_offset, byte_offset + len(data))
|
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:
|
def close(self) -> None:
|
||||||
if self._closed:
|
if self._closed:
|
||||||
@@ -238,11 +283,19 @@ class StreamingSafeTensorFile:
|
|||||||
self.abort()
|
self.abort()
|
||||||
raise ArtifactWriterError(f"tensor payload closed with incomplete writes: {missing}")
|
raise ArtifactWriterError(f"tensor payload closed with incomplete writes: {missing}")
|
||||||
self._closed = True
|
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:
|
try:
|
||||||
os.fsync(self._fd)
|
os.fchmod(self._fd, 0o444)
|
||||||
finally:
|
self._owner._closed(self, digests, self._fd)
|
||||||
|
except BaseException:
|
||||||
os.close(self._fd)
|
os.close(self._fd)
|
||||||
self._owner._closed(self)
|
raise
|
||||||
|
|
||||||
def abort(self) -> None:
|
def abort(self) -> None:
|
||||||
if self._closed:
|
if self._closed:
|
||||||
@@ -281,20 +334,25 @@ class FilesystemArtifactWriter:
|
|||||||
self._capacity_check = capacity_check
|
self._capacity_check = capacity_check
|
||||||
self._active: set[StreamingSafeTensorFile] = set()
|
self._active: set[StreamingSafeTensorFile] = set()
|
||||||
self._files: dict[str, tuple[TensorWriteSpec, ...]] = {}
|
self._files: dict[str, tuple[TensorWriteSpec, ...]] = {}
|
||||||
|
self._digests: dict[str, _PayloadDigests] = {}
|
||||||
self._reserved_bytes = 0
|
self._reserved_bytes = 0
|
||||||
self._finalized = False
|
self._finalized = False
|
||||||
|
self._fsync_executor: ThreadPoolExecutor | None = None
|
||||||
|
self._fsync_futures: list[Future[None]] = []
|
||||||
|
|
||||||
def open_tensor_file(
|
def open_tensor_file(
|
||||||
self,
|
self,
|
||||||
relative_name: str,
|
relative_name: str,
|
||||||
specs: tuple[TensorWriteSpec, ...],
|
specs: tuple[TensorWriteSpec, ...],
|
||||||
|
*,
|
||||||
|
ordered: bool = False,
|
||||||
) -> StreamingSafeTensorFile:
|
) -> StreamingSafeTensorFile:
|
||||||
if self._finalized:
|
if self._finalized:
|
||||||
raise ArtifactWriterError("cannot add payloads after artifact finalization")
|
raise ArtifactWriterError("cannot add payloads after artifact finalization")
|
||||||
relative_name = _safe_relative_name(relative_name)
|
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):
|
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")
|
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:
|
def _reserve_file(self, size: int) -> None:
|
||||||
self._capacity_check(self._reserved_bytes + size, size)
|
self._capacity_check(self._reserved_bytes + size, size)
|
||||||
@@ -306,9 +364,14 @@ class FilesystemArtifactWriter:
|
|||||||
def _opened(self, writer: StreamingSafeTensorFile) -> None:
|
def _opened(self, writer: StreamingSafeTensorFile) -> None:
|
||||||
self._active.add(writer)
|
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._active.remove(writer)
|
||||||
self._files[writer.relative_name] = writer.specs
|
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:
|
def _aborted(self, writer: StreamingSafeTensorFile, size: int) -> None:
|
||||||
self._active.discard(writer)
|
self._active.discard(writer)
|
||||||
@@ -318,6 +381,19 @@ class FilesystemArtifactWriter:
|
|||||||
for writer in tuple(self._active):
|
for writer in tuple(self._active):
|
||||||
with contextlib.suppress(Exception):
|
with contextlib.suppress(Exception):
|
||||||
writer.abort()
|
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(
|
def finalize(
|
||||||
self,
|
self,
|
||||||
@@ -335,46 +411,50 @@ class FilesystemArtifactWriter:
|
|||||||
if metadata.restorer_schema != identity.producer.restorer_schema:
|
if metadata.restorer_schema != identity.producer.restorer_schema:
|
||||||
raise ArtifactWriterError("production metadata does not match the identity's restorer schema")
|
raise ArtifactWriterError("production metadata does not match the identity's restorer schema")
|
||||||
|
|
||||||
|
self._wait_for_payload_fsync()
|
||||||
|
|
||||||
tensor_entries: list[TensorManifestEntry] = []
|
tensor_entries: list[TensorManifestEntry] = []
|
||||||
file_entries: list[FileManifestEntry] = []
|
file_entries: list[FileManifestEntry] = []
|
||||||
seen_tensors: set[str] = set()
|
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
|
path = self.staging_dir / file_name
|
||||||
os.chmod(path, 0o444)
|
specs = self._files[file_name]
|
||||||
with path.open("rb") as handle:
|
digests = self._digests[file_name]
|
||||||
os.fsync(handle.fileno())
|
|
||||||
file_entries.append(
|
file_entries.append(
|
||||||
FileManifestEntry(
|
FileManifestEntry(
|
||||||
file_name=file_name,
|
file_name=file_name,
|
||||||
size=path.stat().st_size,
|
size=path.stat().st_size,
|
||||||
sha256=_file_sha256(path),
|
sha256=digests.file,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
expected = {spec.name: spec for spec in self._files[file_name]}
|
for spec in specs:
|
||||||
with safe_open(str(path), framework="pt", device="cpu") as handle:
|
if spec.name in seen_tensors:
|
||||||
if set(handle.keys()) != set(expected):
|
raise ArtifactWriterError(f"artifact tensor {spec.name!r} occurs in multiple payload files")
|
||||||
raise ArtifactWriterError(f"published tensor keys differ in {file_name!r}")
|
seen_tensors.add(spec.name)
|
||||||
for name in sorted(expected):
|
tensor_entries.append(
|
||||||
if name in seen_tensors:
|
TensorManifestEntry(
|
||||||
raise ArtifactWriterError(f"artifact tensor {name!r} occurs in multiple payload files")
|
name=spec.name,
|
||||||
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,
|
|
||||||
kind=spec.kind,
|
kind=spec.kind,
|
||||||
role=spec.role,
|
role=spec.role,
|
||||||
file_name=file_name,
|
file_name=file_name,
|
||||||
storage_key=name,
|
storage_key=spec.name,
|
||||||
shape=tuple(tensor.shape),
|
shape=spec.shape,
|
||||||
dtype=str(tensor.dtype),
|
dtype=str(spec.dtype),
|
||||||
stride=tuple(tensor.stride()),
|
stride=tuple(torch.empty(spec.shape, dtype=spec.dtype, device="meta").stride()),
|
||||||
nbytes=tensor.numel() * tensor.element_size(),
|
nbytes=spec.nbytes,
|
||||||
sha256=_tensor_sha256(tensor),
|
sha256=digests.tensors[spec.name],
|
||||||
)
|
)
|
||||||
tensor_entries.append(entry)
|
)
|
||||||
|
|
||||||
ordered_tensor_entries = tuple(sorted(tensor_entries, key=lambda entry: entry.name))
|
ordered_tensor_entries = tuple(sorted(tensor_entries, key=lambda entry: entry.name))
|
||||||
content_digest = hashlib.sha256()
|
content_digest = hashlib.sha256()
|
||||||
|
|||||||
@@ -88,6 +88,8 @@ class ArtifactWriter(Protocol):
|
|||||||
self,
|
self,
|
||||||
relative_name: str,
|
relative_name: str,
|
||||||
specs: tuple[TensorWriteSpec, ...],
|
specs: tuple[TensorWriteSpec, ...],
|
||||||
|
*,
|
||||||
|
ordered: bool = False,
|
||||||
) -> TensorFileWriter: ...
|
) -> TensorFileWriter: ...
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -470,10 +470,14 @@ def resolve_quant_config_from_disk(
|
|||||||
)
|
)
|
||||||
return build_quant_config(qc_method, **qc_kwargs)
|
return build_quant_config(qc_method, **qc_kwargs)
|
||||||
|
|
||||||
# AutoRound MXFP8 checkpoints use data_type="mx_fp" instead of
|
# AutoRound MXFP checkpoints use data_type="mx_fp" instead of
|
||||||
# is_checkpoint_*_serialized; rebuild so the offline path is selected.
|
# 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":
|
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)
|
return build_quant_config(qc_method, **qc_kwargs)
|
||||||
|
|
||||||
if (
|
if (
|
||||||
|
|||||||
Reference in new issue
Block a user