Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 hsliuustc0106
|
1843f05ced |
No files matched your search
@@ -169,6 +169,7 @@ steps:
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
|
||||
@@ -83,7 +83,7 @@ steps:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and full_model"
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
@@ -375,7 +375,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion · Quantization Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pip install "vllm-gguf-plugin==0.0.4"
|
||||
- pip install "vllm-gguf-plugin>=0.0.3"
|
||||
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
|
||||
@@ -208,6 +208,7 @@ steps:
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
|
||||
@@ -74,19 +74,6 @@ steps:
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Omni · Perf Test · Async Chunk · Random"
|
||||
key: weekly-omni-performance-async-chunk-random
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and slow"
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Omni · Perf Test · Multi-Replica"
|
||||
key: weekly-omni-performance-multi-replicas
|
||||
timeout_in_minutes: 180
|
||||
|
||||
@@ -105,7 +105,7 @@ For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py`
|
||||
|
||||
## Path B: Custom/Private Repo Model
|
||||
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
|
||||
|
||||
### B1. Understand the reference repo
|
||||
|
||||
@@ -126,9 +126,9 @@ This is the key design decision for custom models. Follow these placement rules:
|
||||
| **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
|
||||
| **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
|
||||
| **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
|
||||
| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | original-repo package via git clone |
|
||||
| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone |
|
||||
| **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
|
||||
| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_<name>.py` |
|
||||
| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` |
|
||||
| **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
|
||||
|
||||
### B3. Handle external dependencies
|
||||
@@ -157,7 +157,7 @@ Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and ada
|
||||
|
||||
Custom models have two common patterns for weight loading:
|
||||
|
||||
**Pattern 1: Bypass standard loader** (eager custom init)
|
||||
**Pattern 1: Bypass standard loader** (DreamID-Omni style)
|
||||
|
||||
When the original model has complex custom init functions that load weights in `__init__`:
|
||||
|
||||
|
||||
@@ -16,19 +16,20 @@ vllm_omni/diffusion/models/wan2_2/
|
||||
|
||||
The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
|
||||
|
||||
### Custom model with external deps
|
||||
### Custom model with external deps (e.g., DreamID-Omni)
|
||||
|
||||
```
|
||||
vllm_omni/diffusion/models/<name>/
|
||||
vllm_omni/diffusion/models/dreamid_omni/
|
||||
├── __init__.py # Exports pipeline only
|
||||
├── pipeline_<name>.py # Pipeline: loads ALL weights in __init__ via custom helpers
|
||||
└── <backbone>.py # Custom transformer / fusion architecture
|
||||
├── pipeline_dreamid_omni.py # Pipeline: loads ALL weights in __init__ via custom helpers
|
||||
├── fusion.py # Custom fusion architecture (video + audio cross-attention)
|
||||
└── wan2_2.py # Re-implemented Wan backbone with split API
|
||||
|
||||
examples/offline_inference/<task>/
|
||||
└── download_<name>.py # Downloads weights and/or clones code repo
|
||||
examples/offline_inference/x_to_video_audio/
|
||||
└── download_dreamid_omni.py # Downloads weights from 3 HF repos + clones code repo
|
||||
```
|
||||
|
||||
All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency imported with try/except.
|
||||
All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
|
||||
|
||||
### Custom model with ported code (e.g., BAGEL)
|
||||
|
||||
@@ -75,7 +76,7 @@ self.weights_sources = [
|
||||
]
|
||||
```
|
||||
|
||||
### Pattern 3: Fully custom loading
|
||||
### Pattern 3: Fully custom loading (DreamID-Omni)
|
||||
|
||||
```
|
||||
init → load ALL weights eagerly via custom helpers → load_weights() = no-op
|
||||
@@ -107,8 +108,8 @@ Standard diffusers `model_index.json`:
|
||||
Custom model `model_index.json` (minimal):
|
||||
```json
|
||||
{
|
||||
"_class_name": "YourCustomPipeline",
|
||||
"fusion": "weights/model.safetensors"
|
||||
"_class_name": "DreamIDOmniPipeline",
|
||||
"fusion": "DreamID-Omni/dreamid_omni.safetensors"
|
||||
}
|
||||
```
|
||||
|
||||
@@ -116,7 +117,7 @@ The only **required** field is `_class_name` — it must match a key in `_DIFFUS
|
||||
|
||||
## External Dependency Management
|
||||
|
||||
### Git clone + .pth injection
|
||||
### Git clone + .pth injection (DreamID-Omni pattern)
|
||||
|
||||
```python
|
||||
def download_dependency():
|
||||
@@ -162,7 +163,7 @@ The engine checks `isinstance(pipeline, SupportImageInput)` at startup to config
|
||||
|
||||
Diffusers models use `config.json` in each subfolder. Custom models often use:
|
||||
|
||||
**Module-level config dicts**:
|
||||
**Module-level config dicts** (DreamID-Omni):
|
||||
```python
|
||||
VIDEO_CONFIG = {
|
||||
"patch_size": [1, 2, 2], "model_type": "ti2v",
|
||||
@@ -180,7 +181,7 @@ vae_cfg = bagel_cfg.get("vae_config", {})
|
||||
|
||||
## Custom Architecture Patterns
|
||||
|
||||
### Split forward API
|
||||
### Split forward API (DreamID-Omni)
|
||||
|
||||
When a fusion model needs to interleave blocks from two backbones:
|
||||
|
||||
|
||||
@@ -1,151 +0,0 @@
|
||||
"""Benchmark CosyVoice3-style CUDA stream handoff strategies."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import statistics
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def run_case(
|
||||
*,
|
||||
device: torch.device,
|
||||
steps: int,
|
||||
warmup_steps: int,
|
||||
producer_cycles: int,
|
||||
estimator_cycles: int,
|
||||
use_stream_dependencies: bool,
|
||||
) -> dict[str, float | int | str]:
|
||||
caller_stream = torch.cuda.Stream(device=device)
|
||||
estimator_stream = torch.cuda.Stream(device=device)
|
||||
input_tensor = torch.zeros(1024, device=device)
|
||||
output_tensor = torch.empty_like(input_tensor)
|
||||
submission_ms = []
|
||||
|
||||
def step(measure: bool) -> None:
|
||||
with torch.cuda.stream(caller_stream):
|
||||
torch.cuda._sleep(producer_cycles)
|
||||
input_tensor.add_(1)
|
||||
|
||||
start = time.perf_counter_ns()
|
||||
if use_stream_dependencies:
|
||||
estimator_stream.wait_stream(caller_stream)
|
||||
else:
|
||||
caller_stream.synchronize()
|
||||
|
||||
with torch.cuda.stream(estimator_stream):
|
||||
torch.cuda._sleep(estimator_cycles)
|
||||
output_tensor.copy_(input_tensor)
|
||||
|
||||
if use_stream_dependencies:
|
||||
caller_stream.wait_stream(estimator_stream)
|
||||
else:
|
||||
estimator_stream.synchronize()
|
||||
|
||||
if measure:
|
||||
submission_ms.append((time.perf_counter_ns() - start) / 1e6)
|
||||
|
||||
for _ in range(warmup_steps):
|
||||
step(measure=False)
|
||||
caller_stream.synchronize()
|
||||
estimator_stream.synchronize()
|
||||
input_tensor.zero_()
|
||||
torch.accelerator.reset_peak_memory_stats(device)
|
||||
initial_allocated = torch.accelerator.memory_allocated(device)
|
||||
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
wall_start = time.perf_counter()
|
||||
with torch.cuda.stream(caller_stream):
|
||||
start_event.record()
|
||||
for _ in range(steps):
|
||||
step(measure=True)
|
||||
with torch.cuda.stream(caller_stream):
|
||||
end_event.record()
|
||||
end_event.synchronize()
|
||||
wall_ms = (time.perf_counter() - wall_start) * 1e3
|
||||
peak_delta_bytes = torch.accelerator.max_memory_allocated(device) - initial_allocated
|
||||
|
||||
expected = torch.full_like(output_tensor, steps)
|
||||
torch.testing.assert_close(output_tensor, expected, rtol=0, atol=0)
|
||||
return {
|
||||
"mode": "wait_stream" if use_stream_dependencies else "synchronize",
|
||||
"steps": steps,
|
||||
"submission_median_ms": statistics.median(submission_ms),
|
||||
"submission_p95_ms": statistics.quantiles(submission_ms, n=20)[18],
|
||||
"submission_mean_ms": statistics.mean(submission_ms),
|
||||
"gpu_elapsed_ms": start_event.elapsed_time(end_event),
|
||||
"wall_elapsed_ms": wall_ms,
|
||||
"peak_allocation_delta_bytes": peak_delta_bytes,
|
||||
"final_value": int(output_tensor[0].item()),
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument("--steps", type=int, default=100)
|
||||
parser.add_argument("--warmup-steps", type=int, default=10)
|
||||
parser.add_argument("--repeats", type=int, default=7)
|
||||
parser.add_argument("--producer-cycles", type=int, default=1_000_000)
|
||||
parser.add_argument("--estimator-cycles", type=int, default=2_000_000)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
device = torch.device(args.device)
|
||||
current_omni_platform.set_device(device)
|
||||
torch.manual_seed(0)
|
||||
|
||||
old_results = []
|
||||
new_results = []
|
||||
for repeat in range(args.repeats):
|
||||
common_args = {
|
||||
"device": device,
|
||||
"steps": args.steps,
|
||||
"warmup_steps": args.warmup_steps,
|
||||
"producer_cycles": args.producer_cycles,
|
||||
"estimator_cycles": args.estimator_cycles,
|
||||
}
|
||||
if repeat % 2 == 0:
|
||||
old_results.append(run_case(**common_args, use_stream_dependencies=False))
|
||||
new_results.append(run_case(**common_args, use_stream_dependencies=True))
|
||||
else:
|
||||
new_results.append(run_case(**common_args, use_stream_dependencies=True))
|
||||
old_results.append(run_case(**common_args, use_stream_dependencies=False))
|
||||
|
||||
old_submission_ms = statistics.median(result["submission_median_ms"] for result in old_results)
|
||||
new_submission_ms = statistics.median(result["submission_median_ms"] for result in new_results)
|
||||
old_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in old_results)
|
||||
new_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in new_results)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"device": torch.cuda.get_device_name(device),
|
||||
"torch": torch.__version__,
|
||||
"cuda": torch.version.cuda,
|
||||
"repeats": args.repeats,
|
||||
"warmup_steps": args.warmup_steps,
|
||||
"producer_cycles": args.producer_cycles,
|
||||
"estimator_cycles": args.estimator_cycles,
|
||||
"old_submission_median_ms": old_submission_ms,
|
||||
"new_submission_median_ms": new_submission_ms,
|
||||
"submission_median_speedup": old_submission_ms / new_submission_ms,
|
||||
"old_gpu_elapsed_median_ms": old_gpu_ms,
|
||||
"new_gpu_elapsed_median_ms": new_gpu_ms,
|
||||
"gpu_elapsed_change_percent": (new_gpu_ms - old_gpu_ms) / old_gpu_ms * 100,
|
||||
"parity": all(result["final_value"] == args.steps for result in old_results + new_results),
|
||||
"old_runs": old_results,
|
||||
"new_runs": new_results,
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+6
-21
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# ruff: noqa
|
||||
# code borrowed from https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py
|
||||
@@ -494,16 +494,6 @@ def is_xnnpack_available():
|
||||
|
||||
|
||||
def get_env_vars():
|
||||
try:
|
||||
from vllm_omni.config.environment_variable_inventory import ENVIRONMENT_VARIABLE_INVENTORY
|
||||
except Exception:
|
||||
# Keep environment collection usable when an incomplete/broken Omni
|
||||
# installation is the reason the user is running this diagnostic.
|
||||
# Importing an Omni submodule executes package initialization, which
|
||||
# may fail with RuntimeError, AssertionError, or OSError as well as
|
||||
# ImportError when vLLM or the accelerator stack is broken.
|
||||
ENVIRONMENT_VARIABLE_INVENTORY = {}
|
||||
|
||||
env_vars = ""
|
||||
secret_terms = ("secret", "token", "api", "access", "password")
|
||||
report_prefix = (
|
||||
@@ -517,17 +507,12 @@ def get_env_vars():
|
||||
"MKL_",
|
||||
"NVIDIA",
|
||||
)
|
||||
for k, v in sorted(os.environ.items()):
|
||||
omni_classification = ENVIRONMENT_VARIABLE_INVENTORY.get(k)
|
||||
if (omni_classification is not None and omni_classification.redact_value) or any(
|
||||
term in k.lower() for term in secret_terms
|
||||
):
|
||||
for k, v in os.environ.items():
|
||||
if any(term in k.lower() for term in secret_terms):
|
||||
continue
|
||||
if (
|
||||
k in environment_variables
|
||||
or (omni_classification is not None and omni_classification.is_public_omni)
|
||||
or k.startswith(report_prefix)
|
||||
):
|
||||
if k in environment_variables:
|
||||
env_vars = env_vars + "{}={}".format(k, v) + "\n"
|
||||
if k.startswith(report_prefix):
|
||||
env_vars = env_vars + "{}={}".format(k, v) + "\n"
|
||||
|
||||
return env_vars
|
||||
|
||||
@@ -24,6 +24,7 @@ nav:
|
||||
- Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
|
||||
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
|
||||
- X-To-Text: user_guide/examples/offline_inference/x_to_text.md
|
||||
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
|
||||
- Online Serving:
|
||||
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
|
||||
|
||||
+1
-7
@@ -2,11 +2,6 @@
|
||||
|
||||
The CLI for vLLM-Omni inherits from vllm with some additional arguments.
|
||||
|
||||
Environment variables are process configuration rather than CLI-only options.
|
||||
See the [Environment Variables](../configuration/environment_variables.md)
|
||||
reference for ownership, precedence, evaluation time, and per-stage `env`
|
||||
scoping.
|
||||
|
||||
## serve
|
||||
|
||||
Starts the vLLM-Omni OpenAI Compatible API server.
|
||||
@@ -44,5 +39,4 @@ vllm bench serve --omni \
|
||||
--num-prompts 5
|
||||
```
|
||||
|
||||
See [vllm bench serve](./bench/serve.md) for serving benchmark arguments and
|
||||
dataset-specific examples, including OmniInteract native-duplex sessions.
|
||||
See [vllm bench serve](./bench/serve.md) for the full reference of all available arguments.
|
||||
+3
-41
@@ -2,7 +2,7 @@
|
||||
The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
|
||||
|
||||
## Notes
|
||||
vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, `daily-omni`, and `openai-realtime-duplex` serving benchmark backends. It also adds the `omniinteract` dataset.
|
||||
vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` serving benchmark backends.
|
||||
|
||||
## Basic Parameter Description
|
||||
You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
|
||||
@@ -10,7 +10,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
|
||||
Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
|
||||
|
||||
- `--backend`
|
||||
Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, `daily-omni`, and `openai-realtime-duplex` to the upstream vLLM backend choices.
|
||||
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.
|
||||
|
||||
- `--model`
|
||||
The model identifier to load, filled according to the models supported by vLLM-Omni.
|
||||
@@ -19,7 +19,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
|
||||
The API endpoint exposed externally, to which clients send their requests.
|
||||
|
||||
- `--dataset-name`
|
||||
The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio), while `omniinteract` replays official OmniInteract videos.
|
||||
The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
|
||||
|
||||
- `--num-prompts`
|
||||
The total number of requests to send, an integer.
|
||||
@@ -269,44 +269,6 @@ We use audio generation time / audio duration to calculate RTF.
|
||||
|
||||
</details>
|
||||
|
||||
### OmniInteract Realtime Benchmark
|
||||
|
||||
OmniInteract runs each video as one native-duplex WebSocket sample in the standard serving-benchmark lifecycle:
|
||||
|
||||
```bash
|
||||
vllm bench serve --omni \
|
||||
--backend openai-realtime-duplex \
|
||||
--dataset-name omniinteract \
|
||||
--dataset-path /path/to/OmniInteract \
|
||||
--model openbmb/MiniCPM-o-4_5 \
|
||||
--base-url http://127.0.0.1:8000 \
|
||||
--endpoint /v1/realtime \
|
||||
--omniinteract-ref-audio /path/to/reference.wav \
|
||||
--omniinteract-output-dir ./omniinteract-artifacts \
|
||||
--num-prompts 3 \
|
||||
--num-warmups 0
|
||||
```
|
||||
|
||||
`--dataset-path` accepts an extracted directory, `data.tar[.gz]`, or a Hugging Face dataset ID; omitting it uses
|
||||
`lucky-lance/OmniInteract`. `--num-prompts` is the total across subsets and defaults to 3 for OmniInteract; explicit `0`
|
||||
selects all and oversize values use all available cases. Reference audio is required, and OmniInteract uses the
|
||||
`/v1/realtime` endpoint.
|
||||
|
||||
Audio is replayed as 16 kHz PCM16 in 200 ms chunks and video at 1 FPS with real-time pacing. All selected media is decoded
|
||||
before timing and remains in client memory for the run, so `--max-concurrency` does not limit media preparation memory; use
|
||||
explicit `--num-prompts 0` only when the client has enough RAM for the full dataset. Media commands are bounded by
|
||||
`--omniinteract-media-timeout-s`, and concurrency defaults to 1. Standard request-rate, warmup, result-saving, and summary
|
||||
options apply. Use `--omniinteract-require-response` only for functional E2E cases; LISTEN is a valid benchmark result.
|
||||
|
||||
Each completed case writes `output.wav`, `wav_transcript.json`, `events.json`, `result.json`, and a final `.done` marker under
|
||||
`--omniinteract-output-dir`. The root also contains `batch_summary.json` and `official_eval_manifest.jsonl`; failed cases write
|
||||
`.failed.json`. Runs sharing one output root are serialized. Completion validates transport, response lifecycle, and artifacts,
|
||||
not answer accuracy. Transcript timestamps are serialized playback-queue times. Clipped or cancelled outputs are ineligible and
|
||||
omitted from the official manifest; `audio_clipped_bytes` records output beyond the rounded video horizon.
|
||||
|
||||
TTFT, TTFP, and RTF start at client receipt of `response.created`. TPOT/ITL use engine stage-0 timing; ITL is emitted only when
|
||||
every token interval is present.
|
||||
|
||||
### Multi-Modal Benchmark
|
||||
|
||||
<details class="admonition abstract" markdown="1">
|
||||
|
||||
@@ -7,9 +7,6 @@ For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration g
|
||||
Each model defines fixed topology in a registered `PipelineConfig` and runtime
|
||||
overrides in a deploy YAML.
|
||||
|
||||
For process-level settings shared by the CLI, deploy YAML, workers, and model
|
||||
integrations, see [Environment Variables](./environment_variables.md).
|
||||
|
||||
For a specific example, see the [Qwen2.5-Omni deploy config](gh-file:vllm_omni/deploy/qwen2_5_omni.yaml). The matching frozen pipeline topology lives at [vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py](gh-file:vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py).
|
||||
|
||||
For an introduction, see [Pipeline and deploy configurations](./stage_configs.md).
|
||||
|
||||
@@ -1,237 +0,0 @@
|
||||
# Environment Variables
|
||||
|
||||
Environment variables are process-level configuration shared by the CLI,
|
||||
deploy YAML, Python APIs, worker processes, and model integrations. They are
|
||||
therefore documented under **User Guide → Configuration**, rather than under
|
||||
CLI Reference.
|
||||
|
||||
This page describes the public vLLM-Omni environment-variable contract on the
|
||||
`latest` documentation branch. Use the documentation for your installed
|
||||
vLLM-Omni release when behavior differs.
|
||||
|
||||
!!! important "Public contract"
|
||||
Only names in [Public vLLM-Omni variables](#public-vllm-omni-variables) are
|
||||
public Omni configuration. A name found in source code is not automatically
|
||||
supported: model-specific, benchmark, platform, and internal variables have
|
||||
separate ownership and lifecycle rules.
|
||||
|
||||
## Naming and implementation boundary
|
||||
|
||||
New public variables owned by vLLM-Omni must use the `VLLM_OMNI_` prefix.
|
||||
Existing public names with older prefixes remain supported for compatibility,
|
||||
but they do not establish naming precedent for new settings. Inherited vLLM
|
||||
and third-party names retain the prefix chosen by their owning project.
|
||||
|
||||
The source [environment-variable inventory](gh-file:vllm_omni/config/environment_variable_inventory.py)
|
||||
records ownership, migration disposition, and redaction metadata. Unlike
|
||||
`vllm.envs`, it is not an executable value resolver: existing consumers still
|
||||
own parsing, defaults, precedence, and evaluation time. The public tables below
|
||||
describe those live consumer contracts. Stable model or stage behavior should
|
||||
move to typed configuration, and request-varying behavior should move to a
|
||||
request schema instead of adding another environment variable.
|
||||
|
||||
## Precedence and evaluation time
|
||||
|
||||
There is no repository-wide rule that makes environment variables override (or
|
||||
lose to) every CLI, YAML, or request option. The tables below state precedence
|
||||
for each public setting. In general:
|
||||
|
||||
- request fields should control behavior that can vary per request;
|
||||
- typed deploy or stage configuration should control stable model and stage
|
||||
behavior;
|
||||
- environment variables are appropriate for process launch, platform
|
||||
integration, emergency diagnostics, and compatibility fallbacks.
|
||||
|
||||
An environment variable can be read while a module is imported, when server or
|
||||
model configuration is constructed, when a worker is spawned, or whenever a
|
||||
property is accessed. Set import- and startup-time variables before invoking
|
||||
`vllm serve` or importing vLLM-Omni. Changing the parent shell after workers
|
||||
start does not update existing worker environments.
|
||||
|
||||
The lifecycle labels used below are:
|
||||
|
||||
- **Stable** — reviewed public configuration; changes require compatibility
|
||||
consideration.
|
||||
- **Experimental** — public for evaluation, but its contract may still evolve.
|
||||
- **Diagnostic** — supported for investigation or temporary mitigation, not as
|
||||
a normal tuning interface.
|
||||
- **Deprecated** — compatibility alias; migrate to the named replacement.
|
||||
|
||||
## Public vLLM-Omni variables
|
||||
|
||||
### Diffusion backend and cache selection
|
||||
|
||||
| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `DIFFUSION_ATTENTION_BACKEND` | Backend name or `auto`; default `auto` (platform selection) | Diffusion stages; read when `OmniDiffusionConfig` is constructed | `diffusion_attention_config.default` wins. An unknown backend fails during backend resolution. | Stable fallback |
|
||||
| `DIFFUSION_CACHE_BACKEND` | `none`, `cache_dit`, `tea_cache`, `mag_cache`, `step_cache`, `stepcache`, or `step_cache_dit`; default `none` | Diffusion runner startup | Explicit `cache_backend` in config wins. Otherwise this name wins over the deprecated alias. Unsupported values raise `ValueError` during runner setup. | Stable fallback |
|
||||
| `DIFFUSION_CACHE_ADAPTER` | Same values as `DIFFUSION_CACHE_BACKEND`; default `none` | Diffusion runner startup | Used only when neither explicit `cache_backend` nor `DIFFUSION_CACHE_BACKEND` is set. Unsupported values raise `ValueError`. | Deprecated; use `DIFFUSION_CACHE_BACKEND` |
|
||||
| `OMNI_DIFFUSION_PROMPT_EMBED_CACHE` | Boolean: `1`, `true`, `yes`, `on`, `0`, `false`, `no`, or `off`; default disabled | Each diffusion runner; resolved during model setup | A recognized environment value overrides the explicit enable setting. An unrecognized value is ignored. | Experimental |
|
||||
| `OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE` | Positive integer; default `32` entries | Each diffusion runner; resolved during model setup | A valid environment value overrides the explicit cache size. A non-integer logs a warning; a non-positive value is ignored. | Experimental |
|
||||
| `OMNI_DIFFUSION_SESSION_STATE_MANAGER` | Boolean with the same accepted spellings as the prompt cache; default disabled | Experimental diffusion session manager; model setup | A recognized environment value overrides the explicit enable setting. An unrecognized value is ignored. | Experimental |
|
||||
| `OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS` | Positive integer; default `64` | Experimental diffusion session manager; model setup | A valid environment value overrides the explicit maximum. A non-integer logs a warning; a non-positive value is ignored. | Experimental |
|
||||
|
||||
Backend names for `DIFFUSION_ATTENTION_BACKEND` are the members of
|
||||
`DiffusionAttentionBackendEnum`, such as `FLASH_ATTN`, `TORCH_SDPA`,
|
||||
`SAGE_ATTN`, `FLASHINFER_ATTN`, and `TRTLLM_ATTN`. Platform support still
|
||||
depends on the installed kernels and model path.
|
||||
|
||||
### Serving and runtime
|
||||
|
||||
| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable |
|
||||
| `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
|
||||
| `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
|
||||
| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
|
||||
| `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic |
|
||||
| `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
|
||||
| `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental |
|
||||
| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental |
|
||||
|
||||
### Server storage
|
||||
|
||||
Storage names use Pydantic's nested-settings delimiter (`__`). They are read
|
||||
when `ServerSettings` is initialized, normally during server import/startup.
|
||||
Invalid integer values raise a Pydantic validation error.
|
||||
|
||||
| Name | Type and default | Meaning | Lifecycle |
|
||||
| --- | --- | --- | --- |
|
||||
| `VLLM_OMNI_SERVER_STORAGE__PATH` | Path; default `/tmp/storage` | Directory for completed server files. | Stable |
|
||||
| `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY` | Integer; default `4` | Maximum concurrent file operations. No additional range validation is currently applied. | Stable |
|
||||
| `VLLM_OMNI_SERVER_STORAGE__FILE_TTL` | Integer seconds or unset; default unset | Optional lifetime for locally stored files. | Stable |
|
||||
| `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL` | Integer seconds or unset | Sweep frequency. If file TTL is set and this is unset, the effective default is `300`. | Stable |
|
||||
| `VLLM_OMNI_STORAGE_PATH` | Path | Deprecated alias for `VLLM_OMNI_SERVER_STORAGE__PATH`. The new name wins when both are set. | Deprecated |
|
||||
| `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` | Integer | Deprecated alias for `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`. The new name wins when both are set. | Deprecated |
|
||||
|
||||
The deprecated aliases emit `DeprecationWarning`. They do not affect file TTL
|
||||
settings.
|
||||
|
||||
### Quantization diagnostics and performance
|
||||
|
||||
| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary |
|
||||
| `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control |
|
||||
|
||||
`QUACK_CACHE_DIR` is owned by the external quack library and is not an
|
||||
Omni-owned variable, even though vLLM-Omni supplies a persistent default when
|
||||
it is unset.
|
||||
|
||||
## Per-stage environment
|
||||
|
||||
Deploy configurations can set arbitrary environment keys for one stage:
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
env:
|
||||
HCCL_IF_BASE_PORT: 23000
|
||||
MY_MODEL_CACHE: /mnt/model-cache
|
||||
```
|
||||
|
||||
The legacy nested form under `runtime.env` is also accepted. vLLM-Omni converts
|
||||
keys and values to strings, applies them while the stage engine or child process
|
||||
is launched, logs keys but not values, and restores the parent process's prior
|
||||
environment afterward. A child process inherits the scoped values. This keeps
|
||||
sibling stages from permanently changing one another's parent environment.
|
||||
|
||||
This scoping has two consequences:
|
||||
|
||||
1. A setting read inside stage/model initialization can see the stage value.
|
||||
2. A setting already cached by an import that occurred before stage launch
|
||||
cannot be changed retroactively with stage `env`.
|
||||
|
||||
Stage `env` is an escape hatch for process and third-party settings, not a
|
||||
replacement for typed stage configuration. Do not put request-varying behavior
|
||||
there. Secret values may reach child processes; logs and diagnostics must emit
|
||||
their keys only.
|
||||
|
||||
## Inherited vLLM variables
|
||||
|
||||
vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
|
||||
the [vLLM 0.26 environment-variable reference](https://docs.vllm.ai/en/v0.26.0/configuration/env_vars.html)
|
||||
for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
|
||||
XPU, ModelScope, and FlashInfer workspace settings.
|
||||
|
||||
An inherited name does not become Omni-owned because Omni reads or forwards it.
|
||||
Where behavior differs, the Omni-specific page should document only that
|
||||
difference. For example, device-visibility values can be further narrowed by a
|
||||
stage's `devices` setting.
|
||||
|
||||
## Platform, external, and secret variables
|
||||
|
||||
Hardware runtimes, distributed launchers, and third-party libraries retain
|
||||
ownership of their variables. Common groups include:
|
||||
|
||||
- device visibility (`ASCEND_RT_VISIBLE_DEVICES`, `HIP_VISIBLE_DEVICES`, and
|
||||
`MUSA_VISIBLE_DEVICES`);
|
||||
- distributed launch (`MASTER_ADDR`, `MASTER_PORT`, `RANK`, and `WORLD_SIZE`);
|
||||
- Hugging Face cache and authentication (`HF_HOME`, `HF_TOKEN`, and related
|
||||
names);
|
||||
- accelerator libraries such as NCCL, FlashInfer, MindIE, MORI, and quack.
|
||||
|
||||
Consult the owning project's documentation for accepted values and support.
|
||||
`HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, and `OPENAI_API_KEY` are explicitly marked
|
||||
for redaction in the reviewed inventory. Never print their values in environment
|
||||
collection, examples, bug reports, or logs.
|
||||
|
||||
## Model-specific variables
|
||||
|
||||
The audit found 56 variables read by a single model or pipeline family. They are
|
||||
not listed as public usage options here because doing so would turn implementation
|
||||
escape hatches into an accidental compatibility contract.
|
||||
|
||||
Every audited model-specific name has a migration disposition in the
|
||||
[environment-variable inventory](gh-file:vllm_omni/config/environment_variable_inventory.py):
|
||||
|
||||
| Disposition | Count | Required outcome |
|
||||
| --- | ---: | --- |
|
||||
| Promote | 33 | Move a stable setting into typed stage or model configuration. |
|
||||
| Request scope | 6 | Move request-varying behavior into a declared request-option schema. |
|
||||
| External | 0 | Retain only when a supported third-party library owns the contract. |
|
||||
| Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. |
|
||||
| Deprecate/remove | 6 | Remove a compatibility escape hatch that has no continuing contract. |
|
||||
|
||||
The disposition is a migration target, not a statement that the existing
|
||||
environment switch is stable. Promote or request-scope work should land in
|
||||
model-owner-reviewed follow-up changes with validation and model-page updates.
|
||||
The family-sized migration packages and their ownership gate are tracked in
|
||||
[#6232](https://github.com/vllm-project/vllm-omni/issues/6232).
|
||||
|
||||
## Environment collection and troubleshooting
|
||||
|
||||
Run the repository collector when preparing a bug report:
|
||||
|
||||
```bash
|
||||
python collect_env.py
|
||||
```
|
||||
|
||||
The environment section includes the safe public Omni allowlist, registered
|
||||
vLLM variables, and selected platform prefixes. It omits names marked for
|
||||
redaction and names containing common secret terms. Review the output before
|
||||
sharing it. A key supplied through stage `env` is not collected merely because
|
||||
it appears in deploy configuration; the same allowlist and prefix rules apply.
|
||||
|
||||
If a setting appears ineffective, first check its read time in the tables
|
||||
above. Restart processes after changing import- or startup-time values, verify
|
||||
that the variable reaches the intended worker, and check whether an explicit
|
||||
typed configuration field takes precedence.
|
||||
|
||||
## Benchmark and internal variables
|
||||
|
||||
The reviewed inventory marks 20 benchmark variables as transitional while
|
||||
benchmark configuration moves to CLI/request schemas in
|
||||
[issue #5376](https://github.com/vllm-project/vllm-omni/issues/5376). They are not
|
||||
part of the serving/runtime contract.
|
||||
|
||||
Internal process-coordination and debug variables are deliberately excluded
|
||||
from this page. The CI scanner resolves string literals and module-level string
|
||||
constants used by direct `os` access, local wrapper helpers, aliased `os`
|
||||
imports, and environment membership checks. Server-storage names generated by
|
||||
Pydantic settings are checked separately against their model fields.
|
||||
|
||||
Names assembled from cross-module metadata and arbitrary keys supplied through
|
||||
stage `env` cannot be inferred reliably. Contributors must classify those names
|
||||
explicitly. CI fails when a statically resolved name has no classification or
|
||||
when a reviewed public name is absent from this page.
|
||||
@@ -14,7 +14,7 @@ describe fixed model topology and are not accepted in deploy YAMLs.
|
||||
Common `PipelineConfig` fields include:
|
||||
|
||||
| Field | Description |
|
||||
| ------- | ------------- |
|
||||
|-------|-------------|
|
||||
| `model_type` | Pipeline identifier used during model and config resolution. |
|
||||
| `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. |
|
||||
| `model_arch` | Default Hugging Face architecture for the pipeline. |
|
||||
@@ -26,7 +26,7 @@ Common `PipelineConfig` fields include:
|
||||
Common `StagePipelineConfig` fields include:
|
||||
|
||||
| Field | Description |
|
||||
| ------- | ------------- |
|
||||
|-------|-------------|
|
||||
| `stage_id` | Stable stage identifier. |
|
||||
| `model_stage` | Logical stage name used by runtime and strategy resolution. |
|
||||
| `execution_type` | `LLM_AR`, `LLM_GENERATION`, or `DIFFUSION`. |
|
||||
@@ -49,10 +49,10 @@ deployment overrides.
|
||||
The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
|
||||
|
||||
| Field | Type | Required | Default | Description |
|
||||
| ------- | ------ | ---------- | --------- | ------------- |
|
||||
|-------|------|----------|---------|-------------|
|
||||
| `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. |
|
||||
| `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. |
|
||||
| `session_mode` | str | optional | `"turn"` | Session behavior. MiniCPM-o 4.5 deploy YAMLs set `"duplex"` so `/v1/realtime` and chat share the same profile. |
|
||||
| `session_mode` | str | optional | `"turn"` | Session behavior. Use `"duplex"` only with a pipeline that enables duplex control. |
|
||||
| `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. |
|
||||
| `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. |
|
||||
| `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). |
|
||||
@@ -84,7 +84,7 @@ subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet.
|
||||
Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine. Frequently used fields are listed below; the source-of-truth schema is `StageDeployConfig` in `vllm_omni/config/stage_config.py`.
|
||||
|
||||
| Field | Type | Required | Default | Description |
|
||||
| ------- | ------ | ---------- | --------- | ------------- |
|
||||
|-------|------|----------|---------|-------------|
|
||||
| `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
|
||||
| `max_num_seqs` | int \| null | optional | `null` | Max concurrent sequences per stage. |
|
||||
| `gpu_memory_utilization` | float \| null | optional | `null` | Per-stage total memory target; used for automatic KV-cache sizing. |
|
||||
@@ -113,9 +113,9 @@ connectors:
|
||||
# Additional connector-specific options
|
||||
```
|
||||
|
||||
| Connector class | Use case | `extra` keys |
|
||||
|--------------------------|-----------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------|
|
||||
| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. |
|
||||
| Connector class | Use case | `extra` keys |
|
||||
|-----------------|----------|--------------|
|
||||
| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. |
|
||||
| `MooncakeStoreConnector` | Cross-host KV transfer over TCP. Required for multi-node deployments. | `host`, `metadata_server`, `master`, `segment` (int bytes), `localbuf` (int bytes), `proto` (`"tcp"` / `"rdma"`). |
|
||||
|
||||
A stage references a connector by name in its `input_connectors` / `output_connectors`:
|
||||
@@ -135,7 +135,7 @@ stages:
|
||||
### CLI flags
|
||||
|
||||
| Flag | Description |
|
||||
| ------ | ------------- |
|
||||
|------|-------------|
|
||||
| `--deploy-config PATH` | Load a deploy YAML. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
|
||||
| `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
|
||||
| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
|
||||
@@ -249,7 +249,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
Effective config per stage after the merge:
|
||||
|
||||
| Stage | Field | Final value | Source |
|
||||
| ------- | ------- | ------------- | -------- |
|
||||
|-------|-------|-------------|--------|
|
||||
| 0 | `gpu_memory_utilization` | `0.9` | base YAML (overlay didn't touch stage 0) |
|
||||
| 0 | `max_num_seqs` | `8` | per-stage CLI (`--stage-overrides`) — wins over base `32` |
|
||||
| 0 | `max_model_len` | `16384` | global CLI |
|
||||
@@ -273,14 +273,12 @@ when loading a deploy configuration.
|
||||
Examples:
|
||||
|
||||
For offline inference (assuming the necessary dependencies have been imported):
|
||||
|
||||
```python
|
||||
model_name = "Qwen/Qwen2.5-Omni-7B"
|
||||
omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml")
|
||||
```
|
||||
|
||||
For online serving:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
@@ -4,8 +4,8 @@ When you want to add L4-level ***performance test*** cases, add entries to JSON
|
||||
|
||||
| Model type | Runner | Example JSON files |
|
||||
| ---------- | ------ | ------------------ |
|
||||
| Omni (nightly) | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json` (`full_model` without `slow` in `mark`) |
|
||||
| Omni (weekly) | `run_benchmark.py` | `test_qwen3_omni_async_chunk.json` (CUDA only), `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` (`slow` in `mark`; **Perf Test** in `test-weekly.yml`) |
|
||||
| Omni (nightly) | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json` (`full_model` in `mark`) |
|
||||
| Omni (weekly) | `run_benchmark.py` | `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` (`slow` in `mark`; **Perf Test** in `test-weekly.yml`) |
|
||||
| TTS | `run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` |
|
||||
| Diffusion | `run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`, `test_wan22_i2v_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, … |
|
||||
|
||||
@@ -20,7 +20,7 @@ Diffusion cases are detected when the JSON has `server_type` (typically `"vllm-o
|
||||
|
||||
#### Running perf cases
|
||||
|
||||
Pass **`--test-config-file`** to 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.
|
||||
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).
|
||||
|
||||
```JSON
|
||||
{
|
||||
@@ -75,7 +75,7 @@ Omit `mark` only for configs not meant to be filtered by `-m`. `server_type` is
|
||||
|
||||
Optional top-level field on each perf JSON **case object** (one per `test_name`). `run_benchmark.py` and `run_diffusion_benchmark.py` read it via `tests.dfx.conftest.resolve_pytest_marks` and attach the marks to the corresponding `pytest.param`, so you can filter locally with `-m`.
|
||||
|
||||
When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `slow`, `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`, `omni`, `tts`, `diffusion`, or `local_model`.
|
||||
|
||||
Supported form:
|
||||
|
||||
@@ -142,18 +142,11 @@ Examples:
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and H100 and diffusion"
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100"
|
||||
|
||||
# Single file (same selectors as the CI Perf steps)
|
||||
# Single file (same as nightly CI Perf steps)
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
|
||||
--test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
|
||||
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \
|
||||
-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"
|
||||
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
|
||||
```
|
||||
|
||||
<!-- markdownlint-disable-next-line MD051 -->
|
||||
|
||||
@@ -63,7 +63,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
|
||||
|
||||
### N-Branch CFG (3+ branches)
|
||||
|
||||
Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches.
|
||||
Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches.
|
||||
|
||||
`predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`):
|
||||
|
||||
|
||||
@@ -14,12 +14,12 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti
|
||||
|
||||
## Status
|
||||
|
||||
The implementation provides contracts and a CPU local-filesystem store. The
|
||||
diffusion consumer additionally defines typed,
|
||||
The first implementation provides contracts and a CPU local-filesystem store.
|
||||
The initial diffusion consumer contract additionally defines typed,
|
||||
representation-independent final-layout identity/restoration mechanics plus a
|
||||
concrete BF16-with-preserved-FP32 policy for MiniMax H3 and
|
||||
`black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration
|
||||
selects, publishes, restores, and transfers these artifacts.
|
||||
concrete BF16-with-preserved-FP32 policy for MiniMax H3. It remains
|
||||
library-only: no loader or DLO path selects, publishes, restores, or transports
|
||||
that artifact yet.
|
||||
|
||||
V1 includes:
|
||||
|
||||
@@ -33,12 +33,13 @@ V1 includes:
|
||||
|
||||
V1 does not include:
|
||||
|
||||
- default-on activation or consumers outside no-AllGather DLO;
|
||||
- online FP8, quantized, merged-adaptation, or additional model producers;
|
||||
- HWR interaction with DLO AllGather;
|
||||
- a public CLI, loader activation, or default-on model integration;
|
||||
- lease handoff to DLO or another transport consumer;
|
||||
- FP8, quantized, merged-adaptation, or additional model producers;
|
||||
- CUDA registration, pinned staging, H2D scheduling, or GPU kernels;
|
||||
- a remote artifact provider or cross-node coordination;
|
||||
- automatic eviction; or
|
||||
- a change to DLO collective or execution behavior.
|
||||
- a change to DLO AllGather or no-AllGather behavior.
|
||||
|
||||
## Motivation and use cases
|
||||
|
||||
@@ -254,8 +255,7 @@ The contract is intentionally separate from loader activation:
|
||||
- `FinalLayoutBF16Producer` accepts only the matching identity context and a
|
||||
finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP
|
||||
coordinate. Its BF16 policy preserves model-declared FP32 parameters and
|
||||
buffers, revalidates MiniMax H3 mixed-precision invariants, and revalidates
|
||||
FLUX.2-klein's two block stacks, packed QKV mapping, and BF16 base layout.
|
||||
buffers and revalidates MiniMax H3 mixed-precision invariants.
|
||||
|
||||
Other representations reuse source identity, typed parallel identity, tensor
|
||||
ownership, and exact restoration only when their policy proves those semantics.
|
||||
|
||||
@@ -196,15 +196,16 @@ adopt it.
|
||||
|
||||
### Final-layout Host Weight Runtime consumer
|
||||
|
||||
> **Status:** PR2 landed in
|
||||
> [PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). PR3 adds
|
||||
> registered direct H2D plus generic HWR publication/source-digest
|
||||
> optimizations without changing AllGather behavior.
|
||||
> **Status:** PR2 implementation. The representation-independent identity,
|
||||
> producer, and restoration contracts landed in
|
||||
> [PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445); this consumer
|
||||
> wires the final-layout BF16 path into eligible no-AllGather DLO.
|
||||
|
||||
Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
|
||||
only to no-AllGather DLO. Enable it with
|
||||
`--host-weight-runtime-mode preferred` (or `required`) and
|
||||
`--host-weight-runtime-root <node-local-root>`.
|
||||
`--host-weight-runtime-root <node-local-root>`. Registration of shared mappings
|
||||
and direct asynchronous H2D remain a separate PR3 transport change.
|
||||
|
||||
The modes express operator fallback policy, not different artifact formats:
|
||||
|
||||
@@ -220,15 +221,9 @@ model revision and parallel layout with `required`. TP coordinates have distinct
|
||||
identities, while equivalent DP replicas share them.
|
||||
|
||||
The current producer/consumer boundary is the model-declared final-layout BF16
|
||||
contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein
|
||||
contract covers both transformer block stacks, constructor-stable packed QKV
|
||||
mapping state, BF16 parameters, and any persistent loader-owned `beta`/`eps`
|
||||
buffers. The shared HWR machinery supports ordinary-loader final layouts for
|
||||
TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence
|
||||
in this change is TP1-only; its TP2/SP layouts rely on the existing
|
||||
per-coordinate identity mechanics and remain unmeasured. Online quantization,
|
||||
HSDP, LoRA/adapted weights, and non-default load formats remain ineligible.
|
||||
HWR mode
|
||||
contract (MiniMax H3 today). It supports ordinary-loader final layouts for TP1
|
||||
and TP2 rank identities plus SP layout identities; online quantization, HSDP,
|
||||
LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode
|
||||
`disabled`, DLO-disabled, and DLO AllGather configurations stop before source
|
||||
identity or store construction and retain the existing checkpoint-mmap or
|
||||
ordinary-loader path.
|
||||
@@ -259,44 +254,6 @@ outcomes are authoritative:
|
||||
Checkpoint mmap remains an unchanged control path whenever final-layout HWR is
|
||||
not selected.
|
||||
|
||||
Local source identity and publication avoid repeated full-payload passes while
|
||||
preserving the same semantic contract:
|
||||
|
||||
- Canonical files without a trusted immutable Hub blob identity are hashed in
|
||||
parallel. Their content IDs are cached per storage domain under the HWR root
|
||||
and reused only when the path, inode, size, timestamps, symlink target, and
|
||||
cache-record checksum still match. Corrupt entries are rebuilt; cache-lock or
|
||||
cache-I/O failure falls back to hashing the canonical file directly.
|
||||
- Producers that emit payload bytes in canonical storage-key order may declare
|
||||
`ordered=True`. The filesystem writer then computes file and tensor SHA256
|
||||
values during the write and overlaps payload `fsync` with later producer
|
||||
work. Producers without that guarantee retain unordered writes and use a
|
||||
parallel readback checksum fallback.
|
||||
- Manifest and `READY` publication still wait for every payload `fsync` and
|
||||
checksum to complete. These optimizations do not weaken artifact validation
|
||||
or change lease ownership.
|
||||
|
||||
#### Registered mmap transport
|
||||
|
||||
After DLO takes an HWR lease, the no-AllGather backend may register the lease's
|
||||
complete immutable mapped ranges under the existing `pin_cpu_memory` policy.
|
||||
Registration is transport state: it does not change artifact identity, the
|
||||
store, tensor ownership, or H2D payload. On success, each tensor view copies
|
||||
directly into the existing rotating HBM block buffers and the two private host
|
||||
staging slots are not allocated.
|
||||
|
||||
`--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling
|
||||
over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned
|
||||
memory policy, unsupported platform/capability, over-budget mapping, or fully
|
||||
rolled-back registration error selects the existing two-slot staging path. A
|
||||
partial registration that cannot be rolled back aborts startup because closing
|
||||
the lease would unmap memory still owned by the platform.
|
||||
|
||||
Direct checkpoint mmap remains unchanged and continues to use staging. It may
|
||||
require loader-owned per-block transforms, while the HWR artifact already
|
||||
contains final runtime bytes. DLO AllGather never receives an HWR final-layout
|
||||
lease and therefore never enters this registration path.
|
||||
|
||||
#### Pre-service transaction boundary
|
||||
|
||||
The startup transaction does not end at restore commit. A lease-backed model is
|
||||
@@ -360,7 +317,6 @@ loader owns
|
||||
-> backend takes once before asynchronous work
|
||||
-> backend drains pending H2D work
|
||||
-> backend releases hooks, staging, and model references
|
||||
-> backend unregisters every mapped range
|
||||
-> backend closes lease
|
||||
```
|
||||
|
||||
@@ -370,15 +326,13 @@ it. Once the backend takes ownership, only backend abort or teardown may drain
|
||||
work and close it.
|
||||
|
||||
The implementation keeps the final-layout lease through backend setup and
|
||||
initial prefetch. Successful registration copies directly from immutable HWR
|
||||
views; otherwise the backend uses the existing two bounded rank-local staging
|
||||
slots. Transactional cleanup drains device work, removes source references,
|
||||
unregisters mappings, and only then closes the lease. An unregistration failure
|
||||
retains both registration and lease for retry/process teardown. Preferred mode
|
||||
then uses the runner's fresh canonical retry; required mode propagates the
|
||||
failure.
|
||||
initial prefetch. The backend uses the existing two bounded rank-local staging
|
||||
slots for the immutable HWR tensors, and its transactional `enable()` cleanup
|
||||
removes partial hooks and closes the lease before reporting a setup failure.
|
||||
Preferred mode then uses the runner's fresh canonical retry; required mode
|
||||
propagates the failure.
|
||||
|
||||
#### PR2 and PR3 promotion gates
|
||||
#### PR2 promotion gates
|
||||
|
||||
- Warm hit performs zero ordinary DiT materialization and zero producer calls.
|
||||
- Shared warm finalization changes neither restored bytes nor backing pointers.
|
||||
@@ -390,10 +344,6 @@ failure.
|
||||
- Disabled mode and AllGather emit zero final-layout HWR interaction.
|
||||
- The lease carrier rejects duplicate take and serialization.
|
||||
- Backend setup or prefetch failure drains asynchronous work before lease close.
|
||||
- Registration is all-or-nothing; rollback/unregistration failure never closes
|
||||
a mapping still owned by the platform.
|
||||
- Registered HWR transport bypasses host staging while preserving output and
|
||||
H2D payload; unsupported registration retains the bounded staging path.
|
||||
- Compatible checkpoint mmap remains an unchanged benchmark and control path.
|
||||
- Prewarm uses one matching producer cohort per TP/SP topology and storage
|
||||
domain.
|
||||
@@ -444,12 +394,11 @@ reconstructs those tensors with their recorded layouts. Other online methods
|
||||
must use `--dlo-no-use-allgather` or disable online quantization until their
|
||||
runtime layouts are validated.
|
||||
|
||||
The Host Weight Runtime representation, publication, and no-AllGather consumer
|
||||
contracts are merged; see
|
||||
[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414),
|
||||
[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445), and
|
||||
[PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). Registered
|
||||
direct H2D is an optional transport layer over that merged lease contract.
|
||||
The Host Weight Runtime representation and publication contracts are merged;
|
||||
see [RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414) and
|
||||
[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445). The planned
|
||||
no-AllGather consumer above is the PR2 implementation; registration and direct
|
||||
asynchronous H2D remain a separate transport follow-up.
|
||||
|
||||
## Validation coverage
|
||||
|
||||
@@ -462,12 +411,8 @@ Current source-level validation includes:
|
||||
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
|
||||
sharding of finalized weights and scales;
|
||||
- exact loader-to-backend plan transfer and ordinary-loader fallback;
|
||||
- ordered publication hashing, overlapped durability, unordered parallel
|
||||
checksum fallback, and node-local source-digest reuse/invalidation;
|
||||
- rank-local mmap source retention, bounded two-slot staging, and adapter
|
||||
transforms without parameter-side flags;
|
||||
- read-only registration preflight, direct source-to-device copies, safe
|
||||
fallback, partial rollback, retryable unregistration, and lease ordering;
|
||||
- resident-layer requests requiring no-AllGather;
|
||||
- DP request-wave validation for denoising-step compatibility;
|
||||
- bounded component allocator-cache retention, conservative/forced release,
|
||||
|
||||
+30
-76
@@ -6,7 +6,7 @@ This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage
|
||||
|
||||
- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
|
||||
- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
|
||||
- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss, plus image and diffusion workload, queueing, execution, and memory signals.
|
||||
- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
|
||||
- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
|
||||
- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
|
||||
|
||||
@@ -20,7 +20,7 @@ vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` col
|
||||
|
||||
### The Problem
|
||||
|
||||
vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm_omni:` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
|
||||
vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
||||
|
||||
Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
|
||||
|
||||
@@ -38,7 +38,7 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
|
||||
|
|
||||
+--------+--------+--------+--------+--------+
|
||||
| |
|
||||
vllm_omni:* vllm:*
|
||||
vllm:omni_* vllm:*
|
||||
collectors collectors
|
||||
| |
|
||||
+----+--------+ +-----------+ +----------+ +-----------+
|
||||
@@ -58,30 +58,17 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
|
||||
|
||||
### Data Flow
|
||||
|
||||
There are five independent paths for metric collection.
|
||||
There are four independent paths for metric collection.
|
||||
|
||||
**Path 1: Pipeline-level metrics (`vllm_omni:*`)**
|
||||
**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
|
||||
|
||||
`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
|
||||
|
||||
- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
|
||||
- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm_omni:requests_success_total{finished_reason}`.
|
||||
- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` on the existing completion Counter.
|
||||
- `inc_requests_failed(reason)` — increments the dedicated failure-attribution Counter once per request. Reasons are normalized to the bounded set `client_abort`, `client_disconnect`, `stage_error`, and `unknown`.
|
||||
- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
|
||||
- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
|
||||
|
||||
**Path 2: Image and diffusion metrics**
|
||||
|
||||
Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetrics` owns pipeline/stage workload families, while `OmniModalityMetrics` owns per-replica diffusion breakdowns. Finished stage messages are consumed once in `OmniBase._process_single_result`, which prevents replayed terminal messages from incrementing Counters or Histograms twice.
|
||||
|
||||
- `stage_gen_time_s`, `stage_in_queue_s`, image workload, inference-step, and per-stage peak-memory values are emitted from `StageRequestStats` when a stage finishes. `image_pixels` and `image_count_total` are restricted to image outputs; `num_inference_steps` is restricted to diffusion stages.
|
||||
- Diffusion engine timings are accumulated per request by `OrchestratorAggregator.accumulate_diffusion_metrics`. Engine-side millisecond values are converted to seconds before `observe_modality_at_finalize(...)` dispatches the per-replica Histogram observations.
|
||||
- `request_queue_wait_s` is emitted once at final request completion from `pipeline_timings["queue_wait_ms"]`. A present value of `0` is a valid observation; a missing key produces no sample.
|
||||
- `stage_waiting_requests` is the sum of the latest scheduler snapshots for all live replicas in a stage. AR snapshots arrive through `SchedulerStats`; diffusion snapshots ride normal diffusion output metrics. A dead replica's snapshot is removed. Because this deliberately reuses existing message paths, the gauge can retain its last value while a replica produces no output.
|
||||
- `kv_wait_s` measures the AR scheduler interval from entering the KV-transfer-free wait state until extraction acknowledgement. The sample is carried to the Orchestrator through the existing engine output and labeled by connector type. Abort/error/removal paths discard incomplete timers rather than emitting partial waits.
|
||||
|
||||
Zero and missing values have different meanings. Valid zero-duration queue observations are recorded, while unavailable optional measurements are omitted instead of being exported as synthetic zeroes. `peak_memory_mb`, `diffusion_forward_s`, `diffusion_kv_load_s`, `vae_decode_s`, and `kv_wait_s` therefore appear only when their data source is active.
|
||||
|
||||
**Path 3: Audio modality metrics (`vllm_omni:audio_*`)**
|
||||
**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
|
||||
|
||||
`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
|
||||
|
||||
@@ -89,7 +76,7 @@ Zero and missing values have different meanings. Valid zero-duration queue obser
|
||||
- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
|
||||
- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
|
||||
|
||||
**Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`)**
|
||||
**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
|
||||
|
||||
`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
|
||||
|
||||
@@ -100,7 +87,7 @@ The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transf
|
||||
|
||||
Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
|
||||
|
||||
**Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap)**
|
||||
**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
|
||||
|
||||
The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
|
||||
|
||||
@@ -110,9 +97,9 @@ The Orchestrator runs in a background thread. The API server (OmniBase) runs in
|
||||
|
||||
### Metric Registration and Lifecycle
|
||||
|
||||
All `vllm_omni:*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
|
||||
All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
|
||||
|
||||
The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm_omni:*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
|
||||
The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
|
||||
|
||||
## OmniPrometheusStatLogger Wrap
|
||||
|
||||
@@ -143,11 +130,11 @@ A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestr
|
||||
|
||||
## Gating: `--log-stats`
|
||||
|
||||
All metrics — both the omni-specific `vllm_omni:*` families and the upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
|
||||
All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
|
||||
|
||||
Behavior with `--log-stats=off` (default):
|
||||
|
||||
- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm_omni:*` stay materialized but never have data written to them.
|
||||
- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
|
||||
- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
|
||||
- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
|
||||
|
||||
@@ -175,43 +162,10 @@ second.
|
||||
|
||||
| Metric | Type | Labels | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `vllm_omni:num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
|
||||
| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
|
||||
| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
||||
| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
|
||||
|
||||
### Image and diffusion service-level metrics
|
||||
|
||||
| Metric | Type | Labels | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `vllm_omni:stage_gen_time_s` | Histogram | `model_name`, `stage`, `stage_type` | Stage submit to finished output; includes in-stage queueing |
|
||||
| `vllm_omni:request_queue_wait_s` | Histogram | `model_name` | Orchestration-layer queue wait; a present zero is recorded |
|
||||
| `vllm_omni:stage_waiting_requests` | Gauge | `model_name`, `stage` | Sum of the latest waiting snapshots across live replicas in the stage |
|
||||
| `vllm_omni:stage_in_queue_s` | Histogram | `model_name`, `stage` | Diffusion scheduler wait from enqueue to initial execution admission |
|
||||
| `vllm_omni:num_inference_steps` | Histogram | `model_name` | Denoising-step distribution for diffusion stages |
|
||||
| `vllm_omni:image_count_total` | Counter | `model_name` | Cumulative number of generated image outputs |
|
||||
| `vllm_omni:image_pixels` | Histogram | `model_name` | Pixel-count distribution for image outputs |
|
||||
| `vllm_omni:peak_memory_mb` | Gauge | `model_name`, `stage` | Positive per-stage peak-memory samples reported by engine outputs |
|
||||
| `vllm_omni:requests_failed_total` | Counter | `model_name`, `reason` | Request failures with bounded reason values (`client_abort`, `client_disconnect`, `stage_error`, `unknown`) |
|
||||
| `vllm_omni:kv_wait_s` | Histogram | `model_name`, `connector_type` | Completed AR KV-transfer wait intervals, grouped by connector backend |
|
||||
|
||||
### Diffusion execution breakdown
|
||||
|
||||
Labels: `{model_name, stage, replica}`.
|
||||
|
||||
| Metric | Type | Description |
|
||||
|--------|------|-------------|
|
||||
| `vllm_omni:diffusion_exec_s` | Histogram | Core diffusion-step execution time per request |
|
||||
| `vllm_omni:diffusion_exec_per_step_s` | Histogram | Core diffusion execution divided by inference-step count |
|
||||
| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time |
|
||||
| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing time |
|
||||
| `vllm_omni:vae_decode_s` | Histogram | VAE decode time when the engine reports it |
|
||||
| `vllm_omni:diffusion_forward_s` | Histogram | Forward-only denoise-loop time when pipeline profiling is enabled |
|
||||
| `vllm_omni:diffusion_kv_load_s` | Histogram | Diffusion-side KV receive/load time when upstream KV is used |
|
||||
| `vllm_omni:image_ttfp_s` | Histogram | Image-stage submit to first materialized image output |
|
||||
| `vllm_omni:denoise_step_latency_s` | Histogram | Mean diffusion forward time per denoise step for image output stages |
|
||||
|
||||
The optional breakdown and memory families are sparse by design: when an engine does not produce the source measurement, no sample is emitted. This differs from queue timings, where an explicitly measured zero is meaningful and is retained.
|
||||
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
|
||||
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
|
||||
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
||||
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
|
||||
|
||||
### Audio (7)
|
||||
|
||||
@@ -219,13 +173,13 @@ Labels: `{model_name, stage, replica}` plus the listed extra label.
|
||||
|
||||
| Metric | Type | Extra label | Description |
|
||||
|--------|------|-------------|-------------|
|
||||
| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
||||
| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
||||
| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
|
||||
| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frames generated |
|
||||
| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
||||
| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
||||
| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
||||
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
||||
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
||||
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
|
||||
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
|
||||
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
||||
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
||||
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
||||
|
||||
### Cross-stage transfer (4)
|
||||
|
||||
@@ -233,10 +187,10 @@ Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
|
||||
|
||||
| Metric | Type | Description |
|
||||
|--------|------|-------------|
|
||||
| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
||||
| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
||||
| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
||||
| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
||||
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
||||
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
||||
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
||||
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
||||
|
||||
### LLM stage-level (wrapped `vllm:*`)
|
||||
|
||||
@@ -249,7 +203,7 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency,
|
||||
- `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
|
||||
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
|
||||
- Sizes use the `_bytes` suffix.
|
||||
- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
|
||||
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
||||
|
||||
## Logging vs. Prometheus
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ validation_paths:
|
||||
- tests/host_weight_runtime/**
|
||||
upstream_refs:
|
||||
- https://github.com/vllm-project/vllm-omni/issues/6414
|
||||
last_reviewed: 2026-08-25
|
||||
last_reviewed: 2026-08-21
|
||||
---
|
||||
|
||||
# Host Weight Runtime
|
||||
@@ -84,13 +84,6 @@ another. `WeightArtifactIdentity` includes:
|
||||
- producer version, implementation fingerprint, manifest schema, and restorer
|
||||
schema.
|
||||
|
||||
Loader integrations may persist content digests for canonical source shards in
|
||||
the selected storage domain. Reuse requires the same path, inode, size,
|
||||
timestamps, symlink target, and cache-record checksum. Corrupt records are
|
||||
rebuilt, and coordination or cache-I/O failure falls back to hashing the
|
||||
canonical source directly. This is an identity-computation optimization, not a
|
||||
second source of truth or an artifact substitution rule.
|
||||
|
||||
The identity excludes transfer policy such as registered mmap versus private
|
||||
pinned staging. Those paths move the same host representation and must not
|
||||
produce duplicate artifacts.
|
||||
@@ -244,12 +237,6 @@ Every publication computes full file and tensor hashes. Local lookup supports:
|
||||
and
|
||||
- `full_checksum`, which additionally reads and hashes every payload file.
|
||||
|
||||
An ordered producer may compute file and tensor digests while payload bytes are
|
||||
written and overlap each closed payload's `fsync` with later producer work. An
|
||||
unordered producer retains parallel readback hashing. In both cases manifest
|
||||
and `READY` publication wait for every payload digest and `fsync`; the atomic
|
||||
publication and validation contract is unchanged.
|
||||
|
||||
Filesystem verity is modeled but unsupported. Invalid artifacts receive an
|
||||
external deny marker before their shared lookup lock is released; a build owner
|
||||
then quarantines the immutable directory under the exclusive lock.
|
||||
|
||||
@@ -37,6 +37,7 @@ GENERAL_EXAMPLE_SLUGS = frozenset(
|
||||
"text_to_speech",
|
||||
"text_to_video",
|
||||
"x_to_text",
|
||||
"x_to_video_audio",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -58,6 +58,7 @@ th {
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
@@ -93,6 +94,7 @@ th {
|
||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
|
||||
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
|
||||
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
|
||||
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
|
||||
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
|
||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
|
||||
|
||||
@@ -65,7 +65,6 @@ omni = Omni(
|
||||
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
|
||||
| `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
|
||||
| `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
|
||||
| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
|
||||
|
||||
## Host-weight loading
|
||||
|
||||
@@ -112,12 +111,6 @@ perform a collective and uses the same rank-local transfer behavior.
|
||||
HWR is an opt-in startup optimization for models that declare the final-layout
|
||||
BF16 restore contract. Use it only with no-AllGather DLO:
|
||||
|
||||
The validated BF16 model contracts currently cover MiniMax H3 and
|
||||
`black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model
|
||||
class but has not been validated against this contract. FLUX.2-dev, online
|
||||
FP8, HSDP, LoRA/adapted weights, and non-default load formats remain outside
|
||||
the validated scope.
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
@@ -185,48 +178,12 @@ returning to `required`.
|
||||
|
||||
On a cold start, the canonical loader remains authoritative and publishes a
|
||||
validated final-layout artifact for later workers. A warm start restores the
|
||||
DiT final tensors without ordinary DiT materialization. DLO then attempts to
|
||||
register the immutable mapping for direct H2D; if registration is unavailable,
|
||||
it streams through the same two bounded host staging slots. The artifact
|
||||
identity includes the TP rank/size and SP layout, so TP1, TP2 rank-local shards,
|
||||
and distinct SP layouts do not alias one another. Publication failure remains
|
||||
separate from a `preferred` serving startup; the next `required` startup
|
||||
provides the explicit artifact-availability check.
|
||||
|
||||
For local canonical checkpoints, the first eligible worker may hash source
|
||||
shards to establish immutable identity. HWR caches those digests in the same
|
||||
node-local storage domain and validates file metadata before reuse, so later
|
||||
workers normally avoid repeating that read. Cold BF16 publication also hashes
|
||||
ordered payloads as they are written and overlaps payload durability work with
|
||||
later shards; this changes startup work only, not artifact contents or runtime
|
||||
H2D behavior.
|
||||
|
||||
#### Registered direct H2D
|
||||
|
||||
Registration is attempted automatically only for an eligible warm HWR hit when
|
||||
`pin_cpu_memory` is enabled. A successful path is:
|
||||
|
||||
```text
|
||||
shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel
|
||||
```
|
||||
|
||||
It removes the recurrent mmap-to-private-staging CPU copy and does not allocate
|
||||
the two host staging slots. It does not reduce the H2D payload or make GPU
|
||||
kernels access host memory directly.
|
||||
|
||||
CUDA requires read-only host-registration capability for the immutable HWR
|
||||
mapping. Unsupported capability, a positive registration limit smaller than
|
||||
the complete page-aligned mapping, or a safely rolled-back registration error
|
||||
falls back to bounded staging. Programmatic configurations can set
|
||||
`pin_cpu_memory=False` to disable registration explicitly. A successful
|
||||
registration locks the complete mapped range in host memory for that worker's
|
||||
lifetime, so use
|
||||
`--dlo-host-registration-limit-gib` when an operator-enforced ceiling is
|
||||
required.
|
||||
|
||||
Shutdown drains H2D work, releases hook/source references, unregisters every
|
||||
range, and only then closes the HWR lease. AllGather and direct checkpoint mmap
|
||||
retain their existing transfer paths.
|
||||
DiT final tensors without ordinary DiT materialization, then DLO streams them
|
||||
through the same two bounded host staging slots. The artifact identity includes
|
||||
the TP rank/size and SP layout, so TP1, TP2 rank-local shards, and distinct SP
|
||||
layouts do not alias one another. Publication failure remains separate from a
|
||||
`preferred` serving startup; the next `required` startup provides the explicit
|
||||
artifact-availability check.
|
||||
|
||||
When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
|
||||
not resolve HWR sources or construct its store.
|
||||
|
||||
@@ -129,6 +129,7 @@ The following tables show which models support each feature:
|
||||
| **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -146,7 +147,7 @@ The following tables show which models support each feature:
|
||||
|
||||
> Notes:
|
||||
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
|
||||
> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
|
||||
> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
|
||||
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
|
||||
> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
|
||||
> 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
|
||||
@@ -164,6 +165,7 @@ The following tables show which models support each feature:
|
||||
| **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
# X-To-Video-Audio
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/x_to_video_audio>.
|
||||
|
||||
|
||||
The `DreamID-Omni` pipeline generates short videos from text, image and video.
|
||||
|
||||
## Local CLI Usage
|
||||
### Download the Model locally
|
||||
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
|
||||
|
||||
```bash
|
||||
python download_dreamid_omni.py --output-dir ./dreamid_omni
|
||||
```
|
||||
After download, the model directory will look like this:
|
||||
|
||||
```
|
||||
dreamid_omni/
|
||||
├── DreamID-Omni/
|
||||
│ ├── dreamid_omni.safetensors
|
||||
├── MMAudio/
|
||||
│ ├── ext_weights/
|
||||
│ │ ├── best_netG.pt
|
||||
│ │ ├── v1-16.pth
|
||||
├── Wan2.2-TI2V-5B/
|
||||
│ ├── google/*
|
||||
│ ├── models_t5_umt5-xxl-enc-bf16.pth
|
||||
│ ├── Wan2.2_VAE.pth
|
||||
│
|
||||
├── model_index.json # create by download_dreamid_omni.py
|
||||
```
|
||||
|
||||
### Run the Inference
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
|
||||
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
|
||||
For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
|
||||
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--use-hsdp \
|
||||
--hsdp-shard-size 2 \
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
|
||||
|
||||
You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
|
||||
|
||||
For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
|
||||
|
||||
```python
|
||||
# Example usage for oneip, ref media from the official repo DreamID-Omni
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
|
||||
--image-path 9.png \
|
||||
--audio-path 9.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output out_dreamid_omni_oneip.mp4
|
||||
```
|
||||
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--model`: path to the model local directory.
|
||||
- `--height/--width`: output resolution (defaults 704 * 1024).
|
||||
- `--image-path`: path to the input image list.
|
||||
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
|
||||
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
|
||||
- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
|
||||
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
|
||||
- `--hsdp-replicate-size`: number of HSDP replica groups.
|
||||
- `--num-inference-steps`: number of denoising steps (defaults 45).
|
||||
- `--video-negative-prompt`: negative prompt for video generation.
|
||||
- `--audio-negative-prompt`: negative prompt for audio generation.
|
||||
- `--enable-cpu-offload`: enable CPU offload (defaults False).
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "download_dreamid_omni.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/x_to_video_audio/download_dreamid_omni.py"
|
||||
``````
|
||||
??? abstract "x_to_video_audio.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/x_to_video_audio/x_to_video_audio.py"
|
||||
``````
|
||||
@@ -0,0 +1,193 @@
|
||||
# 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"
|
||||
``````
|
||||
@@ -0,0 +1,108 @@
|
||||
import argparse
|
||||
import fcntl
|
||||
import os
|
||||
import site
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
|
||||
DEPENDENCY_BRANCH = "omni"
|
||||
CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
|
||||
LOCK_FILE = CACHE_DIR / ".install.lock"
|
||||
DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
|
||||
|
||||
|
||||
def download_dependency():
|
||||
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(LOCK_FILE, "w") as f:
|
||||
fcntl.flock(f, fcntl.LOCK_EX)
|
||||
if not DEPENDENCY_DIR.exists():
|
||||
print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
|
||||
subprocess.run(
|
||||
["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
|
||||
check=True,
|
||||
)
|
||||
print("Download finished.")
|
||||
fcntl.flock(f, fcntl.LOCK_UN)
|
||||
|
||||
# write .pth to site-packages
|
||||
site_packages = Path(site.getsitepackages()[0])
|
||||
pth_file = site_packages / "vllm_omni_dependency.pth"
|
||||
pth_file.write_text(str(DEPENDENCY_DIR))
|
||||
print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
|
||||
|
||||
|
||||
def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
|
||||
"""Download files from HF repo and log time + destination."""
|
||||
if os.path.exists(local_dir):
|
||||
print(f"Directory {local_dir} already exists. Skipping download.")
|
||||
return
|
||||
print(f"Starting download from {repo_id} into {local_dir}")
|
||||
start_time = time.time()
|
||||
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
local_dir=local_dir,
|
||||
local_dir_use_symlinks=False,
|
||||
allow_patterns=allow_patterns,
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
|
||||
|
||||
|
||||
def main(output_dir: str):
|
||||
# Wan2.2
|
||||
wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
|
||||
timed_download(
|
||||
repo_id="Wan-AI/Wan2.2-TI2V-5B",
|
||||
local_dir=wan_dir,
|
||||
allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
|
||||
)
|
||||
|
||||
# MMAudio
|
||||
mm_audio_dir = os.path.join(output_dir, "MMAudio")
|
||||
timed_download(
|
||||
repo_id="hkchengrex/MMAudio",
|
||||
local_dir=mm_audio_dir,
|
||||
allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
|
||||
)
|
||||
|
||||
dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
|
||||
|
||||
timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
|
||||
|
||||
# Now we construct the config file
|
||||
import json
|
||||
|
||||
data = {
|
||||
"_class_name": "DreamIDOmniPipeline",
|
||||
}
|
||||
|
||||
with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
|
||||
|
||||
transformer_dir = os.path.join(output_dir, "transformer")
|
||||
os.makedirs(transformer_dir, exist_ok=True)
|
||||
with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
|
||||
json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
|
||||
print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}")
|
||||
|
||||
# now we download the dependency code
|
||||
download_dependency()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Download models from Hugging Face")
|
||||
parser.add_argument(
|
||||
"--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
main(args.output_dir)
|
||||
@@ -0,0 +1,163 @@
|
||||
# X-To-Video-Audio
|
||||
|
||||
The `DreamID-Omni` pipeline generates short videos from text, image and video.
|
||||
|
||||
## Local CLI Usage
|
||||
### Download the Model locally
|
||||
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
|
||||
|
||||
```bash
|
||||
python download_dreamid_omni.py --output-dir ./dreamid_omni
|
||||
```
|
||||
After download, the model directory will look like this:
|
||||
|
||||
```
|
||||
dreamid_omni/
|
||||
├── DreamID-Omni/
|
||||
│ ├── dreamid_omni.safetensors
|
||||
├── MMAudio/
|
||||
│ ├── ext_weights/
|
||||
│ │ ├── best_netG.pt
|
||||
│ │ ├── v1-16.pth
|
||||
├── Wan2.2-TI2V-5B/
|
||||
│ ├── google/*
|
||||
│ ├── models_t5_umt5-xxl-enc-bf16.pth
|
||||
│ ├── Wan2.2_VAE.pth
|
||||
│
|
||||
├── model_index.json
|
||||
└── transformer/
|
||||
└── config.json # create by download_dreamid_omni.py
|
||||
```
|
||||
|
||||
### Run the Inference
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
|
||||
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
|
||||
For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
|
||||
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--use-hsdp \
|
||||
--hsdp-shard-size 2 \
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
|
||||
|
||||
You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
|
||||
|
||||
For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
|
||||
|
||||
```python
|
||||
# Example usage for oneip, ref media from the official repo DreamID-Omni
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
|
||||
--image-path 9.png \
|
||||
--audio-path 9.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output out_dreamid_omni_oneip.mp4
|
||||
```
|
||||
|
||||
|
||||
## MagiHuman (`--model-type magi-human`)
|
||||
|
||||
MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param
|
||||
T5Gemma text encoder. A detailed text prompt is the only required input; an optional
|
||||
image and/or audio file may be supplied for conditioning. Natively supports Tensor
|
||||
Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard
|
||||
the MoE weights and the text encoder.
|
||||
|
||||
> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct
|
||||
> attention-kernel behaviour (the pipeline otherwise falls back to stubs).
|
||||
|
||||
Text-only generation:
|
||||
|
||||
```bash
|
||||
python x_to_video_audio.py \
|
||||
--model-type magi-human \
|
||||
--model /path/to/daVinci-MagiHuman \
|
||||
--prompt "A young woman with long, wavy golden blonde hair... <dialogue and background sound>" \
|
||||
--tensor-parallel-size 4 \
|
||||
--height 256 --width 448 \
|
||||
--num-inference-steps 8 \
|
||||
--seed 52 \
|
||||
--extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
With optional image and audio conditioning:
|
||||
|
||||
```bash
|
||||
python x_to_video_audio.py \
|
||||
--model-type magi-human \
|
||||
--model /path/to/daVinci-MagiHuman \
|
||||
--prompt "A young woman..." \
|
||||
--tensor-parallel-size 4 \
|
||||
--height 256 --width 448 \
|
||||
--num-inference-steps 8 \
|
||||
--seed 52 \
|
||||
--extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared
|
||||
in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
|
||||
|
||||
- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set,
|
||||
because the audio length then determines the number of frames).
|
||||
- `audio_path`: path to an audio file for audio-to-video conditioning. When provided,
|
||||
the audio drives the frame count and is encoded as a latent condition into the DiT
|
||||
(`is_a2v` mode). Omit for pure text-to-video+audio generation.
|
||||
- `image_path`: path to an image file for visual conditioning. Applied at both the
|
||||
base-resolution (BR) and super-resolution (SR) stages when SR is enabled.
|
||||
- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped
|
||||
when these are omitted.
|
||||
- `sr_num_inference_steps`: denoising steps for the SR stage.
|
||||
|
||||
## DreamID-Omni arguments
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--model`: path to the model local directory.
|
||||
- `--height/--width`: output resolution (defaults 704 * 1024).
|
||||
- `--image-path`: path to the input image list.
|
||||
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
|
||||
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
|
||||
- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
|
||||
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
|
||||
- `--hsdp-replicate-size`: number of HSDP replica groups.
|
||||
- `--num-inference-steps`: number of denoising steps (defaults 45).
|
||||
- `--video-negative-prompt`: negative prompt for video generation.
|
||||
- `--audio-negative-prompt`: negative prompt for audio generation.
|
||||
- `--enable-cpu-offload`: enable CPU offload (defaults False).
|
||||
- `--cache-backend`: enable `cache_dit` for acceleration.
|
||||
- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
|
||||
(VAEs, the T5 text encoder, norms and modulation stay bf16.)
|
||||
@@ -0,0 +1,279 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
|
||||
AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
|
||||
parser.add_argument("--model", required=True, help="Model ckpt root directory.")
|
||||
parser.add_argument(
|
||||
"--model-type",
|
||||
default="dreamid-omni",
|
||||
choices=["dreamid-omni", "magi-human"],
|
||||
help="Model type.",
|
||||
)
|
||||
parser.add_argument("--prompt", default=None, help="Text prompt.")
|
||||
|
||||
parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
|
||||
parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
|
||||
parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=str,
|
||||
default=None,
|
||||
help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
|
||||
'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
|
||||
)
|
||||
|
||||
parser.add_argument("--height", type=int, default=704, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=1280, help="Video width.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
|
||||
parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
|
||||
parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
|
||||
parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
|
||||
parser.add_argument(
|
||||
"--cfg-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
choices=[1, 2, 3, 4],
|
||||
help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-negative-prompt",
|
||||
default="jitter, bad hands, blur, distortion",
|
||||
help="Negative prompt for video.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-negative-prompt",
|
||||
default="robotic, muffled, echo, distorted",
|
||||
help="Negative prompt for audio.",
|
||||
)
|
||||
parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "int8"],
|
||||
help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-cpu-offload",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable CPU offloading for diffusion models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-layerwise-offload",
|
||||
action="store_true",
|
||||
help="Enable layerwise (blockwise) offloading on DiT modules.",
|
||||
)
|
||||
parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit"], help="Cache backend.")
|
||||
parser.add_argument(
|
||||
"--use-hsdp",
|
||||
action="store_true",
|
||||
help="Enable HSDP for DreamID-Omni fused transformer blocks.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hsdp-shard-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for HSDP sharding when HSDP is enabled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hsdp-replicate-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of HSDP replica groups. Default 1 means pure sharding.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_image_and_audio(image_paths, audio_paths):
|
||||
image = []
|
||||
audio = []
|
||||
|
||||
for path in image_paths:
|
||||
with Image.open(path) as img:
|
||||
img = img.convert("RGB")
|
||||
image.append(img)
|
||||
|
||||
for path in audio_paths:
|
||||
audio_array, sr = load_audio(path, sr=16000)
|
||||
audio_array = audio_array[int(sr * 1) : int(sr * 3)]
|
||||
audio.append(audio_array)
|
||||
return image, audio
|
||||
|
||||
|
||||
def _extract_peak_memory_mb(result: Any) -> float:
|
||||
"""Pull worker-reported peak VRAM (MiB) generation result.
|
||||
|
||||
Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
|
||||
"""
|
||||
if isinstance(result, list):
|
||||
result = result[0] if result else None
|
||||
if result is None:
|
||||
return 0.0
|
||||
val = getattr(result, "peak_memory_mb", 0.0)
|
||||
if not val:
|
||||
inner = result
|
||||
if isinstance(inner, list):
|
||||
inner = inner[0] if inner else None
|
||||
val = getattr(inner, "peak_memory_mb", 0.0)
|
||||
try:
|
||||
return float(val or 0.0)
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.prompt is None and args.prompt_file is None:
|
||||
raise ValueError("Either --prompt or --prompt-file must be provided.")
|
||||
|
||||
text_prompt = args.prompt
|
||||
if args.prompt_file:
|
||||
with open(args.prompt_file) as f:
|
||||
text_prompt = json.load(f)
|
||||
text_prompt = re.sub(
|
||||
r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
|
||||
).strip()
|
||||
text_prompt = re.sub(
|
||||
r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
|
||||
).strip()
|
||||
text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
|
||||
text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
|
||||
|
||||
if args.model_type == "magi-human":
|
||||
prompt = text_prompt
|
||||
extra_args = json.loads(args.extra_body) if args.extra_body else {}
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args=extra_args,
|
||||
)
|
||||
else:
|
||||
image, audio = load_image_and_audio(args.image_path, args.audio_path)
|
||||
|
||||
prompt = {
|
||||
"prompt": text_prompt,
|
||||
"video_negative_prompt": args.video_negative_prompt,
|
||||
"audio_negative_prompt": args.audio_negative_prompt,
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
}
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args={
|
||||
"solver_name": args.solver_name,
|
||||
"shift": args.shift,
|
||||
},
|
||||
)
|
||||
|
||||
cache_config = None
|
||||
if args.cache_backend == "cache_dit":
|
||||
cache_config = {
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"max_cached_steps": 20,
|
||||
"residual_diff_threshold": 0.24,
|
||||
"max_continuous_cached_steps": 3,
|
||||
"enable_taylorseer": False,
|
||||
"taylorseer_order": 1,
|
||||
"scm_steps_mask_policy": None,
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
|
||||
omni_kwargs = dict(
|
||||
model=args.model,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
model_type=args.model_type,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
enable_layerwise_offload=args.enable_layerwise_offload,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
)
|
||||
if args.quantization is not None:
|
||||
omni_kwargs["quantization"] = args.quantization
|
||||
omni = Omni(**omni_kwargs)
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(prompt, sampling_params)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
if not outputs:
|
||||
raise RuntimeError("No output returned from the model.")
|
||||
result = outputs[0]
|
||||
if not result.images:
|
||||
raise RuntimeError("No video frames found in model output.")
|
||||
generated_video = result.images[0]
|
||||
mm = result.multimodal_output or {}
|
||||
generated_audio = mm.get("audio")
|
||||
|
||||
if args.model_type == "magi-human":
|
||||
# MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
|
||||
fps = float(mm.get("fps", 25))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 24000))
|
||||
frames = generated_video
|
||||
else:
|
||||
fps = float(mm.get("fps", 24))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 16000))
|
||||
# DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
|
||||
# mux_video_audio_bytes expects (F, H, W, C) uint8.
|
||||
if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
|
||||
raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
|
||||
frames = generated_video.transpose(1, 2, 3, 0)
|
||||
frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
|
||||
|
||||
audio_np = None
|
||||
if generated_audio is not None:
|
||||
audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32)
|
||||
|
||||
output_path = args.output
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
frames,
|
||||
audio_np,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=sample_rate,
|
||||
)
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(video_bytes)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
print(f"Total time: {elapsed:.2f}s")
|
||||
peak_mb = _extract_peak_memory_mb(outputs)
|
||||
if peak_mb:
|
||||
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -32,14 +32,12 @@ GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That
|
||||
profile admits at most four concurrent sequences per stage.
|
||||
|
||||
| deploy config | GPUs | Notes |
|
||||
| --- | --- | --- |
|
||||
|---|---|---|
|
||||
| `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. |
|
||||
| `minicpmo_4_5_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
|
||||
| `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. |
|
||||
| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
|
||||
|
||||
Every profile sets `session_mode: duplex`, so native full-duplex is served
|
||||
from the same process as `/v1/chat/completions`.
|
||||
| `minicpmo_4_5_duplex.yaml` | 1 | Experimental native full-duplex overlay. |
|
||||
|
||||
The split pipeline preserves native-duplex epoch/turn identity, segment text,
|
||||
turn completion, reference voice, and terminal-audio metadata through
|
||||
@@ -57,8 +55,8 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \
|
||||
```
|
||||
|
||||
For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the
|
||||
checkpoint path. For native full-duplex, connect `/v1/realtime?duplex=1` on
|
||||
this same server.
|
||||
checkpoint path. To start the experimental native duplex backend, use
|
||||
`vllm_omni/deploy/minicpmo_4_5_duplex.yaml`.
|
||||
|
||||
### Per-stage overrides
|
||||
|
||||
@@ -147,7 +145,7 @@ assistant template, so it is not an apples-to-apples accuracy run.
|
||||
|
||||
## Run the Realtime duplex CLI demo
|
||||
|
||||
After the server is running, stream one WAV through the Realtime
|
||||
After the duplex backend is running, stream one WAV through the Realtime
|
||||
WebSocket endpoint:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Minimal single-input MiniCPM-o 4.5 Realtime duplex demo.
|
||||
|
||||
Run this after starting the duplex server. Strict lifecycle, overlap, and
|
||||
@@ -22,8 +19,8 @@ if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||
PCM16_BYTES_PER_SAMPLE, # noqa: F401 - compatibility export
|
||||
PCM16_SAMPLE_RATE, # noqa: F401 - compatibility export
|
||||
PCM16_BYTES_PER_SAMPLE,
|
||||
PCM16_SAMPLE_RATE,
|
||||
RealtimeDuplexClient,
|
||||
RealtimeEventCollector,
|
||||
build_realtime_url,
|
||||
@@ -31,13 +28,6 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||
wait_for,
|
||||
write_pcm16_wav,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.client import chunk_period_ms as _chunk_period_ms # noqa: E402
|
||||
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||
has_residual_model_unit as _has_residual_model_unit,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||
reference_audio_data_url as _ref_audio_data_url,
|
||||
)
|
||||
|
||||
|
||||
class _StreamingOutputWriter:
|
||||
@@ -143,6 +133,25 @@ def _latest_model_decision(
|
||||
return decision
|
||||
|
||||
|
||||
def _chunk_period_ms(events: list[dict[str, object]]) -> int:
|
||||
for event in reversed(events):
|
||||
session = event.get("session")
|
||||
if not isinstance(session, dict):
|
||||
continue
|
||||
capabilities = session.get("capabilities")
|
||||
if not isinstance(capabilities, dict):
|
||||
continue
|
||||
chunk_period_ms = capabilities.get("chunk_period_ms")
|
||||
if isinstance(chunk_period_ms, int) and chunk_period_ms > 0:
|
||||
return chunk_period_ms
|
||||
return 1000
|
||||
|
||||
|
||||
def _has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool:
|
||||
unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000
|
||||
return bool(unit_bytes > 0 and len(pcm16) % unit_bytes)
|
||||
|
||||
|
||||
def _response_in_progress(events: list[dict[str, object]]) -> bool:
|
||||
return sum(event.get("type") == "response.created" for event in events) > sum(
|
||||
event.get("type") == "response.done" for event in events
|
||||
@@ -159,6 +168,13 @@ def _event_count_after(
|
||||
return sum(event.get("type") == event_type for event in events[index + 1 :])
|
||||
|
||||
|
||||
def _ref_audio_data_url(path: str | None) -> str | None:
|
||||
if path is None:
|
||||
return None
|
||||
ref_path = Path(path).expanduser()
|
||||
return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii")
|
||||
|
||||
|
||||
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
input_pcm16 = read_pcm16_wav(Path(args.input_wav))
|
||||
if not input_pcm16:
|
||||
|
||||
@@ -0,0 +1,234 @@
|
||||
# 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)
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
#!/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!"
|
||||
+170
@@ -0,0 +1,170 @@
|
||||
#!/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()
|
||||
@@ -839,7 +839,7 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
|
||||
export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
|
||||
```
|
||||
|
||||
Start from a non-offloaded or DLO FL2VA server command and add
|
||||
Start from a non-offloaded FL2VA server command and add
|
||||
`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
|
||||
`--lora-path` preloads the adapter; each request still activates it and uses
|
||||
the published sampling settings:
|
||||
@@ -853,10 +853,9 @@ the published sampling settings:
|
||||
|
||||
For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
|
||||
ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
|
||||
dynamic-only and does not support prefusion or LoRA composition. DLO is
|
||||
supported by keeping the request-switchable LoRA A/B buffers resident on the
|
||||
accelerator while DLO streams only the base blocks; budget for this additional
|
||||
fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
|
||||
dynamic-only and does not support prefusion, DLO, or LoRA composition.
|
||||
It also rejects model-level, layerwise, and distributed layerwise offload;
|
||||
the legacy dynamic LoRA tensors do not participate in those weight lifecycles.
|
||||
The five requested sigma points produce the four denoiser evaluations expected
|
||||
by the Turbo artifact.
|
||||
|
||||
|
||||
@@ -85,7 +85,6 @@ recipes/
|
||||
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
|
||||
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
|
||||
| [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
|
||||
| [`XiaomiMiMo/MiMo-Audio.md`](./XiaomiMiMo/MiMo-Audio.md) | Offline + online omni audio (TTS, ASR, dialogue) | 1x RTX 5090 / 5090D 32GB |
|
||||
| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
|
||||
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
|
||||
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
|
||||
|
||||
@@ -23,6 +23,8 @@ offline validation sample with TeaCache and online serving through the
|
||||
- Model: <https://huggingface.co/stabilityai/stable-audio-open-1.0>
|
||||
- Offline example:
|
||||
[`examples/offline_inference/text_to_audio`](../../examples/offline_inference/text_to_audio)
|
||||
- Online example:
|
||||
[`examples/online_serving/stable_audio`](../../examples/online_serving/stable_audio)
|
||||
- Audio generation API:
|
||||
[`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md)
|
||||
- Related issue:
|
||||
@@ -101,7 +103,7 @@ curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
"seed": 42,
|
||||
"response_format": "wav"
|
||||
}' \
|
||||
--output piano_10s.wav
|
||||
--output examples/online_serving/stable_audio/piano_10s.wav
|
||||
```
|
||||
|
||||
##### Verification
|
||||
@@ -135,12 +137,12 @@ PY
|
||||
Validate the online output:
|
||||
|
||||
```bash
|
||||
ls -lh piano_10s.wav
|
||||
ls -lh examples/online_serving/stable_audio/piano_10s.wav
|
||||
|
||||
python - <<'PY'
|
||||
import soundfile as sf
|
||||
|
||||
path = "piano_10s.wav"
|
||||
path = "examples/online_serving/stable_audio/piano_10s.wav"
|
||||
audio, sample_rate = sf.read(path)
|
||||
print("sample_rate:", sample_rate)
|
||||
print("shape:", audio.shape)
|
||||
|
||||
@@ -1,161 +0,0 @@
|
||||
# MiMo-Audio for omni audio generation and understanding
|
||||
|
||||
> Offline + online TTS, ASR, audio understanding, and multi-turn dialogue
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Xiaomi MiMo
|
||||
- Model: `XiaomiMiMo/MiMo-Audio-7B-Instruct`
|
||||
- Task: Text-to-speech, speech-to-text, audio understanding, spoken dialogue
|
||||
- Mode: Offline end-to-end scripts and online OpenAI-compatible chat serving
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want a known-good starting point for running
|
||||
MiMo-Audio on a single **32 GB Blackwell consumer GPU** (RTX 5090 / 5090D).
|
||||
The bundled deploy profile [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
|
||||
splits VRAM between the AR thinker (stage 0) and code2wav (stage 1), avoids
|
||||
FlashAttention-2 on SM 12.x, and keeps the audio tokenizer on CPU to reduce
|
||||
Stage-1 GPU pressure.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model card: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-7B-Instruct>
|
||||
- Audio tokenizer: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-Tokenizer>
|
||||
- Offline example:
|
||||
[`examples/offline_inference/mimo_audio/end2end.py`](../../examples/offline_inference/mimo_audio/end2end.py)
|
||||
- Online example:
|
||||
[`examples/online_serving/mimo_audio/README.md`](../../examples/online_serving/mimo_audio/README.md)
|
||||
- Default deploy config (generic single GPU):
|
||||
[`vllm_omni/deploy/mimo_audio.yaml`](../../vllm_omni/deploy/mimo_audio.yaml)
|
||||
- 5090D-tuned deploy config:
|
||||
[`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
|
||||
- Pipeline / deploy docs:
|
||||
[Stage configs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/)
|
||||
- Related issue or discussion:
|
||||
[RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
|
||||
- Tokenizer device alignment fix (required for CPU tokenizer on stage 1):
|
||||
[PR #6539](https://github.com/vllm-project/vllm-omni/pull/6539)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x NVIDIA GeForce RTX 5090 / 5090D 32GB
|
||||
|
||||
Two-stage async-chunk pipeline (thinker + code2wav) on a single GPU. Tested
|
||||
offline with `tts_sft`.
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu (Linux x86_64)
|
||||
- Python: 3.12
|
||||
- PyTorch: 2.11.0+cu129
|
||||
- Driver / runtime: NVIDIA driver 570.x, CUDA 12.9 runtime (Blackwell SM 12.x)
|
||||
- GPU: NVIDIA GeForce RTX 5090 D, 32 GB
|
||||
- vLLM: 0.26.0+cu129
|
||||
- vLLM-Omni: `main` (tested while [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) was open for CPU-tokenizer support)
|
||||
|
||||
Install vLLM-Omni from source and align the vLLM wheel with your CUDA stack.
|
||||
On Blackwell, prefer the **cu129** vLLM build over cu13 defaults when using
|
||||
PyTorch 2.11+cu129.
|
||||
|
||||
Download or cache:
|
||||
|
||||
- `XiaomiMiMo/MiMo-Audio-7B-Instruct`
|
||||
- `XiaomiMiMo/MiMo-Audio-Tokenizer` (or a local checkout of the tokenizer repo)
|
||||
|
||||
Export the runtime knobs below before offline or online runs:
|
||||
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="${MIMO_AUDIO_TOKENIZER_PATH:-XiaomiMiMo/MiMo-Audio-Tokenizer}"
|
||||
export MIMO_AUDIO_TOKENIZER_DEVICE=cpu
|
||||
export MIMO_AUDIO_TOKENIZER_CUDA_GRAPH=0
|
||||
export VLLM_USE_FLASHINFER_SAMPLER=0
|
||||
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
```
|
||||
|
||||
Optional (mirror / download stability):
|
||||
|
||||
```bash
|
||||
export HF_HUB_DISABLE_XET=1
|
||||
export HF_ENDPOINT=https://hf-mirror.com # if you use a Hugging Face mirror
|
||||
```
|
||||
|
||||
Do **not** set `VLLM_ATTENTION_BACKEND` on vLLM 0.26.x; configure
|
||||
`attention_backend: TRITON_ATTN` in the deploy YAML instead.
|
||||
|
||||
#### Command
|
||||
|
||||
Offline TTS validation from the repository root:
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/mimo_audio
|
||||
|
||||
python3 -u end2end.py \
|
||||
--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--text "The weather is so nice today."
|
||||
```
|
||||
|
||||
Online serving (requires MiMo chat template):
|
||||
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
|
||||
vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
|
||||
--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
|
||||
--served-model-name "MiMo-Audio-7B-Instruct" \
|
||||
--port 18091 \
|
||||
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
|
||||
```
|
||||
|
||||
Pass the same environment exports to the server process when using a CPU
|
||||
tokenizer or disabling FlashInfer sampling.
|
||||
|
||||
#### Verification
|
||||
|
||||
Offline:
|
||||
|
||||
```bash
|
||||
ls -lh output_audio/tts_sft/*.wav
|
||||
```
|
||||
|
||||
Expect a non-empty `.wav` under `output_audio/tts_sft/` and a companion `.txt`
|
||||
with the decoded text. A successful run completes both stages without a
|
||||
Stage-1 `decode_vq` device mismatch.
|
||||
|
||||
Online:
|
||||
|
||||
```bash
|
||||
curl -s -o /tmp/mimo_health.txt -w '%{http_code}' http://127.0.0.1:18091/health
|
||||
```
|
||||
|
||||
Then use the client under
|
||||
[`examples/online_serving/mimo_audio/`](../../examples/online_serving/mimo_audio/).
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage (observed offline): Stage-0 weights ~16.6 GiB; KV cache ~4.9 GiB
|
||||
with `gpu_memory_utilization: 0.78` on stage 0 and `0.12` on stage 1.
|
||||
- Key flags:
|
||||
- `--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml` for 5090D memory
|
||||
and `TRITON_ATTN`.
|
||||
- `MIMO_AUDIO_TOKENIZER_PATH` is mandatory.
|
||||
- `MIMO_AUDIO_TOKENIZER_DEVICE=cpu` saves VRAM; requires
|
||||
[#6539](https://github.com/vllm-project/vllm-omni/pull/6539) (or equivalent
|
||||
fix) so code2wav decode uses the tokenizer device.
|
||||
- `VLLM_USE_FLASHINFER_SAMPLER=0` avoids FlashInfer top-p/top-k sampler
|
||||
issues on consumer GPUs.
|
||||
- Known limitations:
|
||||
- Blackwell may log `SM 12.x requires CUDA >= 12.9` from some optional
|
||||
tooling; the tested path uses TRITON attention and cu129 wheels.
|
||||
- Without a working `flash-attn` build, audio quality may be slightly
|
||||
metallic; install FlashAttention when a compatible wheel is available.
|
||||
- First request triggers Triton JIT for rotary / attention kernels (latency
|
||||
spike); subsequent requests are faster.
|
||||
- For longer outputs or batching, you may need to raise `max_model_len` in
|
||||
the deploy YAML and match `max_position_embeddings` in the model config
|
||||
(see the offline example README).
|
||||
@@ -1,12 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Unit tests for metrics.py
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
from vllm.benchmarks.serve import TaskType
|
||||
|
||||
@@ -68,22 +66,6 @@ def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutp
|
||||
return output
|
||||
|
||||
|
||||
def _calculate_test_metrics(outputs, goodput=None):
|
||||
return calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=outputs,
|
||||
dur_s=1.0,
|
||||
tokenizer=None,
|
||||
selected_percentiles=[50.0],
|
||||
goodput_config_dict=goodput or {},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=[],
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=1.0,
|
||||
)[0]
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# total_input Tests
|
||||
# ============================================================================
|
||||
@@ -144,166 +126,6 @@ def test_audio_continuity_aggregation():
|
||||
assert p99 is not None and p99 > 0.4
|
||||
|
||||
|
||||
def test_unmeasured_duplex_latency_does_not_add_zero_samples():
|
||||
measured = _make_output(100, output_tokens=2)
|
||||
measured.ttft = 0.1
|
||||
measured.audio_ttfp = 0.2
|
||||
measured.audio_rtf = 0.5
|
||||
measured.duplex_session_metrics = {
|
||||
"mean_ttft_ms": 100.0,
|
||||
"mean_ttfp_ms": 200.0,
|
||||
"mean_rtf": 0.5,
|
||||
}
|
||||
listen_only = _make_output(100, output_tokens=0)
|
||||
listen_only.ttft = listen_only.audio_ttfp = listen_only.audio_rtf = 0.0
|
||||
listen_only.duplex_session_metrics = {
|
||||
"mean_ttft_ms": None,
|
||||
"mean_ttfp_ms": None,
|
||||
"mean_rtf": None,
|
||||
}
|
||||
|
||||
metrics = _calculate_test_metrics([measured, listen_only], {"ttft": 150.0})
|
||||
|
||||
assert metrics.mean_ttft_ms == 100.0
|
||||
assert metrics.mean_audio_ttfp_ms == 200.0
|
||||
assert metrics.mean_audio_rtf == 0.5
|
||||
assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (1, 1, 1)
|
||||
assert metrics.request_goodput == 1.0
|
||||
|
||||
|
||||
def test_all_unmeasured_duplex_latency_is_not_reported_as_zero():
|
||||
listen_only = _make_output(100, output_tokens=0)
|
||||
listen_only.duplex_session_metrics = {
|
||||
"mean_ttft_ms": None,
|
||||
"mean_ttfp_ms": None,
|
||||
"mean_rtf": None,
|
||||
}
|
||||
|
||||
metrics = _calculate_test_metrics(
|
||||
[listen_only],
|
||||
{"ttft": float("inf"), "audio_ttft": float("inf")},
|
||||
)
|
||||
|
||||
assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (0, 0, 0)
|
||||
assert math.isnan(metrics.mean_ttft_ms)
|
||||
assert math.isnan(metrics.mean_audio_ttfp_ms)
|
||||
assert math.isnan(metrics.mean_audio_rtf)
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_unmeasured_duplex_tpot_does_not_add_zero_or_misalign_goodput():
|
||||
missing_tpot = _make_output(100, output_tokens=5)
|
||||
missing_tpot.itl = []
|
||||
missing_tpot.text_latency = missing_tpot.ttft = 0.1
|
||||
missing_tpot.tpot_measured = False
|
||||
missing_tpot.duplex_session_metrics = {"mean_ttft_ms": 100.0}
|
||||
|
||||
slow_ttft = _make_output(100, output_tokens=5)
|
||||
slow_ttft.ttft = 1.0
|
||||
slow_ttft.text_latency = 1.4
|
||||
slow_ttft.itl = [0.1] * 4
|
||||
slow_ttft.duplex_session_metrics = {"mean_ttft_ms": 1000.0}
|
||||
|
||||
metrics = _calculate_test_metrics(
|
||||
[missing_tpot, slow_ttft],
|
||||
{"ttft": 500.0, "tpot": 200.0},
|
||||
)
|
||||
|
||||
assert metrics.num_tpot_samples == 1
|
||||
assert metrics.mean_tpot_ms == pytest.approx(100.0)
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_all_unmeasured_duplex_token_timing_is_not_reported_as_zero():
|
||||
output = _make_output(100, output_tokens=5)
|
||||
output.itl = []
|
||||
output.text_latency = output.ttft
|
||||
output.tpot_measured = False
|
||||
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
|
||||
|
||||
metrics = _calculate_test_metrics([output], {"tpot": float("inf")})
|
||||
|
||||
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 0)
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert math.isnan(metrics.mean_itl_ms)
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_duplex_response_timings_do_not_build_a_session_token_timeline():
|
||||
output = _make_output(100, output_tokens=5)
|
||||
output.latency = 101.0
|
||||
output.duplex_request_metrics = [
|
||||
{"response_id": "r1", "stage0_tokens": {"itls_ms": [100.0, 100.0]}},
|
||||
{"response_id": "r2", "stage0_tokens": {"itls_ms": [100.0, 100.0]}},
|
||||
]
|
||||
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
|
||||
|
||||
metrics = _calculate_test_metrics([output])
|
||||
|
||||
assert math.isnan(metrics.max_output_tokens_per_s)
|
||||
assert metrics.max_concurrent_requests == 1
|
||||
assert metrics.mean_tpot_ms == pytest.approx(100.0)
|
||||
assert metrics.mean_itl_ms == pytest.approx(100.0)
|
||||
|
||||
|
||||
def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback():
|
||||
output = _make_output(100, output_tokens=0)
|
||||
output.generated_text = "timing metadata missing"
|
||||
output.itl = []
|
||||
output.text_latency = output.ttft
|
||||
output.tpot_measured = False
|
||||
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
|
||||
|
||||
def tokenizer(text, *, add_special_tokens):
|
||||
assert text == output.generated_text
|
||||
assert add_special_tokens is False
|
||||
return type("Tokenized", (), {"input_ids": [1, 2, 3]})()
|
||||
|
||||
metrics, _ = calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=[output],
|
||||
dur_s=1.0,
|
||||
tokenizer=tokenizer,
|
||||
selected_percentiles=[50.0],
|
||||
goodput_config_dict={"tpot": float("inf")},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=[],
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=1.0,
|
||||
)
|
||||
|
||||
assert metrics.total_output == 3
|
||||
assert metrics.num_tpot_samples == 0
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_measured_zero_itl_is_not_treated_as_missing():
|
||||
output = _make_output(100, output_tokens=3)
|
||||
output.itl = [0.0, 0.0]
|
||||
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
|
||||
|
||||
metrics = _calculate_test_metrics([output], {"tpot": 1.0})
|
||||
|
||||
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (1, 2)
|
||||
assert metrics.mean_tpot_ms == 0.0
|
||||
assert metrics.mean_itl_ms == 0.0
|
||||
assert metrics.request_goodput == 1.0
|
||||
|
||||
|
||||
def test_duplex_goodput_does_not_pair_measurements_from_different_requests():
|
||||
text_only, audio_only = _make_output(100), _make_output(100)
|
||||
text_only.ttft, text_only.audio_ttfp = 0.1, 0.0
|
||||
text_only.duplex_session_metrics = {"mean_ttft_ms": 100.0, "mean_ttfp_ms": None, "mean_rtf": None}
|
||||
audio_only.ttft, audio_only.audio_ttfp = 0.0, 0.2
|
||||
audio_only.duplex_session_metrics = {"mean_ttft_ms": None, "mean_ttfp_ms": 200.0, "mean_rtf": None}
|
||||
|
||||
metrics = _calculate_test_metrics([text_only, audio_only], {"ttft": 500.0, "audio_ttft": 500.0})
|
||||
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# TTFT suppression for pure-audio (TTS) benchmarks
|
||||
# ============================================================================
|
||||
|
||||
@@ -1,778 +0,0 @@
|
||||
# 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,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
@@ -25,7 +22,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
"dataset_name",
|
||||
[
|
||||
"daily-omni",
|
||||
"omniinteract",
|
||||
"seed-tts",
|
||||
"seed-tts-text",
|
||||
"seed-tts-design",
|
||||
@@ -59,7 +55,6 @@ def test_extend_omni_choices_updates_tracking_parser_shadow(dataset_name: str) -
|
||||
[
|
||||
(["--print-stage"], "print_stage", True),
|
||||
(["--daily-omni-input-mode", "audio"], "daily_omni_input_mode", "audio"),
|
||||
(["--omniinteract-subsets", "1qna"], "omniinteract_subsets", ["1qna"]),
|
||||
(["--seed-tts-locale", "zh"], "seed_tts_locale", "zh"),
|
||||
],
|
||||
)
|
||||
@@ -85,14 +80,12 @@ def test_add_omni_args_preserves_implicit_defaults() -> None:
|
||||
args = parser.parse_args([])
|
||||
assert args.print_stage is False
|
||||
assert args.daily_omni_input_mode == "all"
|
||||
assert args.omniinteract_subsets == ["1q1a", "1q1a_math", "1qna"]
|
||||
assert args.seed_tts_locale == "en"
|
||||
assert args.explicit_keys == set()
|
||||
|
||||
|
||||
def test_update_omni_help_updates_upstream_actions() -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
parser.add_argument("--num-prompts", default=1000, help="Number of prompts.")
|
||||
parser.add_argument("--percentile-metrics", help="Upstream percentile help.")
|
||||
parser.add_argument("--random-mm-limit-mm-per-prompt", help="Upstream limit help.")
|
||||
parser.add_argument("--random-mm-bucket-config", help="Upstream bucket help.")
|
||||
@@ -100,7 +93,6 @@ def test_update_omni_help_updates_upstream_actions() -> None:
|
||||
update_omni_help(parser)
|
||||
|
||||
actions = {action.dest: action for action in parser._actions}
|
||||
assert "OmniInteract uses 3" in actions["num_prompts"].help
|
||||
assert all(metric in actions["percentile_metrics"].help for metric in ("ttfc", "tpop", "audio_rtf"))
|
||||
assert "probabilities are renormalized" in actions["random_mm_limit_mm_per_prompt"].help
|
||||
assert "Currently allows for 3 modalities" in actions["random_mm_bucket_config"].help
|
||||
@@ -128,31 +120,6 @@ def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body(
|
||||
assert args.extra_body == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("argv", "expected"),
|
||||
[
|
||||
([], 3),
|
||||
(["--num-prompts", "0"], 0),
|
||||
(["--num-prompts", "12"], 12),
|
||||
],
|
||||
)
|
||||
def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default(
|
||||
argv: list[str],
|
||||
expected: int,
|
||||
) -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
parser.add_argument("--dataset-name", default="omniinteract")
|
||||
parser.add_argument("--backend", default="openai-realtime-duplex")
|
||||
parser.add_argument("--endpoint", default="/v1/realtime")
|
||||
parser.add_argument("--omniinteract-ref-audio", default="reference.wav")
|
||||
parser.add_argument("--num-prompts", type=int, default=1000)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
preprocess_serve_args(args)
|
||||
|
||||
assert args.num_prompts == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("argv", "expected_extra_body", "expected_explicit"),
|
||||
[
|
||||
|
||||
@@ -1109,35 +1109,8 @@ class TestDeployConfigLoading:
|
||||
with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
|
||||
load_deploy_config(deploy_path)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "max_sessions"),
|
||||
[
|
||||
("minicpmo_4_5.yaml", 4),
|
||||
("minicpmo_4_5_2gpu.yaml", 4),
|
||||
("minicpmo_4_5_3gpu.yaml", 4),
|
||||
("minicpmo_4_5_8x4090.yaml", 1),
|
||||
("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
|
||||
("minicpmo_4_5_4gpu_stage1_replicas.yaml", 4),
|
||||
("minicpmo_4_5_8x4090_stage1_replicas.yaml", 1),
|
||||
],
|
||||
)
|
||||
def test_minicpmo_deploy_configs_enable_duplex(self, filename: str, max_sessions: int):
|
||||
"""Every MiniCPM-o profile serves /v1/realtime, so the retired
|
||||
minicpmo_4_5_duplex.yaml overlay has no replacement to fall back on.
|
||||
"""
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert deploy.active_stream_window == 1
|
||||
assert deploy.duplex_session.max_sessions == max_sessions
|
||||
assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
|
||||
def test_load_minicpmo_default_deploy_config(self):
|
||||
deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
|
||||
def test_load_minicpmo_duplex_deploy_config(self):
|
||||
deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
@@ -1145,15 +1118,17 @@ class TestDeployConfigLoading:
|
||||
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert deploy.async_chunk is True
|
||||
assert stages[0].yaml_engine_args["async_scheduling"] is True
|
||||
assert stages[1].yaml_engine_args["async_scheduling"] is True
|
||||
assert "Async" in (stages[0].scheduler_cls or "")
|
||||
assert "Async" in (stages[1].scheduler_cls or "")
|
||||
assert deploy.active_stream_window == 1
|
||||
assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False]
|
||||
assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
|
||||
assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"]
|
||||
assert deploy.stages[1].enforce_eager is False
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 50
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561]
|
||||
assert "codec_sampling_params" not in stages[1].yaml_engine_args
|
||||
@@ -2141,8 +2116,7 @@ class TestBaseConfigInheritance:
|
||||
def test_minicpmo_overlays_inherit_talker_sampling_params(self):
|
||||
"""Overlays must keep the base Talker codec Sampler knobs."""
|
||||
for filename in (
|
||||
"minicpmo_4_5_2gpu.yaml",
|
||||
"minicpmo_4_5_3gpu.yaml",
|
||||
"minicpmo_4_5_duplex.yaml",
|
||||
"minicpmo_4_5_3gpu_stage1_replicas.yaml",
|
||||
"minicpmo_4_5_4gpu_stage1_replicas.yaml",
|
||||
"minicpmo_4_5_8x4090_stage1_replicas.yaml",
|
||||
@@ -2154,6 +2128,9 @@ class TestBaseConfigInheritance:
|
||||
sampling = deploy.stages[1].default_sampling_params
|
||||
assert sampling is not None, f"{filename} stage 1 lost default_sampling_params"
|
||||
assert sampling["temperature"] == 0.8, filename
|
||||
# Upstream utils.TTSSamplingParams. min_tokens is not checked here:
|
||||
# the duplex overlay drops it to 0 because the model budgets each
|
||||
# generate_chunk itself.
|
||||
assert sampling["top_k"] == 25, filename
|
||||
assert sampling["top_p"] == 0.85, filename
|
||||
assert sampling["repetition_penalty"] == 1.05, filename
|
||||
|
||||
@@ -66,8 +66,6 @@ def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
|
||||
runtime_src = src[runtime_start:]
|
||||
assert "apply_minimax_h3_qwen3vl_patch" not in init_src
|
||||
assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
|
||||
assert "apply_minimax_h3_qwen3vl_sdpa_patch" not in init_src
|
||||
assert "apply_minimax_h3_qwen3vl_sdpa_patch" in runtime_src
|
||||
assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
|
||||
|
||||
|
||||
|
||||
@@ -1,317 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Drift checks for the reviewed environment-variable inventory."""
|
||||
|
||||
import ast
|
||||
import builtins
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.environment_variable_inventory import (
|
||||
ENVIRONMENT_VARIABLE_INVENTORY,
|
||||
EnvironmentVariableCategory,
|
||||
ModelEnvironmentVariableDisposition,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_REPO_ROOT = Path(__file__).parents[2]
|
||||
_PACKAGE_ROOT = _REPO_ROOT / "vllm_omni"
|
||||
_INVENTORY_MODULE = _PACKAGE_ROOT / "config" / "environment_variable_inventory.py"
|
||||
_REFERENCE_PAGE = _REPO_ROOT / "docs" / "configuration" / "environment_variables.md"
|
||||
|
||||
|
||||
def _module_string_constants(tree: ast.Module) -> dict[str, str]:
|
||||
constants: dict[str, str] = {}
|
||||
for statement in tree.body:
|
||||
targets: list[ast.expr] = []
|
||||
value: ast.expr | None = None
|
||||
if isinstance(statement, ast.Assign):
|
||||
targets = statement.targets
|
||||
value = statement.value
|
||||
elif isinstance(statement, ast.AnnAssign):
|
||||
targets = [statement.target]
|
||||
value = statement.value
|
||||
if isinstance(value, ast.Constant) and isinstance(value.value, str):
|
||||
for target in targets:
|
||||
if isinstance(target, ast.Name):
|
||||
constants[target.id] = value.value
|
||||
return constants
|
||||
|
||||
|
||||
def _os_import_aliases(tree: ast.Module) -> tuple[set[str], set[str], set[str]]:
|
||||
os_names: set[str] = set()
|
||||
getenv_names: set[str] = set()
|
||||
environ_names: set[str] = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Import):
|
||||
for imported in node.names:
|
||||
if imported.name == "os":
|
||||
os_names.add(imported.asname or "os")
|
||||
elif isinstance(node, ast.ImportFrom) and node.module == "os":
|
||||
for imported in node.names:
|
||||
if imported.name == "getenv":
|
||||
getenv_names.add(imported.asname or imported.name)
|
||||
elif imported.name == "environ":
|
||||
environ_names.add(imported.asname or imported.name)
|
||||
return os_names, getenv_names, environ_names
|
||||
|
||||
|
||||
def _is_environ(node: ast.expr, os_names: set[str], environ_names: set[str]) -> bool:
|
||||
return (
|
||||
isinstance(node, ast.Attribute)
|
||||
and isinstance(node.value, ast.Name)
|
||||
and node.value.id in os_names
|
||||
and node.attr == "environ"
|
||||
) or (isinstance(node, ast.Name) and node.id in environ_names)
|
||||
|
||||
|
||||
def _environment_key_expressions(
|
||||
node: ast.AST,
|
||||
os_names: set[str],
|
||||
getenv_names: set[str],
|
||||
environ_names: set[str],
|
||||
) -> list[ast.expr]:
|
||||
"""Return key expressions used by one direct environment access."""
|
||||
if isinstance(node, ast.Call) and node.args:
|
||||
func = node.func
|
||||
if (
|
||||
isinstance(func, ast.Attribute)
|
||||
and isinstance(func.value, ast.Name)
|
||||
and func.value.id in os_names
|
||||
and func.attr == "getenv"
|
||||
) or (isinstance(func, ast.Name) and func.id in getenv_names):
|
||||
return [node.args[0]]
|
||||
if (
|
||||
isinstance(func, ast.Attribute)
|
||||
and _is_environ(func.value, os_names, environ_names)
|
||||
and func.attr in {"get", "pop", "setdefault"}
|
||||
):
|
||||
return [node.args[0]]
|
||||
elif isinstance(node, ast.Subscript) and _is_environ(node.value, os_names, environ_names):
|
||||
return [node.slice]
|
||||
elif isinstance(node, ast.Compare):
|
||||
expressions: list[ast.expr] = []
|
||||
left_values = [node.left, *node.comparators[:-1]]
|
||||
for left, operator, right in zip(left_values, node.ops, node.comparators):
|
||||
if isinstance(operator, (ast.In, ast.NotIn)) and _is_environ(right, os_names, environ_names):
|
||||
expressions.append(left)
|
||||
return expressions
|
||||
return []
|
||||
|
||||
|
||||
def _resolve_environment_name(expression: ast.expr, constants: dict[str, str]) -> str | None:
|
||||
value: str | None = None
|
||||
if isinstance(expression, ast.Constant) and isinstance(expression.value, str):
|
||||
value = expression.value
|
||||
elif isinstance(expression, ast.Name):
|
||||
value = constants.get(expression.id)
|
||||
return value or None
|
||||
|
||||
|
||||
def _environment_wrapper_parameters(
|
||||
tree: ast.Module,
|
||||
os_names: set[str],
|
||||
getenv_names: set[str],
|
||||
environ_names: set[str],
|
||||
) -> dict[str, set[tuple[str, int | None]]]:
|
||||
"""Find local helpers that forward one parameter as an environment key."""
|
||||
wrappers: dict[str, set[tuple[str, int | None]]] = {}
|
||||
for function in (node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))):
|
||||
positional = [*function.args.posonlyargs, *function.args.args]
|
||||
parameter_indexes: dict[str, int | None] = {parameter.arg: index for index, parameter in enumerate(positional)}
|
||||
parameter_indexes.update({parameter.arg: None for parameter in function.args.kwonlyargs})
|
||||
forwarded: set[tuple[str, int | None]] = set()
|
||||
for node in ast.walk(function):
|
||||
for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
|
||||
if isinstance(expression, ast.Name) and expression.id in parameter_indexes:
|
||||
forwarded.add((expression.id, parameter_indexes[expression.id]))
|
||||
if forwarded:
|
||||
wrappers[function.name] = forwarded
|
||||
return wrappers
|
||||
|
||||
|
||||
def _environment_accesses(path: Path) -> set[str]:
|
||||
"""Return statically resolvable names accessed through ``os``."""
|
||||
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
||||
constants = _module_string_constants(tree)
|
||||
os_names, getenv_names, environ_names = _os_import_aliases(tree)
|
||||
wrappers = _environment_wrapper_parameters(tree, os_names, getenv_names, environ_names)
|
||||
names: set[str] = set()
|
||||
|
||||
for node in ast.walk(tree):
|
||||
for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
|
||||
if name := _resolve_environment_name(expression, constants):
|
||||
names.add(name)
|
||||
|
||||
if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
|
||||
for parameter_name, position in wrappers.get(node.func.id, set()):
|
||||
argument = node.args[position] if position is not None and position < len(node.args) else None
|
||||
if argument is None:
|
||||
argument = next(
|
||||
(keyword.value for keyword in node.keywords if keyword.arg == parameter_name),
|
||||
None,
|
||||
)
|
||||
if argument is not None and (name := _resolve_environment_name(argument, constants)):
|
||||
names.add(name)
|
||||
|
||||
return names
|
||||
|
||||
|
||||
def test_inventory_matches_reviewed_snapshot_counts():
|
||||
"""Make an inventory expansion an explicit review decision."""
|
||||
category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
|
||||
assert category_counts == {
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 23,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
|
||||
EnvironmentVariableCategory.INTERNAL: 2,
|
||||
}
|
||||
|
||||
disposition_counts = Counter(
|
||||
item.model_disposition
|
||||
for item in ENVIRONMENT_VARIABLE_INVENTORY.values()
|
||||
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
|
||||
)
|
||||
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 33,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
|
||||
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
|
||||
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6,
|
||||
}
|
||||
|
||||
|
||||
def test_new_public_omni_names_use_project_prefix():
|
||||
"""Grandfather legacy names without allowing more prefix exceptions."""
|
||||
legacy_public_names = {
|
||||
"DIFFUSION_ATTENTION_BACKEND",
|
||||
"DIFFUSION_CACHE_ADAPTER",
|
||||
"DIFFUSION_CACHE_BACKEND",
|
||||
"OMNI_DIFFUSION_PROMPT_EMBED_CACHE",
|
||||
"OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE",
|
||||
"OMNI_DIFFUSION_SESSION_STATE_MANAGER",
|
||||
"OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS",
|
||||
"SPEAKER_MAX_UPLOADED",
|
||||
"SPEAKER_SAMPLES_DIR",
|
||||
"VLLM_VIDEO_ASYNC_CHUNK",
|
||||
"VLLM_VIDEO_AUDIO_DELTA_MODE",
|
||||
}
|
||||
public_names = {
|
||||
name
|
||||
for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
|
||||
if item.category is EnvironmentVariableCategory.PUBLIC_OMNI
|
||||
}
|
||||
|
||||
assert {name for name in public_names if not name.startswith("VLLM_OMNI_")} == legacy_public_names
|
||||
|
||||
|
||||
def test_statically_resolvable_environment_accesses_are_classified():
|
||||
discovered: set[str] = set()
|
||||
for path in _PACKAGE_ROOT.rglob("*.py"):
|
||||
discovered.update(_environment_accesses(path))
|
||||
|
||||
assert discovered - ENVIRONMENT_VARIABLE_INVENTORY.keys() == set()
|
||||
|
||||
|
||||
def test_model_and_benchmark_inventory_entries_are_still_referenced():
|
||||
"""Reject transitional rows after their implementation has disappeared."""
|
||||
package_source = "\n".join(
|
||||
path.read_text(encoding="utf-8") for path in _PACKAGE_ROOT.rglob("*.py") if path != _INVENTORY_MODULE
|
||||
)
|
||||
transitional_categories = {
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL,
|
||||
}
|
||||
|
||||
stale_names = {
|
||||
name
|
||||
for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
|
||||
if item.category in transitional_categories and name not in package_source
|
||||
}
|
||||
|
||||
assert stale_names == set()
|
||||
|
||||
|
||||
def test_environment_scanner_covers_indirection_aliases_membership_and_casing():
|
||||
expected_by_path = {
|
||||
"metrics/definitions.py": {
|
||||
"VLLM_OMNI_BENCH_AUDIO_CHANNELS",
|
||||
"VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE",
|
||||
},
|
||||
"entrypoints/openai/video_stream_envs.py": {
|
||||
"VLLM_VIDEO_ASYNC_CHUNK",
|
||||
"VLLM_VIDEO_AUDIO_DELTA_MODE",
|
||||
},
|
||||
"model_executor/models/mimo_audio/mimo_audio.py": {"model_stage"},
|
||||
"model_executor/models/moss_tts/modeling_moss_tts_local_depth.py": {"MOSS_TTS_DEBUG_STOP"},
|
||||
"distributed/ray_utils/utils.py": {"RAY_RAYLET_PID"},
|
||||
}
|
||||
for relative_path, expected in expected_by_path.items():
|
||||
assert expected <= _environment_accesses(_PACKAGE_ROOT / relative_path)
|
||||
|
||||
|
||||
def test_generated_server_storage_environment_names_are_classified():
|
||||
from vllm_omni.config.server_settings import FileBackend
|
||||
|
||||
generated_names = {
|
||||
f"VLLM_OMNI_SERVER_STORAGE__{field_name.upper()}"
|
||||
for field_name in FileBackend.model_fields
|
||||
if field_name != "type"
|
||||
}
|
||||
assert generated_names <= ENVIRONMENT_VARIABLE_INVENTORY.keys()
|
||||
|
||||
|
||||
def test_public_omni_variables_are_in_the_reference_page():
|
||||
reference = _REFERENCE_PAGE.read_text(encoding="utf-8")
|
||||
missing = {
|
||||
name
|
||||
for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
|
||||
if item.is_public_omni and f"`{name}`" not in reference
|
||||
}
|
||||
assert missing == set()
|
||||
|
||||
|
||||
def test_secret_values_are_marked_for_redaction():
|
||||
assert {name for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items() if item.redact_value} == {
|
||||
"HF_TOKEN",
|
||||
"HUGGINGFACE_HUB_TOKEN",
|
||||
"OPENAI_API_KEY",
|
||||
}
|
||||
|
||||
|
||||
def test_collect_env_reports_only_safe_public_omni_values(monkeypatch):
|
||||
from collect_env import get_env_vars
|
||||
|
||||
monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "tea_cache")
|
||||
monkeypatch.setenv("HF_TOKEN", "hf-secret")
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "openai-secret")
|
||||
monkeypatch.setenv("VLLM_OMNI_REPLICA_ID", "internal-replica")
|
||||
|
||||
report = get_env_vars()
|
||||
|
||||
assert "DIFFUSION_CACHE_BACKEND=tea_cache" in report.splitlines()
|
||||
assert "hf-secret" not in report
|
||||
assert "openai-secret" not in report
|
||||
assert "VLLM_OMNI_REPLICA_ID" not in report
|
||||
|
||||
|
||||
@pytest.mark.parametrize("error_type", [ImportError, OSError, AssertionError, RuntimeError])
|
||||
def test_collect_env_survives_inventory_import_error(monkeypatch, error_type):
|
||||
from collect_env import get_env_vars
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
def fail_inventory_import(name, *args, **kwargs):
|
||||
if name == "vllm_omni.config.environment_variable_inventory":
|
||||
raise error_type("broken vllm-omni installation")
|
||||
return real_import(name, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", fail_inventory_import)
|
||||
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
|
||||
|
||||
assert "CUDA_VISIBLE_DEVICES=0" in get_env_vars().splitlines()
|
||||
@@ -252,37 +252,6 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
|
||||
sched.chunk_transfer_adapter.save_async.assert_not_called()
|
||||
|
||||
|
||||
def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:
|
||||
"""Native duplex + async: the next unit is usually already queued when the
|
||||
current one stops, so _handle_stopped_request applies the update and the
|
||||
stop site below it fences the same in-flight decode a second time. An
|
||||
accumulated residue outlives the drain and swallows the next unit's
|
||||
listen/speak, which the client sees as silence until its timeout.
|
||||
"""
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.RUNNING
|
||||
session.resumable = True
|
||||
session.append_output_token_ids([7])
|
||||
session.num_computed_tokens = 4
|
||||
session.num_output_placeholders = 1
|
||||
# This step's frame (settled to 0 by update_from_output) plus one extra
|
||||
# async decode still in flight.
|
||||
session.num_in_flight_tokens = 2
|
||||
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched._enqueue_waiting_request = MagicMock()
|
||||
|
||||
def handle_stopped(request: Request) -> bool:
|
||||
sched._update_request_as_session(request, _make_update([10, 20]))
|
||||
return False
|
||||
|
||||
_run_resumable_segment_stop(session, handle_stopped=handle_stopped)
|
||||
|
||||
# Exactly the one unreported decode, so the drain reaches zero before the
|
||||
# new segment's first frame arrives.
|
||||
assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1
|
||||
|
||||
|
||||
def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
session = _make_request()
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"extra_cli_args": [
|
||||
"--deploy-config",
|
||||
"vllm_omni/deploy/minicpmo_4_5.yaml",
|
||||
"vllm_omni/deploy/minicpmo_4_5_duplex.yaml",
|
||||
"--trust-remote-code"
|
||||
]
|
||||
},
|
||||
|
||||
@@ -21,6 +21,47 @@
|
||||
]
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [
|
||||
4,
|
||||
16,
|
||||
32,
|
||||
64,
|
||||
128
|
||||
],
|
||||
"max_concurrency": [
|
||||
1,
|
||||
4,
|
||||
8,
|
||||
16,
|
||||
32
|
||||
],
|
||||
"random_input_len": 2500,
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
11701.3818,
|
||||
27510.3956,
|
||||
31907.9613,
|
||||
42173.9832,
|
||||
63970.0751
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1175,
|
||||
0.1647,
|
||||
0.2003,
|
||||
0.2726,
|
||||
0.4195
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
"backend": "openai-chat-omni",
|
||||
@@ -129,6 +170,45 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [
|
||||
4,
|
||||
16,
|
||||
32,
|
||||
64,
|
||||
128
|
||||
],
|
||||
"max_concurrency": [
|
||||
1,
|
||||
4,
|
||||
8,
|
||||
16,
|
||||
32
|
||||
],
|
||||
"random_input_len": 2500,
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"extra_body": {
|
||||
"modalities": [
|
||||
"text"
|
||||
]
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
4354.8945,
|
||||
7672.4248,
|
||||
9861.1687,
|
||||
13086.8865,
|
||||
18716.709
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
"backend": "openai-chat-omni",
|
||||
@@ -244,61 +324,5 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"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,12 +29,16 @@
|
||||
"endpoint": "/v1/completions",
|
||||
"num_prompts": [
|
||||
4,
|
||||
16,
|
||||
32,
|
||||
64,
|
||||
128
|
||||
],
|
||||
"max_concurrency": [
|
||||
1,
|
||||
4,
|
||||
8,
|
||||
16,
|
||||
32
|
||||
],
|
||||
"random_input_len": 2500,
|
||||
@@ -50,7 +54,9 @@
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
4619.0068,
|
||||
7868.3523,
|
||||
9760.6175,
|
||||
13422.3491,
|
||||
18007.6918
|
||||
]
|
||||
}
|
||||
@@ -93,6 +99,45 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
"backend": "vllm",
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"endpoint": "/v1/completions",
|
||||
"num_prompts": [
|
||||
40
|
||||
],
|
||||
"request_rate": [
|
||||
0.5
|
||||
],
|
||||
"random_input_len": 100,
|
||||
"random_output_len": 100,
|
||||
"random_range_ratio": 0.0,
|
||||
"ignore_eos": true,
|
||||
"extra_body": {
|
||||
"modalities": [
|
||||
"text"
|
||||
]
|
||||
},
|
||||
"random_mm_base_items_per_request": 2,
|
||||
"random_mm_num_mm_items_range_ratio": 0.5,
|
||||
"random_mm_limit_mm_per_prompt": {
|
||||
"image": 1,
|
||||
"video": 1
|
||||
},
|
||||
"random_mm_bucket_config": {
|
||||
"(256, 256, 1)": 0.5,
|
||||
"(720, 1280, 2)": 0.5
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
524.5321
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
"backend": "vllm",
|
||||
|
||||
@@ -74,14 +74,13 @@
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device step execution: sequential 512/1536 plus 512 concurrency sweep. One server (max-num-seqs 8).",
|
||||
"description": "Single-device baseline (no parallelism) with step execution",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"step-execution": true,
|
||||
"max-num-seqs": 8
|
||||
"step-execution": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
@@ -120,27 +119,6 @@
|
||||
"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]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
@@ -265,5 +243,53 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"step-execution": true,
|
||||
"max-num-seqs": 8
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20_high_concurrency",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": [20, 40, 80, 160],
|
||||
"max-concurrency": [1, 2, 4, 8],
|
||||
"warmup-requests": 8,
|
||||
"warmup-concurrency": 8,
|
||||
"warmup-num-inference-steps": 20,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": [0.3937, 0.6606, 0.6188, 0.7752],
|
||||
"latency_mean": [2.5466, 3.0235, 6.3948, 10.1944],
|
||||
"peak_memory_mb_mean": [55149.0, 55150.0, 55153.1458, 55156.325]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -157,8 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
|
||||
"ref_audio": None,
|
||||
"ref_text": None,
|
||||
},
|
||||
("Invalid voice", "vivian"),
|
||||
id="customvoice_invalid_voice",
|
||||
("CustomVoice", "does not support"),
|
||||
id="customvoice_task_not_supported",
|
||||
),
|
||||
pytest.param(
|
||||
{"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
|
||||
@@ -181,17 +181,13 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
|
||||
("mutually exclusive", "speaker_embedding", "ref_audio"),
|
||||
id="speaker_embedding_with_ref_audio_mutually_exclusive",
|
||||
),
|
||||
pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "greater_than_equal"), id="max_new_tokens_below_min"),
|
||||
pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
|
||||
pytest.param(
|
||||
{"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
|
||||
("max_new_tokens", "exceed 4096"),
|
||||
id="max_new_tokens_above_max",
|
||||
),
|
||||
pytest.param(
|
||||
{"seed": -(2**63) - 1},
|
||||
("seed", "greater_than_equal", "-9223372036854775808"),
|
||||
id="seed_below_min",
|
||||
),
|
||||
pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
|
||||
pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
|
||||
pytest.param(
|
||||
{"initial_codec_chunk_frames": -1},
|
||||
@@ -279,9 +275,7 @@ def test_speech_invalid_field_values(
|
||||
pytest.param({"input": "", "response_format": "wav"}, ("input", "empty"), id="input_empty"),
|
||||
pytest.param({"input": " ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"),
|
||||
pytest.param(
|
||||
{"input": "Hello world.", "max_new_tokens": 0},
|
||||
("max_new_tokens", "greater_than_equal"),
|
||||
id="max_new_tokens_below_min",
|
||||
{"input": "Hello world.", "max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
@@ -560,14 +554,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, online_client: Onli
|
||||
pytest.param(
|
||||
"batch",
|
||||
{"max_new_tokens": 0},
|
||||
("max_new_tokens", "greater_than_equal"),
|
||||
("max_new_tokens", "least 1"),
|
||||
id="batch_max_new_tokens_below_min",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
"item",
|
||||
{"max_new_tokens": 0},
|
||||
("max_new_tokens", "greater_than_equal"),
|
||||
("max_new_tokens", "least 1"),
|
||||
id="item_max_new_tokens_below_min",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
|
||||
@@ -72,9 +72,6 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
body["top_logprobs"] = 5
|
||||
elif case_id == "speaker_unknown":
|
||||
body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz"
|
||||
elif case_id == "audio_format_unsupported":
|
||||
body["modalities"] = ["text", "audio"]
|
||||
body["audio"] = {"voice": "alloy", "format": "aac"}
|
||||
else:
|
||||
raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
|
||||
return {"json": body, "timeout": 120}
|
||||
@@ -122,7 +119,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
pytest.param(
|
||||
"response_format_json_schema_incomplete",
|
||||
400,
|
||||
("response_format", "BadRequestError", "json_schema"),
|
||||
("response_format", "value_error", "json_schema"),
|
||||
id="invalid_response_format_json_schema",
|
||||
),
|
||||
pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
|
||||
@@ -138,12 +135,6 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
("Invalid speaker", "Supported"),
|
||||
id="speaker_unknown_preset",
|
||||
),
|
||||
pytest.param(
|
||||
"audio_format_unsupported",
|
||||
400,
|
||||
("Invalid audio format", "aac", "Supported formats"),
|
||||
id="unsupported_audio_format",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
|
||||
|
||||
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
|
||||
) -> None:
|
||||
"""Run the same prompts sequentially then in parallel and print a comparison."""
|
||||
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
|
||||
omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
|
||||
try:
|
||||
await warmup(omni, WARMUP_PROMPTS)
|
||||
single_time = await run_single(omni, prompts)
|
||||
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
|
||||
await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
|
||||
return
|
||||
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
|
||||
omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
|
||||
try:
|
||||
await warmup(omni, WARMUP_PROMPTS)
|
||||
|
||||
@@ -524,12 +524,7 @@ if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="E2E diffusion concurrent benchmark / validation")
|
||||
parser.add_argument("--model", type=str, required=True, help="Model name or path")
|
||||
parser.add_argument("--num-prompts", type=int, default=8, help="Number of prompts to run")
|
||||
parser.add_argument(
|
||||
"--batch-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size",
|
||||
)
|
||||
parser.add_argument("--batch-size", type=int, default=1, help="Diffusion batch size (1 = no batching)")
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=["batch", "single", "compare", "validate"],
|
||||
|
||||
+6
-3
@@ -365,10 +365,13 @@ class TestCacheDiTBackend:
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
|
||||
def test_enable_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Generic cache path should pick up ``transformer.fused_blocks``."""
|
||||
def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Test DreamID uses pipeline.transformer for cache enable/refresh.
|
||||
|
||||
NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
|
||||
"""
|
||||
mock_pipeline = Mock()
|
||||
mock_pipeline.__class__.__name__ = "FakeFusedBlocksPipeline"
|
||||
mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
|
||||
mock_pipeline.transformer = Mock()
|
||||
mock_pipeline.transformer.fused_blocks = Mock()
|
||||
mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
|
||||
|
||||
+10
@@ -6,6 +6,7 @@ Model specific tests for CacheDiT enablement.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
@@ -24,9 +25,18 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import L
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect
|
||||
# the class even though the dependency may not be set up correctly; this is ok for these
|
||||
# tests because we just inspect it and never initialize the model.
|
||||
for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"):
|
||||
sys.modules.setdefault(mod, Mock())
|
||||
# isort: split
|
||||
from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel # noqa: E402
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
SEPARATE_CFG_TRANSFORMERS = [
|
||||
DreamIdOmniModel,
|
||||
HeliosTransformer3DModel,
|
||||
LongCatImageTransformer2DModel,
|
||||
Cosmos3VFMTransformer,
|
||||
|
||||
@@ -340,97 +340,6 @@ def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.Mo
|
||||
assert captured["epsilon"].item() == pytest.approx(1e-5)
|
||||
|
||||
|
||||
# ── RMSNorm residual tests ──
|
||||
|
||||
|
||||
def test_rmsnorm_matches_native_residual_contract():
|
||||
"""The non-NPU path preserves separate residual-add and RMSNorm math."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
eps = 1e-6
|
||||
norm = RMSNorm(4, eps=eps, dtype=torch.bfloat16)
|
||||
norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
|
||||
x = torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.bfloat16)
|
||||
residual = torch.tensor([[0.00390625, 0.0078125, -0.0078125, 0.015625]], dtype=torch.bfloat16)
|
||||
|
||||
expected_residual = residual + x
|
||||
expected_fp32 = expected_residual.float()
|
||||
expected_output = expected_fp32 * torch.rsqrt(expected_fp32.square().mean(-1, keepdim=True) + eps)
|
||||
expected_output = (expected_output * norm.weight.float()).to(x.dtype)
|
||||
output, updated_residual = norm.forward_native(x, residual)
|
||||
|
||||
torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
|
||||
torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"forward_name",
|
||||
["forward_cuda", "forward_hip", "forward_musa", "forward_xpu"],
|
||||
)
|
||||
def test_rmsnorm_non_npu_residual_delegates_to_native(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
forward_name: str,
|
||||
):
|
||||
"""Platform wrappers keep residual arithmetic in forward_native."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
|
||||
x = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
residual = torch.randn_like(x)
|
||||
expected_output = torch.randn_like(x)
|
||||
expected_residual = torch.randn_like(x)
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def forward_native(native_x: torch.Tensor, native_residual: torch.Tensor | None = None):
|
||||
captured.update(x=native_x, residual=native_residual)
|
||||
return expected_output, expected_residual
|
||||
|
||||
monkeypatch.setattr(norm, "forward_native", forward_native)
|
||||
|
||||
output, updated_residual = getattr(norm, forward_name)(x, residual)
|
||||
|
||||
assert output is expected_output
|
||||
assert updated_residual is expected_residual
|
||||
assert captured["x"] is x
|
||||
assert captured["residual"] is residual
|
||||
|
||||
|
||||
def test_rmsnorm_npu_residual_uses_torch_npu_fused_op(monkeypatch: pytest.MonkeyPatch):
|
||||
"""NPU keeps the direct fused AddRMSNorm path instead of IR fallback."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
fused_output = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
fused_residual = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
|
||||
def npu_add_rms_norm(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
gamma: torch.Tensor,
|
||||
epsilon: float,
|
||||
):
|
||||
captured.update(x=x, residual=residual, gamma=gamma, epsilon=epsilon)
|
||||
return fused_output, torch.empty(0), fused_residual
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"torch_npu",
|
||||
types.SimpleNamespace(npu_add_rms_norm=npu_add_rms_norm),
|
||||
)
|
||||
norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
|
||||
x = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
residual = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
|
||||
output, updated_residual = norm.forward_npu(x, residual)
|
||||
|
||||
assert output is fused_output
|
||||
assert updated_residual is fused_residual
|
||||
assert captured["x"] is x
|
||||
assert captured["residual"] is residual
|
||||
assert captured["gamma"] is norm.weight
|
||||
assert captured["epsilon"] == 1e-5
|
||||
|
||||
|
||||
# ── RMSNorm compile-path regression tests ──
|
||||
|
||||
|
||||
|
||||
@@ -1,34 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.cuda,
|
||||
pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA platform required"),
|
||||
]
|
||||
|
||||
|
||||
def test_rmsnorm_cuda_residual_matches_native_contract() -> None:
|
||||
eps = 1e-6
|
||||
device = torch.device(current_omni_platform.device_type)
|
||||
torch.manual_seed(0)
|
||||
norm = RMSNorm(64, eps=eps, dtype=torch.bfloat16).to(device)
|
||||
x = torch.randn(2, 4, 64, device=device, dtype=torch.bfloat16)
|
||||
residual = torch.randn_like(x)
|
||||
x_before = x.clone()
|
||||
residual_before = residual.clone()
|
||||
|
||||
expected_output, expected_residual = norm.forward_native(x_before, residual_before)
|
||||
output, updated_residual = norm.forward_cuda(x, residual)
|
||||
|
||||
torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
|
||||
torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
|
||||
torch.testing.assert_close(x, x_before, atol=0, rtol=0)
|
||||
torch.testing.assert_close(residual, residual_before, atol=0, rtol=0)
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -89,29 +89,6 @@ def test_diffusion_base_linear_apply_multi_slice():
|
||||
assert torch.allclose(out, expected)
|
||||
|
||||
|
||||
def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
|
||||
from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
|
||||
|
||||
layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
|
||||
torch.nn.Module.__init__(layer)
|
||||
layer.base_layer = torch.nn.Module()
|
||||
layer.n_slices = 1
|
||||
layer.lora_a_stacked = (torch.zeros(2),)
|
||||
layer.lora_b_stacked = (torch.zeros(3),)
|
||||
|
||||
layer._set_diffusion_lora_buffer_device(torch.device("meta"))
|
||||
assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
|
||||
|
||||
def recreate_on_cpu(self, max_loras, lora_config, model_config):
|
||||
self.lora_a_stacked = (torch.zeros(4),)
|
||||
self.lora_b_stacked = (torch.zeros(5),)
|
||||
|
||||
monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu)
|
||||
layer.create_lora_weights(max_loras=1, lora_config=object())
|
||||
|
||||
assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
|
||||
|
||||
|
||||
def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
|
||||
# Verify that after reset_lora(), apply() skips LoRA matmuls even if the
|
||||
# LoRA tensors are still allocated and non-empty.
|
||||
|
||||
@@ -3,8 +3,6 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.lora.lora_weights import LoRALayerWeights
|
||||
@@ -145,37 +143,6 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
|
||||
assert replace_calls == ["foo"]
|
||||
|
||||
|
||||
def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch):
|
||||
import vllm_omni.diffusion.lora.manager as manager_mod
|
||||
|
||||
class _DLOLoRALayer(DummyBaseLayerWithLoRA):
|
||||
lora_buffer_device = None
|
||||
|
||||
def _set_diffusion_lora_buffer_device(self, device):
|
||||
self.lora_buffer_device = device
|
||||
|
||||
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer)
|
||||
monkeypatch.setattr(
|
||||
manager_mod,
|
||||
"from_layer_diffusion",
|
||||
lambda *, layer, **_kwargs: _DLOLoRALayer(layer),
|
||||
)
|
||||
|
||||
pipeline = torch.nn.Module()
|
||||
pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True)
|
||||
pipeline.transformer = torch.nn.Module()
|
||||
pipeline.transformer.foo = _FakeLinearBase()
|
||||
manager = DiffusionLoRAManager(
|
||||
pipeline=pipeline,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
manager._replace_layers_with_lora(type("_PH", (), {"r": 1})())
|
||||
|
||||
assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu")
|
||||
|
||||
|
||||
def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
|
||||
import vllm_omni.diffusion.lora.manager as manager_mod
|
||||
|
||||
|
||||
@@ -24,7 +24,6 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
|
||||
HostWeightPlanResult,
|
||||
TensorBinding,
|
||||
)
|
||||
from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
|
||||
from vllm_omni.diffusion.models.helios import HeliosPipeline
|
||||
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
|
||||
from vllm_omni.diffusion.registry import initialize_model
|
||||
@@ -174,15 +173,6 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
|
||||
str(canonical_root / "model.safetensors"),
|
||||
)
|
||||
store_root = tmp_path / "hwr-store"
|
||||
hash_calls = 0
|
||||
original_sha256 = source_identity_module._sha256_file
|
||||
|
||||
def counted_sha256(path: Path, state: object) -> str:
|
||||
nonlocal hash_calls
|
||||
hash_calls += 1
|
||||
return original_sha256(path, state) # type: ignore[arg-type]
|
||||
|
||||
monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
|
||||
|
||||
def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
|
||||
@@ -215,8 +205,6 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
|
||||
assert warm_plan is not None
|
||||
assert warm_plan.lease_carrier is not None
|
||||
warm_plan.lease_carrier.close()
|
||||
assert hash_calls == 1
|
||||
assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1
|
||||
|
||||
|
||||
def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch):
|
||||
@@ -249,30 +237,6 @@ def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_p
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
|
||||
tmp_path: Path,
|
||||
monkeypatch,
|
||||
):
|
||||
canonical_root = tmp_path / "canonical"
|
||||
canonical_root.mkdir()
|
||||
save_file(
|
||||
{"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
|
||||
str(canonical_root / "model.safetensors"),
|
||||
)
|
||||
loader = DiffusersPipelineLoader(
|
||||
LoadConfig(),
|
||||
_hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
|
||||
)
|
||||
pipeline = _HWRPipeline(canonical_root)
|
||||
monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
|
||||
loader.load_model(load_device="cpu", device=torch.device("cpu"))
|
||||
|
||||
assert pipeline.load_count == 0
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
|
||||
return OmniDiffusionConfig(
|
||||
model="",
|
||||
@@ -285,11 +249,10 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("dist_offload", "use_allgather", "mode"),
|
||||
("dist_offload", "use_allgather"),
|
||||
[
|
||||
(False, False, "preferred"),
|
||||
(True, True, "preferred"),
|
||||
(True, False, "disabled"),
|
||||
(False, False),
|
||||
(True, True),
|
||||
],
|
||||
)
|
||||
def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
|
||||
@@ -297,13 +260,12 @@ def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
|
||||
tmp_path,
|
||||
dist_offload,
|
||||
use_allgather,
|
||||
mode,
|
||||
):
|
||||
"""Disabled and AllGather paths must never construct or probe HWR."""
|
||||
from vllm_omni.host_weight_runtime import HostWeightRuntime
|
||||
|
||||
root = tmp_path / "must-not-be-touched"
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root))
|
||||
model = _DummyPipelineModel(source_prefix="transformer.")
|
||||
modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
|
||||
|
||||
|
||||
@@ -6,8 +6,6 @@ from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import gc
|
||||
import json
|
||||
import threading
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
@@ -26,12 +24,10 @@ from vllm_omni.diffusion.model_loader.host_weights import (
|
||||
FinalLayoutRequest,
|
||||
FinalLayoutTensorRestorer,
|
||||
ImplementationIdentity,
|
||||
NodeSourceDigestCache,
|
||||
PreparedWeightSource,
|
||||
WeightSourceKind,
|
||||
build_final_layout_identity,
|
||||
)
|
||||
from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
|
||||
from vllm_omni.diffusion.model_loader.host_weights.contracts import (
|
||||
FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
|
||||
FinalLayoutArtifactSpec,
|
||||
@@ -303,83 +299,6 @@ class _FakeH3Attention(nn.Module):
|
||||
super().__init__()
|
||||
|
||||
|
||||
class _FakeFluxLinear(_FakeH3Linear):
|
||||
def __init__(
|
||||
self,
|
||||
*args: object,
|
||||
bias: bool = False,
|
||||
params_dtype: torch.dtype | None = None,
|
||||
total_num_heads: int | None = None,
|
||||
total_num_kv_heads: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
*args,
|
||||
bias=bias,
|
||||
params_dtype=params_dtype,
|
||||
total_num_heads=total_num_heads,
|
||||
total_num_kv_heads=total_num_kv_heads,
|
||||
**kwargs,
|
||||
)
|
||||
self.weight.weight_loader = self._load_weight
|
||||
|
||||
@staticmethod
|
||||
def _load_weight(
|
||||
parameter: nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
shard_id: str | None = None,
|
||||
) -> None:
|
||||
del shard_id
|
||||
parameter.data.copy_(loaded_weight.to(parameter).reshape_as(parameter))
|
||||
|
||||
|
||||
class _FakeFluxAttention(nn.Module):
|
||||
def __init__(self, **kwargs: object) -> None:
|
||||
del kwargs
|
||||
super().__init__()
|
||||
self.register_buffer("beta", torch.tensor([0.5]))
|
||||
self.register_buffer("eps", torch.tensor([-1e-6]))
|
||||
|
||||
|
||||
class _Flux2KleinPipeline(nn.Module):
|
||||
def __init__(self, transformer: nn.Module) -> None:
|
||||
super().__init__()
|
||||
self.transformer = transformer
|
||||
self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
|
||||
self.vae = nn.Module()
|
||||
self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))
|
||||
|
||||
|
||||
def _reduced_flux2_klein(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
device: torch.device | str = "cpu",
|
||||
) -> nn.Module:
|
||||
from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as flux2
|
||||
|
||||
monkeypatch.setattr(flux2, "ColumnParallelLinear", _FakeFluxLinear)
|
||||
monkeypatch.setattr(flux2, "MergedColumnParallelLinear", _FakeFluxLinear)
|
||||
monkeypatch.setattr(flux2, "QKVParallelLinear", _FakeFluxLinear)
|
||||
monkeypatch.setattr(flux2, "RowParallelLinear", _FakeFluxLinear)
|
||||
monkeypatch.setattr(flux2, "Attention", _FakeFluxAttention)
|
||||
|
||||
transformer = flux2.Flux2Transformer2DModel(
|
||||
patch_size=1,
|
||||
in_channels=8,
|
||||
out_channels=8,
|
||||
num_layers=1,
|
||||
num_single_layers=1,
|
||||
attention_head_dim=8,
|
||||
num_attention_heads=1,
|
||||
joint_attention_dim=8,
|
||||
timestep_guidance_channels=8,
|
||||
mlp_ratio=1.0,
|
||||
axes_dims_rope=(2, 2, 2, 2),
|
||||
guidance_embeds=False,
|
||||
)
|
||||
return transformer.to(device=device, dtype=torch.bfloat16)
|
||||
|
||||
|
||||
class _H3Pipeline(nn.Module):
|
||||
def __init__(self, transformer: nn.Module) -> None:
|
||||
super().__init__()
|
||||
@@ -461,7 +380,6 @@ def _identity(
|
||||
*,
|
||||
request: FinalLayoutRequest | None = None,
|
||||
policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY,
|
||||
source_digest_cache: NodeSourceDigestCache | None = None,
|
||||
) -> FinalLayoutIdentityContext:
|
||||
transformer = model.get_submodule("transformer")
|
||||
return build_final_layout_identity(
|
||||
@@ -470,7 +388,6 @@ def _identity(
|
||||
prepared_sources=(source,),
|
||||
request=request or _request(),
|
||||
policy=policy,
|
||||
source_digest_cache=source_digest_cache,
|
||||
)
|
||||
|
||||
|
||||
@@ -647,131 +564,6 @@ def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path
|
||||
assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID
|
||||
|
||||
|
||||
def test_node_source_digest_cache_reuses_recovers_and_invalidates(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
model = _TinyPipeline()
|
||||
source = _prepared_source(tmp_path)
|
||||
cache_root = tmp_path / "store"
|
||||
calls = 0
|
||||
original_sha256 = source_identity_module._sha256_file
|
||||
|
||||
def counted_sha256(path: Path, state: object) -> str:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return original_sha256(path, state) # type: ignore[arg-type]
|
||||
|
||||
monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
|
||||
first = _identity(
|
||||
model,
|
||||
source,
|
||||
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
|
||||
)
|
||||
second = _identity(
|
||||
model,
|
||||
source,
|
||||
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
|
||||
)
|
||||
|
||||
assert second.identity == first.identity
|
||||
assert calls == 1
|
||||
|
||||
cache_entry = next((cache_root / "source-digests-v1" / "entries").glob("*.json"))
|
||||
cache_entry.write_text("not-json", encoding="utf-8")
|
||||
recovered = _identity(
|
||||
model,
|
||||
source,
|
||||
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
|
||||
)
|
||||
assert recovered.identity == first.identity
|
||||
assert calls == 2
|
||||
|
||||
cache_document = json.loads(cache_entry.read_text(encoding="utf-8"))
|
||||
cache_document["content_id"] = f"sha256:{'0' * 64}"
|
||||
cache_entry.write_text(json.dumps(cache_document), encoding="utf-8")
|
||||
checksum_recovered = _identity(
|
||||
model,
|
||||
source,
|
||||
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
|
||||
)
|
||||
assert checksum_recovered.identity == first.identity
|
||||
assert calls == 3
|
||||
|
||||
source.weight_files[0].write_bytes(b"changed-canonical-source-with-a-new-size")
|
||||
changed = _identity(
|
||||
model,
|
||||
source,
|
||||
source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
|
||||
)
|
||||
assert changed.identity != first.identity
|
||||
assert calls == 4
|
||||
|
||||
|
||||
def test_source_identity_hashes_checkpoint_shards_in_parallel(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
model = _TinyPipeline()
|
||||
source = _prepared_source(tmp_path)
|
||||
second_file = source.resolved_root / "model-00002.safetensors"
|
||||
second_file.write_bytes(b"second-checkpoint-shard")
|
||||
source = dataclasses.replace(source, weight_files=(*source.weight_files, second_file))
|
||||
barrier = threading.Barrier(2)
|
||||
thread_ids: set[int] = set()
|
||||
thread_ids_lock = threading.Lock()
|
||||
original_sha256 = source_identity_module._sha256_file
|
||||
|
||||
def synchronized_sha256(path: Path, state: object) -> str:
|
||||
with thread_ids_lock:
|
||||
thread_ids.add(threading.get_ident())
|
||||
barrier.wait(timeout=5)
|
||||
return original_sha256(path, state) # type: ignore[arg-type]
|
||||
|
||||
monkeypatch.setattr(source_identity_module, "_sha256_file", synchronized_sha256)
|
||||
_identity(model, source)
|
||||
|
||||
assert len(thread_ids) == 2
|
||||
|
||||
|
||||
def test_source_digest_cache_lock_timeout_falls_back_to_direct_hashing(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
model = _TinyPipeline()
|
||||
source = _prepared_source(tmp_path)
|
||||
cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
|
||||
|
||||
def unavailable_lock(*_args: object, **_kwargs: object) -> object:
|
||||
raise source_identity_module.FileLockTimeoutError("busy")
|
||||
|
||||
monkeypatch.setattr(source_identity_module, "FileLock", unavailable_lock)
|
||||
context = _identity(model, source, source_digest_cache=cache)
|
||||
|
||||
assert context.identity.source.revision.startswith("content-")
|
||||
assert not tuple(cache.entries.glob("*.json"))
|
||||
|
||||
|
||||
def test_source_digest_cache_directory_failure_falls_back_to_direct_hashing(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
model = _TinyPipeline()
|
||||
source = _prepared_source(tmp_path)
|
||||
original_mkdir = Path.mkdir
|
||||
|
||||
def failing_cache_mkdir(path: Path, *args: object, **kwargs: object) -> None:
|
||||
if "source-digests-v1" in path.parts:
|
||||
raise OSError("injected cache directory failure")
|
||||
original_mkdir(path, *args, **kwargs) # type: ignore[arg-type]
|
||||
|
||||
monkeypatch.setattr(Path, "mkdir", failing_cache_mkdir)
|
||||
cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
|
||||
context = _identity(model, source, source_digest_cache=cache)
|
||||
|
||||
assert context.identity.source.revision.startswith("content-")
|
||||
|
||||
|
||||
def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None:
|
||||
model = _TinyPipeline()
|
||||
source_root = tmp_path / "local-symlink-source"
|
||||
@@ -1277,152 +1069,3 @@ def test_minimax_h3_declares_the_final_layout_restore_contract() -> None:
|
||||
assert contract.implementation_id == "minimax-h3-dit"
|
||||
assert callable(MiniMaxH3DiTModel.validate_restored_host_weights)
|
||||
assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__
|
||||
|
||||
|
||||
def test_reduced_flux2_klein_declares_complete_final_layout_contract(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
|
||||
Flux2Transformer2DModel,
|
||||
)
|
||||
|
||||
transformer = _reduced_flux2_klein(monkeypatch)
|
||||
pipeline = _Flux2KleinPipeline(transformer)
|
||||
source = _prepared_source(tmp_path, directory="flux2-klein-source")
|
||||
context = _identity(pipeline, source)
|
||||
targets = collect_final_layout_targets(
|
||||
pipeline,
|
||||
(("transformer", transformer),),
|
||||
policy=FINAL_LAYOUT_BF16_POLICY,
|
||||
require_materialized=True,
|
||||
)
|
||||
target_names = {target.name for target in targets}
|
||||
|
||||
contract = Flux2Transformer2DModel.host_weight_restore_contract
|
||||
assert isinstance(contract, FinalLayoutModelContract)
|
||||
assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
|
||||
assert contract.implementation_id == "flux2-klein-dit"
|
||||
assert contract.version == "1"
|
||||
assert context.model_contracts == (contract,)
|
||||
assert transformer.stacked_params_mapping == [
|
||||
(".to_qkv.", ".to_q.", "q"),
|
||||
(".to_qkv.", ".to_k.", "k"),
|
||||
(".to_qkv.", ".to_v.", "v"),
|
||||
(".add_kv_proj", ".add_q_proj", "q"),
|
||||
(".add_kv_proj", ".add_k_proj", "k"),
|
||||
(".add_kv_proj", ".add_v_proj", "v"),
|
||||
]
|
||||
loaded = transformer.load_weights(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.weight", torch.tensor([1.0], dtype=torch.bfloat16)),
|
||||
("transformer_blocks.0.attn.add_q_proj.weight", torch.tensor([2.0], dtype=torch.bfloat16)),
|
||||
(
|
||||
"single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
|
||||
torch.tensor([3.0], dtype=torch.bfloat16),
|
||||
),
|
||||
]
|
||||
)
|
||||
assert loaded == {
|
||||
"transformer_blocks.0.attn.to_qkv.weight",
|
||||
"transformer_blocks.0.attn.add_kv_proj.weight",
|
||||
"single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
|
||||
}
|
||||
assert any(name.startswith("transformer.transformer_blocks.0.") for name in target_names)
|
||||
assert any(name.startswith("transformer.single_transformer_blocks.0.") for name in target_names)
|
||||
assert "transformer.transformer_blocks.0.attn.to_qkv.weight" in target_names
|
||||
assert "transformer.transformer_blocks.0.attn.add_kv_proj.weight" in target_names
|
||||
assert "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight" in target_names
|
||||
assert "transformer.transformer_blocks.0.attn.attn.beta" in target_names
|
||||
assert "transformer.transformer_blocks.0.attn.attn.eps" in target_names
|
||||
assert all("text_encoder" not in name and "vae" not in name for name in target_names)
|
||||
transformer.validate_restored_host_weights()
|
||||
|
||||
|
||||
def test_flux2_klein_restore_validator_rejects_non_ready_state(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
transformer = _reduced_flux2_klein(monkeypatch)
|
||||
|
||||
mapping = transformer.stacked_params_mapping
|
||||
transformer.stacked_params_mapping = None
|
||||
with pytest.raises(ValueError, match="packed QKV mapping"):
|
||||
transformer.validate_restored_host_weights()
|
||||
transformer.stacked_params_mapping = mapping
|
||||
|
||||
weight = transformer.x_embedder.weight
|
||||
weight.data = weight.data.float()
|
||||
with pytest.raises(ValueError, match="must stay bf16"):
|
||||
transformer.validate_restored_host_weights()
|
||||
weight.data = weight.data.to(torch.bfloat16)
|
||||
|
||||
weight.data = weight.data.T
|
||||
with pytest.raises(ValueError, match="contiguous strided layout"):
|
||||
transformer.validate_restored_host_weights()
|
||||
weight.data = weight.data.contiguous()
|
||||
|
||||
attention = transformer.transformer_blocks[0].attn.attn
|
||||
attention._non_persistent_buffers_set.add("beta")
|
||||
with pytest.raises(ValueError, match="must be persistent"):
|
||||
transformer.validate_restored_host_weights()
|
||||
attention._non_persistent_buffers_set.remove("beta")
|
||||
|
||||
attention.beta.data = attention.beta.data.float()
|
||||
with pytest.raises(ValueError, match="scalar bf16"):
|
||||
transformer.validate_restored_host_weights()
|
||||
|
||||
|
||||
def test_flux2_klein_publication_restore_fallback_and_teardown(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
source = _prepared_source(tmp_path, directory="flux2-klein-source")
|
||||
cold_transformer = _reduced_flux2_klein(monkeypatch)
|
||||
cold_model = _Flux2KleinPipeline(cold_transformer)
|
||||
with torch.no_grad():
|
||||
for index, parameter in enumerate(cold_transformer.parameters(), start=1):
|
||||
parameter.fill_(index)
|
||||
for name, buffer in cold_transformer.named_buffers():
|
||||
if name.endswith((".beta", ".eps")):
|
||||
buffer.fill_(0.5 if name.endswith(".beta") else -1e-6)
|
||||
|
||||
context = _identity(cold_model, source)
|
||||
runtime = _runtime(tmp_path / "flux2-klein-store")
|
||||
miss = runtime.resolve(context.identity)
|
||||
assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
|
||||
|
||||
producer = FinalLayoutBF16Producer(
|
||||
context,
|
||||
cold_model,
|
||||
(("transformer", cold_transformer),),
|
||||
max_shard_bytes=64,
|
||||
)
|
||||
publication = runtime.publish_after_load(context.identity, producer=producer)
|
||||
assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
|
||||
|
||||
warm_transformer = _reduced_flux2_klein(monkeypatch, device="meta")
|
||||
warm_model = _Flux2KleinPipeline(warm_transformer)
|
||||
warm_context = _identity(warm_model, source)
|
||||
assert warm_context.identity == context.identity
|
||||
hit = runtime.resolve(warm_context.identity)
|
||||
assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
|
||||
assert hit.lease is not None
|
||||
|
||||
plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
|
||||
assert all(tensor.is_meta for name, tensor in warm_model.named_parameters() if name.startswith("transformer."))
|
||||
plan.commit()
|
||||
warm_transformer.validate_restored_host_weights()
|
||||
|
||||
cold_tensors = dict(cold_model.named_parameters()) | dict(cold_model.named_buffers())
|
||||
warm_tensors = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
|
||||
for name in warm_context.tensor_names:
|
||||
assert torch.equal(warm_tensors[name], cold_tensors[name])
|
||||
assert warm_tensors[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()
|
||||
assert warm_transformer.stacked_params_mapping == cold_transformer.stacked_params_mapping
|
||||
|
||||
del warm_tensors, warm_model, warm_transformer, plan
|
||||
gc.collect()
|
||||
hit.lease.close()
|
||||
assert hit.lease.closed
|
||||
assert isinstance(runtime.store, FilesystemHostWeightStore)
|
||||
assert runtime.store.cleanup(context.identity) is None
|
||||
@@ -0,0 +1,27 @@
|
||||
import pytest
|
||||
import torch.nn as nn
|
||||
|
||||
pytest.importorskip("dreamid_omni.modules.model")
|
||||
|
||||
from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks():
|
||||
model = object.__new__(FusionModel)
|
||||
nn.Module.__init__(model)
|
||||
model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
|
||||
model.attn = nn.Identity()
|
||||
|
||||
conditions = getattr(model, "_hsdp_shard_conditions", None)
|
||||
|
||||
assert conditions is not None
|
||||
assert len(conditions) == 1
|
||||
|
||||
matched = []
|
||||
for name, module in model.named_modules():
|
||||
if any(cond(name, module) for cond in conditions):
|
||||
matched.append(name)
|
||||
|
||||
assert matched == ["fused_blocks.0", "fused_blocks.1"]
|
||||
@@ -964,80 +964,6 @@ def test_minimax_h3_advertises_the_official_ref2va_image_limit():
|
||||
assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9
|
||||
|
||||
|
||||
def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
|
||||
|
||||
class FakeEncoderGroup:
|
||||
rank_in_group = 0
|
||||
world_size = 2
|
||||
|
||||
def all_reduce(self, tensor):
|
||||
del tensor
|
||||
|
||||
config = SimpleNamespace(
|
||||
hidden_size=8,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
head_dim=2,
|
||||
rms_norm_eps=1e-6,
|
||||
)
|
||||
with (
|
||||
patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
|
||||
patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
|
||||
):
|
||||
attention = encoder_module.MiniMaxH3Qwen3VLTextAttention(FakeEncoderGroup(), config, torch.float32)
|
||||
# This CPU contract test only needs projection shapes and head reshaping.
|
||||
# Avoid dispatching platform-specific norm and linear kernels (for example,
|
||||
# rocm_unquantized_gemm when the test suite is collected on ROCm).
|
||||
attention.q_norm = nn.Identity()
|
||||
attention.k_norm = nn.Identity()
|
||||
with torch.no_grad():
|
||||
attention.qkv_proj.weight.zero_()
|
||||
attention.o_proj = nn.Linear(
|
||||
attention.qkv_proj.local_num_heads * attention.head_dim,
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
)
|
||||
attn_call = {}
|
||||
|
||||
def fake_attention(query, key, value):
|
||||
attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape)
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", fake_attention)
|
||||
hidden_states = torch.randn(1, 3, config.hidden_size)
|
||||
cos = torch.ones(1, 3, config.head_dim)
|
||||
sin = torch.zeros_like(cos)
|
||||
|
||||
output = attention(hidden_states, (cos, sin))
|
||||
|
||||
assert attn_call["query_shape"] == (1, config.num_attention_heads // 2, 3, config.head_dim)
|
||||
assert attn_call["key_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
|
||||
assert attn_call["value_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
|
||||
assert output.shape == hidden_states.shape
|
||||
|
||||
|
||||
def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
|
||||
|
||||
sdpa_call = {}
|
||||
|
||||
def fake_sdpa(query, key, value, **kwargs):
|
||||
sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs)
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
|
||||
output = encoder_module._scaled_dot_product_attention(
|
||||
torch.randn(1, 4, 3, 8),
|
||||
torch.randn(1, 2, 3, 8),
|
||||
torch.randn(1, 2, 3, 8),
|
||||
)
|
||||
|
||||
assert sdpa_call["query_shape"] == sdpa_call["key_shape"] == sdpa_call["value_shape"] == (1, 4, 3, 8)
|
||||
assert sdpa_call["kwargs"] == {"dropout_p": 0.0, "is_causal": True}
|
||||
assert output.shape == (1, 4, 3, 8)
|
||||
|
||||
|
||||
def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import (
|
||||
MiniMaxH3Qwen3VLEncoder,
|
||||
|
||||
@@ -197,6 +197,7 @@ def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
|
||||
[
|
||||
"model-level CPU offload (--enable-cpu-offload)",
|
||||
"layerwise offload (--enable-layerwise-offload)",
|
||||
"distributed layerwise offload (--enable-distributed-layerwise-offload)",
|
||||
],
|
||||
)
|
||||
def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
|
||||
@@ -213,38 +214,6 @@ def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
|
||||
)
|
||||
|
||||
|
||||
def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
from vllm_omni.diffusion.models.minimax_h3 import (
|
||||
pipeline_minimax_h3 as pipeline_module,
|
||||
)
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.partition = "fl2va"
|
||||
pipeline._turbo_lora_adapter_ids = set()
|
||||
pipeline.od_config = SimpleNamespace(
|
||||
enable_cpu_offload=False,
|
||||
enable_layerwise_offload=False,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
)
|
||||
captured = {}
|
||||
|
||||
def load_turbo(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return object(), object()
|
||||
|
||||
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
|
||||
loaded = pipeline._load_diffusion_lora_adapter(
|
||||
lora_request=_request("turbo"),
|
||||
lora_path="turbo",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
assert loaded is not None
|
||||
assert captured["unsupported_offload_mode"] is None
|
||||
|
||||
|
||||
def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
|
||||
unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
|
||||
_write_tiny_turbo(unsupported)
|
||||
|
||||
@@ -1,132 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
def test_minimax_h3_npu_sdpa_patch_is_independent_and_idempotent(monkeypatch) -> None:
|
||||
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
def sdpa_sentinel(query, key, value):
|
||||
del key, value
|
||||
return query
|
||||
|
||||
def rope_sentinel(q, k, cos, sin):
|
||||
del cos, sin
|
||||
return q, k
|
||||
|
||||
monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
|
||||
monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
|
||||
monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
|
||||
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
|
||||
|
||||
npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
|
||||
|
||||
assert encoder_module._scaled_dot_product_attention is npu_minimax_h3._scaled_dot_product_attention_npu
|
||||
assert encoder_module._apply_rotary_pos_emb is rope_sentinel
|
||||
assert not npu_minimax_h3._ROPE_PATCHED
|
||||
assert npu_minimax_h3._SDPA_PATCHED
|
||||
|
||||
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
|
||||
npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
|
||||
|
||||
assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
@pytest.mark.parametrize(
|
||||
("num_heads", "num_kv_heads", "expected_enable_gqa"),
|
||||
[(2, 2, False), (4, 2, True)],
|
||||
)
|
||||
def test_minimax_h3_npu_sdpa_preserves_compressed_heads_and_selects_gqa(
|
||||
monkeypatch,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
expected_enable_gqa: bool,
|
||||
) -> None:
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_sdpa(query, key, value, **kwargs):
|
||||
captured.update(query=query, key=key, value=value, kwargs=kwargs)
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(F, "scaled_dot_product_attention", fake_sdpa)
|
||||
query = torch.randn(1, num_heads, 3, 8)
|
||||
key = torch.randn(1, num_kv_heads, 3, 8)
|
||||
value = torch.randn_like(key)
|
||||
|
||||
output = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
|
||||
|
||||
assert captured["query"] is query
|
||||
assert captured["key"] is key
|
||||
assert captured["value"] is value
|
||||
assert captured["kwargs"] == {
|
||||
"dropout_p": 0.0,
|
||||
"is_causal": True,
|
||||
"enable_gqa": expected_enable_gqa,
|
||||
}
|
||||
assert output is query
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
@pytest.mark.parametrize(
|
||||
("num_heads", "num_key_heads", "num_value_heads", "message"),
|
||||
[
|
||||
(4, 2, 1, "key and value to have the same number of heads"),
|
||||
(4, 0, 0, "at least one KV head"),
|
||||
(3, 2, 2, "query heads to be a multiple of KV heads"),
|
||||
],
|
||||
)
|
||||
def test_minimax_h3_npu_sdpa_rejects_invalid_head_configurations(
|
||||
num_heads: int,
|
||||
num_key_heads: int,
|
||||
num_value_heads: int,
|
||||
message: str,
|
||||
) -> None:
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
query = torch.randn(1, num_heads, 3, 8)
|
||||
key = torch.randn(1, num_key_heads, 3, 8)
|
||||
value = torch.randn(1, num_value_heads, 3, 8)
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
def test_minimax_h3_npu_causal_gqa_matches_expanded_kv() -> None:
|
||||
"""Exercise the real torch-npu causal GQA operator."""
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
torch.manual_seed(0)
|
||||
batch_size, seq_len = 1, 16
|
||||
num_heads, num_kv_heads, head_size = 8, 2, 128
|
||||
device = torch.device("npu")
|
||||
dtype = torch.bfloat16
|
||||
|
||||
query = torch.randn(batch_size, num_heads, seq_len, head_size, device=device, dtype=dtype)
|
||||
key = torch.randn(batch_size, num_kv_heads, seq_len, head_size, device=device, dtype=dtype)
|
||||
value = torch.randn_like(key)
|
||||
|
||||
actual = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
|
||||
|
||||
repeat_num = num_heads // num_kv_heads
|
||||
expected = F.scaled_dot_product_attention(
|
||||
query,
|
||||
key.repeat_interleave(repeat_num, dim=1),
|
||||
value.repeat_interleave(repeat_num, dim=1),
|
||||
dropout_p=0.0,
|
||||
is_causal=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
|
||||
@@ -6,38 +6,7 @@ import torch
|
||||
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
def test_qwen3_vl_npu_rope_patch_is_independent_and_idempotent(monkeypatch) -> None:
|
||||
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
def rope_sentinel(q, k, cos, sin):
|
||||
del cos, sin
|
||||
return q, k
|
||||
|
||||
def sdpa_sentinel(query, key, value):
|
||||
del key, value
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
|
||||
monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
|
||||
monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
|
||||
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
|
||||
|
||||
npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
|
||||
|
||||
assert encoder_module._apply_rotary_pos_emb is npu_minimax_h3._apply_rotary_pos_emb_npu
|
||||
assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
|
||||
assert npu_minimax_h3._ROPE_PATCHED
|
||||
assert not npu_minimax_h3._SDPA_PATCHED
|
||||
|
||||
monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
|
||||
npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
|
||||
|
||||
assert encoder_module._apply_rotary_pos_emb is rope_sentinel
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
|
||||
@@ -1,20 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import (
|
||||
MiniMaxH3Qwen3VLRMSNorm,
|
||||
MiniMaxH3Qwen3VLTextDecoderLayer,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm
|
||||
|
||||
|
||||
def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
|
||||
@@ -34,43 +24,3 @@ def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
|
||||
assert norm.weight.dtype == torch.bfloat16
|
||||
assert set(norm.state_dict()) == {"weight"}
|
||||
torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)
|
||||
|
||||
|
||||
def test_qwen3_vl_decoder_preserves_post_attention_residual_contract() -> None:
|
||||
class ConstantAttention(nn.Module):
|
||||
def forward(self, hidden_states, position_embeddings):
|
||||
del position_embeddings
|
||||
return torch.full_like(hidden_states, 2.0)
|
||||
|
||||
config = SimpleNamespace(
|
||||
hidden_size=8,
|
||||
intermediate_size=16,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
head_dim=2,
|
||||
rms_norm_eps=1e-6,
|
||||
)
|
||||
group = SimpleNamespace(rank_in_group=0, world_size=1)
|
||||
with (
|
||||
patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
|
||||
patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
|
||||
):
|
||||
layer = MiniMaxH3Qwen3VLTextDecoderLayer(group, config, torch.float32)
|
||||
layer.input_layernorm = nn.Identity()
|
||||
layer.self_attn = ConstantAttention()
|
||||
layer.mlp = nn.Identity()
|
||||
|
||||
hidden_states = torch.arange(24, dtype=torch.float32).reshape(1, 3, 8)
|
||||
attention_output = torch.full_like(hidden_states, 2.0)
|
||||
expected_residual = hidden_states + attention_output
|
||||
expected_normalized = RMSNorm.forward_native(layer.post_attention_layernorm, expected_residual)
|
||||
|
||||
with patch.object(
|
||||
layer.post_attention_layernorm,
|
||||
"forward",
|
||||
wraps=layer.post_attention_layernorm.forward_native,
|
||||
):
|
||||
output = layer(hidden_states, (torch.empty(0), torch.empty(0)))
|
||||
|
||||
assert isinstance(layer.post_attention_layernorm, MiniMaxH3Qwen3VLRMSNorm)
|
||||
torch.testing.assert_close(output, expected_residual + expected_normalized)
|
||||
@@ -1,250 +0,0 @@
|
||||
# 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,7 +5,6 @@
|
||||
|
||||
import gc
|
||||
import json
|
||||
import weakref
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -18,7 +17,6 @@ from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
|
||||
|
||||
import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
|
||||
from tests.helpers.runtime import get_distributed_init_method
|
||||
from vllm_omni.diffusion.data import validate_dlo_host_registration_options
|
||||
from vllm_omni.diffusion.model_loader.host_weight_plan import (
|
||||
HostWeightPlan,
|
||||
TensorBinding,
|
||||
@@ -35,17 +33,12 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
|
||||
DistributedLayerwiseOffloadHook,
|
||||
PinnedResidentLayerGroup,
|
||||
)
|
||||
from vllm_omni.diffusion.offloader.host_registration import (
|
||||
HostRegistrationCleanupError,
|
||||
HostRegistrationError,
|
||||
)
|
||||
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
|
||||
from vllm_omni.diffusion.offloader.offload_plan import (
|
||||
OffloadPlan,
|
||||
get_offload_plan,
|
||||
)
|
||||
from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
|
||||
from vllm_omni.host_weight_runtime import MappedHostRegion
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
@@ -710,15 +703,11 @@ class _HWRPipeline(nn.Module):
|
||||
|
||||
|
||||
class _FakeHostWeightLease:
|
||||
def __init__(self, resolution_id: str, events: list[str] | None = None):
|
||||
def __init__(self, resolution_id: str):
|
||||
self.closed = False
|
||||
self.provenance = SimpleNamespace(resolution_id=resolution_id)
|
||||
self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),)
|
||||
self._events = events
|
||||
|
||||
def close(self):
|
||||
if self._events is not None:
|
||||
self._events.append("lease")
|
||||
self.closed = True
|
||||
|
||||
|
||||
@@ -768,206 +757,6 @@ def _hwr_backend():
|
||||
return _HWRPipeline(), backend, carrier, lease
|
||||
|
||||
|
||||
def test_registered_mmap_hook_bypasses_host_staging(patched_offload_runtime):
|
||||
current_block = nn.Linear(2, 2, bias=False)
|
||||
next_block = nn.Linear(2, 2, bias=False)
|
||||
expected = torch.arange(4, dtype=torch.float32).view(2, 2)
|
||||
next_block.weight.data.copy_(expected)
|
||||
hook = DistributedLayerwiseOffloadHook(
|
||||
next_block=next_block,
|
||||
device=torch.device("cpu"),
|
||||
dp_group=None,
|
||||
dp_size=1,
|
||||
rank=0,
|
||||
pin_memory=False,
|
||||
rank_local_mmap=True,
|
||||
)
|
||||
hook.initialize_hook(current_block)
|
||||
hook.registered_mmap = True
|
||||
hook._stage_mmap_sources = lambda _slot: pytest.fail("registered mmap must bypass host staging") # type: ignore[method-assign]
|
||||
|
||||
hook.prefetch_layer(slot=0, non_blocking=True)
|
||||
|
||||
assert torch.equal(next_block.weight, expected)
|
||||
|
||||
|
||||
def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_runtime):
|
||||
block = nn.Linear(2, 2, bias=False)
|
||||
expected = torch.arange(4, dtype=torch.float32).view(2, 2)
|
||||
block.weight.data.copy_(expected)
|
||||
group = PinnedResidentLayerGroup(
|
||||
[block],
|
||||
device=torch.device("cpu"),
|
||||
copy_stream=DummyStream(),
|
||||
pin_memory=False,
|
||||
rank_local_mmap=True,
|
||||
)
|
||||
group.registered_mmap = True
|
||||
group._cpu_staging_buffers.clear()
|
||||
|
||||
group.load()
|
||||
|
||||
assert torch.equal(block.weight, expected)
|
||||
|
||||
|
||||
def test_hwr_registration_uses_transport_budget(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
|
||||
plan, _, lease = _fake_hwr_plan()
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=True,
|
||||
dp_size=1,
|
||||
dlo_use_allgather=False,
|
||||
dlo_host_registration_limit_gib=1.5,
|
||||
),
|
||||
torch.device("cuda"),
|
||||
host_weight_plan=plan,
|
||||
)
|
||||
backend._host_weight_lease = lease # type: ignore[assignment]
|
||||
calls: list[tuple[object, int | None]] = []
|
||||
|
||||
class Registration:
|
||||
total_bytes = 4096
|
||||
region_count = 1
|
||||
|
||||
@staticmethod
|
||||
def close() -> tuple[str, ...]:
|
||||
return ()
|
||||
|
||||
def register(regions, *, device, max_bytes):
|
||||
assert device == torch.device("cuda")
|
||||
calls.append((regions, max_bytes))
|
||||
return Registration()
|
||||
|
||||
monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
|
||||
pinned_source = SimpleNamespace(numel=lambda: 1, is_pinned=lambda: True)
|
||||
|
||||
assert backend._try_register_hwr_mmap((pinned_source,)) # type: ignore[arg-type]
|
||||
assert calls == [(lease.mapped_regions, int(1.5 * 1024**3))]
|
||||
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
|
||||
|
||||
backend._release_registered_mmap()
|
||||
|
||||
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
|
||||
|
||||
|
||||
def test_hwr_registration_budget_validation_is_transport_scoped():
|
||||
assert (
|
||||
validate_dlo_host_registration_options(
|
||||
limit_gib=1.5,
|
||||
enable_dlo=True,
|
||||
use_allgather=False,
|
||||
hwr_mode="preferred",
|
||||
)
|
||||
== 1.5
|
||||
)
|
||||
invalid = (
|
||||
dict(limit_gib=-1, enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
|
||||
dict(limit_gib=float("inf"), enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
|
||||
dict(limit_gib=1, enable_dlo=False, use_allgather=False, hwr_mode="preferred"),
|
||||
dict(limit_gib=1, enable_dlo=True, use_allgather=True, hwr_mode="preferred"),
|
||||
dict(limit_gib=1, enable_dlo=True, use_allgather=False, hwr_mode="disabled"),
|
||||
)
|
||||
for options in invalid:
|
||||
with pytest.raises(ValueError):
|
||||
validate_dlo_host_registration_options(**options)
|
||||
|
||||
|
||||
def test_unregistration_precedes_lease_close_and_retries_failures(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
|
||||
events: list[str] = []
|
||||
responses = iter([("busy",), ()])
|
||||
|
||||
class Registration:
|
||||
@staticmethod
|
||||
def close() -> tuple[str, ...]:
|
||||
events.append("unregister")
|
||||
return next(responses)
|
||||
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=True,
|
||||
dp_size=1,
|
||||
dlo_use_allgather=False,
|
||||
),
|
||||
torch.device("cpu"),
|
||||
)
|
||||
lease = _FakeHostWeightLease("retry", events)
|
||||
backend._host_weight_lease = lease # type: ignore[assignment]
|
||||
backend._host_registration = Registration()
|
||||
backend._using_rank_local_mmap = True
|
||||
backend._using_registered_mmap = True
|
||||
backend.enabled = True
|
||||
monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
|
||||
|
||||
with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
|
||||
backend.disable()
|
||||
assert not lease.closed
|
||||
assert backend._host_registration is not None
|
||||
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
|
||||
|
||||
backend.disable()
|
||||
|
||||
assert lease.closed
|
||||
assert backend._host_registration is None
|
||||
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
|
||||
assert events == ["unregister", "unregister", "lease"]
|
||||
|
||||
|
||||
def test_failed_unregistration_retains_lease_after_backend_is_dropped(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
events: list[str] = []
|
||||
|
||||
class Registration:
|
||||
@staticmethod
|
||||
def close() -> tuple[str, ...]:
|
||||
events.append("unregister")
|
||||
return ("busy",)
|
||||
|
||||
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=True,
|
||||
dp_size=1,
|
||||
dlo_use_allgather=False,
|
||||
),
|
||||
torch.device("cpu"),
|
||||
)
|
||||
lease = _FakeHostWeightLease("retained", events)
|
||||
registration = Registration()
|
||||
backend._host_weight_lease = lease # type: ignore[assignment]
|
||||
backend._host_registration = registration
|
||||
backend._using_rank_local_mmap = True
|
||||
backend.enabled = True
|
||||
monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
|
||||
lease_ref = weakref.ref(lease)
|
||||
|
||||
with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
|
||||
backend.disable()
|
||||
del backend
|
||||
del lease
|
||||
gc.collect()
|
||||
|
||||
retained_lease = dist_backend_module._ACTIVE_HWR_REGISTRATIONS[0][1]
|
||||
assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(registration, retained_lease)]
|
||||
assert lease_ref() is retained_lease
|
||||
assert not retained_lease.closed
|
||||
|
||||
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
|
||||
retained_lease.close()
|
||||
|
||||
|
||||
class _MultiBlockModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
|
||||
|
||||
@@ -1131,51 +920,6 @@ class TestMmapWeightLoading:
|
||||
backend.disable()
|
||||
assert lease.closed
|
||||
|
||||
def test_hwr_registration_failure_falls_back_to_bounded_staging(
|
||||
self,
|
||||
monkeypatch,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=True,
|
||||
dp_size=1,
|
||||
dlo_use_allgather=False,
|
||||
),
|
||||
torch.device("cpu"),
|
||||
host_weight_plan=plan,
|
||||
)
|
||||
pipeline = _HWRPipeline()
|
||||
original_staging_allocator = backend._allocate_shared_cpu_staging_buffers
|
||||
|
||||
def fail_registration(*args, **kwargs):
|
||||
del args, kwargs
|
||||
raise HostRegistrationError("registration unavailable in CPU test")
|
||||
|
||||
def allocate_unpinned_staging(hooks, resident_group=None):
|
||||
# Registration selection still observes pin_cpu_memory=True. The
|
||||
# CPU-only test avoids asking the host for CUDA-pinned allocations.
|
||||
for hook in hooks:
|
||||
hook.pin_memory = False
|
||||
return original_staging_allocator(hooks, resident_group)
|
||||
|
||||
monkeypatch.setattr(dist_backend_module, "register_host_mappings", fail_registration)
|
||||
monkeypatch.setattr(backend, "_allocate_shared_cpu_staging_buffers", allocate_unpinned_staging)
|
||||
|
||||
backend.enable(pipeline)
|
||||
|
||||
assert carrier.taken
|
||||
assert backend._using_rank_local_mmap
|
||||
assert not backend._using_registered_mmap
|
||||
assert backend._host_registration is None
|
||||
assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
|
||||
assert not lease.closed
|
||||
|
||||
backend.disable()
|
||||
assert lease.closed
|
||||
|
||||
def test_hwr_backend_failure_drains_partial_setup_before_lease_close(
|
||||
self,
|
||||
monkeypatch,
|
||||
@@ -2183,9 +1927,8 @@ class TestDynamicSlotTracking:
|
||||
prev_idx = (i - 1) % num_blocks
|
||||
|
||||
# Dynamic slot tracking: read from prev's prefetched slot
|
||||
prefetched_slot = prefetched_slots[prev_idx]
|
||||
if prefetched_slot is not None:
|
||||
current_slots[i] = prefetched_slot
|
||||
if prefetched_slots[prev_idx] is not None:
|
||||
current_slots[i] = prefetched_slots[prev_idx]
|
||||
|
||||
read_slot = current_slots[i]
|
||||
|
||||
|
||||
@@ -1,90 +0,0 @@
|
||||
# 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-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Source-level regression tests for diffusion output/engine helpers.
|
||||
|
||||
These tests verify naming conventions and patterns by inspecting source code
|
||||
@@ -37,50 +37,10 @@ _FORMATTER_PATH = os.path.normpath(
|
||||
"output_formatter.py",
|
||||
)
|
||||
)
|
||||
_METRICS_PATH = os.path.normpath(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"metrics",
|
||||
"utils.py",
|
||||
)
|
||||
)
|
||||
_INLINE_CLIENT_PATH = os.path.normpath(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"diffusion",
|
||||
"inline_stage_diffusion_client.py",
|
||||
)
|
||||
)
|
||||
_STAGE_PROC_PATH = os.path.normpath(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"diffusion",
|
||||
"stage_diffusion_proc.py",
|
||||
)
|
||||
)
|
||||
_ORCHESTRATOR_PATH = os.path.normpath(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"engine",
|
||||
"orchestrator.py",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _read_source(path: str) -> str:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
with open(path) as f:
|
||||
return f.read()
|
||||
|
||||
|
||||
@@ -140,99 +100,9 @@ class TestMetricKeys:
|
||||
raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()")
|
||||
|
||||
assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source
|
||||
# step_total_ms is no longer emitted as a metric (no family consumes
|
||||
# it); it lives only in the debug log breakdown below.
|
||||
assert '"diffusion_engine_total_time_ms"' not in step_streaming_source
|
||||
assert '"diffusion_engine_total_time_ms": step_total_ms' in step_streaming_source
|
||||
assert '"postprocess_time_ms": postprocess_time * 1000' in step_streaming_source
|
||||
|
||||
def test_step_streaming_piggybacks_scheduler_waiting_snapshot(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "diffusion_scheduler_waiting_metrics(" in step_streaming_source
|
||||
assert "_scheduler_num_waiting_reqs" in step_streaming_source
|
||||
|
||||
def test_step_streaming_preserves_snapshot_on_postprocess_error(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "scheduler_metrics = diffusion_scheduler_waiting_metrics(" in step_streaming_source
|
||||
assert 'setattr(exc, "diffusion_metrics", scheduler_metrics)' in step_streaming_source
|
||||
assert step_streaming_source.index("scheduler_metrics =") < step_streaming_source.index(
|
||||
"self.postprocess_output("
|
||||
)
|
||||
|
||||
def test_abort_keeps_output_consumers_alive_for_terminal_snapshot(self) -> None:
|
||||
inline_source = _read_source(_INLINE_CLIENT_PATH)
|
||||
abort_source = _get_function_source(inline_source, "InlineStageDiffusionClient", "abort_requests_async")
|
||||
assert ".cancel(" not in abort_source
|
||||
|
||||
proc_source = _read_source(_STAGE_PROC_PATH)
|
||||
run_loop_source = _get_function_source(proc_source, "StageDiffusionProc", "run_loop")
|
||||
abort_branch = run_loop_source[run_loop_source.index('elif msg_type == "abort":') :]
|
||||
abort_branch = abort_branch[: abort_branch.index('elif msg_type == "collective_rpc":')]
|
||||
assert ".cancel(" not in abort_branch
|
||||
|
||||
def test_orchestrator_consumes_metrics_only_output_without_routing(self) -> None:
|
||||
source = _read_source(_ORCHESTRATOR_PATH)
|
||||
loop_source = _get_function_source(source, "Orchestrator", "_orchestration_loop")
|
||||
snapshot_pos = loop_source.index("_update_stage_replica_waiting(")
|
||||
sentinel_pos = loop_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID")
|
||||
route_pos = loop_source.index("pool.record_output_timestamps([diffusion_output])")
|
||||
assert snapshot_pos < sentinel_pos < route_pos
|
||||
|
||||
|
||||
class TestVaeDecodeEmit:
|
||||
"""VAE decode timing is sourced from ``DiffusionOutput.stage_durations``
|
||||
(populated by the diffusion pipeline profiler) and emitted as
|
||||
``vae_decode_time_ms`` so the ``_MS_TO_S`` accumulator picks it up.
|
||||
"""
|
||||
|
||||
def test_extract_vae_decode_ms_helper_exists(self) -> None:
|
||||
source = _read_source(_METRICS_PATH)
|
||||
helper_src = _get_function_source(source, None, "extract_diffusion_vae_decode_ms")
|
||||
assert ".vae.decode" in helper_src, "helper must key on the '.vae.decode' suffix"
|
||||
|
||||
def test_step_streaming_emits_vae_decode_time_ms(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "extract_diffusion_vae_decode_ms(" in step_streaming_source, (
|
||||
"step_streaming must call extract_diffusion_vae_decode_ms to source VAE decode timing"
|
||||
)
|
||||
assert '"vae_decode_time_ms"' in step_streaming_source, (
|
||||
"step_streaming must emit vae_decode_time_ms key for the _MS_TO_S accumulator"
|
||||
)
|
||||
|
||||
|
||||
class TestDiffuseForwardEmit:
|
||||
"""Forward-only timing → forward_time_ms, mirroring TestVaeDecodeEmit."""
|
||||
|
||||
def test_extract_diffuse_ms_helper_exists(self) -> None:
|
||||
source = _read_source(_METRICS_PATH)
|
||||
helper_src = _get_function_source(source, None, "extract_diffusion_denoise_ms")
|
||||
assert ".diffuse" in helper_src, "helper must key on the '.diffuse' suffix"
|
||||
|
||||
def test_step_streaming_emits_forward_time_ms(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "extract_diffusion_denoise_ms(" in step_streaming_source, (
|
||||
"step_streaming must call extract_diffusion_denoise_ms to source forward-only timing"
|
||||
)
|
||||
assert '"forward_time_ms"' in step_streaming_source, (
|
||||
"step_streaming must emit forward_time_ms key for the _MS_TO_S accumulator"
|
||||
)
|
||||
|
||||
|
||||
class TestKvRecvEmit:
|
||||
"""KV-recv timing → kv_recv_time_ms, sourced from DiffusionOutput.kv_recv_ms
|
||||
(set by the runner's _prepare_request_for_forward, not the profiler)."""
|
||||
|
||||
def test_step_streaming_emits_kv_recv_time_ms(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "kv_recv_ms" in step_streaming_source, "step_streaming must read output.kv_recv_ms"
|
||||
assert '"kv_recv_time_ms"' in step_streaming_source, (
|
||||
"step_streaming must emit kv_recv_time_ms key for the _MS_TO_S accumulator"
|
||||
)
|
||||
|
||||
|
||||
class TestDummyRunAllocation:
|
||||
"""Verify _dummy_run generates exact-sized audio arrays."""
|
||||
|
||||
@@ -181,7 +181,7 @@ def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
|
||||
|
||||
|
||||
def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None:
|
||||
# LTX2 returns (video, audio) tuples as DiffusionOutput.output.
|
||||
# LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output.
|
||||
video = torch.arange(300_000, dtype=torch.float32)
|
||||
audio = torch.arange(300_000, dtype=torch.float32) * 3
|
||||
output = DiffusionOutput(output=(video, audio))
|
||||
|
||||
@@ -773,19 +773,6 @@ class TestRequestScheduler:
|
||||
assert third.num_running_reqs == 1
|
||||
assert third.num_waiting_reqs == 0
|
||||
|
||||
def test_records_initial_scheduler_queue_wait(self, mocker: MockerFixture) -> None:
|
||||
perf_counter = mocker.patch(
|
||||
"vllm_omni.diffusion.sched.base_scheduler.time.perf_counter",
|
||||
side_effect=[10.0, 10.125],
|
||||
)
|
||||
request = _make_step_request("queue-wait", num_inference_steps=2)
|
||||
|
||||
self.scheduler.add_request(request)
|
||||
self.scheduler.schedule()
|
||||
|
||||
assert request.scheduler_queue_wait_ms == pytest.approx(125.0)
|
||||
assert perf_counter.call_count == 2
|
||||
|
||||
def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None:
|
||||
scheduler = RequestScheduler()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import multiprocessing as mp
|
||||
@@ -1079,42 +1079,6 @@ class TestStageDiffusionClientErrorPropagation:
|
||||
|
||||
assert client.get_diffusion_output_nowait() is None
|
||||
|
||||
def test_error_response_preserves_scheduler_metrics(self):
|
||||
from vllm_omni.metrics import definitions as metric_defs
|
||||
|
||||
client = self._make_client()
|
||||
client._response_socket.recv.side_effect = [b"message", zmq.Again()]
|
||||
client._decoder.decode.return_value = {
|
||||
"type": "error",
|
||||
"request_id": "req-error",
|
||||
"error": "gpu fault",
|
||||
"metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
|
||||
}
|
||||
|
||||
output = client.get_diffusion_output_nowait()
|
||||
|
||||
assert output is not None
|
||||
assert output.error == "gpu fault"
|
||||
assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
|
||||
|
||||
def test_metrics_response_creates_metrics_only_output(self):
|
||||
from vllm_omni.metrics import definitions as metric_defs
|
||||
from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID
|
||||
|
||||
client = self._make_client()
|
||||
client._response_socket.recv.side_effect = [b"message", zmq.Again()]
|
||||
client._decoder.decode.return_value = {
|
||||
"type": "metrics",
|
||||
"metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
|
||||
}
|
||||
|
||||
output = client.get_diffusion_output_nowait()
|
||||
|
||||
assert output is not None
|
||||
assert output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID
|
||||
assert output.error is None
|
||||
assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
|
||||
|
||||
def test_check_health_raises_when_proc_dead(self):
|
||||
"""``check_health`` detects a dead subprocess via the manager's proc
|
||||
and raises ``EngineDeadError``, setting ``_engine_dead`` as a
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Unit tests for WorkerWrapperBase class.
|
||||
@@ -12,7 +12,6 @@ This module tests the WorkerWrapperBase implementation:
|
||||
- Dynamic worker class extension
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
@@ -205,27 +204,6 @@ class TestWorkerWrapperBaseDelegation:
|
||||
wrapper.worker.shutdown.assert_called_once()
|
||||
assert result is None
|
||||
|
||||
def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture):
|
||||
events: list[str] = []
|
||||
offload_backend = mocker.Mock()
|
||||
offload_backend.disable.side_effect = lambda: events.append("offload")
|
||||
kv_manager = mocker.Mock()
|
||||
kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv")
|
||||
destroy = mocker.patch(
|
||||
"vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
|
||||
side_effect=lambda: events.append("distributed"),
|
||||
)
|
||||
worker = DiffusionWorker.__new__(DiffusionWorker)
|
||||
worker.model_runner = SimpleNamespace(
|
||||
offload_backend=offload_backend,
|
||||
kv_transfer_manager=kv_manager,
|
||||
)
|
||||
|
||||
worker.shutdown()
|
||||
|
||||
assert events == ["offload", "kv", "distributed"]
|
||||
destroy.assert_called_once_with()
|
||||
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Tests: execute_method
|
||||
|
||||
@@ -1,7 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import wave
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
@@ -24,7 +20,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
def test_realtime_client_builds_explicit_native_duplex_url():
|
||||
url = build_realtime_url(
|
||||
"ws://localhost:8099/v1/realtime?custom=1&duplex=0&model=stale&minicpmo45_native_duplex=0&session_id=stale",
|
||||
"ws://localhost:8099/v1/realtime?custom=1",
|
||||
"openbmb/MiniCPM-o-4_5",
|
||||
session_id="session-a",
|
||||
)
|
||||
@@ -54,7 +50,7 @@ def test_seed_tts_initial_text_is_part_of_native_duplex_context():
|
||||
|
||||
def test_realtime_client_builds_resume_only_url_when_autostart_disabled():
|
||||
url = build_realtime_url(
|
||||
"ws://localhost:8099/v1/realtime?duplex=1&autostart=1",
|
||||
"ws://localhost:8099/v1/realtime?duplex=1",
|
||||
"openbmb/MiniCPM-o-4_5",
|
||||
autostart=False,
|
||||
)
|
||||
@@ -77,15 +73,10 @@ async def test_realtime_client_configure_omits_ref_audio_by_default():
|
||||
|
||||
client = Client()
|
||||
|
||||
await client.configure(
|
||||
"openbmb/MiniCPM-o-4_5",
|
||||
idle_timeout_s=900,
|
||||
timeout_s=1,
|
||||
)
|
||||
await client.configure("openbmb/MiniCPM-o-4_5", timeout_s=1)
|
||||
|
||||
session = client.sent[0]["session"]
|
||||
assert "ref_audio" not in session
|
||||
assert session["idle_timeout_s"] == 900
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -109,74 +100,6 @@ async def test_realtime_client_configure_sends_explicit_ref_audio():
|
||||
|
||||
session = client.sent[0]["session"]
|
||||
assert session["ref_audio"] == "data:audio/wav;base64,AAAA"
|
||||
assert "idle_timeout_s" not in session
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_client_reports_reader_exit_to_event_waiters():
|
||||
client = RealtimeDuplexClient("ws://unused")
|
||||
|
||||
async def stopped():
|
||||
return None
|
||||
|
||||
client._reader_task = asyncio.create_task(stopped())
|
||||
await client._reader_task
|
||||
|
||||
with pytest.raises(ConnectionError, match="closed before"):
|
||||
client.raise_if_reader_stopped()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_client_configure_reports_reader_failure_immediately():
|
||||
class Client(RealtimeDuplexClient):
|
||||
async def send(self, event):
|
||||
async def fail():
|
||||
raise RuntimeError("reader failed")
|
||||
|
||||
self._reader_task = asyncio.create_task(fail())
|
||||
await asyncio.sleep(0)
|
||||
|
||||
client = Client("ws://unused")
|
||||
with pytest.raises(ConnectionError, match="reader failed"):
|
||||
await client.configure("model", timeout_s=10.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_client_close_reports_reader_cancellation_immediately():
|
||||
class Client(RealtimeDuplexClient):
|
||||
async def send(self, event):
|
||||
async def wait_forever():
|
||||
await asyncio.Future()
|
||||
|
||||
self._reader_task = asyncio.create_task(wait_forever())
|
||||
self._reader_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await self._reader_task
|
||||
|
||||
client = Client("ws://unused")
|
||||
with pytest.raises(ConnectionError, match="reader was cancelled"):
|
||||
await client.close_session(timeout_s=10.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_client_close_waits_for_a_new_session_closed_event():
|
||||
class Client(RealtimeDuplexClient):
|
||||
async def send(self, event):
|
||||
assert event["type"] == "session.close"
|
||||
|
||||
async def acknowledge_close():
|
||||
await asyncio.sleep(0.02)
|
||||
self.events.add({"type": "session.closed"})
|
||||
|
||||
asyncio.create_task(acknowledge_close())
|
||||
|
||||
client = Client("ws://unused")
|
||||
client.events.add({"type": "session.closed", "reason": "old"})
|
||||
started_at = asyncio.get_running_loop().time()
|
||||
|
||||
await client.close_session(timeout_s=0.2)
|
||||
|
||||
assert asyncio.get_running_loop().time() - started_at >= 0.01
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -312,7 +235,6 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals():
|
||||
"output_token_count": 4,
|
||||
"ttft_ms": 120.0,
|
||||
"tpot_ms": 15.0,
|
||||
"itls_ms": [10.0, 14.0, 18.0],
|
||||
"inter_token_interval_ms": {
|
||||
"count": 3,
|
||||
"mean": 14.0,
|
||||
|
||||
@@ -0,0 +1,290 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""E2E test for DreamID-Omni offline inference with Cache-Dit.
|
||||
|
||||
Uses the official DreamID-Omni ``test_case/twoip`` example assets instead of
|
||||
synthetic inputs so the coverage matches the upstream demo structure.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import soundfile as sf
|
||||
from PIL import Image
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow]
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
DEFAULT_MODEL_PATH = REPO_ROOT / "dreamid_omni"
|
||||
MODEL_PATH_ENV = "DREAMID_OMNI_MODEL_PATH"
|
||||
EXAMPLE_ROOT_ENV = "DREAMID_OMNI_EXAMPLE_ROOT"
|
||||
NUM_INFERENCE_STEPS = 45
|
||||
|
||||
CACHE_CONFIG = {
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"max_cached_steps": 20,
|
||||
"residual_diff_threshold": 0.24,
|
||||
"max_continuous_cached_steps": 3,
|
||||
"enable_taylorseer": False,
|
||||
"taylorseer_order": 1,
|
||||
"scm_steps_mask_policy": None,
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
|
||||
|
||||
class DreamIDGenerationResult(NamedTuple):
|
||||
video: np.ndarray
|
||||
audio: np.ndarray
|
||||
fps: int
|
||||
duration_s: float
|
||||
elapsed_s: float
|
||||
|
||||
|
||||
class DreamIDExampleCase(NamedTuple):
|
||||
prompt: str
|
||||
image_paths: tuple[Path, ...]
|
||||
audio_paths: tuple[Path, ...]
|
||||
|
||||
|
||||
def _require_model_path() -> str:
|
||||
model_path = os.environ.get(MODEL_PATH_ENV, str(DEFAULT_MODEL_PATH))
|
||||
if not Path(model_path).exists():
|
||||
pytest.skip(f"DreamID-Omni model bundle not found at {model_path}. Set {MODEL_PATH_ENV}.")
|
||||
|
||||
if importlib.util.find_spec("dreamid_omni") is None:
|
||||
pytest.skip(
|
||||
"dreamid_omni dependency package is not importable. "
|
||||
"Run examples/offline_inference/x_to_video_audio/download_dreamid_omni.py first."
|
||||
)
|
||||
|
||||
return model_path
|
||||
|
||||
|
||||
def _dreamid_repo_root_candidates() -> list[Path]:
|
||||
candidates: list[Path] = []
|
||||
example_root = os.environ.get(EXAMPLE_ROOT_ENV)
|
||||
if example_root:
|
||||
candidates.append(Path(example_root).expanduser().resolve())
|
||||
|
||||
spec = importlib.util.find_spec("dreamid_omni")
|
||||
if spec is not None and spec.origin:
|
||||
candidates.append(Path(spec.origin).resolve().parents[1] / "test_case" / "twoip")
|
||||
|
||||
return candidates
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_example_case() -> DreamIDExampleCase:
|
||||
for candidate_root in _dreamid_repo_root_candidates():
|
||||
caption_path = candidate_root / "captions" / "20.json"
|
||||
image_paths = (
|
||||
candidate_root / "imgs" / "20" / "0.png",
|
||||
candidate_root / "imgs" / "20" / "1.png",
|
||||
)
|
||||
audio_paths = (
|
||||
candidate_root / "audios" / "20" / "0.wav",
|
||||
candidate_root / "audios" / "20" / "1.wav",
|
||||
)
|
||||
if caption_path.exists() and all(path.exists() for path in image_paths + audio_paths):
|
||||
prompt = _load_prompt_from_json(caption_path)
|
||||
return DreamIDExampleCase(
|
||||
prompt=prompt,
|
||||
image_paths=image_paths,
|
||||
audio_paths=audio_paths,
|
||||
)
|
||||
|
||||
searched = ", ".join(str(path) for path in _dreamid_repo_root_candidates()) or "<none>"
|
||||
pytest.skip(
|
||||
"DreamID-Omni official example assets not found. "
|
||||
f"Set {EXAMPLE_ROOT_ENV} to the local 'test_case/twoip' directory. "
|
||||
f"Searched: {searched}"
|
||||
)
|
||||
|
||||
|
||||
def _load_prompt_from_json(path: Path) -> str:
|
||||
raw_text = path.read_text(encoding="utf-8").strip()
|
||||
parsed = json.loads(raw_text)
|
||||
if isinstance(parsed, str):
|
||||
prompt = parsed
|
||||
elif isinstance(parsed, dict) and isinstance(parsed.get("content"), str):
|
||||
prompt = parsed["content"]
|
||||
else:
|
||||
raise AssertionError(f"Unexpected DreamID caption format in {path}")
|
||||
|
||||
prompt = re.sub(
|
||||
r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]",
|
||||
"",
|
||||
prompt,
|
||||
flags=re.DOTALL,
|
||||
).strip()
|
||||
prompt = re.sub(
|
||||
r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]",
|
||||
"",
|
||||
prompt,
|
||||
flags=re.DOTALL,
|
||||
).strip()
|
||||
prompt = re.sub(r"\[[A-Z_]+\]", "", prompt)
|
||||
prompt = re.sub(r"\n\s*\n", "\n", prompt).strip()
|
||||
return prompt
|
||||
|
||||
|
||||
def _load_example_images(example_case: DreamIDExampleCase) -> list[Image.Image]:
|
||||
images: list[Image.Image] = []
|
||||
for path in example_case.image_paths:
|
||||
with Image.open(path) as image:
|
||||
images.append(image.convert("RGB"))
|
||||
return images
|
||||
|
||||
|
||||
def _load_example_audios(example_case: DreamIDExampleCase) -> list[np.ndarray]:
|
||||
audios: list[np.ndarray] = []
|
||||
for path in example_case.audio_paths:
|
||||
audio_array, sample_rate = sf.read(str(path), dtype="float32")
|
||||
audio_np = np.asarray(audio_array, dtype=np.float32)
|
||||
if audio_np.ndim == 2:
|
||||
audio_np = np.mean(audio_np, axis=1)
|
||||
audio_np = audio_np[int(sample_rate * 1) : int(sample_rate * 3)]
|
||||
audios.append(audio_np.reshape(-1))
|
||||
return audios
|
||||
|
||||
|
||||
def _sampling_params(num_inference_steps: int) -> OmniDiffusionSamplingParams:
|
||||
return OmniDiffusionSamplingParams(
|
||||
height=704,
|
||||
width=1280,
|
||||
num_inference_steps=num_inference_steps,
|
||||
seed=42,
|
||||
extra_args={
|
||||
"solver_name": "unipc",
|
||||
"shift": 5.0,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _request_payload(example_case: DreamIDExampleCase) -> dict[str, object]:
|
||||
return {
|
||||
"prompt": example_case.prompt,
|
||||
"video_negative_prompt": "jitter, bad hands, blur, distortion",
|
||||
"audio_negative_prompt": "robotic, muffled, echo, distorted",
|
||||
"multi_modal_data": {
|
||||
"image": _load_example_images(example_case),
|
||||
"audio": _load_example_audios(example_case),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _extract_generation(outputs: list[object], elapsed_s: float) -> DreamIDGenerationResult:
|
||||
"""Extract video and audio from raw Omni outputs"""
|
||||
assert outputs, "DreamID-Omni returned no outputs"
|
||||
|
||||
result = outputs[0]
|
||||
assert hasattr(result, "images") and result.images, "not return video frames"
|
||||
|
||||
video = result.images[0]
|
||||
assert isinstance(video, np.ndarray), f"Expected video ndarray, got {type(video)}"
|
||||
assert video.ndim == 4, f"Expected 4D DreamID video tensor, got {video.shape}"
|
||||
|
||||
multimodal_output = getattr(result, "multimodal_output", None) or {}
|
||||
assert "audio" in multimodal_output, "not return audio"
|
||||
audio = np.asarray(multimodal_output["audio"], dtype=np.float32).reshape(-1)
|
||||
assert audio.size > 0, "returned empty audio"
|
||||
|
||||
audio_sample_rate = int(multimodal_output.get("audio_sample_rate", 16000))
|
||||
fps = multimodal_output.get("fps", 24)
|
||||
duration_s = float(audio.size) / float(audio_sample_rate)
|
||||
|
||||
generation = DreamIDGenerationResult(
|
||||
video=video.astype(np.float32, copy=False),
|
||||
audio=audio,
|
||||
fps=fps,
|
||||
duration_s=duration_s,
|
||||
elapsed_s=elapsed_s,
|
||||
)
|
||||
return generation
|
||||
|
||||
|
||||
def _run_generation(
|
||||
*,
|
||||
model_path: str,
|
||||
example_case: DreamIDExampleCase,
|
||||
use_cache_dit: bool,
|
||||
) -> DreamIDGenerationResult:
|
||||
runner_kwargs = {
|
||||
"model_type": "dreamid-omni",
|
||||
"parallel_config": DiffusionParallelConfig(cfg_parallel_size=1),
|
||||
}
|
||||
if use_cache_dit:
|
||||
runner_kwargs["cache_backend"] = "cache_dit"
|
||||
runner_kwargs["cache_config"] = CACHE_CONFIG
|
||||
|
||||
with OmniRunner(model_path, **runner_kwargs) as runner:
|
||||
warmup_payload = _request_payload(example_case)
|
||||
_ = runner.omni.generate(warmup_payload, _sampling_params(2))
|
||||
|
||||
elapsed_samples: list[float] = []
|
||||
extracted_result: DreamIDGenerationResult | None = None
|
||||
|
||||
request_payload = _request_payload(example_case)
|
||||
start = time.perf_counter()
|
||||
outputs = runner.omni.generate(request_payload, _sampling_params(NUM_INFERENCE_STEPS))
|
||||
elapsed_s = time.perf_counter() - start
|
||||
elapsed_samples.append(elapsed_s)
|
||||
if extracted_result is None:
|
||||
extracted_result = _extract_generation(outputs, elapsed_s)
|
||||
|
||||
assert extracted_result is not None
|
||||
median_elapsed_s = float(np.median(np.asarray(elapsed_samples, dtype=np.float64)))
|
||||
return extracted_result._replace(elapsed_s=median_elapsed_s)
|
||||
|
||||
|
||||
def _video_mae(a: np.ndarray, b: np.ndarray) -> float:
|
||||
assert a.shape == b.shape, f"Video shapes differ: {a.shape} vs {b.shape}"
|
||||
return float(np.mean(np.abs(a - b)))
|
||||
|
||||
|
||||
def _audio_rms(audio: np.ndarray) -> float:
|
||||
audio = np.asarray(audio, dtype=np.float32).reshape(-1)
|
||||
return float(np.sqrt(np.mean(audio**2)))
|
||||
|
||||
|
||||
def _audio_rms_ratio(a: DreamIDGenerationResult, b: DreamIDGenerationResult) -> float:
|
||||
rms_a = _audio_rms(a.audio)
|
||||
rms_b = _audio_rms(b.audio)
|
||||
return rms_b / max(rms_a, 1e-6)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_dreamid_omni_cache_dit_e2e() -> None:
|
||||
model_path = _require_model_path()
|
||||
example_case = _load_example_case()
|
||||
|
||||
baseline = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=False)
|
||||
cached = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=True)
|
||||
|
||||
video_mae = _video_mae(baseline.video, cached.video)
|
||||
audio_rms_ratio = _audio_rms_ratio(baseline, cached)
|
||||
speedup = baseline.elapsed_s / cached.elapsed_s if cached.elapsed_s > 0 else float("inf")
|
||||
|
||||
assert video_mae <= 0.15, f"video drift too much: {video_mae:.6f}"
|
||||
assert 0.4 <= audio_rms_ratio <= 2.5, f"audio quality regressed too much: rms_ratio={audio_rms_ratio:.6f}"
|
||||
assert speedup >= 1.05, f"cache-dit improve not enough: {speedup:.3f}x"
|
||||
@@ -0,0 +1,162 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end tests for MagiHuman pipeline via vLLM-Omni."""
|
||||
|
||||
import io
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
MODEL = "SII-GAIR/daVinci-MagiHuman-Base-1080p"
|
||||
|
||||
# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
|
||||
_OMNI_RUNNER_PARAM = (
|
||||
MODEL,
|
||||
None,
|
||||
{
|
||||
"init_timeout": 1200,
|
||||
"tensor_parallel_size": 2,
|
||||
"vae_use_tiling": True,
|
||||
},
|
||||
)
|
||||
|
||||
_OMNI_RUNNER_PARAM_OFFLOAD = (
|
||||
MODEL,
|
||||
None,
|
||||
{
|
||||
"init_timeout": 1200,
|
||||
"tensor_parallel_size": 2,
|
||||
"enable_cpu_offload": True,
|
||||
"vae_use_tiling": True,
|
||||
},
|
||||
)
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.parametrize(
|
||||
"omni_runner",
|
||||
[_OMNI_RUNNER_PARAM, _OMNI_RUNNER_PARAM_OFFLOAD],
|
||||
indirect=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
|
||||
"""Validate that the MP4 contains meaningful video and audio tracks."""
|
||||
container = av.open(io.BytesIO(video_bytes))
|
||||
|
||||
v_streams = [s for s in container.streams if s.type == "video"]
|
||||
assert len(v_streams) >= 1, "No video stream found in MP4"
|
||||
|
||||
a_streams = [s for s in container.streams if s.type == "audio"]
|
||||
assert len(a_streams) >= 1, "No audio stream found in MP4"
|
||||
|
||||
v_stream = v_streams[0]
|
||||
assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}"
|
||||
assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}"
|
||||
|
||||
frame_count = 0
|
||||
for frame in container.decode(video=0):
|
||||
frame_count += 1
|
||||
if frame_count >= min_frames:
|
||||
break
|
||||
assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})"
|
||||
|
||||
container.close()
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
def test_magi_human_e2e(omni_runner: OmniRunner):
|
||||
"""End-to-end test for MagiHuman generating video and audio."""
|
||||
prompt = (
|
||||
"A young woman with long, wavy golden blonde hair and bright blue eyes, "
|
||||
"wearing a fitted ivory silk blouse with a delicate lace collar, sits "
|
||||
"stationary in front of a softly lit, blurred warm-toned interior. Her "
|
||||
"overall disposition is warm, composed, and gently confident. The camera "
|
||||
"holds a static medium close-up, framing her from the shoulders up, "
|
||||
"with shallow depth of field keeping her face in sharp focus. Soft "
|
||||
"directional key light falls from the upper left, casting a gentle "
|
||||
"highlight along her cheekbone and nose bridge. She draws a quiet breath, "
|
||||
"the levator labii superiors relaxing as her lips part. She speaks in "
|
||||
"clear, warm, unhurried American English: "
|
||||
"\"The most beautiful things in life aren't things at all — "
|
||||
"they're moments, feelings, and the people who make you feel truly alive.\" "
|
||||
"Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
|
||||
"shapes each vowel with precision. A faint, genuine smile engages the "
|
||||
"zygomaticus major, lifting her lip corners fractionally. Her brows rest "
|
||||
"in a soft, neutral arch throughout. She maintains steady, forward-facing "
|
||||
"eye contact. Head position remains level; no torso displacement occurs.\n\n"
|
||||
"Dialogue:\n"
|
||||
"<Young blonde woman, American English>: "
|
||||
"\"The most beautiful things in life aren't things at all — "
|
||||
"they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
|
||||
"Background Sound:\n"
|
||||
"<Soft, warm indoor ambience with a faint distant piano melody>"
|
||||
)
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=256,
|
||||
width=448,
|
||||
num_inference_steps=8,
|
||||
seed=52,
|
||||
extra_args={
|
||||
"seconds": 5,
|
||||
"sr_height": 1080,
|
||||
"sr_width": 1920,
|
||||
"sr_num_inference_steps": 5,
|
||||
},
|
||||
)
|
||||
|
||||
omni = omni_runner.omni
|
||||
outputs = list(
|
||||
omni.generate(
|
||||
prompts=[prompt],
|
||||
sampling_params_list=[sampling_params],
|
||||
)
|
||||
)
|
||||
|
||||
assert len(outputs) > 0, "No outputs returned"
|
||||
first = outputs[0]
|
||||
|
||||
assert hasattr(first, "images") and first.images, "No video frames in output"
|
||||
video_frames = first.images[0]
|
||||
assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
|
||||
assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
|
||||
|
||||
mm = first.multimodal_output
|
||||
assert mm, "multimodal_output is empty or missing"
|
||||
|
||||
audio_waveform = mm.get("audio")
|
||||
assert audio_waveform is not None, "No audio waveform in multimodal_output"
|
||||
|
||||
audio_sample_rate = mm.get("audio_sample_rate")
|
||||
assert audio_sample_rate is not None, (
|
||||
"audio_sample_rate not found in multimodal_output; model post-process must propagate it"
|
||||
)
|
||||
assert isinstance(audio_sample_rate, (int, float)), (
|
||||
f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
|
||||
)
|
||||
assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"
|
||||
|
||||
fps = mm.get("fps")
|
||||
assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
|
||||
assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
|
||||
assert int(fps) > 0, f"fps must be positive, got {fps}"
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
video_frames,
|
||||
audio_waveform,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=int(audio_sample_rate),
|
||||
)
|
||||
assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
|
||||
assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
|
||||
|
||||
_validate_mp4(video_bytes)
|
||||
@@ -1,51 +0,0 @@
|
||||
# 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
|
||||
@@ -15,10 +15,20 @@ import pytest
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from tests.helpers.runtime import OmniServerParams, get_model_prefix
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
|
||||
MODEL = "openbmb/MiniCPM-o-4_5"
|
||||
DEPLOY_CONFIG = get_deploy_config_path("minicpmo_4_5.yaml")
|
||||
DEPLOY_CONFIG = modify_stage_config(
|
||||
get_deploy_config_path("minicpmo_4_5_duplex.yaml"),
|
||||
updates={
|
||||
"base_config": get_deploy_config_path("minicpmo_4_5.yaml"),
|
||||
# Talker context is 4096 (tts_config.max_position_embeddings); KV sizing
|
||||
# is left automatic so duplex matches the simplex deploy profiles.
|
||||
"stages": {
|
||||
1: {"max_model_len": 4096},
|
||||
},
|
||||
},
|
||||
)
|
||||
ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5"
|
||||
RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav"
|
||||
RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941"
|
||||
|
||||
@@ -58,14 +58,18 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||
read_pcm16_wav,
|
||||
summarize_session_request_metrics,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
|
||||
reference_audio_data_url as _ref_audio_data_url,
|
||||
)
|
||||
|
||||
_url_with_model = build_realtime_url
|
||||
_read_wav_pcm16 = read_pcm16_wav
|
||||
|
||||
|
||||
def _ref_audio_data_url(path: str | None) -> str | None:
|
||||
if path is None:
|
||||
return None
|
||||
ref_path = Path(path).expanduser()
|
||||
return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii")
|
||||
|
||||
|
||||
@dataclass
|
||||
class DemoState:
|
||||
events: list[dict[str, object]] = field(default_factory=list)
|
||||
|
||||
@@ -7,6 +7,7 @@ E2E Online tests for Qwen3-Omni model.
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from openai import BadRequestError
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
|
||||
@@ -88,7 +89,7 @@ test_params = [
|
||||
),
|
||||
]
|
||||
|
||||
reduced_token_params = [
|
||||
test_token_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
@@ -96,7 +97,7 @@ reduced_token_params = [
|
||||
use_stage_cli=True,
|
||||
server_args=["--async-chunk"],
|
||||
),
|
||||
id="batch_token_2048",
|
||||
id="batch_token_64",
|
||||
)
|
||||
]
|
||||
|
||||
@@ -137,7 +138,29 @@ def get_max_batch_size(size_type="few"):
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_to_audio_001(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text
|
||||
Output Modal: audio
|
||||
Input Setting: stream=True
|
||||
Datasets: single request
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"modalities": ["audio"],
|
||||
"stream": True,
|
||||
"key_words": {"text": ["beijing"]},
|
||||
}
|
||||
|
||||
online_client.send_omni_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
|
||||
def test_text_to_text_audio_001(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text
|
||||
@@ -181,13 +204,36 @@ def test_text_video_to_text_001(omni_server, online_client) -> None:
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
|
||||
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
|
||||
def test_text_audio_to_text_audio_001(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text, audio
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(
|
||||
audio_data_url=audio_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_audio")
|
||||
)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"key_words": {"audio": AUDIO_KEY},
|
||||
}
|
||||
|
||||
online_client.send_omni_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
|
||||
def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text, long-duration audio (~LONG_AUDIO_DURATION_SEC s WAV)
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=False
|
||||
Datasets: few requests
|
||||
Datasets: single request
|
||||
"""
|
||||
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(LONG_AUDIO_DURATION_SEC, 1)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(
|
||||
@@ -206,13 +252,37 @@ def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
|
||||
def test_text_image_to_text_audio_001(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text, image
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(
|
||||
image_data_url=image_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_image")
|
||||
)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"key_words": {"image": IMAGE_KEY},
|
||||
}
|
||||
|
||||
online_client.send_omni_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
|
||||
def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text, high-resolution image (1080p-class JPEG)
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=False
|
||||
Datasets: few requests
|
||||
Datasets: single request
|
||||
"""
|
||||
image_data_url = (
|
||||
f"data:image/jpeg;base64,{generate_synthetic_image(LARGE_IMAGE_WIDTH, LARGE_IMAGE_HEIGHT)['base64']}"
|
||||
@@ -234,7 +304,7 @@ def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
|
||||
def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text, video
|
||||
@@ -258,6 +328,36 @@ def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
|
||||
online_client.send_omni_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
|
||||
def test_mix_to_text_audio_001(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text, audio, image, video
|
||||
Output Modal: text, audio
|
||||
Input Setting: stream=True
|
||||
Datasets: few requests
|
||||
"""
|
||||
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
|
||||
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
video_data_url=video_data_url,
|
||||
image_data_url=image_data_url,
|
||||
audio_data_url=audio_data_url,
|
||||
content_text=get_prompt("mix"),
|
||||
)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
|
||||
}
|
||||
online_client.send_omni_request(request_config, request_num=get_max_batch_size())
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_audio_in_video_001(omni_server, online_client) -> None:
|
||||
@@ -407,22 +507,22 @@ def test_one_word_prompt_001(omni_server, online_client) -> None:
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_speaker_001(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: Chinese text.
|
||||
Output Modal: Chinese text and audio using an uppercase preset speaker name.
|
||||
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"),
|
||||
content_text=get_prompt("text"),
|
||||
)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"speaker": "CHELSIE",
|
||||
"key_words": {"text": ["北京"]},
|
||||
"speaker": "Chelsie",
|
||||
"key_words": {"text": ["beijing"]},
|
||||
}
|
||||
|
||||
online_client.send_omni_request(request_config)
|
||||
@@ -463,6 +563,55 @@ def test_speaker_002(omni_server, online_client) -> None:
|
||||
print(f"Gender assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_speaker_003(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text only (one-word answer constraint).
|
||||
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
|
||||
Input Setting: stream=True
|
||||
Datasets: single request
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
content_text=get_prompt("text"),
|
||||
)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"speaker": "CHELSIE",
|
||||
"key_words": {"text": ["beijing"]},
|
||||
}
|
||||
|
||||
online_client.send_omni_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_language_001(omni_server, online_client) -> None:
|
||||
"""
|
||||
Input Modal: text only (one-word answer constraint).
|
||||
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
|
||||
Input Setting: stream=True
|
||||
Datasets: single request
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
content_text=get_prompt("text_chinese"),
|
||||
)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"key_words": {"text": ["北京"]},
|
||||
}
|
||||
|
||||
online_client.send_omni_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
|
||||
@@ -470,7 +619,7 @@ def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
|
||||
Input Modal: text only (long-form generation prompt).
|
||||
Output Modal: text, audio (default ``modalities``);
|
||||
Input Setting: stream=True
|
||||
Datasets: few requests
|
||||
Datasets: single request
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
@@ -482,3 +631,21 @@ def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
|
||||
text = responses[0].text_content if responses else ""
|
||||
word_count = len(text.split())
|
||||
assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True)
|
||||
def test_invalid_audio_format_rejected(omni_server, openai_client) -> None:
|
||||
"""Ensure invalid audio format is rejected with 400 before streaming starts."""
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
content_text=get_prompt(),
|
||||
)
|
||||
with pytest.raises(BadRequestError, match="audio format"):
|
||||
openai_client.client.chat.completions.create(
|
||||
model=omni_server.model,
|
||||
messages=messages,
|
||||
modalities=["text", "audio"],
|
||||
audio={"voice": "alloy", "format": "aac"},
|
||||
stream=True,
|
||||
)
|
||||
@@ -8,8 +8,9 @@ and are supported by the following text-to-image models:
|
||||
- Qwen-Image-2512
|
||||
|
||||
One feature per test case, matching the Test Plan in PR #1682 (Qwen-Image-Edit).
|
||||
Nightly covers each feature once, alternating models (5× Qwen-Image, 4× Qwen-Image-2512).
|
||||
Ulysses stays on Qwen-Image. See docs/user_guide/diffusion_acceleration.md.
|
||||
Supported features for Qwen-Image series: TeaCache, Cache-DiT, Ulysses-SP, Ring-Attention,
|
||||
CFG-Parallel, Tensor-Parallel, VAE-Patch-Parallel, CPU offload (layerwise).
|
||||
See docs/user_guide/diffusion_acceleration.md.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
@@ -24,116 +25,141 @@ NEGATIVE_PROMPT = "blurry, low quality"
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
|
||||
|
||||
MODEL_IMAGE = "Qwen/Qwen-Image"
|
||||
MODEL_2512 = "Qwen/Qwen-Image-2512"
|
||||
|
||||
# One server per feature. Alternate models; keep feature pytest ids unchanged.
|
||||
FEATURE_CASES = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL_IMAGE,
|
||||
server_args=["--enable-cpu-offload"],
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=["--enable-cpu-offload"],
|
||||
),
|
||||
id="cpu_offload",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
id="cpu_offload",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(model=MODEL_2512, server_args=["--step-execution"]),
|
||||
id="step_execution",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(model=MODEL_IMAGE, server_args=["--cache-backend", "tea_cache"]),
|
||||
id="cache_tea_cache",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL_2512,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"cache_dit",
|
||||
"--enable-layerwise-offload",
|
||||
],
|
||||
pytest.param(
|
||||
OmniServerParams(model=model, server_args=["--step-execution"]),
|
||||
id="step_execution",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
id="layerwise_offload",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL_IMAGE,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"cache_dit",
|
||||
"--ulysses-degree",
|
||||
"2",
|
||||
],
|
||||
pytest.param(
|
||||
OmniServerParams(model=model, server_args=["--cache-backend", "tea_cache"]),
|
||||
id="cache_tea_cache",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
id="ulysses_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL_2512,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"cache_dit",
|
||||
"--ring",
|
||||
"2",
|
||||
],
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"cache_dit",
|
||||
"--enable-layerwise-offload",
|
||||
],
|
||||
),
|
||||
id="layerwise_offload",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
id="ring_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL_IMAGE,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"tea_cache",
|
||||
"--cfg-parallel-size",
|
||||
"2",
|
||||
],
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"cache_dit",
|
||||
"--ulysses-degree",
|
||||
"2",
|
||||
],
|
||||
),
|
||||
id="ulysses_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
id="cfg_parallel_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL_2512,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"cache_dit",
|
||||
"--tensor-parallel-size",
|
||||
"2",
|
||||
"--vae-patch-parallel-size",
|
||||
"2",
|
||||
"--vae-use-tiling",
|
||||
"--diffusion-quantization-config",
|
||||
'{"method":"fp8"}',
|
||||
],
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"cache_dit",
|
||||
"--ring",
|
||||
"2",
|
||||
],
|
||||
),
|
||||
id="ring_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
id="vae_patch_parallel_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL_IMAGE,
|
||||
server_args=[
|
||||
"--use-hsdp",
|
||||
"--hsdp-shard-size",
|
||||
"2",
|
||||
],
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"tea_cache",
|
||||
"--cfg-parallel-size",
|
||||
"2",
|
||||
],
|
||||
),
|
||||
id="cfg_parallel_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
id="parallel_hsdp",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"cache_dit",
|
||||
"--tensor-parallel-size",
|
||||
"2",
|
||||
"--vae-patch-parallel-size",
|
||||
"2",
|
||||
"--vae-use-tiling",
|
||||
"--diffusion-quantization-config",
|
||||
'{"method":"fp8"}',
|
||||
],
|
||||
),
|
||||
id="vae_patch_parallel_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=[
|
||||
"--use-hsdp",
|
||||
"--hsdp-shard-size",
|
||||
"2",
|
||||
],
|
||||
),
|
||||
id="parallel_hsdp",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True)
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_qwen_image(omni_server: OmniServer, online_client: OnlineOmniClient):
|
||||
"""One diffusion feature per case; model is chosen in FEATURE_CASES."""
|
||||
"""Test each diffusion feature with Qwen-Image (text-to-image), one feature per case."""
|
||||
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 2,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
"true_cfg_scale": 4.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
online_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-2512"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_qwen_image_2512(omni_server: OmniServer, online_client: OnlineOmniClient):
|
||||
"""Test each diffusion feature with Qwen-Image-2512 (text-to-image), one feature per case."""
|
||||
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
|
||||
@@ -1,71 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E online serving test for Stable Audio Open text-to-audio diffusion.
|
||||
|
||||
Stable Audio Open is served through the OpenAI-compatible
|
||||
`POST /v1/audio/generate` endpoint (JSON in, binary WAV out).
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers import skip_if_gated_repo_inaccessible
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
|
||||
# Stable Audio Open is gated; override to a local bundle locally if needed.
|
||||
STABLE_AUDIO_TEST_MODEL = os.environ.get("STABLE_AUDIO_TEST_MODEL", "stabilityai/stable-audio-open-1.0")
|
||||
T2A_PROMPT = "A piano playing a gentle melody with soft room ambience."
|
||||
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
|
||||
|
||||
|
||||
def _stable_audio_server_cases(model: str):
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(model=model),
|
||||
id="t2a",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _require_stable_audio_access() -> None:
|
||||
"""Skip cleanly (before the server boots) if the gated checkpoint is inaccessible."""
|
||||
skip_if_gated_repo_inaccessible(STABLE_AUDIO_TEST_MODEL)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize("omni_server", _stable_audio_server_cases(STABLE_AUDIO_TEST_MODEL), indirect=True)
|
||||
def test_stable_audio_t2a_online(
|
||||
_require_stable_audio_access: None,
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
) -> None:
|
||||
"""Stable Audio Open text-to-audio: `/v1/audio/generate` returns a non-empty WAV.
|
||||
|
||||
Uses tiny steps / short duration to keep CI light, matching the offline smoke
|
||||
test in `tests/e2e/offline_inference/test_stable_audio_expansion.py`.
|
||||
"""
|
||||
request_config = {
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"input": T2A_PROMPT,
|
||||
"audio_length": 2.0,
|
||||
"num_inference_steps": 4,
|
||||
"guidance_scale": 7.0,
|
||||
"negative_prompt": "Low quality.",
|
||||
"seed": 42,
|
||||
"response_format": "wav",
|
||||
},
|
||||
"timeout": 300,
|
||||
}
|
||||
responses = openai_client.send_audio_generate_http_request(request_config)
|
||||
assert responses, "no response from /v1/audio/generate"
|
||||
resp = responses[0]
|
||||
assert resp.success, f"audio generate failed: {resp.status_code} {resp.error_message}"
|
||||
assert resp.status_code == 200
|
||||
@@ -474,7 +474,7 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
|
||||
assert os.environ[runtime_env_var] == "parent-value"
|
||||
|
||||
|
||||
def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_seqs(monkeypatch):
|
||||
def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypatch):
|
||||
import vllm_omni.diffusion.stage_diffusion_client as client_mod
|
||||
import vllm_omni.engine.stage_init_utils as init_mod
|
||||
|
||||
@@ -506,7 +506,7 @@ def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_
|
||||
use_inline=True,
|
||||
)
|
||||
|
||||
assert od_config.max_num_seqs == 1
|
||||
assert od_config.max_num_seqs == 4
|
||||
assert captured == {
|
||||
"model": "dummy-model",
|
||||
"config": od_config,
|
||||
@@ -517,7 +517,7 @@ def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_
|
||||
}
|
||||
|
||||
|
||||
def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_size(monkeypatch):
|
||||
def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch):
|
||||
import vllm_omni.diffusion.stage_diffusion_client as client_mod
|
||||
import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
|
||||
import vllm_omni.engine.stage_engine_startup as startup_mod
|
||||
@@ -565,84 +565,10 @@ def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_s
|
||||
)
|
||||
|
||||
assert result is sentinel_client
|
||||
assert od_config.max_num_seqs == 1
|
||||
assert od_config.max_num_seqs == 4
|
||||
assert resources.manager is proc_manager
|
||||
|
||||
|
||||
def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch):
|
||||
import vllm_omni.diffusion.stage_diffusion_client as client_mod
|
||||
import vllm_omni.engine.stage_init_utils as init_mod
|
||||
|
||||
od_config = types.SimpleNamespace(max_num_seqs=8, step_execution=True)
|
||||
monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
|
||||
monkeypatch.setattr(client_mod, "create_diffusion_client", lambda *args: object())
|
||||
metadata = _make_diffusion_metadata(0)
|
||||
|
||||
init_mod.initialize_diffusion_stage(
|
||||
0,
|
||||
"dummy-model",
|
||||
types.SimpleNamespace(),
|
||||
metadata,
|
||||
stage_init_timeout=12,
|
||||
batch_size=4,
|
||||
use_inline=True,
|
||||
)
|
||||
|
||||
assert od_config.max_num_seqs == 8
|
||||
|
||||
|
||||
def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(monkeypatch):
|
||||
import vllm_omni.diffusion.stage_diffusion_client as client_mod
|
||||
import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
|
||||
import vllm_omni.engine.stage_engine_startup as startup_mod
|
||||
|
||||
od_config = types.SimpleNamespace(
|
||||
max_num_seqs=8,
|
||||
step_execution=True,
|
||||
parallel_config=types.SimpleNamespace(world_size=1),
|
||||
)
|
||||
monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
|
||||
monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"register_stage_with_omni_master",
|
||||
lambda **kwargs: types.SimpleNamespace(
|
||||
handshake_address="tcp://127.0.0.1:26001",
|
||||
input_address="tcp://127.0.0.1:26002",
|
||||
output_address="tcp://127.0.0.1:26003",
|
||||
),
|
||||
)
|
||||
omni_master_server = types.SimpleNamespace(
|
||||
address="127.0.0.1",
|
||||
port=25000,
|
||||
release_route_port_reservations=lambda *args, **kwargs: None,
|
||||
)
|
||||
proc_manager = types.SimpleNamespace(
|
||||
addresses=types.SimpleNamespace(
|
||||
inputs=["tcp://127.0.0.1:26002"],
|
||||
outputs=["tcp://127.0.0.1:26003"],
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager)
|
||||
monkeypatch.setattr(
|
||||
client_mod.StageDiffusionClient,
|
||||
"from_addresses",
|
||||
lambda metadata, **kwargs: object(),
|
||||
)
|
||||
|
||||
startup_mod.launch_diffusion_stage_replica(
|
||||
model="dummy-model",
|
||||
stage_config=types.SimpleNamespace(),
|
||||
metadata=types.SimpleNamespace(stage_id=0),
|
||||
stage_init_timeout=12,
|
||||
batch_size=4,
|
||||
use_inline=False,
|
||||
omni_master_server=omni_master_server,
|
||||
)
|
||||
|
||||
assert od_config.max_num_seqs == 8
|
||||
|
||||
|
||||
def test_stage_runtime_initializes_stage_pools(monkeypatch):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -47,7 +47,6 @@ class FakePool:
|
||||
self.invalidated = []
|
||||
self.hub = None
|
||||
self.lb = None
|
||||
self.replica_ids = {}
|
||||
|
||||
def attach_hub(self, hub):
|
||||
self.hub = hub
|
||||
@@ -58,9 +57,7 @@ class FakePool:
|
||||
def add_client(self, input_addr, client, *, replica_id=None):
|
||||
self.added.append((input_addr, client, replica_id))
|
||||
self.clients.append(client)
|
||||
resolved_replica_id = len(self.clients) - 1 if replica_id is None else replica_id
|
||||
self.replica_ids[input_addr] = resolved_replica_id
|
||||
return resolved_replica_id
|
||||
return len(self.clients) - 1 if replica_id is None else replica_id
|
||||
|
||||
def invalidate_addr(self, input_addr):
|
||||
self.invalidated.append(input_addr)
|
||||
@@ -70,12 +67,6 @@ class FakePool:
|
||||
self.removed.append(input_addr)
|
||||
return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown"))
|
||||
|
||||
def get_client_by_addr(self, input_addr):
|
||||
return SimpleNamespace() if input_addr in self.replica_ids else None
|
||||
|
||||
def get_replica_id_by_addr(self, input_addr):
|
||||
return self.replica_ids.get(input_addr)
|
||||
|
||||
|
||||
def _snapshot(*replicas):
|
||||
return SimpleNamespace(replicas=list(replicas))
|
||||
@@ -218,16 +209,6 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey
|
||||
return replacement
|
||||
|
||||
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=factory)
|
||||
removed_replicas = []
|
||||
|
||||
async def cleanup(_request_ids):
|
||||
return None
|
||||
|
||||
controller.install_unregister_handlers(
|
||||
output_queue=asyncio.Queue(),
|
||||
cleanup_callback=cleanup,
|
||||
replica_removed_callback=lambda stage_id, replica_id: removed_replicas.append((stage_id, replica_id)),
|
||||
)
|
||||
await controller.handle_unregister(0, "tcp://old-0")
|
||||
await controller.handle_unregister(0, "tcp://old-1")
|
||||
|
||||
@@ -238,7 +219,6 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey
|
||||
assert pool.get_replica_id_by_addr("tcp://new-1") == 1
|
||||
assert pool.get_replica_id_by_addr("tcp://old-1") is None
|
||||
assert pool.available_replica_ids() == [1]
|
||||
assert removed_replicas == [(0, 0), (0, 1)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
@@ -2066,28 +2063,6 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None:
|
||||
assert metrics.output_unit_count == 3
|
||||
|
||||
|
||||
def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None:
|
||||
stage = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
|
||||
pool = StagePool(
|
||||
1,
|
||||
[stage],
|
||||
output_processor=FakeOutputProcessor(),
|
||||
stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
|
||||
)
|
||||
output = SimpleNamespace(request_id="req-image", _custom_output={"image": "non-empty"})
|
||||
pool.record_output_timestamps([output], output_ts=132.0)
|
||||
|
||||
metrics = pool.build_stage_metrics(
|
||||
[output],
|
||||
submit_ts=130.0,
|
||||
request_timestamp=120.0,
|
||||
replica_id=0,
|
||||
)
|
||||
|
||||
assert metrics.serving_time_to_first_output_ms == pytest.approx(12000.0)
|
||||
assert metrics.image_time_to_first_output_ms == pytest.approx(2000.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None:
|
||||
class FailingStageClient(FakeStageClient):
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for error propagation paths within the Orchestrator.
|
||||
|
||||
Covers:
|
||||
@@ -16,9 +16,7 @@ from types import SimpleNamespace
|
||||
|
||||
import janus
|
||||
import pytest
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine.exceptions import EngineDeadError
|
||||
from vllm.v1.serial_utils import MsgpackEncoder
|
||||
|
||||
from vllm_omni.engine.messages import (
|
||||
AddCompanionRequestMessage,
|
||||
@@ -48,6 +46,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _sampling_params(max_tokens: int = 4):
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
return SamplingParams(max_tokens=max_tokens)
|
||||
|
||||
|
||||
@@ -188,45 +188,6 @@ class _UnavailableDiffusionStageClient(FakeStageClient):
|
||||
raise StageUnavailableError(f"stage {self.stage_id} has no live replica")
|
||||
|
||||
|
||||
class _OverflowOnEncodeStageClient(FakeStageClient):
|
||||
"""Msgpack-encodes the request's sampling params on dispatch, as the real
|
||||
StageEngineCoreClient does. We use this to make sure overflow sampling params
|
||||
do not break the orchestrator."""
|
||||
|
||||
async def add_request_async(self, request, *args, **kwargs) -> None:
|
||||
MsgpackEncoder().encode(request.sampling_params)
|
||||
await super().add_request_async(request, *args, **kwargs) # unreachable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_encode_overflow_should_fail_request_not_loop(orchestrator_factory) -> None:
|
||||
"""Regression test for over-range seed handling."""
|
||||
stage0 = _OverflowOnEncodeStageClient(stage_type="llm", final_output=True)
|
||||
orchestrator_fixture = orchestrator_factory([stage0])
|
||||
|
||||
# seed > 2**63-1 cannot be msgpack-serialized -> the malicious input under test.
|
||||
overflow_params = SamplingParams(seed=2**70)
|
||||
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-overflow",
|
||||
prompt=SimpleNamespace(
|
||||
request_id="req-overflow",
|
||||
prompt_token_ids=[1, 2],
|
||||
sampling_params=overflow_params,
|
||||
),
|
||||
original_prompt={"prompt": "hi"},
|
||||
sampling_params_list=[overflow_params],
|
||||
final_stage_id=0,
|
||||
)
|
||||
|
||||
error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-overflow")
|
||||
assert error_msg.fatal is False
|
||||
assert error_msg.status_code == 400
|
||||
assert orchestrator_fixture.thread.is_alive()
|
||||
assert "req-overflow" not in orchestrator_fixture.orchestrator.request_states
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None:
|
||||
"""Multi-replica stage: one replica dying must fail only the requests bound
|
||||
|
||||
@@ -1175,7 +1175,7 @@ class TestSingleStageReplicaInitialization:
|
||||
os.environ[device_env_var] = prev_device_env
|
||||
|
||||
assert result is sentinel_client
|
||||
assert od_config.max_num_seqs is None
|
||||
assert od_config.max_num_seqs == 4
|
||||
mock_register.assert_called_once_with(
|
||||
omni_master_address="127.0.0.1",
|
||||
omni_master_port=25000,
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
@@ -108,16 +105,16 @@ class FakeEngineClient:
|
||||
self.close_result = close_result
|
||||
self.opened: list[str] = []
|
||||
self.opened_fences: list[DuplexFence | None] = []
|
||||
self.appended: list[tuple[str, str, Any, bool]] = []
|
||||
self.appended: list[tuple[str, str, object, bool]] = []
|
||||
self.append_operation_ids: list[str | None] = []
|
||||
self.appended_fences: list[DuplexFence | None] = []
|
||||
self.opened_configs: list[dict[str, Any]] = []
|
||||
self.opened_runtime_configs: list[dict[str, Any]] = []
|
||||
self.opened_configs: list[dict[str, object]] = []
|
||||
self.opened_runtime_configs: list[dict[str, object]] = []
|
||||
self.signals: list[tuple[str, str]] = []
|
||||
self.signal_fences: list[DuplexFence | None] = []
|
||||
self.signal_next_fences: list[DuplexFence | None] = []
|
||||
self.signal_session_configs: list[dict[str, Any] | None] = []
|
||||
self.signal_runtime_configs: list[dict[str, Any] | None] = []
|
||||
self.signal_session_configs: list[dict[str, object] | None] = []
|
||||
self.signal_runtime_configs: list[dict[str, object] | None] = []
|
||||
self.closed: list[tuple[str, str]] = []
|
||||
self.touched: list[tuple[str, DuplexLeaseActivity]] = []
|
||||
self.touch_fences: list[DuplexFence] = []
|
||||
@@ -139,7 +136,7 @@ class FakeEngineClient:
|
||||
runtime_config: dict[str, object] | None = None,
|
||||
timeout: float | None = None,
|
||||
fence: DuplexFence | None = None,
|
||||
) -> Any:
|
||||
) -> None:
|
||||
del timeout
|
||||
if self.fail_open:
|
||||
raise RuntimeError("open failed")
|
||||
@@ -160,7 +157,7 @@ class FakeEngineClient:
|
||||
timeout: float | None = None,
|
||||
collect_outputs: bool = True,
|
||||
fence: DuplexFence | None = None,
|
||||
) -> Any:
|
||||
) -> None:
|
||||
del timeout, collect_outputs
|
||||
self.appended.append((session_id, mode, payload, final))
|
||||
self.append_operation_ids.append(operation_id)
|
||||
@@ -194,7 +191,7 @@ class FakeEngineClient:
|
||||
next_fence: DuplexFence | None = None,
|
||||
session_config: dict[str, object] | None = None,
|
||||
runtime_config: dict[str, object] | None = None,
|
||||
) -> Any:
|
||||
) -> None:
|
||||
del timeout
|
||||
if self.fail_signal or event in self.fail_signal_events:
|
||||
raise RuntimeError("signal failed")
|
||||
@@ -212,7 +209,7 @@ class FakeEngineClient:
|
||||
reason: str = "client_close",
|
||||
timeout: float | None = None,
|
||||
fence: DuplexFence | None = None,
|
||||
) -> Any:
|
||||
) -> None:
|
||||
del timeout
|
||||
if self.fail_close:
|
||||
raise RuntimeError("close failed")
|
||||
@@ -347,7 +344,7 @@ def _project_data_plane(
|
||||
result: object,
|
||||
*,
|
||||
session: DuplexSession | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> list[dict[str, object]]:
|
||||
return list(data_plane.project(result, context=_data_plane_context(session)))
|
||||
|
||||
|
||||
@@ -4197,7 +4194,6 @@ async def test_realtime_resume_preserves_append_tail_order_across_connections():
|
||||
await asyncio.sleep(0.01)
|
||||
else:
|
||||
pytest.fail("resumable session did not open")
|
||||
assert created is not None
|
||||
session = handler._registry.get("sid-resume-append-order")
|
||||
assert session is not None
|
||||
native_state = handler._minicpmo_session_state(session)
|
||||
@@ -4327,7 +4323,6 @@ async def test_realtime_takeover_does_not_clear_new_attachment_pending_turns():
|
||||
await asyncio.sleep(0.001)
|
||||
else:
|
||||
pytest.fail("first attachment did not create the session")
|
||||
assert created is not None
|
||||
|
||||
second = TimedWebSocket(receive_timeout_s=1.0)
|
||||
second.put(
|
||||
@@ -4391,7 +4386,6 @@ async def test_realtime_takeover_drops_input_queued_by_replaced_attachment():
|
||||
await asyncio.sleep(0.001)
|
||||
else:
|
||||
pytest.fail("first attachment did not create the session")
|
||||
assert created is not None
|
||||
|
||||
second = TimedWebSocket(receive_timeout_s=1.0)
|
||||
second.put(
|
||||
@@ -5008,12 +5002,11 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
|
||||
*,
|
||||
mode: str,
|
||||
payload: object,
|
||||
operation_id: str | None = None,
|
||||
final: bool = False,
|
||||
timeout: float | None = None,
|
||||
collect_outputs: bool = True,
|
||||
fence: DuplexFence | None = None,
|
||||
) -> Any:
|
||||
):
|
||||
self.append_started.set()
|
||||
try:
|
||||
await self.release_append.wait()
|
||||
@@ -5024,7 +5017,6 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
|
||||
session_id,
|
||||
mode=mode,
|
||||
payload=payload,
|
||||
operation_id=operation_id,
|
||||
final=final,
|
||||
timeout=timeout,
|
||||
collect_outputs=collect_outputs,
|
||||
@@ -5136,9 +5128,7 @@ async def test_duplex_handler_session_update_control_failure_rolls_back_config()
|
||||
error = next(m for m in ws.sent if m.get("type") == "error")
|
||||
assert error["code"] == "runtime_signal_failed"
|
||||
assert error["runtime_control"]["error_count"] == 1
|
||||
signaled_config = engine.signal_session_configs[-1]
|
||||
assert signaled_config is not None
|
||||
assert signaled_config["instructions"] == "must roll back"
|
||||
assert engine.signal_session_configs[-1]["instructions"] == "must roll back"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -5439,7 +5429,7 @@ async def test_minicpmo_native_duplex_open_unsupported_fails_session_create():
|
||||
@pytest.mark.asyncio
|
||||
async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract():
|
||||
class LegacyEngineClient(FakeEngineClient):
|
||||
async def open_duplex_session_async( # type: ignore[override]
|
||||
async def open_duplex_session_async(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
@@ -5447,7 +5437,7 @@ async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract():
|
||||
capabilities: dict[str, object] | None = None,
|
||||
session_config: dict[str, object] | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> Any:
|
||||
) -> None:
|
||||
return await super().open_duplex_session_async(
|
||||
session_id,
|
||||
session_mode=session_mode,
|
||||
@@ -5544,9 +5534,6 @@ async def test_minicpmo_native_duplex_separates_public_and_runtime_config(monkey
|
||||
assert opened_runtime_config["duplex_stage_max_tokens"] == {"0": 20, "1": 8192}
|
||||
assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_k"] == 20
|
||||
assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_p"] == 0.8
|
||||
# The Talker YAML carries upstream's min_new_token=50 for chat TTS; a duplex
|
||||
# chunk is 26 codec samples, so check_stop would never release the request.
|
||||
assert opened_runtime_config["duplex_stage_sampling_params"]["1"] == {"min_tokens": 0}
|
||||
assert base64.b64decode(opened_runtime_config["ref_audio_data"]) == struct.pack("<1600f", *([0.25] * 1600))
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
# tests/entrypoints/openai/test_serving_speech.py
|
||||
import asyncio
|
||||
import base64
|
||||
@@ -4150,15 +4147,16 @@ class TestCosyVoice3Serving:
|
||||
error = cosyvoice3_server._validate_tts_request(request)
|
||||
assert error is None
|
||||
|
||||
def test_validate_cosyvoice3_max_new_tokens_range(self):
|
||||
"""Ensure max_new_tokens below the minimum is rejected during request validation."""
|
||||
with pytest.raises(ValidationError, match="max_new_tokens"):
|
||||
OpenAICreateSpeechRequest(
|
||||
input="Hello",
|
||||
ref_audio="data:audio/wav;base64,abc",
|
||||
ref_text="hello",
|
||||
max_new_tokens=0,
|
||||
)
|
||||
def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server):
|
||||
request = OpenAICreateSpeechRequest(
|
||||
input="Hello",
|
||||
ref_audio="data:audio/wav;base64,abc",
|
||||
ref_text="hello",
|
||||
max_new_tokens=0,
|
||||
)
|
||||
error = cosyvoice3_server._validate_tts_request(request)
|
||||
assert error is not None
|
||||
assert "max_new_tokens" in error
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),
|
||||
|
||||
@@ -378,8 +378,6 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
|
||||
"preferred",
|
||||
"--host-weight-runtime-root",
|
||||
"/var/cache/vllm-omni/hwr",
|
||||
"--dlo-host-registration-limit-gib",
|
||||
"80",
|
||||
]
|
||||
)
|
||||
|
||||
@@ -389,10 +387,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
|
||||
|
||||
assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
|
||||
assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
|
||||
assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80
|
||||
assert engine_args["host_weight_runtime_mode"] == "preferred"
|
||||
assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
|
||||
assert engine_args["dlo_host_registration_limit_gib"] == 80
|
||||
|
||||
|
||||
def test_serve_cli_accepts_diffusion_compile_controls():
|
||||
|
||||
@@ -611,13 +611,7 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat
|
||||
# external ID to avoid collisions, so this also tests mapping
|
||||
external_req_id = "req-1"
|
||||
req_id = "req-1-12345678"
|
||||
recorded_failures = []
|
||||
try:
|
||||
monkeypatch.setattr(
|
||||
app,
|
||||
"_record_request_failure_once",
|
||||
lambda request_id, reason: recorded_failures.append((request_id, reason)),
|
||||
)
|
||||
app.request_states[req_id] = ClientRequestState(
|
||||
request_id=req_id,
|
||||
external_request_id=external_req_id,
|
||||
@@ -628,7 +622,6 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat
|
||||
|
||||
assert engine.aborted == [[req_id]]
|
||||
assert external_req_id not in app.request_states
|
||||
assert recorded_failures == [(req_id, "client_abort")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -821,13 +814,7 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
|
||||
_patch_engine(monkeypatch, engine)
|
||||
|
||||
app = Omni("dummy-model")
|
||||
recorded_failures = []
|
||||
try:
|
||||
monkeypatch.setattr(
|
||||
app,
|
||||
"_record_request_failure_once",
|
||||
lambda request_id, reason: recorded_failures.append((request_id, reason)),
|
||||
)
|
||||
app.request_states["req-1"] = object()
|
||||
app.abort("req-1")
|
||||
finally:
|
||||
@@ -835,7 +822,6 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
assert engine.aborted == [["req-1"]]
|
||||
assert "req-1" not in app.request_states
|
||||
assert recorded_failures == [("req-1", "client_abort")]
|
||||
|
||||
|
||||
def test_omni_forces_final_only_on_llm_stages(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
@@ -57,9 +57,7 @@ from vllm_omni.host_weight_runtime.filesystem import (
|
||||
detect_storage_class,
|
||||
)
|
||||
from vllm_omni.host_weight_runtime.filesystem import store as filesystem_store_module
|
||||
from vllm_omni.host_weight_runtime.filesystem import writer as filesystem_writer_module
|
||||
from vllm_omni.host_weight_runtime.filesystem.locks import FileLock
|
||||
from vllm_omni.host_weight_runtime.filesystem.writer import ArtifactWriterError, FilesystemArtifactWriter
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -112,150 +110,6 @@ def _open_fd_targets() -> tuple[str, ...]:
|
||||
return tuple(targets)
|
||||
|
||||
|
||||
def _direct_writer(tmp_path: Path) -> FilesystemArtifactWriter:
|
||||
staging = tmp_path / "staging"
|
||||
staging.mkdir()
|
||||
return FilesystemArtifactWriter(staging, capacity_check=lambda _projected, _additional: None)
|
||||
|
||||
|
||||
def test_ordered_writer_hashes_payload_during_writes(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
writer = _direct_writer(tmp_path)
|
||||
weight = _weight()
|
||||
scale = torch.tensor([1.0, 2.0])
|
||||
|
||||
def fail_rescan(_path: Path, _specs: tuple[TensorWriteSpec, ...]) -> object:
|
||||
raise AssertionError("ordered payload must not be rescanned")
|
||||
|
||||
monkeypatch.setattr(filesystem_writer_module, "_scan_payload", fail_rescan)
|
||||
with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
|
||||
# The writer's ordered contract follows canonical storage-key order.
|
||||
output.write_tensor("layer.scale", scale)
|
||||
output.write_rows("layer.weight", 0, weight[:2])
|
||||
output.write_rows("layer.weight", 2, weight[2:])
|
||||
|
||||
manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
|
||||
|
||||
payload = writer.staging_dir / "weights.safetensors"
|
||||
assert manifest.files[0].sha256 == hashlib.sha256(payload.read_bytes()).hexdigest()
|
||||
expected_tensor_digests = {
|
||||
"layer.scale": hashlib.sha256(memoryview(scale.view(torch.uint8).numpy())).hexdigest(),
|
||||
"layer.weight": hashlib.sha256(memoryview(weight.contiguous().view(torch.uint8).numpy())).hexdigest(),
|
||||
}
|
||||
assert {entry.name: entry.sha256 for entry in manifest.tensors} == expected_tensor_digests
|
||||
|
||||
|
||||
def test_ordered_writer_rejects_out_of_order_payload(tmp_path: Path) -> None:
|
||||
writer = _direct_writer(tmp_path)
|
||||
output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
|
||||
|
||||
with pytest.raises(ArtifactWriterError, match="ordered tensor payload expected offset"):
|
||||
output.write_tensor("layer.weight", _weight())
|
||||
|
||||
output.abort()
|
||||
assert not list(writer.staging_dir.iterdir())
|
||||
|
||||
|
||||
def test_payload_fsync_overlaps_producer_work(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
writer = _direct_writer(tmp_path)
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
closed = threading.Event()
|
||||
original_fsync = filesystem_writer_module._fsync_payload
|
||||
|
||||
def blocking_fsync(fd: int) -> None:
|
||||
started.set()
|
||||
if not release.wait(timeout=5):
|
||||
raise TimeoutError("test did not release payload fsync")
|
||||
original_fsync(fd)
|
||||
|
||||
monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", blocking_fsync)
|
||||
output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
|
||||
output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
|
||||
output.write_tensor("layer.weight", _weight())
|
||||
close_errors: list[BaseException] = []
|
||||
|
||||
def close_output() -> None:
|
||||
try:
|
||||
output.close()
|
||||
except BaseException as exc:
|
||||
close_errors.append(exc)
|
||||
finally:
|
||||
closed.set()
|
||||
|
||||
close_thread = threading.Thread(target=close_output)
|
||||
close_thread.start()
|
||||
try:
|
||||
assert started.wait(timeout=2)
|
||||
assert closed.wait(timeout=1), "payload close waited for fsync instead of overlapping it"
|
||||
finally:
|
||||
release.set()
|
||||
close_thread.join(timeout=5)
|
||||
|
||||
assert not close_errors
|
||||
writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
|
||||
|
||||
|
||||
def test_payload_fsync_failure_prevents_publication(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
writer = _direct_writer(tmp_path)
|
||||
|
||||
def failing_fsync(fd: int) -> None:
|
||||
try:
|
||||
raise OSError(errno.EIO, "injected fsync failure")
|
||||
finally:
|
||||
os.close(fd)
|
||||
|
||||
monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", failing_fsync)
|
||||
with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
|
||||
output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
|
||||
output.write_tensor("layer.weight", _weight())
|
||||
|
||||
with pytest.raises(OSError, match="injected fsync failure"):
|
||||
writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
|
||||
|
||||
assert not (writer.staging_dir / "manifest.json").exists()
|
||||
assert not (writer.staging_dir / "READY.json").exists()
|
||||
|
||||
|
||||
def test_unordered_payload_fallback_hashes_files_in_parallel(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
writer = _direct_writer(tmp_path)
|
||||
specs = (
|
||||
TensorWriteSpec("a", (2,), torch.float32),
|
||||
TensorWriteSpec("b", (2,), torch.float32),
|
||||
)
|
||||
for name, spec in zip(("a.safetensors", "b.safetensors"), specs, strict=True):
|
||||
with writer.open_tensor_file(name, (spec,)) as output:
|
||||
output.write_tensor(spec.name, torch.tensor([1.0, 2.0]))
|
||||
|
||||
original_scan = filesystem_writer_module._scan_payload
|
||||
barrier = threading.Barrier(2)
|
||||
thread_ids: set[int] = set()
|
||||
thread_ids_lock = threading.Lock()
|
||||
|
||||
def synchronized_scan(path: Path, file_specs: tuple[TensorWriteSpec, ...]) -> object:
|
||||
with thread_ids_lock:
|
||||
thread_ids.add(threading.get_ident())
|
||||
barrier.wait(timeout=5)
|
||||
return original_scan(path, file_specs)
|
||||
|
||||
monkeypatch.setattr(filesystem_writer_module, "_scan_payload", synchronized_scan)
|
||||
manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
|
||||
|
||||
assert len(manifest.files) == 2
|
||||
assert len(thread_ids) == 2
|
||||
|
||||
|
||||
class FakeProducer:
|
||||
def __init__(
|
||||
self,
|
||||
@@ -1172,9 +1026,7 @@ def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path:
|
||||
|
||||
free_root = tmp_path / "free-limit"
|
||||
free_root.mkdir(mode=0o700)
|
||||
# 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)
|
||||
unavailable_reserve = shutil.disk_usage(free_root).free + 1
|
||||
free_limited = _make_store(free_root, capacity=CapacityPolicy(min_free_bytes=unavailable_reserve))
|
||||
free_result = free_limited.get_or_build(
|
||||
BuildRequest(identity),
|
||||
|
||||
@@ -1,499 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Verify emit-call wiring in production code paths.
|
||||
|
||||
Uses behavior and Prometheus exposition to verify production call semantics.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.metrics import OmniPrometheusMetrics
|
||||
from vllm_omni.metrics import definitions as defs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_MODEL = "emit-calls-test"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Behavioral pins — verify call semantics with mock OmniPrometheusMetrics.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_omni_base_with_mock_prom(mocker):
|
||||
"""Build a minimal OmniBase shell wired with a mock prom_metrics.
|
||||
|
||||
Mirrors the pattern in ``test_prometheus.py::TestRequestLifecycleGauges``
|
||||
— ``object.__new__`` skips ``__init__`` so we don't need a real engine.
|
||||
"""
|
||||
from vllm_omni.entrypoints.omni_base import OmniBase
|
||||
|
||||
obj = object.__new__(OmniBase)
|
||||
obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics)
|
||||
obj.request_states = {}
|
||||
obj._consumed_metric_messages = {}
|
||||
obj.log_stats = True
|
||||
return obj, obj.prom_metrics
|
||||
|
||||
|
||||
class TestFailureCounterWiring:
|
||||
"""All failure paths share one request-idempotent counter helper."""
|
||||
|
||||
def test_failure_reason_taxonomy_is_bounded(self) -> None:
|
||||
from vllm_omni.metrics.utils import normalize_failure_reason
|
||||
|
||||
assert normalize_failure_reason("client_abort") == "client_abort"
|
||||
assert normalize_failure_reason("client_disconnect") == "client_disconnect"
|
||||
assert normalize_failure_reason("stage_error") == "stage_error"
|
||||
assert normalize_failure_reason("unknown") == "unknown"
|
||||
assert normalize_failure_reason("timeout for request req-1") == "unknown"
|
||||
assert normalize_failure_reason("") == "unknown"
|
||||
assert normalize_failure_reason(None) == "unknown"
|
||||
|
||||
def test_fire_failure_counter_passes_reason_through(self, mocker) -> None:
|
||||
obj, prom = _make_omni_base_with_mock_prom(mocker)
|
||||
obj.request_states["req-1"] = SimpleNamespace(
|
||||
metrics=SimpleNamespace(e2e_done=set()),
|
||||
failure_recorded=False,
|
||||
)
|
||||
|
||||
obj._record_request_failure_once("req-1", reason="client_disconnect")
|
||||
|
||||
prom.request_failed.assert_called_once()
|
||||
prom.inc_requests_failed.assert_called_once_with("client_disconnect")
|
||||
|
||||
def test_fire_failure_counter_normalizes_unrecognized_reason(self, mocker) -> None:
|
||||
obj, prom = _make_omni_base_with_mock_prom(mocker)
|
||||
obj.request_states["req-1"] = SimpleNamespace(
|
||||
metrics=SimpleNamespace(e2e_done=set()),
|
||||
failure_recorded=False,
|
||||
)
|
||||
|
||||
obj._record_request_failure_once("req-1", reason="CUDA error for req-1")
|
||||
|
||||
prom.inc_requests_failed.assert_called_once_with("unknown")
|
||||
|
||||
def test_fire_failure_counter_skips_when_request_already_succeeded(self, mocker) -> None:
|
||||
# When the request is in e2e_done (i.e. finalize already fired
|
||||
# request_succeeded), the failure path must NOT double-count.
|
||||
obj, prom = _make_omni_base_with_mock_prom(mocker)
|
||||
obj.request_states["req-1"] = SimpleNamespace(
|
||||
metrics=SimpleNamespace(e2e_done={"req-1"}),
|
||||
failure_recorded=False,
|
||||
)
|
||||
|
||||
obj._record_request_failure_once("req-1", reason="client_disconnect")
|
||||
|
||||
prom.request_failed.assert_not_called()
|
||||
prom.inc_requests_failed.assert_not_called()
|
||||
|
||||
def test_fire_failure_counter_skips_when_request_state_missing(self, mocker) -> None:
|
||||
# No request_states entry — already popped by abort path. Fail-safe
|
||||
# must NOT raise.
|
||||
obj, prom = _make_omni_base_with_mock_prom(mocker)
|
||||
|
||||
obj._record_request_failure_once("missing-req", reason="oom")
|
||||
|
||||
prom.request_failed.assert_not_called()
|
||||
prom.inc_requests_failed.assert_not_called()
|
||||
|
||||
def test_failure_counter_is_idempotent_across_abort_and_cleanup(self, mocker) -> None:
|
||||
obj, prom = _make_omni_base_with_mock_prom(mocker)
|
||||
obj.request_states["req-1"] = SimpleNamespace(
|
||||
metrics=SimpleNamespace(
|
||||
e2e_done=set(),
|
||||
build_and_log_summary=lambda: None,
|
||||
),
|
||||
failure_recorded=False,
|
||||
)
|
||||
|
||||
obj._record_request_failure_once("req-1", reason="client_disconnect")
|
||||
obj._log_summary_and_cleanup("req-1", reason="stage_error")
|
||||
|
||||
prom.request_failed.assert_called_once()
|
||||
prom.inc_requests_failed.assert_called_once_with("client_disconnect")
|
||||
|
||||
def test_log_summary_and_cleanup_default_reason(self, mocker) -> None:
|
||||
obj, prom = _make_omni_base_with_mock_prom(mocker)
|
||||
obj.request_states["req-1"] = SimpleNamespace(
|
||||
metrics=SimpleNamespace(
|
||||
e2e_done=set(),
|
||||
build_and_log_summary=lambda: None,
|
||||
),
|
||||
failure_recorded=False,
|
||||
)
|
||||
|
||||
obj._log_summary_and_cleanup("req-1")
|
||||
|
||||
prom.inc_requests_failed.assert_called_once_with("stage_error")
|
||||
|
||||
|
||||
class TestEarlyReturnOnLogStatsOff:
|
||||
"""When --log-stats is off (default), observe methods must be silent no-ops.
|
||||
|
||||
This is the gating contract — emit sites can call unconditionally and the
|
||||
helper short-circuits.
|
||||
"""
|
||||
|
||||
def test_observe_methods_silent_when_log_stats_false(self) -> None:
|
||||
prom = OmniPrometheusMetrics(model_name=_MODEL, log_stats=False)
|
||||
# None of these should raise or write to the registry.
|
||||
prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=1.5)
|
||||
prom.observe_stage_in_queue(stage=0, in_queue_s=0.2)
|
||||
prom.observe_queue_wait(queue_wait_s=0.5)
|
||||
prom.set_stage_waiting_requests(stage=0, n_waiting=3)
|
||||
prom.observe_num_inference_steps(n_steps=20)
|
||||
prom.inc_image_count(n_images=1)
|
||||
prom.observe_image_pixels(n_pixels=512 * 512)
|
||||
prom.set_peak_memory(stage=0, peak_memory_mb=1024.0)
|
||||
prom.inc_requests_failed(reason="oom")
|
||||
prom.observe_kv_wait(connector_type="shm", kv_wait_s=0.01)
|
||||
|
||||
|
||||
class TestStageGenerationLabels:
|
||||
def test_stage_gen_time_exposes_bounded_stage_type(self) -> None:
|
||||
from prometheus_client import REGISTRY, generate_latest
|
||||
|
||||
model = _MODEL + "-stage-type"
|
||||
prom = OmniPrometheusMetrics(model_name=model, log_stats=True)
|
||||
prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=0.25)
|
||||
prom.observe_stage_gen_time(stage=1, stage_type="diffusion", gen_time_s=1.5)
|
||||
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="0",stage_type="llm"}} 1.0' in out
|
||||
assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="1",stage_type="diffusion"}} 1.0' in out
|
||||
|
||||
|
||||
class TestCounterExpositionNames:
|
||||
def test_counter_families_have_exactly_one_total_suffix(self) -> None:
|
||||
from prometheus_client import REGISTRY, generate_latest
|
||||
|
||||
model = _MODEL + "-counter-suffix"
|
||||
prom = OmniPrometheusMetrics(model_name=model, log_stats=True)
|
||||
prom.inc_image_count(2)
|
||||
prom.inc_requests_failed("stage_error")
|
||||
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
assert f'{defs.IMAGE_COUNT_METRIC}_total{{model_name="{model}"}} 2.0' in out
|
||||
assert f'{defs.REQUESTS_FAILED}_total{{model_name="{model}",reason="stage_error"}} 1.0' in out
|
||||
assert f"{defs.IMAGE_COUNT_METRIC}_total_total" not in out
|
||||
assert f"{defs.REQUESTS_FAILED}_total_total" not in out
|
||||
|
||||
|
||||
class TestQueueWaitExtraction:
|
||||
def test_present_zero_is_a_valid_observation(self) -> None:
|
||||
from vllm_omni.metrics.utils import extract_queue_wait_s
|
||||
|
||||
assert extract_queue_wait_s({"queue_wait_ms": 0.0}) == 0.0
|
||||
|
||||
def test_missing_queue_wait_is_not_synthetic_zero(self) -> None:
|
||||
from vllm_omni.metrics.utils import extract_queue_wait_s
|
||||
|
||||
assert extract_queue_wait_s({"preprocess_ms": 2.0}) is None
|
||||
assert extract_queue_wait_s(None) is None
|
||||
|
||||
|
||||
class TestStageInQueueObservation:
|
||||
def test_prometheus_records_valid_zero_wait(self) -> None:
|
||||
from prometheus_client import REGISTRY, generate_latest
|
||||
|
||||
model = _MODEL + "-zero-stage-wait"
|
||||
prom = OmniPrometheusMetrics(model_name=model, log_stats=True)
|
||||
|
||||
prom.observe_stage_in_queue(stage=1, in_queue_s=0.0)
|
||||
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
sample = f'{defs.STAGE_IN_QUEUE_S}_count{{model_name="{model}",stage="1"}} 1.0'
|
||||
assert sample in out
|
||||
|
||||
|
||||
class TestStageWorkloadMetricScope:
|
||||
@staticmethod
|
||||
def _observe(mocker, *, stage_type: str, output_unit_type: str):
|
||||
from vllm_omni.metrics.utils import observe_stage_workload_metrics
|
||||
|
||||
prom = mocker.Mock(spec=OmniPrometheusMetrics)
|
||||
stage_metrics = SimpleNamespace(
|
||||
num_inference_steps=50,
|
||||
output_unit_type=output_unit_type,
|
||||
image_pixels=1024 * 1024,
|
||||
output_unit_count=2,
|
||||
)
|
||||
observe_stage_workload_metrics(
|
||||
prom,
|
||||
stage_type=stage_type,
|
||||
stage_metrics=stage_metrics,
|
||||
)
|
||||
return prom
|
||||
|
||||
def test_diffusion_image_observes_steps_and_image_workload(self, mocker) -> None:
|
||||
prom = self._observe(mocker, stage_type="diffusion", output_unit_type="image")
|
||||
|
||||
prom.observe_num_inference_steps.assert_called_once_with(50)
|
||||
prom.observe_image_pixels.assert_called_once_with(1024 * 1024)
|
||||
prom.inc_image_count.assert_called_once_with(2)
|
||||
|
||||
def test_diffusion_video_observes_steps_only(self, mocker) -> None:
|
||||
prom = self._observe(mocker, stage_type="diffusion", output_unit_type="video")
|
||||
|
||||
prom.observe_num_inference_steps.assert_called_once_with(50)
|
||||
prom.observe_image_pixels.assert_not_called()
|
||||
prom.inc_image_count.assert_not_called()
|
||||
|
||||
def test_llm_text_observes_no_workload_metrics(self, mocker) -> None:
|
||||
prom = self._observe(mocker, stage_type="llm", output_unit_type="text")
|
||||
|
||||
prom.observe_num_inference_steps.assert_not_called()
|
||||
prom.observe_image_pixels.assert_not_called()
|
||||
prom.inc_image_count.assert_not_called()
|
||||
|
||||
def test_non_diffusion_image_observes_image_workload_only(self, mocker) -> None:
|
||||
prom = self._observe(mocker, stage_type="llm", output_unit_type="image")
|
||||
|
||||
prom.observe_num_inference_steps.assert_not_called()
|
||||
prom.observe_image_pixels.assert_called_once_with(1024 * 1024)
|
||||
prom.inc_image_count.assert_called_once_with(2)
|
||||
|
||||
@pytest.mark.parametrize("peak_memory_mb", [2048.0, 0.0])
|
||||
def test_same_finished_image_message_is_observed_exactly_once(self, mocker, peak_memory_mb: float) -> None:
|
||||
from vllm_omni.entrypoints.omni_base import OmniBase
|
||||
|
||||
obj = object.__new__(OmniBase)
|
||||
obj._enable_ar_profiler = False
|
||||
obj._consumed_metric_messages = {}
|
||||
obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics)
|
||||
obj.mod_metrics = mocker.Mock()
|
||||
obj.engine = SimpleNamespace(
|
||||
get_stage_metadata=lambda _stage_id: SimpleNamespace(
|
||||
stage_type="diffusion",
|
||||
final_output=False,
|
||||
final_output_type="image",
|
||||
)
|
||||
)
|
||||
stage_metrics = SimpleNamespace(
|
||||
stage_gen_time_ms=1000.0,
|
||||
diffusion_metrics={"diffusion_engine_exec_time_s": 0.8, "scheduler_queue_wait_s": 0.1},
|
||||
num_inference_steps=20,
|
||||
output_unit_type="image",
|
||||
image_pixels=1024 * 1024,
|
||||
output_unit_count=1,
|
||||
serving_time_to_first_output_ms=250.0,
|
||||
image_time_to_first_output_ms=250.0,
|
||||
pipeline_timings={},
|
||||
)
|
||||
result = SimpleNamespace(
|
||||
request_id="req-replay",
|
||||
stage_id=1,
|
||||
replica_id=0,
|
||||
stage_submit_ts=10.0,
|
||||
metrics=stage_metrics,
|
||||
engine_outputs=SimpleNamespace(
|
||||
stage_durations={},
|
||||
peak_memory_mb=peak_memory_mb,
|
||||
finished=True,
|
||||
final_output_type="image",
|
||||
),
|
||||
)
|
||||
aggregator = mocker.Mock()
|
||||
aggregator.stage_events = {}
|
||||
aggregator.stage_first_ts = [None, None]
|
||||
aggregator.stage_last_ts = [None, None]
|
||||
|
||||
for _ in range(2):
|
||||
assert obj._process_single_result(result, 1, aggregator, {}, 0.0, 1) is None
|
||||
|
||||
aggregator.on_stage_metrics.assert_called_once_with(1, "req-replay", stage_metrics, "image")
|
||||
obj.prom_metrics.inc_image_count.assert_called_once_with(1)
|
||||
obj.prom_metrics.observe_image_pixels.assert_called_once_with(1024 * 1024)
|
||||
obj.prom_metrics.observe_num_inference_steps.assert_called_once_with(20)
|
||||
obj.prom_metrics.observe_stage_in_queue.assert_called_once_with(1, 0.1)
|
||||
if peak_memory_mb > 0:
|
||||
obj.prom_metrics.set_peak_memory.assert_called_once_with(1, peak_memory_mb)
|
||||
else:
|
||||
obj.prom_metrics.set_peak_memory.assert_not_called()
|
||||
obj.mod_metrics.observe_image_ttfp.assert_called_once_with("1", "0", 0.25)
|
||||
|
||||
|
||||
class TestStageWaitingAggregation:
|
||||
@staticmethod
|
||||
def _make_orchestrator(mocker):
|
||||
from vllm_omni.engine.orchestrator import Orchestrator
|
||||
|
||||
orchestrator = object.__new__(Orchestrator)
|
||||
orchestrator._prom_metrics = mocker.Mock()
|
||||
orchestrator._stage_replica_waiting = {}
|
||||
return orchestrator
|
||||
|
||||
def test_multiple_replicas_are_summed(self, mocker) -> None:
|
||||
orchestrator = self._make_orchestrator(mocker)
|
||||
|
||||
orchestrator._update_stage_replica_waiting(1, 0, 3)
|
||||
orchestrator._update_stage_replica_waiting(1, 1, 5)
|
||||
|
||||
orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 8)
|
||||
|
||||
def test_zero_update_and_stage_isolation(self, mocker) -> None:
|
||||
orchestrator = self._make_orchestrator(mocker)
|
||||
|
||||
orchestrator._update_stage_replica_waiting(1, 0, 3)
|
||||
orchestrator._update_stage_replica_waiting(2, 0, 7)
|
||||
orchestrator._update_stage_replica_waiting(1, 0, 0)
|
||||
|
||||
orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 0)
|
||||
assert orchestrator._stage_replica_waiting[(2, 0)] == 7
|
||||
|
||||
def test_dead_replica_snapshot_is_removed(self, mocker) -> None:
|
||||
orchestrator = self._make_orchestrator(mocker)
|
||||
orchestrator._update_stage_replica_waiting(1, 0, 3)
|
||||
orchestrator._update_stage_replica_waiting(1, 1, 5)
|
||||
|
||||
orchestrator._remove_stage_replica_waiting(1, 1)
|
||||
|
||||
orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 3)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Behavioral pin: call sites now scope workload observations by stage/output
|
||||
# type. Keep the Prometheus wrapper's <= 0 guard as a second line of defense.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestZeroGuardContract:
|
||||
"""Zero-valued observations must not bump histogram ``_count``."""
|
||||
|
||||
def test_zero_image_pixels_not_observed(self) -> None:
|
||||
from prometheus_client import REGISTRY, generate_latest
|
||||
|
||||
prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-guard")
|
||||
prom.observe_image_pixels(n_pixels=0)
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
# ``.labels()`` in OmniPrometheusMetrics.__init__ already creates the
|
||||
# child sample, so the ``_count`` line exists with value 0.0 even
|
||||
# without observations. The guard prevents ``_count`` from bumping
|
||||
# to 1.0 — assert the value stays at 0.
|
||||
needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-zero-guard"}}'
|
||||
count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)]
|
||||
assert count_lines, "expected image_pixels_count line to exist after OmniPrometheusMetrics construction"
|
||||
assert float(count_lines[0].split()[-1]) == 0.0, (
|
||||
"zero-pixel observation leaked to registry — guard should early-return on <= 0"
|
||||
)
|
||||
|
||||
def test_zero_num_inference_steps_not_observed(self) -> None:
|
||||
from prometheus_client import REGISTRY, generate_latest
|
||||
|
||||
prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-steps")
|
||||
prom.observe_num_inference_steps(n_steps=0)
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
needle = f'vllm_omni:num_inference_steps_count{{model_name="{_MODEL}-zero-steps"}}'
|
||||
count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)]
|
||||
assert count_lines, "expected num_inference_steps_count line to exist after construction"
|
||||
assert float(count_lines[0].split()[-1]) == 0.0, (
|
||||
"zero-step observation leaked to registry — guard should early-return on <= 0"
|
||||
)
|
||||
|
||||
def test_positive_image_pixels_observed(self) -> None:
|
||||
# Sanity check: positive values DO get observed.
|
||||
from prometheus_client import REGISTRY, generate_latest
|
||||
|
||||
prom = OmniPrometheusMetrics(model_name=_MODEL + "-pos-guard")
|
||||
prom.observe_image_pixels(n_pixels=512)
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-pos-guard"}}'
|
||||
count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)]
|
||||
assert count_lines, "expected image_pixels_count line after positive observation"
|
||||
assert float(count_lines[0].split()[-1]) == 1.0, "positive-pixel observation did not increment count to 1"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# KV-wait emit wiring — scheduler ENTER/EXIT lifecycle + orchestrator dispatch.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_scheduler_shell() -> OmniARScheduler:
|
||||
"""Minimal OmniARScheduler shell — skips upstream __init__ via object.__new__."""
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
|
||||
obj = object.__new__(OmniARScheduler)
|
||||
obj.requests = {}
|
||||
obj.running = []
|
||||
obj.waiting = []
|
||||
obj._kv_wait_start_ts = {}
|
||||
obj._omni_kv_config = None
|
||||
return obj
|
||||
|
||||
|
||||
class TestKvWaitSchedulerEmit:
|
||||
"""Scheduler-side half of kv_wait_s: pops start ts, carries the wait across."""
|
||||
|
||||
def test_skips_when_no_start_ts_recorded(self) -> None:
|
||||
sched = _make_scheduler_shell()
|
||||
outputs: dict[int, list] = {}
|
||||
|
||||
sched._emit_kv_wait_output(outputs, "req-no-wait", req=SimpleNamespace(client_index=0))
|
||||
|
||||
assert outputs == {}, "emit must not append output when no start ts recorded"
|
||||
assert "req-no-wait" not in sched._kv_wait_start_ts
|
||||
|
||||
def test_emits_wait_duration_and_connector_type(self) -> None:
|
||||
sched = _make_scheduler_shell()
|
||||
|
||||
sched._kv_wait_start_ts["req-1"] = time.monotonic() - 0.25
|
||||
outputs: dict[int, list] = {}
|
||||
live_req = SimpleNamespace(client_index=3)
|
||||
|
||||
sched._emit_kv_wait_output(outputs, "req-1", req=live_req)
|
||||
|
||||
assert "req-1" not in sched._kv_wait_start_ts, "start ts must be popped after emit"
|
||||
assert 3 in outputs and len(outputs[3]) == 1
|
||||
params = outputs[3][0].kv_transfer_params
|
||||
assert "kv_wait_s" in params and "connector_type" in params
|
||||
assert 0.0 < params["kv_wait_s"] < 1.0
|
||||
assert params["connector_type"] == "unknown"
|
||||
|
||||
def test_connector_type_resolved_from_omni_kv_config(self) -> None:
|
||||
sched = _make_scheduler_shell()
|
||||
sched._omni_kv_config = {"connector_config": {"type": "SharedMemoryConnector"}}
|
||||
|
||||
sched._kv_wait_start_ts["req-3"] = time.monotonic() - 0.01
|
||||
outputs: dict[int, list] = {}
|
||||
|
||||
sched._emit_kv_wait_output(outputs, "req-3", req=SimpleNamespace(client_index=0))
|
||||
|
||||
eco = outputs[0][0]
|
||||
assert eco.kv_transfer_params["connector_type"] == "SharedMemoryConnector"
|
||||
|
||||
|
||||
class TestKvWaitTerminalCleanup:
|
||||
def test_finish_requests_clears_requested_wait_only(self) -> None:
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
scheduler = _make_scheduler_shell()
|
||||
scheduler._kv_wait_start_ts = {"req-1": 10.0, "req-2": 20.0}
|
||||
|
||||
finished = scheduler.finish_requests(
|
||||
(request_id for request_id in ["req-1"]),
|
||||
RequestStatus.FINISHED_ABORTED,
|
||||
)
|
||||
|
||||
assert finished == []
|
||||
assert "req-1" not in scheduler._kv_wait_start_ts
|
||||
assert scheduler._kv_wait_start_ts["req-2"] == 20.0
|
||||
|
||||
def test_cleanup_is_idempotent_when_timestamp_is_already_gone(self) -> None:
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
scheduler = _make_scheduler_shell()
|
||||
|
||||
scheduler.finish_requests("req-already-emitted", RequestStatus.FINISHED_ABORTED)
|
||||
|
||||
assert scheduler._kv_wait_start_ts == {}
|
||||
@@ -1,10 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from prometheus_client import REGISTRY, generate_latest
|
||||
|
||||
@@ -15,7 +10,6 @@ from vllm_omni.metrics.modality import (
|
||||
observe_audio_streaming_finalize,
|
||||
observe_modality_at_finalize,
|
||||
)
|
||||
from vllm_omni.metrics.stats import OrchestratorAggregator
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -52,11 +46,6 @@ _EXPECTED_FAMILIES = [
|
||||
defs.DIFFUSION_EXEC_PER_STEP_S,
|
||||
defs.DIFFUSION_PREPROCESS_S,
|
||||
defs.DIFFUSION_POSTPROCESS_S,
|
||||
defs.VAE_DECODE_S,
|
||||
defs.DIFFUSION_FORWARD_S,
|
||||
defs.DIFFUSION_KV_LOAD_S,
|
||||
defs.IMAGE_TTFP_S,
|
||||
defs.DENOISE_STEP_LATENCY_S,
|
||||
]
|
||||
|
||||
|
||||
@@ -74,11 +63,6 @@ class TestRegistration:
|
||||
mod.observe_diffusion_exec_per_step("s", "r", 0.01)
|
||||
mod.observe_diffusion_preprocess("s", "r", 0.01)
|
||||
mod.observe_diffusion_postprocess("s", "r", 0.2)
|
||||
mod.observe_vae_decode("s", "r", 0.3)
|
||||
mod.observe_diffusion_forward("s", "r", 0.8)
|
||||
mod.observe_diffusion_kv_load("s", "r", 0.04)
|
||||
mod.observe_image_ttfp("s", "r", 2.5)
|
||||
mod.observe_denoise_step_latency("s", "r", 0.05)
|
||||
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
for name in _EXPECTED_FAMILIES:
|
||||
@@ -263,21 +247,6 @@ class _StubModMetrics:
|
||||
def observe_diffusion_postprocess(self, s, r, v):
|
||||
self.calls.append(("observe_diffusion_postprocess", s, r, v))
|
||||
|
||||
def observe_vae_decode(self, s, r, seconds):
|
||||
self.calls.append(("observe_vae_decode", s, r, seconds))
|
||||
|
||||
def observe_diffusion_forward(self, s, r, seconds):
|
||||
self.calls.append(("observe_diffusion_forward", s, r, seconds))
|
||||
|
||||
def observe_diffusion_kv_load(self, s, r, seconds):
|
||||
self.calls.append(("observe_diffusion_kv_load", s, r, seconds))
|
||||
|
||||
def observe_image_ttfp(self, s, r, seconds):
|
||||
self.calls.append(("observe_image_ttfp", s, r, seconds))
|
||||
|
||||
def observe_denoise_step_latency(self, s, r, seconds):
|
||||
self.calls.append(("observe_denoise_step_latency", s, r, seconds))
|
||||
|
||||
|
||||
class _Bag:
|
||||
def __init__(self, **kw):
|
||||
@@ -337,13 +306,11 @@ class TestObserveModalityAtFinalize:
|
||||
def test_diffusion_path_full(self):
|
||||
stub = _StubModMetrics()
|
||||
stage_metrics = _Bag(
|
||||
output_unit_type="image",
|
||||
stage_gen_time_ms=1000.0,
|
||||
diffusion_metrics={
|
||||
"preprocess_time_s": 0.05,
|
||||
"diffusion_engine_exec_time_s": 1.2,
|
||||
"postprocess_time_s": 0.3,
|
||||
"forward_time_s": 1.0,
|
||||
"image_num": 1,
|
||||
"resolution": 1024,
|
||||
"num_inference_steps": 50,
|
||||
@@ -361,71 +328,6 @@ class TestObserveModalityAtFinalize:
|
||||
assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
|
||||
assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls
|
||||
assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls
|
||||
assert ("observe_denoise_step_latency", "2", "0", pytest.approx(0.02)) in stub.calls
|
||||
|
||||
def test_denoise_step_latency_requires_forward_time(self):
|
||||
stub = _StubModMetrics()
|
||||
stage_metrics = _Bag(
|
||||
output_unit_type="image",
|
||||
num_inference_steps=20,
|
||||
diffusion_metrics={"diffusion_engine_exec_time_s": 1.2},
|
||||
)
|
||||
observe_modality_at_finalize(
|
||||
stub,
|
||||
output_type="image",
|
||||
stage_id=2,
|
||||
replica_id=0,
|
||||
stage_metrics=stage_metrics,
|
||||
engine_outputs=_Bag(),
|
||||
)
|
||||
assert not any(c[0] == "observe_denoise_step_latency" for c in stub.calls)
|
||||
|
||||
def test_stage_snapshot_keeps_forward_based_denoise_latency(self):
|
||||
aggregator = OrchestratorAggregator(
|
||||
num_stages=1,
|
||||
log_stats=False,
|
||||
wall_start_ts=0.0,
|
||||
final_stage_id_for_e2e=0,
|
||||
)
|
||||
aggregator.diffusion_metrics["req-image"]["forward_time_s"] = 0.8
|
||||
stage_metrics = _Bag(
|
||||
stage_id=None,
|
||||
request_id=None,
|
||||
final_output_type=None,
|
||||
num_tokens_in=0,
|
||||
num_tokens_out=0,
|
||||
stage_gen_time_ms=1200.0,
|
||||
rx_transfer_bytes=0,
|
||||
rx_decode_time_ms=0.0,
|
||||
rx_in_flight_time_ms=0.0,
|
||||
stage_stats=_Bag(),
|
||||
diffusion_metrics=None,
|
||||
num_inference_steps=20,
|
||||
output_unit_type="image",
|
||||
denoise_step_latency_ms=0.0,
|
||||
)
|
||||
|
||||
stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image")
|
||||
|
||||
assert stats.denoise_step_latency_ms == pytest.approx(40.0)
|
||||
|
||||
def test_stage_snapshot_falls_back_to_stage_time_without_profiler(self):
|
||||
aggregator = object.__new__(OrchestratorAggregator)
|
||||
stage_metrics = _Bag(
|
||||
stage_id=None,
|
||||
request_id=None,
|
||||
final_output_type=None,
|
||||
stage_gen_time_ms=1200.0,
|
||||
diffusion_metrics=None,
|
||||
num_inference_steps=20,
|
||||
output_unit_type="image",
|
||||
denoise_step_latency_ms=0.0,
|
||||
)
|
||||
|
||||
stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image")
|
||||
|
||||
assert stats.diffusion_metrics is None
|
||||
assert stats.denoise_step_latency_ms == pytest.approx(60.0)
|
||||
|
||||
def test_diffusion_per_step_skipped_without_num_inference_steps(self):
|
||||
stub = _StubModMetrics()
|
||||
@@ -520,80 +422,6 @@ class TestObserveModalityAtFinalize:
|
||||
assert stub.calls == []
|
||||
|
||||
|
||||
class TestDiffusionUnitConversion:
|
||||
def test_accumulator_converts_all_engine_milliseconds_to_seconds(self):
|
||||
aggregator = OrchestratorAggregator(
|
||||
num_stages=1,
|
||||
log_stats=False,
|
||||
wall_start_ts=0.0,
|
||||
final_stage_id_for_e2e=0,
|
||||
)
|
||||
engine_output = SimpleNamespace(
|
||||
metrics={
|
||||
"preprocess_time_ms": 50.0,
|
||||
"diffusion_engine_exec_time_ms": 1500.0,
|
||||
"postprocess_time_ms": 20.0,
|
||||
"vae_decode_time_ms": 300.0,
|
||||
"forward_time_ms": 800.0,
|
||||
"scheduler_queue_wait_ms": 125.0,
|
||||
"kv_recv_time_ms": 40.0,
|
||||
}
|
||||
)
|
||||
|
||||
aggregator.accumulate_diffusion_metrics("diffusion", "req-units", engine_output)
|
||||
|
||||
assert dict(aggregator.diffusion_metrics["req-units"]) == {
|
||||
"preprocess_time_s": pytest.approx(0.05),
|
||||
"diffusion_engine_exec_time_s": pytest.approx(1.5),
|
||||
"postprocess_time_s": pytest.approx(0.02),
|
||||
"vae_decode_time_s": pytest.approx(0.3),
|
||||
"forward_time_s": pytest.approx(0.8),
|
||||
"scheduler_queue_wait_s": pytest.approx(0.125),
|
||||
"kv_recv_time_s": pytest.approx(0.04),
|
||||
}
|
||||
|
||||
|
||||
class TestObserveDiffusionFinalize:
|
||||
@pytest.mark.parametrize(
|
||||
("metric_key", "call_name", "value"),
|
||||
[
|
||||
("vae_decode_time_s", "observe_vae_decode", 0.3),
|
||||
("forward_time_s", "observe_diffusion_forward", 0.8),
|
||||
("kv_recv_time_s", "observe_diffusion_kv_load", 0.04),
|
||||
],
|
||||
)
|
||||
def test_emits_optional_breakdown_metric(self, metric_key, call_name, value):
|
||||
stub = _StubModMetrics()
|
||||
observe_modality_at_finalize(
|
||||
stub,
|
||||
output_type="image",
|
||||
stage_id=2,
|
||||
replica_id=0,
|
||||
stage_metrics=_Bag(diffusion_metrics={metric_key: value}),
|
||||
engine_outputs=_Bag(),
|
||||
)
|
||||
|
||||
assert (call_name, "2", "0", value) in stub.calls
|
||||
|
||||
def test_audio_diffusion_stage_fires_both_paths(self):
|
||||
stub = _StubModMetrics()
|
||||
observe_modality_at_finalize(
|
||||
stub,
|
||||
output_type="audio",
|
||||
stage_id=1,
|
||||
replica_id=0,
|
||||
stage_metrics=_Bag(
|
||||
stage_gen_time_ms=1500.0,
|
||||
audio_generated_frames=24000,
|
||||
diffusion_metrics={"diffusion_engine_exec_time_s": 1.2},
|
||||
),
|
||||
engine_outputs=_Bag(multimodal_output={"audio_sample_rate": 24000}),
|
||||
)
|
||||
|
||||
assert ("observe_diffusion_exec", "1", "0", 1.2) in stub.calls
|
||||
assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
|
||||
|
||||
|
||||
class TestObserveAudioFirstPacket:
|
||||
def test_observes_with_valid_inputs(self):
|
||||
stub = _StubModMetrics()
|
||||
@@ -713,39 +541,3 @@ class TestBucketSelection:
|
||||
# SECONDS_FAST_BUCKETS includes le=0.001 which is absent from SECONDS_BUCKETS.
|
||||
fast_marker = f'{defs.AUDIO_UNDERRUN_S}_bucket{{le="0.001"'
|
||||
assert fast_marker in out, "audio_underrun_s should use SECONDS_FAST_BUCKETS containing le=0.001"
|
||||
|
||||
def test_vae_decode_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None:
|
||||
# VAE decode can stretch into seconds for high-resolution latents;
|
||||
# stays on the wide bucket so the 300 s ceiling covers video VAE.
|
||||
stage, replica = "vae_decode_sec", "0"
|
||||
mod.observe_vae_decode(stage, replica, 1.5)
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
sec_marker = f'{defs.VAE_DECODE_S}_bucket{{le="300.0"'
|
||||
assert sec_marker in out, "vae_decode_s should use SECONDS_BUCKETS containing le=300.0"
|
||||
|
||||
def test_diffusion_forward_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None:
|
||||
# Denoise loop can exceed the 60s fast-bucket ceiling for video gen.
|
||||
stage, replica = "diff_forward_sec", "0"
|
||||
mod.observe_diffusion_forward(stage, replica, 2.0)
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
sec_marker = f'{defs.DIFFUSION_FORWARD_S}_bucket{{le="300.0"'
|
||||
assert sec_marker in out, "diffusion_forward_s should use SECONDS_BUCKETS containing le=300.0"
|
||||
|
||||
def test_denoise_step_latency_uses_seconds_fast_buckets(self, mod: OmniModalityMetrics) -> None:
|
||||
# Per-step denoise is sub-second; stays on the fast bucket for ms-level
|
||||
# resolution.
|
||||
stage, replica = "denoise_fast", "0"
|
||||
mod.observe_denoise_step_latency(stage, replica, 0.075)
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
fast_marker = f'{defs.DENOISE_STEP_LATENCY_S}_bucket{{le="0.001"'
|
||||
assert fast_marker in out, "denoise_step_latency_s should use SECONDS_FAST_BUCKETS containing le=0.001"
|
||||
|
||||
|
||||
class TestDenoiseStepLatencyGuard:
|
||||
@pytest.mark.parametrize(("value", "expected_count"), [(0.0, None), (-0.1, None), (0.05, 1.0)])
|
||||
def test_nonpositive_values_are_skipped(self, mod: OmniModalityMetrics, value, expected_count) -> None:
|
||||
stage, replica = f"denoise_{value}", "0"
|
||||
mod.observe_denoise_step_latency(stage, replica, value)
|
||||
out = generate_latest(REGISTRY).decode()
|
||||
prefix = f'{defs.DENOISE_STEP_LATENCY_S}_count{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
|
||||
assert _sample_value(out, prefix) == expected_count
|
||||
@@ -281,153 +281,6 @@ class TestCFM:
|
||||
|
||||
assert out.shape == mu.shape
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_trt_estimator_uses_stream_dependencies(self, monkeypatch):
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ConditionalCFM
|
||||
|
||||
class FakeStream:
|
||||
def __init__(self, name, state):
|
||||
self.name = name
|
||||
self.state = state
|
||||
self.cuda_stream = hash(name)
|
||||
self.synchronize_calls = 0
|
||||
self.waited_on = []
|
||||
|
||||
def synchronize(self):
|
||||
self.synchronize_calls += 1
|
||||
|
||||
def wait_stream(self, stream):
|
||||
self.waited_on.append(stream)
|
||||
|
||||
class FakeStreamContext:
|
||||
def __init__(self, stream):
|
||||
self.stream = stream
|
||||
|
||||
def __enter__(self):
|
||||
self.previous = self.stream.state.current
|
||||
self.stream.state.current = self.stream
|
||||
return self.stream
|
||||
|
||||
def __exit__(self, exc_type, exc, traceback):
|
||||
self.stream.state.current = self.previous
|
||||
|
||||
state = SimpleNamespace(current=None)
|
||||
caller_stream = FakeStream("caller", state)
|
||||
estimator_stream = FakeStream("estimator", state)
|
||||
state.current = caller_stream
|
||||
monkeypatch.setattr(torch.cuda, "current_stream", lambda *args, **kwargs: state.current)
|
||||
stream_contexts = []
|
||||
|
||||
def stream_context(stream):
|
||||
stream_contexts.append(stream)
|
||||
return FakeStreamContext(stream)
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "stream", stream_context)
|
||||
|
||||
class FakeContext:
|
||||
def __init__(self):
|
||||
self.execute_stream = None
|
||||
|
||||
def set_input_shape(self, name, shape):
|
||||
pass
|
||||
|
||||
def set_tensor_address(self, name, address):
|
||||
pass
|
||||
|
||||
def execute_async_v3(self, stream):
|
||||
self.execute_stream = stream
|
||||
return True
|
||||
|
||||
class FakeEngine:
|
||||
@staticmethod
|
||||
def get_tensor_name(index):
|
||||
return f"tensor_{index}"
|
||||
|
||||
context = FakeContext()
|
||||
|
||||
class FakeEstimatorPool:
|
||||
io_dtype = torch.float32
|
||||
|
||||
def __init__(self):
|
||||
self.released = []
|
||||
|
||||
def acquire_estimator(self):
|
||||
return [context, estimator_stream], FakeEngine()
|
||||
|
||||
def release_estimator(self, released_context, released_stream):
|
||||
self.released.append((released_context, released_stream))
|
||||
|
||||
estimator_pool = FakeEstimatorPool()
|
||||
cfm = ConditionalCFM(
|
||||
in_channels=80,
|
||||
cfm_params=DictConfig(
|
||||
{
|
||||
"sigma_min": 1e-6,
|
||||
"solver": "euler",
|
||||
"t_scheduler": "cosine",
|
||||
"training_cfg_rate": 0.2,
|
||||
"inference_cfg_rate": 0.7,
|
||||
}
|
||||
),
|
||||
n_spks=1,
|
||||
spk_emb_dim=80,
|
||||
estimator=estimator_pool,
|
||||
)
|
||||
|
||||
x = torch.randn(2, 80, 4)
|
||||
mask = torch.ones(2, 1, 4)
|
||||
mu = torch.randn(2, 80, 4)
|
||||
timestep = torch.randn(2)
|
||||
speakers = torch.randn(2, 80)
|
||||
condition = torch.randn(2, 80, 4)
|
||||
|
||||
output = cfm.forward_estimator(x, mask, mu, timestep, speakers, condition)
|
||||
|
||||
assert output.shape == x.shape
|
||||
assert caller_stream.synchronize_calls == 0
|
||||
assert estimator_stream.synchronize_calls == 0
|
||||
assert estimator_stream.waited_on == [caller_stream]
|
||||
assert caller_stream.waited_on == [estimator_stream]
|
||||
assert stream_contexts == [estimator_stream]
|
||||
assert context.execute_stream == estimator_stream.cuda_stream
|
||||
assert estimator_pool.released == [(context, estimator_stream)]
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_trt_context_pool_stores_cuda_stream(self, monkeypatch):
|
||||
from vllm_omni.model_executor.models.cosyvoice3.flow_estimator_trt import TrtContextWrapper
|
||||
|
||||
execution_context = object()
|
||||
|
||||
class FakeEngine:
|
||||
@staticmethod
|
||||
def create_execution_context():
|
||||
return execution_context
|
||||
|
||||
cuda_stream = object()
|
||||
monkeypatch.setattr(torch.cuda, "Stream", lambda device: cuda_stream)
|
||||
|
||||
def reject_stream_context(stream):
|
||||
pytest.fail("TrtContextWrapper must store the CUDA stream, not a StreamContext")
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "stream", reject_stream_context)
|
||||
|
||||
engine = FakeEngine()
|
||||
wrapper = TrtContextWrapper(engine, device="cuda:0")
|
||||
[context, stream], acquired_engine = wrapper.acquire_estimator()
|
||||
|
||||
assert context is execution_context
|
||||
assert stream is cuda_stream
|
||||
assert acquired_engine is engine
|
||||
|
||||
wrapper.release_estimator(context, stream)
|
||||
[reused_context, reused_stream], _ = wrapper.acquire_estimator()
|
||||
assert reused_context is context
|
||||
assert reused_stream is stream
|
||||
|
||||
|
||||
class TestSDPAFallback:
|
||||
"""Test SDPA fallback for float32 inputs."""
|
||||
|
||||
@@ -177,6 +177,23 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na
|
||||
assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_magi_human_extra_registry_declares_request_and_response_params() -> None:
|
||||
assert get_extra_body_params("MagiHumanPipeline") == frozenset(
|
||||
{
|
||||
"seconds",
|
||||
"audio_path",
|
||||
"image_path",
|
||||
"sr_height",
|
||||
"sr_width",
|
||||
"sr_num_inference_steps",
|
||||
}
|
||||
)
|
||||
assert get_extra_output_params("MagiHumanPipeline") == frozenset()
|
||||
assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_ltx_extra_registry_declares_official_guidance_params() -> None:
|
||||
|
||||
@@ -57,6 +57,7 @@ EXCLUDED_MODELS = [
|
||||
"OmniGen2Pipeline",
|
||||
"HeliosPipeline",
|
||||
"HeliosPyramidPipeline",
|
||||
"DreamIDOmniPipeline",
|
||||
"SenseNovaU1Pipeline",
|
||||
"HunyuanVideo15Pipeline",
|
||||
"HunyuanVideo15ImageToVideoPipeline",
|
||||
@@ -70,6 +71,7 @@ EXCLUDED_MODELS = [
|
||||
# Requires camera/action assets and persistent AR state; covered by its
|
||||
# dedicated contract tests and real-checkpoint E2E smoke instead.
|
||||
"LingBotWorldCausalDMDPipeline",
|
||||
"MagiHumanPipeline",
|
||||
"OmniVoicePipeline",
|
||||
"OmniVoice",
|
||||
"Cosmos3OmniDiffusersPipeline",
|
||||
|
||||
@@ -1,306 +0,0 @@
|
||||
# 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,9 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import warnings
|
||||
from collections import defaultdict
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import field, make_dataclass
|
||||
|
||||
import numpy as np
|
||||
@@ -19,11 +15,6 @@ _FLOAT_LIST_TYPE = list[float]
|
||||
_INT_LIST_TYPE = list[int]
|
||||
|
||||
_MULTIMODAL_BENCHMARK_FIELDS = [
|
||||
("num_ttft_samples", int, field(default=0)),
|
||||
("num_tpot_samples", int, field(default=0)),
|
||||
("num_itl_samples", int, field(default=0)),
|
||||
("num_audio_ttfp_samples", int, field(default=0)),
|
||||
("num_audio_rtf_samples", int, field(default=0)),
|
||||
(defs.MEAN_AUDIO_TTFP_MS, float, field(default=0.0)),
|
||||
(defs.MEDIAN_AUDIO_TTFP_MS, float, field(default=0.0)),
|
||||
(defs.STD_AUDIO_TTFP_MS, float, field(default=0.0)),
|
||||
@@ -185,17 +176,6 @@ def _stage_modality_flags(
|
||||
)
|
||||
|
||||
|
||||
def has_metric_samples(metrics: object, metric_name: str) -> bool:
|
||||
sample_count_attr = {
|
||||
"ttft": "num_ttft_samples",
|
||||
"tpot": "num_tpot_samples",
|
||||
"itl": "num_itl_samples",
|
||||
defs.AUDIO_TTFP: "num_audio_ttfp_samples",
|
||||
defs.AUDIO_RTF: "num_audio_rtf_samples",
|
||||
}.get(metric_name)
|
||||
return sample_count_attr is None or getattr(metrics, sample_count_attr, 0) > 0
|
||||
|
||||
|
||||
def print_metrics(
|
||||
task_type,
|
||||
selected_percentile_metrics,
|
||||
@@ -328,9 +308,6 @@ def process_one_metric(
|
||||
metric_attribute_name: str,
|
||||
metrics: MultiModalsBenchmarkMetrics,
|
||||
):
|
||||
if not has_metric_samples(metrics, metric_attribute_name):
|
||||
return
|
||||
|
||||
metric_header_map = {
|
||||
"ttft": "Time to First Token",
|
||||
"tpot": "Time per Output Token (excl. 1st token)",
|
||||
@@ -741,12 +718,10 @@ def calculate_metrics(
|
||||
good_completed = 0
|
||||
itls: list[float] = []
|
||||
tpots: list[float] = []
|
||||
all_tpots: list[float | None] = []
|
||||
all_tpots: list[float] = []
|
||||
ttfts: list[float] = []
|
||||
goodput_ttfts: list[float | None] = []
|
||||
e2els: list[float] = []
|
||||
audio_ttfps: list[float] = []
|
||||
goodput_audio_ttfps: list[float | None] = []
|
||||
audio_rtfs: list[float] = []
|
||||
audio_duration: list[float] = []
|
||||
audio_frames: list[int] = []
|
||||
@@ -773,7 +748,7 @@ def calculate_metrics(
|
||||
output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
|
||||
actual_output_lens.append(output_len)
|
||||
total_input += outputs[i].prompt_len
|
||||
tpot: float | None = 0.0
|
||||
tpot = 0
|
||||
if output_len > 1:
|
||||
if outputs[i].itl:
|
||||
# Use mean(ITL) directly so per-request TPOT == mean(ITL).
|
||||
@@ -781,33 +756,19 @@ def calculate_metrics(
|
||||
# bundle multiple tokens per chunk, so len(itl)+1 != output_len.
|
||||
# Using mean(itl) keeps TPOT and ITL on the same footing.
|
||||
tpot = sum(outputs[i].itl) / len(outputs[i].itl)
|
||||
elif getattr(outputs[i], "tpot_measured", True):
|
||||
else:
|
||||
try:
|
||||
latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft
|
||||
except Exception:
|
||||
latency_minus_ttft = outputs[i].latency - outputs[i].ttft
|
||||
tpot = latency_minus_ttft / (output_len - 1)
|
||||
else:
|
||||
tpot = None
|
||||
if tpot is not None:
|
||||
tpots.append(tpot)
|
||||
tpots.append(tpot)
|
||||
# Note: if output_len <= 1, we regard tpot as 0 for goodput
|
||||
all_tpots.append(tpot)
|
||||
itls += outputs[i].itl
|
||||
session_metrics = getattr(outputs[i], "duplex_session_metrics", None)
|
||||
ttft_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_ttft_ms") is not None
|
||||
audio_ttfp_measured = (
|
||||
not isinstance(session_metrics, dict) or session_metrics.get("mean_ttfp_ms") is not None
|
||||
)
|
||||
audio_rtf_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_rtf") is not None
|
||||
if ttft_measured:
|
||||
ttfts.append(outputs[i].ttft)
|
||||
if audio_ttfp_measured:
|
||||
audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
|
||||
if audio_rtf_measured:
|
||||
audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
|
||||
goodput_ttfts.append(outputs[i].ttft if ttft_measured else None)
|
||||
goodput_audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0) if audio_ttfp_measured else None)
|
||||
ttfts.append(outputs[i].ttft)
|
||||
audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
|
||||
audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
|
||||
audio_duration.append(getattr(outputs[i], defs.AUDIO_DURATION, 0.0))
|
||||
audio_frames.append(getattr(outputs[i], defs.AUDIO_FRAMES, 0.0))
|
||||
image_count = int(getattr(outputs[i], defs.IMAGE_COUNT, 0) or 0)
|
||||
@@ -828,14 +789,14 @@ def calculate_metrics(
|
||||
actual_output_lens.append(0)
|
||||
|
||||
if goodput_config_dict:
|
||||
valid_metrics: list[Sequence[float | None]] = []
|
||||
valid_metrics = []
|
||||
slo_values = []
|
||||
|
||||
if "ttft" in goodput_config_dict:
|
||||
valid_metrics.append(goodput_ttfts)
|
||||
valid_metrics.append(ttfts)
|
||||
slo_values.append(goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
|
||||
if "audio_ttft" in goodput_config_dict:
|
||||
valid_metrics.append(goodput_audio_ttfps)
|
||||
valid_metrics.append(audio_ttfps)
|
||||
slo_values.append(goodput_config_dict["audio_ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
|
||||
if "tpot" in goodput_config_dict:
|
||||
valid_metrics.append(all_tpots)
|
||||
@@ -845,7 +806,7 @@ def calculate_metrics(
|
||||
slo_values.append(goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION)
|
||||
|
||||
for req_metric in zip(*valid_metrics):
|
||||
is_good_req = all(r is not None and s >= r for s, r in zip(slo_values, req_metric))
|
||||
is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)])
|
||||
if is_good_req:
|
||||
good_completed += 1
|
||||
|
||||
@@ -861,7 +822,6 @@ def calculate_metrics(
|
||||
# Calculate max output tokens per second metric
|
||||
max_output_tokens_per_s = 0.0
|
||||
max_concurrent_requests = 0
|
||||
token_timeline_available = True
|
||||
|
||||
# Find the time range across all successful requests
|
||||
successful_outputs = [output for output in outputs if output.success]
|
||||
@@ -876,26 +836,19 @@ def calculate_metrics(
|
||||
concurrent_requests_per_second = np.zeros(duration_seconds)
|
||||
|
||||
for i, output in enumerate(successful_outputs):
|
||||
session_metrics = getattr(output, "duplex_session_metrics", None)
|
||||
if isinstance(session_metrics, dict):
|
||||
# Duplex TTFT/ITL samples are response-local. Without each
|
||||
# response.created offset they cannot form a session-global
|
||||
# token timeline, so do not publish a misleading peak rate.
|
||||
token_timeline_available = False
|
||||
else:
|
||||
# Calculate token generation timestamp using
|
||||
# start_time, ttft, and itl
|
||||
token_times = [output.start_time + output.ttft]
|
||||
current_time = token_times[0]
|
||||
for itl_value in output.itl:
|
||||
current_time += itl_value
|
||||
token_times.append(current_time)
|
||||
# Calculate token generation timestamp using
|
||||
# start_time, ttft, and itl
|
||||
token_times = [output.start_time + output.ttft]
|
||||
current_time = token_times[0]
|
||||
for itl_value in output.itl:
|
||||
current_time += itl_value
|
||||
token_times.append(current_time)
|
||||
|
||||
# Add tokens to second buckets
|
||||
for token_time in token_times:
|
||||
second_bucket = int(token_time - min_start_time)
|
||||
if 0 <= second_bucket < duration_seconds:
|
||||
tokens_per_second[second_bucket] += 1
|
||||
# Add tokens to second buckets
|
||||
for token_time in token_times:
|
||||
second_bucket = int(token_time - min_start_time)
|
||||
if 0 <= second_bucket < duration_seconds:
|
||||
tokens_per_second[second_bucket] += 1
|
||||
|
||||
# Track concurrent requests for each second this request was active
|
||||
request_start_second = int(output.start_time - min_start_time)
|
||||
@@ -907,22 +860,18 @@ def calculate_metrics(
|
||||
# Find the maximum tokens per second and corresponding
|
||||
# concurrent requests
|
||||
if len(tokens_per_second) > 0:
|
||||
if token_timeline_available:
|
||||
max_output_tokens_per_s = float(np.max(tokens_per_second))
|
||||
else:
|
||||
max_output_tokens_per_s = float("nan")
|
||||
max_output_tokens_per_s = float(np.max(tokens_per_second))
|
||||
max_concurrent_requests = int(np.max(concurrent_requests_per_second))
|
||||
|
||||
if TERM_PLOTLIB_AVAILABLE:
|
||||
import termplotlib as tpl
|
||||
|
||||
fig = tpl.figure()
|
||||
if token_timeline_available:
|
||||
fig.plot(
|
||||
np.arange(len(tokens_per_second)),
|
||||
tokens_per_second,
|
||||
title="Output tokens per second",
|
||||
)
|
||||
fig.plot(
|
||||
np.arange(len(tokens_per_second)),
|
||||
tokens_per_second,
|
||||
title="Output tokens per second",
|
||||
)
|
||||
fig.plot(
|
||||
np.arange(len(concurrent_requests_per_second)),
|
||||
concurrent_requests_per_second,
|
||||
@@ -932,10 +881,6 @@ def calculate_metrics(
|
||||
else:
|
||||
print("tip: install termplotlib and gnuplot to plot the metrics")
|
||||
|
||||
duplex_metrics_present = any(
|
||||
isinstance(getattr(output, "duplex_session_metrics", None), dict) for output in outputs
|
||||
)
|
||||
missing_duplex_value = float("nan") if duplex_metrics_present else 0
|
||||
metrics = MultiModalsBenchmarkMetrics(
|
||||
completed=completed,
|
||||
failed=len(failed_outputs),
|
||||
@@ -945,18 +890,18 @@ def calculate_metrics(
|
||||
request_goodput=good_completed / dur_s,
|
||||
output_throughput=sum(actual_output_lens) / dur_s,
|
||||
total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,
|
||||
mean_ttft_ms=np.mean(ttfts or missing_duplex_value) * 1000,
|
||||
std_ttft_ms=np.std(ttfts or missing_duplex_value) * 1000,
|
||||
median_ttft_ms=np.median(ttfts or missing_duplex_value) * 1000,
|
||||
percentiles_ttft_ms=[(p, np.percentile(ttfts or missing_duplex_value, p) * 1000) for p in selected_percentiles],
|
||||
mean_tpot_ms=np.mean(tpots or missing_duplex_value) * 1000,
|
||||
std_tpot_ms=np.std(tpots or missing_duplex_value) * 1000,
|
||||
median_tpot_ms=np.median(tpots or missing_duplex_value) * 1000,
|
||||
percentiles_tpot_ms=[(p, np.percentile(tpots or missing_duplex_value, p) * 1000) for p in selected_percentiles],
|
||||
mean_itl_ms=np.mean(itls or missing_duplex_value) * 1000,
|
||||
std_itl_ms=np.std(itls or missing_duplex_value) * 1000,
|
||||
median_itl_ms=np.median(itls or missing_duplex_value) * 1000,
|
||||
percentiles_itl_ms=[(p, np.percentile(itls or missing_duplex_value, p) * 1000) for p in selected_percentiles],
|
||||
mean_ttft_ms=np.mean(ttfts or 0) * 1000, # ttfts is empty if streaming is not supported by the endpoint
|
||||
std_ttft_ms=np.std(ttfts or 0) * 1000,
|
||||
median_ttft_ms=np.median(ttfts or 0) * 1000,
|
||||
percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles],
|
||||
mean_tpot_ms=np.mean(tpots or 0) * 1000,
|
||||
std_tpot_ms=np.std(tpots or 0) * 1000,
|
||||
median_tpot_ms=np.median(tpots or 0) * 1000,
|
||||
percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles],
|
||||
mean_itl_ms=np.mean(itls or 0) * 1000,
|
||||
std_itl_ms=np.std(itls or 0) * 1000,
|
||||
median_itl_ms=np.median(itls or 0) * 1000,
|
||||
percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles],
|
||||
mean_e2el_ms=np.mean(e2els or 0) * 1000,
|
||||
std_e2el_ms=np.std(e2els or 0) * 1000,
|
||||
median_e2el_ms=np.median(e2els or 0) * 1000,
|
||||
@@ -965,16 +910,11 @@ def calculate_metrics(
|
||||
max_concurrent_requests=max_concurrent_requests,
|
||||
rtfx=input_audio_duration / dur_s,
|
||||
**{
|
||||
"num_ttft_samples": len(ttfts),
|
||||
"num_tpot_samples": len(tpots),
|
||||
"num_itl_samples": len(itls),
|
||||
"num_audio_ttfp_samples": len(audio_ttfps),
|
||||
"num_audio_rtf_samples": len(audio_rtfs),
|
||||
defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or missing_duplex_value) * 1000,
|
||||
defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or missing_duplex_value) * 1000,
|
||||
defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or missing_duplex_value) * 1000,
|
||||
defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or 0) * 1000,
|
||||
defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or 0) * 1000,
|
||||
defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or 0) * 1000,
|
||||
defs.PERCENTILES_AUDIO_TTFP_MS: [
|
||||
(p, np.percentile(audio_ttfps or missing_duplex_value, p) * 1000) for p in selected_percentiles
|
||||
(p, np.percentile(audio_ttfps or 0, p) * 1000) for p in selected_percentiles
|
||||
],
|
||||
defs.MEAN_AUDIO_DURATION_S: np.mean(audio_duration or 0),
|
||||
defs.STD_AUDIO_DURATION_S: np.std(audio_duration or 0),
|
||||
@@ -985,12 +925,10 @@ def calculate_metrics(
|
||||
defs.TOTAL_AUDIO_DURATION_S: sum(audio_duration),
|
||||
defs.TOTAL_AUDIO_FRAMES: sum(audio_frames),
|
||||
defs.AUDIO_THROUGHPUT: sum(audio_duration) / dur_s,
|
||||
defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or missing_duplex_value),
|
||||
defs.STD_AUDIO_RTF: np.std(audio_rtfs or missing_duplex_value),
|
||||
defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or missing_duplex_value),
|
||||
defs.PERCENTILES_AUDIO_RTF: [
|
||||
(p, np.percentile(audio_rtfs or missing_duplex_value, p)) for p in selected_percentiles
|
||||
],
|
||||
defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or 0),
|
||||
defs.STD_AUDIO_RTF: np.std(audio_rtfs or 0),
|
||||
defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or 0),
|
||||
defs.PERCENTILES_AUDIO_RTF: [(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles],
|
||||
defs.TOTAL_IMAGES: total_images,
|
||||
defs.IMAGE_THROUGHPUT: total_images / dur_s,
|
||||
defs.AVERAGE_PIXELS_PER_IMAGE: (total_image_pixels / total_images) if total_images > 0 else 0.0,
|
||||
|
||||
Loaded 100 of 186 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user