Compare commits

...
Author SHA1 Message Date
Sy03 04ae6aa009 Merge main and adapt multi-API startup to structured configs
Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-17 12:24:06 +08:00
Sy03 79c629885a Merge main and preserve multi-API serving mode selection
Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-17 11:49:47 +08:00
Sy03 d8458b8417 Revert unrelated ROCm CI test and timeout changes
Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-16 21:02:33 +08:00
Sy03 ec0feea2bd Fix ROCm diffusion CI platform guard and startup deadlines
Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-16 20:18:31 +08:00
Sy03 daf91b4d7a Merge main and preserve multi-API engine and duplex guards
Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-16 16:40:16 +08:00
Sy03 40a000de05 Document multi-API serving tradeoffs in TTS profile
Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-16 16:00:45 +08:00
Sy03 dc3c7dbe16 Rebase multi-API serving on latest main
Signed-off-by: Sy03 <1370724210@qq.com>
2026-09-16 12:13:04 +08:00
26 changed files with 1974 additions and 204 deletions

No files matched your search

+69
View File
@@ -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.
+10
View File
@@ -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
+20
View File
@@ -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
+12 -1
View File
@@ -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
@@ -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)]
)
+26 -13
View File
@@ -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)
+289
View File
@@ -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"
+18
View File
@@ -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:
+17 -19
View File
@@ -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,
+24 -7
View File
@@ -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,
)
+10 -3
View File
@@ -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
View File
@@ -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,
)
+262 -17
View File
@@ -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)
+130 -8
View File
@@ -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):
+61 -1
View File
@@ -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:
+18 -2
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
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]
+3 -3
View File
@@ -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]