Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
49fa488951 | ||
|
|
3b84cc4a7f | ||
|
|
f9253ae1ad | ||
|
|
5eb51516db | ||
|
|
52ec7e4d61 | ||
|
|
00b2eb9399 | ||
|
|
6fe1ff7709 | ||
|
|
cb439f3e84 | ||
|
|
2aac8bd15a | ||
|
|
8c3d3407d1 | ||
|
|
224152566a | ||
|
|
4200768f59 | ||
|
|
a769d9357b | ||
|
|
7a2a16f8af |
No files matched your search
@@ -1,5 +1,67 @@
|
||||
# vllm-omni serve
|
||||
|
||||
## Multiple API frontends
|
||||
|
||||
For local EngineCore pipelines with one or more stages, `--api-server-count` starts
|
||||
multiple API frontend processes that share one parent-owned set of stage
|
||||
engines:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--api-server-count 2 \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
This mode is intended for workloads whose request parsing, multimodal input
|
||||
processing, or response encoding can bottleneck a single frontend. It
|
||||
currently requires local multi-process EngineCore stages with one local
|
||||
process group per replica. A stage may still use multiple `num_replicas`.
|
||||
It cannot be combined with diffusion, headless or
|
||||
remote stages, intra-stage data parallelism, Ray, fault tolerance, elastic
|
||||
expert parallelism, sleep mode, or runtime LoRA updating. Runtime voice upload
|
||||
and deletion are also disabled because those mutations are process-local;
|
||||
built-in voices, inline reference audio, and voices restored at startup remain
|
||||
available. The `/v1/omni/sleep` and `/v1/omni/wakeup` control routes return HTTP
|
||||
409 for the same reason: their bookkeeping is process-local while stage engines
|
||||
are shared.
|
||||
|
||||
A single-stage model pipeline is supported through the normal launch command;
|
||||
this is distinct from the distributed `--stage-id` / `--headless` launch mode.
|
||||
Every stage in the pipeline must be an EngineCore stage. Mixed AR-to-DiT
|
||||
pipelines, including MiniMax-H3, remain unsupported even when their first
|
||||
stage is an EngineCore stage. Diffusion support requires per-frontend channels
|
||||
and shared-engine lifecycle handling in the diffusion backend.
|
||||
|
||||
### Workloads and performance validation
|
||||
|
||||
Each API process owns its frontend processor, Orchestrator, request state,
|
||||
and per-replica EngineCore channels. Requests return through the channel of
|
||||
the originating API process; model weights and GPU stage engines are shared.
|
||||
This distributes frontend and orchestration work across CPU processes.
|
||||
|
||||
| Workload | Potential benefit | Limitation |
|
||||
| ---------- | ------------------- | ------------ |
|
||||
| Concurrent cold image/video inputs on supported EngineCore pipelines | Parallel frontend processing when one API process is CPU-limited | A single cold request and GPU computation may be unchanged |
|
||||
| Many replicas producing frequent text/audio outputs | Less API/Orchestrator work and GIL contention per process | Routing remains serial within each Orchestrator |
|
||||
| Warm inputs or low-concurrency, GPU-limited generation | Additional frontend capacity | More API processes may add overhead without improving throughput |
|
||||
|
||||
Here, a cold input means media or preprocessing caches miss, not that model
|
||||
weights must be loaded or GPU kernels compiled. Each frontend has its own
|
||||
processor/cache state, so cache duplication, CPU thread contention, and memory
|
||||
usage also matter. Multi-API serving does not itself change the output queue,
|
||||
move response encoding to an executor, or change the model's numerical scheduler.
|
||||
|
||||
Compare `--api-server-count 1`, `2`, and higher counts with the same hardware,
|
||||
model/stage configuration, media, sampling parameters, output lengths, and
|
||||
offered load. Keep cold-cache and warmed-cache runs separate and report repeated
|
||||
measurements. Record CPU and request counts for every API worker, input and
|
||||
output queue delay, GPU utilization, first-output latency, and end-to-end latency.
|
||||
Confirm requests reach every worker: connection reuse can concentrate traffic
|
||||
on one process. Queue-only or single-API experiments do not establish a
|
||||
multi-API throughput gain. The multi-replica workload in
|
||||
[#4680](https://github.com/vllm-project/vllm-omni/issues/4680) remains a relevant
|
||||
validation target.
|
||||
|
||||
## Stage-based CLI quickstart
|
||||
|
||||
The stage-based CLI is designed for deployments that require launching each pipeline stage in an isolated process
|
||||
@@ -61,3 +123,10 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
## Arguments
|
||||
|
||||
--8<-- "docs/generated/argparse_omni/omni_serve.inc.md"
|
||||
|
||||
Async video jobs use process-local metadata and task handles. With multiple API
|
||||
workers, creating, listing, retrieving, downloading, and deleting these jobs
|
||||
returns HTTP 409. This restriction is independent of the current diffusion
|
||||
stage launch restriction: supporting multi-worker diffusion also requires a
|
||||
shared job lifecycle. Missing frontend topology state returns HTTP 503 for
|
||||
operations that require process-local state protection.
|
||||
@@ -271,3 +271,111 @@ The required evidence is:
|
||||
|
||||
Audio that is empty, all zero, non-finite, or only a protocol `listen` event is
|
||||
not a successful end-to-end result.
|
||||
|
||||
### Paced multi-session load driver
|
||||
|
||||
The existing driver keeps its two-session admission and slot-reuse validation
|
||||
when `--sessions` is omitted. Passing `--sessions N` selects a separate,
|
||||
bounded load mode; it does not replace the lifecycle test or change server
|
||||
admission capacity. Start an existing PersonaPlex server configured to admit
|
||||
the requested number of sessions, then run from a source checkout with the
|
||||
vLLM-Omni client dependencies installed:
|
||||
|
||||
```bash
|
||||
python -m tests.e2e.online_serving.personaplex_realtime_duplex \
|
||||
--model nvidia/personaplex-7b-v1 \
|
||||
--input-wav /path/to/input.wav \
|
||||
--sessions 2 --load-frames 1000 --drain-s 2 \
|
||||
--server-revision SERVER_COMMIT --server-hardware A100-80GB \
|
||||
--output-dir /tmp/personaplex-load-n2
|
||||
```
|
||||
|
||||
Use a new output directory for every run. The driver refuses to overwrite an
|
||||
existing `load-result.json`. The server labels are supplied by the operator,
|
||||
not detected or verified by the client. Record the server command, package
|
||||
versions, checkpoint revision and hardware alongside a published measurement.
|
||||
|
||||
All attempted sessions finish admission before admitted sessions share a
|
||||
common input start. A rejected session remains a failed row in the report;
|
||||
it is not removed from the denominator. Each session receives at most
|
||||
`--load-frames` input frames, without looping the WAV, and uses 80 ms pacing.
|
||||
A stalled send cannot cause a catch-up burst. The timeline records send
|
||||
schedule lateness so an overloaded client is visible rather than mistaken
|
||||
for a faster server. This mode measures steady concurrency, not join/leave
|
||||
interference; the original lifecycle mode still checks slot reuse.
|
||||
Load mode does not probe an additional N+1 session for `resource_exhausted`;
|
||||
the lifecycle mode owns that admission-limit check. A per-session
|
||||
`client_pacing_warning` and the summary's `client_pacing_warning_sessions`
|
||||
identify send-lateness p99 above one nominal frame (80 ms). This is a warning,
|
||||
not an extra acceptance gate: read it alongside server-facing measurements
|
||||
so a slow client is not mistaken for a slow server.
|
||||
|
||||
The fixed drain window is part of the measurement configuration. Audio
|
||||
arriving during cleanup is excluded, so closing cannot rescue a session that
|
||||
missed its observation deadline. Cleanup is attempted for admitted sessions
|
||||
even when audio validation fails. Protocol errors, missing/invalid audio,
|
||||
shared response IDs and cleanup failures make the report fail and the CLI
|
||||
exit nonzero. Cancellation closes transports and propagates to the caller.
|
||||
`--cleanup-timeout-s` (default 5 s) bounds session-close send and acknowledgement
|
||||
together, independently of `--timeout-s`. The subsequent WebSocket close
|
||||
handshake has the same separate bound; an unresponsive close acknowledgement
|
||||
cannot hold cancellation for the full inference timeout.
|
||||
|
||||
`load-result.json` retains per-session send and audio-packet timelines using
|
||||
the client's monotonic clock. It does not export handshake credentials or
|
||||
raw audio payloads. Metric definitions are explicit:
|
||||
|
||||
Both modes validate each nonempty audio packet's strict base64 encoding, PCM16
|
||||
sample alignment, 24 kHz sample rate and response identity. Valid aggregate
|
||||
audio cannot hide a malformed packet. In load mode, rejected packets make the session fail
|
||||
and appear in `invalid_audio_packets` with their event index, receive time and
|
||||
fixed diagnostic code; they do not contribute samples or timing intervals.
|
||||
Empty audio deltas remain legal and contribute no samples. This check does
|
||||
not change the public client's decoding. The lifecycle scenario and dispatch
|
||||
remain the same, but malformed packets that aggregate checks once tolerated
|
||||
now fail that mode as well.
|
||||
|
||||
The probe accepts both the current OpenAI Realtime audio event name
|
||||
`response.output_audio.delta` and the legacy `response.audio.delta`. Validation
|
||||
uses the wire payload directly rather than relying on a particular client
|
||||
library's alias table, so a server-side event-name migration cannot silently
|
||||
turn real audio into a zero-output measurement.
|
||||
|
||||
| Field | Meaning |
|
||||
| --- | --- |
|
||||
| `client_audio_packet_interval_ms` | Consecutive nonempty packet arrivals: median, linear p99, max, and zero-based index of the first maximal interval. Interval index 0 ends at packet 1. One packet can contain multiple codec frames; this is not server tick latency. |
|
||||
| `client_first_audio_after_stream_start_ms` | First audio receipt minus first input send start; admission is excluded. |
|
||||
| `client_stream_rtf` | First input send start to last audio receipt, divided by received audio duration. This includes real-time input pacing and is not an inference-only RTF. |
|
||||
| `client_send_lateness_ms` | Send start minus its original shared schedule, clipped at zero. |
|
||||
| `frame_deficit` | Integer input-frame minus output-frame count from the shared session validation; missing when valid whole-frame accounting is unavailable. |
|
||||
| `output_frames`, `voiced_frames`, `silent_frames`, `frame_coverage_ratio`, `audio_rms`, `audio_chunks` | Session-validation statistics retained on success and on acceptance failure once computed. |
|
||||
| `acceptance_check` | Failed local check, with numeric diagnostics in `error`; other exceptions retain phase/type only, without remote handshake or metadata text. |
|
||||
| `audio_underrun_s`, `audio_underrun_event_count`, `audio_continuity_ok` | Shared `compute_continuity_stats` on valid PCM16 packet arrivals, with its default 100 ms underrun threshold; absent audio gives missing values rather than a continuity pass. |
|
||||
|
||||
No output gives missing latency/RTF values, not zero. By default RTF is
|
||||
diagnostic; `--max-client-rtf` adds an explicit ceiling. The existing frame
|
||||
coverage, audible-frame and scheduler-data-plane checks still apply. A longer
|
||||
`--drain-s` can demonstrate eventual completion under backlog, but it must not
|
||||
be reported as real-time capacity when the fixed-window run failed or client
|
||||
RTF exceeded the stated target. Likewise, the first live session can pay
|
||||
one-time model/codec/JIT costs; either use the server's configured duplex
|
||||
warmup path or label that run as cold-start evidence before comparing steady
|
||||
capacity. Passing client-only metrics does not certify model quality or a
|
||||
real-time capacity target on another GPU. Begin with N=1 and N=2; higher N is a
|
||||
capacity experiment, not a promised property of an A100 or of this patch.
|
||||
Continuity is a client-arrival diagnostic, not proof of speaker playback and
|
||||
not an additional pass/fail gate. Packet-interval and send-lateness summaries
|
||||
retain their documented linear p99; the older client interval helper reports
|
||||
rounded nearest-rank p95 instead and is not substituted for that schema.
|
||||
|
||||
Deterministic driver and localhost WebSocket checks need no model or GPU:
|
||||
|
||||
```bash
|
||||
python -m pytest tests/e2e/online_serving/test_personaplex_load_driver.py \
|
||||
-m 'core_model and cpu' --run-level core_model -q
|
||||
```
|
||||
|
||||
These checks test the driver against controlled protocol responses. Run the
|
||||
load-driver command against real weights separately before reporting model-serving
|
||||
performance. Server tick timestamps and generic TTFT/TPOT aggregation remain
|
||||
the responsibility of the existing duplex metrics paths, not this driver.
|
||||
@@ -94,6 +94,13 @@ This document owns offline API semantics, CLI and serve composition,
|
||||
supported OpenAI-compatible routes, request validation and normalization,
|
||||
response conversion, streaming/session behavior, and engine handoff.
|
||||
|
||||
When `--api-server-count` is greater than one, the serve composition root
|
||||
starts the frontend process manager while a parent-owned `StageRuntime`
|
||||
launches the shared local stage engines. Each frontend receives only its own
|
||||
stage-client channel configuration and does not launch or retire backend
|
||||
processes. CLI validation rejects unsupported distributed, diffusion,
|
||||
fault-tolerant, elastic-EP, Ray, and runtime-LoRA combinations before startup.
|
||||
|
||||
It does not own configuration precedence, cross-stage routing, stage
|
||||
placement, payload implementation, or semantic error classification.
|
||||
`entrypoints/openai/errors.py` is an explicit primary-path exception owned by
|
||||
@@ -141,6 +148,9 @@ stage so `wake_up(stage_ids=[0])` does not skip a later `wake_up(stage_ids=[1])`
|
||||
Streaming input pumps take an admission slot immediately before each EngineCore
|
||||
ADD or update, not while waiting for the next client chunk. Frontend abort
|
||||
keeps `request_states` until `generate()` consumes the terminal output.
|
||||
Multi-API serving must preserve one client rank and one channel set per
|
||||
frontend, include every frontend in readiness/failure observation, and keep
|
||||
stage lifecycle ownership in the parent composition root.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
|
||||
@@ -77,6 +77,22 @@ readiness, one pool per logical stage, replica identity and membership,
|
||||
selection and affinity, stage-client/process lifecycle, liveness, draining,
|
||||
and shutdown.
|
||||
|
||||
Local EngineCore initialization and multi-API serving use
|
||||
`StageRuntime.launch_stage_engines(num_api_servers=1)` as their common launch
|
||||
entry. It resolves plans when not supplied by the runtime's device-group
|
||||
scheduler, coordinates device locks and environment overlays, starts engines,
|
||||
and rolls back resources on startup or attachment failure. Single-API clients
|
||||
attach locally after readiness and take ownership of the engine resources;
|
||||
there is no additional API subprocess. Diffusion and remote attachment retain
|
||||
their backend-specific initialization.
|
||||
|
||||
For local multi-API serving, the parent runtime owns and launches each stage
|
||||
engine once. It supplies a distinct input/output channel for every API
|
||||
frontend and stage replica; frontend workers attach clients to those engines
|
||||
without acquiring stage-process ownership. Readiness observation and shutdown
|
||||
remain parent-owned. This topology currently excludes diffusion and
|
||||
remote/headless stages, intra-stage data parallelism, and Ray backends.
|
||||
|
||||
It does not own cross-stage request policy, model scheduler policy, connector
|
||||
transfer semantics, configuration precedence, or public error rendering.
|
||||
|
||||
@@ -114,6 +130,10 @@ affinity, membership changes, abort routing, process failure, draining, and
|
||||
repeated shutdown. Cross-stage policy changes belong in
|
||||
`engine_orchestration.md`.
|
||||
|
||||
For multi-API changes, additionally exercise per-frontend channel assignment,
|
||||
shared-engine startup handshakes, frontend failure observation, and parent
|
||||
cleanup without duplicate stage shutdown.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
- Reconcile terminology and ownership after #5441 reaches a final state.
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni>.
|
||||
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
|
||||
@@ -15,6 +14,19 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
For a local deployment with multiple API frontend processes sharing one set
|
||||
of stage engines, add `--api-server-count`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--api-server-count 2 --port 8091
|
||||
```
|
||||
|
||||
This mode supports local EngineCore stages with one local process group per
|
||||
replica, including stages configured with multiple `num_replicas`. Headless,
|
||||
remote, diffusion, intra-replica data-parallel, and Ray stage deployments
|
||||
remain single-frontend.
|
||||
|
||||
The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
|
||||
automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
|
||||
The bundled Qwen3-Omni setup explicitly sets `moe_backend: triton` for the Thinker and Talker stages. This keeps
|
||||
@@ -27,6 +39,7 @@ on the Thinker stage. Keep Talker and Code2Wav at `pipeline_parallel_size: 1`. T
|
||||
default asynchronous scheduling and audio chunk streaming can remain enabled.
|
||||
|
||||
To explicitly utilize a custom deployment YAML, specify the configuration path:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config /path/to/deploy_config_file
|
||||
@@ -50,7 +63,7 @@ Adopt the stage-based CLI architecture to independently instantiate execution pr
|
||||
The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via
|
||||
`CUDA_VISIBLE_DEVICES`.
|
||||
|
||||
**1. Stage 0 (Thinker + API server)**
|
||||
#### 1. Stage 0 (Thinker + API server)
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
@@ -60,7 +73,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
**2. Stage 1 (Talker)**
|
||||
#### 2. Stage 1 (Talker)
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
@@ -70,7 +83,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
**3. Stage 2 (Code2Wav)**
|
||||
#### 3. Stage 2 (Code2Wav)
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
@@ -95,6 +108,9 @@ To experiment with another supported MoE backend, set `moe_backend` in the
|
||||
deployment YAML for both the Thinker and Talker stages. For the stage-based
|
||||
CLI, pass `--moe-backend` directly on each stage command.
|
||||
|
||||
For other stage settings, pass the option on that stage command. For example,
|
||||
to change its GPU memory utilization:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
@@ -107,11 +123,12 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
### Send Multi-modal Request
|
||||
|
||||
Get into the example folder
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/qwen3_omni
|
||||
```
|
||||
|
||||
#### Send request via python
|
||||
#### Send request via python
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image --port 8091 --host "localhost"
|
||||
@@ -126,7 +143,6 @@ The Python client supports the following command-line arguments:
|
||||
- `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type is `use_audio`, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
|
||||
- `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"`
|
||||
|
||||
|
||||
For example, to use a local video file with custom prompt:
|
||||
|
||||
```bash
|
||||
@@ -136,19 +152,20 @@ python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--prompt "What are the main activities shown in this video?"
|
||||
```
|
||||
|
||||
#### Send request via curl
|
||||
#### Send request via curl
|
||||
|
||||
```bash
|
||||
bash run_curl_multimodal_generation.sh use_image
|
||||
```
|
||||
|
||||
## Modality control
|
||||
|
||||
You can control output modalities to specify which types of output the model should generate. This is useful when you only need text output and want to skip audio generation stages for better performance.
|
||||
|
||||
### Supported modalities
|
||||
|
||||
| Modalities | Output |
|
||||
|------------|--------|
|
||||
| ------------ | -------- |
|
||||
| `["text"]` | Text only |
|
||||
| `["audio"]` | Audio only |
|
||||
| `["text", "audio"]` | Text + Audio |
|
||||
@@ -226,7 +243,9 @@ print(response.choices[1].message.audio) # Audio response
|
||||
```
|
||||
|
||||
## Streaming Output
|
||||
|
||||
If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally.
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type use_image \
|
||||
@@ -259,12 +278,14 @@ The convenience script launches both the vLLM server and Gradio demo together:
|
||||
```
|
||||
|
||||
This script will:
|
||||
|
||||
1. Start the vLLM server in the background
|
||||
2. Wait for the server to be ready
|
||||
3. Launch the Gradio demo
|
||||
4. Handle cleanup when you press Ctrl+C
|
||||
|
||||
The script supports the following arguments:
|
||||
|
||||
- `--model`: Model name/path (default: Qwen/Qwen3-Omni-30B-A3B-Instruct)
|
||||
- `--server-port`: Port for vLLM server (default: 8091)
|
||||
- `--gradio-port`: Port for Gradio demo (default: 7861)
|
||||
@@ -275,18 +296,19 @@ The script supports the following arguments:
|
||||
|
||||
#### Option 2: Manual Launch (Two-Step Process)
|
||||
|
||||
**Step 1: Launch the vLLM API server**
|
||||
##### Step 1: Launch the vLLM API server
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
If you have a custom deploy config file:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file
|
||||
```
|
||||
|
||||
**Step 2: Run the Gradio demo**
|
||||
##### Step 2: Run the Gradio demo
|
||||
|
||||
In a separate terminal:
|
||||
|
||||
@@ -299,7 +321,7 @@ Then open `http://localhost:7861/` on your local browser to interact with the we
|
||||
The gradio script supports the following arguments:
|
||||
|
||||
- `--model`: Model name/path (should match the server model)
|
||||
- `--api-base`: Base URL for the vLLM API server (default: http://localhost:8091/v1)
|
||||
- `--api-base`: Base URL for the vLLM API server (default: <http://localhost:8091/v1>)
|
||||
- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
|
||||
- `--port`: Port for Gradio server (default: 7861)
|
||||
- `--share`: Share the Gradio demo publicly (creates a public link)
|
||||
|
||||
@@ -413,6 +413,22 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
|
||||
./qwen3_tts/run_server.sh Base
|
||||
```
|
||||
|
||||
For a local deployment with multiple API frontend processes sharing one set
|
||||
of TTS stage engines, add `--api-server-count`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni --api-server-count 2 --port 8091
|
||||
```
|
||||
|
||||
This mode supports local EngineCore stages with one local process group per
|
||||
replica, including stages configured with multiple `num_replicas`. Headless,
|
||||
remote, intra-replica data-parallel, and Ray stage deployments are not
|
||||
supported. Runtime voice upload and deletion are disabled with multiple
|
||||
API frontends because their registries are process-local. Built-in voices,
|
||||
inline `ref_audio`, and voices restored from `custom_voice_dir` at startup are
|
||||
supported.
|
||||
|
||||
### Executor backend
|
||||
|
||||
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
|
||||
@@ -478,6 +494,10 @@ checkpoint. When a request names an uploaded voice, the server infers
|
||||
`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
|
||||
CustomVoice speakers and require a CustomVoice checkpoint.
|
||||
|
||||
The runtime upload and delete routes require a single API frontend; with
|
||||
`--api-server-count > 1`, use inline `ref_audio` or restore precomputed voices
|
||||
from `custom_voice_dir` at startup.
|
||||
|
||||
### Precomputed custom voices
|
||||
|
||||
For reused Base voice-cloning speakers, precompute the reference artifacts once
|
||||
|
||||
@@ -14,6 +14,17 @@ Please refer to [README.md](../../../README.md)
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
For a local deployment with multiple API frontend processes sharing one set
|
||||
of stage engines, add `--api-server-count`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --api-server-count 2 --port 8091
|
||||
```
|
||||
|
||||
This mode currently supports local EngineCore stages. It is not available
|
||||
with headless or remote stage deployments, and pure diffusion stages remain
|
||||
single-frontend.
|
||||
|
||||
The default deployment configuration, situated at `vllm_omni/deploy/qwen3_omni_moe.yaml`, is resolved and loaded
|
||||
automatically via the model registry, obviating the `--deploy-config` flag in standard deployment topologies.
|
||||
Asynchronous chunk streaming operates as **enabled by default** within this bundled configuration.
|
||||
@@ -156,7 +167,7 @@ parser defaults). If you don't pass a flag, the YAML value wins.
|
||||
> chunked vs end-to-end modes (e.g. qwen3_tts code2wav) dispatch
|
||||
> automatically based on that bool — no extra flag or variant yaml is
|
||||
> needed.
|
||||
|
||||
>
|
||||
> ⚠️ **For multi-stage models that share GPUs (qwen3_omni_moe by default
|
||||
> shares cuda:1 between stages 1 and 2), avoid using global memory flags.**
|
||||
> A global `--gpu-memory-utilization 0.85` would apply to every stage and
|
||||
|
||||
@@ -13,7 +13,7 @@ For the full list of supported architectures across all modalities, see
|
||||
## Supported Models
|
||||
|
||||
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
|
||||
|---|---|---|---|---|---|
|
||||
| --- | --- | --- | --- | --- | --- |
|
||||
| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | speaker tags (`S0`..`S9`, default `S0`) | — |
|
||||
| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
@@ -781,6 +781,20 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
|
||||
./qwen3_tts/run_server.sh Base
|
||||
```
|
||||
|
||||
For a local deployment with multiple API frontend processes sharing one set
|
||||
of TTS stage engines, add `--api-server-count`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni --api-server-count 2 --port 8091
|
||||
```
|
||||
|
||||
This mode supports local EngineCore stages only; headless or remote stage
|
||||
deployments are not supported. Runtime voice upload and deletion are disabled
|
||||
with multiple API frontends because their registries are process-local.
|
||||
Built-in voices, inline `ref_audio`, and voices restored from `custom_voice_dir`
|
||||
at startup are supported.
|
||||
|
||||
### Executor backend
|
||||
|
||||
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
|
||||
@@ -835,6 +849,10 @@ checkpoint. When a request names an uploaded voice, the server infers
|
||||
`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
|
||||
CustomVoice speakers and require a CustomVoice checkpoint.
|
||||
|
||||
The runtime upload and delete routes require a single API frontend; with
|
||||
`--api-server-count > 1`, use inline `ref_audio` or restore precomputed voices
|
||||
from `custom_voice_dir` at startup.
|
||||
|
||||
### Precomputed custom voices
|
||||
|
||||
For reused Base voice-cloning speakers, precompute the reference artifacts once
|
||||
|
||||
@@ -1106,6 +1106,68 @@ generations move the end-to-end figure toward the denoising one. Fusing the
|
||||
adapter does not measurably change startup: weight loading took 77.3 s with it
|
||||
against 85.8 s without.
|
||||
|
||||
### FastH3 8-Step V2 full checkpoint
|
||||
|
||||
[FastH3 8-Step V2](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2)
|
||||
ships a full Diffusers-format transformer, including learned VSA compression
|
||||
gates. Serve the Hugging Face model ID or a downloaded snapshot directly:
|
||||
|
||||
```bash
|
||||
export FASTH3_V2_MODEL=FastVideo/FastVideo-FastH3-8-Step-V2
|
||||
```
|
||||
|
||||
Use a Hugging Face account with access to the release, and pin `--revision`
|
||||
when reproducing a result. The existing component loader reads its safetensors
|
||||
shards; the H3 weight loader maps names and loads Q/K/V and MLP projections into
|
||||
native parameters in memory. No conversion command or second transformer
|
||||
checkpoint is needed.
|
||||
|
||||
The transformer and text components come from the V2 release. To retain Omni's
|
||||
native tiled and parallel VAE execution, the loader fetches only the video/audio
|
||||
VAEs from the `MiniMaxAI/MiniMax-H3` revision pinned in V2's `provenance.json`.
|
||||
Those components use the normal Hugging Face cache. No separate base-model
|
||||
argument or full base-transformer download is required. Download time is a
|
||||
preparation cost, separate from warmup and request latency.
|
||||
|
||||
Install the optional `fastvideo-kernel` build required by the
|
||||
`FASTVIDEO_VSA` backend, including its H3 block-map entry point. A four-GPU
|
||||
configuration is:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve "${FASTH3_V2_MODEL}" --omni \
|
||||
--host 127.0.0.1 --port 8095 --trust-remote-code --task-type fl2va \
|
||||
--served-model-name FastH3-V2 \
|
||||
--num-gpus 4 --usp 4 --ring 1 \
|
||||
--vae-patch-parallel-size 4 --vae-parallel-mode tile --vae-use-tiling \
|
||||
--diffusion-attention-backend FASTVIDEO_VSA
|
||||
|
||||
curl --fail-with-body -sS http://127.0.0.1:8095/v1/videos/sync \
|
||||
-F 'model=FastH3-V2' \
|
||||
-F 'prompt=A red fox runs through fresh snow at dawn, with fast pawsteps, winter wind and distant birds.' \
|
||||
-F 'width=1344' -F 'height=768' -F 'aspect_ratio=16:9' -F 'fps=24' \
|
||||
-F 'num_inference_steps=8' -F 'guidance_scale=1' -F 'flow_shift=10' \
|
||||
-F 'seed=1101' \
|
||||
-F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \
|
||||
-o fasth3_v2.mp4
|
||||
```
|
||||
|
||||
V2 pins video/audio shifts **10/3**, guidance **1**, and the unshifted ladder
|
||||
`[0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0]`.
|
||||
Use **`num_inference_steps=8`** in Omni: these nine nodes bound eight model
|
||||
evaluations. FastVideo's `--steps 9` counts nodes instead. The existing H3 ODE
|
||||
update is reused without stochastic re-noising.
|
||||
|
||||
The trained attention policy is **80% sparsity with 64-token video tiles**.
|
||||
Omni derives the number of selected video tiles from each request's geometry;
|
||||
do not supply a fixed `--fastvideo-vsa-topk`. Text/condition/audio prefix keys
|
||||
remain available to every query, and prefix queries remain dense. Missing VSA
|
||||
geometry or a failed H3 kernel raises an error instead of serving the student
|
||||
with dense attention.
|
||||
|
||||
This release supports T2VA only, with local or pure Ulysses attention. Additional
|
||||
LoRA adapters, Ref2VA and FL2VA conditioning are unsupported. The legacy
|
||||
four-step adapter keeps its original sampling and fixed-top-k behavior.
|
||||
|
||||
## Key parameters
|
||||
|
||||
| Parameter | Recommended value | Notes |
|
||||
|
||||
@@ -0,0 +1,679 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Iterable
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint import (
|
||||
FASTH3_BASE_MODEL,
|
||||
FASTH3_V2_BASE_SCHEDULE,
|
||||
FASTH3_V2_MODEL_ID,
|
||||
FastH3CheckpointSpec,
|
||||
)
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||
MiniMaxH3DiTArchConfig,
|
||||
MiniMaxH3DiTModel,
|
||||
)
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
resolve_minimax_h3_diffusion_model_path,
|
||||
)
|
||||
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
|
||||
is_minimax_h3_modular,
|
||||
resolve_minimax_h3_model_root,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _fastvideo_metadata() -> dict[str, object]:
|
||||
"""Return the release-owned metadata read from fastvideo_inference.json."""
|
||||
return {
|
||||
"model_id": FASTH3_V2_MODEL_ID,
|
||||
"schema_version": "fasth3-inference-contract-v1",
|
||||
"dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125],
|
||||
"transformer_forwards": 8,
|
||||
"video_scheduler_shift": 10.0,
|
||||
"audio_scheduler_shift": 3.0,
|
||||
"guidance_scale": 1.0,
|
||||
"attention_backend": "VIDEO_SPARSE_ATTN_H3",
|
||||
"vsa_sparsity": 0.8,
|
||||
"vsa_tile_size": 64,
|
||||
"task": "t2av",
|
||||
}
|
||||
|
||||
|
||||
def _sampling(
|
||||
*,
|
||||
num_inference_steps: int | str | None = 8,
|
||||
guidance_scale: float | None = 1.0,
|
||||
extra_args: dict[str, object] | None = None,
|
||||
lora_request: LoRARequest | None = None,
|
||||
timesteps: torch.Tensor | None = None,
|
||||
sigmas: list[float] | None = None,
|
||||
) -> OmniDiffusionSamplingParams:
|
||||
return OmniDiffusionSamplingParams(
|
||||
# Keep the string form to exercise malformed client input validation.
|
||||
num_inference_steps=num_inference_steps, # type: ignore[arg-type]
|
||||
guidance_scale=guidance_scale,
|
||||
extra_args=extra_args or {},
|
||||
lora_request=lora_request,
|
||||
timesteps=timesteps,
|
||||
sigmas=sigmas,
|
||||
)
|
||||
|
||||
|
||||
def _od_config(
|
||||
*,
|
||||
backend: str = "FASTVIDEO_VSA",
|
||||
topk: int | None = None,
|
||||
per_role: dict[str, AttentionSpec] | None = None,
|
||||
lora_path: str | None = None,
|
||||
ring_degree: int = 1,
|
||||
allgather_degree: int = 1,
|
||||
) -> OmniDiffusionConfig:
|
||||
attention_config = AttentionConfig(
|
||||
default=AttentionSpec(backend=backend, fastvideo_vsa_topk=topk),
|
||||
per_role=per_role or {},
|
||||
)
|
||||
return OmniDiffusionConfig(
|
||||
diffusion_attention_config=attention_config,
|
||||
lora_path=lora_path,
|
||||
parallel_config=DiffusionParallelConfig(ring_degree=ring_degree, allgather_degree=allgather_degree),
|
||||
)
|
||||
|
||||
|
||||
def test_v2_release_has_nine_sigma_nodes_and_eight_intervals():
|
||||
assert FASTH3_V2_BASE_SCHEDULE.base_schedule == (
|
||||
0.999,
|
||||
0.874,
|
||||
0.749,
|
||||
0.624,
|
||||
0.5,
|
||||
0.375,
|
||||
0.25,
|
||||
0.125,
|
||||
0.0,
|
||||
)
|
||||
assert len(FASTH3_V2_BASE_SCHEDULE.base_schedule) == 9
|
||||
assert FASTH3_V2_BASE_SCHEDULE.num_inference_steps == 8
|
||||
assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(10.0) == pytest.approx(
|
||||
[
|
||||
0.9998999099189271,
|
||||
0.985788405143244,
|
||||
0.9675752486758817,
|
||||
0.94316807738815,
|
||||
0.9090909090909091,
|
||||
0.8571428571428571,
|
||||
0.7692307692307693,
|
||||
0.5882352941176471,
|
||||
0.0,
|
||||
],
|
||||
abs=1e-12,
|
||||
)
|
||||
assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(3.0) == pytest.approx(
|
||||
[
|
||||
0.9996664442961973,
|
||||
0.9541484716157204,
|
||||
0.8995196156925539,
|
||||
0.8327402135231315,
|
||||
0.75,
|
||||
0.6428571428571429,
|
||||
0.5,
|
||||
0.3,
|
||||
0.0,
|
||||
],
|
||||
abs=1e-12,
|
||||
)
|
||||
|
||||
|
||||
def test_fastvideo_metadata_builds_release_sampling_metadata():
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
|
||||
assert spec.vsa_sparsity == 0.8
|
||||
assert spec.release_metadata() == {
|
||||
"partition": "fl2va",
|
||||
"tasks": ["t2va"],
|
||||
"sigma_shift_scales": {"video": 10.0, "audio": 3.0},
|
||||
"base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
("model_id", "FastVideo/FastH3-other"),
|
||||
("schema_version", "fasth3-inference-contract-v0"),
|
||||
("dmd_denoising_steps", [999, 749, 500, 250]),
|
||||
("transformer_forwards", 4),
|
||||
("video_scheduler_shift", 12.0),
|
||||
("audio_scheduler_shift", 4.0),
|
||||
("guidance_scale", 2.0),
|
||||
("attention_backend", "FLASH_ATTN"),
|
||||
("vsa_sparsity", 0.9),
|
||||
("vsa_tile_size", 32),
|
||||
("task", "t2va"),
|
||||
],
|
||||
)
|
||||
def test_fastvideo_metadata_rejects_contract_drift(field, value):
|
||||
metadata = _fastvideo_metadata()
|
||||
metadata[field] = value
|
||||
|
||||
with pytest.raises(ValueError, match=field):
|
||||
FastH3CheckpointSpec.from_metadata(metadata)
|
||||
|
||||
|
||||
def test_fastvideo_metadata_rejects_missing_contract_field():
|
||||
metadata = _fastvideo_metadata()
|
||||
metadata.pop("task")
|
||||
|
||||
with pytest.raises(ValueError, match="task"):
|
||||
FastH3CheckpointSpec.from_metadata(metadata)
|
||||
|
||||
|
||||
def test_legacy_h3_metadata_is_not_treated_as_fastvideo_release():
|
||||
with pytest.raises(ValueError, match="model_id"):
|
||||
FastH3CheckpointSpec.from_metadata(
|
||||
{
|
||||
"schema_version": 1,
|
||||
"partition": "fl2va",
|
||||
"tasks": ["t2va", "fl2va"],
|
||||
"sigma_shift_scales": {"video": 12.0, "audio": 3.0},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_v2_serving_contract_accepts_backend_selected_for_h3_self_role():
|
||||
config = _od_config(
|
||||
backend="TORCH_SDPA",
|
||||
per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA")},
|
||||
)
|
||||
FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).check_serving_contract(
|
||||
partition="fl2va", od_config=config
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "partition", "match"),
|
||||
[
|
||||
({}, "ref2va", "task-type fl2va"),
|
||||
({"lora_path": "/tmp/adapter"}, "fl2va", "LoRA"),
|
||||
({"backend": "TORCH_SDPA"}, "fl2va", "FASTVIDEO_VSA"),
|
||||
({"topk": 2}, "fl2va", "fixed fastvideo_vsa_topk"),
|
||||
({"ring_degree": 2}, "fl2va", "local attention"),
|
||||
({"allgather_degree": 2}, "fl2va", "local attention"),
|
||||
],
|
||||
)
|
||||
def test_v2_serving_contract_rejects_incompatible_runtime(kwargs, partition, match):
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
with pytest.raises(ValueError, match=match):
|
||||
spec.check_serving_contract(partition=partition, od_config=_od_config(**kwargs))
|
||||
|
||||
|
||||
def test_v2_serving_contract_reads_a_per_role_fixed_topk_override():
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
config = _od_config(
|
||||
per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=1)},
|
||||
)
|
||||
with pytest.raises(ValueError, match="fixed fastvideo_vsa_topk"):
|
||||
spec.check_serving_contract(partition="fl2va", od_config=config)
|
||||
|
||||
|
||||
def test_v2_request_contract_accepts_omitted_steps_and_default_shifts():
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
spec.check_request(_sampling(num_inference_steps=None))
|
||||
spec.check_request(_sampling(extra_args={"flow_shift": 10, "audio_flow_shift": 3}, guidance_scale=1))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sampling",
|
||||
[
|
||||
_sampling(num_inference_steps=None),
|
||||
_sampling(timesteps=torch.tensor([0.9, 0.5, 0.0])),
|
||||
_sampling(sigmas=[1.0, 0.5, 0.0]),
|
||||
],
|
||||
)
|
||||
def test_v2_step_execution_requires_pinned_count_and_schedule(sampling):
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
with pytest.raises(OmniClientError, match="step execution"):
|
||||
spec.check_request(sampling, step_execution=True)
|
||||
|
||||
|
||||
def test_v2_step_execution_accepts_only_explicit_pinned_count():
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
spec.check_request(_sampling(num_inference_steps=8), step_execution=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "match"),
|
||||
[
|
||||
({"num_inference_steps": 7}, "num_inference_steps=8"),
|
||||
({"num_inference_steps": 9}, "num_inference_steps=8"),
|
||||
({"num_inference_steps": "8"}, "num_inference_steps=8"),
|
||||
({"extra_args": {"flow_shift": 9}}, "flow_shift=10"),
|
||||
({"extra_args": {"audio_flow_shift": 4}}, "audio_flow_shift=3"),
|
||||
({"extra_args": {"flow_shift": "bad"}}, "flow_shift=10"),
|
||||
({"guidance_scale": 2.0}, "guidance_scale=1"),
|
||||
(
|
||||
{"lora_request": LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/test.safetensors")},
|
||||
"per-request LoRA",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_v2_request_contract_rejects_sampling_drift(kwargs, match):
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
with pytest.raises(OmniClientError, match=match):
|
||||
spec.check_request(_sampling(**kwargs))
|
||||
|
||||
|
||||
def test_resolve_native_vaes_uses_pinned_revision_and_only_vae_patterns(tmp_path, monkeypatch):
|
||||
model_root = tmp_path / "FastH3"
|
||||
model_root.mkdir()
|
||||
revision = "base-revision-123"
|
||||
(model_root / "provenance.json").write_text(
|
||||
json.dumps({"base_model": f"hf://{FASTH3_BASE_MODEL}@{revision}"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
snapshot = tmp_path / "hf-snapshot"
|
||||
calls: dict[str, object] = {}
|
||||
|
||||
def fake_download(**kwargs: object) -> str:
|
||||
calls.update(kwargs)
|
||||
snapshot.mkdir()
|
||||
return str(snapshot)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint.download_weights_from_hf_specific",
|
||||
fake_download,
|
||||
)
|
||||
|
||||
result = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
|
||||
|
||||
assert result == snapshot / "FL2VA"
|
||||
assert calls == {
|
||||
"model_name_or_path": FASTH3_BASE_MODEL,
|
||||
"cache_dir": None,
|
||||
"allow_patterns": ["FL2VA/video_vae/**", "FL2VA/audio_vae/**"],
|
||||
"revision": revision,
|
||||
"require_all": True,
|
||||
}
|
||||
assert sorted(path.name for path in model_root.iterdir()) == ["provenance.json"]
|
||||
assert not (model_root / "model_index.json").exists()
|
||||
assert not (model_root / "conversion.json").exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"base_model",
|
||||
[
|
||||
"MiniMaxAI/MiniMax-H3@main",
|
||||
"hf://MiniMaxAI/MiniMax-H3@",
|
||||
"hf://Other/MiniMax-H3@revision",
|
||||
"",
|
||||
None,
|
||||
],
|
||||
)
|
||||
def test_resolve_native_vaes_rejects_unpinned_provenance(tmp_path, base_model):
|
||||
model_root = tmp_path / "FastH3"
|
||||
model_root.mkdir()
|
||||
(model_root / "provenance.json").write_text(json.dumps({"base_model": base_model}), encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="provenance"):
|
||||
FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
|
||||
|
||||
|
||||
def test_modular_local_model_root_uses_text_encoder_without_partition_files(tmp_path):
|
||||
root = tmp_path / "FastH3"
|
||||
(root / "text_encoder").mkdir(parents=True)
|
||||
(root / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
|
||||
(root / "modular_model_index.json").write_text(
|
||||
json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert is_minimax_h3_modular(str(root)) is True
|
||||
assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "text_encoder")
|
||||
assert not (root / "FL2VA").exists()
|
||||
|
||||
|
||||
def test_modular_hub_model_root_downloads_only_text_encoder(monkeypatch, tmp_path):
|
||||
from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
|
||||
|
||||
snapshot = tmp_path / "modular-snapshot"
|
||||
calls: dict[str, object] = {}
|
||||
|
||||
def fake_index(model: str, *, revision: str | None = None) -> dict[str, str]:
|
||||
assert model == "FastVideo/FastVideo-FastH3-8-Step-V2"
|
||||
assert revision == "release-revision"
|
||||
return {"_class_name": "MiniMaxH3ModularPipeline"}
|
||||
|
||||
def fake_download(**kwargs: object) -> str:
|
||||
calls.update(kwargs)
|
||||
snapshot.mkdir()
|
||||
return str(snapshot)
|
||||
|
||||
monkeypatch.setattr(checkpoint_module, "get_diffusion_model_index", fake_index)
|
||||
monkeypatch.setattr(checkpoint_module, "download_weights_from_hf_specific", fake_download)
|
||||
|
||||
result = resolve_minimax_h3_model_root(
|
||||
"FastVideo/FastVideo-FastH3-8-Step-V2",
|
||||
"release-revision",
|
||||
"fl2va",
|
||||
)
|
||||
|
||||
assert result == str(snapshot / "text_encoder")
|
||||
assert calls == {
|
||||
"model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
|
||||
"cache_dir": None,
|
||||
"allow_patterns": ["text_encoder/**"],
|
||||
"revision": "release-revision",
|
||||
"require_all": True,
|
||||
}
|
||||
|
||||
|
||||
def test_native_hub_model_with_modular_class_keeps_partition_layout(monkeypatch):
|
||||
from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
checkpoint_module,
|
||||
"get_diffusion_model_index",
|
||||
lambda _model, *, revision=None: {"_class_name": "MiniMaxH3ModularPipeline"},
|
||||
)
|
||||
|
||||
assert checkpoint_module.is_minimax_h3_modular("MiniMaxAI/MiniMax-H3") is False
|
||||
|
||||
|
||||
def test_native_local_model_root_keeps_partition_layout(tmp_path):
|
||||
root = tmp_path / "MiniMax-H3"
|
||||
(root / "FL2VA" / "text_encoder").mkdir(parents=True)
|
||||
(root / "FL2VA" / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
|
||||
|
||||
assert is_minimax_h3_modular(str(root)) is False
|
||||
assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
|
||||
|
||||
|
||||
def test_diffusion_resolver_returns_local_modular_root_without_partition_files(tmp_path):
|
||||
root = tmp_path / "FastH3"
|
||||
root.mkdir()
|
||||
(root / "modular_model_index.json").write_text(
|
||||
json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(root)
|
||||
assert not (root / "FL2VA").exists()
|
||||
assert not (root / "Ref2VA").exists()
|
||||
|
||||
|
||||
def test_diffusion_resolver_hub_modular_downloads_transformer_contract_only(monkeypatch, tmp_path):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
|
||||
|
||||
snapshot = tmp_path / "modular-snapshot"
|
||||
calls: dict[str, object] = {}
|
||||
|
||||
def fake_modular(_model: str, _revision: str | None = None) -> bool:
|
||||
return True
|
||||
|
||||
def fake_download(**kwargs: object) -> str:
|
||||
calls.update(kwargs)
|
||||
snapshot.mkdir()
|
||||
return str(snapshot)
|
||||
|
||||
monkeypatch.setattr(pipeline_module, "is_minimax_h3_modular", fake_modular)
|
||||
monkeypatch.setattr(pipeline_module, "download_weights_from_hf_specific", fake_download)
|
||||
|
||||
result = resolve_minimax_h3_diffusion_model_path(
|
||||
"FastVideo/FastVideo-FastH3-8-Step-V2", "release-revision", "fl2va"
|
||||
)
|
||||
|
||||
assert result == str(snapshot)
|
||||
assert calls == {
|
||||
"model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
|
||||
"cache_dir": None,
|
||||
"allow_patterns": [
|
||||
"modular_model_index.json",
|
||||
"fastvideo_inference.json",
|
||||
"provenance.json",
|
||||
"transformer/**",
|
||||
],
|
||||
"revision": "release-revision",
|
||||
"require_all": True,
|
||||
}
|
||||
patterns = calls["allow_patterns"]
|
||||
assert isinstance(patterns, list)
|
||||
assert all("video_vae" not in pattern and "audio_vae" not in pattern for pattern in patterns)
|
||||
|
||||
|
||||
def test_diffusers_arch_config_normalizes_all_aliases():
|
||||
arch = MiniMaxH3DiTArchConfig.from_mapping(
|
||||
{
|
||||
"num_refiner_layers": 3,
|
||||
"ffn_dim": 12,
|
||||
"in_channels": 4,
|
||||
"audio_in_channels": 6,
|
||||
"freq_dim": 8,
|
||||
"time_embed_hidden_dim": 10,
|
||||
"rope_freq_dim": 5,
|
||||
"patch_size": [1, 2, 2],
|
||||
}
|
||||
)
|
||||
|
||||
assert arch.token_refiner_num_layers == 3
|
||||
assert arch.ffn_hidden_size == 12
|
||||
assert arch.latents_dim == 4
|
||||
assert arch.audio_latents_dim == 6
|
||||
assert arch.timestep_input_dim == 8
|
||||
assert arch.time_embed_hidden_size == 10
|
||||
assert arch.rope_inv_freq_len == 5
|
||||
assert arch.patch_size == (1, 2, 2)
|
||||
|
||||
|
||||
class _WeightTarget(nn.Module):
|
||||
def __init__(self, shape: tuple[int, ...], loader: Callable[..., None]) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.zeros(shape))
|
||||
self.weight.weight_loader = loader
|
||||
|
||||
|
||||
class _PipelineLoadTransformer(nn.Module):
|
||||
def __init__(self, loaded: set[str]) -> None:
|
||||
super().__init__()
|
||||
self.arch = MiniMaxH3DiTArchConfig(num_layers=1)
|
||||
self._loaded = loaded
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
list(weights)
|
||||
return self._loaded
|
||||
|
||||
def post_load_weights(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def test_pipeline_load_weights_rejects_missing_fast_h3_compression_gate():
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.transformer = _PipelineLoadTransformer({"blocks.0.attn.qkv_proj.weight"})
|
||||
pipeline.video_vae = None
|
||||
pipeline.audio_vae = None
|
||||
pipeline._fasth3 = None
|
||||
pipeline._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
|
||||
with pytest.raises(ValueError, match="missing compression gates"):
|
||||
pipeline.load_weights([("transformer.blocks.0.attn.qkv_proj.weight", torch.zeros(1))])
|
||||
|
||||
|
||||
def _load_fixture(
|
||||
*, diffusers_weights: bool
|
||||
) -> tuple[MiniMaxH3DiTModel, list[tuple[str | None, torch.Tensor]], list[tuple[int, torch.Tensor]]]:
|
||||
model = object.__new__(MiniMaxH3DiTModel)
|
||||
nn.Module.__init__(model)
|
||||
model.arch = MiniMaxH3DiTArchConfig(
|
||||
num_layers=1,
|
||||
hidden_size=3,
|
||||
num_attention_heads=2,
|
||||
attention_head_dim=2,
|
||||
ffn_hidden_size=6,
|
||||
rope_inv_freq_len=2,
|
||||
)
|
||||
model._diffusers_weights = diffusers_weights
|
||||
model._rope_theta = 10000.0
|
||||
qkv_calls: list[tuple[str | None, torch.Tensor]] = []
|
||||
fc1_calls: list[tuple[int, torch.Tensor]] = []
|
||||
|
||||
def qkv_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str | None = None) -> None:
|
||||
if shard_id is None:
|
||||
assert loaded_weight.shape == (12, 3)
|
||||
param.data.copy_(loaded_weight)
|
||||
qkv_calls.append((None, loaded_weight.clone()))
|
||||
return
|
||||
assert loaded_weight.shape == (4, 3)
|
||||
assert shard_id in {"q", "k", "v"}
|
||||
offset = {"q": 0, "k": 4, "v": 8}[shard_id]
|
||||
param.data[offset : offset + 4].copy_(loaded_weight)
|
||||
qkv_calls.append((shard_id, loaded_weight.clone()))
|
||||
|
||||
def fc1_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: int) -> None:
|
||||
assert loaded_weight.shape == (3, 3)
|
||||
param.data[shard_id * 3 : (shard_id + 1) * 3].copy_(loaded_weight)
|
||||
fc1_calls.append((shard_id, loaded_weight.clone()))
|
||||
|
||||
model.blocks = nn.ModuleList([nn.Module()])
|
||||
model.blocks[0].attn = nn.Module()
|
||||
model.blocks[0].attn.qkv_proj = _WeightTarget((12, 3), qkv_loader)
|
||||
model.blocks[0].mlp = nn.Module()
|
||||
model.blocks[0].mlp.fc1 = _WeightTarget((6, 3), fc1_loader)
|
||||
model.rope = nn.Module()
|
||||
model.rope.register_buffer("inv_freq", torch.zeros(2))
|
||||
return model, qkv_calls, fc1_calls
|
||||
|
||||
|
||||
def test_diffusers_load_weights_uses_shape_aware_qkv_and_mlp_slices():
|
||||
model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=True)
|
||||
q = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 10
|
||||
k = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100
|
||||
v = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 200
|
||||
mlp = torch.arange(18, dtype=torch.float32).reshape(6, 3) + 300
|
||||
|
||||
loaded = model.load_weights(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.weight", q),
|
||||
("transformer_blocks.0.attn.to_k.weight", k),
|
||||
("transformer_blocks.0.attn.to_v.weight", v),
|
||||
("transformer_blocks.0.ff.net.0.proj.weight", mlp),
|
||||
]
|
||||
)
|
||||
|
||||
assert loaded == {
|
||||
"blocks.0.attn.qkv_proj.weight",
|
||||
"blocks.0.mlp.fc1.weight",
|
||||
"rope.inv_freq",
|
||||
}
|
||||
assert [shard_id for shard_id, _ in qkv_calls] == ["q", "k", "v"]
|
||||
assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
|
||||
torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, torch.cat((q, k, v)))
|
||||
torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, torch.cat((mlp[3:], mlp[:3])))
|
||||
torch.testing.assert_close(
|
||||
model.rope.inv_freq,
|
||||
1.0 / (10000.0 ** (torch.arange(0, 4, 2, dtype=torch.float32) / 4)),
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
|
||||
|
||||
def test_native_load_weights_keeps_grouped_qkv_and_gate_first_layouts():
|
||||
model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=False)
|
||||
grouped_qkv = torch.arange(36, dtype=torch.float32).reshape(12, 3)
|
||||
fc1 = torch.arange(18, dtype=torch.float32).reshape(6, 3)
|
||||
|
||||
loaded = model.load_weights(
|
||||
[
|
||||
("blocks.0.attn.qkv_proj.weight", grouped_qkv),
|
||||
("blocks.0.mlp.fc1.weight", fc1),
|
||||
]
|
||||
)
|
||||
|
||||
expected_qkv = torch.cat(
|
||||
(
|
||||
grouped_qkv[:2],
|
||||
grouped_qkv[6:8],
|
||||
grouped_qkv[2:4],
|
||||
grouped_qkv[8:10],
|
||||
grouped_qkv[4:6],
|
||||
grouped_qkv[10:],
|
||||
)
|
||||
)
|
||||
assert loaded == {"blocks.0.attn.qkv_proj.weight", "blocks.0.mlp.fc1.weight"}
|
||||
assert [shard_id for shard_id, _ in qkv_calls] == [None]
|
||||
assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
|
||||
torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, expected_qkv)
|
||||
torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, fc1)
|
||||
assert torch.equal(model.rope.inv_freq, torch.zeros(2))
|
||||
|
||||
|
||||
def test_diffusers_load_weights_rejects_unknown_name():
|
||||
model, _, _ = _load_fixture(diffusers_weights=True)
|
||||
|
||||
with pytest.raises(ValueError, match="unsupported Diffusers H3 weight"):
|
||||
model.load_weights([("transformer_blocks.0.attn.not_a_projection.weight", torch.zeros(4, 3))])
|
||||
|
||||
|
||||
def test_diffusers_load_weights_rejects_duplicate_source_name():
|
||||
model, _, _ = _load_fixture(diffusers_weights=True)
|
||||
weight = torch.zeros(4, 3)
|
||||
|
||||
with pytest.raises(ValueError, match="duplicate Diffusers H3 weight"):
|
||||
model.load_weights(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.weight", weight),
|
||||
("transformer_blocks.0.attn.to_q.weight", weight),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_diffusers_load_weights_rejects_incomplete_qkv_group():
|
||||
model, _, _ = _load_fixture(diffusers_weights=True)
|
||||
|
||||
with pytest.raises(ValueError, match="incomplete Diffusers H3 QKV"):
|
||||
model.load_weights(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.weight", torch.zeros(4, 3)),
|
||||
("transformer_blocks.0.attn.to_k.weight", torch.zeros(4, 3)),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_diffusers_load_weights_rejects_missing_compression_gate_parameter():
|
||||
model, _, _ = _load_fixture(diffusers_weights=True)
|
||||
|
||||
with pytest.raises(ValueError, match="has no model parameter"):
|
||||
model.load_weights([("transformer_blocks.0.attn.to_gate_compress.weight", torch.zeros(4, 3))])
|
||||
|
||||
|
||||
def test_resolve_sigma_positions_uses_release_metadata_schedule():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline._fasth3 = None
|
||||
pipeline._lora_sigma_schedules = {}
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
pipeline._base_schedule_by_partition = {
|
||||
"fl2va": DMD2SigmaSchedule.from_metadata(spec.release_metadata()),
|
||||
}
|
||||
|
||||
positions, steps = pipeline._resolve_sigma_positions("t2va", _sampling())
|
||||
assert positions == FASTH3_V2_BASE_SCHEDULE.base_schedule
|
||||
assert steps == 8 == len(positions) - 1
|
||||
with pytest.raises(OmniClientError, match="must be 8"):
|
||||
pipeline._resolve_sigma_positions("t2va", _sampling(num_inference_steps=9))
|
||||
@@ -423,6 +423,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership(
|
||||
assert pipeline.partition == expected_partition
|
||||
assert pipeline.supported_tasks == expected_tasks
|
||||
assert len(created["dit"]) == expected_dits
|
||||
assert all(kwargs["diffusers_weights"] is False for _, kwargs in created["dit"])
|
||||
component_path = tmp_path / component_partition
|
||||
assert created["video_vae"] == [str(component_path / "video_vae")]
|
||||
assert created["audio_vae"] == [str(component_path / "audio_vae")]
|
||||
|
||||
@@ -262,7 +262,10 @@ def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path):
|
||||
|
||||
def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))})
|
||||
_write_adapter(
|
||||
path,
|
||||
tensors={"transformer_blocks.0.attn.not_a_projection.lora_A.weight": torch.ones((_RANK, _HIDDEN))},
|
||||
)
|
||||
# Dropping it silently would load a model that is not the distilled student.
|
||||
with pytest.raises(FastH3AdapterError, match="name no known"):
|
||||
_load(path.parent)
|
||||
|
||||
@@ -10,6 +10,7 @@ from vllm_omni.diffusion.layers.indexed_modulation import (
|
||||
_MAX_1D_GRID_SIZE,
|
||||
_iter_row_chunks,
|
||||
_launch_row_chunks,
|
||||
_use_hopper_bf16_affine_semantics,
|
||||
indexed_gate,
|
||||
indexed_gate_rms_norm_scale_shift,
|
||||
indexed_scale_shift_,
|
||||
@@ -190,3 +191,55 @@ def test_fused_modulation_preserves_bf16_residual_boundary() -> None:
|
||||
)
|
||||
|
||||
assert torch.equal(modulated_out, expected)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.skipif(
|
||||
torch.cuda.is_available() and torch.cuda.get_device_capability()[0] != 9,
|
||||
reason="SM90-specific precision path",
|
||||
)
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
|
||||
def test_hopper_fused_modulation_matches_pytorch_reference() -> None:
|
||||
torch.manual_seed(42)
|
||||
rows, hidden_size, conditions = 128, 3072, 4
|
||||
residual = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
gate = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
branch = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
shift = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
scale = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
indices = torch.arange(rows, device="cuda") % conditions
|
||||
eps = 1e-6
|
||||
|
||||
reference_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16)
|
||||
normalized = reference_residual.float()
|
||||
variance = normalized.pow(2).mean(-1, keepdim=True)
|
||||
normalized = (weight.float() * (normalized * torch.rsqrt(variance + eps))).to(torch.bfloat16)
|
||||
reference_modulated = (
|
||||
normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)
|
||||
).to(torch.bfloat16)
|
||||
|
||||
actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift(
|
||||
residual,
|
||||
gate,
|
||||
branch,
|
||||
weight,
|
||||
shift,
|
||||
scale,
|
||||
indices,
|
||||
eps,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual_residual, reference_residual, atol=5e-2, rtol=5e-2)
|
||||
torch.testing.assert_close(actual_modulated, reference_modulated, atol=5e-2, rtol=5e-2)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
def test_hopper_precision_dispatch_is_architecture_scoped(monkeypatch) -> None:
|
||||
tensor = torch.empty(0, device="cuda")
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0))
|
||||
assert _use_hopper_bf16_affine_semantics(tensor)
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
|
||||
assert not _use_hopper_bf16_affine_semantics(tensor)
|
||||
@@ -1,4 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
@@ -41,7 +42,7 @@ class _FakeAttention(nn.Module):
|
||||
super().__init__()
|
||||
|
||||
|
||||
def _small_od_config():
|
||||
def _small_od_config(*, class_name=None):
|
||||
arch = {
|
||||
"num_layers": 1,
|
||||
"token_refiner_num_layers": 1,
|
||||
@@ -60,6 +61,8 @@ def _small_od_config():
|
||||
"final_adaln_out_features": 2 * 8,
|
||||
"rope_inv_freq_len": 2,
|
||||
}
|
||||
if class_name is not None:
|
||||
arch["_class_name"] = class_name
|
||||
return SimpleNamespace(
|
||||
tf_model_config=arch,
|
||||
parallel_config=SimpleNamespace(ulysses_degree=1),
|
||||
@@ -133,6 +136,24 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
|
||||
)
|
||||
|
||||
|
||||
def test_explicit_native_weight_format_overrides_diffusers_class_name(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
|
||||
|
||||
monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "Attention", _FakeAttention)
|
||||
monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
|
||||
|
||||
model = h3.MiniMaxH3DiTModel(
|
||||
_small_od_config(class_name="MiniMaxH3Transformer3DModel"),
|
||||
diffusers_weights=False,
|
||||
)
|
||||
|
||||
assert model._diffusers_weights is False
|
||||
|
||||
|
||||
class _WeightTarget(nn.Module):
|
||||
def __init__(self, loader):
|
||||
super().__init__()
|
||||
|
||||
@@ -187,7 +187,11 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
|
||||
"model": model_ref,
|
||||
"enforce_eager": True,
|
||||
# The fused BF16 baseline only fits on H100-80GB with encoder TP.
|
||||
"stage_0_tensor_parallel_size": 2,
|
||||
# This test loads the direct FL2VA diffusion pipeline (not the
|
||||
# disaggregated two-stage MiniMax-H3 pipeline), so configure both the
|
||||
# DiT world and the Qwen3-VL encoder TP explicitly.
|
||||
"tensor_parallel_size": 2,
|
||||
"text_encoder_tp_size": 2,
|
||||
"vae_use_tiling": True,
|
||||
# MiniMax H3's VAE is implemented by code shipped with the checkpoint.
|
||||
"trust_remote_code": True,
|
||||
|
||||
@@ -6,6 +6,7 @@ import pytest
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OfflineOmniClient
|
||||
from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
|
||||
@@ -21,17 +22,19 @@ def test_structured_diffusion_config_reaches_runtime(omni_runner) -> None:
|
||||
"""The default WAN resolver launches the typed diffusion config itself."""
|
||||
omni = omni_runner.omni
|
||||
resolved_stages = omni.engine.stage_configs
|
||||
stage_clients = omni.engine.stage_clients
|
||||
assert len(resolved_stages) == 1
|
||||
assert len(stage_clients) == 1
|
||||
stage_config = resolved_stages[0]
|
||||
od_config = stage_clients[0].od_config
|
||||
|
||||
assert isinstance(stage_config, VllmOmniDiffusionStageConfig)
|
||||
assert isinstance(od_config, OmniDiffusionConfig)
|
||||
assert stage_config.stage_id == 0
|
||||
assert stage_config.model_stage == "dit"
|
||||
assert stage_config.final_output_type == "video"
|
||||
|
||||
assert omni.engine.stage_vllm_configs == [None]
|
||||
od_config = omni.get_diffusion_od_config()
|
||||
assert od_config is not None
|
||||
assert od_config.stage_id == 0
|
||||
assert od_config.model == MODEL
|
||||
assert od_config.model_class_name == "WanPipeline"
|
||||
|
||||
@@ -10,24 +10,24 @@ import base64
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
import uuid
|
||||
import wave
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
import websockets
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
from websockets.exceptions import ConnectionClosed, WebSocketException
|
||||
except ImportError as exc: # pragma: no cover - driver dependency.
|
||||
raise SystemExit("Install websockets first: pip install websockets") from exc
|
||||
|
||||
from vllm_omni.clients.duplex import (
|
||||
EventCollector,
|
||||
write_pcm16_wav,
|
||||
)
|
||||
from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats
|
||||
from vllm_omni.clients.duplex import EventCollector, write_pcm16_wav
|
||||
from vllm_omni.clients.duplex import (
|
||||
wait_for_condition as wait_for,
|
||||
)
|
||||
@@ -35,6 +35,24 @@ from vllm_omni.clients.duplex import (
|
||||
SAMPLE_RATE_HZ = 24_000
|
||||
FRAME_SAMPLES = 1_920
|
||||
FRAME_PERIOD_S = FRAME_SAMPLES / SAMPLE_RATE_HZ
|
||||
AUDIO_DELTA_EVENT_TYPES = frozenset({"response.audio.delta", "response.output_audio.delta"})
|
||||
|
||||
|
||||
class _AudioAcceptanceError(AssertionError):
|
||||
"""A local acceptance failure with numeric diagnostics safe to report."""
|
||||
|
||||
def __init__(self, check: str, message: str, stats: dict[str, object]) -> None:
|
||||
super().__init__(message)
|
||||
self.check = check
|
||||
self.stats = dict(stats)
|
||||
|
||||
|
||||
class _ProbeTransport(Protocol):
|
||||
async def recv(self) -> str | bytes: ...
|
||||
|
||||
async def send(self, message: str) -> None: ...
|
||||
|
||||
async def close(self) -> None: ...
|
||||
|
||||
|
||||
class RawRealtimeProbe:
|
||||
@@ -46,29 +64,33 @@ class RawRealtimeProbe:
|
||||
keeps a transport-only client and reads everything from ``.events``.
|
||||
"""
|
||||
|
||||
def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None:
|
||||
def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024, close_timeout_s: float = 10.0) -> None:
|
||||
self.url = url
|
||||
self.max_size = max_size
|
||||
self.close_timeout_s = close_timeout_s
|
||||
self.events = EventCollector()
|
||||
self._ws = None
|
||||
self._ws: _ProbeTransport | None = None
|
||||
self._reader_task: asyncio.Task[None] | None = None
|
||||
|
||||
async def __aenter__(self) -> RawRealtimeProbe:
|
||||
self._ws = await websockets.connect(self.url, max_size=self.max_size)
|
||||
self._ws = await websockets.connect(self.url, max_size=self.max_size, close_timeout=self.close_timeout_s)
|
||||
self._reader_task = asyncio.create_task(self._read_events())
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, traceback) -> None:
|
||||
if self._ws is not None:
|
||||
await self._ws.close()
|
||||
if self._reader_task is not None:
|
||||
self._reader_task.cancel()
|
||||
try:
|
||||
await self._reader_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
try:
|
||||
if self._ws is not None:
|
||||
await self._ws.close()
|
||||
finally:
|
||||
if self._reader_task is not None:
|
||||
self._reader_task.cancel()
|
||||
try:
|
||||
await self._reader_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def _read_events(self) -> None:
|
||||
assert self._ws is not None
|
||||
try:
|
||||
while True:
|
||||
raw = await self._ws.recv()
|
||||
@@ -81,6 +103,7 @@ class RawRealtimeProbe:
|
||||
return
|
||||
|
||||
async def send(self, event: dict[str, object]) -> None:
|
||||
assert self._ws is not None
|
||||
await self._ws.send(json.dumps(event))
|
||||
|
||||
|
||||
@@ -142,6 +165,38 @@ def _events(client: RawRealtimeProbe, event_type: str) -> list[dict[str, object]
|
||||
return [event for event in client.events.events if event.get("type") == event_type]
|
||||
|
||||
|
||||
def _audio_events(client: RawRealtimeProbe) -> list[dict[str, object]]:
|
||||
"""Return audio deltas from both current and legacy Realtime event names."""
|
||||
return [event for event in client.events.events if event.get("type") in AUDIO_DELTA_EVENT_TYPES]
|
||||
|
||||
|
||||
def _validated_audio_event(event: dict[str, object]) -> bytes:
|
||||
"""Decode one wire audio delta without depending on client event aliases."""
|
||||
encoded = event.get("delta")
|
||||
if not isinstance(encoded, str):
|
||||
raise ValueError("missing_audio_delta")
|
||||
try:
|
||||
chunk = base64.b64decode(encoded, validate=True)
|
||||
except ValueError:
|
||||
raise ValueError("invalid_audio_base64") from None
|
||||
if chunk:
|
||||
if len(chunk) % 2:
|
||||
raise ValueError("unaligned_pcm16_packet")
|
||||
if event.get("sample_rate_hz") != SAMPLE_RATE_HZ:
|
||||
raise ValueError("invalid_audio_sample_rate")
|
||||
if not isinstance(event.get("response_id"), str) or not event.get("response_id"):
|
||||
raise ValueError("missing_audio_response_id")
|
||||
return chunk
|
||||
|
||||
|
||||
def _validated_audio_chunks(client: RawRealtimeProbe) -> list[tuple[dict[str, object], bytes]]:
|
||||
return [(event, _validated_audio_event(event)) for event in _audio_events(client)]
|
||||
|
||||
|
||||
def _audio_bytes(client: RawRealtimeProbe) -> bytes:
|
||||
return b"".join(chunk for _, chunk in _validated_audio_chunks(client) if chunk)
|
||||
|
||||
|
||||
async def _open_session(
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
@@ -151,22 +206,7 @@ async def _open_session(
|
||||
) -> tuple[RawRealtimeProbe, dict[str, object]]:
|
||||
client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id))
|
||||
await client.__aenter__()
|
||||
await client.send(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"session_id": session_id,
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"input_audio_format": "pcm_f32le",
|
||||
"output_audio_format": "pcm16",
|
||||
"voice": args.voice,
|
||||
"instructions": persona,
|
||||
"turn_detection": None,
|
||||
"extra_body": {"auto_response": True},
|
||||
},
|
||||
}
|
||||
)
|
||||
await client.send(_session_update(args, session_id=session_id, persona=persona))
|
||||
event_type = "error" if expect_error else "session.created"
|
||||
await wait_for(
|
||||
lambda: client.events.count(event_type) > 0,
|
||||
@@ -176,6 +216,23 @@ async def _open_session(
|
||||
return client, _events(client, event_type)[-1]
|
||||
|
||||
|
||||
def _session_update(args: argparse.Namespace, *, session_id: str, persona: str) -> dict[str, object]:
|
||||
return {
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"session_id": session_id,
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"input_audio_format": "pcm_f32le",
|
||||
"output_audio_format": "pcm16",
|
||||
"voice": args.voice,
|
||||
"instructions": persona,
|
||||
"turn_detection": None,
|
||||
"extra_body": {"auto_response": True},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
async def _close_session(client: RawRealtimeProbe, *, timeout_s: float) -> None:
|
||||
await client.send({"type": "session.close"})
|
||||
await wait_for(
|
||||
@@ -200,23 +257,24 @@ async def _stream_frames(
|
||||
frame_count = min(frame_count, max_frames)
|
||||
for seq in range(frame_count):
|
||||
frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
|
||||
frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size))
|
||||
frame = np.ascontiguousarray(frame, dtype="<f4")
|
||||
await client.send(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": base64.b64encode(frame).decode("ascii"),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": SAMPLE_RATE_HZ,
|
||||
"duration_ms": 80,
|
||||
"audio_end_ms": (seq + 1) * 80,
|
||||
"is_speech": bool(np.sqrt(np.mean(np.square(frame))) >= 1e-4),
|
||||
}
|
||||
)
|
||||
await client.send(_frame_event(frame, seq))
|
||||
await asyncio.sleep(FRAME_PERIOD_S)
|
||||
return frame_count
|
||||
|
||||
|
||||
def _frame_event(frame: np.ndarray, seq: int) -> dict[str, object]:
|
||||
frame = np.ascontiguousarray(np.pad(frame, (0, FRAME_SAMPLES - frame.size)), dtype="<f4")
|
||||
return {
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": base64.b64encode(frame).decode("ascii"),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": SAMPLE_RATE_HZ,
|
||||
"duration_ms": 80,
|
||||
"audio_end_ms": (seq + 1) * 80,
|
||||
"is_speech": bool(np.sqrt(np.mean(np.square(frame))) >= 1e-4),
|
||||
}
|
||||
|
||||
|
||||
def _audio_frame_stats(
|
||||
raw: bytes,
|
||||
*,
|
||||
@@ -242,11 +300,7 @@ def _audio_frame_stats(
|
||||
|
||||
|
||||
def _response_ids(client: RawRealtimeProbe) -> set[str]:
|
||||
return {
|
||||
str(event["response_id"])
|
||||
for event in _events(client, "response.output_audio.delta")
|
||||
if isinstance(event.get("response_id"), str)
|
||||
}
|
||||
return {str(event["response_id"]) for event in _audio_events(client) if isinstance(event.get("response_id"), str)}
|
||||
|
||||
|
||||
def _session_result(
|
||||
@@ -256,40 +310,66 @@ def _session_result(
|
||||
args: argparse.Namespace,
|
||||
minimum_chunks: int,
|
||||
) -> tuple[bytes, set[str], dict[str, object]]:
|
||||
chunks = sum(len(items) for items in client.events.response_audio.values())
|
||||
raw = client.events.audio_bytes()
|
||||
if chunks < minimum_chunks or not raw or len(raw) % 2:
|
||||
raise AssertionError(f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}")
|
||||
if isinstance(client, RawRealtimeProbe):
|
||||
audio_chunks = _validated_audio_chunks(client)
|
||||
chunks = sum(bool(chunk) for _, chunk in audio_chunks)
|
||||
raw = b"".join(chunk for _, chunk in audio_chunks if chunk)
|
||||
else:
|
||||
# Unit fixtures may provide an already-decoded EventCollector without
|
||||
# retaining the base64 wire field. Real websocket probes always take
|
||||
# the strict wire-validation path above.
|
||||
chunks = sum(len(items) for items in client.events.response_audio.values())
|
||||
raw = client.events.audio_bytes()
|
||||
stats: dict[str, object] = {"audio_chunks": chunks, "audio_samples": len(raw) // 2}
|
||||
if not raw or len(raw) % 2:
|
||||
raise _AudioAcceptanceError("pcm16_output", f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}", stats)
|
||||
pcm = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
|
||||
rms = float(np.sqrt(np.mean(np.square(pcm))))
|
||||
stats["audio_rms"] = rms
|
||||
# This low floor only catches corrupt, non-finite, or effectively zero PCM.
|
||||
# Audible speech is checked per frame below with the higher 1e-3 default.
|
||||
if not np.isfinite(pcm).all() or rms <= 1e-5:
|
||||
raise AssertionError(f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}")
|
||||
rates = {
|
||||
int(event["sample_rate_hz"])
|
||||
for event in _events(client, "response.output_audio.delta")
|
||||
if isinstance(event.get("sample_rate_hz"), int)
|
||||
}
|
||||
raise _AudioAcceptanceError(
|
||||
"audio_rms", f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}", stats
|
||||
)
|
||||
rates = {rate for event in _audio_events(client) if isinstance(rate := event.get("sample_rate_hz"), int)}
|
||||
if rates != {SAMPLE_RATE_HZ}:
|
||||
raise AssertionError(f"unexpected output sample rates: {sorted(rates)}")
|
||||
stats = _audio_frame_stats(
|
||||
raw,
|
||||
input_frames=input_frames,
|
||||
voiced_frame_rms_threshold=args.voiced_frame_rms_threshold,
|
||||
)
|
||||
if not 0 <= int(stats["frame_deficit"]) <= args.max_frame_deficit:
|
||||
raise AssertionError(f"invalid synchronous frame accounting: input={input_frames}, stats={stats}")
|
||||
raise _AudioAcceptanceError("sample_rate", f"unexpected output sample rates: {sorted(rates)}", stats)
|
||||
if input_frames < 1:
|
||||
raise _AudioAcceptanceError("input_frames", f"audio received without input frames: input={input_frames}", stats)
|
||||
try:
|
||||
frame_stats = _audio_frame_stats(
|
||||
raw,
|
||||
input_frames=input_frames,
|
||||
voiced_frame_rms_threshold=args.voiced_frame_rms_threshold,
|
||||
)
|
||||
except AssertionError as exc:
|
||||
# _audio_frame_stats only formats local sample counts, not wire data.
|
||||
raise _AudioAcceptanceError("whole_codec_frames", str(exc), stats) from exc
|
||||
stats.update(frame_stats)
|
||||
if chunks < minimum_chunks:
|
||||
raise _AudioAcceptanceError(
|
||||
"audio_chunks", f"insufficient audio chunks: chunks={chunks}, minimum={minimum_chunks}", stats
|
||||
)
|
||||
if not 0 <= int(frame_stats["frame_deficit"]) <= args.max_frame_deficit:
|
||||
raise _AudioAcceptanceError(
|
||||
"frame_deficit",
|
||||
f"invalid synchronous frame accounting: input={input_frames}, maximum_deficit={args.max_frame_deficit}, "
|
||||
f"stats={stats}",
|
||||
stats,
|
||||
)
|
||||
min_voiced_frames = int(args.min_voiced_frames)
|
||||
if min_voiced_frames < 1:
|
||||
raise ValueError(f"min_voiced_frames must be positive, got {min_voiced_frames}")
|
||||
if int(stats["voiced_frames"]) < min_voiced_frames:
|
||||
raise AssertionError(
|
||||
f"output contained insufficient audible speech: minimum_voiced_frames={min_voiced_frames}, stats={stats}"
|
||||
if int(frame_stats["voiced_frames"]) < min_voiced_frames:
|
||||
raise _AudioAcceptanceError(
|
||||
"voiced_frames",
|
||||
f"output contained insufficient audible speech: minimum_voiced_frames={min_voiced_frames}, stats={stats}",
|
||||
stats,
|
||||
)
|
||||
runtime_metadata = [
|
||||
metadata["vllm_omni"]
|
||||
for event in _events(client, "response.output_audio.delta")
|
||||
for event in _audio_events(client)
|
||||
if isinstance((metadata := event.get("metadata")), dict) and isinstance(metadata.get("vllm_omni"), dict)
|
||||
]
|
||||
if not any(
|
||||
@@ -298,18 +378,19 @@ def _session_result(
|
||||
and item.get("runner_kv_backed") is True
|
||||
for item in runtime_metadata
|
||||
):
|
||||
raise AssertionError(f"audio events did not prove the scheduler data plane: {runtime_metadata}")
|
||||
raise _AudioAcceptanceError(
|
||||
"scheduler_data_plane", "audio events did not prove the scheduler data plane", stats
|
||||
)
|
||||
response_ids = _response_ids(client)
|
||||
if len(response_ids) != 1:
|
||||
raise AssertionError(f"continuous stream used multiple responses: {sorted(response_ids)}")
|
||||
raise _AudioAcceptanceError(
|
||||
"response_identity", f"continuous stream requires one response: count={len(response_ids)}", stats
|
||||
)
|
||||
return (
|
||||
raw,
|
||||
response_ids,
|
||||
{
|
||||
**stats,
|
||||
"audio_chunks": chunks,
|
||||
"audio_samples": len(raw) // 2,
|
||||
"audio_rms": rms,
|
||||
"response_ids": sorted(response_ids),
|
||||
},
|
||||
)
|
||||
@@ -332,6 +413,8 @@ def _save(output_dir: Path, name: str, client: RawRealtimeProbe, raw: bytes) ->
|
||||
|
||||
|
||||
async def run(args: argparse.Namespace) -> dict[str, object]:
|
||||
if getattr(args, "sessions", None) is not None:
|
||||
return await run_load(args)
|
||||
input_path = Path(args.input_wav)
|
||||
input_identity = _input_identity(
|
||||
input_path,
|
||||
@@ -381,14 +464,12 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
|
||||
)
|
||||
await asyncio.gather(
|
||||
wait_for(
|
||||
lambda: len(primary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
|
||||
lambda: len(_audio_bytes(primary)) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
|
||||
timeout_s=args.timeout_s,
|
||||
label="primary frame coverage",
|
||||
),
|
||||
wait_for(
|
||||
lambda: (
|
||||
len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit
|
||||
),
|
||||
lambda: (len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit),
|
||||
timeout_s=args.timeout_s,
|
||||
label="secondary frame coverage",
|
||||
),
|
||||
@@ -426,16 +507,14 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
|
||||
await asyncio.gather(
|
||||
wait_for(
|
||||
lambda: (
|
||||
len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES)
|
||||
>= secondary_total_frames - args.max_frame_deficit
|
||||
len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_total_frames - args.max_frame_deficit
|
||||
),
|
||||
timeout_s=args.timeout_s,
|
||||
label="survivor frame coverage after slot recycle",
|
||||
),
|
||||
wait_for(
|
||||
lambda: (
|
||||
len(replacement.events.audio_bytes()) // (2 * FRAME_SAMPLES)
|
||||
>= replacement_frames - args.max_frame_deficit
|
||||
len(_audio_bytes(replacement)) // (2 * FRAME_SAMPLES) >= replacement_frames - args.max_frame_deficit
|
||||
),
|
||||
timeout_s=args.timeout_s,
|
||||
label="replacement frame coverage",
|
||||
@@ -493,6 +572,284 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
|
||||
return result
|
||||
|
||||
|
||||
def _distribution(values: list[float]) -> dict[str, int | float | None]:
|
||||
if not values:
|
||||
return {"count": 0, "median": None, "p99": None, "max": None, "argmax_index": None}
|
||||
return {
|
||||
"count": len(values),
|
||||
"median": float(np.median(values)),
|
||||
"p99": float(np.quantile(values, 0.99)),
|
||||
"max": max(values),
|
||||
"argmax_index": int(np.argmax(values)),
|
||||
}
|
||||
|
||||
|
||||
def _load_metrics(client: RawRealtimeProbe, sends: list[tuple[float, float, float]]) -> dict[str, object]:
|
||||
# Use the collector's local clock, never a timestamp supplied by the server.
|
||||
packets: list[dict[str, object]] = []
|
||||
response_ids: set[str] = set()
|
||||
packet_times: list[float] = []
|
||||
packet_bytes: list[int] = []
|
||||
invalid_packets: list[dict[str, object]] = []
|
||||
samples = 0
|
||||
for index, (raw, received) in enumerate(zip(client.events.events, client.events.event_received_at_s, strict=True)):
|
||||
if raw.get("type") not in AUDIO_DELTA_EVENT_TYPES:
|
||||
continue
|
||||
try:
|
||||
chunk = _validated_audio_event(raw)
|
||||
except ValueError as exc:
|
||||
# Only fixed diagnostic codes from the validator are exported,
|
||||
# never audio payloads or details from the remote endpoint.
|
||||
invalid_packets.append({"event_index": index, "received_at_s": received, "code": str(exc)})
|
||||
continue
|
||||
if chunk:
|
||||
samples += len(chunk) // 2
|
||||
packet_times.append(received)
|
||||
packet_bytes.append(len(chunk))
|
||||
response_id = raw.get("response_id")
|
||||
packets.append({"received_at_s": received, "samples": len(chunk) // 2, "response_id": response_id})
|
||||
if isinstance(response_id, str) and response_id:
|
||||
response_ids.add(response_id)
|
||||
intervals = [(b - a) * 1000 for a, b in zip(packet_times, packet_times[1:])]
|
||||
first_send = sends[0][1] if sends else None
|
||||
audio_duration = samples / SAMPLE_RATE_HZ
|
||||
first_audio_ms = (packet_times[0] - first_send) * 1000 if packet_times and first_send is not None else None
|
||||
client_rtf = (packet_times[-1] - first_send) / audio_duration if audio_duration and first_send is not None else None
|
||||
continuity = compute_continuity_stats(packet_times, packet_bytes, SAMPLE_RATE_HZ) if packet_times else None
|
||||
lateness = _distribution([max(0.0, sent - planned) * 1000 for planned, sent, _ in sends])
|
||||
p99_lateness = lateness["p99"]
|
||||
return {
|
||||
"input_frames": len(sends),
|
||||
"output_samples": samples,
|
||||
# Populate frame accounting from _session_result, including on failure.
|
||||
"frame_deficit": None,
|
||||
"response_ids": sorted(response_ids),
|
||||
"client_first_audio_after_stream_start_ms": first_audio_ms,
|
||||
"client_stream_rtf": client_rtf,
|
||||
"client_audio_packet_interval_ms": _distribution(intervals),
|
||||
"client_send_lateness_ms": lateness,
|
||||
"client_pacing_warning": isinstance(p99_lateness, float) and p99_lateness > FRAME_PERIOD_S * 1000,
|
||||
"audio_underrun_s": continuity.max_underrun_s if continuity else None,
|
||||
"audio_underrun_event_count": continuity.underrun_event_count if continuity else None,
|
||||
"audio_continuity_ok": continuity.is_continuous if continuity else None,
|
||||
"input_timeline": [
|
||||
{"frame_index": i, "planned_at_s": planned, "send_started_at_s": sent, "send_completed_at_s": completed}
|
||||
for i, (planned, sent, completed) in enumerate(sends)
|
||||
],
|
||||
"audio_packet_timeline": packets,
|
||||
"invalid_audio_packets": invalid_packets,
|
||||
}
|
||||
|
||||
|
||||
def _check_load_connection(client: RawRealtimeProbe) -> None:
|
||||
if client.events.errors():
|
||||
raise RuntimeError("server emitted a Realtime error (see the session error count)")
|
||||
if client._reader_task is not None and client._reader_task.done():
|
||||
client._reader_task.result() # Propagate malformed JSON/reader failures.
|
||||
raise RuntimeError("connection closed before session completion")
|
||||
|
||||
|
||||
async def _wait_load_event(client: RawRealtimeProbe, event_type: str, timeout_s: float) -> None:
|
||||
def received() -> bool:
|
||||
if client.events.count(event_type):
|
||||
return True
|
||||
_check_load_connection(client)
|
||||
return False
|
||||
|
||||
await wait_for(received, timeout_s=timeout_s, label=event_type)
|
||||
|
||||
|
||||
async def _paced_load_frames(
|
||||
client: RawRealtimeProbe,
|
||||
pcm: np.ndarray,
|
||||
*,
|
||||
epoch: float,
|
||||
timeout_s: float,
|
||||
sends: list[tuple[float, float, float]],
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
) -> None:
|
||||
next_send = epoch
|
||||
for seq in range(math.ceil(pcm.size / FRAME_SAMPLES)):
|
||||
await sleep(max(0.0, next_send - clock()))
|
||||
_check_load_connection(client)
|
||||
event = _frame_event(pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES], seq)
|
||||
started = clock()
|
||||
await asyncio.wait_for(client.send(event), timeout=timeout_s)
|
||||
completed = clock()
|
||||
sends.append((epoch + seq * FRAME_PERIOD_S, started, completed))
|
||||
# Do not burst to catch up after a stalled send; report accumulated lag.
|
||||
next_send = max(epoch + (seq + 1) * FRAME_PERIOD_S, completed + FRAME_PERIOD_S)
|
||||
|
||||
|
||||
async def _request_load_close(client: RawRealtimeProbe, timeout_s: float) -> None:
|
||||
await client.send({"type": "session.close"})
|
||||
await _wait_load_event(client, "session.closed", timeout_s)
|
||||
|
||||
|
||||
async def _run_load_session(
|
||||
args: argparse.Namespace,
|
||||
pcm: np.ndarray,
|
||||
index: int,
|
||||
ready: asyncio.Future[None],
|
||||
start: asyncio.Future[float],
|
||||
) -> dict[str, object]:
|
||||
session_id = f"personaplex-load-{index}-{uuid.uuid4().hex}"
|
||||
client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id), close_timeout_s=args.cleanup_timeout_s)
|
||||
sends: list[tuple[float, float, float]] = []
|
||||
error: str | None = None
|
||||
cleanup_error: str | None = None
|
||||
metrics: dict[str, object] | None = None
|
||||
session_stats: dict[str, object] = {}
|
||||
acceptance_check: str | None = None
|
||||
phase = "admission"
|
||||
try:
|
||||
async with client:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
client.send(_session_update(args, session_id=session_id, persona=args.persona)), args.timeout_s
|
||||
)
|
||||
await _wait_load_event(client, "session.created", args.timeout_s)
|
||||
capabilities = _capabilities(_events(client, "session.created")[-1])
|
||||
if (
|
||||
capabilities.get("chunk_period_ms") != 80
|
||||
or capabilities.get("supports_multi_session_same_replica") is not True
|
||||
):
|
||||
raise AssertionError("server did not advertise PersonaPlex multi-session capabilities")
|
||||
ready.set_result(None)
|
||||
phase = "stream"
|
||||
epoch = await asyncio.shield(start)
|
||||
await _paced_load_frames(client, pcm, epoch=epoch, timeout_s=args.timeout_s, sends=sends)
|
||||
phase = "drain"
|
||||
# Fixed, reported observation window. Closing must not extend
|
||||
# it and let a slow session pass by flushing after the cutoff.
|
||||
await asyncio.sleep(args.drain_s)
|
||||
_check_load_connection(client)
|
||||
_, _, session_stats = _session_result(
|
||||
client, input_frames=len(sends), args=args, minimum_chunks=args.minimum_audio_chunks
|
||||
)
|
||||
finally:
|
||||
metrics = _load_metrics(client, sends)
|
||||
if client.events.count("session.created") and not client.events.count("session.closed"):
|
||||
try:
|
||||
# Send and acknowledgement share one short cleanup budget,
|
||||
# independent of the potentially long inference timeout.
|
||||
await asyncio.wait_for(
|
||||
_request_load_close(client, args.cleanup_timeout_s), args.cleanup_timeout_s
|
||||
)
|
||||
except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, WebSocketException) as exc:
|
||||
cleanup_error = f"cleanup: {type(exc).__name__}"
|
||||
except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, AssertionError, WebSocketException) as exc:
|
||||
if isinstance(exc, _AudioAcceptanceError):
|
||||
acceptance_check = exc.check
|
||||
session_stats = exc.stats
|
||||
error = f"{phase}: AssertionError: {exc}"
|
||||
else:
|
||||
# Other assertions/errors may embed handshake credentials or wire
|
||||
# metadata. Only our local numeric acceptance diagnostics are safe.
|
||||
error = f"{phase}: {type(exc).__name__}"
|
||||
finally:
|
||||
# Admission failure must not leave the other sessions at the barrier.
|
||||
if not ready.done():
|
||||
ready.set_result(None)
|
||||
if metrics is None:
|
||||
metrics = _load_metrics(client, sends)
|
||||
if cleanup_error:
|
||||
error = f"{error}; {cleanup_error}" if error else cleanup_error
|
||||
if client.events.errors():
|
||||
error = error or "server emitted a Realtime error"
|
||||
if metrics["invalid_audio_packets"]:
|
||||
error = error or "invalid audio packets"
|
||||
first_audio = metrics["client_first_audio_after_stream_start_ms"]
|
||||
if isinstance(first_audio, float) and first_audio < 0:
|
||||
error = error or "audio arrived before the first input frame"
|
||||
rtf = metrics["client_stream_rtf"]
|
||||
if args.max_client_rtf is not None and isinstance(rtf, float) and rtf > args.max_client_rtf:
|
||||
error = error or f"client_stream_rtf exceeds {args.max_client_rtf}"
|
||||
return {
|
||||
"index": index,
|
||||
"session_id": session_id,
|
||||
"ok": error is None,
|
||||
"error": error,
|
||||
"acceptance_check": acceptance_check,
|
||||
"server_error_count": len(client.events.errors()),
|
||||
"server_error_codes": [
|
||||
str(body.get("code", "unknown")) if isinstance(body := event.get("error"), dict) else "unknown"
|
||||
for event in client.events.errors()
|
||||
],
|
||||
**metrics,
|
||||
**session_stats,
|
||||
}
|
||||
|
||||
|
||||
async def run_load(args: argparse.Namespace) -> dict[str, object]:
|
||||
input_path = Path(args.input_wav)
|
||||
identity = _input_identity(input_path, expected_sha256=args.expected_input_sha256)
|
||||
pcm = _read_wav_as_float32(input_path)
|
||||
if args.tail_s:
|
||||
pcm = np.concatenate([pcm, np.zeros(round(args.tail_s * SAMPLE_RATE_HZ), dtype=np.float32)])
|
||||
pcm = pcm[: args.load_frames * FRAME_SAMPLES]
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
# Refuse to silently overwrite evidence from a previous sweep.
|
||||
result_path = output_dir / "load-result.json"
|
||||
if result_path.exists():
|
||||
raise FileExistsError(result_path)
|
||||
loop = asyncio.get_running_loop()
|
||||
start: asyncio.Future[float] = loop.create_future()
|
||||
ready: list[asyncio.Future[None]] = [loop.create_future() for _ in range(args.sessions)]
|
||||
tasks = [asyncio.create_task(_run_load_session(args, pcm, i, r, start)) for i, r in enumerate(ready)]
|
||||
try:
|
||||
await asyncio.gather(*ready)
|
||||
start.set_result(time.monotonic())
|
||||
rows = await asyncio.gather(*tasks)
|
||||
finally:
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
seen: set[str] = set()
|
||||
shared: set[str] = set()
|
||||
for row in rows:
|
||||
assert isinstance(row["response_ids"], list)
|
||||
response_ids = set(row["response_ids"])
|
||||
shared.update(seen & response_ids)
|
||||
seen.update(response_ids)
|
||||
for row in rows:
|
||||
assert isinstance(row["response_ids"], list)
|
||||
if shared.intersection(row["response_ids"]):
|
||||
row.update(ok=False, error="concurrent sessions shared a response id")
|
||||
result = {
|
||||
"schema_version": 1,
|
||||
"mode": "paced_multi_session",
|
||||
"model": args.model,
|
||||
"requested_sessions": args.sessions,
|
||||
"passed_sessions": sum(row["ok"] is True for row in rows),
|
||||
"ok": all(row["ok"] is True for row in rows),
|
||||
"input": {**identity, "tail_s": args.tail_s, "load_frames": args.load_frames},
|
||||
"frame_period_s": FRAME_PERIOD_S,
|
||||
"drain_s": args.drain_s,
|
||||
"cleanup_timeout_s": args.cleanup_timeout_s,
|
||||
"client_pacing_warning_threshold_ms": FRAME_PERIOD_S * 1000,
|
||||
"client_pacing_warning_sessions": [row["index"] for row in rows if row["client_pacing_warning"]],
|
||||
"max_client_rtf": args.max_client_rtf,
|
||||
"acceptance": {
|
||||
"max_frame_deficit": args.max_frame_deficit,
|
||||
"min_voiced_frames": args.min_voiced_frames,
|
||||
"minimum_audio_chunks": args.minimum_audio_chunks,
|
||||
"voiced_frame_rms_threshold": args.voiced_frame_rms_threshold,
|
||||
},
|
||||
"shared_response_ids": sorted(shared),
|
||||
"sessions": rows,
|
||||
"reported_server": {"revision": args.server_revision, "hardware": args.server_hardware},
|
||||
"measurement_origin": "client monotonic clock; admission excluded; streaming and drain retained",
|
||||
"interval_definition": "nonempty audio packet arrivals, not server ticks or per-frame inference latency",
|
||||
}
|
||||
with result_path.open("x", encoding="utf-8") as stream:
|
||||
json.dump(result, stream, ensure_ascii=False, indent=2, allow_nan=False)
|
||||
return result
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1")
|
||||
@@ -510,6 +867,12 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser.add_argument("--tail-s", type=float, default=2.0)
|
||||
parser.add_argument("--drain-s", type=float, default=2.0)
|
||||
parser.add_argument("--timeout-s", type=float, default=180.0)
|
||||
parser.add_argument(
|
||||
"--cleanup-timeout-s",
|
||||
type=float,
|
||||
default=5.0,
|
||||
help="Load mode: shared session-close/ack budget; also bounds the websocket close handshake.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-frame-deficit",
|
||||
type=int,
|
||||
@@ -536,11 +899,39 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
default=0,
|
||||
help="Frames sent after slot reuse; 0 replays the full input so the replacement must produce audible output.",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
parser.add_argument(
|
||||
"--sessions", type=int, help="Opt into paced load mode with N sessions; omit for the lifecycle test."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--load-frames", type=int, default=1000, help="Maximum input frames per load session (no repetition)."
|
||||
)
|
||||
parser.add_argument("--max-client-rtf", type=float, help="Optional client stream-window RTF acceptance ceiling.")
|
||||
parser.add_argument("--server-revision", help="User-reported server commit/version; not verified by this client.")
|
||||
parser.add_argument(
|
||||
"--server-hardware", help="User-reported server hardware; not inferred from the client machine."
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
if args.sessions is not None:
|
||||
if not 1 <= args.sessions <= 64 or not 1 <= args.load_frames <= 10000:
|
||||
parser.error("load mode requires sessions in [1,64] and load-frames in [1,10000]")
|
||||
for name in ("timeout_s", "cleanup_timeout_s", "drain_s", "tail_s"):
|
||||
value = getattr(args, name)
|
||||
if not math.isfinite(value) or value < 0 or (name in ("timeout_s", "cleanup_timeout_s") and value == 0):
|
||||
parser.error(f"{name} must be finite and non-negative (timeout must be positive)")
|
||||
if args.max_client_rtf is not None and (not math.isfinite(args.max_client_rtf) or args.max_client_rtf <= 0):
|
||||
parser.error("max-client-rtf must be finite and positive")
|
||||
if args.max_frame_deficit < 0 or args.min_voiced_frames < 1 or args.minimum_audio_chunks < 1:
|
||||
parser.error("frame-deficit must be non-negative; voiced frames and audio chunks must be positive")
|
||||
if not math.isfinite(args.voiced_frame_rms_threshold) or args.voiced_frame_rms_threshold < 0:
|
||||
parser.error("voiced-frame-rms-threshold must be finite and non-negative")
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2))
|
||||
result = asyncio.run(run(parse_args()))
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False))
|
||||
if result.get("mode") == "paced_multi_session" and not result["ok"]:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,568 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import wave
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import TypedDict
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import websockets
|
||||
|
||||
from tests.e2e.online_serving import personaplex_realtime_duplex as driver
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _ServerState(TypedDict):
|
||||
index: int
|
||||
closed: bool
|
||||
frames: int
|
||||
close_requested: bool
|
||||
|
||||
|
||||
def _args(tmp_path, *extra):
|
||||
wav_path = tmp_path / "input.wav"
|
||||
with wave.open(str(wav_path), "wb") as stream:
|
||||
stream.setnchannels(1)
|
||||
stream.setsampwidth(2)
|
||||
stream.setframerate(driver.SAMPLE_RATE_HZ)
|
||||
stream.writeframes(np.full(driver.FRAME_SAMPLES * 2, 2000, dtype="<i2").tobytes())
|
||||
return driver.parse_args(
|
||||
[
|
||||
"--model",
|
||||
"fixture-only",
|
||||
"--input-wav",
|
||||
str(wav_path),
|
||||
"--output-dir",
|
||||
str(tmp_path / "out"),
|
||||
"--sessions",
|
||||
"2",
|
||||
"--load-frames",
|
||||
"2",
|
||||
"--tail-s",
|
||||
"0",
|
||||
"--drain-s",
|
||||
"0.04",
|
||||
"--timeout-s",
|
||||
"3",
|
||||
"--cleanup-timeout-s",
|
||||
"1",
|
||||
"--min-voiced-frames",
|
||||
"1",
|
||||
"--minimum-audio-chunks",
|
||||
"1",
|
||||
"--max-frame-deficit",
|
||||
"0",
|
||||
*extra,
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _audio(response_id="r", samples=1920, event_type="response.output_audio.delta"):
|
||||
return {
|
||||
"type": event_type,
|
||||
"response_id": response_id,
|
||||
"delta": base64.b64encode(np.full(samples, 2000, dtype="<i2").tobytes()).decode("ascii"),
|
||||
"sample_rate_hz": driver.SAMPLE_RATE_HZ,
|
||||
"metadata": {
|
||||
"vllm_omni": {
|
||||
"runtime_impl": "scheduler_data_plane",
|
||||
"uses_model_runner_scheduler": True,
|
||||
"runner_kv_backed": True,
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _server(mode="success"):
|
||||
states: list[_ServerState] = []
|
||||
|
||||
async def handle(ws):
|
||||
index = len(states)
|
||||
state: _ServerState = {"index": index, "closed": False, "frames": 0, "close_requested": False}
|
||||
states.append(state)
|
||||
response_id = "shared" if mode == "shared_response" else f"r-{index}"
|
||||
try:
|
||||
async for raw in ws:
|
||||
message = json.loads(raw)
|
||||
if message["type"] == "session.update":
|
||||
if mode == "reject_second" and index == 1:
|
||||
await ws.send(json.dumps({"type": "error", "error": {"code": "resource_exhausted"}}))
|
||||
continue
|
||||
if mode == "missing_capabilities":
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.created",
|
||||
"session": {"resume_token": "must-not-be-exported"},
|
||||
}
|
||||
)
|
||||
)
|
||||
continue
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.created",
|
||||
"session": {
|
||||
"resume_token": "must-not-be-exported",
|
||||
"capabilities": {
|
||||
"chunk_period_ms": 80,
|
||||
"supports_multi_session_same_replica": True,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
await ws.send(json.dumps({"type": "response.created", "response": {"id": response_id}}))
|
||||
elif message["type"] == "input_audio_buffer.append":
|
||||
state["frames"] += 1
|
||||
if mode == "short_audio" and state["frames"] == 2:
|
||||
continue
|
||||
if mode == "disconnect":
|
||||
await ws.close()
|
||||
return
|
||||
if mode == "bad_json":
|
||||
await ws.send("{")
|
||||
elif mode not in ("silent", "late_on_close"):
|
||||
event_type = (
|
||||
"response.audio.delta" if mode == "legacy_audio_event" else "response.output_audio.delta"
|
||||
)
|
||||
packet = _audio(response_id, event_type=event_type)
|
||||
if mode == "private_metadata":
|
||||
packet["metadata"] = {"vllm_omni": {"private": "remote-detail-must-not-be-exported"}}
|
||||
if mode == "odd_packet":
|
||||
raw_audio = base64.b64decode(packet["delta"])
|
||||
raw_audio = raw_audio + b"\x00" if state["frames"] == 1 else raw_audio[:-1]
|
||||
packet["delta"] = base64.b64encode(raw_audio).decode("ascii")
|
||||
if state["frames"] == 2:
|
||||
if mode == "missing_packet_rate":
|
||||
packet.pop("sample_rate_hz")
|
||||
elif mode == "string_packet_rate":
|
||||
packet["sample_rate_hz"] = "24000"
|
||||
elif mode == "invalid_base64":
|
||||
packet["delta"] = "!" + packet["delta"]
|
||||
elif mode == "orphan_audio":
|
||||
orphan = _audio(response_id)
|
||||
orphan.pop("response_id")
|
||||
await ws.send(json.dumps(orphan))
|
||||
await ws.send(json.dumps(packet))
|
||||
elif message["type"] == "session.close":
|
||||
state["close_requested"] = True
|
||||
if mode == "late_on_close":
|
||||
await ws.send(json.dumps(_audio(response_id, samples=driver.FRAME_SAMPLES * state["frames"])))
|
||||
if mode == "error_on_close":
|
||||
await ws.send(json.dumps({"type": "error", "error": {"code": "cleanup_failed"}}))
|
||||
if mode != "no_close_ack":
|
||||
await ws.send(json.dumps({"type": "session.closed"}))
|
||||
finally:
|
||||
state["closed"] = True
|
||||
|
||||
async with websockets.serve(handle, "127.0.0.1", 0) as server:
|
||||
port = server.sockets[0].getsockname()[1]
|
||||
yield f"ws://127.0.0.1:{port}/v1/realtime", states
|
||||
|
||||
|
||||
def test_omitted_sessions_preserves_lifecycle_mode():
|
||||
args = driver.parse_args(["--model", "unused", "--input-wav", "unused.wav"])
|
||||
assert args.sessions is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"options",
|
||||
[
|
||||
["--sessions", "0"],
|
||||
["--sessions", "65"],
|
||||
["--load-frames", "0"],
|
||||
["--load-frames", "10001"],
|
||||
["--drain-s", "nan"],
|
||||
["--tail-s", "-1"],
|
||||
["--timeout-s", "0"],
|
||||
["--timeout-s", "inf"],
|
||||
["--cleanup-timeout-s", "0"],
|
||||
["--cleanup-timeout-s", "nan"],
|
||||
["--max-client-rtf", "-1"],
|
||||
["--max-frame-deficit", "-1"],
|
||||
["--minimum-audio-chunks", "0"],
|
||||
["--min-voiced-frames", "0"],
|
||||
["--voiced-frame-rms-threshold", "nan"],
|
||||
],
|
||||
)
|
||||
def test_load_rejects_invalid_configuration(tmp_path, options):
|
||||
with pytest.raises(SystemExit):
|
||||
_args(tmp_path, *options)
|
||||
|
||||
|
||||
def test_metric_clock_and_packet_frame_distinction():
|
||||
client = driver.RawRealtimeProbe("ws://unused")
|
||||
client.events.add({"type": "session.created", "session": {"resume_token": "secret"}}, received_at_s=7)
|
||||
for when in (10.4, 10.8, 11.6):
|
||||
event = _audio(samples=driver.FRAME_SAMPLES * 5)
|
||||
event["_client_received_at_s"] = -999.0
|
||||
client.events.add(event, received_at_s=when)
|
||||
report = driver._load_metrics(client, [(10, 10, 10.01)])
|
||||
intervals = report["client_audio_packet_interval_ms"]
|
||||
assert isinstance(intervals, dict)
|
||||
assert intervals["count"] == 2
|
||||
assert intervals["median"] == pytest.approx(600)
|
||||
assert intervals["p99"] == pytest.approx(796)
|
||||
assert intervals["max"] == pytest.approx(800)
|
||||
assert intervals["argmax_index"] == 1
|
||||
assert report["client_first_audio_after_stream_start_ms"] == pytest.approx(400)
|
||||
assert report["output_samples"] == 15 * driver.FRAME_SAMPLES
|
||||
assert report["audio_underrun_s"] == pytest.approx(0.4)
|
||||
assert report["audio_continuity_ok"] is False
|
||||
underruns = report["audio_underrun_event_count"]
|
||||
assert isinstance(underruns, int) and underruns >= 1
|
||||
assert report["client_pacing_warning"] is False
|
||||
assert "secret" not in json.dumps(report)
|
||||
assert "-999" not in json.dumps(report)
|
||||
|
||||
|
||||
def test_empty_output_has_missing_not_zero_latency():
|
||||
report = driver._load_metrics(driver.RawRealtimeProbe("ws://unused"), [])
|
||||
assert report["client_stream_rtf"] is None
|
||||
assert report["audio_continuity_ok"] is None
|
||||
assert report["audio_underrun_s"] is None
|
||||
intervals = report["client_audio_packet_interval_ms"]
|
||||
assert isinstance(intervals, dict)
|
||||
assert intervals["p99"] is None
|
||||
|
||||
|
||||
def test_failed_acceptance_retains_frame_counts_without_remote_data():
|
||||
client = driver.RawRealtimeProbe("ws://unused")
|
||||
client.events.add(_audio(), received_at_s=1.0)
|
||||
args = driver.parse_args(
|
||||
["--model", "unused", "--input-wav", "unused.wav", "--max-frame-deficit", "0", "--min-voiced-frames", "1"]
|
||||
)
|
||||
with pytest.raises(AssertionError) as raised:
|
||||
driver._session_result(client, input_frames=2, args=args, minimum_chunks=1)
|
||||
assert getattr(raised.value, "check", None) == "frame_deficit"
|
||||
stats = getattr(raised.value, "stats", {})
|
||||
assert stats["frame_deficit"] == 1 and type(stats["frame_deficit"]) is int
|
||||
assert stats["output_frames"] == stats["voiced_frames"] == 1
|
||||
assert stats["frame_coverage_ratio"] == 0.5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("event_type", sorted(driver.AUDIO_DELTA_EVENT_TYPES))
|
||||
def test_session_result_accepts_current_and_legacy_audio_event_names(tmp_path, event_type):
|
||||
client = driver.RawRealtimeProbe("ws://unused")
|
||||
client.events.add(_audio(samples=driver.FRAME_SAMPLES, event_type=event_type), received_at_s=1.0)
|
||||
raw, response_ids, stats = driver._session_result(
|
||||
client,
|
||||
input_frames=1,
|
||||
args=_args(tmp_path, "--sessions", "1"),
|
||||
minimum_chunks=1,
|
||||
)
|
||||
assert len(raw) == driver.FRAME_SAMPLES * 2
|
||||
assert response_ids == {"r"}
|
||||
assert stats["frame_deficit"] == 0
|
||||
|
||||
|
||||
def test_legacy_audio_event_name_contributes_load_metrics():
|
||||
client = driver.RawRealtimeProbe("ws://unused")
|
||||
client.events.add(_audio(samples=driver.FRAME_SAMPLES * 5, event_type="response.audio.delta"), received_at_s=10.4)
|
||||
report = driver._load_metrics(client, [(10.0, 10.0, 10.01)])
|
||||
assert report["output_samples"] == driver.FRAME_SAMPLES * 5
|
||||
assert report["response_ids"] == ["r"]
|
||||
assert report["invalid_audio_packets"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sessions", [1, 2, 4])
|
||||
async def test_real_websocket_load_success(tmp_path, sessions):
|
||||
args = _args(tmp_path, "--sessions", str(sessions))
|
||||
async with _server() as (url, states):
|
||||
args.url = url
|
||||
result = await driver.run(args)
|
||||
assert result["ok"] is True
|
||||
assert result["passed_sessions"] == sessions
|
||||
rows = result["sessions"]
|
||||
assert isinstance(rows, list)
|
||||
assert len(rows) == sessions
|
||||
assert all(row["frame_deficit"] == 0 for row in rows)
|
||||
for row in rows:
|
||||
assert type(row["frame_deficit"]) is int
|
||||
assert row["output_frames"] == row["voiced_frames"] == row["audio_chunks"] == 2
|
||||
assert row["silent_frames"] == 0
|
||||
assert row["frame_coverage_ratio"] == 1.0
|
||||
assert row["audio_rms"] == pytest.approx(2000 / 32768)
|
||||
assert row["acceptance_check"] is None
|
||||
assert isinstance(row["audio_continuity_ok"], bool)
|
||||
assert all(s["closed"] and s["close_requested"] for s in states)
|
||||
assert "must-not-be-exported" not in json.dumps(result)
|
||||
saved = json.loads((tmp_path / "out/load-result.json").read_text())
|
||||
assert saved == result
|
||||
origins = [row["input_timeline"][0]["planned_at_s"] for row in rows]
|
||||
assert len(set(origins)) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_websocket_load_accepts_legacy_audio_event_name(tmp_path):
|
||||
args = _args(tmp_path, "--sessions", "2")
|
||||
async with _server("legacy_audio_event") as (url, states):
|
||||
args.url = url
|
||||
result = await driver.run(args)
|
||||
assert result["ok"] is True
|
||||
assert result["passed_sessions"] == 2
|
||||
assert result["shared_response_ids"] == []
|
||||
assert all(s["closed"] and s["close_requested"] for s in states)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"mode",
|
||||
[
|
||||
"reject_second",
|
||||
"silent",
|
||||
"disconnect",
|
||||
"no_close_ack",
|
||||
"shared_response",
|
||||
"late_on_close",
|
||||
"error_on_close",
|
||||
"bad_json",
|
||||
"missing_capabilities",
|
||||
],
|
||||
)
|
||||
async def test_real_websocket_failures_remain_in_results(tmp_path, mode):
|
||||
args = _args(tmp_path)
|
||||
async with _server(mode) as (url, states):
|
||||
args.url = url
|
||||
result = await asyncio.wait_for(driver.run(args), timeout=4)
|
||||
assert result["ok"] is False
|
||||
assert "must-not-be-exported" not in json.dumps(result)
|
||||
assert result["requested_sessions"] == 2
|
||||
rows = result["sessions"]
|
||||
assert isinstance(rows, list)
|
||||
assert len(rows) == 2
|
||||
assert all(s["closed"] for s in states)
|
||||
if mode == "reject_second":
|
||||
assert result["passed_sessions"] == 1
|
||||
rejected = next(row for row in rows if not row["ok"])
|
||||
assert rejected["input_frames"] == 0
|
||||
assert rejected["client_stream_rtf"] is None
|
||||
assert rejected["server_error_codes"] == ["resource_exhausted"]
|
||||
if mode == "silent":
|
||||
# A failed measurement must still release its admitted session.
|
||||
assert all(s["close_requested"] for s in states)
|
||||
if mode == "late_on_close":
|
||||
assert all(row["output_samples"] == 0 for row in rows)
|
||||
if mode == "shared_response":
|
||||
assert result["passed_sessions"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_existing_report_is_not_overwritten(tmp_path):
|
||||
args = _args(tmp_path)
|
||||
path = tmp_path / "out/load-result.json"
|
||||
path.parent.mkdir()
|
||||
path.write_text("previous measurement")
|
||||
with pytest.raises(FileExistsError):
|
||||
await driver.run(args)
|
||||
assert path.read_text() == "previous measurement"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["success", "no_close_ack"])
|
||||
async def test_cancellation_closes_all_connections(tmp_path, mode):
|
||||
args = _args(tmp_path, "--drain-s", "10", "--timeout-s", "30")
|
||||
async with _server(mode) as (url, states):
|
||||
args.url = url
|
||||
task = asyncio.create_task(driver.run(args))
|
||||
|
||||
async def started():
|
||||
while len(states) < 2 or not all(state["frames"] for state in states):
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
await asyncio.wait_for(started(), timeout=3)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=3)
|
||||
assert all(s["closed"] for s in states)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_rtf_ceiling_marks_slow_sessions_failed(tmp_path):
|
||||
args = _args(tmp_path, "--max-client-rtf", "0.001")
|
||||
async with _server() as (url, _):
|
||||
args.url = url
|
||||
result = await driver.run(args)
|
||||
assert result["passed_sessions"] == 0
|
||||
rows = result["sessions"]
|
||||
assert isinstance(rows, list)
|
||||
assert all("client_stream_rtf exceeds" in row["error"] for row in rows)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stalled_send_does_not_cause_catchup_burst():
|
||||
class Clock:
|
||||
now = 10.0
|
||||
|
||||
def monotonic(self):
|
||||
return self.now
|
||||
|
||||
async def sleep(self, delay):
|
||||
self.now += delay
|
||||
|
||||
class Client(driver.RawRealtimeProbe):
|
||||
async def send(self, event):
|
||||
if event["audio_end_ms"] == 80:
|
||||
clock.now += 0.25
|
||||
|
||||
clock = Clock()
|
||||
sends: list[tuple[float, float, float]] = []
|
||||
await driver._paced_load_frames(
|
||||
Client("ws://unused"),
|
||||
np.zeros(driver.FRAME_SAMPLES * 3, dtype="<f4"),
|
||||
epoch=10.0,
|
||||
timeout_s=1.0,
|
||||
sends=sends,
|
||||
clock=clock.monotonic,
|
||||
sleep=clock.sleep,
|
||||
)
|
||||
assert len(sends) == 3
|
||||
assert sends[1][1] >= sends[0][2] + driver.FRAME_PERIOD_S
|
||||
assert sends[1][1] - sends[1][0] == pytest.approx(0.25)
|
||||
assert sends[2][1] >= sends[1][2] + driver.FRAME_PERIOD_S
|
||||
metrics = driver._load_metrics(Client("ws://unused"), sends)
|
||||
assert metrics["client_pacing_warning"] is True
|
||||
assert asyncio.sleep != clock.sleep
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_frame_acceptance_preserves_numeric_diagnostics(tmp_path):
|
||||
args = _args(tmp_path, "--sessions", "1")
|
||||
async with _server("short_audio") as (url, states):
|
||||
args.url = url
|
||||
result = await driver.run(args)
|
||||
rows = result["sessions"]
|
||||
assert isinstance(rows, list)
|
||||
row = rows[0]
|
||||
assert row["ok"] is False
|
||||
assert row["acceptance_check"] == "frame_deficit"
|
||||
assert row["frame_deficit"] == 1 and type(row["frame_deficit"]) is int
|
||||
assert row["output_frames"] == row["voiced_frames"] == row["audio_chunks"] == 1
|
||||
assert row["frame_coverage_ratio"] == 0.5
|
||||
assert row["audio_rms"] == pytest.approx(2000 / 32768)
|
||||
assert "input=2" in row["error"] and "maximum_deficit=0" in row["error"]
|
||||
assert all(state["closed"] and state["close_requested"] for state in states)
|
||||
assert json.loads((tmp_path / "out/load-result.json").read_text()) == result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acceptance_failure_does_not_export_remote_metadata(tmp_path):
|
||||
args = _args(tmp_path, "--sessions", "1")
|
||||
async with _server("private_metadata") as (url, _):
|
||||
args.url = url
|
||||
result = await driver.run(args)
|
||||
rows = result["sessions"]
|
||||
assert isinstance(rows, list)
|
||||
row = rows[0]
|
||||
assert row["ok"] is False
|
||||
assert row["acceptance_check"] == "scheduler_data_plane"
|
||||
assert "remote-detail" not in json.dumps(result)
|
||||
assert "must-not-be-exported" not in json.dumps(result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ok", [True, False])
|
||||
def test_load_main_exit_status(tmp_path, monkeypatch, ok):
|
||||
args = _args(tmp_path)
|
||||
|
||||
async def completed(_args):
|
||||
return {"mode": "paced_multi_session", "ok": ok}
|
||||
|
||||
monkeypatch.setattr(driver, "parse_args", lambda: args)
|
||||
monkeypatch.setattr(driver, "run", completed)
|
||||
if ok:
|
||||
driver.main()
|
||||
else:
|
||||
with pytest.raises(SystemExit) as exc:
|
||||
driver.main()
|
||||
assert exc.value.code == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handshake_rejection_is_retained_without_server_details(tmp_path, monkeypatch):
|
||||
from websockets.exceptions import InvalidHandshake
|
||||
|
||||
async def rejected(_self):
|
||||
raise InvalidHandshake("credentials-in-remote-handshake-detail")
|
||||
|
||||
monkeypatch.setattr(driver.RawRealtimeProbe, "__aenter__", rejected)
|
||||
result = await asyncio.wait_for(driver.run(_args(tmp_path)), timeout=3)
|
||||
assert result["ok"] is False
|
||||
rows = result["sessions"]
|
||||
assert isinstance(rows, list) and len(rows) == 2
|
||||
assert all(row["error"] == "admission: InvalidHandshake" for row in rows)
|
||||
assert "credentials-in-remote-handshake-detail" not in json.dumps(result)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sessions", [1, 2])
|
||||
@pytest.mark.parametrize(
|
||||
"mode",
|
||||
[
|
||||
"missing_packet_rate",
|
||||
"string_packet_rate",
|
||||
"invalid_base64",
|
||||
"orphan_audio",
|
||||
"odd_packet",
|
||||
],
|
||||
)
|
||||
async def test_load_rejects_invalid_individual_audio_packets(tmp_path, sessions, mode):
|
||||
# The concatenated PCM can still meet coverage/voicing thresholds, so the
|
||||
# validator must also check each packet before accepting the whole run.
|
||||
args = _args(tmp_path, "--sessions", str(sessions))
|
||||
async with _server(mode) as (url, states):
|
||||
args.url = url
|
||||
result = await asyncio.wait_for(driver.run(args), timeout=4)
|
||||
assert result["ok"] is False
|
||||
assert result["passed_sessions"] == 0
|
||||
rows = result["sessions"]
|
||||
assert isinstance(rows, list) and len(rows) == sessions
|
||||
assert all(row["error"] is not None for row in rows)
|
||||
expected_code = {
|
||||
"missing_packet_rate": "invalid_audio_sample_rate",
|
||||
"string_packet_rate": "invalid_audio_sample_rate",
|
||||
"invalid_base64": "invalid_audio_base64",
|
||||
"orphan_audio": "missing_audio_response_id",
|
||||
"odd_packet": "unaligned_pcm16_packet",
|
||||
}[mode]
|
||||
assert all({packet["code"] for packet in row["invalid_audio_packets"]} == {expected_code} for row in rows)
|
||||
assert all(
|
||||
sum(packet["samples"] for packet in row["audio_packet_timeline"]) == row["output_samples"] for row in rows
|
||||
)
|
||||
assert all(state["closed"] and state["close_requested"] for state in states)
|
||||
assert json.loads((tmp_path / "out/load-result.json").read_text()) == result
|
||||
|
||||
|
||||
def test_load_allows_empty_and_multi_frame_audio_packets():
|
||||
client = driver.RawRealtimeProbe("ws://unused")
|
||||
client.events.add(_audio(samples=0), received_at_s=1.0)
|
||||
client.events.add(_audio(samples=driver.FRAME_SAMPLES * 5), received_at_s=1.4)
|
||||
metrics = driver._load_metrics(client, [(1.0, 1.0, 1.001)])
|
||||
assert metrics["invalid_audio_packets"] == []
|
||||
assert metrics["output_samples"] == driver.FRAME_SAMPLES * 5
|
||||
packets = metrics["audio_packet_timeline"]
|
||||
assert isinstance(packets, list) and len(packets) == 1
|
||||
|
||||
|
||||
def test_load_packet_diagnostics_do_not_copy_remote_payload():
|
||||
client = driver.RawRealtimeProbe("ws://unused")
|
||||
event = _audio()
|
||||
event["delta"] = "private-remote-payload!"
|
||||
client.events.add(event, received_at_s=1.4)
|
||||
metrics = driver._load_metrics(client, [])
|
||||
assert metrics["output_samples"] == 0
|
||||
assert metrics["invalid_audio_packets"] == [
|
||||
{
|
||||
"event_index": 0,
|
||||
"received_at_s": 1.4,
|
||||
"code": "invalid_audio_base64",
|
||||
}
|
||||
]
|
||||
assert "private-remote-payload" not in json.dumps(metrics)
|
||||
@@ -8,14 +8,17 @@ import json
|
||||
import os
|
||||
import time
|
||||
import types
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import pytest
|
||||
from omegaconf import OmegaConf
|
||||
from vllm.v1.engine.utils import EngineZmqAddresses
|
||||
|
||||
from vllm_omni.config.omni_config import OmniStageRuntimeConfig
|
||||
from vllm_omni.diffusion.data import AttentionConfig
|
||||
from vllm_omni.engine import omni_engine_base as async_omni_engine_module
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
||||
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
|
||||
from vllm_omni.engine.stage_init_utils import (
|
||||
LogicalStageInitPlan,
|
||||
ReplicaInitPlan,
|
||||
@@ -29,6 +32,32 @@ from vllm_omni.engine.stage_runtime import StageRuntime
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeParallelConfig:
|
||||
enable_fault_tolerance: bool = False
|
||||
enable_elastic_ep: bool = False
|
||||
data_parallel_size: int = 1
|
||||
use_ray: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeVllmConfig:
|
||||
parallel_config: _FakeParallelConfig = field(default_factory=_FakeParallelConfig)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeRuntimeConfig:
|
||||
devices: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeStageConfig:
|
||||
stage_id: int
|
||||
stage_type: str
|
||||
engine_args: dict[str, object]
|
||||
runtime: _FakeRuntimeConfig
|
||||
|
||||
|
||||
def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch):
|
||||
env_key = "VLLM_OMNI_TEST_TYPED_STAGE_ENV"
|
||||
monkeypatch.delenv(env_key, raising=False)
|
||||
@@ -200,7 +229,7 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
|
||||
)
|
||||
|
||||
|
||||
def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch):
|
||||
def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):
|
||||
import vllm_omni.engine.omni_engine_base as engine_mod
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
@@ -218,6 +247,7 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
|
||||
engine._omni_lb_policy = "random"
|
||||
engine.request_queue = types.SimpleNamespace()
|
||||
engine._log_stats = True
|
||||
engine._client_config = engine_mod.OmniClientConfig(client_count=2, client_index=1, stage_addresses={})
|
||||
engine._parallel_stage_init = False
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
@@ -233,6 +263,7 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
|
||||
|
||||
assert captured["stage_init_timeout"] == 7
|
||||
assert captured["log_stats"] is True
|
||||
assert captured["client_config"] is engine._client_config
|
||||
|
||||
|
||||
def test_async_omni_engine_initialize_stages_retains_stage0_prompt_transform(monkeypatch):
|
||||
@@ -833,6 +864,289 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
|
||||
assert captured["client_log_stats"] is True
|
||||
|
||||
|
||||
def test_stage_runtime_attaches_external_llm_client_without_launching_engine(monkeypatch):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
client_addresses = {
|
||||
"input_address": "ipc://stage0-input-1",
|
||||
"output_address": "ipc://stage0-output-1",
|
||||
}
|
||||
client_config = {
|
||||
"client_count": 2,
|
||||
"client_index": 1,
|
||||
"stage_addresses": {0: {0: client_addresses}},
|
||||
}
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=1,
|
||||
async_chunk=False,
|
||||
client_config=client_config,
|
||||
)
|
||||
parallel_config = _FakeParallelConfig()
|
||||
vllm_config = _FakeVllmConfig(parallel_config=parallel_config)
|
||||
plan = _make_llm_plan(0, stage_id=0, vllm_config=vllm_config).replicas[0]
|
||||
captured: dict[str, object] = {}
|
||||
stage_client = object()
|
||||
|
||||
def _capture_client(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return stage_client
|
||||
|
||||
monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", _capture_client)
|
||||
monkeypatch.setattr(
|
||||
runtime_mod,
|
||||
"launch_stage_replica",
|
||||
lambda **_kwargs: pytest.fail("external client must not launch a stage engine"),
|
||||
)
|
||||
|
||||
assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client
|
||||
assert captured["client_addresses"] == client_addresses
|
||||
assert captured["client_count"] == 2
|
||||
assert captured["client_index"] == 1
|
||||
assert not hasattr(parallel_config, "_api_process_count")
|
||||
assert not hasattr(parallel_config, "_api_process_rank")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage_ids", [(0,), (0, 1)], ids=["single-stage", "multi-stage"])
|
||||
@pytest.mark.parametrize("client_count", [1, 2])
|
||||
def test_stage_runtime_launches_shared_engines_with_per_client_addresses(monkeypatch, stage_ids, client_count):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
runtime = _make_stage_runtime()
|
||||
stage_plans = []
|
||||
for stage_id in stage_ids:
|
||||
parallel_config = _FakeParallelConfig()
|
||||
plan = _make_llm_plan(
|
||||
stage_id,
|
||||
stage_id=stage_id,
|
||||
vllm_config=_FakeVllmConfig(parallel_config=parallel_config),
|
||||
)
|
||||
plan.replicas[0].engine_args_dict = {}
|
||||
stage_plans.append(plan)
|
||||
stage_plans[0].replicas[0].metadata.runtime_cfg = {"env": {"VLLM_OMNI_TEST_STAGE_RUNTIME_ENV": "enabled"}}
|
||||
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
|
||||
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda stage_id, _cfg: str(stage_id))
|
||||
monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
|
||||
|
||||
events: list[tuple[str, int]] = []
|
||||
captured_launch_kwargs: list[dict[str, object]] = []
|
||||
captured_launch_env: list[str | None] = []
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _fake_launch_stage_replica(**kwargs):
|
||||
stage_id = kwargs["stage_id"]
|
||||
events.append(("enter", stage_id))
|
||||
captured_launch_kwargs.append(kwargs)
|
||||
captured_launch_env.append(os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV"))
|
||||
yield StageReplicaResources(
|
||||
addresses=EngineZmqAddresses(
|
||||
inputs=[f"ipc://stage{stage_id}-input-{idx}" for idx in range(client_count)],
|
||||
outputs=[f"ipc://stage{stage_id}-output-{idx}" for idx in range(client_count)],
|
||||
),
|
||||
)
|
||||
events.append(("exit", stage_id))
|
||||
|
||||
monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
|
||||
|
||||
with runtime.launch_stage_engines(client_count) as launch:
|
||||
expected = (
|
||||
[("enter", stage_id) for stage_id in stage_ids]
|
||||
if client_count > 1
|
||||
else [(event, stage_id) for stage_id in stage_ids for event in ("enter", "exit")]
|
||||
)
|
||||
assert events == expected
|
||||
for client_index, config in enumerate(launch.client_configs):
|
||||
assert config["client_count"] == client_count
|
||||
assert config["client_index"] == client_index
|
||||
for stage_id in stage_ids:
|
||||
assert config["stage_addresses"][stage_id][0] == {
|
||||
"input_address": f"ipc://stage{stage_id}-input-{client_index}",
|
||||
"output_address": f"ipc://stage{stage_id}-output-{client_index}",
|
||||
}
|
||||
|
||||
assert events == (
|
||||
[(event, stage_id) for event in ("enter", "exit") for stage_id in stage_ids] if client_count > 1 else expected
|
||||
)
|
||||
assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_count")
|
||||
assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_rank")
|
||||
assert all(
|
||||
kwargs["watched_frontend_processes"] is (launch.watched_frontend_processes if client_count > 1 else None)
|
||||
for kwargs in captured_launch_kwargs
|
||||
)
|
||||
assert captured_launch_env == ["enabled" if stage_id == 0 else None for stage_id in stage_ids]
|
||||
assert os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client_count", [1, 2])
|
||||
def test_stage_runtime_overlapping_devices_acquire_real_locks_once(monkeypatch, tmp_path, client_count):
|
||||
import fcntl
|
||||
|
||||
import vllm_omni.engine.stage_init_utils as init_utils
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
runtime = _make_stage_runtime()
|
||||
plans = [_make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig()) for stage_id in (0, 1)]
|
||||
for stage_id, plan in enumerate(plans):
|
||||
plan.replicas[0].engine_args_dict = {"tensor_parallel_size": stage_id + 1}
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: plans)
|
||||
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda sid, _cfg: "0" if sid == 0 else "0,1")
|
||||
monkeypatch.setattr(
|
||||
init_utils, "device_init_lock_path", lambda device_id: str(tmp_path / f"device-{device_id}.lock")
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
init_utils.time, "sleep", lambda _: pytest.fail("Overlapping stages must not wait on their own lock")
|
||||
)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def launch(**kwargs):
|
||||
# Real flock on a separate descriptor proves the runtime holds each lock.
|
||||
for device_id in range(kwargs["stage_id"] + 1):
|
||||
with open(tmp_path / f"device-{device_id}.lock", "a") as probe:
|
||||
with pytest.raises(BlockingIOError):
|
||||
fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
yield StageReplicaResources(
|
||||
addresses=EngineZmqAddresses(
|
||||
inputs=[f"ipc://input-{i}" for i in range(client_count)],
|
||||
outputs=[f"ipc://output-{i}" for i in range(client_count)],
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
|
||||
with runtime.launch_stage_engines(client_count):
|
||||
pass
|
||||
for device_id in (0, 1):
|
||||
with open(tmp_path / f"device-{device_id}.lock", "a") as probe:
|
||||
fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
fcntl.flock(probe, fcntl.LOCK_UN)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("diffusion_stage_id", [0, 1], ids=["diffusion-only", "enginecore-to-diffusion"])
|
||||
def test_stage_runtime_multi_api_rejects_diffusion_before_launch(monkeypatch, diffusion_stage_id):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
runtime = _make_stage_runtime()
|
||||
stage_plans = [
|
||||
_make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig())
|
||||
for stage_id in range(diffusion_stage_id)
|
||||
]
|
||||
stage_plans.append(_make_diffusion_plan(diffusion_stage_id, stage_id=diffusion_stage_id))
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
|
||||
monkeypatch.setattr(
|
||||
runtime_mod,
|
||||
"acquire_device_locks",
|
||||
lambda *_args: pytest.fail("unsupported topology must fail before acquiring devices"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
runtime_mod,
|
||||
"launch_stage_replica",
|
||||
lambda **_kwargs: pytest.fail("unsupported topology must not partially launch EngineCore stages"),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="diffusion stage\\(s\\) are not supported"):
|
||||
with runtime.launch_stage_engines(2):
|
||||
pytest.fail("unsupported topology must not yield a launch handle")
|
||||
|
||||
|
||||
def test_stage_runtime_multi_api_maps_stage_devices_from_launcher_visibility(monkeypatch):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
runtime = _make_stage_runtime()
|
||||
stage_plan = _make_llm_plan(
|
||||
0,
|
||||
stage_id=0,
|
||||
vllm_config=_FakeVllmConfig(),
|
||||
)
|
||||
replica = stage_plan.replicas[0]
|
||||
replica.engine_args_dict = {}
|
||||
replica.metadata.runtime_cfg = {"devices": "0"}
|
||||
|
||||
device_env = runtime_mod.current_omni_platform.device_control_env_var
|
||||
monkeypatch.setenv(device_env, "5")
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
|
||||
monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
|
||||
captured_devices: list[str | None] = []
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _fake_launch_stage_replica(**kwargs):
|
||||
captured_devices.append(kwargs["stage_visible_devices"])
|
||||
yield StageReplicaResources(
|
||||
addresses=EngineZmqAddresses(
|
||||
inputs=["ipc://input-0", "ipc://input-1"],
|
||||
outputs=["ipc://output-0", "ipc://output-1"],
|
||||
),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
|
||||
|
||||
with runtime.launch_stage_engines(2):
|
||||
pass
|
||||
|
||||
assert captured_devices == ["5"]
|
||||
assert os.environ[device_env] == "5"
|
||||
|
||||
|
||||
def test_stage_runtime_multi_api_rejects_elastic_ep(monkeypatch):
|
||||
runtime = _make_stage_runtime()
|
||||
stage_plan = _make_llm_plan(
|
||||
0,
|
||||
stage_id=0,
|
||||
vllm_config=_FakeVllmConfig(parallel_config=_FakeParallelConfig(enable_elastic_ep=True)),
|
||||
)
|
||||
stage_plan.replicas[0].engine_args_dict = {}
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
|
||||
|
||||
with pytest.raises(ValueError, match="enable-elastic-ep"):
|
||||
with runtime.launch_stage_engines(2):
|
||||
pass
|
||||
|
||||
|
||||
def test_stage_runtime_multi_api_failure_shuts_down_before_exceptional_context_exit(monkeypatch):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
runtime = _make_stage_runtime()
|
||||
stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig())
|
||||
stage_plan.replicas[0].engine_args_dict = {}
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
|
||||
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_args: None)
|
||||
monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
|
||||
|
||||
events: list[str] = []
|
||||
|
||||
class _Manager:
|
||||
def shutdown(self) -> None:
|
||||
events.append("shutdown")
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _fake_launch_stage_replica(**_kwargs):
|
||||
try:
|
||||
yield StageReplicaResources(
|
||||
manager=_Manager(),
|
||||
addresses=EngineZmqAddresses(
|
||||
inputs=["tcp://127.0.0.1:0", "tcp://127.0.0.1:1"],
|
||||
outputs=["tcp://127.0.0.1:2", "tcp://127.0.0.1:3"],
|
||||
),
|
||||
)
|
||||
except RuntimeError as exc:
|
||||
events.append(f"exceptional-exit:{exc}")
|
||||
raise
|
||||
else:
|
||||
events.append("normal-exit")
|
||||
|
||||
monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
|
||||
|
||||
with pytest.raises(RuntimeError, match="deferred TCP addresses"):
|
||||
with runtime.launch_stage_engines(2):
|
||||
pass
|
||||
|
||||
assert events == [
|
||||
"shutdown",
|
||||
"exceptional-exit:Stage 0 returned deferred TCP addresses; multi-API launch requires fixed ports or IPC addresses",
|
||||
]
|
||||
|
||||
|
||||
def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
@@ -1079,7 +1393,7 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
|
||||
num_replicas=1,
|
||||
launch_mode="local",
|
||||
stage_cfg=types.SimpleNamespace(engine_args={}, runtime=types.SimpleNamespace(devices="0")),
|
||||
metadata=types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}),
|
||||
metadata=types.SimpleNamespace(stage_id=0, stage_type="llm", runtime_cfg={"devices": "0"}),
|
||||
stage_connector_spec={},
|
||||
omni_kv_connector=(None, None, None),
|
||||
stage_vllm_config=fake_vllm_config,
|
||||
@@ -1933,6 +2247,102 @@ def test_port_from_zmq_address_parsing():
|
||||
assert _port_from_zmq_address("tcp://host:not-a-port") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("parallel", [False, True])
|
||||
@pytest.mark.parametrize("failure", ["ready", "attach", None])
|
||||
def test_single_api_common_launch_ownership_and_rollback(monkeypatch, parallel, failure):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
runtime = _make_stage_runtime()
|
||||
runtime._parallel_stage_init = parallel
|
||||
plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig()).replicas[0]
|
||||
plan.engine_args_dict = {}
|
||||
events = []
|
||||
|
||||
class Manager:
|
||||
def shutdown(self):
|
||||
events.append("shutdown")
|
||||
|
||||
manager = Manager()
|
||||
addresses = EngineZmqAddresses(inputs=["ipc://input"], outputs=["ipc://output"])
|
||||
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0")
|
||||
|
||||
def acquire(*_):
|
||||
events.append("lock")
|
||||
return [42]
|
||||
|
||||
monkeypatch.setattr(runtime_mod, "acquire_device_locks", acquire)
|
||||
monkeypatch.setattr(runtime_mod, "release_device_locks", lambda _: events.append("unlock"))
|
||||
|
||||
@contextlib.contextmanager
|
||||
def launch(**kwargs):
|
||||
assert kwargs["num_api_servers"] == 1
|
||||
assert kwargs["omni_parallel_stage_init"] is parallel
|
||||
events.append("spawn")
|
||||
yield StageReplicaResources(manager=manager, addresses=addresses)
|
||||
assert runtime._replica_launch_lock.locked() is False
|
||||
events.append("ready")
|
||||
if failure == "ready":
|
||||
raise RuntimeError("ready")
|
||||
|
||||
def attach(**kwargs):
|
||||
assert "ready" in events
|
||||
assert kwargs["engine_manager"] is manager
|
||||
events.append("attach")
|
||||
if failure == "attach":
|
||||
raise RuntimeError("attach")
|
||||
return manager
|
||||
|
||||
monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
|
||||
monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", attach)
|
||||
if failure:
|
||||
with pytest.raises(RuntimeError, match=failure):
|
||||
runtime._initialize_local_llm_replica(plan, 1)
|
||||
assert events.count("shutdown") == 1
|
||||
else:
|
||||
assert runtime._initialize_local_llm_replica(plan, 1) is manager
|
||||
assert "shutdown" not in events
|
||||
assert events.count("lock") == events.count("unlock") == (0 if parallel else 1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client_count", [1, 2])
|
||||
def test_common_launch_parallel_admission_before_spawn(monkeypatch, client_count):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
runtime = _make_stage_runtime()
|
||||
runtime._parallel_stage_init = True
|
||||
plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig())
|
||||
plan.replicas[0].engine_args_dict = {}
|
||||
events = []
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [plan])
|
||||
monkeypatch.setattr(runtime, "_reject_unguardable_executors", lambda _: events.append("guard"))
|
||||
monkeypatch.setattr(runtime, "_run_stage_admission", lambda _: events.append("admit"))
|
||||
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0")
|
||||
|
||||
def forbidden_lock(*_):
|
||||
pytest.fail("parallel initialization must use child phase locks")
|
||||
|
||||
monkeypatch.setattr(runtime_mod, "acquire_device_locks", forbidden_lock)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def launch(**kwargs):
|
||||
assert events == ["guard", "admit"]
|
||||
assert kwargs["omni_parallel_stage_init"] is True
|
||||
events.append("spawn")
|
||||
yield StageReplicaResources(
|
||||
addresses=EngineZmqAddresses(
|
||||
inputs=[f"ipc://in{i}" for i in range(client_count)],
|
||||
outputs=[f"ipc://out{i}" for i in range(client_count)],
|
||||
)
|
||||
)
|
||||
assert not runtime._replica_launch_lock.locked()
|
||||
events.append("ready")
|
||||
|
||||
monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
|
||||
with runtime.launch_stage_engines(client_count):
|
||||
pass
|
||||
assert events == ["guard", "admit", "spawn", "ready"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"stage_modes,expected", [([False, True, True], True), ([False, False], False), ([True, False], True)]
|
||||
)
|
||||
|
||||
@@ -18,6 +18,7 @@ import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from vllm_omni.engine import stage_init_utils
|
||||
from vllm_omni.engine.stage_init_utils import (
|
||||
@@ -240,7 +241,7 @@ def test_build_vllm_config_fails_before_engine_config_on_mismatch():
|
||||
get_class.assert_not_called()
|
||||
|
||||
|
||||
def test_build_vllm_config_proceeds_on_consistent_layout():
|
||||
def test_build_vllm_config_proceeds_on_consistent_layout(mocker: MockerFixture) -> None:
|
||||
"""The guard must not false-positive: a consistent single-GPU stage flows past
|
||||
it and reaches ``create_engine_config`` / ``Executor.get_class`` as usual."""
|
||||
stage = _stage(1, devices="1")
|
||||
@@ -249,18 +250,30 @@ def test_build_vllm_config_proceeds_on_consistent_layout():
|
||||
model_config=types.SimpleNamespace(hf_config=types.SimpleNamespace()),
|
||||
)
|
||||
sentinel_executor = object()
|
||||
with (
|
||||
mock.patch.object(
|
||||
stage_init_utils.OmniEngineArgs, "create_engine_config", return_value=fake_config
|
||||
) as create_engine_config,
|
||||
mock.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor),
|
||||
mock.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant),
|
||||
):
|
||||
vllm_config, executor_class = build_vllm_config(
|
||||
stage,
|
||||
model="dummy-model",
|
||||
engine_args_dict={"tensor_parallel_size": 1},
|
||||
)
|
||||
captured_api_process_config: dict[str, int] = {}
|
||||
|
||||
def capture_create_engine_config(engine_args, *args, **kwargs):
|
||||
captured_api_process_config["count"] = engine_args._api_process_count
|
||||
captured_api_process_config["rank"] = engine_args._api_process_rank
|
||||
return fake_config
|
||||
|
||||
create_engine_config = mocker.patch.object(
|
||||
stage_init_utils.OmniEngineArgs,
|
||||
"create_engine_config",
|
||||
autospec=True,
|
||||
side_effect=capture_create_engine_config,
|
||||
)
|
||||
mocker.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor)
|
||||
mocker.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant)
|
||||
|
||||
vllm_config, executor_class = build_vllm_config(
|
||||
stage,
|
||||
model="dummy-model",
|
||||
engine_args_dict={"tensor_parallel_size": 1},
|
||||
api_process_count=3,
|
||||
api_process_rank=2,
|
||||
)
|
||||
create_engine_config.assert_called_once()
|
||||
assert vllm_config is fake_config
|
||||
assert executor_class is sentinel_executor
|
||||
assert captured_api_process_config == {"count": 3, "rank": 2}
|
||||
@@ -42,7 +42,10 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import signal
|
||||
from argparse import Namespace
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -59,6 +62,66 @@ from vllm_omni.entrypoints.serve.utils import errors as serve_errors
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeParallelConfig:
|
||||
_api_process_rank: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeVllmConfig:
|
||||
parallel_config: _FakeParallelConfig
|
||||
|
||||
|
||||
def test_omni_api_worker_sets_parent_death_signal_before_validation(mocker):
|
||||
args = _minimal_args(_omni_stage_client_configs=[{"stage_addresses": {}}])
|
||||
death_signal = mocker.patch.object(api_server, "set_death_signal")
|
||||
|
||||
with pytest.raises(RuntimeError, match="client configuration"):
|
||||
api_server.run_omni_api_server_worker_proc(
|
||||
"127.0.0.1:8000",
|
||||
mocker.Mock(),
|
||||
args,
|
||||
{"client_count": 2, "client_index": -1},
|
||||
)
|
||||
|
||||
death_signal.assert_called_once_with(signal.SIGTERM)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_omni_run_server_worker_forwards_client_config(monkeypatch):
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_build_async_omni(*_args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
yield _FakeEngineClient(vllm_config=_FakeVllmConfig(_FakeParallelConfig(_api_process_rank=1)))
|
||||
|
||||
monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
|
||||
monkeypatch.setattr(api_server, "build_openai_app", lambda *_args: FastAPI())
|
||||
monkeypatch.setattr(api_server, "omni_init_app_state", lambda *_args: asyncio.sleep(0))
|
||||
monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_args: None)
|
||||
|
||||
async def fake_storage_start() -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start)
|
||||
|
||||
async def fake_serve_http(app, **_kwargs):
|
||||
assert app.state.api_server_count == 2
|
||||
task = asyncio.create_task(asyncio.sleep(0))
|
||||
await asyncio.sleep(0)
|
||||
return task
|
||||
|
||||
monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
|
||||
|
||||
args = _minimal_args(model="dummy")
|
||||
config = {"client_count": 2, "client_index": 1, "stage_addresses": {}}
|
||||
|
||||
await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), args, config)
|
||||
assert captured["client_config"] is config
|
||||
|
||||
|
||||
# FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
|
||||
# would invent ("HEAD", path) entries for every GET and fail the current census.
|
||||
# TODO: extend this set and the expected census if @router.head / @router.options appear.
|
||||
@@ -626,6 +689,48 @@ async def test_realtime_route_defaults_to_configured_duplex_handler(
|
||||
assert calls == [expected_handler]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime", "/v1/realtime?duplex=1"])
|
||||
def test_multi_api_duplex_reconnect_rejected_on_each_worker(path: str) -> None:
|
||||
class SessionHandler:
|
||||
async def handle_realtime_session(self, _websocket) -> None:
|
||||
pytest.fail("Multi-API requests must not access process-local sessions")
|
||||
|
||||
class PendingWarmup:
|
||||
def is_set(self) -> bool:
|
||||
raise AssertionError("Multi-API rejection must not wait for duplex warmup")
|
||||
|
||||
# Separate apps represent workers with independent frontend registries.
|
||||
# Opening on A then reconnecting to B must reject at the same boundary.
|
||||
for worker_index in (0, 1):
|
||||
app = FastAPI()
|
||||
app.state.api_server_count = 2
|
||||
app.state.api_server_index = worker_index
|
||||
app.state.openai_serving_duplex = SessionHandler()
|
||||
app.state.duplex_warmup_done = PendingWarmup()
|
||||
app.include_router(api_server.router)
|
||||
with TestClient(app) as client, client.websocket_connect(path) as websocket:
|
||||
assert websocket.receive_json()["code"] == "multi_api_duplex_unsupported"
|
||||
with pytest.raises(WebSocketDisconnect) as exc:
|
||||
websocket.receive_text()
|
||||
assert exc.value.code == 1008
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime?duplex=1"])
|
||||
def test_single_api_duplex_session_still_reaches_handler(path: str) -> None:
|
||||
class SessionHandler:
|
||||
async def handle_realtime_session(self, websocket) -> None:
|
||||
await websocket.accept()
|
||||
await websocket.send_json({"type": "session.created"})
|
||||
await websocket.close()
|
||||
|
||||
app = FastAPI()
|
||||
app.state.api_server_count = 1
|
||||
app.state.openai_serving_duplex = SessionHandler()
|
||||
app.include_router(api_server.router)
|
||||
with TestClient(app) as client, client.websocket_connect(path) as websocket:
|
||||
assert websocket.receive_json() == {"type": "session.created"}
|
||||
|
||||
|
||||
def test_health_without_engine_returns_stable_unhealthy_response() -> None:
|
||||
"""Lock ``/health`` with no engine initialized.
|
||||
|
||||
@@ -715,6 +820,38 @@ def test_speech_without_handler_preserves_not_found_http_error() -> None:
|
||||
assert exc_info.value.detail == "The model does not support Speech API"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multi_api_rejects_runtime_voice_upload() -> None:
|
||||
app = FastAPI()
|
||||
app.state.api_server_count = 2
|
||||
raw_request = _request_for(app, method="POST", path="/v1/audio/voices")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await api_server.upload_voice(
|
||||
raw_request,
|
||||
audio_sample=None,
|
||||
speaker_embedding=None,
|
||||
consent="consent-id",
|
||||
name="probe",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 409
|
||||
assert "process-local frontend state" in exc_info.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multi_api_rejects_sleep_route() -> None:
|
||||
app = FastAPI()
|
||||
app.state.api_server_count = 2
|
||||
raw_request = _request_for(app, method="POST", path="/v1/omni/sleep")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await api_server.omni_sleep(api_server.OmniSleepRequest(stage_ids=[0]), raw_request)
|
||||
|
||||
assert exc_info.value.status_code == 409
|
||||
assert "process-local frontend state" in exc_info.value.detail
|
||||
|
||||
|
||||
def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
|
||||
"""Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
|
||||
|
||||
@@ -937,3 +1074,57 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
|
||||
must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
|
||||
must_be_none=_MULTISTAGE_MUST_BE_NONE,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("count", [None, 0, -1, "2", True])
|
||||
def test_process_local_guard_rejects_unknown_topology(count):
|
||||
app = FastAPI()
|
||||
if count is not None:
|
||||
app.state.api_server_count = count
|
||||
# CLI args cannot substitute for missing/invalid actual runtime state.
|
||||
app.state.args = Namespace(api_server_count=1)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
api_server._reject_process_local_state_with_multiple_api_workers(_request_for(app), "Sleep")
|
||||
assert exc.value.status_code == 503
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"requested,config,expected",
|
||||
[(None, None, 1), (1, None, 1), (2, {"client_count": 2}, 2), (None, {"client_count": 2}, 2)],
|
||||
)
|
||||
def test_api_topology_resolved_from_launch(requested, config, expected):
|
||||
assert api_server._resolve_api_server_count(Namespace(api_server_count=requested), config) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("requested,config", [(2, None), (2, {}), (2, {"client_count": 1}), (2, {"client_count": "2"})])
|
||||
def test_api_topology_rejects_inconsistent_launch(requested, config):
|
||||
with pytest.raises(ValueError):
|
||||
api_server._resolve_api_server_count(Namespace(api_server_count=requested), config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("count,status", [(2, 409), (None, 503)])
|
||||
@pytest.mark.parametrize(
|
||||
"method,path",
|
||||
[
|
||||
("POST", "/v1/videos"),
|
||||
("GET", "/v1/videos"),
|
||||
("GET", "/v1/videos/job"),
|
||||
("GET", "/v1/videos/job/content"),
|
||||
("DELETE", "/v1/videos/job"),
|
||||
],
|
||||
)
|
||||
def test_video_lifecycle_guard_precedes_processing(monkeypatch, count, status, method, path):
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
app = FastAPI()
|
||||
if count is not None:
|
||||
app.state.api_server_count = count
|
||||
app.include_router(api_server.router)
|
||||
|
||||
async def forbidden_form():
|
||||
pytest.fail("video form processing must not run before topology rejection")
|
||||
|
||||
app.dependency_overrides[api_server._parse_video_form] = forbidden_form
|
||||
with TestClient(app) as client:
|
||||
response = client.request(method, path)
|
||||
assert response.status_code == status
|
||||
@@ -20,6 +20,7 @@ def _make_app(engine_client):
|
||||
from vllm_omni.entrypoints.openai.api_server import router
|
||||
|
||||
app = FastAPI()
|
||||
app.state.api_server_count = 1
|
||||
app.include_router(router)
|
||||
app.state.engine_client = engine_client
|
||||
app.state.sleeping_stages = set()
|
||||
@@ -253,6 +254,7 @@ def pure_diffusion_app(pure_diffusion_engine, mocker):
|
||||
)
|
||||
|
||||
app = FastAPI()
|
||||
app.state.api_server_count = 1
|
||||
app.include_router(router)
|
||||
|
||||
mocker.patch(
|
||||
|
||||
@@ -3826,6 +3826,7 @@ def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixt
|
||||
def _make_api_server_request(handler, *, method: str = "POST", path: str = "/v1/audio/voices") -> Request:
|
||||
app = FastAPI()
|
||||
app.state.openai_serving_speech = handler
|
||||
app.state.api_server_count = 1
|
||||
scope = {
|
||||
"type": "http",
|
||||
"app": app,
|
||||
|
||||
@@ -463,6 +463,7 @@ async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch)
|
||||
@pytest.fixture
|
||||
def test_client():
|
||||
app = FastAPI()
|
||||
app.state.api_server_count = 1
|
||||
app.include_router(router)
|
||||
app.state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion(
|
||||
diffusion_engine=FakeAsyncOmni(),
|
||||
@@ -1804,6 +1805,7 @@ def test_action_extraction_accepts_multimodal_actions_payload():
|
||||
|
||||
def test_missing_handler_returns_503():
|
||||
app = FastAPI()
|
||||
app.state.api_server_count = 1
|
||||
app.include_router(router)
|
||||
app.state.openai_serving_video = None
|
||||
client = TestClient(app)
|
||||
@@ -2790,6 +2792,7 @@ def test_cosmos3_control_upload_rejects_invalid_size(control_bytes, message, tes
|
||||
|
||||
def test_sync_missing_handler_returns_503():
|
||||
app = FastAPI()
|
||||
app.state.api_server_count = 1
|
||||
app.include_router(router)
|
||||
app.state.openai_serving_video = None
|
||||
client = TestClient(app)
|
||||
|
||||
@@ -6,14 +6,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from dataclasses import dataclass
|
||||
from multiprocessing.reduction import ForkingPickler
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.v1.engine.utils import EngineZmqAddresses
|
||||
|
||||
from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
|
||||
from vllm_omni.config.resolver import OmniConfigResolution
|
||||
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
|
||||
from vllm_omni.engine.stage_runtime import StageEngineLaunch
|
||||
from vllm_omni.entrypoints.cli.serve import (
|
||||
OmniServeCommand,
|
||||
_parse_stage_overrides,
|
||||
@@ -25,6 +30,31 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingName
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _FakeProcess:
|
||||
sentinel: str
|
||||
exitcode: int | None
|
||||
name: str
|
||||
pid: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeProcessOwner:
|
||||
processes: list[_FakeProcess]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeStageEngineArgs:
|
||||
async_chunk: bool
|
||||
enable_sleep_mode: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeStageConfig:
|
||||
stage_id: int
|
||||
engine_args: _FakeStageEngineArgs
|
||||
|
||||
|
||||
def _resolved(*stages: SimpleNamespace) -> OmniConfigResolution:
|
||||
return OmniConfigResolution(config_path="/fake/stages.yaml", stage_configs=tuple(stages))
|
||||
|
||||
@@ -44,6 +74,8 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
|
||||
explicit = args.get_explicit_kwargs_dict()
|
||||
assert args.get_explicit_kwargs_dict()
|
||||
assert not explicit["async_chunk"]
|
||||
assert (args.api_server_count or 1) == 1
|
||||
assert "api_server_count" not in explicit
|
||||
|
||||
|
||||
def _parse_serve_args(argv: list[str]) -> TrackingNamespace:
|
||||
@@ -169,6 +201,244 @@ def test_serve_parser_rejects_invalid_video_output_transport(value: str) -> None
|
||||
parser.parse_args(["serve", "fake-model", "--omni", "--video-output-transport", value])
|
||||
|
||||
|
||||
def test_tracking_namespace_is_picklable_for_spawned_api_workers() -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
subparsers = parser.add_subparsers(dest="subcommand")
|
||||
OmniServeCommand().subparser_init(subparsers)
|
||||
args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2"])
|
||||
args._omni_stage_client_configs = [{"stage_addresses": {0: {0: {"input_address": "ipc://input"}}}}]
|
||||
|
||||
restored = ForkingPickler.loads(ForkingPickler.dumps(args))
|
||||
|
||||
assert restored.api_server_count == 2
|
||||
assert restored.get_explicit_kwargs_dict()["api_server_count"] == 2
|
||||
assert restored._omni_stage_client_configs == args._omni_stage_client_configs
|
||||
|
||||
|
||||
def test_serve_validate_rejects_multiple_api_servers_for_diffusion(mocker: MockerFixture) -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
subparsers = parser.add_subparsers(dest="subcommand")
|
||||
cmd = OmniServeCommand()
|
||||
cmd.subparser_init(subparsers)
|
||||
args = parser.parse_args(["serve", "fake-diffusion-model", "--omni", "--api-server-count", "2"])
|
||||
|
||||
mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=True)
|
||||
validate = mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
|
||||
|
||||
with pytest.raises(ValueError, match="not supported for diffusion"):
|
||||
cmd.validate(args)
|
||||
|
||||
validate.assert_not_called()
|
||||
|
||||
|
||||
def test_serve_validate_rejects_sleep_mode_with_multiple_api_servers(mocker: MockerFixture) -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
subparsers = parser.add_subparsers(dest="subcommand")
|
||||
cmd = OmniServeCommand()
|
||||
cmd.subparser_init(subparsers)
|
||||
args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2", "--enable-sleep-mode"])
|
||||
|
||||
mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=False)
|
||||
mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
|
||||
|
||||
with pytest.raises(ValueError, match="enable-sleep-mode"):
|
||||
cmd.validate(args)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("typed", [False, True])
|
||||
@pytest.mark.parametrize("downstream_async_chunk", [False, True])
|
||||
def test_build_multi_api_stage_runtime_matches_current_constructor(
|
||||
mocker: MockerFixture, typed: bool, downstream_async_chunk: bool
|
||||
) -> None:
|
||||
from vllm_omni.entrypoints.cli import serve as serve_module
|
||||
|
||||
args = TrackingNamespace(
|
||||
argparse.Namespace(
|
||||
model="dummy-model",
|
||||
model_tag=None,
|
||||
stage_init_timeout=1,
|
||||
tokenizer=None,
|
||||
disable_log_stats=False,
|
||||
),
|
||||
frozenset({"model"}),
|
||||
)
|
||||
if typed:
|
||||
from vllm_omni.config.omni_config import VllmOmniARStageConfig
|
||||
from vllm_omni.config.stage_config import StagePipelineConfig
|
||||
|
||||
stages = [
|
||||
VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=i, model_stage="ar"))
|
||||
for i in range(2)
|
||||
]
|
||||
stages[1].connector_config.async_chunk = downstream_async_chunk
|
||||
else:
|
||||
stages = [
|
||||
_FakeStageConfig(
|
||||
stage_id=i,
|
||||
engine_args=_FakeStageEngineArgs(
|
||||
async_chunk=bool(i and downstream_async_chunk), enable_sleep_mode=False
|
||||
),
|
||||
)
|
||||
for i in range(2)
|
||||
]
|
||||
mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
|
||||
mocker.patch(
|
||||
"vllm_omni.config.config_factory.with_trust_remote_code_override",
|
||||
side_effect=lambda kwargs, _trust: kwargs,
|
||||
)
|
||||
mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={})
|
||||
mocker.patch(
|
||||
"vllm_omni.config.resolver.resolve_omni_config",
|
||||
return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=tuple(stages)),
|
||||
)
|
||||
|
||||
runtime = serve_module._build_multi_api_stage_runtime(args, 2)
|
||||
|
||||
assert runtime._client_count == 1
|
||||
assert runtime._client_index == 0
|
||||
assert runtime._async_chunk is downstream_async_chunk
|
||||
|
||||
|
||||
@pytest.mark.parametrize("typed", [False, True])
|
||||
def test_build_multi_api_stage_runtime_rejects_sleep_enabled_in_stage_config(
|
||||
mocker: MockerFixture, typed: bool
|
||||
) -> None:
|
||||
from vllm_omni.entrypoints.cli import serve as serve_module
|
||||
|
||||
args = TrackingNamespace(
|
||||
argparse.Namespace(
|
||||
model="dummy-model",
|
||||
model_tag=None,
|
||||
stage_init_timeout=1,
|
||||
tokenizer=None,
|
||||
disable_log_stats=False,
|
||||
),
|
||||
frozenset({"model"}),
|
||||
)
|
||||
stage_config = _FakeStageConfig(
|
||||
stage_id=3,
|
||||
engine_args=_FakeStageEngineArgs(async_chunk=False, enable_sleep_mode=True),
|
||||
)
|
||||
if typed:
|
||||
from vllm_omni.config.omni_config import VllmOmniARStageConfig
|
||||
from vllm_omni.config.stage_config import StagePipelineConfig
|
||||
|
||||
typed_stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=3, model_stage="ar"))
|
||||
typed_stage.model_config.enable_sleep_mode = True
|
||||
stage_config = typed_stage
|
||||
mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
|
||||
mocker.patch(
|
||||
"vllm_omni.config.config_factory.with_trust_remote_code_override",
|
||||
side_effect=lambda kwargs, _trust: kwargs,
|
||||
)
|
||||
mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={})
|
||||
mocker.patch(
|
||||
"vllm_omni.config.resolver.resolve_omni_config",
|
||||
return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=(stage_config,)),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=r"sleep mode.*stage\(s\) \[3\]"):
|
||||
serve_module._build_multi_api_stage_runtime(args, 2)
|
||||
|
||||
|
||||
def test_run_multi_api_server_omni_starts_workers_after_shared_engine_launch(mocker: MockerFixture) -> None:
|
||||
from vllm_omni.entrypoints.cli import serve as serve_module
|
||||
|
||||
args = TrackingNamespace(
|
||||
argparse.Namespace(api_server_count=2, shutdown_timeout=1),
|
||||
frozenset({"api_server_count", "shutdown_timeout"}),
|
||||
)
|
||||
socket = mocker.Mock()
|
||||
primary_addresses = EngineZmqAddresses(
|
||||
inputs=["ipc://input-0", "ipc://input-1"],
|
||||
outputs=["ipc://output-0", "ipc://output-1"],
|
||||
)
|
||||
engine_launch = StageEngineLaunch(
|
||||
client_configs=[{"client_count": 2, "client_index": index, "stage_addresses": {}} for index in range(2)],
|
||||
resources=[StageReplicaResources(addresses=primary_addresses)],
|
||||
)
|
||||
runtime = mocker.MagicMock()
|
||||
runtime.launch_stage_engines.return_value.__enter__.return_value = engine_launch
|
||||
runtime.launch_stage_engines.return_value.__exit__.return_value = False
|
||||
|
||||
manager = mocker.Mock()
|
||||
manager.processes = [mocker.Mock()]
|
||||
manager.gather_actual_addresses.return_value = (primary_addresses.inputs, primary_addresses.outputs)
|
||||
|
||||
mocker.patch.object(serve_module, "_build_multi_api_stage_runtime", return_value=runtime)
|
||||
mocker.patch.object(serve_module, "_wait_for_multi_api_server_completion")
|
||||
mocker.patch("signal.signal")
|
||||
mocker.patch("vllm.entrypoints.openai.api_server.setup_server", return_value=("127.0.0.1:8000", socket))
|
||||
mocker.patch("vllm.v1.metrics.prometheus.setup_multiprocess_prometheus")
|
||||
start_manager = mocker.patch.object(serve_module, "_start_api_server_process_manager", return_value=manager)
|
||||
|
||||
serve_module.run_multi_api_server_omni(args)
|
||||
|
||||
assert args._omni_stage_client_configs is engine_launch.client_configs
|
||||
start_manager.assert_called_once()
|
||||
manager_kwargs = start_manager.call_args.kwargs
|
||||
assert manager_kwargs["cleanup_timeout"] == 1
|
||||
assert manager_kwargs["num_servers"] == 2
|
||||
assert manager_kwargs["input_addresses"] == ["ipc://input-0", "ipc://input-1"]
|
||||
assert manager_kwargs["target_server_fn"] is serve_module.run_omni_api_server_worker_proc
|
||||
assert engine_launch.watched_frontend_processes == manager.processes
|
||||
runtime.shutdown.assert_called_once_with()
|
||||
socket.close.assert_called_once_with()
|
||||
|
||||
|
||||
def test_start_api_server_process_manager_cleans_up_partial_start(mocker: MockerFixture) -> None:
|
||||
from vllm_omni.entrypoints.cli import serve as serve_module
|
||||
|
||||
first_process = mocker.Mock(name="first_process")
|
||||
second_process = mocker.Mock(name="second_process")
|
||||
first_process.pid = 1234
|
||||
second_process.pid = None
|
||||
second_process.start.side_effect = RuntimeError("spawn failed")
|
||||
parent_pipes = [mocker.Mock(name="parent_pipe_0"), mocker.Mock(name="parent_pipe_1")]
|
||||
child_pipes = [mocker.Mock(name="child_pipe_0"), mocker.Mock(name="child_pipe_1")]
|
||||
spawn_context = mocker.Mock()
|
||||
spawn_context.Process.side_effect = [first_process, second_process]
|
||||
spawn_context.Pipe.side_effect = list(zip(parent_pipes, child_pipes))
|
||||
mocker.patch("vllm.v1.utils.multiprocessing.get_context", return_value=spawn_context)
|
||||
shutdown = mocker.patch("vllm.v1.utils.shutdown")
|
||||
|
||||
with pytest.raises(RuntimeError, match="spawn failed"):
|
||||
serve_module._start_api_server_process_manager(
|
||||
cleanup_timeout=1.5,
|
||||
listen_address="127.0.0.1:8000",
|
||||
sock=mocker.Mock(),
|
||||
args=argparse.Namespace(),
|
||||
num_servers=2,
|
||||
input_addresses=["ipc://input-0", "ipc://input-1"],
|
||||
output_addresses=["ipc://output-0", "ipc://output-1"],
|
||||
target_server_fn=mocker.Mock(),
|
||||
)
|
||||
|
||||
first_process.start.assert_called_once_with()
|
||||
second_process.start.assert_called_once_with()
|
||||
for pipe in parent_pipes:
|
||||
pipe.close.assert_called_once_with()
|
||||
shutdown.assert_called_once_with([first_process], timeout=1.5)
|
||||
|
||||
|
||||
def test_wait_for_multi_api_server_completion_rejects_engine_failure(
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
from vllm_omni.entrypoints.cli import serve as serve_module
|
||||
|
||||
api_process = _FakeProcess(sentinel="api-0", exitcode=None, name="api-0", pid=10)
|
||||
engine_process = _FakeProcess(sentinel="engine-0", exitcode=4, name="engine-0", pid=20)
|
||||
manager = _FakeProcessOwner(processes=[api_process])
|
||||
engine_launch = StageEngineLaunch(
|
||||
client_configs=[],
|
||||
resources=[StageReplicaResources(manager=_FakeProcessOwner(processes=[engine_process]))],
|
||||
)
|
||||
mocker.patch("multiprocessing.connection.wait", return_value=["engine-0"])
|
||||
|
||||
with pytest.raises(RuntimeError, match=r"Shared stage engine process engine-0 .* exited with code 4"):
|
||||
serve_module._wait_for_multi_api_server_completion(manager, engine_launch)
|
||||
|
||||
|
||||
def test_serve_parser_rejects_stage_configs_path() -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
subparsers = parser.add_subparsers(dest="subcommand")
|
||||
@@ -847,3 +1117,22 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
|
||||
|
||||
with pytest.raises(RuntimeError, match=r"exited with code 137"):
|
||||
run_headless(_make_headless_args(stage_id=1))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["unfiltered_ns", "explicit_keys", "missing", "__dict__"])
|
||||
def test_tracking_namespace_uninitialized_access_raises_attribute_error(name):
|
||||
namespace = object.__new__(TrackingNamespace)
|
||||
with pytest.raises(AttributeError):
|
||||
getattr(namespace, name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("deep", [False, True])
|
||||
def test_tracking_namespace_copy_preserves_tracking(deep):
|
||||
import copy
|
||||
|
||||
namespace = TrackingNamespace(
|
||||
argparse.Namespace(model="example", api_server_count=2), frozenset({"api_server_count"})
|
||||
)
|
||||
restored = copy.deepcopy(namespace) if deep else copy.copy(namespace)
|
||||
assert restored.get_explicit_kwargs_dict() == {"api_server_count": 2}
|
||||
assert restored.model == "example"
|
||||
@@ -603,3 +603,21 @@ class TestCumulativeStreamingCoercion:
|
||||
assert all([isinstance(r, SamplingParams) for r in result])
|
||||
assert result[0].output_kind == RequestOutputKind.DELTA
|
||||
assert result[1].output_kind == RequestOutputKind.DELTA
|
||||
|
||||
|
||||
@pytest.mark.parametrize("api_server_count", [1, 2])
|
||||
def test_prepare_stage_config_inputs_filters_frontend_options(api_server_count):
|
||||
from vllm_omni.entrypoints.utils import prepare_stage_config_inputs
|
||||
|
||||
kwargs = {
|
||||
"api_server_count": api_server_count,
|
||||
"disable_log_stats": True,
|
||||
"model_tag": "example",
|
||||
"max_num_seqs": 8,
|
||||
"stage_overrides": '{"0": {"max_num_seqs": 4}}',
|
||||
}
|
||||
inputs = prepare_stage_config_inputs("example", kwargs, trust_remote_code=None)
|
||||
|
||||
assert inputs.kwargs == {"max_num_seqs": 8}
|
||||
assert inputs.stage_overrides == {"0": {"max_num_seqs": 4}}
|
||||
assert kwargs["api_server_count"] == api_server_count
|
||||
@@ -9,6 +9,15 @@
|
||||
# Profile validated for the c64 / PROMPTS=512 performance experiments:
|
||||
# Stage 0 talker on GPU 0 with S0=64, Stage 1 Code2Wav on GPU 1 with S1=10.
|
||||
#
|
||||
# For CPU-bound frontend or orchestration workloads, this profile can be used
|
||||
# with ``--api-server-count 2`` or higher to run multiple API frontends over
|
||||
# one shared set of local stage engines. This can improve request parsing,
|
||||
# multimodal/audio preprocessing, and response handling under concurrency;
|
||||
# it does not duplicate GPU engines or make a single GPU-bound request faster.
|
||||
# Each frontend keeps process-local state, so runtime voice mutations, sleep /
|
||||
# wakeup, async video jobs, and resumable duplex sessions are not supported
|
||||
# with more than one API worker. See docs/cli/serve.md for the full boundary.
|
||||
#
|
||||
async_chunk: true
|
||||
|
||||
connectors:
|
||||
|
||||
@@ -24,6 +24,18 @@ def _launch_row_chunks(kernel, rows: int, device_type: str, *args, **kwargs) ->
|
||||
kernel[(chunk_rows,)](*args, row_offset, **kwargs)
|
||||
|
||||
|
||||
def _use_hopper_bf16_affine_semantics(x: torch.Tensor) -> bool:
|
||||
"""Use the unfused AdaLN dtype boundaries on Hopper.
|
||||
|
||||
The fused FP32 modulation path is stable on Blackwell, but the SM90
|
||||
kernel can accumulate enough extra precision in the repeated RMSNorm and
|
||||
AdaLN operations to move the 50-step I2VA result below its CI threshold.
|
||||
Keep the existing path on other architectures until their numerical
|
||||
baselines are independently validated.
|
||||
"""
|
||||
return x.device.type == "cuda" and torch.cuda.get_device_capability(x.device)[0] == 9
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _indexed_scale_shift_kernel(
|
||||
output_ptr,
|
||||
@@ -102,6 +114,7 @@ def _rms_norm_indexed_scale_shift_kernel(
|
||||
stride_scale_row,
|
||||
stride_indices,
|
||||
row_offset,
|
||||
use_bf16_affine: tl.constexpr,
|
||||
block_n: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0) + row_offset
|
||||
@@ -112,13 +125,20 @@ def _rms_norm_indexed_scale_shift_kernel(
|
||||
x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
variance = tl.sum(x * x, axis=0) / hidden_size
|
||||
normalized = x * tl.rsqrt(variance + eps) * weight
|
||||
|
||||
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
if use_bf16_affine:
|
||||
# Match the unfused RMSNorm output consumed by AdaLN on Hopper.
|
||||
normalized = (x * tl.rsqrt(variance + eps) * weight).to(tl.bfloat16)
|
||||
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0)
|
||||
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0)
|
||||
output = normalized * (scale + 1.0) + shift
|
||||
else:
|
||||
normalized = x * tl.rsqrt(variance + eps) * weight
|
||||
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
output = normalized * (1.0 + scale) + shift
|
||||
tl.store(
|
||||
output_ptr + row * hidden_size + columns,
|
||||
normalized * (1.0 + scale) + shift,
|
||||
output,
|
||||
mask=mask,
|
||||
)
|
||||
|
||||
@@ -143,6 +163,7 @@ def _indexed_gate_rms_norm_scale_shift_kernel(
|
||||
stride_scale_row,
|
||||
stride_indices,
|
||||
row_offset,
|
||||
use_bf16_affine: tl.constexpr,
|
||||
block_n: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0) + row_offset
|
||||
@@ -159,12 +180,20 @@ def _indexed_gate_rms_norm_scale_shift_kernel(
|
||||
|
||||
weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
variance = tl.sum(updated * updated, axis=0) / hidden_size
|
||||
normalized = updated * tl.rsqrt(variance + eps) * weight
|
||||
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
if use_bf16_affine:
|
||||
# Match the unfused RMSNorm output consumed by AdaLN on Hopper.
|
||||
normalized = (updated * tl.rsqrt(variance + eps) * weight).to(tl.bfloat16)
|
||||
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0)
|
||||
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0)
|
||||
output = normalized * (scale + 1.0) + shift
|
||||
else:
|
||||
normalized = updated * tl.rsqrt(variance + eps) * weight
|
||||
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
|
||||
output = normalized * (1.0 + scale) + shift
|
||||
tl.store(
|
||||
modulated_out_ptr + row * hidden_size + columns,
|
||||
normalized * (1.0 + scale) + shift,
|
||||
output,
|
||||
mask=mask,
|
||||
)
|
||||
|
||||
@@ -271,6 +300,7 @@ def rms_norm_indexed_scale_shift(
|
||||
shift.stride(0),
|
||||
scale.stride(0),
|
||||
indices.stride(0),
|
||||
use_bf16_affine=_use_hopper_bf16_affine_semantics(x),
|
||||
block_n=triton.next_power_of_2(hidden_size),
|
||||
num_warps=8,
|
||||
)
|
||||
@@ -324,6 +354,7 @@ def indexed_gate_rms_norm_scale_shift(
|
||||
shift.stride(0),
|
||||
scale.stride(0),
|
||||
indices.stride(0),
|
||||
use_bf16_affine=_use_hopper_bf16_affine_semantics(residual),
|
||||
block_n=triton.next_power_of_2(hidden_size),
|
||||
num_warps=8,
|
||||
)
|
||||
|
||||
@@ -102,6 +102,7 @@ _MODEL_LEVEL_TARGETS = {
|
||||
"time_embedder.linear_2": "time_embedder.proj_out",
|
||||
"norm_out.linear": "final_layer.adaln_proj.linear",
|
||||
"norm_out.norm": "final_layer.norm",
|
||||
"token_refiner.final_norm": "token_refiner.final_norm",
|
||||
}
|
||||
|
||||
# Per-block adapter suffix -> (native suffix, how a delta enters the native
|
||||
@@ -115,6 +116,8 @@ _BLOCK_TARGETS = {
|
||||
"attn.to_v": ("attn.qkv_proj", _QKV_V),
|
||||
"attn.to_out.0": ("attn.out_proj", _PLAIN),
|
||||
"attn.to_gate_compress": ("attn.to_gate_compress", _PLAIN),
|
||||
"attn.norm_q": ("attn.q_norm", _PLAIN),
|
||||
"attn.norm_k": ("attn.k_norm", _PLAIN),
|
||||
"ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES),
|
||||
"ff.net.2": ("mlp.fc2", _PLAIN),
|
||||
"adaln_proj.linear": ("adaln_proj.linear", _PLAIN),
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Sampling contract for FastVideo FastH3 full checkpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
||||
|
||||
from .fasth3 import FASTH3_BASE_MODEL, _resolve_dit_attention_backend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
FASTH3_V2_MODEL_ID = "FastVideo/FastVideo-FastH3-8-Step-V2"
|
||||
FASTH3_V2_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FastH3CheckpointSpec:
|
||||
"""The full V2 release has its own schedule and trained attention policy.
|
||||
|
||||
This is independent of ``FastH3WeightFusion``: the released weights already
|
||||
contain the entire student, including its learned compression gates.
|
||||
"""
|
||||
|
||||
vsa_sparsity: float = 0.8
|
||||
|
||||
@classmethod
|
||||
def from_metadata(cls, metadata: Mapping[str, object]) -> FastH3CheckpointSpec:
|
||||
"""Validate the release's own fastvideo_inference.json."""
|
||||
required = {
|
||||
"model_id": FASTH3_V2_MODEL_ID,
|
||||
"schema_version": "fasth3-inference-contract-v1",
|
||||
"dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125],
|
||||
"transformer_forwards": 8,
|
||||
"video_scheduler_shift": 10.0,
|
||||
"audio_scheduler_shift": 3.0,
|
||||
"guidance_scale": 1.0,
|
||||
"attention_backend": "VIDEO_SPARSE_ATTN_H3",
|
||||
"vsa_sparsity": 0.8,
|
||||
"vsa_tile_size": 64,
|
||||
"task": "t2av",
|
||||
}
|
||||
for key, expected in required.items():
|
||||
if metadata.get(key) != expected:
|
||||
raise ValueError(f"unsupported FastH3 contract: {key}={metadata.get(key)!r}, expected {expected!r}")
|
||||
return cls()
|
||||
|
||||
def release_metadata(self) -> dict[str, Any]:
|
||||
"""Express the sampling policy in the existing H3 pipeline schema."""
|
||||
return {
|
||||
"partition": "fl2va",
|
||||
"tasks": ["t2va"],
|
||||
"sigma_shift_scales": {"video": 10.0, "audio": 3.0},
|
||||
"base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule),
|
||||
}
|
||||
|
||||
def resolve_native_vaes(self, model_root: Path) -> Path:
|
||||
"""Reuse the frozen base VAEs with Omni's native tiled/parallel runtime.
|
||||
|
||||
Only VAE components are fetched; the student and text encoder come
|
||||
directly from the FastVideo release. The release pins the base revision.
|
||||
"""
|
||||
provenance = json.loads((model_root / "provenance.json").read_text(encoding="utf-8"))
|
||||
base = provenance.get("base_model", "")
|
||||
prefix = f"hf://{FASTH3_BASE_MODEL}@"
|
||||
if not isinstance(base, str) or not base.startswith(prefix) or not base[len(prefix) :]:
|
||||
raise ValueError("FastH3 V2 provenance must pin a MiniMaxAI/MiniMax-H3 base revision")
|
||||
return (
|
||||
Path(
|
||||
download_weights_from_hf_specific(
|
||||
model_name_or_path=FASTH3_BASE_MODEL,
|
||||
cache_dir=None,
|
||||
allow_patterns=["FL2VA/video_vae/**", "FL2VA/audio_vae/**"],
|
||||
revision=base[len(prefix) :],
|
||||
require_all=True,
|
||||
)
|
||||
)
|
||||
/ "FL2VA"
|
||||
)
|
||||
|
||||
def check_serving_contract(self, *, partition: str, od_config: OmniDiffusionConfig) -> None:
|
||||
if partition != "fl2va":
|
||||
raise ValueError("FastH3 V2 requires --task-type fl2va (T2VA requests only)")
|
||||
if getattr(od_config, "lora_path", None):
|
||||
raise ValueError("FastH3 V2 is a full checkpoint; additional LoRA adapters are unsupported")
|
||||
if _resolve_dit_attention_backend(od_config) != "FASTVIDEO_VSA":
|
||||
raise ValueError("FastH3 V2 requires --diffusion-attention-backend FASTVIDEO_VSA")
|
||||
attention_config = getattr(od_config, "diffusion_attention_config", None)
|
||||
per_role = getattr(attention_config, "per_role", None) or {}
|
||||
spec = per_role.get("self") or getattr(attention_config, "default", None)
|
||||
if getattr(spec, "fastvideo_vsa_topk", None) is not None:
|
||||
raise ValueError("FastH3 V2 pins VSA sparsity=0.8; remove the fixed fastvideo_vsa_topk override")
|
||||
parallel = getattr(od_config, "parallel_config", None)
|
||||
if any(int(getattr(parallel, key, 1) or 1) != 1 for key in ("ring_degree", "allgather_degree")):
|
||||
raise ValueError("FastH3 V2 supports local attention or pure Ulysses sequence parallelism")
|
||||
|
||||
def check_request(self, sampling: OmniDiffusionSamplingParams, *, step_execution: bool = False) -> None:
|
||||
if step_execution:
|
||||
# StepScheduler admits requests before the pipeline hook runs and
|
||||
# derives its lifetime from these fields. The pinned schedule must
|
||||
# therefore be explicit and cannot be overridden by custom arrays.
|
||||
if sampling.num_inference_steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps:
|
||||
raise OmniClientError("FastH3 V2 step execution requires num_inference_steps=8")
|
||||
if sampling.timesteps is not None or sampling.sigmas is not None:
|
||||
raise OmniClientError("FastH3 V2 step execution does not support custom timesteps or sigmas")
|
||||
if sampling.lora_request is not None:
|
||||
raise OmniClientError("FastH3 V2 does not support per-request LoRA adapters")
|
||||
steps = sampling.num_inference_steps
|
||||
if steps is not None and steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps:
|
||||
raise OmniClientError("FastH3 V2 requires num_inference_steps=8 (nine sigma points), or omitted")
|
||||
extra = sampling.extra_args or {}
|
||||
for key, expected in (("flow_shift", 10.0), ("audio_flow_shift", 3.0)):
|
||||
try:
|
||||
value = float(extra.get(key, expected))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}") from exc
|
||||
if not math.isclose(value, expected):
|
||||
raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}, got {value:g}")
|
||||
if sampling.guidance_scale is not None and sampling.guidance_scale != 1.0:
|
||||
raise OmniClientError("FastH3 V2 requires guidance_scale=1")
|
||||
@@ -55,6 +55,8 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||||
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
from .fasth3 import _resolve_native_target
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
@@ -118,6 +120,18 @@ class MiniMaxH3DiTArchConfig:
|
||||
|
||||
@classmethod
|
||||
def from_mapping(cls, config: Mapping[str, Any]) -> MiniMaxH3DiTArchConfig:
|
||||
# The modular Diffusers checkpoint uses these spellings for the same
|
||||
# architecture. Normalize before constructing the existing native DiT.
|
||||
aliases = {
|
||||
"num_refiner_layers": "token_refiner_num_layers",
|
||||
"ffn_dim": "ffn_hidden_size",
|
||||
"in_channels": "latents_dim",
|
||||
"audio_in_channels": "audio_latents_dim",
|
||||
"freq_dim": "timestep_input_dim",
|
||||
"time_embed_hidden_dim": "time_embed_hidden_size",
|
||||
"rope_freq_dim": "rope_inv_freq_len",
|
||||
}
|
||||
config = {aliases.get(name, name): value for name, value in config.items()}
|
||||
fields = cls.__dataclass_fields__
|
||||
values = {name: config[name] for name in fields if name in config}
|
||||
if "patch_size" in values:
|
||||
@@ -394,6 +408,9 @@ def _sdpa_varlen_attention(
|
||||
|
||||
|
||||
class MiniMaxH3Attention(nn.Module):
|
||||
# Full sparse checkpoints pin a ratio; legacy adapters use backend top-k.
|
||||
vsa_sparsity: float | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
arch: MiniMaxH3DiTArchConfig,
|
||||
@@ -588,6 +605,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
# (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other
|
||||
# backend/path.
|
||||
"laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE,
|
||||
**({"vsa_h3_sparsity": self.vsa_sparsity} if self.vsa_sparsity is not None else {}),
|
||||
# Present only for a VSA artifact; the VSA backend reads it as
|
||||
# the learned compression gate and every other backend ignores it.
|
||||
**({"gate_compress": gate_compress.unsqueeze(0)} if gate_compress is not None else {}),
|
||||
@@ -1127,10 +1145,23 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
self,
|
||||
od_config: OmniDiffusionConfig,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
*,
|
||||
diffusers_weights: bool | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
tf_config = od_config.tf_model_config
|
||||
config_mapping = tf_config.to_dict() if hasattr(tf_config, "to_dict") else dict(tf_config)
|
||||
# The native MiniMax-H3 Hub snapshot advertises the Diffusers
|
||||
# transformer class in its root config, while its FL2VA/Ref2VA
|
||||
# components still contain native weights. The pipeline has already
|
||||
# resolved the actual source format, so let it override the
|
||||
# class-name heuristic. Keep the heuristic for standalone callers.
|
||||
self._diffusers_weights = (
|
||||
config_mapping.get("_class_name") == "MiniMaxH3Transformer3DModel"
|
||||
if diffusers_weights is None
|
||||
else diffusers_weights
|
||||
)
|
||||
self._rope_theta = float(config_mapping.get("rope_theta", 10000.0))
|
||||
arch = MiniMaxH3DiTArchConfig.from_mapping(config_mapping)
|
||||
self.arch = arch
|
||||
self.od_config = od_config
|
||||
@@ -1211,7 +1242,7 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
)
|
||||
self._mark_missing_params_required()
|
||||
|
||||
def enable_vsa_gates(self) -> None:
|
||||
def enable_vsa_gates(self, *, sparsity: float | None = None) -> None:
|
||||
"""Give every DiT block's attention a VSA compression gate.
|
||||
|
||||
A FastH3 VSA artifact assigns these projections rather than adding to
|
||||
@@ -1223,6 +1254,9 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
return
|
||||
for block in self.blocks:
|
||||
block.attn.enable_vsa_gate()
|
||||
block.attn.vsa_sparsity = sparsity
|
||||
if sparsity is not None:
|
||||
block.attn.to_gate_compress.weight.missing_param_init = "error"
|
||||
self.vsa_gates_enabled = True
|
||||
|
||||
def _mark_missing_params_required(self) -> None:
|
||||
@@ -1293,17 +1327,37 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> set[str]:
|
||||
"""Load exact H3 checkpoint names with logical TP-aware loaders."""
|
||||
"""Load native or Diffusers H3 weights with the existing TP-aware loaders."""
|
||||
params = dict(self.named_parameters())
|
||||
params.update(dict(self.named_buffers()))
|
||||
loaded: set[str] = set()
|
||||
diffusers_weights = getattr(self, "_diffusers_weights", False)
|
||||
qkv_parts: dict[str, set[str]] = {}
|
||||
source_names: set[str] = set()
|
||||
for name, loaded_weight in weights:
|
||||
layout = "plain"
|
||||
if diffusers_weights:
|
||||
if name in source_names:
|
||||
raise ValueError(f"duplicate Diffusers H3 weight: {name}")
|
||||
source_names.add(name)
|
||||
module, _, kind = name.rpartition(".")
|
||||
target = _resolve_native_target(module)
|
||||
if target is None or kind not in {"weight", "bias"}:
|
||||
raise ValueError(f"unsupported Diffusers H3 weight: {name}")
|
||||
name, layout = f"{target[0]}.{kind}", target[1]
|
||||
param = params.get(name)
|
||||
if param is None:
|
||||
if diffusers_weights:
|
||||
raise ValueError(f"Diffusers H3 weight has no model parameter: {name}")
|
||||
logger.warning("Skipping MiniMax H3 weight not present in model: %s", name)
|
||||
continue
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
if name.endswith(".attn.qkv_proj.weight"):
|
||||
if layout in {"q", "k", "v"}:
|
||||
# vLLM can load each projection directly into its packed QKV
|
||||
# parameter, including TP slicing and online quantization.
|
||||
weight_loader(param, loaded_weight, layout)
|
||||
qkv_parts.setdefault(name, set()).add(layout)
|
||||
elif name.endswith(".attn.qkv_proj.weight"):
|
||||
# Transform checkpoint layout before entering vLLM's loader so
|
||||
# online FP8 can keep ``online_process_loader`` outermost.
|
||||
loaded_weight = _reorder_grouped_qkv_to_qkv(
|
||||
@@ -1319,12 +1373,26 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
"MiniMax H3 fc1 checkpoint rows must split evenly into "
|
||||
f"gate/up matrices, got {tuple(loaded_weight.shape)}"
|
||||
)
|
||||
gate, up = loaded_weight.chunk(2, dim=0)
|
||||
first, second = loaded_weight.chunk(2, dim=0)
|
||||
gate, up = (second, first) if layout == "swap_halves" else (first, second)
|
||||
weight_loader(param, gate, 0)
|
||||
weight_loader(param, up, 1)
|
||||
else:
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded.add(name)
|
||||
if diffusers_weights:
|
||||
for name, parts in qkv_parts.items():
|
||||
if parts != {"q", "k", "v"}:
|
||||
raise ValueError(f"incomplete Diffusers H3 QKV group {name}: {sorted(parts)}")
|
||||
# Diffusers reconstructs RoPE from config instead of storing this
|
||||
# native checkpoint buffer. Compute on CPU for identical values.
|
||||
freq_dim = self.arch.rope_inv_freq_len
|
||||
rope = 1.0 / (
|
||||
self._rope_theta
|
||||
** (torch.arange(0, 2 * freq_dim, 2, dtype=torch.float32, device="cpu") / (2 * freq_dim))
|
||||
)
|
||||
default_weight_loader(params["rope.inv_freq"], rope)
|
||||
loaded.add("rope.inv_freq")
|
||||
return loaded
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -67,6 +67,7 @@ from vllm_omni.model_executor.model_loader.weight_utils import (
|
||||
download_weights_from_hf_specific,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
|
||||
is_minimax_h3_modular,
|
||||
resolve_minimax_h3_partition,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
|
||||
@@ -106,6 +107,7 @@ from .denoise_loop import (
|
||||
)
|
||||
from .encoder import MiniMaxH3Qwen3VLEncoder
|
||||
from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion
|
||||
from .fasth3_checkpoint import FastH3CheckpointSpec
|
||||
from .lora import TurboSpec, load_minimax_h3_turbo_lora
|
||||
from .minimax_h3_transformer import (
|
||||
MiniMaxH3Attention,
|
||||
@@ -213,7 +215,11 @@ def _resolve_minimax_h3_model_root(
|
||||
if path.name in {"FL2VA", "Ref2VA"}:
|
||||
return path.parent
|
||||
return path
|
||||
if load_text_encoder:
|
||||
if is_minimax_h3_modular(model, revision):
|
||||
allow_patterns = ["modular_model_index.json", "fastvideo_inference.json", "provenance.json", "transformer/**"]
|
||||
if load_text_encoder:
|
||||
allow_patterns += ["text_encoder/**", "tokenizer/**", "processor/**"]
|
||||
elif load_text_encoder:
|
||||
allow_patterns = (
|
||||
MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition]
|
||||
)
|
||||
@@ -314,6 +320,8 @@ def resolve_minimax_h3_diffusion_model_path(
|
||||
partition,
|
||||
load_text_encoder=False,
|
||||
)
|
||||
if is_minimax_h3_modular(str(model_root), revision):
|
||||
return str(model_root)
|
||||
if partition == "combined":
|
||||
return str(model_root)
|
||||
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
|
||||
@@ -575,6 +583,7 @@ class MiniMaxH3Pipeline(
|
||||
_base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {}
|
||||
# Set from --lora-path during construction; absent means no FastH3 adapter.
|
||||
_fasth3: FastH3WeightFusion | None = None
|
||||
_fasth3_checkpoint: FastH3CheckpointSpec | None = None
|
||||
|
||||
def _load_diffusion_lora_adapter(
|
||||
self,
|
||||
@@ -800,10 +809,13 @@ class MiniMaxH3Pipeline(
|
||||
self._PROFILER_TARGETS.remove("encode_prompt")
|
||||
if not self.load_vae_encoder:
|
||||
self._PROFILER_TARGETS.remove("_encode_local_media")
|
||||
modular = is_minimax_h3_modular(str(od_config.model), od_config.revision)
|
||||
self.partition = _minimax_h3_partition_for_task(
|
||||
getattr(od_config, "task_type", None),
|
||||
str(od_config.model),
|
||||
)
|
||||
if modular and str(od_config.task_type or "auto").lower() == "auto":
|
||||
self.partition = "fl2va"
|
||||
self._turbo_lora_specs: dict[int, TurboSpec] = {}
|
||||
self._native_lora_adapter_ids: set[int] = set()
|
||||
self._lora_sigma_schedules: dict[int, DMD2SigmaSchedule] = {}
|
||||
@@ -813,9 +825,19 @@ class MiniMaxH3Pipeline(
|
||||
self.partition,
|
||||
load_text_encoder=self.load_text_encoder,
|
||||
)
|
||||
model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA")
|
||||
model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8"))
|
||||
release = model_index.get("_minimax_h3") or {}
|
||||
if modular:
|
||||
model_path = model_root
|
||||
self._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(
|
||||
json.loads((model_root / "fastvideo_inference.json").read_text(encoding="utf-8"))
|
||||
)
|
||||
self._fasth3_checkpoint.check_serving_contract(partition=self.partition, od_config=od_config)
|
||||
release = self._fasth3_checkpoint.release_metadata()
|
||||
vae_model_path = self._fasth3_checkpoint.resolve_native_vaes(model_root)
|
||||
else:
|
||||
model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA")
|
||||
model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8"))
|
||||
release = model_index.get("_minimax_h3") or {}
|
||||
vae_model_path = model_path
|
||||
partition = str(release.get("partition", "")).lower()
|
||||
expected_partition = "ref2va" if self.partition == "ref2va" else "fl2va"
|
||||
if partition != expected_partition:
|
||||
@@ -876,11 +898,18 @@ class MiniMaxH3Pipeline(
|
||||
self.transformer = MiniMaxH3DiTModel(
|
||||
od_config,
|
||||
quant_config=transformer_quant_config,
|
||||
diffusers_weights=modular,
|
||||
)
|
||||
if self._fasth3_checkpoint is not None:
|
||||
self.transformer.enable_vsa_gates(sparsity=self._fasth3_checkpoint.vsa_sparsity)
|
||||
logger.info(
|
||||
"FastH3 V2 full checkpoint: 8 transformer forwards, video/audio shifts 10/3, VSA sparsity=0.8 tile=64"
|
||||
)
|
||||
if ref2va_model_path is not None:
|
||||
self.transformers_ref = MiniMaxH3DiTModel(
|
||||
od_config,
|
||||
quant_config=transformer_quant_config,
|
||||
diffusers_weights=modular,
|
||||
)
|
||||
|
||||
self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer)
|
||||
@@ -963,14 +992,14 @@ class MiniMaxH3Pipeline(
|
||||
# so VAEs stay resident for new configurations.
|
||||
component_load_device = torch.device("cpu") if legacy_manual_components else self.device
|
||||
self.video_vae = MiniMaxH3VideoVAE(
|
||||
os.path.join(model_path, "video_vae"),
|
||||
os.path.join(vae_model_path, "video_vae"),
|
||||
device=self.device,
|
||||
load_device=component_load_device,
|
||||
decode_only=not self.load_vae_encoder,
|
||||
trust_remote_code=od_config.trust_remote_code,
|
||||
)
|
||||
self.audio_vae = MiniMaxH3AudioVAE(
|
||||
os.path.join(model_path, "audio_vae"),
|
||||
os.path.join(vae_model_path, "audio_vae"),
|
||||
device=self.device,
|
||||
load_device=component_load_device,
|
||||
decode_only=not self.load_vae_encoder,
|
||||
@@ -1046,6 +1075,12 @@ class MiniMaxH3Pipeline(
|
||||
# load_weights only warns on a parameter the model does not have, so
|
||||
# close the adapter against what the DiT actually consumed.
|
||||
self._fasth3.validate_fully_applied(transformer_loaded)
|
||||
if self._fasth3_checkpoint is not None:
|
||||
required_gates = {
|
||||
f"blocks.{i}.attn.to_gate_compress.weight" for i in range(self.transformer.arch.num_layers)
|
||||
}
|
||||
if missing_gates := required_gates - transformer_loaded:
|
||||
raise ValueError(f"FastH3 V2 checkpoint is missing compression gates: {sorted(missing_gates)}")
|
||||
return loaded_with_prefix
|
||||
|
||||
@property
|
||||
@@ -2056,6 +2091,11 @@ class MiniMaxH3Pipeline(
|
||||
self._validate_turbo_sampling(sampling, turbo_spec)
|
||||
if has_native_lora:
|
||||
self._validate_native_sampling(sampling, task=task)
|
||||
if self._fasth3_checkpoint is not None:
|
||||
self._fasth3_checkpoint.check_request(
|
||||
sampling,
|
||||
step_execution=bool(getattr(self.od_config, "step_execution", False)),
|
||||
)
|
||||
if self._fasth3 is not None:
|
||||
self._fasth3.check_request(
|
||||
sampling,
|
||||
|
||||
@@ -61,11 +61,12 @@ from vllm_omni.engine.stage_client import StageClient
|
||||
from vllm_omni.engine.stage_init_utils import build_stage0_input_processor
|
||||
from vllm_omni.engine.stage_pool import StagePool
|
||||
from vllm_omni.engine.stage_runtime import (
|
||||
OmniClientConfig,
|
||||
StageRuntimeInfo,
|
||||
create_stage_runtime,
|
||||
)
|
||||
from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
|
||||
from vllm_omni.entrypoints.utils import parse_stage_overrides
|
||||
from vllm_omni.entrypoints.utils import prepare_stage_config_inputs
|
||||
from vllm_omni.inputs.data import OmniSamplingParams
|
||||
from vllm_omni.metrics.prometheus import OmniRequestCounter
|
||||
|
||||
@@ -122,6 +123,7 @@ class OmniEngineBase:
|
||||
_transfer_emitter: Any = None
|
||||
_prom_metrics: Any = None
|
||||
_enable_orch_monitor: bool = False
|
||||
_client_config: OmniClientConfig | None = None
|
||||
# Lazily created by get_output_blocking_async().
|
||||
_output_drain_executor: concurrent.futures.ThreadPoolExecutor | None = None
|
||||
|
||||
@@ -136,6 +138,7 @@ class OmniEngineBase:
|
||||
log_stats: bool = False,
|
||||
tokenizer: str | None = None,
|
||||
trust_remote_code: bool | None = None,
|
||||
client_config: OmniClientConfig | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self.model = model
|
||||
@@ -155,6 +158,7 @@ class OmniEngineBase:
|
||||
# --log-stats CLI flag set by the user via OmniBase.
|
||||
self._log_stats = log_stats
|
||||
self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False))
|
||||
self._client_config = client_config
|
||||
|
||||
logger.info(f"[OmniEngine] Initializing with model {model}")
|
||||
|
||||
@@ -166,6 +170,8 @@ class OmniEngineBase:
|
||||
_stage_id_kwarg = kwargs.get("stage_id")
|
||||
if isinstance(_stage_id_kwarg, int) and not single_stage_mode:
|
||||
single_stage_mode = True
|
||||
if client_config is not None and int(client_config.get("client_count", 1)) > 1 and single_stage_mode:
|
||||
raise ValueError("Multiple API servers cannot be combined with single-stage distributed mode")
|
||||
|
||||
self.single_stage_mode: bool = single_stage_mode
|
||||
self._single_stage_id_filter: int | None = (
|
||||
@@ -378,6 +384,7 @@ class OmniEngineBase:
|
||||
omni_lb_policy=self._omni_lb_policy,
|
||||
request_queue=self.request_queue,
|
||||
log_stats=self._log_stats,
|
||||
client_config=self._client_config,
|
||||
)
|
||||
self._runtime.initialize()
|
||||
|
||||
@@ -892,24 +899,15 @@ class OmniEngineBase:
|
||||
) -> tuple[str, list[Any]]:
|
||||
"""Resolve stage configs and inject defaults shared by orchestrator/headless."""
|
||||
|
||||
for legacy_arg in ("stage_configs_path", "stage_configs"):
|
||||
if legacy_arg in kwargs:
|
||||
raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.")
|
||||
|
||||
# log_stats is captured by __init__; its CLI-only negative alias must
|
||||
# not cross into per-stage structured config ownership validation.
|
||||
kwargs.pop("disable_log_stats", None)
|
||||
deploy_config_path = kwargs.pop("deploy_config", None)
|
||||
strategy_config_path = kwargs.pop("strategy_config", None)
|
||||
# CLI callers arrive pre-parsed; offline Python callers may use the
|
||||
# JSON-string form documented in recipes.
|
||||
stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None))
|
||||
|
||||
# ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg;
|
||||
# stage configs know the field as ``streaming_output``. Mirror only a
|
||||
# truthy value so the CLI's False default does not override deploy YAML.
|
||||
if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None:
|
||||
kwargs["streaming_output"] = True
|
||||
config_inputs = prepare_stage_config_inputs(
|
||||
model,
|
||||
kwargs,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
kwargs = config_inputs.kwargs
|
||||
deploy_config_path = config_inputs.deploy_config_path
|
||||
strategy_config_path = config_inputs.strategy_config_path
|
||||
stage_overrides = config_inputs.stage_overrides
|
||||
|
||||
resolution = cast(
|
||||
_ConfigResolutionResult,
|
||||
|
||||
@@ -13,6 +13,7 @@ import threading
|
||||
from collections.abc import Callable, Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from multiprocessing import connection
|
||||
from multiprocessing.process import BaseProcess
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
@@ -1138,6 +1139,8 @@ def launch_stage_replica(
|
||||
stage_visible_devices: str | None = None,
|
||||
spawn_device_lock: threading.Lock | None = None,
|
||||
omni_parallel_stage_init: bool = False,
|
||||
num_api_servers: int = 1,
|
||||
watched_frontend_processes: list[BaseProcess] | None = None,
|
||||
) -> Iterator[StageReplicaResources]:
|
||||
"""Launch a local LLM stage replica.
|
||||
|
||||
@@ -1148,6 +1151,8 @@ def launch_stage_replica(
|
||||
keep the same returned resource bundle.
|
||||
"""
|
||||
if omni_master_server is not None:
|
||||
if num_api_servers != 1:
|
||||
raise ValueError("Multiple API servers are not supported with distributed stage launch")
|
||||
with _launch_omni_core_engines(
|
||||
vllm_config=vllm_config,
|
||||
executor_class=executor_class,
|
||||
@@ -1173,7 +1178,16 @@ def launch_stage_replica(
|
||||
|
||||
from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager
|
||||
|
||||
addresses = get_engine_zmq_addresses(vllm_config)
|
||||
if num_api_servers > 1:
|
||||
# Multi-API Omni is intentionally local-only. IPC paths avoid the
|
||||
# bind-after-allocation race of tcp://host:0 and are stable for every
|
||||
# stage, including stages that are not used by APIServerProcessManager.
|
||||
addresses = EngineZmqAddresses(
|
||||
inputs=[get_open_zmq_ipc_path() for _ in range(num_api_servers)],
|
||||
outputs=[get_open_zmq_ipc_path() for _ in range(num_api_servers)],
|
||||
)
|
||||
else:
|
||||
addresses = get_engine_zmq_addresses(vllm_config)
|
||||
handshake_address = get_open_zmq_ipc_path()
|
||||
engines_to_handshake = [CoreEngine(index=0, local=True)]
|
||||
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock, stage_id=stage_id, replica_id=replica_id):
|
||||
@@ -1197,18 +1211,21 @@ def launch_stage_replica(
|
||||
manager=engine_manager,
|
||||
addresses=addresses,
|
||||
)
|
||||
engine_launch = CoreEngineLaunch(
|
||||
engine_manager=engine_manager,
|
||||
coordinator=None,
|
||||
addresses=addresses,
|
||||
tensor_queue=None,
|
||||
)
|
||||
if watched_frontend_processes is not None:
|
||||
engine_launch.watched_frontend_processes = watched_frontend_processes
|
||||
wait_for_engine_startup(
|
||||
handshake_socket,
|
||||
engines_to_handshake,
|
||||
vllm_config.parallel_config,
|
||||
False, # coordinated_dp
|
||||
vllm_config.cache_config,
|
||||
CoreEngineLaunch(
|
||||
engine_manager=engine_manager,
|
||||
coordinator=None,
|
||||
addresses=addresses,
|
||||
tensor_queue=None,
|
||||
),
|
||||
engine_launch,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1475,6 +1475,8 @@ def build_vllm_config(
|
||||
stage_connector_spec: dict[str, Any] | None = None,
|
||||
engine_args_dict: dict[str, Any] | None = None,
|
||||
headless: bool = False,
|
||||
api_process_count: int = 1,
|
||||
api_process_rank: int = 0,
|
||||
) -> tuple[Any, type]:
|
||||
"""Build engine args, then create VllmConfig and executor_class.
|
||||
|
||||
@@ -1491,6 +1493,9 @@ def build_vllm_config(
|
||||
)
|
||||
|
||||
filtered_engine_args_dict = filter_dataclass_kwargs(OmniEngineArgs, engine_args_dict)
|
||||
if api_process_count != 1 or api_process_rank != 0:
|
||||
filtered_engine_args_dict["_api_process_count"] = api_process_count
|
||||
filtered_engine_args_dict["_api_process_rank"] = api_process_rank
|
||||
|
||||
# _to_dict serializes dataclass fields (e.g. StructuredOutputsConfig) into
|
||||
# plain dicts. When OmniEngineArgs is instantiated with the dict, these
|
||||
@@ -1747,6 +1752,7 @@ def acquire_device_locks(
|
||||
stage_id: int,
|
||||
engine_args_dict: dict[str, Any],
|
||||
stage_init_timeout: int,
|
||||
locked_devices: set[int] | None = None,
|
||||
) -> list[int]:
|
||||
"""Acquire exclusive file locks on devices needed by this stage.
|
||||
|
||||
@@ -1801,8 +1807,7 @@ def acquire_device_locks(
|
||||
f"but only {len(physical_devices)} device(s) are available: {physical_devices}"
|
||||
)
|
||||
|
||||
num_devices_to_lock = num_devices_per_stage
|
||||
devices_to_lock = sorted(physical_devices[:num_devices_to_lock])
|
||||
devices_to_lock = sorted(set(physical_devices[:num_devices_per_stage]) - (locked_devices or set()))
|
||||
|
||||
logger.debug(
|
||||
"Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
|
||||
@@ -1812,7 +1817,7 @@ def acquire_device_locks(
|
||||
prefill_context_parallel_size,
|
||||
sequence_parallel_size,
|
||||
cfg_parallel_size,
|
||||
num_devices_to_lock,
|
||||
len(devices_to_lock),
|
||||
devices_to_lock,
|
||||
)
|
||||
|
||||
@@ -1830,6 +1835,8 @@ def acquire_device_locks(
|
||||
record_lock_holder_pid(lock_fd, lock_writable)
|
||||
lock_acquired = True
|
||||
lock_fds.append(lock_fd)
|
||||
if locked_devices is not None:
|
||||
locked_devices.add(device_id)
|
||||
logger.debug("Acquired exclusive lock for device %s", device_id)
|
||||
except BlockingIOError:
|
||||
os.close(lock_fd)
|
||||
|
||||
+314
-115
@@ -6,14 +6,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import concurrent.futures
|
||||
import contextlib
|
||||
import copy
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, cast
|
||||
from contextlib import AbstractContextManager, contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from multiprocessing.process import BaseProcess
|
||||
from typing import Any, TypedDict, cast
|
||||
|
||||
import janus
|
||||
from vllm.logger import init_logger
|
||||
@@ -94,6 +95,44 @@ class StageRemoteFactoryContext:
|
||||
executor_class: type | None = None
|
||||
|
||||
|
||||
class OmniClientConfig(TypedDict, total=False):
|
||||
"""Per-frontend transport/configuration passed across process boundaries."""
|
||||
|
||||
client_count: int
|
||||
client_index: int
|
||||
stage_addresses: dict[int | str, dict[int | str, dict[str, Any]]]
|
||||
input_address: str
|
||||
output_address: str
|
||||
stats_update_address: str
|
||||
actual_address_pipe: Any
|
||||
tensor_queue: Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class StageEngineLaunch:
|
||||
"""Engine resources handed to a local client or the serving parent.
|
||||
|
||||
Successful launch transfers resource ownership to the caller. Startup or
|
||||
attachment failures are rolled back by ``launch_stage_engines``.
|
||||
"""
|
||||
|
||||
client_configs: list[OmniClientConfig]
|
||||
resources: list[StageReplicaResources]
|
||||
watched_frontend_processes: list[BaseProcess] = field(default_factory=list)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
seen: set[int] = set()
|
||||
for stage_resources in reversed(self.resources):
|
||||
for resource in (stage_resources.manager, stage_resources.coordinator):
|
||||
if resource is None or id(resource) in seen:
|
||||
continue
|
||||
seen.add(id(resource))
|
||||
try:
|
||||
resource.shutdown()
|
||||
except Exception:
|
||||
logger.warning("[StageEngineLaunch] resource shutdown failed", exc_info=True)
|
||||
|
||||
|
||||
def _build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]:
|
||||
try:
|
||||
normalized = LoadBalancingPolicy(policy)
|
||||
@@ -125,13 +164,14 @@ class StageRuntime:
|
||||
self,
|
||||
stage_configs: list[Any],
|
||||
model: str,
|
||||
config_path: str,
|
||||
config_path: str | None,
|
||||
*,
|
||||
stage_init_timeout: int,
|
||||
async_chunk: bool,
|
||||
tokenizer: str | None = None,
|
||||
parallel_stage_init: bool = False,
|
||||
log_stats: bool = False,
|
||||
client_config: OmniClientConfig | None = None,
|
||||
) -> None:
|
||||
self._stage_configs = stage_configs
|
||||
self._model = model
|
||||
@@ -146,6 +186,14 @@ class StageRuntime:
|
||||
self._parallel_stage_init = parallel_stage_init
|
||||
self._log_stats = log_stats
|
||||
self._num_stages = len(stage_configs)
|
||||
self._client_count = int((client_config or {}).get("client_count", 1))
|
||||
self._client_index = int((client_config or {}).get("client_index", 0))
|
||||
if self._client_count < 1:
|
||||
raise ValueError(f"client_count must be >= 1, got {self._client_count}")
|
||||
if self._client_index != -1 and not 0 <= self._client_index < self._client_count:
|
||||
raise ValueError(f"client_index must be in [0, {self._client_count}), got {self._client_index}")
|
||||
self._external_stage_addresses = (client_config or {}).get("stage_addresses")
|
||||
self._api_process_rank = self._client_index
|
||||
|
||||
# Populated by initialize()
|
||||
self.stage_pools: list[StagePool] = []
|
||||
@@ -264,6 +312,205 @@ class StageRuntime:
|
||||
self._cleanup_after_initialize_failure()
|
||||
raise exc
|
||||
|
||||
@contextmanager
|
||||
def launch_stage_engines(
|
||||
self,
|
||||
num_api_servers: int = 1,
|
||||
*,
|
||||
stage_plans: Sequence[LogicalStageInitPlan] | None = None,
|
||||
stage_init_timeout: int | None = None,
|
||||
) -> Iterator[StageEngineLaunch]:
|
||||
"""Launch EngineCore stages for local or subprocess client attachment.
|
||||
|
||||
The regular runtime supplies a replica plan from its device-group
|
||||
scheduler; the serving parent resolves the whole pipeline here. Both
|
||||
use the same device locks, spawn, readiness and rollback machinery.
|
||||
One client completes readiness before attachment, preserving the
|
||||
existing in-process frontend. Multiple clients attach in the caller's
|
||||
``with`` body before readiness completes on context exit.
|
||||
"""
|
||||
if self._external_stage_addresses is not None:
|
||||
raise ValueError("A runtime attached to external stage engines cannot launch engines")
|
||||
if num_api_servers < 1:
|
||||
raise ValueError(f"num_api_servers must be >= 1, got {num_api_servers}")
|
||||
timeout = self._stage_init_timeout if stage_init_timeout is None else stage_init_timeout
|
||||
resolve_plans = stage_plans is None
|
||||
if stage_plans is None:
|
||||
self._client_count = num_api_servers
|
||||
self._api_process_rank = -1 if num_api_servers > 1 else self._client_index
|
||||
self._init_visible_devices_baseline = os.environ.get(current_omni_platform.device_control_env_var)
|
||||
stage_plans = self._prepare_stage_plans()
|
||||
|
||||
unsupported = [
|
||||
plan.stage_id
|
||||
for plan in stage_plans
|
||||
if any(replica.metadata.stage_type == "diffusion" for replica in plan.replicas)
|
||||
]
|
||||
if unsupported:
|
||||
raise ValueError(
|
||||
"--api-server-count currently supports EngineCore stages only; "
|
||||
f"diffusion stage(s) are not supported: {unsupported}"
|
||||
)
|
||||
if any(replica.launch_mode != "local" for plan in stage_plans for replica in plan.replicas):
|
||||
raise ValueError("--api-server-count is not supported with remote or headless stages")
|
||||
|
||||
if num_api_servers > 1:
|
||||
parallel_configs: list[Any] = []
|
||||
for plan in stage_plans:
|
||||
for replica in plan.replicas:
|
||||
if replica.stage_vllm_config is None or replica.executor_class is None:
|
||||
raise RuntimeError(f"LLM stage {plan.stage_id} is missing its engine configuration")
|
||||
parallel_configs.append(replica.stage_vllm_config.parallel_config)
|
||||
if any(getattr(parallel_config, "enable_fault_tolerance", False) for parallel_config in parallel_configs):
|
||||
raise ValueError("--api-server-count > 1 cannot be combined with --enable-fault-tolerance")
|
||||
if any(getattr(parallel_config, "enable_elastic_ep", False) for parallel_config in parallel_configs):
|
||||
raise ValueError("--api-server-count > 1 cannot be combined with --enable-elastic-ep")
|
||||
if any(
|
||||
getattr(parallel_config, "data_parallel_size", 1) != 1 or getattr(parallel_config, "use_ray", False)
|
||||
for parallel_config in parallel_configs
|
||||
):
|
||||
raise ValueError(
|
||||
"--api-server-count > 1 currently supports one local process group per replica; "
|
||||
"intra-replica data parallelism and Ray backends are not supported"
|
||||
)
|
||||
|
||||
if resolve_plans and self._parallel_stage_init:
|
||||
self._reject_unguardable_executors(stage_plans)
|
||||
self._run_stage_admission(stage_plans)
|
||||
|
||||
client_configs: list[OmniClientConfig] = [
|
||||
OmniClientConfig(
|
||||
client_count=num_api_servers,
|
||||
client_index=client_index,
|
||||
stage_addresses={},
|
||||
)
|
||||
for client_index in range(num_api_servers)
|
||||
]
|
||||
entered_contexts: list[AbstractContextManager[StageReplicaResources]] = []
|
||||
resources: list[StageReplicaResources] = []
|
||||
watched_frontend_processes: list[BaseProcess] = []
|
||||
# launch_stage_replica performs its readiness handshake in __exit__,
|
||||
# after API workers are started, so initialization locks must remain
|
||||
# held until all launch contexts have exited.
|
||||
lock_fds: list[int] = []
|
||||
locked_devices: set[int] = set()
|
||||
launch: StageEngineLaunch | None = None
|
||||
exited_contexts = 0
|
||||
|
||||
try:
|
||||
for plan in stage_plans:
|
||||
for replica in plan.replicas:
|
||||
if replica.stage_vllm_config is None or replica.executor_class is None:
|
||||
raise RuntimeError(f"LLM stage {plan.stage_id} is missing its engine configuration")
|
||||
if replica.engine_args_dict is None:
|
||||
raise RuntimeError(f"LLM stage {plan.stage_id} is missing engine args")
|
||||
|
||||
physical_devices = self._resolve_replica_physical_devices(
|
||||
replica.metadata.stage_id,
|
||||
replica.metadata.runtime_cfg,
|
||||
)
|
||||
if not self._parallel_stage_init:
|
||||
with self._scoped_spawn_device_env(physical_devices):
|
||||
lock_fds.extend(
|
||||
acquire_device_locks(
|
||||
replica.metadata.stage_id,
|
||||
replica.engine_args_dict,
|
||||
timeout,
|
||||
locked_devices,
|
||||
)
|
||||
)
|
||||
|
||||
launch_context = launch_stage_replica(
|
||||
vllm_config=replica.stage_vllm_config,
|
||||
executor_class=replica.executor_class,
|
||||
log_stats=self._log_stats,
|
||||
stage_id=replica.metadata.stage_id,
|
||||
replica_id=replica.replica_id,
|
||||
stage_config=replica.stage_cfg,
|
||||
omni_master_server=self._get_omni_master_server(),
|
||||
omni_coordinator_address=self._get_coordinator_address(),
|
||||
stage_visible_devices=physical_devices,
|
||||
spawn_device_lock=self._spawn_device_lock,
|
||||
omni_parallel_stage_init=self._parallel_stage_init,
|
||||
num_api_servers=num_api_servers,
|
||||
watched_frontend_processes=watched_frontend_processes if num_api_servers > 1 else None,
|
||||
)
|
||||
# Environment overlays are process-global. Serialize spawn;
|
||||
# parallel initialization waits for READY outside this lock.
|
||||
with self._replica_launch_lock:
|
||||
with stage_runtime_env(replica.metadata.stage_id, replica.metadata.runtime_cfg):
|
||||
stage_resources = launch_context.__enter__()
|
||||
entered_contexts.append(launch_context)
|
||||
if stage_resources is None:
|
||||
raise RuntimeError(f"LLM stage {plan.stage_id} launcher returned no resources")
|
||||
resources.append(stage_resources)
|
||||
if num_api_servers == 1:
|
||||
exited_contexts += 1
|
||||
launch_context.__exit__(None, None, None)
|
||||
addresses = stage_resources.addresses
|
||||
# The generic resource bundle allows no addresses (e.g.
|
||||
# diffusion), but EngineCore attachment requires them.
|
||||
if addresses is None:
|
||||
raise RuntimeError(f"LLM stage {plan.stage_id} launcher returned no addresses")
|
||||
# Unlike the single-client handshake, only stage 0 is
|
||||
# handed to APIServerProcessManager for address updates.
|
||||
# Other stages must already expose connectable endpoints.
|
||||
if num_api_servers > 1 and any(
|
||||
address.startswith("tcp://") and address.rsplit(":", 1)[-1] == "0"
|
||||
for address in (*addresses.inputs, *addresses.outputs)
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Stage {plan.stage_id} returned deferred TCP addresses; "
|
||||
"multi-API launch requires fixed ports or IPC addresses"
|
||||
)
|
||||
if len(addresses.inputs) != num_api_servers or len(addresses.outputs) != num_api_servers:
|
||||
raise RuntimeError(
|
||||
f"Stage {plan.stage_id} returned {len(addresses.inputs)} input and "
|
||||
f"{len(addresses.outputs)} output addresses for {num_api_servers} API servers"
|
||||
)
|
||||
|
||||
if num_api_servers == 1 and lock_fds:
|
||||
release_device_locks(lock_fds)
|
||||
lock_fds.clear()
|
||||
locked_devices.clear()
|
||||
|
||||
for client_index, client_config in enumerate(client_configs):
|
||||
stage_addresses = client_config["stage_addresses"]
|
||||
replica_addresses = stage_addresses.setdefault(plan.stage_id, {})
|
||||
replica_addresses[replica.replica_id] = {
|
||||
"input_address": addresses.inputs[client_index],
|
||||
"output_address": addresses.outputs[client_index],
|
||||
}
|
||||
if addresses.frontend_stats_publish_address is not None:
|
||||
replica_addresses[replica.replica_id]["stats_update_address"] = (
|
||||
addresses.frontend_stats_publish_address
|
||||
)
|
||||
|
||||
launch = StageEngineLaunch(
|
||||
client_configs=client_configs,
|
||||
resources=resources,
|
||||
watched_frontend_processes=watched_frontend_processes,
|
||||
)
|
||||
yield launch
|
||||
|
||||
while exited_contexts < len(entered_contexts):
|
||||
entered_context = entered_contexts[exited_contexts]
|
||||
exited_contexts += 1
|
||||
entered_context.__exit__(None, None, None)
|
||||
except BaseException as exc:
|
||||
if launch is not None:
|
||||
launch.shutdown()
|
||||
else:
|
||||
StageEngineLaunch(client_configs=client_configs, resources=resources).shutdown()
|
||||
exc_info = (type(exc), exc, exc.__traceback__)
|
||||
for pending_context in reversed(entered_contexts[exited_contexts:]):
|
||||
with contextlib.suppress(BaseException):
|
||||
pending_context.__exit__(*exc_info)
|
||||
raise
|
||||
finally:
|
||||
if lock_fds:
|
||||
release_device_locks(lock_fds)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
for pool in self.stage_pools:
|
||||
for client in pool.clients:
|
||||
@@ -406,6 +653,8 @@ class StageRuntime:
|
||||
self._model,
|
||||
stage_connector_spec=stage_connector_spec,
|
||||
engine_args_dict=engine_args_dict,
|
||||
api_process_count=self._client_count,
|
||||
api_process_rank=self._api_process_rank,
|
||||
)
|
||||
|
||||
for replica_id in range(num_replicas):
|
||||
@@ -737,122 +986,70 @@ class StageRuntime:
|
||||
stage_init_timeout: int,
|
||||
) -> StagePoolClient:
|
||||
"""Initialize one local LLM replica using vLLM's launch/attach pattern."""
|
||||
resources: StageReplicaResources | None = None
|
||||
stage_client = None
|
||||
lock_fds: list[int] = []
|
||||
try:
|
||||
physical_devices = self._resolve_replica_physical_devices(
|
||||
plan.metadata.stage_id,
|
||||
plan.metadata.runtime_cfg,
|
||||
)
|
||||
if physical_devices:
|
||||
logger.info(
|
||||
"[stage_init] Stage-%s set runtime devices: %s",
|
||||
plan.metadata.stage_id,
|
||||
physical_devices,
|
||||
)
|
||||
vllm_config = plan.stage_vllm_config
|
||||
executor_class = plan.executor_class
|
||||
if vllm_config is None:
|
||||
raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing vllm_config")
|
||||
if executor_class is None:
|
||||
raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing executor_class")
|
||||
if plan.engine_args_dict is None:
|
||||
raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing engine args")
|
||||
# G3 device locks: in the default (serial) path the parent holds a
|
||||
# per-device LOCK_EX across the whole child init. When parallel stage
|
||||
# init is enabled the engine-core child takes phased SH/EX locks
|
||||
# itself (see stage_phase_lock), so the parent must NOT also hold the
|
||||
# lock here — it would self-deadlock while waiting for the child's
|
||||
# READY handshake.
|
||||
if not self._parallel_stage_init:
|
||||
g3_start = time.perf_counter()
|
||||
with self._scoped_spawn_device_env(physical_devices):
|
||||
lock_fds = acquire_device_locks(
|
||||
plan.metadata.stage_id,
|
||||
plan.engine_args_dict,
|
||||
stage_init_timeout,
|
||||
)
|
||||
logger.debug(
|
||||
"[stage_init] Stage-%s G3 device-lock acquire took %.3fs",
|
||||
plan.metadata.stage_id,
|
||||
time.perf_counter() - g3_start,
|
||||
)
|
||||
|
||||
launch_cm = launch_stage_replica(
|
||||
vllm_config=vllm_config,
|
||||
executor_class=executor_class,
|
||||
log_stats=self._log_stats,
|
||||
stage_id=plan.metadata.stage_id,
|
||||
replica_id=plan.replica_id,
|
||||
stage_config=plan.stage_cfg,
|
||||
omni_master_server=self._get_omni_master_server(),
|
||||
omni_coordinator_address=self._get_coordinator_address(),
|
||||
stage_visible_devices=physical_devices,
|
||||
spawn_device_lock=self._spawn_device_lock,
|
||||
omni_parallel_stage_init=self._parallel_stage_init,
|
||||
)
|
||||
# Only process spawning and the runtime.env overlay need the global
|
||||
# launch lock. The launch context's exit waits for READY: holding the
|
||||
# lock there serializes model loading and compilation even across
|
||||
# different GPUs. Default initialization keeps its per-device EX
|
||||
# locks until READY; parallel_stage_init uses child phase locks.
|
||||
g2_start = time.perf_counter()
|
||||
with ExitStack() as launch_stack:
|
||||
with self._replica_launch_lock:
|
||||
g2_locked = time.perf_counter()
|
||||
with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
|
||||
resources = launch_stack.enter_context(launch_cm)
|
||||
g2_spawned = time.perf_counter()
|
||||
logger.debug(
|
||||
"[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn=%.3fs, READY=%.3fs",
|
||||
plan.metadata.stage_id,
|
||||
g2_locked - g2_start,
|
||||
g2_spawned - g2_locked,
|
||||
time.perf_counter() - g2_spawned,
|
||||
external_addresses = self._get_external_client_addresses(
|
||||
plan.metadata.stage_id,
|
||||
plan.replica_id,
|
||||
)
|
||||
if external_addresses is not None:
|
||||
if plan.stage_vllm_config is None or plan.executor_class is None:
|
||||
raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing its engine configuration")
|
||||
return cast(
|
||||
StagePoolClient,
|
||||
StageEngineCoreClientBase.make_async_mp_client(
|
||||
vllm_config=plan.stage_vllm_config,
|
||||
executor_class=plan.executor_class,
|
||||
log_stats=self._log_stats,
|
||||
metadata=plan.metadata,
|
||||
client_addresses=external_addresses,
|
||||
client_count=self._client_count,
|
||||
client_index=self._client_index,
|
||||
),
|
||||
)
|
||||
|
||||
logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id)
|
||||
if resources is None:
|
||||
raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no resources")
|
||||
if resources.addresses is None:
|
||||
raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no addresses")
|
||||
stage_client = StageEngineCoreClientBase.make_async_mp_client(
|
||||
vllm_config=vllm_config,
|
||||
executor_class=executor_class,
|
||||
log_stats=self._log_stats,
|
||||
metadata=plan.metadata,
|
||||
client_addresses=self._client_addresses_from_zmq(resources.addresses),
|
||||
engine_manager=resources.manager,
|
||||
coordinator=resources.coordinator,
|
||||
if plan.stage_vllm_config is None or plan.executor_class is None:
|
||||
raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing its engine configuration")
|
||||
stage_plan = LogicalStageInitPlan(
|
||||
stage_idx=plan.metadata.stage_id,
|
||||
stage_id=plan.metadata.stage_id,
|
||||
replicas=[plan],
|
||||
)
|
||||
with self.launch_stage_engines(
|
||||
stage_plans=[stage_plan],
|
||||
stage_init_timeout=stage_init_timeout,
|
||||
) as launch:
|
||||
resources = launch.resources[0]
|
||||
return cast(
|
||||
StagePoolClient,
|
||||
StageEngineCoreClientBase.make_async_mp_client(
|
||||
vllm_config=plan.stage_vllm_config,
|
||||
executor_class=plan.executor_class,
|
||||
log_stats=self._log_stats,
|
||||
metadata=plan.metadata,
|
||||
client_addresses=self._client_addresses_from_zmq(resources.addresses),
|
||||
engine_manager=resources.manager,
|
||||
coordinator=resources.coordinator,
|
||||
),
|
||||
)
|
||||
|
||||
logger.info("[StageRuntime] Stage %s initialized", plan.metadata.stage_id)
|
||||
return cast(StagePoolClient, stage_client)
|
||||
except Exception:
|
||||
if stage_client is not None:
|
||||
try:
|
||||
stage_client.shutdown()
|
||||
except Exception as cleanup_error:
|
||||
logger.warning(
|
||||
"[StageRuntime] Failed to cleanup stage %s after init failure: %s",
|
||||
plan.metadata.stage_id,
|
||||
cleanup_error,
|
||||
)
|
||||
else:
|
||||
self._cleanup_launched_resources(
|
||||
stage_id=plan.metadata.stage_id,
|
||||
resources=resources,
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
if lock_fds:
|
||||
release_device_locks(lock_fds)
|
||||
|
||||
def _get_coordinator_address(self) -> str | None:
|
||||
"""Return coordinator router address. Overridden by DistStageRuntime."""
|
||||
return None
|
||||
|
||||
def _get_external_client_addresses(self, stage_id: int, replica_id: int) -> dict[str, Any] | None:
|
||||
if self._external_stage_addresses is None:
|
||||
return None
|
||||
stage_addresses = self._external_stage_addresses.get(stage_id)
|
||||
if stage_addresses is None:
|
||||
stage_addresses = self._external_stage_addresses.get(str(stage_id))
|
||||
if stage_addresses is None:
|
||||
raise ValueError(f"Missing external addresses for stage {stage_id}")
|
||||
replica_addresses = stage_addresses.get(replica_id)
|
||||
if replica_addresses is None:
|
||||
replica_addresses = stage_addresses.get(str(replica_id))
|
||||
if replica_addresses is None:
|
||||
raise ValueError(f"Missing external addresses for stage {stage_id} replica {replica_id}")
|
||||
return dict(replica_addresses)
|
||||
|
||||
def _get_omni_master_server(self) -> OmniMasterServer | None:
|
||||
"""Return the master server for local distributed launches, if any."""
|
||||
return None
|
||||
@@ -991,7 +1188,7 @@ class DistStageRuntime(StageRuntime):
|
||||
self,
|
||||
stage_configs: list[Any],
|
||||
model: str,
|
||||
config_path: str,
|
||||
config_path: str | None,
|
||||
*,
|
||||
stage_init_timeout: int,
|
||||
async_chunk: bool,
|
||||
@@ -1324,7 +1521,7 @@ class DistStageRuntime(StageRuntime):
|
||||
def create_stage_runtime(
|
||||
stage_configs: list[Any],
|
||||
model: str,
|
||||
config_path: str,
|
||||
config_path: str | None,
|
||||
*,
|
||||
single_stage_mode: bool,
|
||||
stage_init_timeout: int,
|
||||
@@ -1340,6 +1537,7 @@ def create_stage_runtime(
|
||||
omni_lb_policy: str = "random",
|
||||
request_queue: janus.Queue[EngineQueueMessage] | None = None,
|
||||
log_stats: bool = False,
|
||||
client_config: OmniClientConfig | None = None,
|
||||
) -> StageRuntime:
|
||||
"""Factory: select StageRuntime or DistStageRuntime."""
|
||||
if single_stage_mode:
|
||||
@@ -1371,4 +1569,5 @@ def create_stage_runtime(
|
||||
tokenizer=tokenizer,
|
||||
parallel_stage_init=parallel_stage_init,
|
||||
log_stats=log_stats,
|
||||
client_config=client_config,
|
||||
)
|
||||
@@ -8,13 +8,16 @@ Supports both multi-stage LLM models (e.g., Qwen2.5-Omni) and
|
||||
diffusion models (e.g., Qwen-Image) through the same CLI interface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import signal
|
||||
from types import FrameType
|
||||
from typing import Any, cast
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
import uvloop
|
||||
from vllm.entrypoints.cli.types import CLISubcommand
|
||||
@@ -23,10 +26,18 @@ from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.entrypoints.cli.logo import log_logo
|
||||
from vllm_omni.entrypoints.openai.api_server import omni_run_server
|
||||
from vllm_omni.entrypoints.utils import parse_stage_overrides
|
||||
from vllm_omni.entrypoints.openai.api_server import (
|
||||
omni_run_server,
|
||||
run_omni_api_server_worker_proc,
|
||||
)
|
||||
from vllm_omni.entrypoints.utils import parse_stage_overrides, prepare_stage_config_inputs
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.v1.utils import APIServerProcessManager
|
||||
|
||||
from vllm_omni.engine.stage_runtime import StageEngineLaunch, StageRuntime
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni models
|
||||
@@ -141,6 +152,8 @@ class OmniServeCommand(CLISubcommand):
|
||||
|
||||
if args.headless:
|
||||
run_headless(args)
|
||||
elif (getattr(args, "api_server_count", None) or 1) > 1:
|
||||
run_multi_api_server_omni(args)
|
||||
else:
|
||||
uvloop.run(omni_run_server(args))
|
||||
|
||||
@@ -198,6 +211,39 @@ class OmniServeCommand(CLISubcommand):
|
||||
if getattr(args, "omni_replica_address", None) is not None and not args.headless:
|
||||
raise ValueError("--omni-replica-address requires --headless to be set")
|
||||
|
||||
api_server_count = getattr(args, "api_server_count", None)
|
||||
if api_server_count is not None and api_server_count < 1 and not args.headless:
|
||||
raise ValueError("--api-server-count must be >= 1 unless --headless is set")
|
||||
if args.headless and api_server_count is not None and api_server_count > 0:
|
||||
raise ValueError("--api-server-count cannot be used with --headless")
|
||||
if api_server_count is not None and api_server_count > 1:
|
||||
distributed_args = (
|
||||
"stage_id",
|
||||
"omni_master_address",
|
||||
"omni_master_port",
|
||||
"omni_replica_address",
|
||||
)
|
||||
if any(getattr(args, name, None) is not None for name in distributed_args):
|
||||
raise ValueError(
|
||||
"--api-server-count > 1 cannot be combined with stage-based "
|
||||
"or distributed stage arguments (--stage-id/--omni-master-* / "
|
||||
"--omni-replica-address)"
|
||||
)
|
||||
if getattr(args, "omni_dp_size_local", 1) != 1:
|
||||
raise ValueError("--api-server-count > 1 requires --omni-dp-size-local=1")
|
||||
if getattr(args, "worker_backend", "multi_process") != "multi_process":
|
||||
raise ValueError("--api-server-count > 1 requires --worker-backend=multi_process")
|
||||
if getattr(args, "enable_fault_tolerance", False):
|
||||
raise ValueError("--api-server-count > 1 cannot be combined with --enable-fault-tolerance")
|
||||
if getattr(args, "enable_elastic_ep", False):
|
||||
raise ValueError("--api-server-count > 1 cannot be combined with --enable-elastic-ep")
|
||||
if getattr(args, "enable_sleep_mode", False):
|
||||
raise ValueError("--api-server-count > 1 cannot be combined with --enable-sleep-mode")
|
||||
from vllm import envs as vllm_envs
|
||||
|
||||
if getattr(vllm_envs, "VLLM_ALLOW_RUNTIME_LORA_UPDATING", False):
|
||||
raise ValueError("--api-server-count > 1 cannot be combined with VLLM_ALLOW_RUNTIME_LORA_UPDATING")
|
||||
|
||||
# --omni-dp-size-local is process-local. A value other than 1 only
|
||||
# makes sense when this process owns a stage (head or headless).
|
||||
omni_dp_size_local = getattr(args, "omni_dp_size_local", None)
|
||||
@@ -221,7 +267,6 @@ class OmniServeCommand(CLISubcommand):
|
||||
"data_parallel_rpc_port": "--data-parallel-rpc-port",
|
||||
"data_parallel_start_rank": "--data-parallel-start-rank",
|
||||
"data_parallel_backend": "--data-parallel-backend",
|
||||
"api_server_count": "--api-server-count",
|
||||
"enable_expert_parallel": "--enable-expert-parallel",
|
||||
}
|
||||
offenders = sorted(flag for dest, flag in prohibited_with_omni.items() if dest in explicit_cli_keys)
|
||||
@@ -261,6 +306,8 @@ class OmniServeCommand(CLISubcommand):
|
||||
|
||||
model = getattr(args, "model_tag", None) or getattr(args, "model", None)
|
||||
if model and is_diffusion_model(model):
|
||||
if api_server_count is not None and api_server_count > 1:
|
||||
raise ValueError("--api-server-count > 1 is not supported for diffusion models")
|
||||
logger.info("Detected diffusion model: %s", model)
|
||||
return
|
||||
validate_parsed_serve_args(args)
|
||||
@@ -1027,6 +1074,207 @@ class OmniServeCommand(CLISubcommand):
|
||||
return serve_parser
|
||||
|
||||
|
||||
def _build_multi_api_stage_runtime(args: TrackingNamespace, num_api_servers: int) -> StageRuntime:
|
||||
"""Resolve the local EngineCore stages that the parent process owns."""
|
||||
from vllm_omni.config.resolver import resolve_omni_config
|
||||
from vllm_omni.engine.stage_runtime import StageRuntime
|
||||
|
||||
kwargs = args.get_explicit_kwargs_dict()
|
||||
model = kwargs.pop("model", None) or args.model
|
||||
|
||||
trust_remote_code = kwargs.get("trust_remote_code")
|
||||
if trust_remote_code is False:
|
||||
trust_remote_code = None
|
||||
config_inputs = prepare_stage_config_inputs(
|
||||
model,
|
||||
kwargs,
|
||||
trust_remote_code=trust_remote_code,
|
||||
snapshot_model=True,
|
||||
)
|
||||
model = config_inputs.model
|
||||
kwargs = config_inputs.kwargs
|
||||
resolved = resolve_omni_config(
|
||||
model,
|
||||
cli_overrides=kwargs,
|
||||
trust_remote_code=config_inputs.trust_remote_code,
|
||||
deploy_config_path=config_inputs.deploy_config_path,
|
||||
stage_overrides=config_inputs.stage_overrides,
|
||||
strategy_config_path=config_inputs.strategy_config_path,
|
||||
)
|
||||
|
||||
config_path = resolved.config_path
|
||||
stage_configs = list(resolved.stage_configs)
|
||||
|
||||
sleep_stages = [
|
||||
int(getattr(stage_config, "stage_id", stage_index))
|
||||
for stage_index, stage_config in enumerate(stage_configs)
|
||||
if bool(
|
||||
getattr(
|
||||
getattr(stage_config, "model_config", getattr(stage_config, "engine_args", None)),
|
||||
"enable_sleep_mode",
|
||||
False,
|
||||
)
|
||||
)
|
||||
]
|
||||
if sleep_stages:
|
||||
raise ValueError(
|
||||
"--api-server-count > 1 cannot be combined with sleep mode; "
|
||||
f"disable enable_sleep_mode for stage(s) {sleep_stages}"
|
||||
)
|
||||
|
||||
async_chunk = any(
|
||||
bool(
|
||||
getattr(
|
||||
getattr(stage, "connector_config", None),
|
||||
"async_chunk",
|
||||
getattr(getattr(stage, "engine_args", None), "async_chunk", False),
|
||||
)
|
||||
)
|
||||
for stage in stage_configs
|
||||
)
|
||||
return StageRuntime(
|
||||
stage_configs=stage_configs,
|
||||
model=model,
|
||||
config_path=config_path,
|
||||
stage_init_timeout=int(getattr(args, "stage_init_timeout", 300)),
|
||||
parallel_stage_init=bool(getattr(args, "parallel_stage_init", False)),
|
||||
async_chunk=async_chunk,
|
||||
tokenizer=getattr(args, "tokenizer", None),
|
||||
log_stats=not bool(getattr(args, "disable_log_stats", False)),
|
||||
)
|
||||
|
||||
|
||||
def _wait_for_multi_api_server_completion(
|
||||
api_server_manager: APIServerProcessManager,
|
||||
engine_launch: StageEngineLaunch,
|
||||
) -> None:
|
||||
"""Wait until API workers complete or any shared stage engine fails."""
|
||||
from multiprocessing import connection
|
||||
|
||||
api_processes = list(api_server_manager.processes)
|
||||
engine_processes = [
|
||||
process
|
||||
for resources in engine_launch.resources
|
||||
if resources.manager is not None
|
||||
for process in resources.manager.processes
|
||||
]
|
||||
api_by_sentinel = {process.sentinel: process for process in api_processes}
|
||||
engine_by_sentinel = {process.sentinel: process for process in engine_processes}
|
||||
|
||||
while api_by_sentinel:
|
||||
ready = connection.wait([*api_by_sentinel, *engine_by_sentinel])
|
||||
for sentinel in ready:
|
||||
if sentinel in engine_by_sentinel:
|
||||
process = engine_by_sentinel[sentinel]
|
||||
raise RuntimeError(
|
||||
f"Shared stage engine process {process.name} (PID: {process.pid}) "
|
||||
f"exited with code {process.exitcode}"
|
||||
)
|
||||
process = api_by_sentinel.pop(sentinel)
|
||||
if process.exitcode != 0:
|
||||
raise RuntimeError(
|
||||
f"API server process {process.name} (PID: {process.pid}) exited with code {process.exitcode}"
|
||||
)
|
||||
|
||||
|
||||
def _start_api_server_process_manager(
|
||||
*,
|
||||
cleanup_timeout: float,
|
||||
**manager_kwargs: object,
|
||||
) -> APIServerProcessManager:
|
||||
"""Construct vLLM's manager with rollback for a partial ``__init__``.
|
||||
|
||||
``APIServerProcessManager`` starts workers one by one in ``__init__`` and
|
||||
installs its finalizer only after every ``Process.start`` succeeds. Keep a
|
||||
reference to the partially initialized object so workers started before a
|
||||
later failure are still terminated.
|
||||
"""
|
||||
from vllm.v1.utils import APIServerProcessManager, shutdown
|
||||
|
||||
manager = APIServerProcessManager.__new__(APIServerProcessManager)
|
||||
try:
|
||||
APIServerProcessManager.__init__(manager, **manager_kwargs)
|
||||
except BaseException:
|
||||
try:
|
||||
for pipe in getattr(manager, "_address_pipes", ()):
|
||||
with contextlib.suppress(Exception):
|
||||
pipe.close()
|
||||
# A Process whose start() raised is still present in the upstream
|
||||
# manager's list, but calling is_alive() on it raises because it
|
||||
# has no Popen object. Only hand successfully started children to
|
||||
# vLLM's shutdown helper.
|
||||
started_processes = [process for process in getattr(manager, "processes", ()) if process.pid is not None]
|
||||
if started_processes:
|
||||
shutdown(started_processes, timeout=cleanup_timeout)
|
||||
except Exception:
|
||||
logger.exception("Failed to clean up partially started API server processes")
|
||||
raise
|
||||
return manager
|
||||
|
||||
|
||||
def run_multi_api_server_omni(args: TrackingNamespace) -> None:
|
||||
"""Launch API subprocesses that share one set of local stage engines."""
|
||||
from vllm.entrypoints.openai.api_server import setup_server
|
||||
from vllm.v1.metrics.prometheus import setup_multiprocess_prometheus
|
||||
|
||||
num_api_servers = int(args.api_server_count)
|
||||
if num_api_servers < 2:
|
||||
raise ValueError(f"api_server_count must be >= 2, got {num_api_servers}")
|
||||
|
||||
setup_multiprocess_prometheus()
|
||||
shutdown_requested = False
|
||||
|
||||
def signal_handler(signum: int, frame: FrameType | None) -> None:
|
||||
nonlocal shutdown_requested
|
||||
logger.debug("Received %d signal.", signum)
|
||||
if not shutdown_requested:
|
||||
shutdown_requested = True
|
||||
raise SystemExit
|
||||
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
listen_address, sock = setup_server(args, reuse_port=True)
|
||||
stage_runtime = None
|
||||
api_server_manager = None
|
||||
engine_launch = None
|
||||
try:
|
||||
stage_runtime = _build_multi_api_stage_runtime(args, num_api_servers)
|
||||
with stage_runtime.launch_stage_engines(num_api_servers) as engine_launch:
|
||||
args._omni_stage_client_configs = engine_launch.client_configs
|
||||
primary_addresses = engine_launch.resources[0].addresses
|
||||
if primary_addresses is None:
|
||||
raise RuntimeError("Primary stage engine returned no addresses")
|
||||
timeout = float(getattr(args, "shutdown_timeout", 5) or 5)
|
||||
api_server_manager = _start_api_server_process_manager(
|
||||
cleanup_timeout=timeout,
|
||||
listen_address=listen_address,
|
||||
sock=sock,
|
||||
args=args,
|
||||
num_servers=num_api_servers,
|
||||
input_addresses=primary_addresses.inputs,
|
||||
output_addresses=primary_addresses.outputs,
|
||||
target_server_fn=run_omni_api_server_worker_proc,
|
||||
)
|
||||
engine_launch.watched_frontend_processes.extend(api_server_manager.processes)
|
||||
actual_inputs, actual_outputs = api_server_manager.gather_actual_addresses()
|
||||
primary_addresses.inputs = actual_inputs
|
||||
primary_addresses.outputs = actual_outputs
|
||||
|
||||
_wait_for_multi_api_server_completion(api_server_manager, engine_launch)
|
||||
finally:
|
||||
timeout = float(getattr(args, "shutdown_timeout", 5) or 5)
|
||||
if api_server_manager is not None:
|
||||
api_server_manager.shutdown(timeout=timeout)
|
||||
if engine_launch is not None:
|
||||
engine_launch.shutdown()
|
||||
if stage_runtime is not None:
|
||||
stage_runtime.shutdown()
|
||||
sock.close()
|
||||
if shutdown_requested:
|
||||
logger.info("Shared API server shutdown completed")
|
||||
|
||||
|
||||
def run_headless(args: TrackingNamespace) -> None:
|
||||
"""Run a single stage in headless mode.
|
||||
|
||||
@@ -1075,15 +1323,6 @@ def run_headless(args: TrackingNamespace) -> None:
|
||||
|
||||
# Filter down to a dict of things explicitly requested by the user
|
||||
args_dict = args.get_explicit_kwargs_dict()
|
||||
# This controls only the API-process WebSocket and is not a stage-engine option.
|
||||
args_dict.pop("robot_openpi_idle_timeout", None)
|
||||
# This CLI-only negative alias is consumed below when selecting the
|
||||
# launcher log_stats value; it is not a per-stage config override.
|
||||
args_dict.pop("disable_log_stats", None)
|
||||
|
||||
deploy_config_path = args_dict.pop("deploy_config", None)
|
||||
strategy_config_path = args_dict.pop("strategy_config", None)
|
||||
stage_overrides = args_dict.pop("stage_overrides", None)
|
||||
|
||||
# ``--replica-id`` is deprecated and ignored — replica ids are
|
||||
# auto-assigned by ``OmniMasterServer`` so headless processes carry
|
||||
@@ -1099,15 +1338,21 @@ def run_headless(args: TrackingNamespace) -> None:
|
||||
)
|
||||
args_dict.pop("replica_id")
|
||||
|
||||
resolved = resolve_omni_config(
|
||||
config_inputs = prepare_stage_config_inputs(
|
||||
model,
|
||||
args_dict,
|
||||
# store_true cannot express an explicit False: absent maps to None
|
||||
# ("not specified") so the deploy yaml's per-stage value applies.
|
||||
trust_remote_code=getattr(args, "trust_remote_code", None) or None,
|
||||
)
|
||||
args_dict = config_inputs.kwargs
|
||||
resolved = resolve_omni_config(
|
||||
model,
|
||||
trust_remote_code=config_inputs.trust_remote_code,
|
||||
cli_overrides=args_dict,
|
||||
deploy_config_path=deploy_config_path,
|
||||
stage_overrides=stage_overrides,
|
||||
strategy_config_path=strategy_config_path,
|
||||
deploy_config_path=config_inputs.deploy_config_path,
|
||||
stage_overrides=config_inputs.stage_overrides,
|
||||
strategy_config_path=config_inputs.strategy_config_path,
|
||||
)
|
||||
config_path = resolved.config_path
|
||||
stage_configs = list(resolved.stage_configs)
|
||||
|
||||
@@ -6,6 +6,7 @@ This module owns app construction, server startup, app-state initialization,
|
||||
and route bodies that have not yet moved to endpoint-owned modules."""
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import dataclasses
|
||||
import json
|
||||
import multiprocessing
|
||||
@@ -14,6 +15,8 @@ import os
|
||||
|
||||
# Image generation API imports
|
||||
import random
|
||||
import signal
|
||||
import socket
|
||||
import time
|
||||
from argparse import Namespace
|
||||
from collections.abc import AsyncIterator
|
||||
@@ -22,6 +25,7 @@ from http import HTTPStatus
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
import uvloop
|
||||
import vllm.envs as envs
|
||||
from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket
|
||||
from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse
|
||||
@@ -81,12 +85,14 @@ from vllm.renderers.online_renderer import OnlineRenderer
|
||||
from vllm.tasks import POOLING_TASKS
|
||||
from vllm.tool_parsers import ToolParserManager
|
||||
from vllm.utils import random_uuid
|
||||
from vllm.utils.system_utils import decorate_logs
|
||||
from vllm.utils.system_utils import decorate_logs, set_process_title
|
||||
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
|
||||
|
||||
from vllm_omni.config.endpoint_policy import (
|
||||
shutdown_unsupported_routes,
|
||||
)
|
||||
from vllm_omni.engine.stage_init_utils import set_death_signal
|
||||
from vllm_omni.engine.stage_runtime import OmniClientConfig
|
||||
from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni
|
||||
from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler
|
||||
from vllm_omni.entrypoints.duplex.warmup import _warmup_duplex_realtime
|
||||
@@ -232,8 +238,55 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
|
||||
await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
|
||||
|
||||
|
||||
async def omni_run_server_worker(listen_address, sock, args, client_config=None, **uvicorn_kwargs) -> None:
|
||||
def run_omni_api_server_worker_proc(
|
||||
listen_address: str,
|
||||
sock: socket.socket,
|
||||
args: TrackingNamespace,
|
||||
client_config: dict[str, Any] | None = None,
|
||||
**uvicorn_kwargs: object,
|
||||
) -> None:
|
||||
"""Entrypoint used by vLLM's API server process manager."""
|
||||
set_death_signal(signal.SIGTERM)
|
||||
manager_config = client_config or {}
|
||||
client_index = int(manager_config.get("client_index", 0))
|
||||
all_client_configs = getattr(args, "_omni_stage_client_configs", None)
|
||||
if not all_client_configs or not 0 <= client_index < len(all_client_configs):
|
||||
raise RuntimeError(f"Missing Omni stage client configuration for API server {client_index}")
|
||||
|
||||
omni_client_config: OmniClientConfig = copy.deepcopy(all_client_configs[client_index])
|
||||
omni_client_config["client_count"] = int(manager_config.get("client_count", 1))
|
||||
omni_client_config["client_index"] = client_index
|
||||
|
||||
stage_addresses = omni_client_config["stage_addresses"]
|
||||
first_stage_id = min(stage_addresses)
|
||||
first_replica_id = min(stage_addresses[first_stage_id])
|
||||
first_addresses = stage_addresses[first_stage_id][first_replica_id]
|
||||
for key in ("input_address", "output_address", "actual_address_pipe", "tensor_queue"):
|
||||
if key in manager_config:
|
||||
first_addresses[key] = manager_config[key]
|
||||
|
||||
set_process_title("APIServer", str(client_index))
|
||||
decorate_logs("APIServer", skip_if_decorated=True)
|
||||
uvloop.run(
|
||||
omni_run_server_worker(
|
||||
listen_address,
|
||||
sock,
|
||||
args,
|
||||
client_config=omni_client_config,
|
||||
**uvicorn_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def omni_run_server_worker(
|
||||
listen_address: str,
|
||||
sock: socket.socket,
|
||||
args: TrackingNamespace,
|
||||
client_config: OmniClientConfig | None = None,
|
||||
**uvicorn_kwargs: object,
|
||||
) -> None:
|
||||
"""Run a single API server worker."""
|
||||
api_server_count = _resolve_api_server_count(args, client_config)
|
||||
|
||||
if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
|
||||
ToolParserManager.import_tool_parser(args.tool_parser_plugin)
|
||||
@@ -273,6 +326,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
|
||||
|
||||
# OMNI: Pass supported_tasks to build_app (required by upstream vLLM)
|
||||
app = build_openai_app(args, supported_tasks)
|
||||
app.state.api_server_count = api_server_count
|
||||
|
||||
# OMNI: Remove upstream routes that we override with omni-specific handlers
|
||||
remove_route_from_app(app, "/v1/chat/completions", {"POST"})
|
||||
@@ -348,6 +402,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
|
||||
if scope["type"] == "http":
|
||||
scope.setdefault("state", {})
|
||||
scope["state"]["request_timestamp"] = time.time()
|
||||
|
||||
await self._inner(scope, receive, send)
|
||||
|
||||
# Startup duplex warmup (duplex_session.warmup_frames in the deploy
|
||||
@@ -406,7 +461,7 @@ async def build_async_omni(
|
||||
args: TrackingNamespace,
|
||||
*,
|
||||
disable_frontend_multiprocessing: bool | None = None,
|
||||
client_config: dict[str, Any] | None = None,
|
||||
client_config: OmniClientConfig | None = None,
|
||||
) -> AsyncIterator[EngineClient]:
|
||||
"""Build an AsyncOmni instance from command-line arguments.
|
||||
|
||||
@@ -437,6 +492,7 @@ async def build_async_omni(
|
||||
async with build_async_omni_from_stage_config(
|
||||
args,
|
||||
disable_frontend_multiprocessing=disable_frontend_multiprocessing,
|
||||
client_config=client_config,
|
||||
) as async_omni:
|
||||
yield async_omni
|
||||
|
||||
@@ -476,6 +532,7 @@ async def build_async_omni_from_stage_config(
|
||||
args: TrackingNamespace,
|
||||
*,
|
||||
disable_frontend_multiprocessing: bool = False,
|
||||
client_config: OmniClientConfig | None = None,
|
||||
) -> AsyncIterator[EngineClient]:
|
||||
"""Create AsyncOmni from stage configuration.
|
||||
|
||||
@@ -532,6 +589,8 @@ async def build_async_omni_from_stage_config(
|
||||
kwargs.pop("robot_openpi_idle_timeout", None)
|
||||
model = kwargs.pop("model", None) or args.model
|
||||
kwargs.setdefault("log_stats", not args.disable_log_stats)
|
||||
if client_config is not None:
|
||||
kwargs["client_config"] = client_config
|
||||
if _should_serve_duplex(model, kwargs):
|
||||
async_omni = DuplexOmni(model=model, **kwargs)
|
||||
else:
|
||||
@@ -1430,6 +1489,42 @@ async def list_voices(raw_request: Request):
|
||||
return JSONResponse(content={"voices": speakers, "uploaded_voices": uploaded_speakers})
|
||||
|
||||
|
||||
def _resolve_api_server_count(args: Namespace, client_config: OmniClientConfig | None) -> int:
|
||||
"""Resolve frontend topology once, from the launcher's client configuration."""
|
||||
requested = getattr(args, "api_server_count", None)
|
||||
actual = 1 if client_config is None else client_config.get("client_count")
|
||||
if type(actual) is not int or actual < 1:
|
||||
raise ValueError("API client configuration requires a positive integer client_count")
|
||||
if requested is not None and (type(requested) is not int or requested < 1 or requested != actual):
|
||||
raise ValueError("API server count does not match the frontend launch configuration")
|
||||
return actual
|
||||
|
||||
|
||||
def _reject_process_local_state_with_multiple_api_workers(raw_request: Request, operation: str) -> None:
|
||||
"""Reject access to frontend state that is not shared across API workers."""
|
||||
api_server_count = getattr(raw_request.app.state, "api_server_count", None)
|
||||
if type(api_server_count) is not int or api_server_count < 1:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
|
||||
detail="API worker topology is not initialized.",
|
||||
)
|
||||
if api_server_count > 1:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.CONFLICT.value,
|
||||
detail=(
|
||||
f"{operation} is not supported with --api-server-count > 1 because "
|
||||
"the operation uses process-local frontend state"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _require_single_api_video_store(raw_request: Request) -> None:
|
||||
# VIDEO_STORE and VIDEO_TASKS are process-local, including reads and
|
||||
# cancellation. Keep this independent of the temporary diffusion launch
|
||||
# restriction: enabling multi-API diffusion alone cannot make jobs shared.
|
||||
_reject_process_local_state_with_multiple_api_workers(raw_request, "Asynchronous video jobs")
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/audio/voices",
|
||||
responses={
|
||||
@@ -1475,6 +1570,7 @@ async def upload_voice(
|
||||
Returns:
|
||||
JSON response with voice information
|
||||
"""
|
||||
_reject_process_local_state_with_multiple_api_workers(raw_request, "Runtime voice upload")
|
||||
handler = Omnispeech(raw_request)
|
||||
if handler is None:
|
||||
return _create_speech_error_json_response(
|
||||
@@ -1540,6 +1636,7 @@ async def delete_voice(name: str, raw_request: Request):
|
||||
Returns:
|
||||
JSON response indicating success or failure
|
||||
"""
|
||||
_reject_process_local_state_with_multiple_api_workers(raw_request, "Runtime voice deletion")
|
||||
handler = Omnispeech(raw_request)
|
||||
if handler is None:
|
||||
return _create_speech_error_json_response(
|
||||
@@ -1633,13 +1730,15 @@ async def streaming_video_output(websocket: WebSocket):
|
||||
@router.websocket("/v1/realtime")
|
||||
async def realtime_websocket(websocket: WebSocket):
|
||||
"""WebSocket endpoint for OpenAI-style realtime interactions."""
|
||||
await _wait_for_duplex_warmup(websocket)
|
||||
duplex_handler = getattr(websocket.app.state, "openai_serving_duplex", None)
|
||||
duplex_query = websocket.query_params.get("duplex")
|
||||
use_duplex_realtime = duplex_handler is not None and (
|
||||
duplex_query is None or (isinstance(duplex_query, str) and duplex_query.lower() in {"1", "true", "on"})
|
||||
)
|
||||
if use_duplex_realtime and duplex_handler is not None:
|
||||
if await _reject_multi_api_duplex(websocket):
|
||||
return
|
||||
await _wait_for_duplex_warmup(websocket)
|
||||
await duplex_handler.handle_realtime_session(websocket)
|
||||
return
|
||||
|
||||
@@ -1669,6 +1768,8 @@ async def _wait_for_duplex_warmup(websocket: WebSocket) -> None:
|
||||
@router.websocket("/v1/duplex")
|
||||
async def duplex_websocket(websocket: WebSocket):
|
||||
"""Alias of ``/v1/realtime?duplex=1``: the same Realtime duplex session protocol."""
|
||||
if await _reject_multi_api_duplex(websocket):
|
||||
return
|
||||
await _wait_for_duplex_warmup(websocket)
|
||||
handler = getattr(websocket.app.state, "openai_serving_duplex", None)
|
||||
if handler is None:
|
||||
@@ -1709,6 +1810,24 @@ async def realtime_robot_openpi(websocket: WebSocket):
|
||||
await connection.handle_connection()
|
||||
|
||||
|
||||
async def _reject_multi_api_duplex(websocket: WebSocket) -> bool:
|
||||
# Session credentials and replay state are local to each frontend; a shared
|
||||
# listening socket cannot route reconnects back to the session owner.
|
||||
count = getattr(websocket.app.state, "api_server_count", None)
|
||||
if count is None or (type(count) is int and count == 1):
|
||||
return False
|
||||
await websocket.accept()
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "error",
|
||||
"error": "Duplex sessions require a single API worker with initialized topology.",
|
||||
"code": "multi_api_duplex_unsupported",
|
||||
}
|
||||
)
|
||||
await websocket.close(code=1008)
|
||||
return True
|
||||
|
||||
|
||||
# RL Rollout serving (RFC #3747, P0)
|
||||
|
||||
|
||||
@@ -2476,6 +2595,7 @@ async def edit_images(
|
||||
|
||||
@router.post(
|
||||
"/v1/videos",
|
||||
dependencies=[Depends(_require_single_api_video_store)],
|
||||
responses={
|
||||
HTTPStatus.OK.value: {"model": VideoResponse},
|
||||
HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
|
||||
@@ -2617,7 +2737,7 @@ async def create_video_sync(
|
||||
)
|
||||
|
||||
|
||||
@router.get("/v1/videos", response_model=VideoListResponse)
|
||||
@router.get("/v1/videos", response_model=VideoListResponse, dependencies=[Depends(_require_single_api_video_store)])
|
||||
async def list_videos(
|
||||
after: str | None = None,
|
||||
limit: int | None = Query(None, ge=0, le=100),
|
||||
@@ -2654,7 +2774,7 @@ async def list_videos(
|
||||
return VideoListResponse(data=jobs, has_more=has_more, first_id=first_id, last_id=last_id)
|
||||
|
||||
|
||||
@router.get("/v1/videos/{video_id}", response_model=None)
|
||||
@router.get("/v1/videos/{video_id}", response_model=None, dependencies=[Depends(_require_single_api_video_store)])
|
||||
async def retrieve_video(video_id: str) -> VideoResponse | JSONResponse:
|
||||
"""Retrieve metadata for a previously created video job.
|
||||
|
||||
@@ -2682,7 +2802,7 @@ async def retrieve_video(video_id: str) -> VideoResponse | JSONResponse:
|
||||
return job
|
||||
|
||||
|
||||
@router.delete("/v1/videos/{video_id}")
|
||||
@router.delete("/v1/videos/{video_id}", dependencies=[Depends(_require_single_api_video_store)])
|
||||
async def delete_video(video_id: str, raw_request: Request) -> VideoDeleteResponse:
|
||||
"""Delete a stored video job and any generated output.
|
||||
|
||||
@@ -2752,7 +2872,7 @@ async def delete_video(video_id: str, raw_request: Request) -> VideoDeleteRespon
|
||||
return VideoDeleteResponse(id=job.id, deleted=True)
|
||||
|
||||
|
||||
@router.get("/v1/videos/{video_id}/content")
|
||||
@router.get("/v1/videos/{video_id}/content", dependencies=[Depends(_require_single_api_video_store)])
|
||||
async def download_video(video_id: str) -> Response:
|
||||
"""Download the generated file for a completed video job.
|
||||
|
||||
@@ -2845,6 +2965,7 @@ async def stop_profile(raw_request: Request, request: ProfileRequest | None = No
|
||||
|
||||
@router.post("/v1/omni/sleep")
|
||||
async def omni_sleep(request: OmniSleepRequest, raw_request: Request):
|
||||
_reject_process_local_state_with_multiple_api_workers(raw_request, "Sleep")
|
||||
engine_client = raw_request.app.state.engine_client
|
||||
sleeping_set = raw_request.app.state.sleeping_stages
|
||||
if not hasattr(engine_client, "sleep"):
|
||||
@@ -2857,6 +2978,7 @@ async def omni_sleep(request: OmniSleepRequest, raw_request: Request):
|
||||
|
||||
@router.post("/v1/omni/wakeup")
|
||||
async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request):
|
||||
_reject_process_local_state_with_multiple_api_workers(raw_request, "Wakeup")
|
||||
engine_client = raw_request.app.state.engine_client
|
||||
sleeping_set = raw_request.app.state.sleeping_stages
|
||||
if not any(sid in sleeping_set for sid in request.stage_ids):
|
||||
|
||||
@@ -4,12 +4,14 @@
|
||||
import json
|
||||
import os
|
||||
import types
|
||||
from dataclasses import fields, is_dataclass
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, fields, is_dataclass
|
||||
from typing import Any, get_args, get_origin
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
|
||||
from vllm_omni.config.config_factory import with_trust_remote_code_override
|
||||
from vllm_omni.entrypoints.stage_utils import _to_dict
|
||||
from vllm_omni.inputs.data import OmniSamplingParams
|
||||
|
||||
@@ -86,6 +88,64 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
|
||||
return parsed
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StageConfigInputs:
|
||||
"""Normalized inputs shared by all stage-config loading paths."""
|
||||
|
||||
model: str
|
||||
kwargs: dict[str, Any]
|
||||
trust_remote_code: bool | None
|
||||
deploy_config_path: str | None
|
||||
stage_overrides: dict[str, dict[str, Any]] | None
|
||||
strategy_config_path: str | None
|
||||
|
||||
|
||||
def prepare_stage_config_inputs(
|
||||
model: str,
|
||||
kwargs: Mapping[str, Any],
|
||||
*,
|
||||
trust_remote_code: bool | None,
|
||||
snapshot_model: bool = False,
|
||||
) -> StageConfigInputs:
|
||||
"""Normalize model/config arguments before resolving stage configs.
|
||||
|
||||
The standard API worker, multi-API parent, and headless entrypoint must
|
||||
apply the same legacy-key filtering, trust-remote-code precedence, and
|
||||
config-file extraction before calling :func:`resolve_omni_config`.
|
||||
"""
|
||||
resolved_model = model
|
||||
if snapshot_model:
|
||||
# Keep model materialization at the CLI boundary for the parent-owned
|
||||
# launch path. The API worker's OmniBase has already done this work.
|
||||
from vllm_omni.entrypoints.omni_base import omni_snapshot_download
|
||||
|
||||
resolved_model = omni_snapshot_download(model)
|
||||
|
||||
resolved_kwargs = dict(kwargs)
|
||||
for legacy_arg in ("stage_configs_path", "stage_configs"):
|
||||
if legacy_arg in resolved_kwargs:
|
||||
raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.")
|
||||
# Frontend process count is not a stage engine override.
|
||||
resolved_kwargs.pop("api_server_count", None)
|
||||
resolved_kwargs.pop("disable_log_stats", None)
|
||||
if resolved_kwargs.get("diffusion_streaming_output") and resolved_kwargs.get("streaming_output") is None:
|
||||
resolved_kwargs["streaming_output"] = True
|
||||
resolved_kwargs.pop("model_tag", None)
|
||||
|
||||
resolved_kwargs = with_trust_remote_code_override(resolved_kwargs, trust_remote_code)
|
||||
deploy_config_path = resolved_kwargs.pop("deploy_config", None)
|
||||
strategy_config_path = resolved_kwargs.pop("strategy_config", None)
|
||||
stage_overrides = parse_stage_overrides(resolved_kwargs.pop("stage_overrides", None))
|
||||
return StageConfigInputs(
|
||||
model=resolved_model,
|
||||
kwargs=resolved_kwargs,
|
||||
trust_remote_code=trust_remote_code,
|
||||
deploy_config_path=deploy_config_path,
|
||||
stage_overrides=stage_overrides,
|
||||
strategy_config_path=strategy_config_path,
|
||||
)
|
||||
|
||||
|
||||
def get_final_stage_id_for_e2e(
|
||||
output_modalities: list[str] | None, default_modalities: list[str], stage_list: list
|
||||
) -> int:
|
||||
|
||||
@@ -1,10 +1,16 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""MiniMax H3 checkpoint resolution for the AR text-encoder stage."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
|
||||
from vllm_omni.model_executor.model_loader.weight_utils import (
|
||||
download_weights_from_hf_specific,
|
||||
)
|
||||
|
||||
_FASTH3_MODULAR_MODEL_PREFIX = "FastVideo/"
|
||||
|
||||
MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = {
|
||||
partition: [
|
||||
f"{subdir}/text_encoder/**",
|
||||
@@ -16,6 +22,22 @@ MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = {
|
||||
MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["combined"] = MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["fl2va"]
|
||||
|
||||
|
||||
def is_minimax_h3_modular(model: str, revision: str | None = None) -> bool:
|
||||
path = Path(model)
|
||||
if path.is_dir():
|
||||
# Native MiniMax-H3 snapshots can advertise the Diffusers modular
|
||||
# class in ``model_index.json`` while retaining the legacy FL2VA /
|
||||
# Ref2VA layout. FastH3 modular releases carry this marker file.
|
||||
return (path / "modular_model_index.json").is_file() or (path / "fastvideo_inference.json").is_file()
|
||||
# The FastVideo release is the modular checkpoint handled by this
|
||||
# resolver. Do not classify the native MiniMax-H3 Hub repository from
|
||||
# its class name alone; its partition files are still required here.
|
||||
if not model.startswith(_FASTH3_MODULAR_MODEL_PREFIX):
|
||||
return False
|
||||
index = get_diffusion_model_index(model, revision=revision) or {}
|
||||
return index.get("_class_name") == "MiniMaxH3ModularPipeline"
|
||||
|
||||
|
||||
def resolve_minimax_h3_partition(
|
||||
model: str,
|
||||
task_type: str | None,
|
||||
@@ -42,10 +64,13 @@ def resolve_minimax_h3_encoder_model_root(
|
||||
) -> str:
|
||||
path = Path(model)
|
||||
partition = resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va")
|
||||
modular = is_minimax_h3_modular(model, revision)
|
||||
|
||||
if path.is_dir():
|
||||
if path.name == "text_encoder" and (path / "config.json").is_file():
|
||||
return str(path)
|
||||
if modular:
|
||||
return str(path / "text_encoder")
|
||||
if path.name in {"FL2VA", "Ref2VA"}:
|
||||
path = path.parent
|
||||
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
|
||||
@@ -53,11 +78,13 @@ def resolve_minimax_h3_encoder_model_root(
|
||||
snapshot = download_weights_from_hf_specific(
|
||||
model_name_or_path=model,
|
||||
cache_dir=None,
|
||||
allow_patterns=MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition],
|
||||
allow_patterns=["text_encoder/**"] if modular else MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition],
|
||||
revision=revision,
|
||||
require_all=True,
|
||||
)
|
||||
|
||||
if modular:
|
||||
return str(Path(snapshot) / "text_encoder")
|
||||
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
|
||||
return str(Path(snapshot) / subdir / "text_encoder")
|
||||
|
||||
|
||||
@@ -177,6 +177,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
|
||||
num_scheduled_tokens_np=np.ones(bsz, dtype=np.int32),
|
||||
max_num_scheduled_tokens=1,
|
||||
use_cascade_attn=False,
|
||||
force_uniform_decode=True,
|
||||
)
|
||||
n = batch_desc.num_tokens
|
||||
ids = self.talker_mtp_input_ids.gpu[:n]
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import argparse
|
||||
from typing import Any
|
||||
@@ -52,12 +52,28 @@ class TrackingNamespace(argparse.Namespace):
|
||||
else:
|
||||
setattr(self.unfiltered_ns, name, value)
|
||||
|
||||
def __getstate__(self) -> dict[str, Any]:
|
||||
"""Preserve wrapper state when API workers use ``spawn``.
|
||||
|
||||
Our __dict__ property exposes the inner namespace, so default state
|
||||
serialization would lose unfiltered_ns and explicit_keys.
|
||||
"""
|
||||
return {
|
||||
"unfiltered_ns": self.unfiltered_ns,
|
||||
"explicit_keys": self.explicit_keys,
|
||||
}
|
||||
|
||||
def __setstate__(self, state: dict[str, Any]) -> None:
|
||||
object.__setattr__(self, "unfiltered_ns", state["unfiltered_ns"])
|
||||
object.__setattr__(self, "explicit_keys", state["explicit_keys"])
|
||||
|
||||
def get_explicit_kwargs_dict(self):
|
||||
"""Return a dict containing only the explicitly passed key-value pairs."""
|
||||
return {k: v for k, v in vars(self.unfiltered_ns).items() if k in self.explicit_keys}
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self.unfiltered_ns, name)
|
||||
namespace = object.__getattribute__(self, "unfiltered_ns")
|
||||
return getattr(namespace, name)
|
||||
|
||||
@property
|
||||
def __dict__(self):
|
||||
|
||||
@@ -812,7 +812,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
|
||||
|
||||
if self.uses_mrope:
|
||||
positions = self.mrope_positions.gpu[:, :num_tokens_padded]
|
||||
elif self.uses_xdrope_dim > 0:
|
||||
elif getattr(self, "uses_xdrope_dim", 0) > 0:
|
||||
positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
|
||||
else:
|
||||
positions = self.positions[:num_tokens_padded]
|
||||
|
||||
@@ -690,7 +690,7 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
self._init_mrope_positions(req_state)
|
||||
|
||||
# Only relevant for models using XD-RoPE (e.g, HunYuan-VL)
|
||||
if self.uses_xdrope_dim > 0:
|
||||
if getattr(self, "uses_xdrope_dim", 0) > 0:
|
||||
self._init_xdrope_positions(req_state)
|
||||
|
||||
reqs_to_add.append(self.requests[req_id])
|
||||
@@ -1164,7 +1164,7 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
|
||||
if self.uses_mrope:
|
||||
positions = self.mrope_positions.gpu[:, :num_tokens_padded]
|
||||
elif self.uses_xdrope_dim > 0:
|
||||
elif getattr(self, "uses_xdrope_dim", 0) > 0:
|
||||
positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
|
||||
else:
|
||||
positions = self.positions[:num_tokens_padded]
|
||||
@@ -1701,7 +1701,7 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
|
||||
if self.uses_mrope:
|
||||
positions = self.mrope_positions.gpu[:, :num_input_tokens]
|
||||
elif self.uses_xdrope_dim > 0:
|
||||
elif getattr(self, "uses_xdrope_dim", 0) > 0:
|
||||
positions = self.xdrope_positions.gpu[:, :num_input_tokens]
|
||||
else:
|
||||
positions = self.positions[:num_input_tokens]
|
||||
|
||||
Reference in new issue
Block a user