From 1e5f288a915494f8ffd9783a4886bbfe9929e65e Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Wed, 6 May 2026 19:33:18 +0800 Subject: [PATCH] [FEAT] support multi-stage deployment (#2396) Signed-off-by: ZhengWG Signed-off-by: Zheng Wengang Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com> Signed-off-by: yinpe <11810305@mail.sustech.edu.cn> Signed-off-by: yinpeiqi Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com> Co-authored-by: yinpe <11810305@mail.sustech.edu.cn> Co-authored-by: yinpeiqi Co-authored-by: Hongsheng Liu Co-authored-by: Gao Han Co-authored-by: Chenguang Zheng <645327136@qq.com> --- .buildkite/test-nightly.yml | 79 +- benchmarks/diffusion/README.md | 31 + .../diffusion/multi_replica_serving.md | 80 + docs/user_guide/diffusion_features.md | 3 +- examples/online_serving/qwen3_omni/README.md | 8 + .../test_inline_stage_diffusion_client.py | 22 + .../test_multiproc_engine_concurrency.py | 23 + .../test_bagel_multi_replicas.py | 55 + .../test_qwen3_omni_multi_replicas.py | 137 ++ tests/engine/test_async_omni_engine_input.py | 2 +- .../test_async_omni_engine_stage_init.py | 734 +++--- tests/engine/test_cfg_companion_tracker.py | 13 - tests/engine/test_orchestrator.py | 524 ++++- .../test_orchestrator_error_handling.py | 10 +- .../test_orchestrator_kv_sender_info.py | 53 +- tests/engine/test_single_stage_mode.py | 1991 ++++++----------- tests/entrypoints/test_serve.py | 52 +- tests/entrypoints/test_utils.py | 41 + tests/helpers/runtime.py | 92 +- tests/helpers/stage_config.py | 73 + tests/helpers/test_runtime_stage_cli.py | 56 + .../test_cosyvoice3_stage_input_processors.py | 12 +- .../stage_input_processors/test_glm_image.py | 46 +- .../test_mimo_audio_llm2code2wav.py | 15 +- .../test_qwen3_tts_async_chunk.py | 4 +- .../test_voxtral_tts_async_chunk.py | 20 +- tests/test_config_factory.py | 21 + tests/worker/test_omni_connector_mixin.py | 6 +- vllm_omni/config/stage_config.py | 23 +- .../deploy/qwen3_omni_moe_multi_replicas.yaml | 101 + .../inline_stage_diffusion_client.py | 21 +- vllm_omni/diffusion/stage_diffusion_client.py | 37 +- vllm_omni/engine/__init__.py | 36 + vllm_omni/engine/async_omni_engine.py | 1223 +++++----- vllm_omni/engine/cfg_companion_tracker.py | 21 +- vllm_omni/engine/orchestrator.py | 1169 +++++----- vllm_omni/engine/output_processor.py | 18 + vllm_omni/engine/stage_engine_core_client.py | 85 +- vllm_omni/engine/stage_engine_startup.py | 150 +- vllm_omni/engine/stage_init_utils.py | 305 ++- vllm_omni/engine/stage_pool.py | 403 ++++ vllm_omni/entrypoints/cli/serve.py | 40 +- .../stage_input_processors/cosyvoice3.py | 8 +- .../stage_input_processors/dynin_omni.py | 28 +- .../stage_input_processors/fish_speech.py | 10 +- .../stage_input_processors/glm_image.py | 25 +- .../stage_input_processors/hunyuan_image3.py | 17 +- .../stage_input_processors/mammoth_moda2.py | 9 +- .../stage_input_processors/mimo_audio.py | 17 +- .../stage_input_processors/omnivoice.py | 13 +- .../stage_input_processors/qwen2_5_omni.py | 12 +- .../stage_input_processors/qwen3_omni.py | 71 +- .../stage_input_processors/qwen3_tts.py | 8 +- .../stage_input_processors/voxcpm.py | 20 +- .../stage_input_processors/voxtral_tts.py | 18 +- 55 files changed, 4684 insertions(+), 3407 deletions(-) create mode 100644 docs/user_guide/diffusion/multi_replica_serving.md create mode 100644 tests/e2e/online_serving/test_bagel_multi_replicas.py create mode 100644 tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py create mode 100644 tests/helpers/test_runtime_stage_cli.py create mode 100644 vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml create mode 100644 vllm_omni/engine/stage_pool.py diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index f3c8027ef..daa7f0fe0 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -215,6 +215,44 @@ steps: type: DirectoryOrCreate + - label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100" + timeout_in_minutes: 45 + commands: + - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "core_model" --run-level "core_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 4 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - group: ":card_index_dividers: TTS Model Test" key: nightly-tts-test-group depends_on: upload-nightly-pipeline @@ -297,7 +335,8 @@ steps: - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100" timeout_in_minutes: 120 commands: - - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy + # Keep Bagel expansion and multi-replica tests in their dedicated H100 jobs below. + - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_bagel_multi_replicas.py agents: queue: "mithril-h100-pool" plugins: @@ -370,6 +409,44 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: ":full_moon: Diffusion X2I(&A&T) · BAGEL Multi-Replica with H100" + timeout_in_minutes: 120 + commands: + - pytest -s -v tests/e2e/online_serving/test_bagel_multi_replicas.py -m "full_model and diffusion and H100" --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 4 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4" timeout_in_minutes: 60 commands: diff --git a/benchmarks/diffusion/README.md b/benchmarks/diffusion/README.md index 16a86c88f..a8b11fe90 100644 --- a/benchmarks/diffusion/README.md +++ b/benchmarks/diffusion/README.md @@ -32,6 +32,8 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - The benchmark talks to `http://:/v1/chat/completions`. - If you run the server on another host or port, pass `--base-url` accordingly. +- To compare diffusion multi-replica serving, restart the server with different + `runtime.num_replicas` values and keep the benchmark parameters unchanged. ## 2. Supported Datasets @@ -93,6 +95,35 @@ You can point to your own trace using `--dataset-path`. ## 3. Benchmark Parameters +### Multi-replica sweep + +For replica scaling, use the same model, prompt count, image/video size, and +step count for every run. Change only the serving stage config: + +```yaml +runtime: + devices: "0,1,2,3" + num_replicas: 4 +``` + +Then run a fixed benchmark and write one JSON file per replica count: + +```bash +python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ + --base-url http://localhost:8099 \ + --model Qwen/Qwen-Image \ + --task t2i \ + --dataset random \ + --num-prompts 32 \ + --max-concurrency 8 \ + --request-rate inf \ + --num-inference-steps 10 \ + --width 512 --height 512 \ + --output-file diffusion_replica_4.json +``` + +Repeat for `num_replicas=1`, `2`, `3`, and `4`, using matching `runtime.devices`. + ### Basic flags - `--base-url`: Server address (the script calls `.../v1/chat/completions`). diff --git a/docs/user_guide/diffusion/multi_replica_serving.md b/docs/user_guide/diffusion/multi_replica_serving.md new file mode 100644 index 000000000..2da8c698f --- /dev/null +++ b/docs/user_guide/diffusion/multi_replica_serving.md @@ -0,0 +1,80 @@ +# Diffusion Multi-Replica Serving + +Diffusion stages can be replicated in online serving by setting `runtime.num_replicas` on a `stage_type: diffusion` stage. Each replica owns an independent diffusion engine process and receives requests through round-robin routing with request affinity. + +## When to Use + +Use multi-replica serving when one diffusion process does not saturate the available GPUs, or when you want higher request throughput for many independent image/video generation requests. It improves concurrency throughput; it does not make a single request faster. + +## Configuration + +A replica consumes the devices required by that diffusion stage. For single-device diffusion pipelines, set `runtime.devices` to one device per replica: + +```yaml +stages: + - stage_id: 0 + stage_type: diffusion + runtime: + devices: "0,1,2,3" + num_replicas: 4 + engine_args: + model: Qwen/Qwen-Image + parallel_config: + world_size: 1 +``` + +For multi-device diffusion pipelines, `parallel_config.world_size` is treated as the number of devices per replica. For example, four 2-GPU replicas need eight devices: + +```yaml +stages: + - stage_id: 0 + stage_type: diffusion + runtime: + devices: "0,1,2,3,4,5,6,7" + num_replicas: 4 + engine_args: + model: Wan-AI/Wan2.2-T2V-A14B-Diffusers + parallel_config: + world_size: 2 +``` + +The runtime splits the device list into contiguous chunks per replica. With the second example, replica 0 uses `0,1`, replica 1 uses `2,3`, replica 2 uses `4,5`, and replica 3 uses `6,7`. + +## Serving + +Start the server with the stage config: + +```bash +vllm serve Qwen/Qwen-Image \ + --omni \ + --port 8099 \ + --stage-configs-path /path/to/diffusion_multi_replica.yaml +``` + +Then send normal OpenAI-compatible diffusion chat requests. The orchestrator routes new requests across replicas and keeps updates/abort calls on the same replica for each request. + +## Benchmarking + +Benchmark each replica count by changing `runtime.num_replicas` and `runtime.devices`, restarting the server, then running: + +```bash +python benchmarks/diffusion/diffusion_benchmark_serving.py \ + --base-url http://localhost:8099 \ + --model Qwen/Qwen-Image \ + --task t2i \ + --dataset random \ + --num-prompts 8 \ + --request-rate inf \ + --num-inference-steps 2 \ + --width 256 --height 256 \ + --output-file diffusion_replica_4.json +``` + +For a quick scaling sweep, test `num_replicas` values `1`, `2`, `3`, and `4` with the same prompt count and generation parameters. Compare throughput and end-to-end latency from the JSON outputs. + +## Notes and Limitations + +- Multi-replica serving is currently supported for diffusion stages. Single-stage mode still rejects non-diffusion replicas. +- `runtime.devices` must provide exactly `num_replicas * devices_per_replica` devices. +- Request-level latency may stay similar or slightly increase because each request still runs on one replica. +- Throughput improves only when enough concurrent requests are issued and each replica has enough GPU memory. diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 69c81737d..643301d1b 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -15,7 +15,7 @@ vLLM-Omni supports various advanced features for diffusion models: - Acceleration: **cache methods**, **parallelism methods**, **startup optimizations** - Memory optimization: **cpu offloading**, **quantization** - Extensions: **LoRA inference**, **frame interpolation** -- Execution modes: **step execution** +- Execution modes: **step execution**, **multi-replica serving** ## Supported Features @@ -43,6 +43,7 @@ Parallelism methods distribute computation across GPUs without quality loss (mat | **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs | | **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP | | **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) | +| **[Multi-Replica Serving](diffusion/multi_replica_serving.md)** | Runs multiple independent diffusion stage replicas and routes requests across them | Online serving throughput with many concurrent requests | #### Startup Optimization diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index b34b83e73..a012cf21f 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -28,6 +28,14 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ --deploy-config /path/to/your_deploy_config.yaml ``` +For the bundled 3x-GPU multi-replica layout (talker/code2wav scale-out), +use: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ + --deploy-config vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml +``` + ### Launch individual stages (stage-based CLI) Use the stage-based CLI when you want to run one stage per process. diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index 385f39b12..c495f0832 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -94,3 +95,24 @@ def test_inline_shutdown(client, mock_engine): assert client._shutting_down mock_engine.close.assert_called_once() + + +def test_inline_client_requires_replica_id(mock_engine): + metadata = SimpleNamespace( + stage_id=0, + final_output=True, + final_output_type="image", + default_sampling_params={}, + requires_multimodal_data=False, + custom_process_input_func=None, + engine_input_source=[], + ) + with patch.object(InlineStageDiffusionClient, "_enrich_config"): + od_config = MagicMock(spec=OmniDiffusionConfig) + with pytest.raises(AttributeError, match="replica_id"): + InlineStageDiffusionClient( + model="test_model", + od_config=od_config, + metadata=metadata, + batch_size=1, + ) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 230838fbb..517fdc4c1 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -621,6 +621,29 @@ class TestStageDiffusionClientErrorPropagation: assert client._shutting_down is True assert client._engine_dead is True + def test_initialize_client_requires_replica_id(self): + from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient + + client = object.__new__(StageDiffusionClient) + metadata = SimpleNamespace( + stage_id=0, + final_output=True, + final_output_type="image", + default_sampling_params=None, + requires_multimodal_data=False, + custom_process_input_func=None, + engine_input_source=[], + ) + + with pytest.raises(AttributeError, match="replica_id"): + client._initialize_client( + metadata, + "tcp://req", + "tcp://resp", + proc=None, + batch_size=1, + ) + # ───────── monitor thread & death sentinel integration tests ───────── diff --git a/tests/e2e/online_serving/test_bagel_multi_replicas.py b/tests/e2e/online_serving/test_bagel_multi_replicas.py new file mode 100644 index 000000000..ba6ae6953 --- /dev/null +++ b/tests/e2e/online_serving/test_bagel_multi_replicas.py @@ -0,0 +1,55 @@ +"""Nightly-only e2e coverage for BAGEL diffusion multi-replica serving. + +This test needs 4 H100 GPUs; keep it out of generic test_bagel_* jobs. +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data +from tests.helpers.stage_config import get_deploy_config_path + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +MODEL = "ByteDance-Seed/BAGEL-7B-MoT" +MULTI_REPLICA_DEPLOY = get_deploy_config_path("ci/bagel_multi_replicas_4gpu.yaml") +ROUTE_STRESS_REQUESTS = 6 + +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + +test_params = [ + OmniServerParams( + model=MODEL, + stage_config_path=MULTI_REPLICA_DEPLOY, + server_args=["--disable-log-stats"], + use_stage_cli=True, + ) +] + + +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_multi_stage_diffusion_uses_multi_replica_dit_stage(omni_server, openai_client) -> None: + request_config = { + "model": omni_server.model, + "messages": dummy_messages_from_mix_data( + content_text="<|im_start|>A small red cube on a white table.<|im_end|>" + ), + "modalities": ["image"], + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 2, + "guidance_scale": 0.0, + "seed": 42, + }, + } + + responses = openai_client.send_diffusion_request(request_config, request_num=ROUTE_STRESS_REQUESTS) + + assert len(responses) == ROUTE_STRESS_REQUESTS + assert all(response.success for response in responses) + assert all(response.images is not None and response.images[0].size == (512, 512) for response in responses) diff --git a/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py b/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py new file mode 100644 index 000000000..5b51615fe --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py @@ -0,0 +1,137 @@ +""" +Core-model CI guard for Qwen3-Omni multi-replica stage-pool routing on 4 GPUs. +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video +from tests.helpers.runtime import OmniResponse, OmniServerParams, dummy_messages_from_mix_data +from tests.helpers.stage_config import get_deploy_config_path + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" +MULTI_REPLICA_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe_multi_replicas_4gpu.yaml") +ROUTE_STRESS_REQUESTS = 6 +MIXED_MODAL_REQUESTS = 4 + +test_params = [ + OmniServerParams( + model=MODEL, + stage_config_path=MULTI_REPLICA_DEPLOY, + server_args=["--disable-log-stats"], + ) +] + + +def _system_prompt() -> dict[str, object]: + return { + "role": "system", + "content": [ + { + "type": "text", + "text": ( + "You are Qwen, a virtual human developed by the Qwen Team, " + "Alibaba Group, capable of perceiving auditory and visual inputs, " + "as well as generating text and speech." + ), + } + ], + } + + +def _text_messages() -> list[dict[str, object]]: + return dummy_messages_from_mix_data( + system_prompt=_system_prompt(), + content_text="What is the capital of China? Answer in one short sentence.", + ) + + +def _mixed_messages() -> list[dict[str, object]]: + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}" + return dummy_messages_from_mix_data( + system_prompt=_system_prompt(), + video_data_url=video_data_url, + image_data_url=image_data_url, + audio_data_url=audio_data_url, + content_text="What is recited in the audio? What is in this image? Describe the video briefly.", + ) + + +def _assert_batch_size(responses: list[OmniResponse], expected: int) -> None: + assert len(responses) == expected, f"Expected {expected} responses, got {len(responses)}" + assert all(resp.success for resp in responses), "At least one request failed" + assert all(resp.e2e_latency is not None and resp.e2e_latency > 0 for resp in responses), "Missing request latency" + + +def _assert_text_outputs(responses: list[OmniResponse]) -> None: + assert all(resp.text_content is not None for resp in responses), "Missing text output" + assert all(resp.audio_bytes is None for resp in responses), "Text-only request unexpectedly produced audio" + + +def _assert_audio_outputs(responses: list[OmniResponse], *, expect_text: bool) -> None: + assert all(resp.audio_bytes is not None and len(resp.audio_bytes) > 128 for resp in responses), ( + "Missing audio output" + ) + if expect_text: + assert all(resp.text_content is not None for resp in responses), "Missing text output" + else: + assert all(not (resp.text_content or "").strip() for resp in responses), ( + "Audio-only request unexpectedly produced text" + ) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_only_batch_uses_multi_replica_talker(omni_server, openai_client) -> None: + request_config = { + "model": omni_server.model, + "messages": _text_messages(), + "stream": False, + "modalities": ["text"], + } + + responses = openai_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS) + _assert_batch_size(responses, ROUTE_STRESS_REQUESTS) + _assert_text_outputs(responses) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, openai_client) -> None: + request_config = { + "model": omni_server.model, + "messages": _text_messages(), + "stream": True, + "modalities": ["audio"], + } + + responses = openai_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS) + _assert_batch_size(responses, ROUTE_STRESS_REQUESTS) + _assert_audio_outputs(responses, expect_text=False) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_mixed_modal_stream_batch_generates_text_and_audio(omni_server, openai_client) -> None: + request_config = { + "model": omni_server.model, + "messages": _mixed_messages(), + "stream": True, + } + + responses = openai_client.send_omni_request(request_config, request_num=MIXED_MODAL_REQUESTS) + _assert_batch_size(responses, MIXED_MODAL_REQUESTS) + _assert_audio_outputs(responses, expect_text=True) diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index 3700e426d..d86483f15 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -59,7 +59,7 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock assert request.additional_information is not None assert request.additional_information.entries["text"].list_data == ["hello world"] assert request.additional_information.entries["speaker"].list_data == ["vivian"] - output_processor.add_request.assert_called_once() + output_processor.add_request.assert_not_called() def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture): diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 5c2a9edb7..68ead868e 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1,15 +1,138 @@ import importlib import os import threading +import time import types import pytest from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.stage_init_utils import ( + LogicalStageInitPlan, + ReplicaInitPlan, + build_stage0_input_processor, + compute_replica_layout, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def _make_llm_metadata( + stage_id: int, + *, + replica_id: int = 0, + final_output: bool = False, + final_output_type: str | None = None, + is_comprehension: bool = False, +): + return types.SimpleNamespace( + stage_id=stage_id, + stage_type="llm", + runtime_cfg={}, + prompt_expand_func=None, + final_output=final_output, + final_output_type=final_output_type, + default_sampling_params=types.SimpleNamespace(name=f"sp-{stage_id}-{replica_id}"), + custom_process_input_func=None, + engine_input_source=[] if stage_id == 0 else [stage_id - 1], + engine_output_type="token_ids", + replica_id=replica_id, + is_comprehension=is_comprehension, + ) + + +def _make_diffusion_metadata(stage_id: int, *, replica_id: int = 0, final_output_type: str = "image"): + return types.SimpleNamespace( + stage_id=stage_id, + stage_type="diffusion", + runtime_cfg={"devices": str(replica_id)}, + prompt_expand_func=None, + final_output=True, + final_output_type=final_output_type, + default_sampling_params=types.SimpleNamespace(name=f"dsp-{stage_id}-{replica_id}"), + custom_process_input_func=None, + engine_input_source=[], + cfg_kv_collect_func=None, + replica_id=replica_id, + ) + + +def _make_llm_plan( + stage_idx: int, + *, + configured_stage_id: int, + vllm_config: object, + num_replicas: int = 1, + final_output: bool = False, + final_output_type: str | None = None, + is_comprehension: bool = False, +): + replicas: list[ReplicaInitPlan] = [] + for replica_id in range(num_replicas): + stage_cfg = types.SimpleNamespace( + stage_id=configured_stage_id, + stage_type="llm", + runtime=types.SimpleNamespace(devices=str(replica_id)), + engine_args={}, + ) + replicas.append( + ReplicaInitPlan( + replica_id=replica_id, + num_replicas=num_replicas, + launch_mode="local", + stage_cfg=stage_cfg, + metadata=_make_llm_metadata( + configured_stage_id, + replica_id=replica_id, + final_output=final_output, + final_output_type=final_output_type, + is_comprehension=is_comprehension and replica_id == 0, + ), + stage_connector_spec={}, + omni_kv_connector=(None, None, None), + stage_vllm_config=vllm_config, + executor_class=object, + ) + ) + return LogicalStageInitPlan( + stage_idx=stage_idx, + configured_stage_id=configured_stage_id, + replicas=replicas, + ) + + +def _make_diffusion_plan( + stage_idx: int, + *, + configured_stage_id: int, + num_replicas: int = 1, +): + replicas: list[ReplicaInitPlan] = [] + for replica_id in range(num_replicas): + stage_cfg = types.SimpleNamespace( + stage_id=configured_stage_id, + stage_type="diffusion", + runtime=types.SimpleNamespace(devices=str(replica_id)), + engine_args={}, + ) + replicas.append( + ReplicaInitPlan( + replica_id=replica_id, + num_replicas=num_replicas, + launch_mode="local", + stage_cfg=stage_cfg, + metadata=_make_diffusion_metadata(configured_stage_id, replica_id=replica_id), + stage_connector_spec={}, + omni_kv_connector=(None, None, None), + ) + ) + return LogicalStageInitPlan( + stage_idx=stage_idx, + configured_stage_id=configured_stage_id, + replicas=replicas, + ) + + def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): """Regression test for forward references in make_async_mp_client.""" import vllm_omni.engine.stage_engine_core_client as client_mod @@ -21,15 +144,279 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): ) -def test_initialize_stages_restores_device_visibility_after_diffusion_init(monkeypatch): - """Regression test for stage device env leakage across stage init. +def test_compute_replica_layout_splits_diffusion_devices_by_world_size(): + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="diffusion", + engine_args={"parallel_config": {"tensor_parallel_size": 2}}, + runtime={"devices": "0,1,2,3", "num_replicas": 2}, + ) - Diffusion init mutates process-level CUDA visibility. Ensure AsyncOmniEngine - restores the previous value after diffusion stage setup. - """ + replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg]) + + assert replicas_per_stage == [2] + assert replica_devices_map == {0: ["0,1", "2,3"]} + + +def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): + shared = types.SimpleNamespace(name="shared") + extra = types.SimpleNamespace(name="extra") + + cleanup_clients = AsyncOmniEngine._collect_initialized_clients_for_cleanup( + stage_pools=[types.SimpleNamespace(clients=[shared, None])], + initialized_clients_by_stage={0: [shared], 1: [extra]}, + ) + + assert cleanup_clients == [shared, extra] + + +def test_initialize_stages_rejects_non_diffusion_replicas_in_single_stage_mode(): + engine = object.__new__(AsyncOmniEngine) + engine.single_stage_mode = True + engine.stage_configs = [types.SimpleNamespace(stage_id=0, runtime={"num_replicas": 2})] + + with pytest.raises( + ValueError, + match="single_stage_mode only supports num_replicas > 1 for diffusion stages", + ): + engine._validate_single_stage_mode_replica_constraints() + + +def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch): import vllm_omni.engine.async_omni_engine as engine_mod from vllm_omni.platforms import current_omni_platform + engine = object.__new__(AsyncOmniEngine) + engine.model = "dummy-model" + engine.num_stages = 1 + engine.diffusion_batch_size = 1 + engine.single_stage_mode = False + engine._omni_master_server = None + engine.stage_configs = [] + + plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0] + + env_var = current_omni_platform.device_control_env_var + old_env = os.environ.get(env_var) + os.environ[env_var] = "0,1" + + def _fake_setup_stage_devices(_stage_id, _runtime_cfg): + current_omni_platform.set_device_control_env_var("1") + + monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices) + monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None) + monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace()) + + try: + engine._initialize_diffusion_replica(plan, stage_init_timeout=1, stage_launch_lock=threading.Lock()) + assert os.environ.get(env_var) == "0,1" + finally: + if old_env is None: + os.environ.pop(env_var, None) + else: + os.environ[env_var] = old_env + + +def test_initialize_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch): + import vllm_omni.engine.async_omni_engine as engine_mod + + engine = object.__new__(AsyncOmniEngine) + engine.model = "dummy-model" + engine.num_stages = 1 + engine.diffusion_batch_size = 4 + engine.single_stage_mode = False + engine._omni_master_server = None + engine.stage_configs = [] + + plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0] + + captured: dict[str, object] = {} + + monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None) + monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None) + + def _capture_initialize_diffusion_stage( + stage_id, _model, _stage_cfg, _metadata, *, stage_init_timeout, batch_size, use_inline + ): + captured["stage_id"] = stage_id + captured["stage_init_timeout"] = stage_init_timeout + captured["batch_size"] = batch_size + captured["use_inline"] = use_inline + return types.SimpleNamespace() + + monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", _capture_initialize_diffusion_stage) + + engine._initialize_diffusion_replica(plan, stage_init_timeout=302, stage_launch_lock=threading.Lock()) + + assert captured == { + "stage_id": 0, + "stage_init_timeout": 302, + "batch_size": 4, + "use_inline": True, + } + + +def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_input_processor(monkeypatch): + import vllm_omni.engine.async_omni_engine as engine_mod + + engine = object.__new__(AsyncOmniEngine) + engine.model = "dummy-model" + engine.config_path = "dummy-config" + engine.num_stages = 2 + engine.async_chunk = False + engine.diffusion_batch_size = 1 + engine.single_stage_mode = False + engine._single_stage_id_filter = None + engine._omni_master_server = None + engine.stage_configs = [types.SimpleNamespace(), types.SimpleNamespace()] + + cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) + cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) + stage_plans = [ + _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True), + _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, final_output=True), + ] + + stage0_client_r0 = types.SimpleNamespace( + stage_type="llm", + is_comprehension=True, + final_output=False, + final_output_type=None, + default_sampling_params=types.SimpleNamespace(name="sp0"), + ) + stage0_client_r1 = types.SimpleNamespace( + stage_type="llm", + is_comprehension=False, + final_output=False, + final_output_type=None, + default_sampling_params=types.SimpleNamespace(name="sp0r1"), + ) + stage1_client_r0 = types.SimpleNamespace( + stage_type="llm", + is_comprehension=False, + final_output=True, + final_output_type=None, + default_sampling_params=types.SimpleNamespace(name="sp1"), + ) + initialized_clients = { + 0: [stage0_client_r0, stage0_client_r1], + 1: [stage1_client_r0], + } + + stage0_output_processor = object() + stage1_output_processor = object() + top_level_input_processor = object() + + monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) + monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) + monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2, 1], {})) + monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None)) + monkeypatch.setattr(engine, "_initialize_stage_replicas", lambda *_: initialized_clients) + monkeypatch.setattr( + engine_mod, + "build_llm_stage_output_processor", + lambda plan, _cfg: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor, + ) + monkeypatch.setattr(engine_mod, "build_stage0_input_processor", lambda _cfg: top_level_input_processor) + + engine._initialize_stages(stage_init_timeout=1) + + assert len(engine.stage_pools) == 2 + assert engine.input_processor is top_level_input_processor + assert engine.stage_clients == [stage0_client_r0, stage1_client_r0] + assert engine.stage_vllm_configs == [cfg0, cfg1] + assert engine.output_processors == [stage0_output_processor, stage1_output_processor] + assert engine.default_sampling_params_list == [ + stage0_client_r0.default_sampling_params, + stage1_client_r0.default_sampling_params, + ] + assert engine.stage_metadata == [ + {"final_output": False, "final_output_type": None, "stage_type": "llm"}, + {"final_output": True, "final_output_type": None, "stage_type": "llm"}, + ] + + +def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch): + import vllm_omni.engine.async_omni_engine as engine_mod + + engine = object.__new__(AsyncOmniEngine) + engine.model = "dummy-model" + engine.async_chunk = False + engine.single_stage_mode = False + engine._single_stage_id_filter = None + engine.stage_configs = [ + types.SimpleNamespace(stage_id=0, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="0")), + types.SimpleNamespace( + stage_id=1, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="1,2,3") + ), + ] + + metadata_by_stage = { + 0: _make_llm_metadata(0), + 1: _make_llm_metadata(1), + } + + monkeypatch.setattr( + engine_mod, + "extract_stage_metadata", + lambda cfg: types.SimpleNamespace(**metadata_by_stage[cfg.stage_id].__dict__), + ) + monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr( + engine_mod, + "build_vllm_config", + lambda stage_cfg, *_args, **_kwargs: (types.SimpleNamespace(tag=f"cfg-{stage_cfg.stage_id}"), object), + ) + + stage_plans, prompt_expand_func = engine._build_logical_stage_init_plans( + omni_transfer_config=None, + replicas_per_stage=[1, 3], + replica_devices_map={1: ["1", "2", "3"]}, + ) + + assert prompt_expand_func is None + assert [plan.configured_stage_id for plan in stage_plans] == [0, 1] + assert [replica.stage_cfg.runtime.devices for replica in stage_plans[1].replicas] == ["1", "2", "3"] + assert [replica.replica_id for replica in stage_plans[1].replicas] == [0, 1, 2] + assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas) + + +def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch): + engine = object.__new__(AsyncOmniEngine) + + cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) + cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) + stage_plans = [ + _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2), + _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, num_replicas=2), + ] + + clients = { + (0, 0): types.SimpleNamespace(name="stage0-replica0"), + (0, 1): types.SimpleNamespace(name="stage0-replica1"), + (1, 0): types.SimpleNamespace(name="stage1-replica0"), + (1, 1): types.SimpleNamespace(name="stage1-replica1"), + } + + def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): + time.sleep(0.02 * (3 - plan.metadata.stage_id - plan.replica_id)) + return clients[(plan.metadata.stage_id, plan.replica_id)] + + monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica) + + initialized_clients = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=123) + + assert initialized_clients == { + 0: [clients[(0, 0)], clients[(0, 1)]], + 1: [clients[(1, 0)], clients[(1, 1)]], + } + + +def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch): + import vllm_omni.engine.async_omni_engine as engine_mod + engine = object.__new__(AsyncOmniEngine) engine.model = "dummy-model" engine.config_path = "dummy-config" @@ -39,158 +426,111 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke engine.single_stage_mode = False engine._single_stage_id_filter = None engine._omni_master_server = None - engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")] + engine.stage_configs = [types.SimpleNamespace()] - env_var = current_omni_platform.device_control_env_var - old_env = os.environ.get(env_var) - os.environ[env_var] = "0,1" - - diffusion_client = types.SimpleNamespace(is_comprehension=False) - - metadata = types.SimpleNamespace( - stage_id=0, - stage_type="diffusion", - runtime_cfg={"devices": "1"}, - prompt_expand_func=None, - ) + cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) + stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)] + initialized_client = types.SimpleNamespace(shutdown=lambda: None) monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) - monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata) - monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) - monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {})) + monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None)) - def _fake_setup_stage_devices(_stage_id, _runtime_cfg): - # Simulate diffusion setup mutating process-global visibility. - current_omni_platform.set_device_control_env_var("1") + def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): + if plan.replica_id == 0: + return initialized_client + time.sleep(0.05) + raise RuntimeError("replica launch failed") - monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices) - monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None) - monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: diffusion_client) - monkeypatch.setattr( - engine_mod, - "finalize_initialized_stages", - lambda stage_clients, _input_processor: ( - stage_clients, - [types.SimpleNamespace()], - [{"final_output_type": "image"}], - ), - ) + monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica) - try: + captured_cleanup: list[list[object]] = [] + + def _capture_shutdown(clients): + captured_cleanup.append(list(clients)) + + monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown) + + with pytest.raises(RuntimeError, match="replica launch failed"): engine._initialize_stages(stage_init_timeout=1) - assert os.environ.get(env_var) == "0,1" - finally: - if old_env is None: - os.environ.pop(env_var, None) - else: - os.environ[env_var] = old_env + + assert captured_cleanup == [[initialized_client]] -def test_initialize_stages_passes_stage_init_timeout_to_diffusion_handshake(monkeypatch): - """Regression test for stage_init_timeout passing to complete_diffusion_handshake - in the diffusion stage path. - """ - import vllm_omni.diffusion.data as diffusion_data_mod - import vllm_omni.diffusion.stage_diffusion_client as client_mod +def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_replica_failure(monkeypatch): import vllm_omni.engine.async_omni_engine as engine_mod - from vllm_omni.platforms import current_omni_platform engine = object.__new__(AsyncOmniEngine) - engine.log_stats = False engine.model = "dummy-model" engine.config_path = "dummy-config" - engine.num_stages = 2 + engine.num_stages = 1 engine.async_chunk = False engine.diffusion_batch_size = 1 engine.single_stage_mode = False + engine._single_stage_id_filter = None engine._omni_master_server = None - engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion", engine_args={})] + engine.stage_configs = [types.SimpleNamespace()] - metadata = types.SimpleNamespace( - stage_id=0, - stage_type="diffusion", - runtime_cfg={"devices": "0"}, - prompt_expand_func=None, - final_output=True, - final_output_type="image", - default_sampling_params=None, - custom_process_input_func=None, - engine_input_source=None, - cfg_kv_collect_func=None, - ) - - captured_timeout = None - device_env_var = current_omni_platform.device_control_env_var - prev_device_env = os.environ.get(device_env_var) - os.environ[device_env_var] = "0" + cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) + stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)] + initialized_client = types.SimpleNamespace(shutdown=lambda: None) + replica1_started = threading.Event() monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) - monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata) - monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None) - monkeypatch.setattr( - engine_mod, - "finalize_initialized_stages", - lambda stage_clients, _input_processor: ( - stage_clients, - [types.SimpleNamespace()], - [{"final_output_type": "image"}], - ), - ) - monkeypatch.setattr( - diffusion_data_mod.OmniDiffusionConfig, - "from_kwargs", - classmethod(lambda cls, **kwargs: types.SimpleNamespace(parallel_config=types.SimpleNamespace(world_size=1))), - ) - monkeypatch.setattr( - client_mod, - "spawn_diffusion_proc", - lambda model, od_cfg: (object(), "ipc://handshake", "ipc://request", "ipc://response"), - ) + monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {})) + monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None)) - def _capture_handshake_timeout(proc, handshake_address, handshake_timeout): - nonlocal captured_timeout - captured_timeout = handshake_timeout + def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): + if plan.replica_id == 0: + assert replica1_started.wait(timeout=1) + raise RuntimeError("replica launch failed") + replica1_started.set() + time.sleep(0.05) + return initialized_client - monkeypatch.setattr(client_mod, "complete_diffusion_handshake", _capture_handshake_timeout) - monkeypatch.setattr( - client_mod.zmq, - "Context", - lambda: types.SimpleNamespace(socket=lambda _: types.SimpleNamespace(connect=lambda _: None)), - ) + monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica) - try: - engine._initialize_stages(stage_init_timeout=302) - finally: - if prev_device_env is None: - os.environ.pop(device_env_var, None) - else: - os.environ[device_env_var] = prev_device_env + captured_cleanup: list[list[object]] = [] - assert captured_timeout == 302 + def _capture_shutdown(clients): + captured_cleanup.append(list(clients)) + + monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown) + + with pytest.raises(RuntimeError, match="replica launch failed"): + engine._initialize_stages(stage_init_timeout=1) + + assert captured_cleanup == [[initialized_client]] -def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch): - """Regression test for stage_init_timeout reaching complete_stage_handshake - in the LLM stage path. - """ +def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch): import vllm_omni.engine.async_omni_engine as engine_mod from vllm_omni.platforms import current_omni_platform engine = object.__new__(AsyncOmniEngine) - engine.log_stats = False engine.model = "dummy-model" engine.single_stage_mode = False engine._omni_master_server = None engine.stage_configs = [] - metadata = types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}) fake_vllm_config = types.SimpleNamespace() - fake_addresses = types.SimpleNamespace() + fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None) fake_proc = types.SimpleNamespace() + captured_timeout: int | None = None - captured_timeout = None + plan = ReplicaInitPlan( + replica_id=0, + 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"}), + stage_connector_spec={}, + omni_kv_connector=(None, None, None), + stage_vllm_config=fake_vllm_config, + executor_class=object, + ) device_env_var = current_omni_platform.device_control_env_var prev_device_env = os.environ.get(device_env_var) @@ -198,28 +538,22 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake( monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None) monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) - monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (fake_vllm_config, object)) monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: []) - monkeypatch.setattr( - engine_mod, - "spawn_stage_core", - lambda **_: (fake_addresses, fake_proc, "ipc://handshake"), - ) + monkeypatch.setattr(engine_mod, "spawn_stage_core", lambda **_: (fake_addresses, fake_proc, "ipc://handshake")) def _capture_stage_timeout(_proc, _handshake_addr, _addresses, _vllm_cfg, handshake_timeout): nonlocal captured_timeout captured_timeout = handshake_timeout monkeypatch.setattr(engine_mod, "complete_stage_handshake", _capture_stage_timeout) + monkeypatch.setattr( + engine_mod.StageEngineCoreClientBase, + "make_async_mp_client", + staticmethod(lambda **_: types.SimpleNamespace(shutdown=lambda: None)), + ) try: - engine._launch_llm_stage( - stage_cfg=types.SimpleNamespace(engine_args={}), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=302, - llm_stage_launch_lock=threading.Lock(), - ) + engine._initialize_llm_replica(plan, 302, threading.Lock()) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) @@ -229,155 +563,27 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake( assert captured_timeout == 302 -def test_launch_llm_stage_releases_launch_lock_before_complete_stage_handshake(monkeypatch): - """Regression test for parallel LLM stage startup during handshake wait.""" - import vllm_omni.engine.async_omni_engine as engine_mod - from vllm_omni.platforms import current_omni_platform - - engine = object.__new__(AsyncOmniEngine) - engine.log_stats = False - engine.model = "dummy-model" - engine.single_stage_mode = False - engine._omni_master_server = None - engine.stage_configs = [] - - fake_vllm_config = types.SimpleNamespace() - fake_addresses = types.SimpleNamespace() - shared_launch_lock = threading.Lock() - counter_lock = threading.Lock() - first_handshake_started = threading.Event() - second_stage_spawned = threading.Event() - allow_first_handshake_to_finish = threading.Event() - launch_errors: list[BaseException] = [] - spawn_count = 0 - - device_env_var = current_omni_platform.device_control_env_var - prev_device_env = os.environ.get(device_env_var) - os.environ[device_env_var] = "0" - - monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None) - monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) - monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (fake_vllm_config, object)) - monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: []) - - def _spawn_stage_core(**_): - nonlocal spawn_count - with counter_lock: - spawn_count += 1 - call_idx = spawn_count - if call_idx == 2: - second_stage_spawned.set() - return fake_addresses, types.SimpleNamespace(), f"ipc://handshake-{call_idx}" - - def _complete_stage_handshake(_proc, handshake_address, _addresses, _vllm_cfg, _timeout): - if handshake_address == "ipc://handshake-1": - first_handshake_started.set() - assert second_stage_spawned.wait(timeout=1), ( - "second stage did not reach spawn_stage_core while first stage waited in handshake" - ) - assert allow_first_handshake_to_finish.wait(timeout=1), ( - "second stage did not enter handshake while first stage was still waiting" - ) - else: - allow_first_handshake_to_finish.set() - - monkeypatch.setattr(engine_mod, "spawn_stage_core", _spawn_stage_core) - monkeypatch.setattr(engine_mod, "complete_stage_handshake", _complete_stage_handshake) - - def _launch_stage(stage_id: int) -> None: - metadata = types.SimpleNamespace(stage_id=stage_id, runtime_cfg={"devices": str(stage_id)}) - try: - engine._launch_llm_stage( - stage_cfg=types.SimpleNamespace(engine_args={}), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=302, - llm_stage_launch_lock=shared_launch_lock, - ) - except BaseException as exc: # pragma: no cover - surfaced through assertion below - launch_errors.append(exc) - - try: - first_thread = threading.Thread(target=_launch_stage, args=(0,)) - first_thread.start() - assert first_handshake_started.wait(timeout=1), "first stage never entered handshake" - - second_thread = threading.Thread(target=_launch_stage, args=(1,)) - second_thread.start() - - first_thread.join(timeout=3) - second_thread.join(timeout=3) - finally: - if prev_device_env is None: - os.environ.pop(device_env_var, None) - else: - os.environ[device_env_var] = prev_device_env - - assert not first_thread.is_alive() - assert not second_thread.is_alive() - assert second_stage_spawned.is_set() - assert not launch_errors - - -def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch): - """Regression test for GLM-Image t2i preprocessing path. - - Stage-0 InputProcessor must use OmniInputPreprocessor so text prompts with - mm_processor_kwargs go through multimodal preprocessing. - """ - import vllm_omni.engine.async_omni_engine as engine_mod - - class DummyStageEngineCoreClient: - def __init__(self, **kwargs): - self.kwargs = kwargs - - def shutdown(self): - return None - - class DummyOutputProcessor: - def __init__(self, **kwargs): - self.kwargs = kwargs +def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch): + import vllm_omni.engine.stage_init_utils as init_mod class DummyInputProcessor: def __init__(self, vllm_config): self.vllm_config = vllm_config self.renderer = object() - self.input_preprocessor = object() + self.input_preprocessor = None class DummyOmniInputPreprocessor: def __init__(self, vllm_config, renderer=None): self.vllm_config = vllm_config self.renderer = renderer - monkeypatch.setattr( - engine_mod.StageEngineCoreClientBase, - "make_async_mp_client", - staticmethod(lambda **kwargs: DummyStageEngineCoreClient(**kwargs)), - ) - monkeypatch.setattr(engine_mod, "MultimodalOutputProcessor", DummyOutputProcessor) - monkeypatch.setattr(engine_mod, "InputProcessor", DummyInputProcessor) - monkeypatch.setattr(engine_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor) + monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) + monkeypatch.setattr(init_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor) - started = types.SimpleNamespace( - stage_id=0, - metadata=types.SimpleNamespace(stage_id=0, engine_output_type="token_ids"), - vllm_config=types.SimpleNamespace(model_config=types.SimpleNamespace(skip_tokenizer_init=True)), - executor_class=object, - engine_manager=object(), - coordinator=object(), - proc=None, - addresses=types.SimpleNamespace( - inputs=["inproc://input"], - outputs=["inproc://output"], - frontend_stats_publish_address=None, - ), + input_processor = build_stage0_input_processor( + types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {})) ) - engine = object.__new__(AsyncOmniEngine) - - _stage_client, _out_proc, _vllm_cfg, input_processor = engine._attach_llm_stage(started) - - assert input_processor is not None assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor) assert input_processor.input_preprocessor.renderer is input_processor.renderer diff --git a/tests/engine/test_cfg_companion_tracker.py b/tests/engine/test_cfg_companion_tracker.py index f856a38c3..c3809edc0 100644 --- a/tests/engine/test_cfg_companion_tracker.py +++ b/tests/engine/test_cfg_companion_tracker.py @@ -1,7 +1,6 @@ import pytest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker -from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -25,18 +24,6 @@ def test_register_companion_and_cleanup(): assert tracker.get_companion_request_ids("req1") == {} -def test_attach_cfg_request_ids_clones_diffusion_params(): - tracker = CfgCompanionTracker() - tracker.register_companion("req1", "cfg_text", "req1__cfg_text") - - params = OmniDiffusionSamplingParams() - updated = tracker.attach_cfg_request_ids("req1", params) - - assert updated is not params - assert params.cfg_kv_request_ids is None - assert updated.cfg_kv_request_ids == {"cfg_text": "req1__cfg_text"} - - def test_abort_parent_expands_to_companions_and_cleans_up_deferred_parent(): tracker = CfgCompanionTracker() tracker.register_companion("req1", "cfg_text", "req1__cfg_text") diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index c33e23b05..93b2761ae 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import concurrent.futures +import logging import queue import threading import time @@ -16,8 +17,9 @@ from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams from vllm.v1.engine.core_client import AsyncMPClient -from vllm_omni.engine.orchestrator import Orchestrator +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.engine.stage_pool import StagePool from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -73,7 +75,7 @@ class FakeStageClient: def set_engine_outputs(self, outputs) -> None: return None - def process_engine_inputs(self, stage_list, prompt=None, streaming_context=None): + def process_engine_inputs(self, source_outputs, prompt=None, streaming_context=None): return list(self.next_inputs) async def abort_requests_async(self, request_ids: list[str]) -> None: @@ -90,12 +92,33 @@ class FakeStageClient: self._diffusion_outputs.put_nowait(output) +class FakeCollectiveRpcStageClient(FakeStageClient): + def __init__(self, *args, rpc_result: Any = None, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.rpc_result = rpc_result + self.collective_rpc_calls: list[tuple[str, float | None, tuple[Any, ...], dict[str, Any]]] = [] + + async def collective_rpc_async( + self, + *, + method: str, + timeout: float | None = None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Any: + normalized_kwargs = dict(kwargs or {}) + self.collective_rpc_calls.append((method, timeout, args, normalized_kwargs)) + return self.rpc_result + + class FakeOutputProcessor: def __init__(self, *, request_outputs: list[object] | None = None) -> None: self.request_outputs = list(request_outputs or []) + self.add_request_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] self.abort_calls: list[list[str]] = [] - def add_request(self, *_args, **_kwargs) -> None: + def add_request(self, *args, **kwargs) -> None: + self.add_request_calls.append((args, kwargs)) return None def process_outputs(self, *_args, **_kwargs): @@ -185,17 +208,60 @@ def _build_request_output( ) +def _build_stage_pools( + stage_clients: list[list[FakeStageClient]], + *, + output_processors: list[FakeOutputProcessor] | None = None, + stage_vllm_configs: list[object] | None = None, +) -> list[StagePool]: + """Build StagePool list from per-stage replica lists. + + ``stage_clients[i]`` is the list of FakeStageClient replicas for stage i. + """ + num_stages = len(stage_clients) + if output_processors is None: + output_processors = [FakeOutputProcessor() for _ in stage_clients] + if stage_vllm_configs is None: + stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) for _ in stage_clients] + + pools: list[StagePool] = [] + for stage_id in range(num_stages): + clients = stage_clients[stage_id] + if clients[0].stage_type == "diffusion": + pools.append(StagePool(stage_id, clients[0])) + else: + pools.append( + StagePool( + stage_id, + clients, + output_processor=output_processors[stage_id], + stage_vllm_config=stage_vllm_configs[stage_id], + ) + ) + return pools + + def _build_harness( stage_clients: list[object], *, output_processors: list[object] | None = None, stage_vllm_configs: list[object] | None = None, async_chunk: bool = False, + stage_pools: list[StagePool] | None = None, ) -> OrchestratorFixture: - if output_processors is None: - output_processors = [FakeOutputProcessor() for _ in stage_clients] - if stage_vllm_configs is None: - stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) for _ in stage_clients] + """Build an Orchestrator test harness. + + Accepts either pre-built ``stage_pools`` or flat lists of single-replica + clients/processors. + """ + if stage_pools is None: + # Wrap flat lists into per-stage single-replica lists. + nested_clients = [[c] for c in stage_clients] + stage_pools = _build_stage_pools( + nested_clients, + output_processors=output_processors, + stage_vllm_configs=stage_vllm_configs, + ) ready_future: concurrent.futures.Future[tuple[Orchestrator, janus.Queue, janus.Queue, janus.Queue]] = ( concurrent.futures.Future() @@ -214,9 +280,7 @@ def _build_harness( request_async_queue=request_queue.async_q, output_async_queue=output_queue.async_q, rpc_async_queue=rpc_queue.async_q, - stage_clients=stage_clients, - output_processors=output_processors, - stage_vllm_configs=stage_vllm_configs, + stage_pools=stage_pools, async_chunk=async_chunk, ) ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue)) @@ -283,6 +347,18 @@ async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, time return msg +async def _get_rpc_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict: + deadline = time.monotonic() + timeout + rpc_sync_q = orchestrator_fixture.queues[2].sync_q + while True: + if time.monotonic() >= deadline: + raise AssertionError("Timed out waiting for orchestrator rpc output") + try: + return rpc_sync_q.get_nowait() + except queue.Empty: + await asyncio.sleep(0.01) + + async def _enqueue_add_request( orchestrator_fixture: OrchestratorFixture, *, @@ -387,6 +463,11 @@ def orchestrator_factory(): q.close() +# --------------------------------------------------------------------------- +# Existing single-replica tests (adapted to StagePool interface) +# --------------------------------------------------------------------------- + + @pytest.mark.asyncio async def test_run_two_stage_llm(orchestrator_factory) -> None: stage0 = FakeStageClient(stage_type="llm", final_output=False) @@ -600,12 +681,425 @@ async def test_run_abort(orchestrator_factory) -> None: await _wait_for(lambda: len(stages[0].add_request_calls) == 1) await _enqueue_abort_request(orchestrator_fixture, ["req-abort"]) - await _wait_for(lambda: all(stage.abort_calls for stage in stages)) + await _wait_for(lambda: bool(stages[0].abort_calls)) - for stage in stages: - assert stage.abort_calls == [["req-abort"]] - for proc in processors: - assert proc.abort_calls == [["req-abort"]] + assert stages[0].abort_calls == [["req-abort"]] + assert stages[1].abort_calls == [] assert "req-abort" not in orchestrator_fixture.orchestrator.request_states finally: await _shutdown_orchestrator(orchestrator_fixture) + + +# --------------------------------------------------------------------------- +# Multi-replica tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_multi_replica_round_robin_distribution(orchestrator_factory) -> None: + """Two replicas at stage-0, single replica at stage-1. + + Send two requests — they should land on different stage-0 replicas + (round-robin), then both forward to the single stage-1 replica. + """ + stage0_r0 = FakeStageClient(stage_type="llm", final_output=False) + stage0_r1 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = FakeStageClient( + stage_type="llm", + final_output=True, + next_inputs=[{"prompt_token_ids": [7, 8]}], + ) + + proc0 = FakeOutputProcessor(request_outputs=[_build_request_output("req-0", token_ids=[3], finished=True)]) + proc1 = FakeOutputProcessor(request_outputs=[_build_request_output("req-0", token_ids=[10], finished=True)]) + + default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) + stage_pools = _build_stage_pools( + [[stage0_r0, stage0_r1], [stage1]], + output_processors=[proc0, proc1], + stage_vllm_configs=[default_vllm_cfg, default_vllm_cfg], + ) + + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools) + + try: + # Request 0 → should land on replica 0 (RR starts at 0) + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-0", + prompt=SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]), + original_prompt={"prompt": "hello 0"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + await _wait_for(lambda: len(stage0_r0.add_request_calls) == 1) + assert len(stage0_r1.add_request_calls) == 0 + + # Request 1 → should land on replica 1 (RR advances) + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-1", + prompt=SimpleNamespace(request_id="req-1", prompt_token_ids=[5, 6]), + original_prompt={"prompt": "hello 1"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + await _wait_for(lambda: len(stage0_r1.add_request_calls) == 1) + assert len(stage0_r0.add_request_calls) == 1 # unchanged + + # Complete req-0 at stage-0 replica-0 → should forward to stage-1 + stage0_r0.push_engine_core_outputs(_engine_core_outputs("s0r0-raw", 1.0)) + await _wait_for(lambda: len(stage1.add_request_calls) == 1) + assert stage1.add_request_calls[0][0].request_id == "req-0" + + # Complete req-0 at stage-1 → final output + proc1.request_outputs = [_build_request_output("req-0", token_ids=[10], finished=True)] + stage1.push_engine_core_outputs(_engine_core_outputs("s1-raw", 2.0)) + output_msg = await _get_output_message(orchestrator_fixture) + + assert output_msg["request_id"] == "req-0" + assert output_msg["stage_id"] == 1 + assert output_msg["finished"] is True + assert "req-0" not in orchestrator_fixture.orchestrator.request_states + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_multi_replica_abort_broadcasts_to_all_replicas(orchestrator_factory) -> None: + """Abort must be sent to every replica across all stages.""" + stage0_r0 = FakeStageClient(stage_type="llm", final_output=False) + stage0_r1 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = FakeStageClient(stage_type="llm", final_output=True) + + proc0 = FakeOutputProcessor() + proc1 = FakeOutputProcessor() + + default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) + stage_pools = _build_stage_pools( + [[stage0_r0, stage0_r1], [stage1]], + output_processors=[proc0, proc1], + stage_vllm_configs=[default_vllm_cfg, default_vllm_cfg], + ) + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-abort-mr", + prompt=SimpleNamespace(request_id="req-abort-mr", prompt_token_ids=[1]), + original_prompt={"prompt": "cancel"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + await _wait_for(lambda: len(stage0_r0.add_request_calls) == 1) + + await _enqueue_abort_request(orchestrator_fixture, ["req-abort-mr"]) + await _wait_for(lambda: bool(stage0_r0.abort_calls)) + + assert stage0_r0.abort_calls == [["req-abort-mr"]] + assert stage0_r1.abort_calls == [] + assert stage1.abort_calls == [] + assert "req-abort-mr" not in orchestrator_fixture.orchestrator.request_states + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_multi_replica_shutdown_all_replicas(orchestrator_factory) -> None: + """Shutdown must shut down every replica across all stages.""" + stage0_r0 = FakeStageClient(stage_type="llm", final_output=False) + stage0_r1 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = FakeStageClient(stage_type="llm", final_output=True) + + default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) + stage_pools = _build_stage_pools( + [[stage0_r0, stage0_r1], [stage1]], + stage_vllm_configs=[default_vllm_cfg, default_vllm_cfg], + ) + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools) + + await _shutdown_orchestrator(orchestrator_fixture) + + assert not orchestrator_fixture.thread.is_alive() + for client in [stage0_r0, stage0_r1, stage1]: + assert client.shutdown_calls == 1 + + +@pytest.mark.asyncio +async def test_stage_pool_submit_update_reuses_existing_binding() -> None: + """A request admitted to one replica must keep using that replica on updates.""" + stage0_r0 = FakeStageClient(stage_type="llm", final_output=False) + stage0_r1 = FakeStageClient(stage_type="llm", final_output=False) + pool = StagePool( + 0, + [stage0_r0, stage0_r1], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + + req0_state = OrchestratorRequestState( + request_id="req-0", + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + req1_state = OrchestratorRequestState( + request_id="req-1", + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + + await pool.submit_initial("req-0", req0_state, SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2])) + await pool.submit_update("req-0", req0_state, SimpleNamespace(request_id="req-0", prompt_token_ids=[3])) + await pool.submit_initial("req-1", req1_state, SimpleNamespace(request_id="req-1", prompt_token_ids=[4, 5])) + await pool.submit_update("req-1", req1_state, SimpleNamespace(request_id="req-1", prompt_token_ids=[6])) + + assert pool.get_bound_replica_id("req-0") == 0 + assert pool.get_bound_replica_id("req-1") == 1 + assert len(stage0_r0.add_request_calls) == 2 + assert len(stage0_r1.add_request_calls) == 2 + assert stage0_r0.add_request_calls[0][0].request_id == "req-0" + assert stage0_r0.add_request_calls[1][0].request_id == "req-0" + assert stage0_r1.add_request_calls[0][0].request_id == "req-1" + assert stage0_r1.add_request_calls[1][0].request_id == "req-1" + + +@pytest.mark.asyncio +async def test_stage_pool_submit_update_refreshes_output_processor_state() -> None: + output_processor = FakeOutputProcessor() + + class AssertingStageClient(FakeStageClient): + async def add_request_async(self, *args, **kwargs) -> None: + if len(self.add_request_calls) == 1: + prompts = [call_kwargs["prompt"] for _, call_kwargs in output_processor.add_request_calls] + assert prompts == ["seg-1", "seg-2"] + await super().add_request_async(*args, **kwargs) + + stage0 = AssertingStageClient(stage_type="llm", final_output=False) + pool = StagePool( + 0, + [stage0], + output_processor=output_processor, + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + req_state = OrchestratorRequestState( + request_id="req-0", + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + + await pool.submit_initial( + "req-0", + req_state, + SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]), + prompt_text="seg-1", + ) + await pool.submit_update( + "req-0", + req_state, + SimpleNamespace(request_id="req-0", prompt_token_ids=[3], resumable=True), + prompt_text="seg-2", + ) + + assert len(output_processor.add_request_calls) == 2 + assert output_processor.add_request_calls[1][1]["prompt"] == "seg-2" + + +@pytest.mark.asyncio +async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None: + class RecordingPool: + def __init__(self) -> None: + self.calls: list[tuple[str, Any]] = [] + + async def submit_update(self, request_id, req_state, request, *, prompt_text=None) -> int: + self.calls.append((request_id, prompt_text)) + return 0 + + pool = RecordingPool() + orchestrator = object.__new__(Orchestrator) + orchestrator.request_states = { + "req-stream": OrchestratorRequestState( + request_id="req-stream", + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + } + orchestrator.stage_pools = [pool] + + await orchestrator._handle_streaming_update( + { + "request_id": "req-stream", + "prompt": SimpleNamespace(request_id="req-stream", prompt_token_ids=[1], resumable=True), + "sampling_params_list": [_sampling_params()], + "output_prompt_text": "segment-2", + } + ) + + assert pool.calls == [("req-stream", "segment-2")] + assert orchestrator.request_states["req-stream"].streaming.enabled is True + + +@pytest.mark.asyncio +async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None: + class FailingStageClient(FakeStageClient): + async def add_request_async(self, *args, **_kwargs) -> None: + raise RuntimeError("submit failed") + + class TrackingOutputProcessor(FakeOutputProcessor): + def __init__(self) -> None: + super().__init__() + self.added_request_ids: list[str] = [] + self.removed_request_ids: list[str] = [] + + def add_request(self, request, *_args, **_kwargs) -> None: + self.added_request_ids.append(request.request_id) + + def remove_request(self, request_id: str) -> None: + self.removed_request_ids.append(request_id) + + client = FailingStageClient(stage_type="llm", final_output=False) + output_processor = TrackingOutputProcessor() + pool = StagePool( + 0, + [client], + output_processor=output_processor, + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + req_state = OrchestratorRequestState( + request_id="req-0", + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + + with pytest.raises(RuntimeError, match="submit failed"): + await pool.submit_initial("req-0", req_state, SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2])) + + assert output_processor.added_request_ids == ["req-0"] + assert output_processor.removed_request_ids == ["req-0"] + assert pool.get_bound_replica_id("req-0") is None + + +@pytest.mark.asyncio +async def test_stage_pool_abort_requests_logs_when_binding_is_missing(caplog) -> None: + stage0 = FakeStageClient(stage_type="llm", final_output=False) + pool = StagePool( + 0, + [stage0], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + + target_logger = logging.getLogger("vllm_omni.engine.stage_pool") + target_logger.addHandler(caplog.handler) + prev_level = target_logger.level + target_logger.setLevel(logging.DEBUG) + try: + await pool.abort_requests(["missing-req"]) + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(prev_level) + + assert not stage0.abort_calls + assert "abort: no binding for req=missing-req in stage-0" in caplog.text + + +@pytest.mark.asyncio +async def test_collective_rpc_ignores_invalid_stage_ids(orchestrator_factory, caplog) -> None: + stage0 = FakeCollectiveRpcStageClient(stage_type="llm", final_output=True, rpc_result={"stage": 0}) + stage1 = FakeCollectiveRpcStageClient(stage_type="llm", final_output=True, rpc_result={"stage": 1}) + stage_pools = _build_stage_pools( + [[stage0], [stage1]], + output_processors=[FakeOutputProcessor(), FakeOutputProcessor()], + stage_vllm_configs=[ + SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ], + ) + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools) + + try: + target_logger = logging.getLogger("vllm_omni.engine.orchestrator") + target_logger.addHandler(caplog.handler) + prev_level = target_logger.level + target_logger.setLevel(logging.WARNING) + try: + orchestrator_fixture.request_sync_q.put_nowait( + { + "type": "collective_rpc", + "rpc_id": "rpc-1", + "method": "list_loras", + "stage_ids": [99, 1], + } + ) + + msg = await _get_rpc_message(orchestrator_fixture) + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(prev_level) + + assert msg["type"] == "collective_rpc_result" + assert msg["rpc_id"] == "rpc-1" + assert msg["stage_ids"] == [1] + assert msg["results"] == [{"stage": 1}] + assert not stage0.collective_rpc_calls + assert len(stage1.collective_rpc_calls) == 1 + assert "collective_rpc: ignoring invalid stage_id 99" in caplog.text + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator_factory) -> None: + """CFG companions should be routed to the same stage-0 replica as their parent.""" + stage0_r0 = FakeStageClient(stage_type="llm", final_output=False) + stage0_r1 = FakeStageClient(stage_type="llm", final_output=False) + default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) + stage_pools = _build_stage_pools( + [[stage0_r0, stage0_r1]], + output_processors=[FakeOutputProcessor()], + stage_vllm_configs=[default_vllm_cfg], + ) + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools) + + try: + # Consume replica-0 first so the parent request binds to replica-1. + await _enqueue_add_request( + orchestrator_fixture, + request_id="warmup", + prompt=SimpleNamespace(request_id="warmup", prompt_token_ids=[0]), + original_prompt={"prompt": "warmup"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(stage0_r0.add_request_calls) == 1) + + await _enqueue_add_request( + orchestrator_fixture, + request_id="parent", + prompt=SimpleNamespace(request_id="parent", prompt_token_ids=[1, 2]), + original_prompt={"prompt": "parent"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(stage0_r1.add_request_calls) == 1) + + orchestrator_fixture.request_sync_q.put_nowait( + { + "type": "add_companion_request", + "companion_id": "parent-neg", + "parent_id": "parent", + "role": "negative", + "prompt": SimpleNamespace(request_id="parent-neg", prompt_token_ids=[9]), + "companion_prompt_text": {"prompt": "negative"}, + "sampling_params_list": [_sampling_params()], + } + ) + await _wait_for(lambda: len(stage0_r1.add_request_calls) == 2) + + assert stage_pools[0].get_bound_replica_id("parent") == 1 + assert stage_pools[0].get_bound_replica_id("parent-neg") == 1 + assert len(stage0_r0.add_request_calls) == 1 + assert stage0_r1.add_request_calls[0][0].request_id == "parent" + assert stage0_r1.add_request_calls[1][0].request_id == "parent-neg" + finally: + await _shutdown_orchestrator(orchestrator_fixture) diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 6131c5e9b..18099c016 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -124,8 +124,8 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa @pytest.mark.asyncio async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -> None: """When a diffusion stage returns an OmniRequestOutput with a non-None - error, the orchestrator must route it as a finished output message and - clean up the request state. + error, the orchestrator must route it as an error message and clean up + the request state. """ stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image") orchestrator_fixture = orchestrator_factory([stage0]) @@ -148,10 +148,10 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) - msg = await _get_any_output_message(orchestrator_fixture) - assert msg["type"] == "output" + assert msg["type"] == "error" assert msg["request_id"] == "req-err" - assert msg["finished"] is True - assert msg["engine_outputs"].error == "gpu fault" + assert msg["stage_id"] == 0 + assert msg["error"] == "gpu fault" # Request state should be cleaned up. await _wait_for(lambda: "req-err" not in orchestrator_fixture.orchestrator.request_states) diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index ec3a42e35..2bd9b0947 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -7,18 +7,18 @@ from vllm import SamplingParams from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.engine.stage_pool import StagePool from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _DummySenderStage: + stage_type = "llm" + final_output = False + def __init__(self, sender_info): self._sender_info = sender_info - self.engine_outputs = None - - def set_engine_outputs(self, outputs): - self.engine_outputs = outputs def get_kv_sender_info(self): return self._sender_info @@ -26,6 +26,7 @@ class _DummySenderStage: class _DummyDiffusionStage: stage_type = "diffusion" + final_output = True custom_process_input_func = None def __init__(self, engine_input_source=None): @@ -43,6 +44,15 @@ class _DummyDiffusionStage: ) +def _build_sender_pool(stage_id: int, sender_info: dict[str, object]) -> StagePool: + return StagePool( + stage_id, + _DummySenderStage(sender_info), + output_processor=object(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + + def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address(): client = object.__new__(StageEngineCoreClient) client.stage_id = 0 @@ -126,14 +136,13 @@ def test_stage_engine_core_client_preserves_explicit_loopback_sender_host(): def test_forward_to_diffusion_attaches_kv_sender_info(): orchestrator = object.__new__(Orchestrator) - sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151}) diffusion_stage = _DummyDiffusionStage(engine_input_source=[0]) + sender_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151}) + diffusion_pool = StagePool(1, diffusion_stage) orchestrator.num_stages = 2 - orchestrator.stage_clients = [sender_stage, diffusion_stage] + orchestrator.stage_pools = [sender_pool, diffusion_pool] orchestrator._cfg_tracker = CfgCompanionTracker() - orchestrator.stage_vllm_configs = [None, None] - orchestrator.output_processors = [None, None] params = OmniDiffusionSamplingParams() req_state = OrchestratorRequestState( @@ -144,9 +153,8 @@ def test_forward_to_diffusion_attaches_kv_sender_info(): ) output = SimpleNamespace(request_id="req-1", finished=True) - asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", 0, output, req_state)) + asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", sender_pool.stage_id, output, req_state)) - assert sender_stage.engine_outputs == [output] assert diffusion_stage.calls[0]["request_id"] == "req-1" assert diffusion_stage.calls[0]["kv_sender_info"] == { 0: {"host": "10.0.0.2", "zmq_port": 50151}, @@ -156,15 +164,14 @@ def test_forward_to_diffusion_attaches_kv_sender_info(): def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): orchestrator = object.__new__(Orchestrator) - source_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151}) - previous_stage = _DummySenderStage({"host": "10.0.0.9", "zmq_port": 59999}) diffusion_stage = _DummyDiffusionStage(engine_input_source=[0]) + source_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151}) + previous_pool = _build_sender_pool(1, {"host": "10.0.0.9", "zmq_port": 59999}) + diffusion_pool = StagePool(2, diffusion_stage) orchestrator.num_stages = 3 - orchestrator.stage_clients = [source_stage, previous_stage, diffusion_stage] + orchestrator.stage_pools = [source_pool, previous_pool, diffusion_pool] orchestrator._cfg_tracker = CfgCompanionTracker() - orchestrator.stage_vllm_configs = [None, None, None] - orchestrator.output_processors = [None, None, None] params = OmniDiffusionSamplingParams() req_state = OrchestratorRequestState( @@ -175,9 +182,8 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): ) output = SimpleNamespace(request_id="req-3", finished=True) - asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", 1, output, req_state)) + asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", previous_pool.stage_id, output, req_state)) - assert previous_stage.engine_outputs == [output] assert diffusion_stage.calls[0]["kv_sender_info"] == { 0: {"host": "10.0.0.2", "zmq_port": 50151}, } @@ -185,9 +191,10 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs(): orchestrator = object.__new__(Orchestrator) - sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151}) diffusion_stage = _DummyDiffusionStage(engine_input_source=[0]) diffusion_stage.custom_process_input_func = lambda *_args, **_kwargs: [] + sender_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151}) + diffusion_pool = StagePool(1, diffusion_stage) class _AsyncQueue: def __init__(self): @@ -197,10 +204,8 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs(): self.items.append(item) orchestrator.num_stages = 2 - orchestrator.stage_clients = [sender_stage, diffusion_stage] + orchestrator.stage_pools = [sender_pool, diffusion_pool] orchestrator._cfg_tracker = CfgCompanionTracker() - orchestrator.stage_vllm_configs = [None, None] - orchestrator.output_processors = [None, None] orchestrator.output_async_queue = _AsyncQueue() orchestrator.request_states = {} orchestrator._pd_kv_params = {} @@ -217,7 +222,6 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs(): output = SimpleNamespace(request_id="req-empty", finished=True) asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-empty", 0, output, req_state)) - assert sender_stage.engine_outputs == [output] assert diffusion_stage.calls == [] assert len(orchestrator.output_async_queue.items) == 1 terminal_msg = orchestrator.output_async_queue.items[0] @@ -231,10 +235,11 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs(): def test_prewarm_diffusion_attaches_kv_sender_info(): orchestrator = object.__new__(Orchestrator) - sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151}) diffusion_stage = _DummyDiffusionStage(engine_input_source=[0]) + sender_pool = _build_sender_pool(0, {"host": "10.0.0.3", "zmq_port": 50151}) + diffusion_pool = StagePool(1, diffusion_stage) - orchestrator.stage_clients = [sender_stage, diffusion_stage] + orchestrator.stage_pools = [sender_pool, diffusion_pool] orchestrator.num_stages = 2 req_state = OrchestratorRequestState( diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index abc42ac9f..bcd910620 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -1,20 +1,8 @@ -"""Unit tests for AsyncOmniEngine single-stage mode and OmniMasterServer. - -These tests cover: -- OmniMasterServer address pre-allocation & ZMQ registration handshake -- AsyncOmniEngine single_stage_mode detection / _single_stage_id_filter setup -- _initialize_stages stage routing (local launch vs. remote-wait) in - single_stage_mode -- _create_remote_llm_stage delegation to connect_remote_engine_cores -- _launch_llm_stage delegation to launch_omni_core_engines in - single_stage_mode - -All tests run without real hardware by mocking ZMQ, vllm_config, and the -heavy initialization helpers. -""" +"""Unit tests for AsyncOmniEngine single-stage mode and OmniMasterServer.""" from __future__ import annotations +import os import threading from contextlib import contextmanager from types import SimpleNamespace @@ -25,6 +13,7 @@ from pytest_mock import MockerFixture from vllm.v1.engine.utils import EngineZmqAddresses from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase from vllm_omni.engine.stage_engine_startup import ( OmniMasterServer, StageAllocation, @@ -32,21 +21,17 @@ from vllm_omni.engine.stage_engine_startup import ( connect_remote_engine_cores, launch_omni_core_engines, ) -from vllm_omni.engine.stage_init_utils import StartedLlmStage +from vllm_omni.engine.stage_init_utils import LogicalStageInitPlan, ReplicaInitPlan pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - def _make_stage_cfg(stage_id: int, stage_type: str = "llm"): """Return a lightweight stage config mock.""" return SimpleNamespace( stage_id=stage_id, stage_type=stage_type, + runtime=SimpleNamespace(devices="0"), engine_args=SimpleNamespace( async_chunk=False, model_stage=None, @@ -55,67 +40,134 @@ def _make_stage_cfg(stage_id: int, stage_type: str = "llm"): ) -def _make_started_llm_stage(stage_id: int) -> StartedLlmStage: - """Return a minimal StartedLlmStage for mocking.""" - addresses = SimpleNamespace( - inputs=["tcp://127.0.0.1:5000"], - outputs=["tcp://127.0.0.1:5001"], - frontend_stats_publish_address=None, +def _make_llm_plan( + stage_idx: int, + *, + configured_stage_id: int, + launch_mode: str, + vllm_config: Any | None = None, +) -> LogicalStageInitPlan: + stage_cfg = _make_stage_cfg(configured_stage_id) + metadata = SimpleNamespace( + stage_id=configured_stage_id, + stage_type="llm", + runtime_cfg={"devices": "0"}, + prompt_expand_func=None, + final_output=False, + final_output_type=None, + default_sampling_params=SimpleNamespace(), + custom_process_input_func=None, + engine_input_source=[], + engine_output_type="token_ids", + replica_id=0, ) - return StartedLlmStage( - stage_id=stage_id, - metadata=SimpleNamespace(stage_id=stage_id), - vllm_config=SimpleNamespace(), - executor_class=SimpleNamespace(), - engine_manager=SimpleNamespace(), - coordinator=SimpleNamespace(), - addresses=addresses, + return LogicalStageInitPlan( + stage_idx=stage_idx, + configured_stage_id=configured_stage_id, + replicas=[ + ReplicaInitPlan( + replica_id=0, + num_replicas=1, + launch_mode=launch_mode, + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + omni_kv_connector=(None, None, None), + stage_vllm_config=vllm_config + or SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)), + executor_class=object, + ) + ], + ) + + +def _make_diffusion_plan( + stage_idx: int, + *, + configured_stage_id: int, + launch_mode: str, +) -> LogicalStageInitPlan: + stage_cfg = _make_stage_cfg(configured_stage_id, stage_type="diffusion") + metadata = SimpleNamespace( + stage_id=configured_stage_id, + stage_type="diffusion", + runtime_cfg={"devices": "0"}, + prompt_expand_func=None, + final_output=True, + final_output_type="image", + default_sampling_params=SimpleNamespace(), + custom_process_input_func=None, + engine_input_source=[], + cfg_kv_collect_func=None, + replica_id=0, + ) + return LogicalStageInitPlan( + stage_idx=stage_idx, + configured_stage_id=configured_stage_id, + replicas=[ + ReplicaInitPlan( + replica_id=0, + num_replicas=1, + launch_mode=launch_mode, + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + omni_kv_connector=(None, None, None), + ) + ], ) # --------------------------------------------------------------------------- -# OmniMasterServer – address pre-allocation +# OmniMasterServer address pre-allocation # --------------------------------------------------------------------------- class TestOmniMasterServerAllocation: - """Test address pre-allocation in OmniMasterServer.__init__.""" - def test_public_address_and_port_properties_expose_registration_endpoint(self): - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=15000, - stage_ids=[0], - ) + server = OmniMasterServer(master_address="127.0.0.1", master_port=15000, stage_ids=[0]) assert server.address == "127.0.0.1" assert server.port == 15000 def test_allocations_created_for_each_stage_id(self): - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=15000, - stage_ids=[0, 1, 2], - ) - assert set(server._allocations.keys()) == {0, 1, 2} + server = OmniMasterServer(master_address="127.0.0.1", master_port=15000, stage_ids=[0, 1, 2]) + assert set(server._stage_routes.keys()) == {(0, 0), (1, 0), (2, 0)} def test_each_allocation_is_stage_allocation(self): + server = OmniMasterServer(master_address="127.0.0.1", master_port=15000, stage_ids=[0, 1]) + for sid in (0, 1): + assert isinstance(server.get_allocation(sid), StageAllocation) + + def test_replica_allocations_are_distinct(self): server = OmniMasterServer( master_address="127.0.0.1", master_port=15000, - stage_ids=[0, 1], + stage_ids=[0], + stage_replica_counts={0: 3}, ) - for sid in (0, 1): - alloc = server._allocations[sid] - assert isinstance(alloc, StageAllocation) + + allocations = [server.get_allocation(0, replica_id) for replica_id in range(3)] + assert len({alloc.handshake_bind_address for alloc in allocations}) == 3 + assert server.get_allocation(0) is allocations[0] + + def test_replica_stage_configs_are_isolated(self): + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15000, + stage_ids=[0], + stage_replica_counts={0: 2}, + ) + + server.register_stage_config(0, {"replica": 0}, replica_id=0) + server.register_stage_config(0, {"replica": 1}, replica_id=1) + + assert server.get_stage_config(0, replica_id=0) == {"replica": 0} + assert server.get_stage_config(0, replica_id=1) == {"replica": 1} def test_allocation_addresses_reference_master_address(self): - server = OmniMasterServer( - master_address="192.168.1.10", - master_port=20000, - stage_ids=[0], - ) - alloc = server._allocations[0] - for addr in ( + server = OmniMasterServer(master_address="192.168.1.10", master_port=20000, stage_ids=[0]) + alloc = server.get_allocation(0) + for address in ( alloc.handshake_bind_address, alloc.handshake_connect_address, alloc.input_bind_address, @@ -123,83 +175,57 @@ class TestOmniMasterServerAllocation: alloc.output_bind_address, alloc.output_connect_address, ): - assert "192.168.1.10" in addr, f"Expected master address in {addr}" + assert "192.168.1.10" in address def test_port_uniqueness_within_single_allocation(self): - """Each allocation uses three distinct ports.""" - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=15001, - stage_ids=[0], - ) - alloc = server._allocations[0] - hs_port = int(alloc.handshake_bind_address.split(":")[-1]) - inp_port = int(alloc.input_bind_address.split(":")[-1]) - out_port = int(alloc.output_bind_address.split(":")[-1]) - assert len({hs_port, inp_port, out_port}) == 3, "Expected three distinct ports per stage allocation" + server = OmniMasterServer(master_address="127.0.0.1", master_port=15001, stage_ids=[0]) + alloc = server.get_allocation(0) + handshake_port = int(alloc.handshake_bind_address.split(":")[-1]) + input_port = int(alloc.input_bind_address.split(":")[-1]) + output_port = int(alloc.output_bind_address.split(":")[-1]) + assert len({handshake_port, input_port, output_port}) == 3 def test_get_zmq_addresses_returns_bind_addresses(self): - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=15002, - stage_ids=[0], - ) - alloc = server._allocations[0] + server = OmniMasterServer(master_address="127.0.0.1", master_port=15002, stage_ids=[0]) + alloc = server.get_allocation(0) zmq_addrs = server.get_zmq_addresses(0) assert zmq_addrs.inputs == [alloc.input_bind_address] assert zmq_addrs.outputs == [alloc.output_bind_address] def test_get_engine_zmq_addresses_returns_connect_addresses(self): - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=15003, - stage_ids=[0], - ) - alloc = server._allocations[0] - engine_addrs = server.get_engine_zmq_addresses(0) - assert engine_addrs.inputs == [alloc.input_connect_address] - assert engine_addrs.outputs == [alloc.output_connect_address] + server = OmniMasterServer(master_address="127.0.0.1", master_port=15003, stage_ids=[0]) + alloc = server.get_allocation(0) + zmq_addrs = server.get_engine_zmq_addresses(0) + assert zmq_addrs.inputs == [alloc.input_connect_address] + assert zmq_addrs.outputs == [alloc.output_connect_address] def test_get_allocation_returns_correct_object(self): - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=15004, - stage_ids=[3], - ) - assert server.get_allocation(3) is server._allocations[3] + server = OmniMasterServer(master_address="127.0.0.1", master_port=15004, stage_ids=[3]) + assert server.get_allocation(3) is server._stage_routes[(3, 0)] # --------------------------------------------------------------------------- -# OmniMasterServer – ZMQ registration flow +# OmniMasterServer registration flow # --------------------------------------------------------------------------- class TestOmniMasterServerRegistration: - """Test that the server correctly handles a stage registration.""" - def test_registration_reply_contains_handshake_address(self): - """A DEALER client that sends a registration msg gets the handshake - address back from the ROUTER registration socket.""" import msgspec import zmq from vllm.utils.network_utils import get_open_port master_port = get_open_port() - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=master_port, - stage_ids=[0], - ) + server = OmniMasterServer(master_address="127.0.0.1", master_port=master_port, stage_ids=[0]) server.start() - expected_hs = server._allocations[0].handshake_connect_address + expected_hs = server.get_allocation(0).handshake_connect_address ctx = zmq.Context() try: sock = ctx.socket(zmq.DEALER) sock.connect(f"tcp://127.0.0.1:{master_port}") sock.send(msgspec.msgpack.encode({"stage_id": 0})) - if not sock.poll(timeout=5_000): - pytest.fail("No reply received from OmniMasterServer within 5 s") + assert sock.poll(timeout=5_000) reply = msgspec.msgpack.decode(sock.recv()) assert reply["handshake_address"] == expected_hs finally: @@ -208,92 +234,65 @@ class TestOmniMasterServerRegistration: server.stop() def test_server_handles_unknown_stage_id_gracefully(self): - """A registration for an unrecognised stage_id must not crash the server.""" import msgspec import zmq from vllm.utils.network_utils import get_open_port master_port = get_open_port() - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=master_port, - stage_ids=[0], - ) + server = OmniMasterServer(master_address="127.0.0.1", master_port=master_port, stage_ids=[0]) server.start() ctx = zmq.Context() + bad_sock = None + good_sock = None try: bad_sock = ctx.socket(zmq.DEALER) bad_sock.connect(f"tcp://127.0.0.1:{master_port}") - # Send unknown stage_id=99 bad_sock.send(msgspec.msgpack.encode({"stage_id": 99})) - # Server should NOT reply for an unknown id; wait briefly - has_reply = bad_sock.poll(timeout=500) - assert not has_reply, "Server should not reply to unknown stage_id" - # Then register the valid stage so the server thread can exit + assert not bad_sock.poll(timeout=500) + good_sock = ctx.socket(zmq.DEALER) good_sock.connect(f"tcp://127.0.0.1:{master_port}") good_sock.send(msgspec.msgpack.encode({"stage_id": 0})) - good_sock.poll(timeout=2_000) + assert good_sock.poll(timeout=2_000) + good_sock.recv() finally: - for s in (bad_sock, good_sock): - try: - s.close(linger=0) - except Exception: - pass + for sock in (bad_sock, good_sock): + if sock is not None: + sock.close(linger=0) ctx.term() server.stop() def test_registration_stores_stage_config(self): - """Stage registration should persist the sender's stage config.""" import msgspec import zmq from vllm.utils.network_utils import get_open_port master_port = get_open_port() - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=master_port, - stage_ids=[0], - ) + server = OmniMasterServer(master_address="127.0.0.1", master_port=master_port, stage_ids=[0]) server.start() - payload = { - "stage_id": 0, - "stage_config": { - "stage_id": 0, - "stage_type": "llm", - "engine_args": {"model": "fake-model"}, - }, - } - + stage_config = {"stage_id": 0, "stage_type": "llm"} ctx = zmq.Context() try: sock = ctx.socket(zmq.DEALER) sock.connect(f"tcp://127.0.0.1:{master_port}") - sock.send(msgspec.msgpack.encode(payload)) + sock.send(msgspec.msgpack.encode({"stage_id": 0, "stage_config": stage_config})) assert sock.poll(timeout=5_000) sock.recv() - - stored = server.get_stage_config(0, timeout_s=0.1) - assert stored == payload["stage_config"] + assert server.get_stage_config(0) == stage_config finally: sock.close(linger=0) ctx.term() server.stop() def test_registration_stores_coordinator_addresses(self): - """Stage registration should persist optional coordinator addresses.""" import msgspec import zmq from vllm.utils.network_utils import get_open_port master_port = get_open_port() - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=master_port, - stage_ids=[0], - ) + server = OmniMasterServer(master_address="127.0.0.1", master_port=master_port, stage_ids=[0]) server.start() payload = { @@ -303,7 +302,6 @@ class TestOmniMasterServerRegistration: "coordinator_output": "tcp://127.0.0.1:31002", "frontend_stats_publish_address": "tcp://127.0.0.1:31003", } - ctx = zmq.Context() try: sock = ctx.socket(zmq.DEALER) @@ -311,9 +309,7 @@ class TestOmniMasterServerRegistration: sock.send(msgspec.msgpack.encode(payload)) assert sock.poll(timeout=5_000) sock.recv() - - stored = server.get_stage_coordinator_addresses(0, timeout_s=0.1) - assert stored == StageCoordinatorAddresses( + assert server.get_stage_coordinator_addresses(0) == StageCoordinatorAddresses( coordinator_input=payload["coordinator_input"], coordinator_output=payload["coordinator_output"], frontend_stats_publish_address=payload["frontend_stats_publish_address"], @@ -327,57 +323,47 @@ class TestOmniMasterServerRegistration: from vllm.utils.network_utils import get_open_port master_port = get_open_port() - server = OmniMasterServer( - master_address="127.0.0.1", - master_port=master_port, - stage_ids=[], # no stages → thread exits immediately - ) + server = OmniMasterServer(master_address="127.0.0.1", master_port=master_port, stage_ids=[]) server.start() + assert server._thread is not None server.stop() - # Thread should have exited (joined with timeout=10 inside stop()) assert not server._thread.is_alive() # --------------------------------------------------------------------------- -# AsyncOmniEngine – single_stage_mode detection in __init__ +# AsyncOmniEngine single_stage_mode detection in __init__ # --------------------------------------------------------------------------- class TestSingleStageModeDetection: - """Test __init__ single_stage_mode / _single_stage_id_filter setup. - - We bypass the real __init__ by patching _resolve_stage_configs and - the orchestrator thread, so no actual engines are started. - """ - - def _make_engine_no_thread(self, mocker: MockerFixture, **kwargs: Any) -> AsyncOmniEngine: - """Create an AsyncOmniEngine without starting the orchestrator thread.""" - stage_cfg = _make_stage_cfg(0) - mock_stage_configs = [stage_cfg] + def _make_engine_no_thread( + self, + mocker: MockerFixture, + *, + stage_cfgs: list[Any] | None = None, + **kwargs: Any, + ) -> AsyncOmniEngine: + mock_stage_configs = stage_cfgs or [_make_stage_cfg(0)] mocker.patch.object( AsyncOmniEngine, "_resolve_stage_configs", return_value=("/fake/path", mock_stage_configs), ) - mocker.patch.object( - AsyncOmniEngine, - "_bootstrap_orchestrator", - ) + mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator") mock_thread_cls = mocker.patch("threading.Thread") mock_future_cls = mocker.patch("concurrent.futures.Future") mock_future = mocker.Mock() - mock_future.result.return_value = mocker.Mock() # simulates a loop + mock_future.result.return_value = mocker.Mock() mock_future_cls.return_value = mock_future mock_thread = mocker.Mock() mock_thread.is_alive.return_value = False mock_thread_cls.return_value = mock_thread - engine = AsyncOmniEngine(model="fake-model", **kwargs) - return engine + return AsyncOmniEngine(model="fake-model", **kwargs) def test_explicit_single_stage_mode_true(self, mocker: MockerFixture): engine = self._make_engine_no_thread( @@ -407,9 +393,7 @@ class TestSingleStageModeDetection: assert engine._single_stage_id_filter == 1 def test_no_stage_id_no_single_stage_mode(self, mocker: MockerFixture): - engine = self._make_engine_no_thread( - mocker, - ) + engine = self._make_engine_no_thread(mocker) assert engine.single_stage_mode is False assert engine._single_stage_id_filter is None @@ -420,6 +404,7 @@ class TestSingleStageModeDetection: omni_master_address="127.0.0.1", omni_master_port=20003, ) + assert engine.single_stage_mode is True assert engine._single_stage_id_filter is None def test_engine_args_create_only_forwards_explicit_fields(self, mocker: MockerFixture): @@ -466,907 +451,569 @@ class TestSingleStageModeDetection: assert engine._omni_master_port == 12345 def test_omni_master_server_starts_as_none(self, mocker: MockerFixture): - engine = self._make_engine_no_thread( - mocker, - ) + engine = self._make_engine_no_thread(mocker) assert engine._omni_master_server is None # --------------------------------------------------------------------------- -# AsyncOmniEngine – _initialize_stages stage routing +# AsyncOmniEngine single-stage initialization paths # --------------------------------------------------------------------------- -class TestInitializeStagesRouting: - """Verify that _initialize_stages routes each stage to the correct launch - function depending on single_stage_mode and _single_stage_id_filter.""" - - _COMMON_PATCHES = [ - "vllm_omni.engine.async_omni_engine.prepare_engine_environment", - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - ] - - def _build_engine_skeleton( - self, - stage_cfgs: list[Any], - single_stage_mode: bool, - stage_id_filter: int | None, - omni_master_address: str = "127.0.0.1", - omni_master_port: int = 25000, +class TestSingleStageInitialization: + def _build_engine( + self, stage_cfgs: list[Any], *, single_stage_mode: bool, stage_id_filter: int | None ) -> AsyncOmniEngine: - """Build a bare AsyncOmniEngine without launching any threads.""" engine = object.__new__(AsyncOmniEngine) engine.model = "fake-model" - engine.config_path = "/fake" - engine.stage_configs = stage_cfgs + engine.config_path = "/fake/stages.yaml" engine.num_stages = len(stage_cfgs) - engine.async_chunk = False + engine.stage_configs = stage_cfgs engine.single_stage_mode = single_stage_mode engine._single_stage_id_filter = stage_id_filter - engine._omni_master_address = omni_master_address - engine._omni_master_port = omni_master_port + engine._omni_master_address = "127.0.0.1" + engine._omni_master_port = 26000 engine._omni_master_server = None - engine._llm_stage_launch_lock = __import__("threading").Lock() - engine.diffusion_batch_size = 1 - engine.stage_clients = [] - engine.stage_vllm_configs = [] - engine.output_processors = [] - engine.input_processor = None - engine.supported_tasks = ("generate",) - engine.default_sampling_params_list = [] - engine.stage_metadata = [] - engine.prompt_expand_func = None + engine.async_chunk = False + engine.diffusion_batch_size = 2 return engine - def _fake_metadata(self, mocker: MockerFixture, stage_id: int, stage_type: str = "llm") -> Any: - meta = mocker.Mock() - meta.stage_id = stage_id - meta.stage_type = stage_type - meta.runtime_cfg = {} - meta.prompt_expand_func = None - meta.engine_output_type = None - meta.is_comprehension = False - meta.final_output = True if stage_id == 0 else False - meta.final_output_type = None - return meta + def test_build_logical_stage_init_plans_marks_non_matching_stage_remote(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod - def _run_initialize_stages_mocked( - self, - mocker: MockerFixture, - engine: AsyncOmniEngine, - stage_cfgs: list[Any], - *, - launch_side_effect: Any = None, - remote_side_effect: Any = None, - attach_result: Any = None, - ) -> tuple[Any, Any]: - """Execute _initialize_stages with all heavy helpers mocked. + stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11)] + engine = self._build_engine(stage_cfgs, single_stage_mode=True, stage_id_filter=7) - Returns (mock_launch_llm_stage, mock_create_remote_llm_stage). - """ - started_by_stage: dict[int, StartedLlmStage] = { - cfg.stage_id: _make_started_llm_stage(cfg.stage_id) - for cfg in stage_cfgs - if getattr(cfg, "stage_type", "llm") != "diffusion" - } - - default_attach = (mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()) - - mock_launch = mocker.Mock( - side_effect=launch_side_effect - or (lambda cfg, meta, spec, timeout, llm_stage_launch_lock, kv: started_by_stage[meta.stage_id]) - ) - mock_remote = mocker.Mock( - side_effect=remote_side_effect or (lambda cfg, meta, spec, timeout, srv: started_by_stage[meta.stage_id]) - ) - mock_attach = mocker.Mock(return_value=attach_result or default_attach) - - mock_oms = mocker.Mock(spec=OmniMasterServer) - mock_oms.get_zmq_addresses.side_effect = lambda sid: mocker.Mock() - - finalized = ( - [mocker.Mock() for _ in stage_cfgs], - [mocker.Mock() for _ in stage_cfgs], - [{"final_output": True, "final_output_type": None, "stage_type": "llm"} for _ in stage_cfgs], - ) - - mocker.patch.object(engine, "_launch_llm_stage", mock_launch) - mocker.patch.object(engine, "_create_remote_llm_stage", mock_remote) - mocker.patch.object(engine, "_attach_llm_stage", mock_attach) - mocker.patch( - "vllm_omni.engine.async_omni_engine.OmniMasterServer", - return_value=mock_oms, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.prepare_engine_environment", - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - return_value={}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata( - mocker, - cfg.stage_id, - getattr(cfg, "stage_type", "llm"), + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr( + engine_mod, + "extract_stage_metadata", + lambda cfg: SimpleNamespace( + stage_id=cfg.stage_id, + stage_type=getattr(cfg, "stage_type", "llm"), + prompt_expand_func=None, + runtime_cfg={}, ), ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - return_value=finalized, - ) + monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (SimpleNamespace(), object)) + try: + stage_plans, _ = engine._build_logical_stage_init_plans(None, [1, 1], {}) + finally: + monkeypatch.undo() - engine._initialize_stages(stage_init_timeout=60) + assert [plan.replicas[0].launch_mode for plan in stage_plans] == ["local", "remote"] - return mock_launch, mock_remote + def test_start_omni_master_server_uses_configured_stage_ids(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod - # -- single-stage mode: stage matches filter → local launch --------------- - - def test_matching_stage_uses_launch_llm_stage(self, mocker: MockerFixture): - """stage_id == _single_stage_id_filter → _launch_llm_stage is called.""" - stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) - - launched_ids = [c.args[1].stage_id for c in mock_launch.call_args_list] - assert 0 in launched_ids, "_launch_llm_stage should be called for stage 0" - - def test_non_matching_stage_uses_create_remote_llm_stage(self, mocker: MockerFixture): - """stage_id != _single_stage_id_filter → _create_remote_llm_stage is called.""" - stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) - - remote_ids = [c.args[1].stage_id for c in mock_remote.call_args_list] - assert 1 in remote_ids, "_create_remote_llm_stage should be called for stage 1" - - def test_filter_1_routes_correctly(self, mocker: MockerFixture): - """With filter=1, stage 0 is remote and stage 1 is local.""" - stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=1) - mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) - - launched_ids = [c.args[1].stage_id for c in mock_launch.call_args_list] - remote_ids = [c.args[1].stage_id for c in mock_remote.call_args_list] - assert 1 in launched_ids, "stage 1 should be launched locally with filter=1" - assert 0 in remote_ids, "stage 0 should use remote path with filter=1" - - def test_no_filter_all_stages_use_launch_path(self, mocker: MockerFixture): - """single_stage_mode=True but no filter → all stages use _launch_llm_stage.""" - stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=None) - mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) - - assert mock_remote.call_count == 0, "No remote launches without a filter" - launched_ids = [c.args[1].stage_id for c in mock_launch.call_args_list] - assert set(launched_ids) == {0, 1} - - def test_non_single_stage_mode_never_calls_create_remote(self, mocker: MockerFixture): - """Outside single_stage_mode, _create_remote_llm_stage must not be called.""" - stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=False, stage_id_filter=None) - mock_launch, mock_remote = self._run_initialize_stages_mocked(mocker, engine, stage_cfgs) - - assert mock_remote.call_count == 0 - - def test_omni_master_server_started_in_single_stage_mode(self, mocker: MockerFixture): - """OmniMasterServer.start() must be called when single_stage_mode=True.""" - stage_cfgs = [_make_stage_cfg(0)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) + engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) mock_oms = mocker.Mock(spec=OmniMasterServer) - mock_oms.get_zmq_addresses.return_value = mocker.Mock() - finalized = ( - [mocker.Mock()], - [mocker.Mock()], - [{"final_output": True, "final_output_type": None, "stage_type": "llm"}], - ) + mocker.patch.object(engine_mod, "OmniMasterServer", return_value=mock_oms) - mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)) - mocker.patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(0)) - mocker.patch.object( - engine, - "_attach_llm_stage", - return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.OmniMasterServer", - return_value=mock_oms, - ) - mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - return_value={}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(mocker, cfg.stage_id), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - return_value=finalized, - ) + stage_plans = [ + _make_llm_plan(0, configured_stage_id=7, launch_mode="local"), + _make_diffusion_plan(1, configured_stage_id=11, launch_mode="remote"), + ] - engine._initialize_stages(stage_init_timeout=60) + engine._start_omni_master_server(stage_plans) + engine_mod.OmniMasterServer.assert_called_once_with( + master_address="127.0.0.1", + master_port=26000, + stage_ids=[7, 11], + stage_replica_counts={7: 1, 11: 1}, + ) mock_oms.start.assert_called_once() - def test_omni_master_server_uses_configured_stage_ids(self, mocker: MockerFixture): - """Configured stage IDs, not list indexes, should drive pre-allocation.""" - stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=7) - mock_oms = mocker.Mock(spec=OmniMasterServer) - mock_oms.get_zmq_addresses.return_value = mocker.Mock() - finalized = ( - [mocker.Mock(), mocker.Mock()], - [mocker.Mock(), mocker.Mock()], - [{"final_output": False, "final_output_type": None, "stage_type": "llm"} for _ in stage_cfgs], - ) + def test_start_omni_master_server_duplicate_stage_ids_raise(self): + engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) + stage_plans = [ + _make_llm_plan(0, configured_stage_id=7, launch_mode="local"), + _make_llm_plan(1, configured_stage_id=7, launch_mode="remote"), + ] - mocker.patch.object( - engine, - "_launch_llm_stage", - side_effect=[_make_started_llm_stage(7), _make_started_llm_stage(11)], - ) - mocker.patch.object( - engine, - "_create_remote_llm_stage", - return_value=_make_started_llm_stage(11), - ) - mocker.patch.object( - engine, - "_attach_llm_stage", - return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), - ) - mock_oms_cls = mocker.patch( - "vllm_omni.engine.async_omni_engine.OmniMasterServer", - return_value=mock_oms, - ) - mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - return_value={}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(mocker, cfg.stage_id), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - return_value=finalized, - ) - - engine._initialize_stages(stage_init_timeout=60) - - mock_oms_cls.assert_called_once_with( - master_address=engine._omni_master_address, - master_port=engine._omni_master_port, - stage_ids=[7, 11], - ) - - def test_single_stage_filter_uses_configured_stage_ids(self, mocker: MockerFixture): - """Local/remote dispatch should compare against configured stage IDs.""" - stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=7) - mock_oms = mocker.Mock(spec=OmniMasterServer) - finalized = ( - [mocker.Mock(), mocker.Mock()], - [mocker.Mock(), mocker.Mock()], - [{"final_output": False, "final_output_type": None, "stage_type": "llm"} for _ in stage_cfgs], - ) - - mock_launch = mocker.patch.object( - engine, - "_launch_llm_stage", - side_effect=[_make_started_llm_stage(7)], - ) - mock_remote = mocker.patch.object( - engine, - "_create_remote_llm_stage", - return_value=_make_started_llm_stage(11), - ) - mocker.patch.object( - engine, - "_attach_llm_stage", - return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.OmniMasterServer", - return_value=mock_oms, - ) - mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - return_value={}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(mocker, cfg.stage_id), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - return_value=finalized, - ) - - engine._initialize_stages(stage_init_timeout=60) - - assert [call.args[1].stage_id for call in mock_launch.call_args_list] == [7] - assert [call.args[1].stage_id for call in mock_remote.call_args_list] == [11] - - def test_omni_master_server_preallocates_diffusion_stage_ids(self, mocker: MockerFixture): - """Diffusion stages should also receive OmniMasterServer allocations.""" - stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11, stage_type="diffusion")] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=7) - mock_oms = mocker.Mock(spec=OmniMasterServer) - finalized = ( - [mocker.Mock(), mocker.Mock()], - [mocker.Mock(), mocker.Mock()], - [ - {"final_output": False, "final_output_type": None, "stage_type": "llm"}, - {"final_output": False, "final_output_type": None, "stage_type": "diffusion"}, - ], - ) - - mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(7)) - mocker.patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(7)) - mocker.patch.object(engine, "_launch_diffusion_stage", return_value=mocker.Mock()) - mocker.patch.object( - engine, - "_create_remote_diffusion_stage", - return_value=mocker.Mock(), - ) - mocker.patch.object( - engine, - "_attach_llm_stage", - return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), - ) - mock_oms_cls = mocker.patch( - "vllm_omni.engine.async_omni_engine.OmniMasterServer", - return_value=mock_oms, - ) - mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - return_value={}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata( - mocker, - cfg.stage_id, - getattr(cfg, "stage_type", "llm"), - ), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - return_value=finalized, - ) - - engine._initialize_stages(stage_init_timeout=60) - - mock_oms_cls.assert_called_once_with( - master_address=engine._omni_master_address, - master_port=engine._omni_master_port, - stage_ids=[7, 11], - ) - - def test_duplicate_llm_stage_ids_raise(self, mocker: MockerFixture): - """Duplicate configured LLM stage IDs should fail fast.""" - stage_cfgs = [_make_stage_cfg(3), _make_stage_cfg(3)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=3) - - mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) with pytest.raises(ValueError, match="Duplicate stage_id"): - engine._initialize_stages(stage_init_timeout=60) + engine._start_omni_master_server(stage_plans) - def test_omni_master_server_not_started_in_normal_mode(self, mocker: MockerFixture): - """OmniMasterServer must NOT be instantiated outside single_stage_mode.""" - stage_cfgs = [_make_stage_cfg(0)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=False, stage_id_filter=None) - finalized = ( - [mocker.Mock()], - [mocker.Mock()], - [{"final_output": True, "final_output_type": None, "stage_type": "llm"}], - ) + def test_start_omni_master_server_missing_address_raises(self): + engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) + engine._omni_master_address = None - mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)) - mocker.patch.object( - engine, - "_attach_llm_stage", - return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), - ) - mock_oms_cls = mocker.patch("vllm_omni.engine.async_omni_engine.OmniMasterServer") - mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - return_value={}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata(mocker, cfg.stage_id), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - return_value=finalized, - ) + with pytest.raises(ValueError, match="requires both"): + engine._start_omni_master_server([_make_llm_plan(0, configured_stage_id=7, launch_mode="local")]) - engine._initialize_stages(stage_init_timeout=60) + def test_build_logical_stage_init_plans_clears_runtime_cfg_in_single_stage_mode(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod - mock_oms_cls.assert_not_called() + engine = self._build_engine([_make_stage_cfg(7)], single_stage_mode=True, stage_id_filter=7) - def test_single_stage_mode_missing_master_address_raises(self, mocker: MockerFixture): - """single_stage_mode without master address/port raises ValueError.""" - stage_cfgs = [_make_stage_cfg(0)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - engine._omni_master_address = None # missing - engine._omni_master_port = None - - mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) - with pytest.raises(ValueError, match="omni_master_address"): - engine._initialize_stages(stage_init_timeout=60) - - def test_matching_diffusion_stage_uses_local_registered_launch(self, mocker: MockerFixture): - """A local diffusion stage should use the registered single-stage launch path.""" - stage_cfgs = [_make_stage_cfg(0, stage_type="diffusion"), _make_stage_cfg(1)] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - mock_oms = mocker.Mock(spec=OmniMasterServer) - diffusion_client = mocker.Mock(stage_type="diffusion") - finalized = ( - [diffusion_client, mocker.Mock()], - [mocker.Mock(), mocker.Mock()], - [ - {"final_output": False, "final_output_type": None, "stage_type": "diffusion"}, - {"final_output": False, "final_output_type": None, "stage_type": "llm"}, - ], - ) - - mock_local_diff = mocker.patch.object( - engine, - "_launch_diffusion_stage", - return_value=diffusion_client, - ) - mock_remote_diff = mocker.patch.object(engine, "_create_remote_diffusion_stage") - mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(1)) - mocker.patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(1)) - mocker.patch.object( - engine, - "_attach_llm_stage", - return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.OmniMasterServer", - return_value=mock_oms, - ) - mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - return_value={}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata( - mocker, - cfg.stage_id, - getattr(cfg, "stage_type", "llm"), + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr( + engine_mod, + "extract_stage_metadata", + lambda cfg: SimpleNamespace( + stage_id=cfg.stage_id, + stage_type="llm", + prompt_expand_func=None, + runtime_cfg={"devices": "0"}, ), ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - return_value=finalized, + monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr( + engine_mod, + "build_vllm_config", + lambda *_, **__: (SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)), object), ) + try: + stage_plans, _ = engine._build_logical_stage_init_plans(None, [1], {}) + finally: + monkeypatch.undo() - engine._initialize_stages(stage_init_timeout=60) + assert stage_plans[0].replicas[0].metadata.runtime_cfg is None - assert mock_local_diff.call_count == 1 - assert mock_local_diff.call_args.args[1].stage_id == 0 - mock_remote_diff.assert_not_called() + def test_validate_single_stage_mode_allows_diffusion_replicas(self): + stage_cfg = _make_stage_cfg(0, stage_type="diffusion") + stage_cfg.runtime.num_replicas = 2 + engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=0) - def test_non_matching_diffusion_stage_uses_remote_diffusion_client(self, mocker: MockerFixture): - """A non-local diffusion stage should attach via the remote diffusion path.""" - stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1, stage_type="diffusion")] - engine = self._build_engine_skeleton(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - mock_oms = mocker.Mock(spec=OmniMasterServer) - remote_diffusion_client = mocker.Mock(stage_type="diffusion") - finalized = ( - [mocker.Mock(), remote_diffusion_client], - [mocker.Mock(), mocker.Mock()], - [ - {"final_output": False, "final_output_type": None, "stage_type": "llm"}, - {"final_output": False, "final_output_type": None, "stage_type": "diffusion"}, - ], - ) + engine._validate_single_stage_mode_replica_constraints() - mock_local_diff = mocker.patch.object(engine, "_launch_diffusion_stage") - mock_remote_diff = mocker.patch.object( - engine, - "_create_remote_diffusion_stage", - return_value=remote_diffusion_client, - ) - mocker.patch.object(engine, "_launch_llm_stage", return_value=_make_started_llm_stage(0)) - mocker.patch.object(engine, "_create_remote_llm_stage", return_value=_make_started_llm_stage(0)) - mocker.patch.object( - engine, - "_attach_llm_stage", - return_value=(mocker.Mock(), mocker.Mock(), mocker.Mock(), mocker.Mock()), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.OmniMasterServer", - return_value=mock_oms, - ) - mocker.patch("vllm_omni.engine.async_omni_engine.prepare_engine_environment") - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_omni_transfer_config_for_model", - return_value=None, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.get_stage_connector_spec", - return_value={}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.extract_stage_metadata", - side_effect=lambda cfg: self._fake_metadata( - mocker, - cfg.stage_id, - getattr(cfg, "stage_type", "llm"), + def test_validate_single_stage_mode_rejects_llm_replicas(self): + stage_cfg = _make_stage_cfg(0, stage_type="llm") + stage_cfg.runtime.num_replicas = 2 + engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=0) + + with pytest.raises(ValueError, match="only supports num_replicas > 1 for diffusion"): + engine._validate_single_stage_mode_replica_constraints() + + def test_build_logical_stage_init_plans_preserves_diffusion_runtime_cfg_in_single_stage_mode( + self, mocker: MockerFixture + ): + import vllm_omni.engine.async_omni_engine as engine_mod + + stage_cfg = _make_stage_cfg(7, stage_type="diffusion") + stage_cfg.runtime.devices = "0,1" + engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=7) + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr( + engine_mod, + "extract_stage_metadata", + lambda cfg: SimpleNamespace( + stage_id=cfg.stage_id, + stage_type="diffusion", + prompt_expand_func=None, + runtime_cfg={"devices": cfg.runtime.devices}, + final_output=True, + final_output_type="image", + default_sampling_params=SimpleNamespace(), + custom_process_input_func=None, + engine_input_source=[], + cfg_kv_collect_func=None, + replica_id=0, ), ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.finalize_initialized_stages", - return_value=finalized, + monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + try: + stage_plans, _ = engine._build_logical_stage_init_plans(None, [2], {0: ["0", "1"]}) + finally: + monkeypatch.undo() + + replicas = stage_plans[0].replicas + assert [replica.replica_id for replica in replicas] == [0, 1] + assert [replica.stage_cfg.runtime.devices for replica in replicas] == ["0", "1"] + assert [replica.metadata.runtime_cfg for replica in replicas] == [{"devices": "0"}, {"devices": "1"}] + + def test_initialize_stages_calls_master_server_only_in_single_stage_mode(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod + + stage_cfgs = [_make_stage_cfg(0)] + engine = self._build_engine(stage_cfgs, single_stage_mode=True, stage_id_filter=0) + stage_plan = _make_llm_plan(0, configured_stage_id=0, launch_mode="local") + client = SimpleNamespace( + stage_type="llm", + is_comprehension=False, + final_output=True, + final_output_type=None, + default_sampling_params=SimpleNamespace(), ) + mocker.patch.object(engine_mod, "prepare_engine_environment") + mocker.patch.object(engine_mod, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(engine_mod, "compute_replica_layout", return_value=([1], {})) + mocker.patch.object(engine, "_build_logical_stage_init_plans", return_value=([stage_plan], None)) + mock_start = mocker.patch.object(engine, "_start_omni_master_server") + mocker.patch.object(engine, "_initialize_stage_replicas", return_value={0: [client]}) + mocker.patch.object(engine_mod, "build_stage0_input_processor", return_value=object()) + mocker.patch.object(engine_mod, "build_llm_stage_output_processor", return_value=object()) + engine._initialize_stages(stage_init_timeout=60) + mock_start.assert_called_once() - mock_local_diff.assert_not_called() - assert mock_remote_diff.call_count == 1 - assert mock_remote_diff.call_args.args[0].stage_id == 1 + engine = self._build_engine(stage_cfgs, single_stage_mode=False, stage_id_filter=None) + mocker.patch.object(engine_mod, "prepare_engine_environment") + mocker.patch.object(engine_mod, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(engine_mod, "compute_replica_layout", return_value=([1], {})) + mocker.patch.object(engine, "_build_logical_stage_init_plans", return_value=([stage_plan], None)) + mock_start = mocker.patch.object(engine, "_start_omni_master_server") + mocker.patch.object(engine, "_initialize_stage_replicas", return_value={0: [client]}) + mocker.patch.object(engine_mod, "build_stage0_input_processor", return_value=object()) + mocker.patch.object(engine_mod, "build_llm_stage_output_processor", return_value=object()) + + engine._initialize_stages(stage_init_timeout=60) + mock_start.assert_not_called() + + def test_initialize_stages_stops_master_server_and_shuts_down_initialized_clients_on_failure( + self, + mocker: MockerFixture, + ): + import vllm_omni.engine.async_omni_engine as engine_mod + + stage_cfgs = [_make_stage_cfg(0)] + engine = self._build_engine(stage_cfgs, single_stage_mode=True, stage_id_filter=0) + stage_plan = _make_llm_plan(0, configured_stage_id=0, launch_mode="local") + initialized_client = mocker.Mock() + mock_master = mocker.Mock(spec=OmniMasterServer) + + mocker.patch.object(engine_mod, "prepare_engine_environment") + mocker.patch.object(engine_mod, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(engine_mod, "compute_replica_layout", return_value=([1], {})) + mocker.patch.object(engine, "_build_logical_stage_init_plans", return_value=([stage_plan], None)) + + def _start_master(_plans): + engine._omni_master_server = mock_master + + mocker.patch.object(engine, "_start_omni_master_server", side_effect=_start_master) + mocker.patch.object(engine, "_initialize_stage_replicas", return_value={0: [initialized_client]}) + mocker.patch.object(engine_mod, "build_stage0_input_processor", return_value=object()) + mocker.patch.object(engine, "_assemble_stage_pools", side_effect=RuntimeError("assemble failed")) + mock_shutdown = mocker.patch.object(engine, "_shutdown_initialized_clients") + + with pytest.raises(RuntimeError, match="assemble failed"): + engine._initialize_stages(stage_init_timeout=60) + + mock_shutdown.assert_called_once_with([initialized_client]) + mock_master.stop.assert_called_once() -# --------------------------------------------------------------------------- -# AsyncOmniEngine – _launch_diffusion_stage -# --------------------------------------------------------------------------- +class TestSingleStageReplicaInitialization: + def test_initialize_llm_replica_remote_uses_connect_remote_engine_cores(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod - -class TestLaunchDiffusionStage: - """Test local diffusion stage launch wiring.""" - - def test_registers_stage_with_public_master_properties(self, mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) engine.model = "fake-model" + engine.single_stage_mode = True + engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) + engine._omni_master_server.get_stage_config.return_value = {"stage_id": 7, "stage_type": "llm"} + + fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)) + fake_addresses = SimpleNamespace( + inputs=["tcp://in"], outputs=["tcp://out"], frontend_stats_publish_address=None + ) + fake_manager = mocker.Mock() + fake_coordinator = mocker.Mock() + events: list[str] = [] + + @contextmanager + def _fake_connect(**kwargs): + events.append("enter") + try: + yield fake_manager, fake_coordinator, fake_addresses, None + finally: + events.append("exit") + + plan = _make_llm_plan( + 0, + configured_stage_id=7, + launch_mode="remote", + vllm_config=fake_vllm_config, + ).replicas[0] + sentinel_client = SimpleNamespace() + + mock_connect = mocker.patch.object(engine_mod, "connect_remote_engine_cores", side_effect=_fake_connect) + mocker.patch.object( + StageEngineCoreClientBase, + "make_async_mp_client", + side_effect=lambda **_: (events.append("attach"), sentinel_client)[1], + ) + + result = engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + + assert result is sentinel_client + engine._omni_master_server.get_stage_config.assert_called_once_with(7, timeout_s=60, replica_id=0) + assert fake_vllm_config.parallel_config.data_parallel_size_local == 0 + assert mock_connect.call_args.kwargs["stage_id"] == 7 + assert mock_connect.call_args.kwargs["replica_id"] == 0 + assert events == ["enter", "exit", "attach"] + + def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine.single_stage_mode = True + engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) + engine._omni_master_server.get_stage_config.return_value = None + + plan = _make_llm_plan(0, configured_stage_id=7, launch_mode="remote").replicas[0] + + with pytest.raises(ValueError, match="registered without stage config"): + engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + + def test_initialize_llm_replica_remote_attach_failure_cleans_up_started_resources(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod + + engine = object.__new__(AsyncOmniEngine) + engine.single_stage_mode = True + engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) + engine._omni_master_server.get_stage_config.return_value = {"stage_id": 7, "stage_type": "llm"} + + fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)) + fake_addresses = SimpleNamespace( + inputs=["tcp://in"], outputs=["tcp://out"], frontend_stats_publish_address=None + ) + fake_manager = mocker.Mock() + fake_coordinator = mocker.Mock() + + @contextmanager + def _fake_connect(**kwargs): + yield fake_manager, fake_coordinator, fake_addresses, None + + plan = _make_llm_plan( + 0, + configured_stage_id=7, + launch_mode="remote", + vllm_config=fake_vllm_config, + ).replicas[0] + mocker.patch.object(engine_mod, "connect_remote_engine_cores", side_effect=_fake_connect) + mocker.patch.object( + StageEngineCoreClientBase, + "make_async_mp_client", + side_effect=RuntimeError("attach failed"), + ) + + with pytest.raises(RuntimeError, match="attach failed"): + engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + + fake_manager.shutdown.assert_called_once() + fake_coordinator.shutdown.assert_called_once() + + def test_initialize_llm_replica_single_stage_local_uses_launch_omni_core_engines(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod + from vllm_omni.platforms import current_omni_platform + + engine = object.__new__(AsyncOmniEngine) + engine.model = "fake-model" + engine.single_stage_mode = True + engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) + engine.stage_configs = [] + + fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace()) + fake_addresses = SimpleNamespace( + inputs=["tcp://in"], outputs=["tcp://out"], frontend_stats_publish_address=None + ) + + @contextmanager + def _fake_launch(**kwargs): + yield mocker.Mock(), None, fake_addresses + + plan = _make_llm_plan( + 0, + configured_stage_id=3, + launch_mode="local", + vllm_config=fake_vllm_config, + ).replicas[0] + sentinel_client = SimpleNamespace() + + device_env_var = current_omni_platform.device_control_env_var + prev_device_env = os.environ.get(device_env_var) + os.environ[device_env_var] = "0" + + mocker.patch.object(engine_mod, "setup_stage_devices") + mocker.patch.object(engine_mod, "build_engine_args_dict", return_value={}) + mocker.patch.object(engine_mod, "acquire_device_locks", return_value=[]) + mocker.patch.object(engine_mod, "release_device_locks") + mock_launch = mocker.patch.object(engine_mod, "launch_omni_core_engines", side_effect=_fake_launch) + mocker.patch.object( + StageEngineCoreClientBase, + "make_async_mp_client", + side_effect=lambda **_: sentinel_client, + ) + try: + result = engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + finally: + if prev_device_env is None: + os.environ.pop(device_env_var, None) + else: + os.environ[device_env_var] = prev_device_env + + assert result is sentinel_client + assert mock_launch.call_args.kwargs["stage_id"] == 3 + assert mock_launch.call_args.kwargs["stage_config"] is plan.stage_cfg + assert mock_launch.call_args.kwargs["replica_id"] == 0 + + def test_initialize_diffusion_replica_remote_uses_from_addresses(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod + + engine = object.__new__(AsyncOmniEngine) + engine.single_stage_mode = True engine.diffusion_batch_size = 4 + engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) + engine._omni_master_server.get_stage_config.return_value = {"stage_id": 11, "stage_type": "diffusion"} + engine._omni_master_server.get_zmq_addresses.return_value = SimpleNamespace( + inputs=["tcp://in"], + outputs=["tcp://out"], + ) - stage_cfg = _make_stage_cfg(5, stage_type="diffusion") - metadata = mocker.Mock(stage_id=5) - omni_master_server = mocker.Mock(spec=OmniMasterServer) - omni_master_server.address = "127.0.0.1" - omni_master_server.port = 25000 + remote_metadata = _make_diffusion_plan(1, configured_stage_id=11, launch_mode="remote").replicas[0].metadata + plan = _make_diffusion_plan(1, configured_stage_id=11, launch_mode="remote").replicas[0] + sentinel_client = SimpleNamespace() + mocker.patch.object(engine_mod, "extract_stage_metadata", return_value=remote_metadata) + mock_from_addresses = mocker.patch.object( + engine_mod.StageDiffusionClient, "from_addresses", return_value=sentinel_client + ) + + result = engine._initialize_diffusion_replica(plan, stage_init_timeout=60, stage_launch_lock=threading.Lock()) + + assert result is sentinel_client + engine._omni_master_server.get_stage_config.assert_called_once_with(11, timeout_s=60, replica_id=0) + engine._omni_master_server.get_zmq_addresses.assert_called_once_with(11, replica_id=0) + mock_from_addresses.assert_called_once() + + def test_initialize_diffusion_replica_single_stage_local_registers_with_master(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod + from vllm_omni.platforms import current_omni_platform + + engine = object.__new__(AsyncOmniEngine) + engine.model = "fake-model" + engine.single_stage_mode = True + engine.diffusion_batch_size = 4 + engine.stage_configs = [] + engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) + engine._omni_master_server.address = "127.0.0.1" + engine._omni_master_server.port = 25000 + + plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0] + sentinel_client = SimpleNamespace() proc = mocker.Mock() - diffusion_client = mocker.Mock() - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_diffusion_config", - return_value="diffusion-config", + device_env_var = current_omni_platform.device_control_env_var + prev_device_env = os.environ.get(device_env_var) + os.environ[device_env_var] = "0" + + mocker.patch.object(engine_mod, "setup_stage_devices") + mocker.patch.object(engine_mod, "inject_kv_stage_info") + mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config") + mock_register = mocker.patch.object( + engine_mod, + "register_stage_with_omni_master", + return_value=("tcp://hs", "tcp://req", "tcp://resp"), ) - mock_register = mocker.patch( - "vllm_omni.engine.async_omni_engine.register_stage_with_omni_master", - return_value=( - "tcp://127.0.0.1:25001", - "tcp://127.0.0.1:25002", - "tcp://127.0.0.1:25003", - ), - ) - mock_spawn = mocker.patch( - "vllm_omni.engine.async_omni_engine.spawn_diffusion_proc", + mock_spawn = mocker.patch.object( + engine_mod, + "spawn_diffusion_proc", return_value=(proc, None, None, None), ) - mock_handshake = mocker.patch("vllm_omni.engine.async_omni_engine.complete_diffusion_handshake") - mock_from_addresses = mocker.patch( - "vllm_omni.engine.async_omni_engine.StageDiffusionClient.from_addresses", - return_value=diffusion_client, - ) - mock_acquire_locks = mocker.patch( - "vllm_omni.engine.async_omni_engine.acquire_diffusion_device_locks", - return_value=[101], - ) - mock_release_locks = mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") - - result = engine._launch_diffusion_stage( - stage_cfg=stage_cfg, - metadata=metadata, - omni_master_server=omni_master_server, - stage_init_timeout=60, + mock_handshake = mocker.patch.object(engine_mod, "complete_diffusion_handshake") + mock_from_addresses = mocker.patch.object( + engine_mod.StageDiffusionClient, + "from_addresses", + return_value=sentinel_client, ) - mock_acquire_locks.assert_called_once_with(5, "diffusion-config", 60) + try: + result = engine._initialize_diffusion_replica( + plan, stage_init_timeout=60, stage_launch_lock=threading.Lock() + ) + finally: + if prev_device_env is None: + os.environ.pop(device_env_var, None) + else: + os.environ[device_env_var] = prev_device_env + + assert result is sentinel_client mock_register.assert_called_once_with( omni_master_address="127.0.0.1", omni_master_port=25000, omni_stage_id=5, - omni_stage_config=stage_cfg, + omni_stage_config=plan.stage_cfg, return_addresses=True, + replica_id=0, ) mock_spawn.assert_called_once_with( "fake-model", "diffusion-config", - handshake_address="tcp://127.0.0.1:25001", - request_address="tcp://127.0.0.1:25002", - response_address="tcp://127.0.0.1:25003", + handshake_address="tcp://hs", + request_address="tcp://req", + response_address="tcp://resp", ) - mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:25001", 60) - mock_release_locks.assert_called_once_with([101]) + mock_handshake.assert_called_once_with(proc, "tcp://hs", 60) mock_from_addresses.assert_called_once_with( - metadata, - request_address="tcp://127.0.0.1:25002", - response_address="tcp://127.0.0.1:25003", + plan.metadata, + request_address="tcp://req", + response_address="tcp://resp", proc=proc, batch_size=4, ) - assert result is diffusion_client + def test_initialize_diffusion_replica_local_failure_terminates_proc(self, mocker: MockerFixture): + import vllm_omni.engine.async_omni_engine as engine_mod + from vllm_omni.platforms import current_omni_platform -# --------------------------------------------------------------------------- -# AsyncOmniEngine – _create_remote_llm_stage -# --------------------------------------------------------------------------- - - -class TestCreateRemoteLlmStage: - """Test _create_remote_llm_stage delegates correctly.""" - - def _engine(self, mocker: MockerFixture) -> AsyncOmniEngine: engine = object.__new__(AsyncOmniEngine) engine.model = "fake-model" engine.single_stage_mode = True - engine._single_stage_id_filter = 0 + engine.diffusion_batch_size = 4 + engine.stage_configs = [] engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.get_zmq_addresses.return_value = mocker.Mock() - engine._omni_master_server.get_allocation.return_value = mocker.Mock() - engine._omni_master_server.get_stage_config.return_value = { - "stage_id": 0, - "stage_type": "llm", - "engine_args": {}, - } - return engine + engine._omni_master_server.address = "127.0.0.1" + engine._omni_master_server.port = 25000 - def _mock_build_and_connect(self, mocker: MockerFixture, stage_id: int): - fake_vllm_config = mocker.Mock() - fake_executor_cls = mocker.Mock() - fake_addresses = mocker.Mock() - fake_addresses.inputs = ["tcp://127.0.0.1:5000"] - fake_addresses.outputs = ["tcp://127.0.0.1:5001"] - fake_addresses.frontend_stats_publish_address = None + plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0] + proc = mocker.Mock() - eng_mgr = mocker.Mock() - coordinator = mocker.Mock() + device_env_var = current_omni_platform.device_control_env_var + prev_device_env = os.environ.get(device_env_var) + os.environ[device_env_var] = "0" - @contextmanager - def fake_connect_cm(*args, **kwargs): - yield eng_mgr, coordinator, fake_addresses, None - - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": stage_id}, + mocker.patch.object(engine_mod, "setup_stage_devices") + mocker.patch.object(engine_mod, "inject_kv_stage_info") + mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config") + mocker.patch.object( + engine_mod, + "register_stage_with_omni_master", + return_value=("tcp://hs", "tcp://req", "tcp://resp"), ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(fake_vllm_config, fake_executor_cls), - ) - mock_connect = mocker.patch( - "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", - return_value=fake_connect_cm(), + mocker.patch.object( + engine_mod, + "spawn_diffusion_proc", + return_value=(proc, None, None, None), ) + mocker.patch.object(engine_mod, "complete_diffusion_handshake", side_effect=RuntimeError("handshake failed")) + mock_terminate = mocker.patch.object(engine_mod, "terminate_alive_proc") - return mock_connect, fake_vllm_config, fake_executor_cls, fake_addresses + try: + with pytest.raises(RuntimeError, match="handshake failed"): + engine._initialize_diffusion_replica(plan, stage_init_timeout=60, stage_launch_lock=threading.Lock()) + finally: + if prev_device_env is None: + os.environ.pop(device_env_var, None) + else: + os.environ[device_env_var] = prev_device_env - def test_returns_started_llm_stage_with_correct_stage_id(self, mocker: MockerFixture): - engine = self._engine(mocker) - stage_cfg = _make_stage_cfg(1) - metadata = mocker.Mock(stage_id=1) - omni_ms = engine._omni_master_server - omni_ms.get_stage_config.return_value = { - "stage_id": 1, - "stage_type": "llm", - "engine_args": {}, - } + mock_terminate.assert_called_once_with(proc) - self._mock_build_and_connect(mocker, 1) - result = engine._create_remote_llm_stage( - stage_cfg=stage_cfg, - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - omni_master_server=omni_ms, - ) - assert isinstance(result, StartedLlmStage) - assert result.stage_id == 1 - def test_connect_remote_engine_cores_called_with_stage_id(self, mocker: MockerFixture): - engine = self._engine(mocker) - stage_cfg = _make_stage_cfg(2) - metadata = mocker.Mock(stage_id=2) - omni_ms = engine._omni_master_server - omni_ms.get_zmq_addresses.return_value = mocker.Mock(inputs=["x"], outputs=["y"]) - omni_ms.get_stage_config.return_value = { - "stage_id": 2, - "stage_type": "llm", - "engine_args": {}, - } - - fake_vllm_config = mocker.Mock() - fake_executor_cls = mocker.Mock() - fake_addresses = mocker.Mock() - fake_addresses.inputs = ["tcp://127.0.0.1:5000"] - fake_addresses.outputs = ["tcp://127.0.0.1:5001"] - fake_addresses.frontend_stats_publish_address = None - - @contextmanager - def fake_connect_cm(*args, **kwargs): - yield mocker.Mock(), mocker.Mock(), fake_addresses, None - - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 2}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(fake_vllm_config, fake_executor_cls), - ) - mock_connect = mocker.patch( - "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", - return_value=fake_connect_cm(), - ) - - engine._create_remote_llm_stage( - stage_cfg=stage_cfg, - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - omni_master_server=omni_ms, - ) - - mock_connect.assert_called_once() - _, kwargs = mock_connect.call_args - assert kwargs.get("stage_id") == 2 or mock_connect.call_args.args[-1] == 2 - omni_ms.get_stage_config.assert_called_once_with(2, timeout_s=60) - - def test_missing_registered_stage_config_raises_value_error(self, mocker: MockerFixture): - engine = self._engine(mocker) - stage_cfg = _make_stage_cfg(3) - metadata = mocker.Mock(stage_id=3) - omni_ms = engine._omni_master_server - omni_ms.get_stage_config.return_value = None - - mock_build_args = mocker.patch("vllm_omni.engine.async_omni_engine.build_engine_args_dict") - with pytest.raises( - ValueError, - match="Remote stage 3 registered without stage config", - ): - engine._create_remote_llm_stage( - stage_cfg=stage_cfg, - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - omni_master_server=omni_ms, - ) - - mock_build_args.assert_not_called() - - def test_exception_during_connect_closes_started_stage(self, mocker: MockerFixture): - """If an error occurs after StartedLlmStage creation, close_started_llm_stage is called.""" - engine = self._engine(mocker) - stage_cfg = _make_stage_cfg(1) - metadata = mocker.Mock(stage_id=1) - omni_ms = engine._omni_master_server - omni_ms.get_stage_config.return_value = { - "stage_id": 1, - "stage_type": "llm", - "engine_args": {}, - } - - @contextmanager - def boom(*args, **kwargs): - yield mocker.Mock(), mocker.Mock(), mocker.Mock(), None - raise RuntimeError("handshake failed") - - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 1}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(mocker.Mock(), mocker.Mock()), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.connect_remote_engine_cores", - return_value=boom(), - ) - mock_close = mocker.patch("vllm_omni.engine.async_omni_engine.close_started_llm_stage") - with pytest.raises(RuntimeError, match="handshake failed"): - engine._create_remote_llm_stage( - stage_cfg=stage_cfg, - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - omni_master_server=omni_ms, - ) - mock_close.assert_called_once() +# --------------------------------------------------------------------------- +# Stage engine startup helpers +# --------------------------------------------------------------------------- class TestConnectRemoteEngineCoresCoordinator: - """Test coordinator launch parity with launch_core_engines.""" - @staticmethod def _build_vllm_config( mocker: MockerFixture, *, dp_rank: int = 0, offline_mode: bool = False, needs_dp_coordinator: bool = True @@ -1388,7 +1035,8 @@ class TestConnectRemoteEngineCoresCoordinator: omni_master_server = mocker.Mock(spec=OmniMasterServer) omni_master_server.get_zmq_addresses.return_value = EngineZmqAddresses( - inputs=["tcp://client-in"], outputs=["tcp://client-out"] + inputs=["tcp://client-in"], + outputs=["tcp://client-out"], ) omni_master_server.get_allocation.return_value = mocker.Mock(handshake_bind_address="tcp://127.0.0.1:26001") omni_master_server.get_stage_coordinator_addresses.return_value = StageCoordinatorAddresses( @@ -1401,35 +1049,31 @@ class TestConnectRemoteEngineCoresCoordinator: def fake_socket_ctx(*args, **kwargs): yield mocker.Mock() - mocker.patch( - "vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", - return_value=fake_socket_ctx(), - ) + mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") with connect_remote_engine_cores( vllm_config=vllm_config, omni_master_server=omni_master_server, stage_id=7, + replica_id=2, ) as (_, yielded_coordinator, yielded_addresses, _tensor_queue): assert yielded_coordinator is None assert yielded_addresses.coordinator_input == "tcp://coord-in" assert yielded_addresses.coordinator_output == "tcp://coord-out" assert yielded_addresses.frontend_stats_publish_address == "tcp://stats" - omni_master_server.get_stage_coordinator_addresses.assert_called_once_with(7) + omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2) + omni_master_server.get_stage_coordinator_addresses.assert_called_once_with(7, replica_id=2) + omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2) mock_wait.assert_called_once() def test_defaults_to_no_coordinator_addresses_when_none_registered(self, mocker: MockerFixture): - vllm_config = self._build_vllm_config( - mocker, - dp_rank=0, - offline_mode=False, - needs_dp_coordinator=True, - ) + vllm_config = self._build_vllm_config(mocker, dp_rank=0, offline_mode=False, needs_dp_coordinator=True) omni_master_server = mocker.Mock(spec=OmniMasterServer) omni_master_server.get_zmq_addresses.return_value = EngineZmqAddresses( - inputs=["tcp://client-in"], outputs=["tcp://client-out"] + inputs=["tcp://client-in"], + outputs=["tcp://client-out"], ) omni_master_server.get_allocation.return_value = mocker.Mock(handshake_bind_address="tcp://127.0.0.1:26001") omni_master_server.get_stage_coordinator_addresses.return_value = StageCoordinatorAddresses() @@ -1438,10 +1082,7 @@ class TestConnectRemoteEngineCoresCoordinator: def fake_socket_ctx(*args, **kwargs): yield mocker.Mock() - mocker.patch( - "vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", - return_value=fake_socket_ctx(), - ) + mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) mocker.patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") with connect_remote_engine_cores( vllm_config=vllm_config, @@ -1455,8 +1096,6 @@ class TestConnectRemoteEngineCoresCoordinator: class TestLaunchOmniCoreEngines: - """Tests for local omni engine launch wiring.""" - def test_registers_stage_once_and_reuses_handshake_for_all_local_engines(self, mocker: MockerFixture): parallel_config = mocker.Mock( data_parallel_size_local=2, @@ -1481,10 +1120,7 @@ class TestLaunchOmniCoreEngines: "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", return_value="tcp://127.0.0.1:26001", ) - mocker.patch( - "vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", - return_value=fake_socket_ctx(), - ) + mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) mock_manager_cls = mocker.patch( "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", return_value=local_engine_manager, @@ -1497,9 +1133,11 @@ class TestLaunchOmniCoreEngines: omni_master_server=omni_master_server, stage_id=7, stage_config=stage_config, + replica_id=2, ) as (yielded_manager, yielded_coordinator, yielded_addresses): assert yielded_manager is local_engine_manager assert yielded_coordinator is None + assert yielded_addresses is not None mock_register.assert_called_once_with( omni_master_address="127.0.0.1", @@ -1507,16 +1145,15 @@ class TestLaunchOmniCoreEngines: omni_stage_id=7, omni_stage_config=stage_config, coordinator=None, + replica_id=2, ) - mock_manager_cls.assert_called_once() + omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2) + omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2) manager_kwargs = mock_manager_cls.call_args.kwargs assert manager_kwargs["local_engine_count"] == 2 assert manager_kwargs["start_index"] == 3 assert manager_kwargs["local_start_index"] == 0 - assert manager_kwargs["vllm_config"] is vllm_config - assert manager_kwargs["local_client"] is True assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" - assert manager_kwargs["executor_class"] is not None def test_registers_stage_with_coordinator_when_started(self, mocker: MockerFixture): parallel_config = mocker.Mock( @@ -1524,15 +1161,19 @@ class TestLaunchOmniCoreEngines: data_parallel_size=2, data_parallel_rank=0, ) - vllm_config = mocker.Mock(parallel_config=parallel_config) - vllm_config.needs_dp_coordinator = True - vllm_config.model_config = mocker.Mock(is_moe=False) + vllm_config = mocker.Mock( + parallel_config=parallel_config, + needs_dp_coordinator=True, + model_config=mocker.Mock(is_moe=False), + cache_config=mocker.Mock(), + ) omni_master_server = mocker.Mock(spec=OmniMasterServer) omni_master_server.address = "127.0.0.1" omni_master_server.port = 26000 omni_master_server.get_zmq_addresses.return_value = EngineZmqAddresses( - inputs=["tcp://client-in"], outputs=["tcp://client-out"] + inputs=["tcp://client-in"], + outputs=["tcp://client-out"], ) omni_master_server.get_allocation.return_value = mocker.Mock(handshake_bind_address="tcp://127.0.0.1:26001") @@ -1550,11 +1191,8 @@ class TestLaunchOmniCoreEngines: "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", return_value="tcp://127.0.0.1:26001", ) + mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) mocker.patch( - "vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", - return_value=fake_socket_ctx(), - ) - mock_manager_cls = mocker.patch( "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager", return_value=mocker.Mock(), ) @@ -1566,8 +1204,12 @@ class TestLaunchOmniCoreEngines: omni_master_server=omni_master_server, stage_id=7, stage_config={"stage_id": 7}, - ): - pass + replica_id=3, + ) as (_, yielded_coordinator, yielded_addresses): + assert yielded_coordinator is coordinator + assert yielded_addresses.coordinator_input == "tcp://coord-in" + assert yielded_addresses.coordinator_output == "tcp://coord-out" + assert yielded_addresses.frontend_stats_publish_address == "tcp://stats" mock_register.assert_called_once_with( omni_master_address="127.0.0.1", @@ -1575,313 +1217,6 @@ class TestLaunchOmniCoreEngines: omni_stage_id=7, omni_stage_config={"stage_id": 7}, coordinator=coordinator, + replica_id=3, ) - manager_kwargs = mock_manager_cls.call_args.kwargs - assert manager_kwargs["log_stats"] is False mock_wait.assert_called_once() - - -# --------------------------------------------------------------------------- -# AsyncOmniEngine – _launch_llm_stage single_stage_mode codepath -# --------------------------------------------------------------------------- - - -class TestLaunchLlmStageSingleStageMode: - """Test that _launch_llm_stage selects launch_omni_core_engines when - single_stage_mode=True and _omni_master_server is set.""" - - def _build_engine_with_oms(self, mocker: MockerFixture) -> AsyncOmniEngine: - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = True - engine._single_stage_id_filter = 0 - engine._llm_stage_launch_lock = threading.Lock() - engine.stage_configs = [] - mock_oms = mocker.Mock(spec=OmniMasterServer) - mock_oms.address = "127.0.0.1" - mock_oms.port = 25000 - alloc = mocker.Mock() - alloc.handshake_bind_address = "tcp://127.0.0.1:25001" - mock_oms.get_allocation.return_value = alloc - fake_addresses = mocker.Mock() - fake_addresses.inputs = ["tcp://127.0.0.1:5000"] - fake_addresses.outputs = ["tcp://127.0.0.1:5001"] - fake_addresses.frontend_stats_publish_address = None - mock_oms.get_zmq_addresses.return_value = fake_addresses - engine._omni_master_server = mock_oms - return engine - - def _mock_launch_omni(self, mocker: MockerFixture, stage_id: int): - fake_vllm_config = mocker.Mock() - fake_executor_cls = mocker.Mock() - fake_addresses = mocker.Mock() - fake_addresses.inputs = ["tcp://127.0.0.1:5000"] - fake_addresses.outputs = ["tcp://127.0.0.1:5001"] - fake_addresses.frontend_stats_publish_address = None - - eng_mgr = mocker.Mock() - - @contextmanager - def fake_launch_omni(*args, **kwargs): - yield eng_mgr, None, fake_addresses - - mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": stage_id}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(fake_vllm_config, fake_executor_cls), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.acquire_device_locks", - return_value=[], - ) - mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") - return mocker.patch( - "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", - return_value=fake_launch_omni(), - ) - - def test_launch_omni_core_engines_used_in_single_stage_mode(self, mocker: MockerFixture): - """single_stage_mode + _omni_master_server → launch_omni_core_engines.""" - engine = self._build_engine_with_oms(mocker) - metadata = mocker.Mock(stage_id=0, runtime_cfg={}) - stage_cfg = _make_stage_cfg(0) - - mock_launch_omni = self._mock_launch_omni(mocker, 0) - result = engine._launch_llm_stage( - stage_cfg=stage_cfg, - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - llm_stage_launch_lock=threading.Lock(), - ) - - mock_launch_omni.assert_called_once() - assert mock_launch_omni.call_args.kwargs["stage_config"] is stage_cfg - assert isinstance(result, StartedLlmStage) - assert result.stage_id == 0 - - def test_spawn_stage_core_used_in_normal_mode(self, mocker: MockerFixture): - """~single_stage_mode → spawn_stage_core + complete_stage_handshake.""" - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = False - engine._omni_master_server = None - engine._llm_stage_launch_lock = threading.Lock() - engine.stage_configs = [] - - fake_vllm_config = mocker.Mock() - fake_executor_cls = mocker.Mock() - fake_addresses = mocker.Mock() - fake_addresses.inputs = ["tcp://127.0.0.1:5000"] - fake_addresses.outputs = ["tcp://127.0.0.1:5001"] - fake_addresses.frontend_stats_publish_address = None - - fake_proc = mocker.Mock() - fake_handshake_address = "ipc:///tmp/fake-handshake" - stage_init_timeout = 60 - - mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 0}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(fake_vllm_config, fake_executor_cls), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.acquire_device_locks", - return_value=[], - ) - mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") - mock_spawn = mocker.patch( - "vllm_omni.engine.async_omni_engine.spawn_stage_core", - return_value=(fake_addresses, fake_proc, fake_handshake_address), - ) - mock_handshake = mocker.patch("vllm_omni.engine.async_omni_engine.complete_stage_handshake") - mock_omni = mocker.patch("vllm_omni.engine.async_omni_engine.launch_omni_core_engines") - metadata = mocker.Mock(stage_id=0, runtime_cfg={}) - result = engine._launch_llm_stage( - stage_cfg=_make_stage_cfg(0), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=stage_init_timeout, - llm_stage_launch_lock=threading.Lock(), - ) - - mock_spawn.assert_called_once_with( - vllm_config=fake_vllm_config, - executor_class=fake_executor_cls, - log_stats=False, - ) - mock_handshake.assert_called_once_with( - fake_proc, - fake_handshake_address, - fake_addresses, - fake_vllm_config, - stage_init_timeout, - ) - mock_omni.assert_not_called() - assert isinstance(result, StartedLlmStage) - assert result.proc is fake_proc - - def test_launch_omni_passes_stage_id_and_master_server(self, mocker: MockerFixture): - """launch_omni_core_engines receives the correct stage_id and omni_master_server.""" - engine = self._build_engine_with_oms(mocker) - metadata = mocker.Mock(stage_id=0, runtime_cfg={}) - - captured_kwargs: dict[str, Any] = {} - - @contextmanager - def capturing_launch(*args, **kwargs): - captured_kwargs.update(kwargs) - fake_addresses = mocker.Mock() - fake_addresses.inputs = ["tcp://127.0.0.1:5000"] - fake_addresses.outputs = ["tcp://127.0.0.1:5001"] - fake_addresses.frontend_stats_publish_address = None - yield mocker.Mock(), None, fake_addresses - - mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 0}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(mocker.Mock(), mocker.Mock()), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.acquire_device_locks", - return_value=[], - ) - mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") - mocker.patch( - "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", - side_effect=capturing_launch, - ) - - engine._launch_llm_stage( - stage_cfg=_make_stage_cfg(0), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - llm_stage_launch_lock=threading.Lock(), - ) - - assert captured_kwargs.get("stage_id") == 0 - assert captured_kwargs.get("omni_master_server") is engine._omni_master_server - - def test_launch_omni_context_exits_before_stage_cleanup_on_error(self, mocker: MockerFixture): - """Errors after entering the omni launch context still unwind it first.""" - engine = self._build_engine_with_oms(mocker) - metadata = mocker.Mock(stage_id=0, runtime_cfg={}) - - fake_addresses = mocker.Mock() - fake_addresses.inputs = ["tcp://127.0.0.1:5000"] - fake_addresses.outputs = ["tcp://127.0.0.1:5001"] - fake_addresses.frontend_stats_publish_address = None - - events: list[str] = [] - - @contextmanager - def fake_launch_omni(*args, **kwargs): - try: - yield mocker.Mock(), None, fake_addresses - finally: - events.append("launch_exit") - - mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 0}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(mocker.Mock(), mocker.Mock()), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.acquire_device_locks", - return_value=[], - ) - mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") - mocker.patch( - "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", - return_value=fake_launch_omni(), - ) - mocker.patch("vllm_omni.engine.async_omni_engine.logger.info", side_effect=RuntimeError("boom")) - mock_close_stage = mocker.patch( - "vllm_omni.engine.async_omni_engine.close_started_llm_stage", - side_effect=lambda _started: events.append("stage_close"), - ) - with pytest.raises(RuntimeError, match="boom"): - engine._launch_llm_stage( - stage_cfg=_make_stage_cfg(0), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - llm_stage_launch_lock=threading.Lock(), - ) - - mock_close_stage.assert_called_once() - assert events == ["launch_exit", "stage_close"] - - def test_base_exception_propagates_without_started_stage_cleanup(self, mocker: MockerFixture): - """BaseException subclasses should bypass the Exception cleanup path.""" - engine = self._build_engine_with_oms(mocker) - metadata = mocker.Mock(stage_id=0, runtime_cfg={}) - - fake_addresses = mocker.Mock() - fake_addresses.inputs = ["tcp://127.0.0.1:5000"] - fake_addresses.outputs = ["tcp://127.0.0.1:5001"] - fake_addresses.frontend_stats_publish_address = None - - events: list[str] = [] - - class FatalLaunchInterrupt(BaseException): - pass - - @contextmanager - def fake_launch_omni(*args, **kwargs): - try: - yield mocker.Mock(), None, fake_addresses - finally: - events.append("launch_exit") - - mocker.patch("vllm_omni.engine.async_omni_engine.setup_stage_devices") - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_engine_args_dict", - return_value={"model": "fake", "stage_id": 0}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.build_vllm_config", - return_value=(mocker.Mock(), mocker.Mock()), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.acquire_device_locks", - return_value=[], - ) - mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks") - mocker.patch( - "vllm_omni.engine.async_omni_engine.launch_omni_core_engines", - return_value=fake_launch_omni(), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.logger.info", - side_effect=FatalLaunchInterrupt("stop"), - ) - mock_close_stage = mocker.patch("vllm_omni.engine.async_omni_engine.close_started_llm_stage") - with pytest.raises(FatalLaunchInterrupt, match="stop"): - engine._launch_llm_stage( - stage_cfg=_make_stage_cfg(0), - metadata=metadata, - stage_connector_spec={}, - stage_init_timeout=60, - llm_stage_launch_lock=threading.Lock(), - ) - - mock_close_stage.assert_not_called() - assert events == ["launch_exit"] diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 999c0a4ea..5f3415bb3 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -38,6 +38,7 @@ def _make_headless_args() -> argparse.Namespace: return argparse.Namespace( model="fake-model", stage_id=3, + replica_id=0, omni_master_address="127.0.0.1", omni_master_port=26000, api_server_count=0, @@ -45,6 +46,7 @@ def _make_headless_args() -> argparse.Namespace: stage_configs_path=None, log_stats=False, disable_log_stats=False, + stage_init_timeout=600, ) @@ -59,6 +61,7 @@ def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker node_rank_within_dp=0, ) vllm_config = mocker.Mock(parallel_config=parallel_config) + vllm_config.needs_dp_coordinator = False executor_class = mocker.Mock() engine_manager = mocker.Mock() @@ -99,6 +102,7 @@ def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker omni_stage_id=3, omni_stage_config=stage_cfg, coordinator=None, + replica_id=0, ) mock_manager_cls.assert_called_once() manager_kwargs = mock_manager_cls.call_args.kwargs @@ -155,13 +159,56 @@ def test_run_headless_honors_explicit_log_stats_flag(mocker: MockerFixture) -> N assert manager_kwargs["log_stats"] is True +def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None: + args = _make_headless_args() + args.replica_id = 2 + stage_cfg = mocker.Mock(stage_id=3) + parallel_config = mocker.Mock( + data_parallel_size_local=1, + data_parallel_rank=0, + data_parallel_rank_local=0, + node_rank_within_dp=0, + ) + vllm_config = mocker.Mock(parallel_config=parallel_config) + vllm_config.needs_dp_coordinator = False + engine_manager = mocker.Mock() + + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", [stage_cfg]), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") + mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock()) + mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}) + mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}) + mocker.patch( + "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.stage_init_utils.build_vllm_config", + return_value=(vllm_config, mocker.Mock()), + ) + mock_register = mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value="tcp://127.0.0.1:26001", + ) + mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) + mocker.patch("signal.signal") + + run_headless(args) + + assert mock_register.call_args.kwargs["replica_id"] == 2 + + def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFixture) -> None: args = _make_headless_args() + args.replica_id = 1 stage_cfg = mocker.Mock(stage_id=3, stage_type="diffusion") stage_cfg.engine_args = mocker.Mock() stage_cfg.engine_input_source = [] stage_cfgs = [stage_cfg] - metadata = mocker.Mock(stage_id=3) + metadata = mocker.Mock(stage_id=3, stage_type="diffusion") od_config = mocker.Mock() proc = mocker.Mock() proc.exitcode = 0 @@ -199,6 +246,7 @@ def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFix omni_stage_id=3, omni_stage_config=stage_cfg, return_addresses=True, + replica_id=1, ) mock_spawn.assert_called_once_with( "fake-model", @@ -207,5 +255,5 @@ def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFix request_address="tcp://127.0.0.1:26002", response_address="tcp://127.0.0.1:26003", ) - mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001") + mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001", 600) proc.join.assert_called_once_with() diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 6e52e4c6c..843b5d5d6 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -360,6 +360,47 @@ class TestLoadAndResolveStageConfigs: assert len(stage_configs) == 1 assert "dtype" in stage_configs[0]["engine_args"] + def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas( + self, tmp_path, mocker: MockerFixture + ): + deploy_path = tmp_path / "qwen3_multi.yaml" + deploy_path.write_text( + 'stages:\n - stage_id: 0\n devices: "0"\n - stage_id: 1\n devices: "1,2,3"\n num_replicas: 3\n', + encoding="utf-8", + ) + + returned_stage_configs = [ + create_config({"stage_id": 0, "runtime": {"devices": "0"}, "engine_args": {"model": "dummy"}}), + create_config( + { + "stage_id": 1, + "runtime": {"devices": "1,2,3", "num_replicas": 3}, + "engine_args": {"model": "dummy"}, + } + ), + ] + load_stage_configs = mocker.patch( + "vllm_omni.entrypoints.utils.load_stage_configs_from_model", + return_value=returned_stage_configs, + ) + + config_path, stage_configs = load_and_resolve_stage_configs( + model="dummy-model", + stage_configs_path=str(deploy_path), + kwargs={}, + ) + + load_stage_configs.assert_called_once_with( + "dummy-model", + base_engine_args={}, + deploy_config_path=str(deploy_path), + stage_overrides=None, + ) + assert config_path == str(deploy_path) + assert len(stage_configs) == 2 + assert stage_configs[1].runtime.num_replicas == 3 + assert stage_configs[1].runtime.devices == "1,2,3" + class TestLoadStageConfigsFromYaml: """Regression tests for stage-config loading and merging.""" diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 79f6af791..0c4f093c3 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -385,7 +385,8 @@ class OmniServerStageCli(OmniServer): self.stage_ids = stage_ids or self._load_stage_ids(resolved_cfg) if 0 not in self.stage_ids: raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}") - self.stage_procs: dict[int, subprocess.Popen] = {} + self.stage_replica_counts = self._load_stage_replica_counts(resolved_cfg) + self.stage_procs: dict[tuple[int, int], subprocess.Popen] = {} self.proc = None @staticmethod @@ -415,6 +416,19 @@ class OmniServerStageCli(OmniServer): runtime_devices[int(stage_id)] = str(devices) return runtime_devices + @staticmethod + def _load_stage_replica_counts(resolved_config: dict) -> dict[int, int]: + replica_counts: dict[int, int] = {} + for stage in OmniServerStageCli._stage_entries(resolved_config): + stage_id = stage.get("stage_id") + if stage_id is None: + continue + replica_counts[int(stage_id)] = max( + 1, + int(stage.get("num_replicas") or stage.get("runtime", {}).get("num_replicas", 1)), + ) + return replica_counts + @classmethod def _parse_device_list(cls, devices: str | int) -> list[str]: if isinstance(devices, int): @@ -452,13 +466,29 @@ class OmniServerStageCli(OmniServer): return ",".join(visible_device_list[idx] for idx in logical_ids) - def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str) -> None: - mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, devices) + def _devices_for_replica(self, stage_id: int, devices: str, replica_id: int) -> str: + replica_count = self.stage_replica_counts.get(stage_id, 1) + if replica_count == 1: + return devices + + device_list = self._parse_device_list(devices) + if len(device_list) % replica_count != 0: + raise ValueError( + f"Stage {stage_id} has {len(device_list)} device(s) for {replica_count} replica(s); " + "device count must be divisible by replica count" + ) + devices_per_replica = len(device_list) // replica_count + start = replica_id * devices_per_replica + return ",".join(device_list[start : start + devices_per_replica]) + + def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str, replica_id: int = 0) -> None: + replica_devices = self._devices_for_replica(stage_id, devices, replica_id) + mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, replica_devices) env_var = getattr(current_omni_platform, "device_control_env_var", None) if env_var: env[env_var] = mapped_devices - def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]: + def _build_stage_cmd(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> list[str]: cmd = [ sys.executable, "-m", @@ -474,6 +504,8 @@ class OmniServerStageCli(OmniServer): self.host, "--omni-master-port", str(self.master_port), + "--replica-id", + str(replica_id), ] if headless: @@ -484,7 +516,7 @@ class OmniServerStageCli(OmniServer): cmd += self.serve_args return cmd - def _launch_stage(self, stage_id: int, *, headless: bool) -> None: + def _launch_stage(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> None: env = os.environ.copy() env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" if self.env_dict is not None: @@ -492,19 +524,20 @@ class OmniServerStageCli(OmniServer): devices = self.stage_runtime_devices.get(stage_id) if devices: - self._set_stage_device_env(stage_id, env, devices) + self._set_stage_device_env(stage_id, env, devices, replica_id=replica_id) - cmd = self._build_stage_cmd(stage_id, headless=headless) - print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}") + cmd = self._build_stage_cmd(stage_id, headless=headless, replica_id=replica_id) + print(f"Launching OmniServerStageCli stage {stage_id} replica {replica_id}: {' '.join(cmd)}") # Capture each subprocess's stdout+stderr to a per-stage log file so # debugging "Stage N exited before API server ready" doesn't rely on # guessing; the file is surfaced in the RuntimeError message. - log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_{self.master_port}.log" + log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_replica_{replica_id}_{self.master_port}.log" self._stage_log_paths = getattr(self, "_stage_log_paths", {}) - self._stage_log_paths[stage_id] = log_path + stage_key = (stage_id, replica_id) + self._stage_log_paths[stage_key] = log_path log_fh = open(log_path, "w", buffering=1) # noqa: SIM115 - closed in __exit__ self._stage_log_files = getattr(self, "_stage_log_files", {}) - self._stage_log_files[stage_id] = log_fh + self._stage_log_files[stage_key] = log_fh proc = subprocess.Popen( cmd, env=env, @@ -512,36 +545,39 @@ class OmniServerStageCli(OmniServer): stdout=log_fh, stderr=subprocess.STDOUT, ) - self.stage_procs[stage_id] = proc - if stage_id == 0: + self.stage_procs[stage_key] = proc + if stage_id == 0 and replica_id == 0: self.proc = proc def _ensure_stage_processes_alive(self) -> None: - for stage_id, proc in self.stage_procs.items(): + for (stage_id, replica_id), proc in self.stage_procs.items(): ret = proc.poll() if ret is not None: - log_path = getattr(self, "_stage_log_paths", {}).get(stage_id) + log_path = getattr(self, "_stage_log_paths", {}).get((stage_id, replica_id)) tail = "" if log_path and log_path.exists(): try: with open(log_path, encoding="utf-8", errors="replace") as f: lines = f.readlines() - tail = "\n=== Last 60 lines of stage {} log ({}) ===\n{}".format( - stage_id, log_path, "".join(lines[-60:]) or "" + tail = "\n=== Last 60 lines of stage {} replica {} log ({}) ===\n{}".format( + stage_id, replica_id, log_path, "".join(lines[-60:]) or "" ) except Exception as exc: # pragma: no cover - diagnostic only tail = f"\n" - raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.{tail}") + raise RuntimeError( + f"Stage {stage_id} replica {replica_id} exited with code {ret} before API server became ready.{tail}" + ) def _start_server(self) -> None: ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]] - self._launch_stage(0, headless=False) + self._launch_stage(0, headless=False, replica_id=0) time.sleep(2) self._ensure_stage_processes_alive() for stage_id in ordered_stage_ids[1:]: - self._launch_stage(stage_id, headless=True) + for replica_id in range(self.stage_replica_counts.get(stage_id, 1)): + self._launch_stage(stage_id, headless=True, replica_id=replica_id) max_wait = 1200 start_time = time.time() @@ -568,15 +604,15 @@ class OmniServerStageCli(OmniServer): whatever state the stage was in when it was torn down. """ log_paths = getattr(self, "_stage_log_paths", {}) or {} - for stage_id in sorted(log_paths): - log_path = log_paths[stage_id] + for stage_id, replica_id in sorted(log_paths): + log_path = log_paths[(stage_id, replica_id)] if not log_path or not log_path.exists(): continue try: with open(log_path, encoding="utf-8", errors="replace") as f: lines = f.readlines() except Exception as exc: # pragma: no cover - diagnostic only - print(f"[OmniServerStageCli] stage {stage_id} log read failed: {exc}", flush=True) + print(f"[OmniServerStageCli] stage {stage_id} replica {replica_id} log read failed: {exc}", flush=True) continue total = len(lines) if total <= head_lines + tail_lines: @@ -587,18 +623,18 @@ class OmniServerStageCli(OmniServer): head_chunk = lines[:head_lines] tail_chunk = lines[-tail_lines:] elided = total - head_lines - tail_lines - print(f"\n=== stage {stage_id} log HEAD ({log_path}) ===", flush=True) + print(f"\n=== stage {stage_id} replica {replica_id} log HEAD ({log_path}) ===", flush=True) print("".join(head_chunk).rstrip("\n"), flush=True) if tail_chunk: print(f"\n... [{elided} lines elided] ...", flush=True) - print(f"\n=== stage {stage_id} log TAIL ({log_path}) ===", flush=True) + print(f"\n=== stage {stage_id} replica {replica_id} log TAIL ({log_path}) ===", flush=True) print("".join(tail_chunk).rstrip("\n"), flush=True) - print(f"=== end stage {stage_id} log ===\n", flush=True) + print(f"=== end stage {stage_id} replica {replica_id} log ===\n", flush=True) def __exit__(self, exc_type, exc_val, exc_tb): self._dump_stage_logs_for_debug() - for stage_id in sorted(self.stage_procs, reverse=True): - proc = self.stage_procs[stage_id] + for stage_key in sorted(self.stage_procs, reverse=True): + proc = self.stage_procs[stage_key] if proc.poll() is None: self._kill_process_tree(proc.pid) run_forced_gpu_cleanup_round() diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py index 23d80c826..4532c7ecf 100644 --- a/tests/helpers/stage_config.py +++ b/tests/helpers/stage_config.py @@ -387,6 +387,79 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = { }, }, }, + "qwen3_omni_moe_multi_replicas_4gpu": { + "base_config": "qwen3_omni_moe.yaml", + "async_chunk": True, + "stages": [ + { + "stage_id": 0, + "devices": "0", + "gpu_memory_utilization": 0.85, + "max_num_seqs": 6, + "max_model_len": 32768, + "mm_processor_cache_gb": 0, + "load_format": "dummy", + "default_sampling_params": {"max_tokens": 150, "ignore_eos": False}, + }, + { + "stage_id": 1, + "devices": "1,2,3", + "num_replicas": 3, + "gpu_memory_utilization": 0.6, + "max_num_seqs": 2, + "max_model_len": 32768, + "load_format": "dummy", + "default_sampling_params": {"max_tokens": 1000}, + }, + { + "stage_id": 2, + "devices": "1,2,3", + "num_replicas": 3, + "gpu_memory_utilization": 0.1, + "max_num_seqs": 2, + "max_num_batched_tokens": 65536, + "load_format": "dummy", + "default_sampling_params": {"max_tokens": 2000}, + }, + ], + }, + "bagel_multi_replicas_4gpu": { + "base_config": "bagel.yaml", + "async_chunk": False, + "stages": [ + { + "stage_id": 0, + "devices": "0", + "max_num_seqs": 6, + "max_num_batched_tokens": 16384, + "gpu_memory_utilization": 0.45, + "load_format": "dummy", + "default_sampling_params": { + "temperature": 0.4, + "top_p": 0.9, + "top_k": 1, + "max_tokens": 256, + "detokenize": False, + }, + }, + { + "stage_id": 1, + "devices": "1,2,3", + "num_replicas": 3, + "max_num_seqs": 1, + "enforce_eager": True, + "gpu_memory_utilization": 0.7, + "load_format": "dummy", + "default_sampling_params": { + "seed": 42, + "num_inference_steps": 2, + "guidance_scale": 0.0, + "height": 512, + "width": 512, + }, + }, + ], + }, "bagel": { "base_config": "bagel.yaml", "stages": [ diff --git a/tests/helpers/test_runtime_stage_cli.py b/tests/helpers/test_runtime_stage_cli.py new file mode 100644 index 000000000..d9b50c85f --- /dev/null +++ b/tests/helpers/test_runtime_stage_cli.py @@ -0,0 +1,56 @@ +"""Tests for OmniServerStageCli process planning helpers.""" + +from __future__ import annotations + +from pathlib import Path + +from tests.helpers.runtime import OmniServerStageCli + + +def _write_stage_config(tmp_path: Path) -> str: + path = tmp_path / "stages.yaml" + path.write_text( + """ +stages: + - stage_id: 0 + devices: "0" + - stage_id: 1 + devices: "1,2,3" + num_replicas: 3 +""".strip(), + encoding="utf-8", + ) + return str(path) + + +def test_stage_cli_builds_headless_replica_cmd(tmp_path: Path) -> None: + server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), ["--disable-log-stats"]) + + cmd = server._build_stage_cmd(1, headless=True, replica_id=2) + + assert "--headless" in cmd + assert cmd[cmd.index("--stage-id") + 1] == "1" + assert cmd[cmd.index("--replica-id") + 1] == "2" + assert cmd[cmd.index("--stage-configs-path") + 1] == server.stage_config_path + + +def test_stage_cli_splits_devices_per_replica(tmp_path: Path) -> None: + server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), []) + + assert server._devices_for_replica(1, "1,2,3", 0) == "1" + assert server._devices_for_replica(1, "1,2,3", 1) == "2" + assert server._devices_for_replica(1, "1,2,3", 2) == "3" + + +def test_stage_cli_maps_visible_devices_after_replica_split(tmp_path: Path) -> None: + server = OmniServerStageCli( + "fake-model", + _write_stage_config(tmp_path), + [], + env_dict={"CUDA_VISIBLE_DEVICES": "4,5,6,7"}, + ) + env: dict[str, str] = {} + + server._set_stage_device_env(1, env, "1,2,3", replica_id=1) + + assert env["CUDA_VISIBLE_DEVICES"] == "6" diff --git a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py index 9c2366646..a762ce415 100644 --- a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py +++ b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py @@ -44,16 +44,12 @@ def _transfer_manager( def test_text2flow_supports_batched_source_outputs(): - stage_list = [ - SimpleNamespace( - engine_outputs=[ - _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}), - _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}), - ] - ) + source_outputs = [ + _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}), + _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}), ] - outputs = text2flow(stage_list=stage_list, engine_input_source=[0], prompt=None) + outputs = text2flow(source_outputs=source_outputs, prompt=None) assert len(outputs) == 2 assert outputs[0]["prompt_token_ids"] == [1, 2, 3] diff --git a/tests/model_executor/stage_input_processors/test_glm_image.py b/tests/model_executor/stage_input_processors/test_glm_image.py index 542cdd5fb..029ba99c2 100644 --- a/tests/model_executor/stage_input_processors/test_glm_image.py +++ b/tests/model_executor/stage_input_processors/test_glm_image.py @@ -32,11 +32,6 @@ def _source_output(token_ids: list[int], mm_output: dict | None = None): ) -def _stage_with_outputs(outputs): - """Create a stage list entry with engine_outputs.""" - return SimpleNamespace(engine_outputs=outputs) - - # ============================================================================= # Tests for _has_source_image # ============================================================================= @@ -285,11 +280,11 @@ class TestAr2Diffusion: """Test basic text-to-image pipeline: AR -> Diffusion.""" # 1024x1024 t2i: small(256) + large(1024) + EOS token_ids = list(range(256)) + list(range(1024)) + [16385] - stage_list = [_stage_with_outputs([_source_output(token_ids)])] + source_outputs = [_source_output(token_ids)] prompt = {"prompt": "a cat", "mm_processor_kwargs": {"target_h": 1024, "target_w": 1024}} - result = ar2diffusion(stage_list, [0], prompt=[prompt]) + result = ar2diffusion(source_outputs, prompt=[prompt]) assert len(result) == 1 assert result[0]["prompt"] == "a cat" assert result[0]["height"] == 1024 @@ -300,7 +295,7 @@ class TestAr2Diffusion: """Test image-to-image with prior_token_image_ids from AR model.""" token_ids = list(range(1024)) + [16385] mm_output = {"ids": {"prior_image": torch.tensor([1, 2, 3])}} - stage_list = [_stage_with_outputs([_source_output(token_ids, mm_output)])] + source_outputs = [_source_output(token_ids, mm_output)] from PIL import Image @@ -311,14 +306,14 @@ class TestAr2Diffusion: "multi_modal_data": {"image": img}, } - result = ar2diffusion(stage_list, [0], prompt=[prompt]) + result = ar2diffusion(source_outputs, prompt=[prompt]) assert len(result) == 1 assert result[0]["extra"]["prior_token_image_ids"] is not None def test_i2i_detected_via_modalities(self): """Test i2i mode detected via modalities field.""" token_ids = list(range(1024)) + [16385] - stage_list = [_stage_with_outputs([_source_output(token_ids)])] + source_outputs = [_source_output(token_ids)] prompt = { "prompt": "edit this", @@ -326,35 +321,26 @@ class TestAr2Diffusion: "modalities": ["img2img"], } - result = ar2diffusion(stage_list, [0], prompt=[prompt]) + result = ar2diffusion(source_outputs, prompt=[prompt]) assert len(result) == 1 - def test_empty_engine_input_source_raises(self): - with pytest.raises(ValueError, match="engine_input_source cannot be empty"): - ar2diffusion([], [], prompt={}) - - def test_invalid_stage_id_raises(self): - with pytest.raises(IndexError, match="Invalid stage_id"): - ar2diffusion([_stage_with_outputs(None)], [5], prompt={}) - - def test_no_outputs_raises(self): - with pytest.raises(RuntimeError, match="has no outputs yet"): - ar2diffusion([SimpleNamespace(engine_outputs=None)], [0], prompt={}) + def test_empty_source_outputs_returns_empty_list(self): + assert ar2diffusion([], prompt={}) == [] def test_default_dimensions(self): """When no height/width in prompt, defaults to 1024x1024.""" token_ids = list(range(256)) + list(range(1024)) + [16385] - stage_list = [_stage_with_outputs([_source_output(token_ids)])] + source_outputs = [_source_output(token_ids)] prompt = {"prompt": "test"} - result = ar2diffusion(stage_list, [0], prompt=[prompt]) + result = ar2diffusion(source_outputs, prompt=[prompt]) assert result[0]["height"] == 1024 assert result[0]["width"] == 1024 def test_requires_multimodal_data_with_pil_image(self): """Test that pil_image is included when requires_multimodal_data=True.""" token_ids = list(range(256)) + list(range(1024)) + [16385] - stage_list = [_stage_with_outputs([_source_output(token_ids)])] + source_outputs = [_source_output(token_ids)] from PIL import Image @@ -364,13 +350,13 @@ class TestAr2Diffusion: "multi_modal_data": {"image": img}, } - result = ar2diffusion(stage_list, [0], prompt=[prompt], requires_multimodal_data=True) + result = ar2diffusion(source_outputs, prompt=[prompt], requires_multimodal_data=True) assert result[0]["pil_image"] is img def test_extra_params_passed_through(self): """Test that seed, num_inference_steps, guidance_scale, negative_prompt are passed.""" token_ids = list(range(256)) + list(range(1024)) + [16385] - stage_list = [_stage_with_outputs([_source_output(token_ids)])] + source_outputs = [_source_output(token_ids)] prompt = { "prompt": "test", @@ -380,7 +366,7 @@ class TestAr2Diffusion: "negative_prompt": "blurry", } - result = ar2diffusion(stage_list, [0], prompt=[prompt]) + result = ar2diffusion(source_outputs, prompt=[prompt]) assert result[0]["seed"] == 42 assert result[0]["num_inference_steps"] == 50 assert result[0]["guidance_scale"] == 7.5 @@ -390,14 +376,14 @@ class TestAr2Diffusion: """Test processing multiple requests in a batch.""" tokens1 = list(range(256)) + list(range(1024)) + [16385] tokens2 = list(range(256)) + list(range(1024)) + [16385] - stage_list = [_stage_with_outputs([_source_output(tokens1), _source_output(tokens2)])] + source_outputs = [_source_output(tokens1), _source_output(tokens2)] prompts = [ {"prompt": "first", "mm_processor_kwargs": {"target_h": 1024, "target_w": 1024}}, {"prompt": "second", "mm_processor_kwargs": {"target_h": 512, "target_w": 512}}, ] - result = ar2diffusion(stage_list, [0], prompt=prompts) + result = ar2diffusion(source_outputs, prompt=prompts) assert len(result) == 2 assert result[0]["prompt"] == "first" assert result[1]["prompt"] == "second" diff --git a/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py b/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py index 1f0ec02f7..d1cc0dbf7 100644 --- a/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py +++ b/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py @@ -17,12 +17,11 @@ from vllm_omni.model_executor.stage_input_processors.mimo_audio import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _make_stage_list(codec_codes: torch.Tensor, request_id: str = "req-0"): - """Build a minimal stage_list[0] with a single talker output carrying codec_codes.""" +def _make_source_outputs(codec_codes: torch.Tensor, request_id: str = "req-0"): output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}}) + """Build minimal source_outputs with one talker output carrying codec codes.""" talker_output = SimpleNamespace(outputs=[output], request_id=request_id) - stage0 = SimpleNamespace(engine_outputs=[talker_output]) - return [stage0] + return [talker_output] def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog): @@ -33,7 +32,7 @@ def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog): frames = (MAX_CODE2WAV_TOKENS // 36) + 100 codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long) - stage_list = _make_stage_list(codec_codes, request_id="req-long") + source_outputs = _make_source_outputs(codec_codes, request_id="req-long") # Attach caplog's handler directly to the module logger so the warning is # captured regardless of propagation (vllm's logger configuration can @@ -43,7 +42,7 @@ def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog): prev_level = target_logger.level target_logger.setLevel(logging.WARNING) try: - prompts = llm2code2wav(stage_list, engine_input_source=[0]) + prompts = llm2code2wav(source_outputs) finally: target_logger.removeHandler(caplog.handler) target_logger.setLevel(prev_level) @@ -58,9 +57,9 @@ def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog): def test_llm2code2wav_short_sequence_unchanged(): """Short codec sequences are returned without truncation.""" codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long) - stage_list = _make_stage_list(codec_codes, request_id="req-short") + source_outputs = _make_source_outputs(codec_codes, request_id="req-short") - prompts = llm2code2wav(stage_list, engine_input_source=[0]) + prompts = llm2code2wav(source_outputs) assert len(prompts) == 1 # 4 frames + 1 pad row, flattened col-major → well below the cap diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 07e343bf0..f5cb77972 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -342,7 +342,7 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context(): engine_outputs=[SimpleNamespace(outputs=[output], finished=True)], ) - prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0]) + prompts = talker2code2wav(stage.engine_outputs) assert len(prompts) == 1 prompt = prompts[0] @@ -388,7 +388,7 @@ def test_non_async_processor_filters_out_of_range_codec_values(): engine_outputs=[SimpleNamespace(outputs=[output], finished=True)], ) - prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0]) + prompts = talker2code2wav(stage.engine_outputs) assert len(prompts) == 1 prompt = prompts[0] diff --git a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py index 1b78b103d..a61c11b5e 100644 --- a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py @@ -38,7 +38,7 @@ class TestGenerator2Tokenizer: audio = torch.tensor([10, 20, 30, 40]) stage = _make_stage([_make_output([audio])]) - result = generator2tokenizer([stage], engine_input_source=[0]) + result = generator2tokenizer(stage.engine_outputs) assert len(result) == 1 assert result[0]["prompt_token_ids"] == [10, 20, 30, 40] @@ -49,27 +49,15 @@ class TestGenerator2Tokenizer: audio = torch.tensor([[1, 2, 3], [4, 5, 6]]) # (2, 3) stage = _make_stage([_make_output([audio])]) - result = generator2tokenizer([stage], engine_input_source=[0]) + result = generator2tokenizer(stage.engine_outputs) assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6] - def test_empty_engine_input_source_raises(self): - """Empty engine_input_source raises ValueError.""" - stage = _make_stage([_make_output([torch.tensor([1])])]) - with pytest.raises(ValueError, match="cannot be empty"): - generator2tokenizer([stage], engine_input_source=[]) - - def test_invalid_stage_id_raises(self): - """Stage id beyond stage_list length raises IndexError.""" - stage = _make_stage([_make_output([torch.tensor([1])])]) - with pytest.raises(IndexError, match="Invalid stage_id"): - generator2tokenizer([stage], engine_input_source=[5]) - def test_no_outputs_yet_raises(self): """Stage with engine_outputs=None raises RuntimeError.""" stage = _make_stage(engine_outputs=None) - with pytest.raises(RuntimeError, match="no outputs yet"): - generator2tokenizer([stage], engine_input_source=[0]) + with pytest.raises(TypeError): + generator2tokenizer(stage.engine_outputs) # ---- Helpers for generator2tokenizer_async_chunk (streaming) ---- diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 6799fb80a..7b620ea6e 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -920,6 +920,27 @@ class TestDeployConfigLoading: assert s0.yaml_engine_args["engine_output_type"] == "latent" assert s0.yaml_extras["default_sampling_params"]["detokenize"] is True + def test_merge_pipeline_deploy_preserves_num_replicas(self, tmp_path): + from pathlib import Path + + import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401 + from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy + + pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] + base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + if not base.exists(): + pytest.skip("Deploy config not found") + + overlay = tmp_path / "multi_replicas.yaml" + overlay.write_text(f'base_config: {base}\nstages:\n - stage_id: 1\n devices: "1,2"\n num_replicas: 2\n') + + deploy = load_deploy_config(overlay) + assert deploy.stages[1].num_replicas == 2 + + stages = merge_pipeline_deploy(pipeline, deploy) + assert stages[1].yaml_runtime["devices"] == "1,2" + assert stages[1].yaml_runtime["num_replicas"] == 2 + class TestQwen3OmniPipeline: def test_registered(self): diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 0e4539cff..3187031bf 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -354,15 +354,15 @@ class TestFinishedLoadReqsDrain(unittest.TestCase): class TestLoadCustomFuncSelection(unittest.TestCase): - def test_skips_legacy_stage_list_processors_for_full_payload_mode(self): - legacy_paths = [ + def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self): + incompatible_paths = [ "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav", "vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit", "vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow", "vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion", ] - for func_path in legacy_paths: + for func_path in incompatible_paths: selected_path, func = MixinHost._load_custom_func( SimpleNamespace( async_chunk=False, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 358725766..8d59338db 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -407,6 +407,7 @@ class StageDeployConfig: # Stage identity and Omni runtime placement. stage_id: int devices: str | None = None + num_replicas: int = 1 # Inter-stage connector wiring and request defaults. output_connectors: dict[str, str] | None = None @@ -476,6 +477,7 @@ _STAGE_NON_ENGINE_KEYS = frozenset( { "stage_id", "devices", + "num_replicas", "output_connectors", "input_connectors", "default_sampling_params", @@ -490,13 +492,20 @@ _STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name n def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: """Parse a single stage entry from deploy YAML into StageDeployConfig.""" if "engine_args" in stage_data: + runtime_cfg = dict(stage_data.get("runtime", {})) engine_args = dict(stage_data["engine_args"]) devices = stage_data.get("runtime", {}).get("devices", stage_data.get("devices")) + num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1)) else: engine_args = {k: v for k, v in stage_data.items() if k not in _STAGE_NON_ENGINE_KEYS and k != "stage_id"} devices = stage_data.get("devices") + num_replicas = stage_data.get("num_replicas", stage_data.get("runtime", {}).get("num_replicas", 1)) - kwargs: dict[str, Any] = {"stage_id": stage_data["stage_id"], "devices": devices} + kwargs: dict[str, Any] = { + "stage_id": stage_data["stage_id"], + "devices": devices, + "num_replicas": int(num_replicas), + } for name, f in _STAGE_DEPLOY_FIELDS.items(): if name in engine_args: kwargs[name] = engine_args.pop(name) @@ -632,7 +641,11 @@ def _extract_platform_overrides(ps: dict[str, Any]) -> tuple[dict[str, Any], str the flat layout. ``devices`` is ``None`` when no override is set. """ if "engine_args" in ps: - return dict(ps["engine_args"]), ps.get("runtime", {}).get("devices") + overrides = dict(ps["engine_args"]) + runtime_cfg = ps.get("runtime", {}) + if "num_replicas" in runtime_cfg: + overrides["num_replicas"] = runtime_cfg["num_replicas"] + return overrides, runtime_cfg.get("devices") overrides = {k: v for k, v in ps.items() if k not in ("stage_id", "devices")} return overrides, ps.get("devices") @@ -837,8 +850,10 @@ def merge_pipeline_deploy( engine_args["async_scheduling"] = sched_cls is OmniARAsyncScheduler extras = _build_extras(ps, ds) runtime: dict[str, Any] = {"process": True} - if ds is not None and ds.devices is not None: - runtime["devices"] = ds.devices + if ds is not None: + if ds.devices is not None: + runtime["devices"] = ds.devices + runtime["num_replicas"] = ds.num_replicas result.append( StageConfig( diff --git a/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml b/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml new file mode 100644 index 000000000..791e29a8d --- /dev/null +++ b/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml @@ -0,0 +1,101 @@ +# Qwen3-Omni-MoE multi-replica deploy, verified on 3x H20-96G +# (stage 0 on cuda:0, stages 1+2 scaled out across cuda:1,2). +# +# Fields omitted from a stage fall back to StageDeployConfig defaults (see +# vllm_omni/config/stage_config.py). Notable implicit defaults for this +# layout: +# * Stages 0/1 (thinker, talker) do not set max_num_batched_tokens: +# chat-sized prefill fits in the 32768 default. +# * Stage 1 (talker) does not set enforce_eager: cudagraph runs by +# default (false). +# * Stage 2 (code2wav) sets enforce_eager=true because its generation +# loop is cudagraph-incompatible. +# * Stages 1/2 use num_replicas=2 and share devices "1,2". Replica +# routing is handled by StagePool with request-level affinity. +async_chunk: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_chunk_frames: 25 + codec_left_context_frames: 25 + +stages: + - stage_id: 0 + gpu_memory_utilization: 0.9 + devices: "0" + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 2048 + seed: 42 + repetition_penalty: 1.05 + + - stage_id: 1 + gpu_memory_utilization: 0.6 + devices: "1,2" + num_replicas: 2 + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + repetition_penalty: 1.05 + + - stage_id: 2 + gpu_memory_utilization: 0.1 + max_num_seqs: 1 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 51200 + devices: "1,2" + num_replicas: 2 + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.1 + +platforms: + npu: + stages: + - stage_id: 0 + gpu_memory_utilization: 0.6 + tensor_parallel_size: 2 + devices: "0,1" + - stage_id: 1 + gpu_memory_utilization: 0.6 + enforce_eager: true + devices: "2" + - stage_id: 2 + gpu_memory_utilization: 0.3 + devices: "2" + + rocm: + stages: + - stage_id: 0 + enforce_eager: true + + xpu: + stages: + - stage_id: 0 + tensor_parallel_size: 4 + enforce_eager: true + max_cudagraph_capture_size: 0 + devices: "0,1,2,3" + - stage_id: 1 + enforce_eager: true + max_cudagraph_capture_size: 0 + devices: "4" + - stage_id: 2 + gpu_memory_utilization: 0.3 + max_cudagraph_capture_size: 0 + devices: "4" diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 1817e7fdc..75dcb247a 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -35,6 +35,7 @@ class InlineStageDiffusionClient: """Runs DiffusionEngine in a thread executor inside the Orchestrator.""" stage_type: str = "diffusion" + replica_id: int = 0 def __init__( self, @@ -46,9 +47,11 @@ class InlineStageDiffusionClient: self.model = model self.od_config = od_config self.stage_id = metadata.stage_id + self.replica_id = metadata.replica_id self.final_output = metadata.final_output self.final_output_type = metadata.final_output_type self.default_sampling_params = metadata.default_sampling_params + self.requires_multimodal_data = metadata.requires_multimodal_data self.custom_process_input_func = metadata.custom_process_input_func self.engine_input_source = metadata.engine_input_source self.batch_size = batch_size @@ -62,8 +65,9 @@ class InlineStageDiffusionClient: self._shutting_down = False logger.info( - "[InlineStageDiffusionClient] Stage-%s initialized inline (batch_size=%d)", + "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (batch_size=%d)", self.stage_id, + self.replica_id, self.batch_size, ) @@ -82,6 +86,12 @@ class InlineStageDiffusionClient: sampling_params: OmniDiffusionSamplingParams, kv_sender_info: dict[int, dict[str, Any]] | None = None, ) -> None: + logger.info( + "[InlineStageDiffusionClient] stage-%s [rep-%s] add request: %s", + self.stage_id, + self.replica_id, + request_id, + ) task = asyncio.create_task( self._dispatch_request( request_id, @@ -134,6 +144,13 @@ class InlineStageDiffusionClient: sampling_params: OmniDiffusionSamplingParams, kv_sender_info: dict[int, dict[str, Any]] | None = None, ) -> None: + logger.info( + "[InlineStageDiffusionClient] stage-%s [rep-%s] add batch request: %s (%d prompts)", + self.stage_id, + self.replica_id, + request_id, + len(prompts), + ) task = asyncio.create_task( self._dispatch_batch( request_id, @@ -252,7 +269,7 @@ class InlineStageDiffusionClient: is_start = args[0] if args else True profile_prefix = args[1] if len(args) > 1 else None if is_start and profile_prefix is None: - profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}" + profile_prefix = f"stage_{self.stage_id}_rep_{self.replica_id}_diffusion_{int(time.time())}" return await loop.run_in_executor( self._executor, self._engine.profile, diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index ff2d9f1e8..e19db4536 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -67,6 +67,7 @@ class StageDiffusionClient: """ stage_type: str = "diffusion" + replica_id: int = 0 def __init__( self, @@ -112,11 +113,13 @@ class StageDiffusionClient: batch_size: int, ) -> None: self.stage_id = metadata.stage_id + self.replica_id = metadata.replica_id self.final_output = metadata.final_output self.final_output_type = metadata.final_output_type self.default_sampling_params = metadata.default_sampling_params - self.custom_process_input_func = metadata.custom_process_input_func - self.engine_input_source = metadata.engine_input_source + self.requires_multimodal_data = getattr(metadata, "requires_multimodal_data", False) + self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None) + self.engine_input_source = getattr(metadata, "engine_input_source", []) self._proc = proc self._owns_process = proc is not None @@ -142,8 +145,9 @@ class StageDiffusionClient: self._start_proc_monitor() logger.info( - "[StageDiffusionClient] Stage-%s initialized (owns_process=%s, batch_size=%d)", + "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s, batch_size=%d)", self.stage_id, + self.replica_id, self._owns_process, batch_size, ) @@ -172,8 +176,9 @@ class StageDiffusionClient: return client._engine_dead = True logger.error( - "[StageDiffusionClient] Stage-%s StageDiffusionProc died unexpectedly (exit code %s).", + "[StageDiffusionClient] stage-%s [rep-%s] StageDiffusionProc died unexpectedly (exit code %s).", client.stage_id, + client.replica_id, proc.exitcode, ) @@ -195,8 +200,9 @@ class StageDiffusionClient: if raw == StageDiffusionProc.DIFFUSION_PROC_DEAD: self._engine_dead = True logger.error( - "[StageDiffusionClient] Stage-%s received DIFFUSION_PROC_DEAD sentinel from subprocess.", + "[StageDiffusionClient] stage-%s [rep-%s] received DIFFUSION_PROC_DEAD sentinel from subprocess.", self.stage_id, + self.replica_id, ) break @@ -212,8 +218,9 @@ class StageDiffusionClient: rpc_id = msg.get("rpc_id") error_msg = msg.get("error") logger.error( - "[StageDiffusionClient] Stage-%s subprocess error for %s: %s", + "[StageDiffusionClient] stage-%s [rep-%s] subprocess error for %s: %s", self.stage_id, + self.replica_id, rpc_id or req_id, error_msg, ) @@ -292,6 +299,12 @@ class StageDiffusionClient: ) -> None: if self._engine_dead: raise EngineDeadError() + logger.info( + "[StageDiffusionClient] stage-%s [rep-%s] add request: %s", + self.stage_id, + self.replica_id, + request_id, + ) self._request_socket.send( self._encoder.encode( { @@ -321,6 +334,13 @@ class StageDiffusionClient: """ if self._engine_dead: raise EngineDeadError() + logger.info( + "[StageDiffusionClient] stage-%s [rep-%s] add batch request: %s (%d prompts)", + self.stage_id, + self.replica_id, + request_id, + len(prompts), + ) task = asyncio.create_task( self._run_batch( request_id, @@ -353,8 +373,9 @@ class StageDiffusionClient: ) except Exception as e: logger.exception( - "[StageDiffusionClient] Stage-%s batch req=%s failed: %s", + "[StageDiffusionClient] stage-%s [rep-%s] batch req=%s failed: %s", self.stage_id, + self.replica_id, request_id, e, ) @@ -414,7 +435,7 @@ class StageDiffusionClient: is_start = args_list[0] if args_list else True profile_prefix = args_list[1] if len(args_list) > 1 else None if is_start and profile_prefix is None: - profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}" + profile_prefix = f"stage_{self.stage_id}_rep_{self.replica_id}_diffusion_{int(time.time())}" if len(args_list) > 1: args_list[1] = profile_prefix else: diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 6c92d7952..55d119ad6 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -76,6 +76,42 @@ class OmniEngineCoreRequest(EngineCoreRequest): # Optional additional information dictionary (serialized) additional_information: AdditionalInformationPayload | None = None + @classmethod + def from_request( + cls, + request: EngineCoreRequest, + *, + prompt_embeds: torch.Tensor | None = None, + additional_information: AdditionalInformationPayload | None = None, + ) -> "OmniEngineCoreRequest": + """Clone an EngineCoreRequest into an OmniEngineCoreRequest with optional payload overrides.""" + + if prompt_embeds is None: + prompt_embeds = request.prompt_embeds + if additional_information is None: + additional_information = getattr(request, "additional_information", None) + + return cls( + request_id=request.request_id, + prompt_token_ids=request.prompt_token_ids, + mm_features=request.mm_features, + sampling_params=request.sampling_params, + pooling_params=request.pooling_params, + arrival_time=request.arrival_time, + lora_request=request.lora_request, + cache_salt=request.cache_salt, + data_parallel_rank=request.data_parallel_rank, + prompt_embeds=prompt_embeds, + client_index=request.client_index, + current_wave=request.current_wave, + priority=request.priority, + trace_headers=request.trace_headers, + resumable=request.resumable, + external_req_id=request.external_req_id, + reasoning_ended=request.reasoning_ended, + additional_information=additional_information, + ) + class OmniEngineCoreOutput(EngineCoreOutput): pooling_output: dict[str, torch.Tensor] | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index a5a17c719..bd98f9beb 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -9,11 +9,11 @@ from __future__ import annotations import asyncio import concurrent.futures +import copy import dataclasses import json import os import queue -import sys import threading import time import uuid @@ -30,7 +30,6 @@ from vllm import envs as vllm_envs from vllm.engine.arg_utils import EngineArgs from vllm.inputs import PromptType from vllm.logger import init_logger -from vllm.tokenizers import cached_tokenizer_from_config from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.input_processor import InputProcessor @@ -47,7 +46,6 @@ from vllm_omni.distributed.omni_connectors.utils.initialization import ( ) from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.orchestrator import Orchestrator -from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.engine.serialization import ( deserialize_additional_information, serialize_additional_information, @@ -64,17 +62,18 @@ from vllm_omni.engine.stage_engine_startup import ( register_stage_with_omni_master, ) from vllm_omni.engine.stage_init_utils import ( - StartedLlmStage, + LogicalStageInitPlan, + ReplicaInitPlan, _inject_inferred_kv_tp_topology, acquire_device_locks, acquire_diffusion_device_locks, build_diffusion_config, build_engine_args_dict, + build_llm_stage_output_processor, + build_stage0_input_processor, build_vllm_config, - cleanup_failed_stage_initialization, - close_started_llm_stage, + compute_replica_layout, extract_stage_metadata, - finalize_initialized_stages, get_stage_connector_spec, initialize_diffusion_stage, inject_kv_stage_info, @@ -84,12 +83,12 @@ from vllm_omni.engine.stage_init_utils import ( setup_stage_devices, terminate_alive_proc, ) +from vllm_omni.engine.stage_pool import StagePool from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import ( inject_omni_kv_config, load_and_resolve_stage_configs, ) -from vllm_omni.inputs.preprocess import OmniInputPreprocessor from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: @@ -134,16 +133,6 @@ _PARENT_ARGS_STRIP: frozenset[str] = frozenset({"stage_configs_path"}) _PARENT_ARGS_NO_WARN: frozenset[str] = frozenset({"model"}) -def _patch_generation_config_if_needed(model_config: Any) -> None: - """Ensure try_get_generation_config won't crash for models whose HF - config.json lacks model_type (e.g. CosyVoice3). We probe it once; - if it raises, we monkey-patch the method to return None.""" - try: - model_config.try_get_generation_config() - except Exception: - model_config.try_get_generation_config = lambda: {} - - def _inject_global_id(target: Any, request_id: str) -> None: """Inject global_request_id into a prompt dict's additional_information.""" if isinstance(target, dict): @@ -176,24 +165,9 @@ def _upgrade_to_omni_request( if prompt_embeds is None and additional_information is None: return request - return OmniEngineCoreRequest( - request_id=request.request_id, - prompt_token_ids=request.prompt_token_ids, - mm_features=request.mm_features, - sampling_params=request.sampling_params, - pooling_params=request.pooling_params, - arrival_time=request.arrival_time, - lora_request=request.lora_request, - cache_salt=request.cache_salt, - data_parallel_rank=request.data_parallel_rank, + return OmniEngineCoreRequest.from_request( + request, prompt_embeds=prompt_embeds, - client_index=request.client_index, - current_wave=request.current_wave, - priority=request.priority, - trace_headers=request.trace_headers, - resumable=request.resumable, - external_req_id=request.external_req_id, - reasoning_ended=request.reasoning_ended, additional_information=additional_information, ) @@ -208,24 +182,8 @@ def _apply_omni_final_stage_metadata( merged = deserialize_additional_information(request.additional_information) merged["omni_final_stage_id"] = final_stage_id payload = serialize_additional_information(merged) - return OmniEngineCoreRequest( - request_id=request.request_id, - prompt_token_ids=request.prompt_token_ids, - mm_features=request.mm_features, - sampling_params=request.sampling_params, - pooling_params=request.pooling_params, - arrival_time=request.arrival_time, - lora_request=request.lora_request, - cache_salt=request.cache_salt, - data_parallel_rank=request.data_parallel_rank, - prompt_embeds=request.prompt_embeds, - client_index=request.client_index, - current_wave=request.current_wave, - priority=request.priority, - trace_headers=request.trace_headers, - resumable=request.resumable, - external_req_id=request.external_req_id, - reasoning_ended=request.reasoning_ended, + return OmniEngineCoreRequest.from_request( + request, additional_information=payload, ) @@ -329,13 +287,13 @@ class AsyncOmniEngine: ) self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs) + self._validate_single_stage_mode_replica_constraints() self.num_stages = len(self.stage_configs) stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None self.async_chunk = bool(getattr(stage0_args, "async_chunk", False)) - self.stage_clients: list[Any] = [] - self.stage_vllm_configs: list[Any] = [] - self.output_processors: list[MultimodalOutputProcessor | None] = [] + self.stage_pools: list[StagePool] = [] + self.stage_clients: list[Any] = [] # logical-stage view for external readers self.input_processor: InputProcessor | None = None self.supported_tasks: tuple[str, ...] = ("generate",) self.default_sampling_params_list: list[Any] = [] @@ -362,7 +320,6 @@ class AsyncOmniEngine: name="orchestrator", ) self.orchestrator_thread.start() - self._wait_for_orchestrator_init(startup_future, startup_timeout) # Stage runtime fields are assigned directly on self by the bootstrap thread. @@ -377,225 +334,520 @@ class AsyncOmniEngine: logger.info(f"[AsyncOmniEngine] Orchestrator ready with {self.num_stages} stages") - def _launch_llm_stage( + @staticmethod + def _cleanup_launched_llm_resources( + *, + stage_id: int, + proc: Any = None, + engine_manager: Any = None, + coordinator: Any = None, + ) -> None: + """Release launch-only LLM resources when client creation never completed.""" + + if proc is not None: + try: + terminate_alive_proc(proc) + except Exception as cleanup_error: + logger.warning( + "[AsyncOmniEngine] Failed to terminate process for stage %s: %s", + stage_id, + cleanup_error, + ) + + for resource, resource_name in ( + (engine_manager, "engine manager"), + (coordinator, "coordinator"), + ): + if resource is None: + continue + shutdown = getattr(resource, "shutdown", None) + close = getattr(resource, "close", None) + try: + if callable(shutdown): + shutdown() + elif callable(close): + close() + except Exception as cleanup_error: + logger.warning( + "[AsyncOmniEngine] Failed to cleanup launched %s for stage %s: %s", + resource_name, + stage_id, + cleanup_error, + ) + + @staticmethod + def _collect_initialized_clients_for_cleanup( + stage_pools: Sequence[Any], + initialized_clients_by_stage: Mapping[int, Sequence[Any | None]], + ) -> list[Any]: + """Collect initialized clients exactly once for failure cleanup.""" + + collected: list[Any] = [] + seen: set[int] = set() + + def _add_client(client: Any) -> None: + if client is None: + return + client_id = id(client) + if client_id in seen: + return + seen.add(client_id) + collected.append(client) + + for pool in stage_pools: + for client in getattr(pool, "clients", ()): + _add_client(client) + + for clients in initialized_clients_by_stage.values(): + for client in clients: + _add_client(client) + + return collected + + @staticmethod + def _shutdown_initialized_clients(clients: Sequence[Any]) -> None: + """Best-effort shutdown for attached clients after init failure.""" + + for client in reversed(list(clients)): + if client is None: + continue + try: + client.shutdown() + except Exception as cleanup_error: + logger.warning( + "[AsyncOmniEngine] Failed to shutdown initialized client after init failure: %s", + cleanup_error, + ) + + def _validate_single_stage_mode_replica_constraints(self) -> None: + """Reject unsupported replica fan-out in single-stage mode.""" + if not self.single_stage_mode: + return + + unsupported: list[tuple[int, int]] = [] + for idx, stage_cfg in enumerate(self.stage_configs): + runtime_cfg = getattr(stage_cfg, "runtime", {}) + num_replicas = int( + runtime_cfg.get("num_replicas", 1) + if hasattr(runtime_cfg, "get") + else getattr(runtime_cfg, "num_replicas", 1) + ) + if num_replicas <= 1: + continue + if getattr(stage_cfg, "stage_type", "llm") == "diffusion": + continue + stage_id = int(getattr(stage_cfg, "stage_id", idx)) + unsupported.append((stage_id, num_replicas)) + + if unsupported: + raise ValueError( + "single_stage_mode only supports num_replicas > 1 for diffusion stages; " + f"found non-diffusion stages {unsupported}" + ) + + def _build_logical_stage_init_plans( self, - stage_cfg: Any, - metadata: Any, - stage_connector_spec: dict[str, Any], + omni_transfer_config: Any, + replicas_per_stage: Sequence[int], + replica_devices_map: Mapping[int, Sequence[str]], + ) -> tuple[list[LogicalStageInitPlan], Any]: + """Build startup plans for every logical stage and replica.""" + + prompt_expand_func = None + stage_plans: list[LogicalStageInitPlan] = [] + + for stage_idx, stage_cfg in enumerate(self.stage_configs): + base_metadata = extract_stage_metadata(stage_cfg) + configured_stage_id = base_metadata.stage_id + if base_metadata.prompt_expand_func is not None: + prompt_expand_func = base_metadata.prompt_expand_func + + stage_connector_spec = get_stage_connector_spec( + omni_transfer_config=omni_transfer_config, + stage_id=configured_stage_id, + async_chunk=self.async_chunk, + ) + omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, configured_stage_id) + num_replicas = replicas_per_stage[stage_idx] + launch_mode = "local" + if ( + self.single_stage_mode + and self._single_stage_id_filter is not None + and configured_stage_id != self._single_stage_id_filter + ): + launch_mode = "remote" + + replicas: list[ReplicaInitPlan] = [] + stage_vllm_config = None + executor_class = None + if base_metadata.stage_type != "diffusion": + engine_args_dict = build_engine_args_dict( + stage_cfg, + self.model, + stage_connector_spec=stage_connector_spec, + ) + omni_conn_cfg, omni_from, omni_to = omni_kv_connector + if omni_conn_cfg: + omni_kv = engine_args_dict.get("omni_kv_config") or {} + if not isinstance(omni_kv, dict): + omni_kv = dict(omni_kv) + omni_kv["connector_config"] = omni_conn_cfg + omni_kv["omni_from_stage"] = omni_from + omni_kv["omni_to_stage"] = omni_to + omni_kv.setdefault("stage_id", configured_stage_id) + engine_args_dict["omni_kv_config"] = omni_kv + if self.stage_configs: + _inject_inferred_kv_tp_topology( + engine_args_dict.get("omni_kv_config"), + configured_stage_id, + self.stage_configs, + ) + stage_vllm_config, executor_class = build_vllm_config( + stage_cfg, + self.model, + stage_connector_spec=stage_connector_spec, + engine_args_dict=engine_args_dict, + ) + + for replica_id in range(num_replicas): + replica_cfg = copy.deepcopy(stage_cfg) if replica_id > 0 else stage_cfg + if stage_idx in replica_devices_map: + replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id] + + replica_metadata = extract_stage_metadata(replica_cfg) + replica_metadata.replica_id = replica_id + if self.single_stage_mode: + if replica_metadata.stage_type != "diffusion": + replica_metadata.runtime_cfg = None + + replicas.append( + ReplicaInitPlan( + replica_id=replica_id, + num_replicas=num_replicas, + launch_mode=launch_mode, + stage_cfg=replica_cfg, + metadata=replica_metadata, + stage_connector_spec=stage_connector_spec, + omni_kv_connector=omni_kv_connector, + stage_vllm_config=stage_vllm_config, + executor_class=executor_class, + ) + ) + + stage_plans.append( + LogicalStageInitPlan( + stage_idx=stage_idx, + configured_stage_id=configured_stage_id, + replicas=replicas, + ) + ) + + return stage_plans, prompt_expand_func + + def _start_omni_master_server(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + """Start OmniMasterServer for single-stage mode.""" + + if not self._omni_master_address or not self._omni_master_port: + raise ValueError( + "AsyncOmniEngine single_stage_mode requires both omni_master_address and omni_master_port to be set." + ) + + all_stage_ids: list[int] = [] + stage_replica_counts: dict[int, int] = {} + seen_stage_ids: set[int] = set() + for plan in stage_plans: + stage_id = plan.configured_stage_id + if stage_id in seen_stage_ids: + raise ValueError( + f"Duplicate stage_id {stage_id!r} detected among configured stages; stage_ids must be unique." + ) + seen_stage_ids.add(stage_id) + all_stage_ids.append(stage_id) + stage_replica_counts[stage_id] = len(plan.replicas) + + self._omni_master_server = OmniMasterServer( + master_address=self._omni_master_address, + master_port=self._omni_master_port, + stage_ids=all_stage_ids, + stage_replica_counts=stage_replica_counts, + ) + self._omni_master_server.start() + logger.info( + "[AsyncOmniEngine] OmniMasterServer started for stages %s", + all_stage_ids, + ) + + def _initialize_llm_replica( + self, + plan: ReplicaInitPlan, stage_init_timeout: int, llm_stage_launch_lock: threading.Lock, - omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None] = (None, None, None), - ) -> StartedLlmStage: - """Launch one LLM stage to READY state in a helper thread.""" - started_stage: StartedLlmStage | None = None + ) -> Any: + """Initialize one LLM replica end-to-end.""" + + proc = None + engine_manager = None + coordinator = None + stage_client = None lock_fds: list[int] = [] device_control_env = current_omni_platform.device_control_env_var + stage_cfg = plan.stage_cfg + try: - proc = None - handshake_address = None - with ExitStack() as launch_stack: - with llm_stage_launch_lock: - previous_visible_devices = os.environ.get(device_control_env) - try: - setup_stage_devices(metadata.stage_id, metadata.runtime_cfg) - engine_args_dict = build_engine_args_dict( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - ) - omni_conn_cfg, omni_from, omni_to = omni_kv_connector - if omni_conn_cfg: - omni_kv = engine_args_dict.get("omni_kv_config") or {} - if not isinstance(omni_kv, dict): - omni_kv = dict(omni_kv) - omni_kv["connector_config"] = omni_conn_cfg - omni_kv["omni_from_stage"] = omni_from - omni_kv["omni_to_stage"] = omni_to - omni_kv.setdefault("stage_id", metadata.stage_id) - engine_args_dict["omni_kv_config"] = omni_kv - if self.stage_configs: - _inject_inferred_kv_tp_topology( - engine_args_dict.get("omni_kv_config"), - metadata.stage_id, - self.stage_configs, + if plan.launch_mode == "remote": + assert self._omni_master_server is not None + raw_stage_cfg = self._omni_master_server.get_stage_config( + plan.metadata.stage_id, + timeout_s=stage_init_timeout, + replica_id=plan.replica_id, + ) + if raw_stage_cfg is None: + raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config") + vllm_config = plan.stage_vllm_config + executor_class = plan.executor_class + assert vllm_config is not None + assert executor_class is not None + vllm_config.parallel_config.data_parallel_size_local = 0 + launch_cm = connect_remote_engine_cores( + vllm_config=vllm_config, + omni_master_server=self._omni_master_server, + stage_id=plan.metadata.stage_id, + replica_id=plan.replica_id, + ) + logger.info( + "[AsyncOmniEngine] Stage %s remote engine handshake started", + plan.metadata.stage_id, + ) + with launch_cm as remote_resources: + engine_manager, coordinator, addresses, _tensor_queue = remote_resources + + logger.info( + "[AsyncOmniEngine] Stage %s remote engine startup completed", + plan.metadata.stage_id, + ) + client_addresses: dict[str, str] = { + "input_address": addresses.inputs[0], + "output_address": addresses.outputs[0], + } + if addresses.frontend_stats_publish_address is not None: + client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address + stage_client = StageEngineCoreClientBase.make_async_mp_client( + vllm_config=vllm_config, + executor_class=executor_class, + metadata=plan.metadata, + client_addresses=client_addresses, + engine_manager=engine_manager, + coordinator=coordinator, + ) + else: + handshake_address = None + with ExitStack() as launch_stack: + with llm_stage_launch_lock: + previous_visible_devices = os.environ.get(device_control_env) + try: + setup_stage_devices(plan.metadata.stage_id, plan.metadata.runtime_cfg) + vllm_config = plan.stage_vllm_config + executor_class = plan.executor_class + assert vllm_config is not None + assert executor_class is not None + engine_args_dict = build_engine_args_dict( + stage_cfg, + self.model, + stage_connector_spec=plan.stage_connector_spec, ) - vllm_config, executor_class = build_vllm_config( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - engine_args_dict=engine_args_dict, - ) - lock_fds = acquire_device_locks( - metadata.stage_id, - engine_args_dict, - stage_init_timeout, - ) - if self.single_stage_mode and self._omni_master_server is not None: - engine_manager, coordinator, addresses = launch_stack.enter_context( - launch_omni_core_engines( + lock_fds = acquire_device_locks( + plan.metadata.stage_id, + engine_args_dict, + stage_init_timeout, + ) + if self.single_stage_mode and self._omni_master_server is not None: + engine_manager, coordinator, addresses = launch_stack.enter_context( + launch_omni_core_engines( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=False, + omni_master_server=self._omni_master_server, + stage_id=plan.metadata.stage_id, + stage_config=stage_cfg, + replica_id=plan.replica_id, + ) + ) + else: + addresses, proc, handshake_address = spawn_stage_core( vllm_config=vllm_config, executor_class=executor_class, log_stats=False, - omni_master_server=self._omni_master_server, - stage_id=metadata.stage_id, - stage_config=stage_cfg, ) + logger.info( + "[AsyncOmniEngine] Stage %s engine launch started", + plan.metadata.stage_id, ) - started_stage = StartedLlmStage( - stage_id=metadata.stage_id, - metadata=metadata, - vllm_config=vllm_config, - executor_class=executor_class, - addresses=addresses, - engine_manager=engine_manager, - coordinator=coordinator, + finally: + if previous_visible_devices is None: + current_omni_platform.unset_device_control_env_var() + else: + current_omni_platform.set_device_control_env_var(previous_visible_devices) + + if self.single_stage_mode and self._omni_master_server is not None: + launch_stack.close() + else: + assert proc is not None + assert handshake_address is not None + complete_stage_handshake(proc, handshake_address, addresses, vllm_config, stage_init_timeout) + logger.info( + "[AsyncOmniEngine] Stage %s engine startup completed", + plan.metadata.stage_id, + ) + + client_addresses: dict[str, str] = { + "input_address": addresses.inputs[0], + "output_address": addresses.outputs[0], + } + if addresses.frontend_stats_publish_address is not None: + client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address + stage_client = StageEngineCoreClientBase.make_async_mp_client( + vllm_config=vllm_config, + executor_class=executor_class, + metadata=plan.metadata, + client_addresses=client_addresses, + proc=proc, + engine_manager=engine_manager, + coordinator=coordinator, + ) + + logger.info("[AsyncOmniEngine] Stage %s initialized", plan.metadata.stage_id) + return stage_client + except Exception: + if stage_client is not None: + try: + stage_client.shutdown() + except Exception as cleanup_error: + logger.warning( + "[AsyncOmniEngine] Failed to cleanup stage %s after attach failure: %s", + plan.metadata.stage_id, + cleanup_error, + ) + else: + self._cleanup_launched_llm_resources( + stage_id=plan.metadata.stage_id, + proc=proc, + engine_manager=engine_manager, + coordinator=coordinator, + ) + raise + finally: + if lock_fds: + release_device_locks(lock_fds) + + def _initialize_diffusion_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + stage_launch_lock: threading.Lock, + ) -> Any: + """Initialize one diffusion replica end-to-end.""" + + client = None + proc = None + lock_fds: list[int] = [] + try: + if plan.launch_mode == "remote": + assert self._omni_master_server is not None + remote_stage_cfg = OmegaConf.create( + self._omni_master_server.get_stage_config( + plan.metadata.stage_id, + timeout_s=stage_init_timeout, + replica_id=plan.replica_id, + ) + ) + remote_metadata = extract_stage_metadata(remote_stage_cfg) + addresses = self._omni_master_server.get_zmq_addresses( + plan.metadata.stage_id, + replica_id=plan.replica_id, + ) + logger.info( + "[AsyncOmniEngine] Stage %s remote diffusion startup completed", + plan.metadata.stage_id, + ) + client = StageDiffusionClient.from_addresses( + remote_metadata, + request_address=addresses.inputs[0], + response_address=addresses.outputs[0], + batch_size=self.diffusion_batch_size, + ) + else: + device_control_env = current_omni_platform.device_control_env_var + with stage_launch_lock: + previous_visible_devices = os.environ.get(device_control_env) + try: + setup_stage_devices(plan.metadata.stage_id, plan.metadata.runtime_cfg) + omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector + if omni_conn_cfg: + inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to) + inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self.stage_configs) + if self.single_stage_mode: + assert self._omni_master_server is not None + od_config = build_diffusion_config(self.model, plan.stage_cfg, plan.metadata) + lock_fds = acquire_diffusion_device_locks( + plan.metadata.stage_id, + od_config, + stage_init_timeout, + ) + handshake_address, request_address, response_address = register_stage_with_omni_master( + omni_master_address=self._omni_master_server.address, + omni_master_port=self._omni_master_server.port, + omni_stage_id=plan.metadata.stage_id, + omni_stage_config=plan.stage_cfg, + return_addresses=True, + replica_id=plan.replica_id, + ) + logger.info( + "[AsyncOmniEngine] Stage %s diffusion registration completed", + plan.metadata.stage_id, + ) + proc, _, _, _ = spawn_diffusion_proc( + self.model, + od_config, + handshake_address=handshake_address, + request_address=request_address, + response_address=response_address, + ) + complete_diffusion_handshake(proc, handshake_address, stage_init_timeout) + logger.info( + "[AsyncOmniEngine] Stage %s diffusion startup completed", + plan.metadata.stage_id, + ) + client = StageDiffusionClient.from_addresses( + plan.metadata, + request_address=request_address, + response_address=response_address, + proc=proc, + batch_size=self.diffusion_batch_size, ) else: - addresses, proc, handshake_address = spawn_stage_core( - vllm_config=vllm_config, - executor_class=executor_class, - log_stats=False, + client = initialize_diffusion_stage( + plan.metadata.stage_id, + self.model, + plan.stage_cfg, + plan.metadata, + stage_init_timeout=stage_init_timeout, + batch_size=self.diffusion_batch_size, + use_inline=self.num_stages == 1 and plan.num_replicas == 1, ) - started_stage = StartedLlmStage( - stage_id=metadata.stage_id, - metadata=metadata, - vllm_config=vllm_config, - executor_class=executor_class, - addresses=addresses, - proc=proc, - ) - logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id) finally: if previous_visible_devices is None: current_omni_platform.unset_device_control_env_var() else: current_omni_platform.set_device_control_env_var(previous_visible_devices) - # After StageEngineCoreProc has been spawned it carries its - # stage-specific device visibility into descendants, so the - # slow HELLO/READY handshake can run without holding the - # process-wide launch lock. - if self.single_stage_mode and self._omni_master_server is not None: - launch_stack.close() - else: - assert proc is not None - assert handshake_address is not None - complete_stage_handshake(proc, handshake_address, addresses, vllm_config, stage_init_timeout) - logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id) - - assert started_stage is not None - return started_stage - except Exception: - if started_stage is not None: - close_started_llm_stage(started_stage) - raise - finally: - if lock_fds: - release_device_locks(lock_fds) - - def _create_remote_llm_stage( - self, - stage_cfg: Any, - metadata: Any, - stage_connector_spec: dict[str, Any], - stage_init_timeout: int, - omni_master_server: OmniMasterServer, - ) -> StartedLlmStage: - """Attach to a remote engine core and wait for its startup handshake.""" - started_stage: StartedLlmStage | None = None - try: - raw_stage_cfg = omni_master_server.get_stage_config( - metadata.stage_id, - timeout_s=stage_init_timeout, - ) - if raw_stage_cfg is None: - raise ValueError(f"Remote stage {metadata.stage_id} registered without stage config") - stage_cfg = OmegaConf.create(raw_stage_cfg) - engine_args_dict = build_engine_args_dict( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - ) - vllm_config, executor_class = build_vllm_config( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - engine_args_dict=engine_args_dict, - ) - vllm_config.parallel_config.data_parallel_size_local = 0 - launch_cm = connect_remote_engine_cores( - vllm_config=vllm_config, - omni_master_server=omni_master_server, - stage_id=metadata.stage_id, - ) - logger.info("[AsyncOmniEngine] Stage %s remote engine handshake started", metadata.stage_id) - (engine_manager, coordinator, addresses, _tensor_queue) = launch_cm.__enter__() - try: - started_stage = StartedLlmStage( - stage_id=metadata.stage_id, - metadata=metadata, - vllm_config=vllm_config, - executor_class=executor_class, - engine_manager=engine_manager, - coordinator=coordinator, - addresses=addresses, - ) - except BaseException: - if not launch_cm.__exit__(*sys.exc_info()): - raise - else: - launch_cm.__exit__(None, None, None) - logger.info("[AsyncOmniEngine] Stage %s remote engine startup completed", metadata.stage_id) - assert started_stage is not None - return started_stage - except Exception: - if started_stage is not None: - close_started_llm_stage(started_stage) - raise - - def _launch_diffusion_stage( - self, - stage_cfg: Any, - metadata: Any, - omni_master_server: OmniMasterServer, - stage_init_timeout: int, - ) -> StageDiffusionClient: - """Launch a local diffusion stage on OmniMasterServer-allocated sockets.""" - proc = None - lock_fds: list[int] = [] - try: - od_config = build_diffusion_config(self.model, stage_cfg, metadata) - lock_fds = acquire_diffusion_device_locks(metadata.stage_id, od_config, stage_init_timeout) - handshake_address, request_address, response_address = register_stage_with_omni_master( - omni_master_address=omni_master_server.address, - omni_master_port=omni_master_server.port, - omni_stage_id=metadata.stage_id, - omni_stage_config=stage_cfg, - return_addresses=True, - ) logger.info( - "[AsyncOmniEngine] Stage %s diffusion registration completed", - metadata.stage_id, - ) - proc, _, _, _ = spawn_diffusion_proc( - self.model, - od_config, - handshake_address=handshake_address, - request_address=request_address, - response_address=response_address, - ) - complete_diffusion_handshake(proc, handshake_address, stage_init_timeout) - logger.info( - "[AsyncOmniEngine] Stage %s diffusion startup completed", - metadata.stage_id, - ) - return StageDiffusionClient.from_addresses( - metadata, - request_address=request_address, - response_address=response_address, - proc=proc, - batch_size=self.diffusion_batch_size, + "[AsyncOmniEngine] Stage %s replica %s initialized (diffusion, batch_size=%d, devices=%s)", + plan.metadata.stage_id, + plan.replica_id, + self.diffusion_batch_size, + getattr(getattr(plan.stage_cfg, "runtime", None), "devices", "default"), ) + return client except Exception: if proc is not None: terminate_alive_proc(proc) @@ -604,305 +856,170 @@ class AsyncOmniEngine: if lock_fds: release_device_locks(lock_fds) - def _create_remote_diffusion_stage( + def _initialize_replica( self, - metadata: Any, + plan: ReplicaInitPlan, stage_init_timeout: int, - omni_master_server: OmniMasterServer, - ) -> StageDiffusionClient: - """Attach to a remote diffusion stage registered with OmniMasterServer.""" - remote_stage_cfg = OmegaConf.create( - omni_master_server.get_stage_config( - metadata.stage_id, - timeout_s=stage_init_timeout, - ) - ) - remote_metadata = extract_stage_metadata(remote_stage_cfg) - addresses = omni_master_server.get_zmq_addresses(metadata.stage_id) - logger.info( - "[AsyncOmniEngine] Stage %s remote diffusion startup completed", - metadata.stage_id, - ) - return StageDiffusionClient.from_addresses( - remote_metadata, - request_address=addresses.inputs[0], - response_address=addresses.outputs[0], - batch_size=self.diffusion_batch_size, - ) + stage_launch_lock: threading.Lock, + ) -> Any: + """Initialize one replica, regardless of backend type.""" - def _attach_llm_stage( + if plan.metadata.stage_type == "diffusion": + return self._initialize_diffusion_replica(plan, stage_init_timeout, stage_launch_lock) + return self._initialize_llm_replica(plan, stage_init_timeout, stage_launch_lock) + + def _initialize_stage_replicas( self, - started: StartedLlmStage, - ) -> tuple[Any, Any, Any, InputProcessor | None]: - """Attach a READY LLM stage to the orchestrator event loop.""" + stage_plans: Sequence[LogicalStageInitPlan], + stage_init_timeout: int, + ) -> dict[int, list[Any | None]]: + """Initialize all stage replicas in parallel.""" - client_addresses: dict[str, str] = { - "input_address": started.addresses.inputs[0], - "output_address": started.addresses.outputs[0], + stage_launch_lock = threading.Lock() + initialized_clients_by_stage: dict[int, list[Any | None]] = { + plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans } - if started.addresses.frontend_stats_publish_address is not None: - client_addresses["stats_update_address"] = started.addresses.frontend_stats_publish_address + total_replicas = sum(len(plan.replicas) for plan in stage_plans) + future_to_replica: dict[concurrent.futures.Future[Any], tuple[int, int]] = {} + primary_exc: Exception | None = None - try: - stage_client = StageEngineCoreClientBase.make_async_mp_client( - vllm_config=started.vllm_config, - executor_class=started.executor_class, - metadata=started.metadata, - client_addresses=client_addresses, - proc=started.proc, - engine_manager=started.engine_manager, - coordinator=started.coordinator, + with concurrent.futures.ThreadPoolExecutor( + max_workers=max(1, total_replicas), + thread_name_prefix="stage-init", + ) as init_executor: + for plan in stage_plans: + for replica in plan.replicas: + future = init_executor.submit( + self._initialize_replica, + replica, + stage_init_timeout, + stage_launch_lock, + ) + future_to_replica[future] = (plan.stage_idx, replica.replica_id) + + for future in concurrent.futures.as_completed(future_to_replica): + stage_idx, replica_id = future_to_replica[future] + try: + initialized_clients_by_stage[stage_idx][replica_id] = future.result() + except concurrent.futures.CancelledError: + continue + except Exception as exc: + if primary_exc is None: + primary_exc = exc + for other_future in future_to_replica: + if other_future is future: + continue + other_future.cancel() + + if primary_exc is not None: + setattr(primary_exc, "_initialized_clients_by_stage", initialized_clients_by_stage) + raise primary_exc + + return initialized_clients_by_stage + + def _assemble_stage_pools( + self, + stage_plans: Sequence[LogicalStageInitPlan], + initialized_clients_by_stage: Mapping[int, Sequence[Any | None]], + ) -> list[StagePool]: + """Assemble logical stage pools and update top-level stage metadata.""" + + stage_pools: list[StagePool] = [] + default_sampling_params_list: list[Any] = [] + stage_metadata_list: list[dict[str, Any]] = [] + + for plan in stage_plans: + replica_clients = initialized_clients_by_stage[plan.stage_idx] + first_client = replica_clients[0] if replica_clients else None + if first_client is None: + raise RuntimeError(f"Stage {plan.stage_idx} initialization completed with a missing client") + + clients = [client for client in replica_clients if client is not None] + stage_vllm_config = None + output_processor = None + if plan.replicas[0].metadata.stage_type != "diffusion": + stage_vllm_config = plan.replicas[0].stage_vllm_config + assert stage_vllm_config is not None + output_processor = build_llm_stage_output_processor(plan, stage_vllm_config) + + stage_pools.append( + StagePool( + plan.stage_idx, + clients, + output_processor=output_processor, + stage_vllm_config=stage_vllm_config, + ) ) - started.proc = None - started.engine_manager = None - started.coordinator = None - except Exception: - close_started_llm_stage(started) - raise - - try: - if started.vllm_config.model_config.skip_tokenizer_init: - tokenizer = None - else: - tokenizer = cached_tokenizer_from_config( - model_config=started.vllm_config.model_config, - ) - output_processor = MultimodalOutputProcessor( - tokenizer=tokenizer, - log_stats=False, - engine_core_output_type=started.metadata.engine_output_type, + default_sampling_params_list.append(first_client.default_sampling_params) + stage_metadata_list.append( + { + "final_output": first_client.final_output, + "final_output_type": first_client.final_output_type, + "stage_type": first_client.stage_type, + } ) - input_processor = None - if started.stage_id == 0: - # Some omni models (e.g. CosyVoice3) have an empty HF - # config.json without model_type, which causes - # try_get_generation_config -> AutoConfig.from_pretrained - # to raise ValueError. Patch it to return None so - # InputProcessor doesn't crash. - _patch_generation_config_if_needed(started.vllm_config.model_config) - input_processor = InputProcessor(vllm_config=started.vllm_config) - # Use omni preprocessor so text-only prompts with - # mm_processor_kwargs (e.g. GLM-Image t2i target_h/target_w) - # still go through multimodal processor path. - input_processor.input_preprocessor = OmniInputPreprocessor( - vllm_config=started.vllm_config, - renderer=input_processor.renderer, - ) - except Exception: - try: - stage_client.shutdown() - except Exception as cleanup_error: - logger.warning( - "[AsyncOmniEngine] Failed to cleanup stage %s after attach failure: %s", - started.stage_id, - cleanup_error, - ) - raise - logger.info("[AsyncOmniEngine] Stage %s initialized", started.stage_id) - return stage_client, output_processor, started.vllm_config, input_processor + self.default_sampling_params_list = list(default_sampling_params_list) + self.stage_metadata = list(stage_metadata_list) + return stage_pools def _initialize_stages(self, stage_init_timeout: int) -> None: - """Initialize stage clients/processors in orchestrator thread and assign to self.""" - device_control_env = current_omni_platform.device_control_env_var + """Initialize stage clients/processors in orchestrator thread and assign to self. - num_stages = self.num_stages - stage_clients: list[Any | None] = [None] * num_stages - output_processors: list[Any | None] = [None] * num_stages - stage_vllm_configs: list[Any | None] = [None] * num_stages - input_processor: InputProcessor | None = None - llm_stage_positions: list[int] = [] - llm_launch_futures: dict[int, concurrent.futures.Future[StartedLlmStage]] = {} - started_llm_stages: dict[int, StartedLlmStage] = {} - llm_stage_launch_lock = threading.Lock() + Phases: + 1. Compute replica layout (counts + device splits). + 2. Build per-stage/per-replica startup plans. + 3. Initialize all replicas in parallel via backend-specific launchers. + 4. Build logical StagePools and finalize runtime metadata. - async_chunk = self.async_chunk - prompt_expand_func = None - llm_stage_count = sum( - 1 for stage_cfg in self.stage_configs if getattr(stage_cfg, "stage_type", "llm") != "diffusion" - ) + TODO(stage-pool): move per-stage launch + attach logic into a + StagePool.build_from_config() classmethod so this method only + iterates stage_configs, collects pools, and finalizes metadata. + """ + num_stages = len(self.stage_configs) + self.num_stages = num_stages + self._validate_single_stage_mode_replica_constraints() + + replicas_per_stage, replica_devices_map = compute_replica_layout(self.stage_configs) prepare_engine_environment() omni_transfer_config = load_omni_transfer_config_for_model(self.model, self.config_path) - - # ------------------------------------------------------------------ # - # Single-stage mode: start OmniMasterServer before launching stages. # - # ------------------------------------------------------------------ # + stage_plans, prompt_expand_func = self._build_logical_stage_init_plans( + omni_transfer_config, + replicas_per_stage, + replica_devices_map, + ) if self.single_stage_mode: - if not self._omni_master_address or not self._omni_master_port: - raise ValueError( - "AsyncOmniEngine single_stage_mode requires both " - "omni_master_address and omni_master_port to be set." - ) - # Collect all configured stage IDs for pre-allocation. - all_stage_ids: list[int] = [] - seen_stage_ids: set[int] = set() - for i, sc in enumerate(self.stage_configs): - stage_id = int(getattr(sc, "stage_id", i)) - if stage_id in seen_stage_ids: - raise ValueError( - f"Duplicate stage_id {stage_id!r} detected among configured stages; stage_ids must be unique." - ) - seen_stage_ids.add(stage_id) - all_stage_ids.append(stage_id) - self._omni_master_server = OmniMasterServer( - master_address=self._omni_master_address, - master_port=self._omni_master_port, - stage_ids=all_stage_ids, - ) - self._omni_master_server.start() - logger.info( - "[AsyncOmniEngine] OmniMasterServer started for stages %s", - all_stage_ids, - ) + self._start_omni_master_server(stage_plans) + + stage_pools: list[StagePool] = [] + input_processor: InputProcessor | None = None + initialized_clients_by_stage: dict[int, list[Any | None]] = { + plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans + } try: - with concurrent.futures.ThreadPoolExecutor( - max_workers=max(1, llm_stage_count), - thread_name_prefix="llm-stage-launch", - ) as launch_executor: - for stage_idx, stage_cfg in enumerate(self.stage_configs): - metadata = extract_stage_metadata(stage_cfg) - configured_stage_id = metadata.stage_id - logger.info("[AsyncOmniEngine] Initializing stage %s", configured_stage_id) - if metadata.prompt_expand_func is not None: - prompt_expand_func = metadata.prompt_expand_func - - if self.single_stage_mode: - metadata.runtime_cfg = None - - stage_connector_spec = get_stage_connector_spec( - omni_transfer_config=omni_transfer_config, - stage_id=configured_stage_id, - async_chunk=async_chunk, - ) - - omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, configured_stage_id) - - if metadata.stage_type == "diffusion": - is_remote_diffusion_stage = ( - self.single_stage_mode - and self._single_stage_id_filter is not None - and configured_stage_id != self._single_stage_id_filter - ) - if is_remote_diffusion_stage: - assert self._omni_master_server is not None - stage_clients[stage_idx] = self._create_remote_diffusion_stage( - metadata, - stage_init_timeout, - self._omni_master_server, - ) - continue - - with llm_stage_launch_lock: - previous_visible_devices = os.environ.get(device_control_env) - try: - setup_stage_devices(configured_stage_id, metadata.runtime_cfg) - omni_conn_cfg, omni_from, omni_to = omni_kv_connector - if omni_conn_cfg: - inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) - inject_kv_stage_info(stage_cfg, configured_stage_id, self.stage_configs) - if self.single_stage_mode: - assert self._omni_master_server is not None - stage_clients[stage_idx] = self._launch_diffusion_stage( - stage_cfg, - metadata, - self._omni_master_server, - stage_init_timeout, - ) - else: - use_inline = True if self.num_stages == 1 else False - stage_clients[stage_idx] = initialize_diffusion_stage( - configured_stage_id, - self.model, - stage_cfg, - metadata, - stage_init_timeout=stage_init_timeout, - batch_size=self.diffusion_batch_size, - use_inline=use_inline, - ) - logger.info( - "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)", - configured_stage_id, - self.diffusion_batch_size, - ) - finally: - if previous_visible_devices is None: - current_omni_platform.unset_device_control_env_var() - else: - current_omni_platform.set_device_control_env_var(previous_visible_devices) - continue - - llm_stage_positions.append(stage_idx) - - # In single-stage mode, stages that don't match the local - # stage_id filter are skipped. - if ( - self.single_stage_mode - and self._single_stage_id_filter is not None - and configured_stage_id != self._single_stage_id_filter - ): - assert self._omni_master_server is not None - llm_launch_futures[stage_idx] = launch_executor.submit( - self._create_remote_llm_stage, - stage_cfg, - metadata, - stage_connector_spec, - stage_init_timeout, - self._omni_master_server, - ) - else: - llm_launch_futures[stage_idx] = launch_executor.submit( - self._launch_llm_stage, - stage_cfg, - metadata, - stage_connector_spec, - stage_init_timeout, - llm_stage_launch_lock, - omni_kv_connector, - ) - - concurrent.futures.wait(list(llm_launch_futures.values())) - - for stage_idx in llm_stage_positions: - started_llm_stages[stage_idx] = llm_launch_futures[stage_idx].result() - - attach_futures: dict[concurrent.futures.Future[tuple[Any, Any, Any, InputProcessor | None]], int] = {} - with concurrent.futures.ThreadPoolExecutor( - max_workers=max(1, len(llm_stage_positions)), - thread_name_prefix="llm-stage-attach", - ) as attach_executor: - for stage_idx in llm_stage_positions: - attach_futures[attach_executor.submit(self._attach_llm_stage, started_llm_stages[stage_idx])] = ( - stage_idx - ) - - for future in concurrent.futures.as_completed(attach_futures): - stage_idx = attach_futures[future] - stage_client, output_processor, vllm_config, stage0_input_processor = future.result() - stage_clients[stage_idx] = stage_client - output_processors[stage_idx] = output_processor - stage_vllm_configs[stage_idx] = vllm_config - if stage0_input_processor is not None: - input_processor = stage0_input_processor - - initialized_stage_clients, default_sampling_params_list, stage_metadata = finalize_initialized_stages( - stage_clients, - input_processor, + initialized_clients_by_stage = self._initialize_stage_replicas(stage_plans, stage_init_timeout) + if stage_plans and stage_plans[0].replicas[0].metadata.stage_type != "diffusion": + stage0_vllm_config = stage_plans[0].replicas[0].stage_vllm_config + assert stage0_vllm_config is not None + input_processor = build_stage0_input_processor(stage0_vllm_config) + stage_pools = self._assemble_stage_pools(stage_plans, initialized_clients_by_stage) + except Exception as exc: + initialized_clients_by_stage = getattr( + exc, + "_initialized_clients_by_stage", + initialized_clients_by_stage, + ) + cleanup_clients = self._collect_initialized_clients_for_cleanup( + stage_pools, + initialized_clients_by_stage, ) - except Exception: - for stage_id, future in llm_launch_futures.items(): - if not future.done() or future.cancelled() or future.exception() is not None: - continue - started_llm_stages.setdefault(stage_id, future.result()) logger.exception( - "[AsyncOmniEngine] Stage initialization failed; shutting down %s initialized stage(s)", - len([stage_client for stage_client in stage_clients if stage_client is not None]), - ) - cleanup_failed_stage_initialization( - stage_clients, - [started_llm_stages[stage_idx] for stage_idx in llm_stage_positions if stage_idx in started_llm_stages], + "[AsyncOmniEngine] Stage initialization failed; shutting down %s initialized client(s)", + len(cleanup_clients), ) + self._shutdown_initialized_clients(cleanup_clients) if self._omni_master_server is not None: try: self._omni_master_server.stop() @@ -910,22 +1027,23 @@ class AsyncOmniEngine: logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during stage-init cleanup") raise - self.stage_clients = initialized_stage_clients - self.output_processors = output_processors - self.stage_vllm_configs = stage_vllm_configs + self.stage_pools = stage_pools self.input_processor = input_processor self.prompt_expand_func = prompt_expand_func + + # Derive logical-stage views for external readers (entrypoints/async_omni.py). + self.stage_clients = [pool.stage_client for pool in self.stage_pools] + self.stage_vllm_configs = [pool.stage_vllm_config for pool in self.stage_pools] + self.output_processors = [pool.output_processor for pool in self.stage_pools] + # TODO(Peiqi): Hack here supported_tasks: set[str] = set() - if any(getattr(stage_client, "is_comprehension", False) for stage_client in initialized_stage_clients): + if any(getattr(pool.stage_client, "is_comprehension", False) for pool in self.stage_pools): supported_tasks.add("generate") - if any(metadata.get("final_output_type") == "audio" for metadata in stage_metadata): + if any(m.get("final_output_type") == "audio" for m in self.stage_metadata): supported_tasks.add("speech") self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",) - self.default_sampling_params_list = default_sampling_params_list - self.stage_metadata = stage_metadata - def _initialize_janus_queues(self) -> None: """Initialize janus queues inside orchestrator thread loop context.""" self.request_queue = janus.Queue() @@ -952,10 +1070,8 @@ class AsyncOmniEngine: request_async_queue=self.request_queue.async_q, output_async_queue=self.output_queue.async_q, rpc_async_queue=self.rpc_output_queue.async_q, + stage_pools=self.stage_pools, async_chunk=self.async_chunk, - stage_clients=self.stage_clients, - output_processors=self.output_processors, - stage_vllm_configs=self.stage_vllm_configs, pd_config=pd_config, ) if not startup_future.done(): @@ -1056,6 +1172,7 @@ class AsyncOmniEngine: original_prompt = prompt stage_type = self.stage_metadata[0].get("stage_type") + output_prompt_text: Any = None _preprocess_ms = 0.0 if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest): # Inject global_request_id into the raw prompt. @@ -1097,17 +1214,12 @@ class AsyncOmniEngine: request.external_req_id = request_id request = _apply_omni_final_stage_metadata(request, final_stage_id) - # Register with stage 0's output processor. + # Registration with stage 0's output processor is deferred to the + # orchestrator thread (see Orchestrator._handle_add_request), which + # now routes admission through StagePool.submit_initial(). output_prompt_text = prompt_text if output_prompt_text is None and isinstance(original_prompt, dict): output_prompt_text = original_prompt.get("prompt") - self.output_processors[0].add_request( - request=request, - prompt=output_prompt_text, - parent_req=None, - request_index=0, - queue=None, - ) prompt = request return { @@ -1115,6 +1227,7 @@ class AsyncOmniEngine: "request_id": request_id, "prompt": prompt, "original_prompt": original_prompt, + "output_prompt_text": output_prompt_text, "sampling_params_list": effective_sampling_params_list, "final_stage_id": final_stage_id, "preprocess_ms": _preprocess_ms, @@ -1155,14 +1268,9 @@ class AsyncOmniEngine: ) request.external_req_id = cid - self.output_processors[0].add_request( - request=request, - prompt=companion_prompt, - parent_req=None, - request_index=0, - queue=None, - ) - + # Registration of this companion on stage-0's output processor is + # deferred to Orchestrator._handle_add_companion, which routes + # admission through StagePool.submit_initial(..., affinity_request_id=...). self.request_queue.sync_q.put_nowait( { "type": "add_companion_request", @@ -1170,6 +1278,7 @@ class AsyncOmniEngine: "parent_id": parent_id, "role": ep.role, "prompt": request, + "companion_prompt_text": companion_prompt, "sampling_params_list": companion_spl, } ) @@ -1475,6 +1584,8 @@ class AsyncOmniEngine: cfg.engine_args.enable_sleep_mode = global_sleep_mode if getattr(cfg, "stage_type", None) != "diffusion": continue + if not hasattr(cfg, "engine_args") or cfg.engine_args is None: + cfg.engine_args = OmegaConf.create({}) if kwargs.get("lora_path") is not None: if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None: cfg.engine_args.lora_path = kwargs["lora_path"] diff --git a/vllm_omni/engine/cfg_companion_tracker.py b/vllm_omni/engine/cfg_companion_tracker.py index b9dfae833..1105e887c 100644 --- a/vllm_omni/engine/cfg_companion_tracker.py +++ b/vllm_omni/engine/cfg_companion_tracker.py @@ -10,8 +10,6 @@ from __future__ import annotations import logging from typing import Any -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - logger = logging.getLogger(__name__) @@ -31,6 +29,10 @@ class CfgCompanionTracker: def has_companions(self, parent_id: str) -> bool: return parent_id in self._companion_map + def get_parent_id(self, req_id: str) -> str | None: + """Return the parent request id for a companion, or None.""" + return self._companion_to_parent.get(req_id) + def all_companions_done(self, parent_id: str) -> bool: role_map = self._companion_map.get(parent_id, {}) done_set = self._done.get(parent_id, set()) @@ -50,21 +52,6 @@ class CfgCompanionTracker: self._companion_ids.add(companion_id) self._companion_to_parent[companion_id] = parent_id - def attach_cfg_request_ids(self, parent_id: str, sampling_params: Any) -> Any: - cfg_ids = self.get_companion_request_ids(parent_id) - if not cfg_ids: - return sampling_params - - if isinstance(sampling_params, OmniDiffusionSamplingParams): - sampling_params = sampling_params.clone() - sampling_params.cfg_kv_request_ids = cfg_ids - logger.info( - "Attaching cfg_kv_request_ids=%s to request %s", - cfg_ids, - parent_id, - ) - return sampling_params - def on_companion_completed(self, companion_id: str) -> str | None: """Mark done. Returns parent_id only if parent is pending and all companions finished.""" parent_id = self._companion_to_parent.get(companion_id) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index b53bdae84..584ea6c5d 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -2,14 +2,13 @@ Orchestrator for vLLM-Omni multi-stage runtime. Runs inside a background thread with its own asyncio event loop. -Owns all StageEngineCoreClient instances, input/output processors, -and handles stage-to-stage transfer logic. +Owns logical request progression across stage pools and handles +stage-to-stage transfer logic. """ from __future__ import annotations import asyncio -import copy import time as _time from dataclasses import dataclass, field from typing import Any @@ -18,21 +17,15 @@ import janus import torch from vllm.config import ModelConfig from vllm.logger import init_logger -from vllm.outputs import RequestOutput from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreOutputs from vllm.v1.engine.exceptions import EngineDeadError -from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length -from vllm_omni.engine import ( - OmniEngineCoreRequest, -) +from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.serialization import serialize_additional_information -from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics -from vllm_omni.metrics.stats import StageStats -from vllm_omni.metrics.utils import count_tokens_from_outputs +from vllm_omni.engine.stage_pool import StagePool from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -47,19 +40,12 @@ def build_engine_core_request_from_tokens( resumable: bool = False, mm_features: list | None = None, ) -> OmniEngineCoreRequest: - """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt. - - Lightweight alternative to the full InputProcessor pipeline - skips - tokenization, multimodal preprocessing, LoRA validation, and platform - validation. Intended for stage 1+ where the upstream stage has already - produced token IDs and optional embeddings. - """ + """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt.""" if arrival_time is None: arrival_time = _time.time() prompt_token_ids = prompt["prompt_token_ids"] - # Clone params and set max_tokens if needed sampling_params = None pooling_params = None if isinstance(params, SamplingParams): @@ -70,8 +56,6 @@ def build_engine_core_request_from_tokens( pooling_params = params.clone() prompt_embeds: torch.Tensor | None = prompt.get("prompt_embeds") - - # Serialize additional_information if present additional_info_payload = serialize_additional_information( prompt.get("additional_information"), log_prefix=f"build_engine_core_request_from_tokens req={request_id}", @@ -93,11 +77,6 @@ def build_engine_core_request_from_tokens( ) -# ============================================================ -# Orchestrator internals (run inside the background thread) -# ============================================================ - - @dataclass class OrchestratorRequestState: """Per-request bookkeeping inside the Orchestrator.""" @@ -107,10 +86,11 @@ class OrchestratorRequestState: sampling_params_list: list[Any] = field(default_factory=list) final_stage_id: int = -1 - # Metrics: timestamp when request was submitted to each stage + # Metrics: timestamp when request was submitted to each stage. stage_submit_ts: dict[int, float] = field(default_factory=dict) mm_processor_kwargs: dict | None = None mm_features: list | None = None + pd_prefill_multimodal_output: dict[str, Any] | None = None streaming: StreamingInputState = field(default_factory=lambda: StreamingInputState()) @@ -131,20 +111,14 @@ class StreamingInputState: class Orchestrator: - """Runs inside a background thread's asyncio event loop. - - Owns all StageEngineCoreClient instances, input/output processors, - and handles stage-to-stage transfer logic. - """ + """Runs inside a background thread's asyncio event loop.""" def __init__( self, request_async_queue: janus.AsyncQueue[dict[str, Any]], output_async_queue: janus.AsyncQueue[dict[str, Any]], rpc_async_queue: janus.AsyncQueue[dict[str, Any]], - stage_clients: list[Any], - output_processors: list[Any], - stage_vllm_configs: list[Any], + stage_pools: list[StagePool], *, async_chunk: bool = False, pd_config: dict[str, Any] | None = None, @@ -153,12 +127,9 @@ class Orchestrator: self.output_async_queue = output_async_queue self.rpc_async_queue = rpc_async_queue - self.num_stages = len(stage_clients) self.async_chunk = bool(async_chunk) - - self.stage_clients: list[Any] = stage_clients - self.output_processors: list[Any] = output_processors - self.stage_vllm_configs: list[Any] = stage_vllm_configs + self.num_stages = len(stage_pools) + self.stage_pools: list[StagePool] = stage_pools # PD disaggregation state self._pd_pair: tuple[int, int] | None = None @@ -169,19 +140,9 @@ class Orchestrator: self._pd_pair = pd_config.get("pd_pair") self._pd_bootstrap_addr = pd_config.get("bootstrap_addr") self._pd_prefill_engine_id = pd_config.get("prefill_engine_id") - - # Per-request state self.request_states: dict[str, OrchestratorRequestState] = {} - - # CFG companion tracking self._cfg_tracker = CfgCompanionTracker() - # Per-stage metrics accumulators. - self._batch_seq: list[int] = [0] * self.num_stages - self._agg_total_tokens: list[int] = [0] * self.num_stages - self._agg_total_gen_time_ms: list[float] = [0.0] * self.num_stages - - # Shutdown coordination self._shutdown_event = asyncio.Event() self._stages_shutdown = False self._fatal_error: str | None = None @@ -198,7 +159,6 @@ class Orchestrator: ) try: - # Run both tasks concurrently; if either fails the other is cancelled. await asyncio.gather(request_task, output_task) except asyncio.CancelledError: raise @@ -207,9 +167,9 @@ class Orchestrator: raise finally: self._shutdown_event.set() - for t in (request_task, output_task): - if not t.done(): - t.cancel() + for task in (request_task, output_task): + if not task.done(): + task.cancel() try: await asyncio.gather(request_task, output_task, return_exceptions=True) except Exception: @@ -223,15 +183,15 @@ class Orchestrator: self._shutdown_stages() - # Cancel any remaining tasks spawned by wait_for / gather so - # the event loop can close cleanly without "pending task" errors. loop = asyncio.get_running_loop() pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task() and not t.done()] - for t in pending: - t.cancel() + for task in pending: + task.cancel() if pending: await asyncio.gather(*pending, return_exceptions=True) + # ---- Request handling ---- + async def _request_handler(self) -> None: """Read messages from the main thread via request_async_queue.""" while True: @@ -254,7 +214,195 @@ class Orchestrator: self._shutdown_stages() break else: - logger.warning(f"[Orchestrator] Unknown message type: {msg_type}") + logger.warning("[Orchestrator] Unknown message type: %s", msg_type) + + async def _handle_add_request(self, msg: dict[str, Any]) -> None: + """Handle an add_request message from the main thread.""" + stage_id = 0 + request_id = msg["request_id"] + prompt = msg["prompt"] + original_prompt = msg.get("original_prompt", prompt) + sampling_params_list = msg["sampling_params_list"] + if not sampling_params_list: + raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.") + final_stage_id = msg["final_stage_id"] + + logger.info( + "[Orchestrator] _handle_add_request: stage=%s req=%s " + "prompt_type=%s original_prompt_type=%s final_stage=%s " + "num_sampling_params=%d", + stage_id, + request_id, + type(prompt).__name__, + type(original_prompt).__name__, + final_stage_id, + len(sampling_params_list), + ) + + req_state = OrchestratorRequestState( + request_id=request_id, + prompt=original_prompt, + sampling_params_list=sampling_params_list, + final_stage_id=final_stage_id, + mm_features=getattr(prompt, "mm_features", None), + ) + self.request_states[request_id] = req_state + req_state.streaming.enabled = bool(getattr(prompt, "resumable", False)) + req_state.stage_submit_ts[stage_id] = _time.time() + enqueue_ts = msg.get("enqueue_ts", 0.0) + if enqueue_ts > 0: + req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - enqueue_ts) * 1000.0 + preprocess_ms = msg.get("preprocess_ms", 0.0) + if preprocess_ms > 0: + req_state.pipeline_timings["preprocess_ms"] = preprocess_ms + await self.stage_pools[stage_id].submit_initial( + request_id, + req_state, + prompt, + prompt_text=msg.get("output_prompt_text"), + ) + + if self.async_chunk and stage_id == 0 and final_stage_id > 0: + await self._prewarm_async_chunk_stages(request_id, prompt, req_state) + + async def _handle_streaming_update(self, msg: dict[str, Any]) -> None: + """Handle a streaming_update message for an existing request.""" + stage_id = 0 + request_id = msg["request_id"] + request = msg["prompt"] + + req_state = self.request_states.get(request_id) + if req_state is None: + logger.warning( + "[Orchestrator] streaming_update for unknown req=%s, falling back to add_request", + request_id, + ) + fallback_msg = dict(msg) + fallback_msg["type"] = "add_request" + await self._handle_add_request(fallback_msg) + return + + if "sampling_params_list" in msg and msg["sampling_params_list"]: + req_state.sampling_params_list = msg["sampling_params_list"] + + req_state.streaming.enabled = True + req_state.stage_submit_ts[stage_id] = _time.time() + await self.stage_pools[stage_id].submit_update( + request_id, + req_state, + request, + prompt_text=msg.get("output_prompt_text"), + ) + + async def _handle_add_companion(self, msg: dict[str, Any]) -> None: + """Handle an add_companion_request message: submit companion to stage 0.""" + companion_id = msg["companion_id"] + parent_id = msg["parent_id"] + role = msg["role"] + companion_prompt = msg["prompt"] + sampling_params_list = msg["sampling_params_list"] + + parent_state = self.request_states.get(parent_id) + if parent_state is None: + logger.info( + "[Orchestrator] Dropping CFG companion %s (role=%s): parent %s is no longer active", + companion_id, + role, + parent_id, + ) + return + + self._cfg_tracker.register_companion(parent_id, role, companion_id) + + companion_state = OrchestratorRequestState( + request_id=companion_id, + prompt=companion_prompt, + sampling_params_list=sampling_params_list, + final_stage_id=0, + ) + self.request_states[companion_id] = companion_state + companion_state.stage_submit_ts[0] = _time.time() + companion_replica_id = await self.stage_pools[0].submit_initial( + companion_id, + companion_state, + companion_prompt, + prompt_text=msg.get("companion_prompt_text"), + affinity_request_id=parent_id, + ) + + logger.info( + "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)", + companion_id, + role, + parent_id, + companion_replica_id, + ) + + async def _handle_abort(self, msg: dict[str, Any]) -> None: + """Handle an abort message from the main thread.""" + request_ids = msg["request_ids"] + await self._cleanup_request_ids( + self._cfg_tracker.abort_parents(request_ids), + abort=True, + ) + logger.info("[Orchestrator] Aborted request(s) %s", request_ids) + + async def _abort_request_ids(self, request_ids: list[str]) -> None: + """Forward abort requests to all stage pools.""" + if not request_ids: + return + for pool in self.stage_pools: + await pool.abort_requests(request_ids) + + def _release_request_bindings(self, request_ids: list[str]) -> None: + """Release all stage-local route bindings for the given request ids.""" + for pool in self.stage_pools: + pool.release_bindings(request_ids) + + async def _handle_collective_rpc(self, msg: dict[str, Any]) -> None: + """Handle a control-plane RPC request from the main thread.""" + rpc_id = msg["rpc_id"] + method = msg["method"] + timeout = msg.get("timeout") + args = tuple(msg.get("args", ())) + kwargs = dict(msg.get("kwargs") or {}) + requested_stage_ids = msg.get("stage_ids") + + target_pools: list[StagePool] = [] + if requested_stage_ids is None: + target_pools.extend(self.stage_pools) + else: + for lid in requested_stage_ids: + if not (0 <= lid < self.num_stages): + logger.warning("[Orchestrator] collective_rpc: ignoring invalid stage_id %s", lid) + continue + target_pools.append(self.stage_pools[lid]) + + results: list[Any] = [] + stage_ids: list[int] = [] + for pool in target_pools: + for replica_id in range(pool.num_replicas): + stage_result = await pool.collective_rpc( + replica_id=replica_id, + method=method, + timeout=timeout, + args=args, + kwargs=kwargs, + ) + stage_ids.append(pool.stage_id) + results.append(stage_result) + + await self.rpc_async_queue.put( + { + "type": "collective_rpc_result", + "rpc_id": rpc_id, + "method": method, + "stage_ids": stage_ids, + "results": results, + } + ) + + # ---- Orchestration loop ---- async def _orchestration_output_handler(self) -> None: """Poll all stages, handle transfers, send final outputs to main.""" @@ -265,137 +413,170 @@ class Orchestrator: return async def _orchestration_loop(self) -> None: - """Inner loop for _orchestration_output_handler (clean cancellation). - - Control flow: poll raw → process through output processor → route. - """ + """Poll stage pools and route logical outputs.""" while not self._shutdown_event.is_set(): idle = True for stage_id in range(self.num_stages): - if self._shutdown_event.is_set(): - return - - # 1) Diffusion stage: poll non-blocking queue - # TODO (Peiqi): the output of diffusion stage is OmniRequestOutput, - # which is different from EngineCoreOutputs (LLM stages). We may want to unify - # the output format in the future to simplify the processing logic in Orchestrator. - stage_client = self.stage_clients[stage_id] - if stage_client.stage_type == "diffusion": - output = stage_client.get_diffusion_output_nowait() - if output is not None: - idle = False - - if getattr(output, "error", None) is not None: - await self.output_async_queue.put( - { - "type": "output", - "request_id": output.request_id, - "stage_id": stage_id, - "engine_outputs": output, - "metrics": None, - "finished": True, - } - ) - self.request_states.pop(output.request_id, None) - continue - - req_state = self.request_states.get(output.request_id) - if req_state is not None: - if getattr(output, "error", None) is not None: - parent_id = self._companion_to_parent.get(output.request_id, output.request_id) - await self.output_async_queue.put( - { - "type": "error", - "request_id": parent_id, - "stage_id": stage_id, - "error": output.error, - } - ) - role_map = self._companion_map.get(parent_id, {}) - for cid in role_map.values(): - self.request_states.pop(cid, None) - self._cleanup_companion_state(parent_id) - self.request_states.pop(parent_id, None) - continue - - stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state) - await self._route_output(stage_id, output, req_state, stage_metrics) - continue - - # 1) Poll raw outputs from the stage - try: - raw_outputs = await asyncio.wait_for(self._poll_stage_raw(stage_id), timeout=0.001) - except asyncio.TimeoutError: - continue - except asyncio.CancelledError: - raise - except EngineDeadError as e: - logger.error( - "[Orchestrator] Stage-%s is dead: %s", - stage_id, - e, - ) - self._fatal_error = str(e) - self._fatal_error_stage_id = stage_id - for req_id, req_state in list(self.request_states.items()): - if stage_id in req_state.stage_submit_ts: - await self.output_async_queue.put( - { - "type": "error", - "error": str(e), - "fatal": True, - "request_id": req_id, - "stage_id": stage_id, - } - ) - self.request_states.pop(req_id, None) - self._shutdown_event.set() - raise - except Exception: + pool = self.stage_pools[stage_id] + for replica_id in range(pool.num_replicas): if self._shutdown_event.is_set(): return - logger.exception( - "[Orchestrator] _poll_stage_raw failed for stage-%s", - stage_id, - ) - raise - if raw_outputs is None: - continue - idle = False + if pool.stage_type == "diffusion": + output = pool.poll_diffusion_output(replica_id) + if output is None: + continue - # Handle prefill-finished KV-ready signals before finished outputs. - await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs) + await self._handle_processed_outputs(stage_id, replica_id, [output]) + idle = False + else: + try: + raw_outputs = await pool.poll_llm_raw_output(replica_id, timeout_s=0.001) + if raw_outputs is None: + continue - # 2) Process raw outputs through the output processor - request_outputs = await self._process_stage_outputs(stage_id, raw_outputs) + await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs) + for eco in raw_outputs.outputs: + req_state = self.request_states.get(getattr(eco, "request_id", None)) + if req_state is None: + continue + req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False)) + req_state.streaming.new_prompt_len_snapshot = getattr( + eco, + "new_prompt_len_snapshot", + None, + ) + raw_output = await pool.process_llm_raw_outputs(replica_id, raw_outputs) + except asyncio.CancelledError: + raise + except EngineDeadError as e: + logger.error( + "[Orchestrator] Stage-%s is dead: %s", + stage_id, + e, + ) + # TODO: Fault handling is intentionally fail-stop at + # the orchestrator level today. If one replica in a + # logical stage dies, we promote it to `_fatal_error`, + # notify requests already admitted to that stage, and + # re-raise so `run()` shuts down all stages. This is + # conservative but means a single unhealthy replica in + # a multi-replica deployment can take down otherwise + # healthy replicas in other stages. Revisit this when + # adding per-replica fault isolation / eviction. + self._fatal_error = str(e) + self._fatal_error_stage_id = stage_id + for req_id, req_state in list(self.request_states.items()): + if stage_id in req_state.stage_submit_ts: + await self.output_async_queue.put( + { + "type": "error", + "error": str(e), + "fatal": True, + "request_id": req_id, + "stage_id": stage_id, + } + ) + self.request_states.pop(req_id, None) + self._shutdown_event.set() + raise + except Exception: + if self._shutdown_event.is_set(): + return + logger.exception( + "[Orchestrator] Stage-%s replica-%s processing failed", + stage_id, + replica_id, + ) + raise - # 3) Route each processed output - for output in request_outputs: - req_state = self.request_states.get(output.request_id) - if req_state is None: - logger.warning( - "[Orchestrator] Dropping output for unknown req %s at stage-%s (known reqs: %s)", - output.request_id, - stage_id, - list(self.request_states.keys()), - ) - continue - stage_metrics = None - if output.finished: - stage_metrics = self._build_stage_metrics( - stage_id, - output.request_id, - [output], - req_state, - ) - await self._route_output(stage_id, output, req_state, stage_metrics) + await self._handle_processed_outputs(stage_id, replica_id, raw_output) + idle = False if idle: await asyncio.sleep(0.001) else: await asyncio.sleep(0) + async def _handle_processed_outputs(self, stage_id: int, replica_id: int, outputs: list[Any]) -> None: + """Route processed stage outputs produced by one stage poll.""" + pool = self.stage_pools[stage_id] + for output in outputs: + req_state = self.request_states.get(output.request_id) + if req_state is None: + logger.warning( + "[Orchestrator] Dropping output for unknown req %s at stage-%s (known reqs: %s)", + output.request_id, + stage_id, + list(self.request_states.keys()), + ) + continue + + if getattr(output, "error", None) is not None: + await self._handle_stage_error(stage_id, output) + continue + + stage_metrics = None + if output.finished: + stage_metrics = pool.build_stage_metrics( + [output], + submit_ts=req_state.stage_submit_ts.get(stage_id, _time.time()), + replica_id=replica_id, + ) + stage_metrics.pipeline_timings = dict(req_state.pipeline_timings) + + await self._route_output(stage_id, output, req_state, stage_metrics) + + async def _handle_stage_error(self, stage_id: int, output: Any) -> None: + """Emit a frontend-visible error and clean up request state.""" + if self._cfg_tracker.is_companion(output.request_id): + parent_id = self._cfg_tracker.get_parent_id(output.request_id) or output.request_id + else: + parent_id = output.request_id + await self.output_async_queue.put( + { + "type": "error", + "request_id": parent_id, + "stage_id": stage_id, + "error": output.error, + } + ) + await self._cleanup_request_ids( + [parent_id, *self._cfg_tracker.cleanup_parent(parent_id)], + abort=True, + ) + + # ---- Shared helpers ---- + + async def _cleanup_request_ids(self, request_ids: list[str], *, abort: bool = False) -> None: + """Release pool bindings and logical request state for the given ids.""" + if not request_ids: + return + + if abort: + await self._abort_request_ids(request_ids) + self._release_request_bindings(request_ids) + for request_id in request_ids: + self._pd_kv_params.pop(request_id, None) + self.request_states.pop(request_id, None) + + def _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any: + """Attach CFG companion ids to diffusion sampling params when needed.""" + companion_request_ids = self._cfg_tracker.get_companion_request_ids(request_id) + if not companion_request_ids: + return params + + import copy + + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + if not isinstance(params, OmniDiffusionSamplingParams): + return params + + params = copy.deepcopy(params) + params.cfg_kv_request_ids = companion_request_ids + return params + async def _route_output( self, stage_id: int, @@ -403,20 +584,17 @@ class Orchestrator: req_state: OrchestratorRequestState, stage_metrics: Any, ) -> None: - """Route a processed output: send to main thread and/or forward to next stage.""" + """Route a processed output: send to frontend and/or forward.""" req_id = output.request_id finished = output.finished submit_ts = req_state.stage_submit_ts.get(stage_id) - stage_client = self.stage_clients[stage_id] - # CFG companion handling: companions don't produce user-visible output - # and don't forward to the next stage directly. if finished and self._cfg_tracker.is_companion(req_id): await self._handle_cfg_companion_ready(req_id) - self.request_states.pop(req_id, None) + await self._cleanup_request_ids([req_id]) return - if stage_client.final_output: + if self.stage_pools[stage_id].final_output: await self.output_async_queue.put( { "type": "output", @@ -439,22 +617,11 @@ class Orchestrator: } ) - # PD disaggregation: extract KV transfer params from prefill stage output if self._pd_pair is not None and finished and stage_id == self._pd_pair[0]: kv_params = getattr(output, "kv_transfer_params", None) if kv_params is not None: self._pd_kv_params[req_id] = kv_params if isinstance(kv_params, dict) else dict(kv_params) - logger.debug( - "[Orchestrator][PD] stored kv_transfer_params for req=%s (keys=%s)", - req_id, - list(self._pd_kv_params[req_id].keys()), - ) - else: - logger.warning( - "[Orchestrator][PD] prefill stage output for req=%s has no kv_transfer_params; " - "KV transfer may fail. Ensure apply_mooncake_connector_patch() was called.", - req_id, - ) + req_state.pd_prefill_multimodal_output = getattr(output, "multimodal_output", None) if ( (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished)) @@ -489,10 +656,7 @@ class Orchestrator: ) if finished and stage_id == req_state.final_stage_id: - # PD: clean up any lingering KV params for this request - self._pd_kv_params.pop(req_id, None) - self._cfg_tracker.cleanup_parent(req_id) - self.request_states.pop(req_id, None) + await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: return (stage_id + 1) in req_state.stage_submit_ts @@ -502,26 +666,40 @@ class Orchestrator: parent_id = self._cfg_tracker.on_companion_completed(req_id) if parent_id is None: return + deferred = self._cfg_tracker.pop_pending_parent(parent_id) if deferred is None: return - parent_state = self.request_states.get(parent_id) - if parent_state is not None and not self._next_stage_already_submitted(deferred["stage_id"], parent_state): - await self._forward_to_next_stage( - parent_id, - deferred["stage_id"], - deferred["engine_outputs"], - parent_state, - ) - async def _handle_kv_ready_raw_outputs(self, stage_id: int, raw_outputs: EngineCoreOutputs) -> None: - """Forward split requests once stage-0 KV is ready, not only when decode fully finishes.""" + parent_state = self.request_states.get(parent_id) + if parent_state is None: + return + + stage_id = deferred["stage_id"] + if (stage_id + 1) in parent_state.stage_submit_ts: + return + + await self._forward_to_next_stage( + parent_id, + stage_id, + deferred["engine_outputs"], + parent_state, + ) + + async def _handle_kv_ready_raw_outputs( + self, + stage_id: int, + raw_outputs: EngineCoreOutputs, + ) -> None: + """Forward split requests once stage-0 KV is ready.""" if self.async_chunk: return + for raw_output in raw_outputs.outputs: kv_params = getattr(raw_output, "kv_transfer_params", None) if not (isinstance(kv_params, dict) and kv_params.get("kv_ready")): continue + req_id = raw_output.request_id req_state = self.request_states.get(req_id) if req_state is None: @@ -531,8 +709,9 @@ class Orchestrator: continue if stage_id >= req_state.final_stage_id: continue - if self._next_stage_already_submitted(stage_id, req_state): + if (stage_id + 1) in req_state.stage_submit_ts: continue + if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id): self._cfg_tracker.defer_parent(req_id, raw_output, stage_id) else: @@ -583,107 +762,33 @@ class Orchestrator: ) return sp - def _build_stage_metrics( - self, - stage_id: int, - req_id: str, - request_outputs: list[RequestOutput], - req_state: OrchestratorRequestState, - ) -> StageRequestMetrics: - """Build StageRequestMetrics for a finished request at a stage. - - Reuses StageRequestMetrics so OrchestratorMetrics and downstream - metric handlers can consume a stable schema. - """ - now = _time.time() - submit_ts = req_state.stage_submit_ts.get(stage_id, now) - stage_gen_time_ms = (now - submit_ts) * 1000.0 - - num_tokens_out = count_tokens_from_outputs(request_outputs) - num_tokens_in = 0 - if stage_id == 0: - for ro in request_outputs: - ptids = getattr(ro, "prompt_token_ids", None) - if ptids is not None: - num_tokens_in += len(ptids) - - # Monotonic batch counter per stage. - self._batch_seq[stage_id] += 1 - batch_id = self._batch_seq[stage_id] - - # Accumulate for running-average stage_stats - self._agg_total_tokens[stage_id] += num_tokens_out - self._agg_total_gen_time_ms[stage_id] += stage_gen_time_ms - - return StageRequestMetrics( - num_tokens_in=num_tokens_in, - num_tokens_out=num_tokens_out, - stage_gen_time_ms=stage_gen_time_ms, - batch_id=batch_id, - batch_size=1, - rx_decode_time_ms=0.0, - rx_transfer_bytes=0, - rx_in_flight_time_ms=0.0, - stage_stats=StageStats( - total_token=self._agg_total_tokens[stage_id], - total_gen_time_ms=self._agg_total_gen_time_ms[stage_id], - ), - pipeline_timings=dict(req_state.pipeline_timings), - ) - - def _build_kv_sender_info(self, sender_stage_ids: list[int]) -> dict[int, dict[str, Any]] | None: - """Build per-request sender info for diffusion KV-transfer receivers.""" - sender_infos: dict[int, dict[str, Any]] = {} - for sender_stage_id in dict.fromkeys(sender_stage_ids): - if sender_stage_id < 0 or sender_stage_id >= self.num_stages: - continue - - sender_stage = self.stage_clients[sender_stage_id] - get_sender_info = getattr(sender_stage, "get_kv_sender_info", None) - if not callable(get_sender_info): - continue - - sender_info = get_sender_info() - if not sender_info: - logger.warning( - "[Orchestrator] Stage-%s has no KV sender info available", - sender_stage_id, - ) - continue - - sender_infos[sender_stage_id] = sender_info - - return sender_infos or None - async def _forward_to_next_stage( self, req_id: str, - stage_id: int, + src_stage_id: int, output: Any, req_state: OrchestratorRequestState, *, is_streaming_session: bool = False, is_final_update: bool = False, ) -> None: - """Forward output from current stage to the next stage. - - Handles the full pipeline: set outputs on current stage, compute - next-stage inputs, build lightweight requests, and submit them. - """ - next_stage_id = stage_id + 1 - next_client = self.stage_clients[next_stage_id] - params = req_state.sampling_params_list[next_stage_id] + """Forward output from the current logical stage to the next one.""" + next_logical = src_stage_id + 1 + next_pool = self.stage_pools[next_logical] + next_client = next_pool.stage_client + params = req_state.sampling_params_list[next_logical] + source_outputs = [output] next_stage_resumable = is_streaming_session and not is_final_update + already_submitted = self._next_stage_already_submitted(src_stage_id, req_state) + requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False) - if next_client.stage_type == "diffusion": - self.stage_clients[stage_id].set_engine_outputs([output]) + if next_pool.stage_type == "diffusion": if next_client.custom_process_input_func is not None: _t_ar2d = _time.perf_counter() diffusion_prompt = next_client.custom_process_input_func( - self.stage_clients, - next_client.engine_input_source, + source_outputs, req_state.prompt, - False, + requires_multimodal_data, ) _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000 req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d @@ -691,66 +796,61 @@ class Orchestrator: "[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d", req_id, _dt_ar2d, - stage_id, - next_stage_id, + src_stage_id, + next_logical, ) if isinstance(diffusion_prompt, list): if not diffusion_prompt: error_output = OmniRequestOutput.from_error( req_id, - f"Stage-{stage_id} produced no valid inputs for diffusion stage-{next_stage_id}", + f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", ) logger.warning( "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " "routing terminal error output", req_id, - stage_id, - next_stage_id, + src_stage_id, + next_logical, ) await self.output_async_queue.put( { "type": "output", "request_id": req_id, - "stage_id": next_stage_id, + "stage_id": next_logical, "engine_outputs": error_output, "metrics": None, "finished": True, } ) - self._pd_kv_params.pop(req_id, None) - self._cfg_tracker.cleanup_parent(req_id) - self.request_states.pop(req_id, None) + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + ) return - diffusion_prompt = diffusion_prompt[0] + if already_submitted and len(diffusion_prompt) == 1: + diffusion_prompt = diffusion_prompt[0] else: diffusion_prompt = req_state.prompt - params = self._cfg_tracker.attach_cfg_request_ids(req_id, params) - - source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [stage_id]) - kv_sender_info = self._build_kv_sender_info(sender_stage_ids=source_stage_ids) - if isinstance(diffusion_prompt, list): - await next_client.add_batch_request_async( - req_id, - diffusion_prompt, - params, - kv_sender_info=kv_sender_info, - ) + if already_submitted: + await next_pool.submit_update(req_id, req_state, diffusion_prompt) else: - await next_client.add_request_async( + await next_pool.submit_initial( req_id, + req_state, diffusion_prompt, - params, - kv_sender_info=kv_sender_info, + submit_kwargs={ + "kv_sender_info": self._build_kv_sender_info( + list(getattr(next_client, "engine_input_source", None) or [src_stage_id]), + request_id=req_id, + ) + }, + params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), ) - req_state.stage_submit_ts[next_stage_id] = _time.time() + req_state.stage_submit_ts[next_logical] = _time.time() return # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params - if self._pd_pair is not None and (stage_id, next_stage_id) == self._pd_pair: - # Save prefill stage outputs so thinker2talker can merge embeddings later - self.stage_clients[stage_id].set_engine_outputs([output]) - + if self._pd_pair is not None and (src_stage_id, next_logical) == self._pd_pair: params = self._build_pd_decode_params(req_id, params) # Use the original user prompt for the decode stage (not processed embeddings) @@ -775,37 +875,36 @@ class Orchestrator: request_id=req_id, prompt=decode_input, params=params, - model_config=self.stage_vllm_configs[next_stage_id].model_config, - mm_features=req_state.mm_features, # Pass mm_features for M-RoPE + model_config=next_pool.stage_vllm_config.model_config, + mm_features=req_state.mm_features, + resumable=next_stage_resumable, ) request.external_req_id = request.request_id + if already_submitted: + await next_pool.submit_update(req_id, req_state, request) + else: + await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) - self.output_processors[next_stage_id].add_request( - request=request, - prompt=None, - parent_req=None, - request_index=0, - queue=None, - ) - await next_client.add_request_async(request) - - req_state.stage_submit_ts[next_stage_id] = _time.time() + req_state.stage_submit_ts[next_logical] = _time.time() return - self.stage_clients[stage_id].set_engine_outputs([output]) + if req_state.pd_prefill_multimodal_output is not None: + req_state.streaming.bridge_states.setdefault( + "pd_prefill_multimodal_output_by_req", + {}, + )[req_id] = req_state.pd_prefill_multimodal_output - # Process inputs for next stage try: next_inputs = next_client.process_engine_inputs( - stage_list=self.stage_clients, - prompt=req_state.prompt, - streaming_context=(req_state.streaming if req_state.streaming.enabled else None), + source_outputs, + req_state.prompt, + streaming_context=req_state.streaming, ) except Exception: logger.exception( "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s", req_id, - next_stage_id, + next_logical, ) raise @@ -813,174 +912,24 @@ class Orchestrator: for next_input in next_inputs: # Only AR thinker stages consume encoder mm_features; downstream # (talker/code2wav/…) must not see them (avoids encoder-cache misses). - _ms = getattr(next_client, "model_stage", None) - _mm_features = req_state.mm_features if _ms == "thinker" else None + model_stage = getattr(next_client, "model_stage", None) + mm_features = req_state.mm_features if model_stage == "thinker" else None request = build_engine_core_request_from_tokens( request_id=req_id, prompt=next_input, params=params, - model_config=self.stage_vllm_configs[next_stage_id].model_config, - mm_features=_mm_features, + model_config=next_pool.stage_vllm_config.model_config, + mm_features=mm_features, resumable=next_stage_resumable, ) - # TODO: Here we directly use the req id to assign. request.external_req_id = request.request_id - - self.output_processors[next_stage_id].add_request( - request=request, - prompt=None, - parent_req=None, - request_index=0, - queue=None, - ) - - await next_client.add_request_async(request) - - # Record submit timestamp for the next stage - req_state.stage_submit_ts[next_stage_id] = _time.time() - - async def _poll_stage_raw(self, stage_id: int) -> EngineCoreOutputs | None: - """Pull raw EngineCoreOutputs from a stage client without processing. - - Returns the raw outputs object, or None when there is nothing - to consume. - """ - outputs = await self.stage_clients[stage_id].get_output_async() - if not outputs.outputs: - return None - return outputs - - async def _process_stage_outputs(self, stage_id: int, raw_outputs: EngineCoreOutputs) -> list[RequestOutput]: - """Run the output processor on raw outputs, returning RequestOutputs. - - Also handles abort forwarding and scheduler stats updates. - """ - processor = self.output_processors[stage_id] - - processed = processor.process_outputs( - raw_outputs.outputs, - raw_outputs.timestamp, - None, - ) - for eco in raw_outputs.outputs: - if not hasattr(eco, "request_id"): - continue - req_state = self.request_states.get(eco.request_id) - if req_state: - req_state.streaming.segment_finished = eco.is_segment_finished - req_state.streaming.new_prompt_len_snapshot = eco.new_prompt_len_snapshot - - if processed.reqs_to_abort: - await self.stage_clients[stage_id].abort_requests_async(processed.reqs_to_abort) - # Same cleanup as _handle_abort — engine-driven aborts (stop - # token, max tokens, etc.) also skip the normal process_outputs() - # completion path, so accumulated CPU tensors in OmniRequestState - # and OutputProcessor would leak. - self.output_processors[stage_id].abort_requests(processed.reqs_to_abort, internal=True) - for req_id in processed.reqs_to_abort: - self.request_states.pop(req_id, None) - - if raw_outputs.scheduler_stats is not None: - processor.update_scheduler_stats(raw_outputs.scheduler_stats) - - return processed.request_outputs - - async def _handle_add_request(self, msg: dict[str, Any]) -> None: - """Handle an add_request message from the main thread.""" - stage_id = 0 - request_id = msg["request_id"] - prompt = msg["prompt"] - original_prompt = msg.get("original_prompt", prompt) - sampling_params_list = msg["sampling_params_list"] - if not sampling_params_list: - raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.") - params = sampling_params_list[0] - final_stage_id = msg["final_stage_id"] - - logger.info( - "[Orchestrator] _handle_add_request: stage=%s req=%s " - "prompt_type=%s original_prompt_type=%s final_stage=%s " - "num_sampling_params=%d", - stage_id, - request_id, - type(prompt).__name__, - type(original_prompt).__name__, - final_stage_id, - len(sampling_params_list), - ) - - # Track request state - use original_prompt so downstream stages - # (e.g. thinker2talker) can access the raw dict with multi_modal_data. - request = prompt - is_streaming = bool(getattr(request, "resumable", False)) - req_state = OrchestratorRequestState( - request_id=request_id, - prompt=original_prompt, - sampling_params_list=sampling_params_list, - final_stage_id=final_stage_id, - mm_features=getattr(prompt, "mm_features", None), # Save mm_features for PD - ) - req_state.streaming.enabled = is_streaming - req_state.stage_submit_ts[stage_id] = _time.time() - - # Per-request pipeline timings from caller thread - _enqueue_ts = msg.get("enqueue_ts", 0.0) - if _enqueue_ts > 0: - req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - _enqueue_ts) * 1000.0 - _preprocess_ms = msg.get("preprocess_ms", 0.0) - if _preprocess_ms > 0: - req_state.pipeline_timings["preprocess_ms"] = _preprocess_ms - - self.request_states[request_id] = req_state - - # Stage-0 prompt is already a fully-formed OmniEngineCoreRequest - # (pre-processed by AsyncOmniEngine.add_request, output processor - # already registered there) - submit directly. - stage_client = self.stage_clients[stage_id] - if stage_client.stage_type == "diffusion": - if isinstance(prompt, list): - await stage_client.add_batch_request_async( - request_id, - prompt, - params, - ) + if already_submitted: + await next_pool.submit_update(req_id, req_state, request) else: - await stage_client.add_request_async(request_id, prompt, params) - else: - await stage_client.add_request_async(request) + await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) - if self.async_chunk and stage_id == 0 and final_stage_id > 0: - await self._prewarm_async_chunk_stages(request_id, request, req_state) - - async def _handle_streaming_update(self, msg: dict[str, Any]) -> None: - """Handle a streaming_update message for an existing request.""" - stage_id = 0 - request_id = msg["request_id"] - request = msg["prompt"] - - req_state = self.request_states.get(request_id) - if req_state is None: - logger.warning( - "[Orchestrator] streaming_update for unknown req=%s, falling back to add_request", - request_id, - ) - fallback_msg = dict(msg) - fallback_msg["type"] = "add_request" - await self._handle_add_request(fallback_msg) - return - - if "sampling_params_list" in msg and msg["sampling_params_list"]: - req_state.sampling_params_list = msg["sampling_params_list"] - req_state.streaming.enabled = True - - req_state.stage_submit_ts[stage_id] = _time.time() - stage_client = self.stage_clients[stage_id] - if stage_client.stage_type == "diffusion": - params = req_state.sampling_params_list[stage_id] - await stage_client.add_request_async(request_id, request, params) - else: - await stage_client.add_request_async(request) + req_state.stage_submit_ts[next_logical] = _time.time() async def _prewarm_async_chunk_stages( self, @@ -988,12 +937,7 @@ class Orchestrator: stage0_request: Any, req_state: OrchestratorRequestState, ) -> None: - """Pre-submit downstream stages for async-chunk mode. - - In async-chunk mode, stages exchange data through connectors/chunk adapters, - so downstream stages should be armed once at request start instead of waiting - for stage-finished forwarding. - """ + """Pre-submit downstream stages for async-chunk mode.""" if req_state.final_stage_id <= 0: return @@ -1005,163 +949,86 @@ class Orchestrator: ) return - # Pre-arm stage-1+ with placeholder prompt IDs. - try: - next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) - except Exception: - next_prompt_len = max(1, len(prompt_token_ids)) - original_prompt = req_state.prompt - if isinstance(original_prompt, dict): - base_input = copy.deepcopy(original_prompt) - else: - base_input = {} - base_input["prompt_token_ids"] = [0] * next_prompt_len - base_input["multi_modal_data"] = None - base_input["mm_processor_kwargs"] = None - for next_stage_id in range(1, req_state.final_stage_id + 1): - next_client = self.stage_clients[next_stage_id] + next_pool = self.stage_pools[next_stage_id] params = req_state.sampling_params_list[next_stage_id] - if next_client.stage_type == "diffusion": - source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [next_stage_id - 1]) - kv_sender_info = self._build_kv_sender_info(sender_stage_ids=source_stage_ids) - await next_client.add_request_async( - request_id, - req_state.prompt, - params, - kv_sender_info=kv_sender_info, - ) - req_state.stage_submit_ts[next_stage_id] = _time.time() - continue - - request = build_engine_core_request_from_tokens( - request_id=request_id, - prompt=base_input, - params=params, - model_config=self.stage_vllm_configs[next_stage_id].model_config, - ) - request.external_req_id = request.request_id - - self.output_processors[next_stage_id].add_request( - request=request, - prompt=None, - parent_req=None, - request_index=0, - queue=None, - ) - await next_client.add_request_async(request) req_state.stage_submit_ts[next_stage_id] = _time.time() - async def _handle_add_companion(self, msg: dict[str, Any]) -> None: - """Handle an add_companion_request message: submit companion to stage 0.""" - companion_id = msg["companion_id"] - parent_id = msg["parent_id"] - role = msg["role"] - companion_prompt = msg["prompt"] - sampling_params_list = msg["sampling_params_list"] + if next_pool.stage_type == "diffusion": + await next_pool.submit_initial( + request_id, + req_state, + req_state.prompt, + submit_kwargs={ + "kv_sender_info": self._build_kv_sender_info( + list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]), + request_id=request_id, + ) + }, + ) + else: + import copy - self._cfg_tracker.register_companion(parent_id, role, companion_id) + from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length - companion_state = OrchestratorRequestState( - request_id=companion_id, - prompt=companion_prompt, - sampling_params_list=sampling_params_list, - final_stage_id=0, - ) - companion_state.stage_submit_ts[0] = _time.time() - self.request_states[companion_id] = companion_state + try: + next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) + except Exception: + next_prompt_len = max(1, len(prompt_token_ids)) - request = companion_prompt # Already a processed OmniEngineCoreRequest - stage_client = self.stage_clients[0] - await stage_client.add_request_async(request) + original_prompt = req_state.prompt + if isinstance(original_prompt, dict): + base_input = copy.deepcopy(original_prompt) + else: + base_input = {} - logger.info( - "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s)", - companion_id, - role, - parent_id, - ) + base_input["prompt_token_ids"] = [0] * next_prompt_len + base_input["multi_modal_data"] = None + base_input["mm_processor_kwargs"] = None - async def _handle_abort(self, msg: dict[str, Any]) -> None: - """Handle an abort message from the main thread.""" - request_ids = msg["request_ids"] - all_ids_to_abort = self._cfg_tracker.abort_parents(request_ids) - for stage_id in range(self.num_stages): - await self.stage_clients[stage_id].abort_requests_async(all_ids_to_abort) - # Clean up OutputProcessor state (e.g. mm_accumulated tensors) that - # would otherwise leak — normal completion cleans up via - # process_outputs(), but aborted requests never produce an - # EngineCoreOutput, so we must purge them explicitly. - self.output_processors[stage_id].abort_requests(all_ids_to_abort, internal=True) - for req_id in all_ids_to_abort: - self.request_states.pop(req_id, None) - logger.info("[Orchestrator] Aborted request(s) %s", request_ids) + request = build_engine_core_request_from_tokens( + request_id=request_id, + prompt=base_input, + params=params, + model_config=next_pool.stage_vllm_config.model_config, + ) + request.external_req_id = request.request_id + await next_pool.submit_initial(request_id, req_state, request, prompt_text=None) - async def _handle_collective_rpc(self, msg: dict[str, Any]) -> None: - """Handle a control-plane RPC request from the main thread. + def _build_kv_sender_info( + self, + sender_stage_ids: list[int], + *, + request_id: str | None = None, + ) -> dict[int, dict[str, Any]] | None: + """Build per-request sender info for diffusion KV-transfer receivers.""" + sender_infos: dict[int, dict[str, Any]] = {} + for sender_stage_id in dict.fromkeys(sender_stage_ids): + if sender_stage_id < 0 or sender_stage_id >= len(self.stage_pools): + continue - TODO(AsyncOmni): parallelize stage dispatch if control latency becomes - noticeable. The current sequential fanout keeps the first version simple - and deterministic. - """ - rpc_id = msg["rpc_id"] - method = msg["method"] - timeout = msg.get("timeout") - args = tuple(msg.get("args", ())) - kwargs = dict(msg.get("kwargs") or {}) - requested_stage_ids = msg.get("stage_ids") - stage_ids = list(range(self.num_stages)) if requested_stage_ids is None else list(requested_stage_ids) + sender_pool = self.stage_pools[sender_stage_id] + sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None + if sender_stage is None: + sender_stage = sender_pool.stage_client + get_sender_info = getattr(sender_stage, "get_kv_sender_info", None) + if not callable(get_sender_info): + continue - results: list[Any] = [] - for stage_id in stage_ids: - if stage_id < 0 or stage_id >= self.num_stages: - results.append( - { - "supported": False, - "todo": True, - "error": f"Invalid stage id {stage_id}", - } + sender_info = get_sender_info() + if not sender_info: + logger.warning( + "[Orchestrator] Stage-%s has no KV sender info available", + sender_stage_id, ) continue - stage_client = self.stage_clients[stage_id] - try: - if hasattr(stage_client, "collective_rpc_async"): - stage_result = await stage_client.collective_rpc_async( - method=method, - timeout=timeout, - args=args, - kwargs=kwargs, - ) - else: - stage_result = { - "supported": False, - "todo": True, - "reason": (f"{stage_client.__class__.__name__}.collective_rpc_async is not implemented yet"), - } - except Exception as exc: - logger.exception( - "[Orchestrator] collective_rpc failed: stage=%s method=%s", - stage_id, - method, - ) - stage_result = { - "supported": False, - "error": str(exc), - } + sender_infos[sender_stage_id] = sender_info - results.append(stage_result) + return sender_infos or None - await self.rpc_async_queue.put( - { - "type": "collective_rpc_result", - "rpc_id": rpc_id, - "method": method, - "stage_ids": stage_ids, - "results": results, - } - ) + # ---- Shutdown / lifecycle ---- async def _drain_pending_requests_on_fatal(self) -> None: """Drain the request queue and broadcast fatal errors for any @@ -1213,15 +1080,13 @@ class Orchestrator: self.request_states.pop(req_id, None) def _shutdown_stages(self) -> None: - """Shutdown all stage clients.""" + """Shutdown all stage pools.""" if self._stages_shutdown: return self._stages_shutdown = True - logger.info("[Orchestrator] Shutting down all stages") - for stage_id, stage_client in enumerate(self.stage_clients): - try: - stage_client.shutdown() - logger.info(f"[Orchestrator] Stage {stage_id} shut down") - except Exception as e: - logger.warning(f"[Orchestrator] Failed to shutdown stage {stage_id}: {e}") + total = sum(pool.num_replicas for pool in self.stage_pools) + logger.info("[Orchestrator] Shutting down all %d client(s)", total) + for pool in self.stage_pools: + for replica_id in range(pool.num_replicas): + pool.shutdown_replica(replica_id) diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 5672f9121..6689ac05b 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -359,6 +359,24 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): self.parent_requests[parent_req.request_id] = parent_req self.external_req_ids[req_state.external_req_id].append(request_id) + def remove_request(self, request_id: str) -> None: + """Rollback one previously registered request if it was never submitted.""" + req_state = self.request_states.pop(request_id, None) + if req_state is None: + return + + external_req_id = getattr(req_state, "external_req_id", None) + if external_req_id is not None: + request_ids = self.external_req_ids.get(external_req_id) + if request_ids is not None: + self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id] + if not self.external_req_ids[external_req_id]: + self.external_req_ids.pop(external_req_id, None) + + parent_req = getattr(req_state, "parent_req", None) + if parent_req is not None: + self.parent_requests.pop(parent_req.request_id, None) + def process_outputs( self, engine_core_outputs: list[EngineCoreOutput], diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 37b0e538a..04a047fb3 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -6,6 +6,7 @@ Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture. from __future__ import annotations +import inspect import multiprocessing.connection import socket import threading @@ -34,6 +35,27 @@ logger = init_logger(__name__) SHUTDOWN_TIMEOUT_S = 5 +def _default_process_engine_inputs( + source_outputs: list[Any], + prompt: Any, + requires_multimodal_data: bool, +) -> list[OmniTokensPrompt]: + from vllm_omni.inputs.data import OmniTokensPrompt + + if not isinstance(prompt, list): + prompt = [prompt] + + mm_data = {so.request_id: p.get("multi_modal_data") for so, p in zip(source_outputs, prompt)} + + return [ + OmniTokensPrompt( + prompt_token_ids=so.outputs[0].token_ids, + multi_modal_data=(mm_data[so.request_id] if requires_multimodal_data else None), + ) + for so in source_outputs + ] + + class StageEngineCoreClientBase: """Shared stage-aware behavior for async EngineCore clients. @@ -51,6 +73,8 @@ class StageEngineCoreClientBase: ``engine_manager`` / ``coordinator`` pair created elsewhere. """ + replica_id: int = 0 + @staticmethod def make_async_mp_client( vllm_config: Any, @@ -109,8 +133,10 @@ class StageEngineCoreClientBase: manage the process lifecycle on shutdown. """ # -------- Stage metadata (public fields used at runtime) -------- + self.replica_id = 0 if metadata is not None: self.stage_id = metadata.stage_id + self.replica_id = getattr(metadata, "replica_id", 0) self.stage_type = metadata.stage_type self.engine_output_type = metadata.engine_output_type self.is_comprehension = metadata.is_comprehension @@ -132,9 +158,10 @@ class StageEngineCoreClientBase: client_name = self.__class__.__name__ logger.info( - "[%s] Stage-%s initializing EngineCore", + "[%s] stage-%s [rep-%s] initializing EngineCore", client_name, self.stage_id, + self.replica_id, ) try: super().__init__( @@ -151,17 +178,19 @@ class StageEngineCoreClientBase: self.resources.coordinator = coordinator except Exception: logger.exception( - "[%s] Stage-%s EngineCore init failed", + "[%s] stage-%s [rep-%s] EngineCore init failed", client_name, self.stage_id, + self.replica_id, ) try: self.shutdown() except Exception as shutdown_error: logger.warning( - "[%s] Stage-%s cleanup after init failure failed: %s", + "[%s] stage-%s [rep-%s] cleanup after init failure failed: %s", client_name, self.stage_id, + self.replica_id, shutdown_error, ) raise @@ -172,9 +201,10 @@ class StageEngineCoreClientBase: self._start_proc_monitor() logger.info( - "[%s] Stage-%s EngineCore running", + "[%s] stage-%s [rep-%s] EngineCore running", client_name, self.stage_id, + self.replica_id, ) def _start_proc_monitor(self) -> None: @@ -188,6 +218,7 @@ class StageEngineCoreClientBase: proc = self._proc resources_ref = weakref.ref(self.resources) stage_id = self.stage_id + replica_id = self.replica_id def _monitor() -> None: try: @@ -199,8 +230,9 @@ class StageEngineCoreClientBase: return resources.engine_dead = True logger.error( - "[StageEngineCoreClient] Stage-%s subprocess died unexpectedly (exit code %s).", + "[StageEngineCoreClient] stage-%s [rep-%s] subprocess died unexpectedly (exit code %s).", stage_id, + replica_id, proc.exitcode, ) @@ -228,9 +260,10 @@ class StageEngineCoreClientBase: async def add_request_async(self, request: EngineCoreRequest) -> None: """Add request to the stage engine core.""" logger.info( - "[%s] Stage-%s adding request: %s", + "[%s] stage-%s [rep-%s] add request: %s", self.__class__.__name__, self.stage_id, + self.replica_id, request.request_id, ) await super().add_request_async(request) @@ -311,9 +344,10 @@ class StageEngineCoreClientBase: sender_port = int(base_port) + KV_TRANSFER_PORT_OFFSET + int(from_stage) except (TypeError, ValueError): logger.warning( - "[StageEngineCoreClient] Stage-%s could not resolve sender_zmq_port " + "[StageEngineCoreClient] stage-%s [rep-%s] could not resolve sender_zmq_port " "from base_port=%s and from_stage=%s", self.stage_id, + self.replica_id, base_port, from_stage, ) @@ -358,48 +392,33 @@ class StageEngineCoreClientBase: def process_engine_inputs( self, - stage_list: list[Any], - prompt: OmniTokensPrompt | list[OmniTokensPrompt] | None = None, + source_outputs: list[Any], + prompt: Any = None, streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: - """Process inputs from upstream stages.""" - from vllm_omni.inputs.data import OmniTokensPrompt + """Process inputs from upstream stages. + Transition planning is expressed in terms of the upstream outputs + and the original prompt. + """ if self.custom_process_input_func is not None: - # Keep legacy arg call for non-streaming processors. - if bool(getattr(streaming_context, "enabled", False)): + signature = inspect.signature(self.custom_process_input_func) + if len(signature.parameters) >= 4: return self.custom_process_input_func( - stage_list, - self.engine_input_source, + source_outputs, prompt, self.requires_multimodal_data, streaming_context, ) return self.custom_process_input_func( - stage_list, - self.engine_input_source, + source_outputs, prompt, self.requires_multimodal_data, ) if not self.engine_input_source: raise ValueError(f"engine_input_source empty for stage {self.stage_id}") - - source_id = self.engine_input_source[0] - source_outputs = stage_list[source_id].engine_outputs - - if not isinstance(prompt, list): - prompt = [prompt] - - mm_data = {so.request_id: p.get("multi_modal_data") for so, p in zip(source_outputs, prompt)} - - return [ - OmniTokensPrompt( - prompt_token_ids=so.outputs[0].cumulative_token_ids, - multi_modal_data=(mm_data[so.request_id] if self.requires_multimodal_data else None), - ) - for so in source_outputs - ] + return _default_process_engine_inputs(source_outputs, prompt, self.requires_multimodal_data) async def collective_rpc_async( self, diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 568fcba20..05bcdf7d1 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -29,6 +29,8 @@ from vllm.v1.executor import Executor logger = init_logger(__name__) +StageRoute = tuple[int, int] + # Poll period (ms) used by the registration/handshake loop. _POLL_PERIOD_MS = 5_000 # Default timeout (s) for a stage to send READY. @@ -107,28 +109,33 @@ class OmniMasterServer: master_address: str, master_port: int, stage_ids: list[int], + stage_replica_counts: dict[int, int] | None = None, ) -> None: self._address = master_address self._port = master_port - self._allocations: dict[int, StageAllocation] = {} - self._stage_configs: dict[int, Any] = {} - self._stage_coordinator_addresses: dict[int, StageCoordinatorAddresses] = {} - self._stage_config_events: dict[int, threading.Event] = {} + self._stage_routes: dict[StageRoute, StageAllocation] = {} + self._stage_configs: dict[StageRoute, Any] = {} + self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {} + self._stage_config_events: dict[StageRoute, threading.Event] = {} self._thread: threading.Thread | None = None self._stop_event = threading.Event() + stage_replica_counts = dict(stage_replica_counts or {}) for sid in stage_ids: - self._stage_config_events[sid] = threading.Event() - self._stage_coordinator_addresses[sid] = StageCoordinatorAddresses() - hs_port, inp_port, out_port = get_open_ports_list(count=3) - self._allocations[sid] = StageAllocation( - handshake_bind_address=f"tcp://{master_address}:{hs_port}", - handshake_connect_address=f"tcp://{master_address}:{hs_port}", - input_bind_address=f"tcp://{master_address}:{inp_port}", - input_connect_address=f"tcp://{master_address}:{inp_port}", - output_bind_address=f"tcp://{master_address}:{out_port}", - output_connect_address=f"tcp://{master_address}:{out_port}", - ) + replica_count = max(1, int(stage_replica_counts.get(sid, 1))) + for replica_id in range(replica_count): + route = (sid, replica_id) + self._stage_config_events[route] = threading.Event() + self._stage_coordinator_addresses[route] = StageCoordinatorAddresses() + hs_port, inp_port, out_port = get_open_ports_list(count=3) + self._stage_routes[route] = StageAllocation( + handshake_bind_address=f"tcp://{master_address}:{hs_port}", + handshake_connect_address=f"tcp://{master_address}:{hs_port}", + input_bind_address=f"tcp://{master_address}:{inp_port}", + input_connect_address=f"tcp://{master_address}:{inp_port}", + output_bind_address=f"tcp://{master_address}:{out_port}", + output_connect_address=f"tcp://{master_address}:{out_port}", + ) logger.info( "[OmniMasterServer] Pre-allocated addresses for stages %s (master=%s:%d)", @@ -150,71 +157,78 @@ class OmniMasterServer: """Return the registration port exposed to stage launchers.""" return self._port - def get_allocation(self, stage_id: int) -> StageAllocation: + def get_allocation(self, stage_id: int, replica_id: int = 0) -> StageAllocation: """Return the full address allocation for *stage_id*.""" - return self._allocations[stage_id] + return self._stage_routes[(stage_id, replica_id)] def register_stage_config( self, stage_id: int, stage_config: Any, coordinator_addresses: StageCoordinatorAddresses | None = None, + replica_id: int = 0, ) -> None: """Store the latest stage registration payload for *stage_id*.""" - if stage_id not in self._allocations: - raise KeyError(stage_id) - self._stage_configs[stage_id] = stage_config + key = (stage_id, replica_id) + if key not in self._stage_routes: + raise KeyError(key) + self._stage_configs[key] = stage_config if coordinator_addresses is not None: - self._stage_coordinator_addresses[stage_id] = coordinator_addresses - self._stage_config_events[stage_id].set() + self._stage_coordinator_addresses[key] = coordinator_addresses + self._stage_config_events[key].set() - def get_stage_config(self, stage_id: int, timeout_s: float | None = None) -> Any: + def get_stage_config(self, stage_id: int, timeout_s: float | None = None, replica_id: int = 0) -> Any: """Return the stage config for *stage_id*, waiting if necessary.""" - if stage_id not in self._allocations: - raise KeyError(stage_id) + key = (stage_id, replica_id) + if key not in self._stage_routes: + raise KeyError(key) - if stage_id in self._stage_configs: - return self._stage_configs[stage_id] + if key in self._stage_configs: + return self._stage_configs[key] - if not self._stage_config_events[stage_id].wait(timeout=timeout_s): - raise TimeoutError(f"Timed out waiting for stage config for stage {stage_id}.") + if not self._stage_config_events[key].wait(timeout=timeout_s): + raise TimeoutError(f"Timed out waiting for stage config for stage {stage_id} replica {replica_id}.") - return self._stage_configs[stage_id] + return self._stage_configs[key] def get_stage_coordinator_addresses( self, stage_id: int, timeout_s: float | None = None, + replica_id: int = 0, ) -> StageCoordinatorAddresses: """Return the registered coordinator addresses for *stage_id*.""" - if stage_id not in self._allocations: - raise KeyError(stage_id) + key = (stage_id, replica_id) + if key not in self._stage_routes: + raise KeyError(key) - if not self._stage_config_events[stage_id].is_set(): - if not self._stage_config_events[stage_id].wait(timeout=timeout_s): - raise TimeoutError(f"Timed out waiting for stage registration for stage {stage_id}.") + if not self._stage_config_events[key].is_set(): + if not self._stage_config_events[key].wait(timeout=timeout_s): + raise TimeoutError( + f"Timed out waiting for stage registration for stage {stage_id} replica {replica_id}." + ) - return self._stage_coordinator_addresses[stage_id] + return self._stage_coordinator_addresses[key] - def get_client_addresses(self, stage_id: int) -> dict[str, str]: + def get_client_addresses(self, stage_id: int, replica_id: int = 0) -> dict[str, str]: """Return the addresses the client-side sockets should *bind* to.""" - alloc = self._allocations[stage_id] + alloc = self.get_allocation(stage_id, replica_id) return { "input_address": alloc.input_bind_address, "output_address": alloc.output_bind_address, } - def get_zmq_addresses(self, stage_id: int) -> EngineZmqAddresses: + def get_zmq_addresses(self, stage_id: int, replica_id: int = 0) -> EngineZmqAddresses: """Return EngineZmqAddresses using the *bind* (client) side addresses.""" - alloc = self._allocations[stage_id] + alloc = self.get_allocation(stage_id, replica_id) return EngineZmqAddresses( inputs=[alloc.input_bind_address], outputs=[alloc.output_bind_address], ) - def get_engine_zmq_addresses(self, stage_id: int) -> EngineZmqAddresses: + def get_engine_zmq_addresses(self, stage_id: int, replica_id: int = 0) -> EngineZmqAddresses: """Return EngineZmqAddresses using the *connect* (engine) addresses.""" - alloc = self._allocations[stage_id] + alloc = self.get_allocation(stage_id, replica_id) return EngineZmqAddresses( inputs=[alloc.input_connect_address], outputs=[alloc.output_connect_address], @@ -266,7 +280,7 @@ class OmniMasterServer: poller = zmq.Poller() poller.register(reg_socket, zmq.POLLIN) - pending: set[int] = set(self._allocations.keys()) + pending: set[StageRoute] = set(self._stage_routes.keys()) while pending and not self._stop_event.is_set(): events: list[tuple[zmq.Socket, int]] = poller.poll(_POLL_PERIOD_MS) # type: ignore[assignment] @@ -276,15 +290,15 @@ class OmniMasterServer: for sock, _ in events: if sock is reg_socket: - sid = self._handle_registration(reg_socket) - if sid is not None: - pending.discard(sid) + route = self._handle_registration(reg_socket) + if route is not None: + pending.discard(route) # Cleanup reg_socket.close(linger=0) logger.info("[OmniMasterServer] All stages registered; server thread exiting.") - def _handle_registration(self, reg_socket: zmq.Socket) -> int | None: # type: ignore[type-arg] + def _handle_registration(self, reg_socket: zmq.Socket) -> StageRoute | None: # type: ignore[type-arg] """Receive a stage registration and reply with the handshake address. Returns the registered stage_id on success, or None on failure. @@ -305,10 +319,13 @@ class OmniMasterServer: return None stage_id: int | None = msg.get("stage_id") - if stage_id not in self._allocations: + replica_id = int(msg.get("replica_id", 0) or 0) + key = (stage_id, replica_id) + if key not in self._stage_routes: logger.warning( - "[OmniMasterServer] Received registration for unknown stage_id=%s", + "[OmniMasterServer] Received registration for unknown stage_id=%s replica_id=%s", stage_id, + replica_id, ) return None @@ -320,9 +337,10 @@ class OmniMasterServer: coordinator_output=msg.get("coordinator_output"), frontend_stats_publish_address=msg.get("frontend_stats_publish_address"), ), + replica_id=replica_id, ) - alloc = self._allocations[stage_id] + alloc = self._stage_routes[key] response = msgspec.msgpack.encode( { "handshake_address": alloc.handshake_connect_address, @@ -333,11 +351,12 @@ class OmniMasterServer: # ROUTER-DEALER: reply is [identity, payload] (no empty delimiter). reg_socket.send_multipart([identity, response]) logger.info( - "[OmniMasterServer] Stage %d registered; assigned handshake=%s", + "[OmniMasterServer] Stage %d replica %d registered; assigned handshake=%s", stage_id, + replica_id, alloc.handshake_connect_address, ) - return stage_id + return key def register_stage_with_omni_master( @@ -348,6 +367,7 @@ def register_stage_with_omni_master( omni_stage_config: Any = None, coordinator: DPCoordinator | None = None, return_addresses: bool = False, + replica_id: int = 0, ) -> str | tuple[str, str, str]: """Register a stage with the omni master server. @@ -363,6 +383,7 @@ def register_stage_with_omni_master( reg_sock.connect(f"tcp://{omni_master_address}:{omni_master_port}") payload = { "stage_id": omni_stage_id, + "replica_id": replica_id, "stage_config": _serialize_stage_config(omni_stage_config), } if coordinator is not None: @@ -469,9 +490,10 @@ def connect_remote_engine_cores( vllm_config: VllmConfig, omni_master_server: OmniMasterServer, stage_id: int, + replica_id: int = 0, ) -> Iterator[tuple[None, DPCoordinator | None, EngineZmqAddresses, None]]: """Wait for remote engine cores to connect through the omni handshake.""" - addresses = omni_master_server.get_zmq_addresses(stage_id) + addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id) parallel_config = vllm_config.parallel_config # Mirror the engine-count logic from launch_omni_core_engines. remote_engine_count = ( @@ -482,7 +504,10 @@ def connect_remote_engine_cores( start_index = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 coordinator = None - registered_coordinator_addresses = omni_master_server.get_stage_coordinator_addresses(stage_id) + registered_coordinator_addresses = omni_master_server.get_stage_coordinator_addresses( + stage_id, + replica_id=replica_id, + ) addresses.coordinator_input = registered_coordinator_addresses.coordinator_input addresses.coordinator_output = registered_coordinator_addresses.coordinator_output addresses.frontend_stats_publish_address = registered_coordinator_addresses.frontend_stats_publish_address @@ -490,12 +515,13 @@ def connect_remote_engine_cores( engines_to_handshake = [CoreEngine(index=start_index + i, local=False) for i in range(remote_engine_count)] logger.info( - "Waiting for %d remote engine(s) for stage %d", + "Waiting for %d remote engine(s) for stage %d replica %d", remote_engine_count, stage_id, + replica_id, ) - handshake_bind_address = omni_master_server.get_allocation(stage_id).handshake_bind_address + handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: yield None, coordinator, addresses, None @@ -516,9 +542,10 @@ def launch_omni_core_engines( omni_master_server: OmniMasterServer, stage_id: int, stage_config: Any = None, + replica_id: int = 0, ) -> Iterator[tuple[CoreEngineProcManager, DPCoordinator | None, EngineZmqAddresses]]: """Launch local engine cores using the omni registration flow.""" - addresses = omni_master_server.get_zmq_addresses(stage_id) + addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id) parallel_config = vllm_config.parallel_config # Determine the number of local engines and their ranks. local_engine_count = ( @@ -546,17 +573,19 @@ def launch_omni_core_engines( addresses.frontend_stats_publish_address = coordinator.get_stats_publish_address() logger.info( - "[omni] Started DP Coordinator process for stage %d (PID: %d)", + "[omni] Started DP Coordinator process for stage %d replica %d (PID: %d)", stage_id, + replica_id, coordinator.proc.pid, ) else: coordinator = None logger.info( - "Starting %d local engine(s) for stage %d (dp_rank=%d)", + "Starting %d local engine(s) for stage %d replica %d (dp_rank=%d)", local_engine_count, stage_id, + replica_id, dp_rank, ) @@ -568,6 +597,7 @@ def launch_omni_core_engines( omni_stage_id=stage_id, omni_stage_config=stage_config, coordinator=coordinator, + replica_id=replica_id, ) # One CoreEngine entry per local engine so wait_for_engine_startup can @@ -575,7 +605,7 @@ def launch_omni_core_engines( engines_to_handshake = [CoreEngine(index=start_index + i, local=True) for i in range(local_engine_count)] # Bind the pre-allocated handshake socket for this stage. - handshake_bind_address = omni_master_server.get_allocation(stage_id).handshake_bind_address + handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: local_engine_manager = CoreEngineProcManager( diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 277edb5eb..2f007335b 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -19,21 +19,48 @@ from typing import Any, Literal from vllm.logger import init_logger from vllm.sampling_params import SamplingParams +from vllm.tokenizers import cached_tokenizer_from_config from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.input_processor import InputProcessor from vllm.v1.executor import Executor from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs +from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams +from vllm_omni.inputs.preprocess import OmniInputPreprocessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig logger = init_logger(__name__) +@dataclass +class ReplicaInitPlan: + """One concrete replica startup unit within a logical stage.""" + + replica_id: int + num_replicas: int + launch_mode: str + stage_cfg: Any + metadata: Any + stage_connector_spec: dict[str, Any] + omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None] + stage_vllm_config: Any | None = None + executor_class: type | None = None + + +@dataclass +class LogicalStageInitPlan: + """Startup plan for one logical stage.""" + + stage_idx: int + configured_stage_id: int + replicas: list[ReplicaInitPlan] + + def _resolve_model_to_local_path(model: str) -> str: """Resolve an HF Hub model ID to a local cache path.""" if os.path.isdir(model): @@ -85,6 +112,14 @@ def terminate_alive_proc(proc, timeout=5): proc.kill() +def patch_generation_config_if_needed(model_config: Any) -> None: + """Guard InputProcessor init for models whose config lacks model_type.""" + try: + model_config.try_get_generation_config() + except Exception: + model_config.try_get_generation_config = lambda: {} + + def resolve_worker_cls(engine_args: dict[str, Any]) -> None: """Resolve worker_cls from worker_type for non-diffusion stages.""" worker_type = engine_args.get("worker_type", None) @@ -259,20 +294,9 @@ class StageMetadata: runtime_cfg: Any prompt_expand_func: Callable | None = None cfg_kv_collect_func: Callable | None = None - - -@dataclass -class StartedLlmStage: - """Resources for an LLM stage that has completed startup.""" - - stage_id: int - metadata: Any - vllm_config: Any - executor_class: type - addresses: Any - proc: Any = None - engine_manager: Any = None - coordinator: Any = None + # Multi-replica: replica_id distinguishes replicas of the same stage. + # For single-replica stages this defaults to 0. + replica_id: int = 0 def extract_stage_metadata(stage_config: Any) -> StageMetadata: @@ -375,6 +399,119 @@ def prepare_engine_environment() -> None: pass +def split_devices_for_replicas( + devices_str: str | None, + num_replicas: int, + tp_size: int, + stage_id: int, +) -> list[str]: + """Split a devices string into per-replica subsets. + + When ``num_replicas`` is 1, returns ``[devices_str]`` unchanged. + Otherwise, the total number of device IDs must equal + ``num_replicas * tp_size``; each replica gets ``tp_size`` consecutive + device IDs. + + Example:: + + split_devices_for_replicas("1,2,3,4", num_replicas=2, tp_size=2, stage_id=1) + # → ["1,2", "3,4"] + """ + if num_replicas <= 1 or devices_str is None: + return [devices_str] if devices_str is not None else [devices_str] + + device_list = [d.strip() for d in devices_str.split(",") if d.strip()] + required = num_replicas * tp_size + if len(device_list) != required: + raise ValueError( + f"Stage {stage_id}: num_replicas={num_replicas}, " + f"tensor_parallel_size={tp_size} requires " + f"{required} devices, got {len(device_list)}: {devices_str}" + ) + + result: list[str] = [] + for r in range(num_replicas): + chunk = device_list[r * tp_size : (r + 1) * tp_size] + result.append(",".join(chunk)) + return result + + +def get_stage_tp_size(stage_cfg: Any) -> int: + """Extract tensor_parallel_size from a stage config object.""" + engine_args = getattr(stage_cfg, "engine_args", {}) + if hasattr(engine_args, "get"): + return int(engine_args.get("tensor_parallel_size", 1) or 1) + return int(getattr(engine_args, "tensor_parallel_size", 1) or 1) + + +def get_stage_devices_per_replica(stage_cfg: Any) -> int: + """Return the number of devices consumed by one replica of *stage_cfg*.""" + if getattr(stage_cfg, "stage_type", "llm") != "diffusion": + return get_stage_tp_size(stage_cfg) + + parallel_config = _get_attr_or_item(getattr(stage_cfg, "engine_args", {}), "parallel_config") + if parallel_config is None: + return 1 + + world_size = _get_attr_or_item(parallel_config, "world_size") + if world_size is not None: + return max(1, int(world_size)) + + try: + from vllm_omni.diffusion.data import DiffusionParallelConfig + + return max(1, int(DiffusionParallelConfig.from_dict(_to_dict(parallel_config)).world_size)) + except Exception: + return 1 + + +def compute_replica_layout( + stage_configs: Sequence[Any], +) -> tuple[list[int], dict[int, list[str]]]: + """Compute per-stage replica counts and device assignments. + + Returns: + replicas_per_stage: num_replicas per logical stage. + replica_devices_map: stage_idx -> per-replica device strings + (only for stages with num_replicas > 1). + """ + replicas_per_stage: list[int] = [] + for stage_cfg in stage_configs: + runtime_cfg = getattr(stage_cfg, "runtime", {}) + num_replicas = int( + runtime_cfg.get("num_replicas", 1) + if hasattr(runtime_cfg, "get") + else getattr(runtime_cfg, "num_replicas", 1) + ) + replicas_per_stage.append(max(1, num_replicas)) + + replica_devices_map: dict[int, list[str]] = {} + for stage_id, stage_cfg in enumerate(stage_configs): + num_replicas = replicas_per_stage[stage_id] + if num_replicas <= 1: + continue + runtime_cfg = getattr(stage_cfg, "runtime", {}) + devices_str = ( + runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + ) + devices_per_replica = get_stage_devices_per_replica(stage_cfg) + replica_devices_map[stage_id] = split_devices_for_replicas( + devices_str, + num_replicas, + devices_per_replica, + stage_id, + ) + logger.info( + "[stage_init] Stage %s: %d replicas, devices_per_replica=%d, devices split: %s", + stage_id, + num_replicas, + devices_per_replica, + replica_devices_map[stage_id], + ) + + return replicas_per_stage, replica_devices_map + + def setup_stage_devices(stage_id: int, runtime_cfg: Any) -> None: """Device mapping via set_stage_devices for a single stage.""" physical_devices = set_stage_devices( @@ -419,6 +556,13 @@ def build_engine_args_dict( if stage_type != "diffusion": resolve_worker_cls(engine_args_dict) + if engine_args_dict.get("worker_type") == "generation": + # Non-AR generation stages (e.g. code2wav) do not benefit from + # prefix caching and can expose hybrid KV-cache layouts that vLLM's + # prefix-cache coordinator does not handle. + engine_args_dict.setdefault("disable_hybrid_kv_cache_manager", True) + engine_args_dict.setdefault("enable_prefix_caching", False) + # Check whether the stage's default_sampling_params defines extra_args. default_sp = _to_dict(getattr(stage_config, "default_sampling_params", {})) engine_args_dict["has_sampling_extra_args"] = bool(default_sp.get("extra_args")) @@ -475,6 +619,35 @@ def build_vllm_config( return vllm_config, executor_class +def build_llm_stage_output_processor(plan: LogicalStageInitPlan, stage_vllm_config: Any) -> Any | None: + """Build one output processor per logical LLM stage.""" + + metadata = plan.replicas[0].metadata + if stage_vllm_config.model_config.skip_tokenizer_init: + tokenizer = None + else: + tokenizer = cached_tokenizer_from_config( + model_config=stage_vllm_config.model_config, + ) + return MultimodalOutputProcessor( + tokenizer=tokenizer, + log_stats=False, + engine_core_output_type=metadata.engine_output_type, + ) + + +def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor: + """Build the shared stage-0 input processor.""" + + patch_generation_config_if_needed(stage_vllm_config.model_config) + input_processor = InputProcessor(vllm_config=stage_vllm_config) + input_processor.input_preprocessor = OmniInputPreprocessor( + vllm_config=stage_vllm_config, + renderer=input_processor.renderer, + ) + return input_processor + + def acquire_device_locks( stage_id: int, engine_args_dict: dict[str, Any], @@ -678,7 +851,6 @@ def build_diffusion_config( metadata: StageMetadata, ) -> Any: """Build diffusion config for a stage.""" - from vllm_omni.diffusion.data import OmniDiffusionConfig engine_args_dict = build_engine_args_dict(stage_cfg, model) od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict) @@ -726,108 +898,5 @@ def initialize_diffusion_stage( """ from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client - engine_args = _to_dict(stage_cfg.engine_args) - engine_args.pop("stage_id", None) - od_config = OmniDiffusionConfig.from_kwargs( - stage_id=stage_id, - model=model, - **engine_args, - ) - if metadata.cfg_kv_collect_func is not None: - od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func od_config = build_diffusion_config(model, stage_cfg, metadata) return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline) - - -def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int) -> None: - """vLLM CoreEngineProcManager / coordinators use ``shutdown()``, not ``close()``.""" - if resource is None: - return - shutdown = getattr(resource, "shutdown", None) - if callable(shutdown): - try: - shutdown() - except Exception as cleanup_error: - logger.warning( - "[stage_init] Failed to shutdown launched %s for stage %s: %s", - resource_name, - stage_id, - cleanup_error, - ) - return - close = getattr(resource, "close", None) - if callable(close): - try: - close() - except Exception as cleanup_error: - logger.warning( - "[stage_init] Failed to close launched %s for stage %s: %s", - resource_name, - stage_id, - cleanup_error, - ) - - -def close_started_llm_stage(started: StartedLlmStage) -> None: - """Release resources owned by a launched stage that never attached.""" - if started.proc is not None: - try: - terminate_alive_proc(started.proc) - except Exception as cleanup_error: - logger.warning( - "[stage_init] Failed to terminate process for stage %s: %s", - started.stage_id, - cleanup_error, - ) - _shutdown_or_close_resource(started.engine_manager, "engine manager", started.stage_id) - _shutdown_or_close_resource(started.coordinator, "coordinator", started.stage_id) - - -def finalize_initialized_stages( - stage_clients: list[Any | None], - input_processor: InputProcessor | None, -) -> tuple[list[Any], list[Any], list[dict[str, Any]]]: - """Validate successful init and build runtime metadata lists.""" - if any(stage_client is None for stage_client in stage_clients): - raise RuntimeError("Stage initialization completed with missing stage clients") - - initialized_stage_clients = [stage_client for stage_client in stage_clients if stage_client is not None] - default_sampling_params_list = [stage_client.default_sampling_params for stage_client in initialized_stage_clients] - stage_metadata = [ - { - "final_output": stage_client.final_output, - "final_output_type": stage_client.final_output_type, - "stage_type": stage_client.stage_type, - } - for stage_client in initialized_stage_clients - ] - - if not isinstance(input_processor, InputProcessor): - has_llm_stage = any(metadata.get("stage_type") != "diffusion" for metadata in stage_metadata) - if has_llm_stage: - raise RuntimeError("Failed to initialize stage-0 InputProcessor for LLM pipeline") - - return initialized_stage_clients, default_sampling_params_list, stage_metadata - - -def cleanup_failed_stage_initialization( - stage_clients: list[Any | None], - started_llm_stages: list[StartedLlmStage], -) -> None: - """Shutdown attached stages and close any launched-but-unattached engines.""" - for cleanup_stage_id, stage_client in reversed(list(enumerate(stage_clients))): - if stage_client is None: - continue - try: - stage_client.shutdown() - except Exception as cleanup_error: - logger.warning( - "[stage_init] Failed to shutdown initialized stage %s after init failure: %s", - cleanup_stage_id, - cleanup_error, - ) - - for started in reversed(started_llm_stages): - if stage_clients[started.stage_id] is not None: - continue - close_started_llm_stage(started) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py new file mode 100644 index 000000000..6f7454271 --- /dev/null +++ b/vllm_omni/engine/stage_pool.py @@ -0,0 +1,403 @@ +"""Unified stage-local runtime abstraction for vLLM-Omni.""" + +from __future__ import annotations + +import asyncio +import time as _time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from vllm.logger import init_logger +from vllm.v1.engine import EngineCoreOutputs + +from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics +from vllm_omni.metrics.stats import StageStats +from vllm_omni.metrics.utils import count_tokens_from_outputs + +if TYPE_CHECKING: + from vllm_omni.engine.orchestrator import OrchestratorRequestState + +logger = init_logger(__name__) + + +@dataclass +class _ReplicaMetrics: + """Per-replica metrics accumulators owned by a stage pool.""" + + batch_seq: int = 0 + agg_total_tokens: int = 0 + agg_total_gen_time_ms: float = 0.0 + + +class StagePool: + """Replicas of one logical stage with RR + affinity selection.""" + + def __init__( + self, + stage_id: int, + clients: Any | list[Any], + *, + output_processor: Any = None, + stage_vllm_config: Any = None, + ) -> None: + if isinstance(clients, list): + normalized_clients = list(clients) + else: + normalized_clients = [clients] + + if not normalized_clients: + raise ValueError(f"StagePool for stage {stage_id} has no replicas") + self.stage_id = stage_id + self.clients: list[Any] = normalized_clients + self._output_processor = output_processor + self._stage_vllm_config = stage_vllm_config + self._next_replica_id = 0 + self._request_bindings: dict[str, int] = {} + self._replica_metrics: list[_ReplicaMetrics] = [_ReplicaMetrics() for _ in self.clients] + + # ---- Stage-level properties ---- + + @property + def num_replicas(self) -> int: + return len(self.clients) + + @property + def stage_type(self) -> str | None: + return getattr(self.stage_client, "stage_type", None) + + @property + def final_output(self) -> bool: + return bool(getattr(self.clients[0], "final_output", False)) + + @property + def stage_client(self) -> Any: + return self.clients[0] + + @property + def stage_vllm_config(self) -> Any: + return self._stage_vllm_config + + @property + def output_processor(self) -> Any: + return self._output_processor + + # ---- Route binding lifecycle ---- + + def get_bound_replica_id(self, request_id: str) -> int | None: + """Return the currently bound replica id for *request_id* if present.""" + return self._request_bindings.get(request_id) + + def get_bound_client(self, request_id: str) -> Any | None: + """Return the currently bound client for *request_id* if present.""" + replica_id = self.get_bound_replica_id(request_id) + if replica_id is None: + return None + return self.clients[replica_id] + + def release_binding(self, request_id: str) -> None: + """Drop the route binding for *request_id* in this stage.""" + self._request_bindings.pop(request_id, None) + + def release_bindings(self, request_ids: list[str]) -> None: + """Drop route bindings for the given request ids in this stage.""" + for request_id in request_ids: + self.release_binding(request_id) + + def select_replica_id( + self, + request_id: str, + *, + affinity_request_id: str | None = None, + ) -> int: + """Pick a replica id for *request_id* and cache the choice.""" + cached = self.get_bound_replica_id(request_id) + if cached is not None: + return cached + + chosen = self.get_bound_replica_id(affinity_request_id) if affinity_request_id is not None else None + if chosen is None: + if self.num_replicas == 1: + chosen = 0 + else: + chosen = self._next_replica_id + self._next_replica_id = (self._next_replica_id + 1) % self.num_replicas + + self._request_bindings[request_id] = chosen + return chosen + + # ---- Metrics ---- + + def build_stage_metrics( + self, + request_outputs: list[Any], + *, + submit_ts: float, + replica_id: int, + ) -> StageRequestMetrics: + """Build stage metrics for outputs produced on one replica.""" + now = _time.time() + stage_gen_time_ms = (now - submit_ts) * 1000.0 + + num_tokens_out = count_tokens_from_outputs(request_outputs) + num_tokens_in = 0 + if self.stage_id == 0: + for ro in request_outputs: + ptids = getattr(ro, "prompt_token_ids", None) + if ptids is not None: + num_tokens_in += len(ptids) + + metrics = self._replica_metrics[replica_id] + metrics.batch_seq += 1 + batch_id = metrics.batch_seq + metrics.agg_total_tokens += num_tokens_out + metrics.agg_total_gen_time_ms += stage_gen_time_ms + + return StageRequestMetrics( + num_tokens_in=num_tokens_in, + num_tokens_out=num_tokens_out, + stage_gen_time_ms=stage_gen_time_ms, + batch_id=batch_id, + batch_size=1, + rx_decode_time_ms=0.0, + rx_transfer_bytes=0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats( + total_token=metrics.agg_total_tokens, + total_gen_time_ms=metrics.agg_total_gen_time_ms, + ), + ) + + # ---- Stage-local admission ---- + + async def submit_initial( + self, + request_id: str, + req_state: OrchestratorRequestState, + request: Any, + *, + prompt_text: Any = None, + affinity_request_id: str | None = None, + submit_kwargs: dict[str, Any] | None = None, + params_override: Any = None, + ) -> int: + """Submit a stage-entry request into this pool.""" + params = params_override if params_override is not None else req_state.sampling_params_list[self.stage_id] + submit_kwargs = dict(submit_kwargs or {}) + if self.stage_type == "diffusion": + replica_id = self.select_replica_id( + request_id, + affinity_request_id=affinity_request_id, + ) + client = self.clients[replica_id] + if isinstance(request, list): + await client.add_batch_request_async(request_id, request, params, **submit_kwargs) + else: + await client.add_request_async(request_id, request, params, **submit_kwargs) + return replica_id + + replica_id = self.select_replica_id( + request_id, + affinity_request_id=affinity_request_id, + ) + try: + self.output_processor.add_request( + request=request, + prompt=prompt_text, + parent_req=None, + request_index=0, + queue=None, + ) + except Exception: + self.release_binding(request_id) + raise + + try: + await self.clients[replica_id].add_request_async(request, **submit_kwargs) + except Exception: + self.release_binding(request_id) + rollback = getattr(self.output_processor, "remove_request", None) + if callable(rollback): + try: + rollback(request_id) + except Exception as rollback_error: + logger.warning( + "[StagePool] Failed to rollback output processor state for req=%s stage-%s: %s", + request_id, + self.stage_id, + rollback_error, + ) + raise + return replica_id + + async def submit_update( + self, + request_id: str, + req_state: OrchestratorRequestState, + request: Any, + *, + prompt_text: Any = None, + ) -> int: + """Submit a streaming update to an already admitted request.""" + params = req_state.sampling_params_list[self.stage_id] + replica_id = self.get_bound_replica_id(request_id) + if replica_id is None: + replica_id = self.select_replica_id(request_id) + + if self.stage_type == "diffusion": + await self.clients[replica_id].add_request_async(request_id, request, params) + else: + # Refresh the shared output-processor state before yielding to the + # stage client so streaming segments are merged against the latest + # prompt/token metadata. + self.output_processor.add_request( + request=request, + prompt=prompt_text, + parent_req=None, + request_index=0, + queue=None, + ) + await self.clients[replica_id].add_request_async(request) + return replica_id + + # ---- Stage-local polling ---- + + async def _poll_stage_raw(self, client: Any) -> EngineCoreOutputs | None: + """Pull raw EngineCoreOutputs from a stage replica without processing.""" + outputs = await client.get_output_async() + if not outputs.outputs: + return None + return outputs + + async def process_llm_raw_outputs( + self, + replica_id: int, + raw_outputs: EngineCoreOutputs, + ) -> list[Any]: + """Run the shared LLM output processor on one raw poll result.""" + client = self.clients[replica_id] + processor = self.output_processor + processed = processor.process_outputs( + raw_outputs.outputs, + raw_outputs.timestamp, + None, + ) + + if processed.reqs_to_abort: + await client.abort_requests_async(processed.reqs_to_abort) + + if raw_outputs.scheduler_stats is not None: + processor.update_scheduler_stats(raw_outputs.scheduler_stats) + + return processed.request_outputs + + async def poll_llm_raw_output( + self, + replica_id: int, + *, + timeout_s: float = 0.001, + ) -> EngineCoreOutputs | None: + """Poll raw EngineCore outputs from one LLM replica once.""" + client = self.clients[replica_id] + try: + return await asyncio.wait_for( + self._poll_stage_raw(client), + timeout=timeout_s, + ) + except asyncio.TimeoutError: + return None + except asyncio.CancelledError: + raise + except Exception: + logger.exception( + "[StagePool] _poll_stage_raw failed for stage-%s replica-%s", + self.stage_id, + replica_id, + ) + raise + + def poll_diffusion_output(self, replica_id: int) -> Any | None: + """Drain one ready diffusion output from the given replica if present.""" + return self.clients[replica_id].get_diffusion_output_nowait() + + # ---- Stage-local control plane ---- + + async def abort_requests(self, request_ids: list[str]) -> None: + """Abort the given requests in this stage pool. + + Request-bound abort routing stays inside the pool because route affinity + (`request_id -> replica_id`) is pool-owned. + """ + if not request_ids: + return + + request_ids_by_replica: dict[int, list[str]] = {} + for request_id in request_ids: + replica_id = self.get_bound_replica_id(request_id) + if replica_id is None: + logger.debug("[StagePool] abort: no binding for req=%s in stage-%s", request_id, self.stage_id) + continue + request_ids_by_replica.setdefault(replica_id, []).append(request_id) + + for replica_id, replica_request_ids in request_ids_by_replica.items(): + await self.clients[replica_id].abort_requests_async(replica_request_ids) + + # Clean up OutputProcessor state (e.g. mm_accumulated tensors) that + # would otherwise leak — aborted requests never produce a final + # EngineCoreOutput, so process_outputs() never fires its cleanup path. + all_aborted = [rid for ids in request_ids_by_replica.values() for rid in ids] + if all_aborted and self._output_processor is not None: + self._output_processor.abort_requests(all_aborted, internal=True) + + async def collective_rpc( + self, + replica_id: int, + method: str, + timeout: float | None = None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Any: + """Dispatch a stage-scoped control-plane RPC to one physical route.""" + kwargs = dict(kwargs or {}) + client = self.clients[replica_id] + try: + if hasattr(client, "collective_rpc_async"): + return await client.collective_rpc_async( + method=method, + timeout=timeout, + args=args, + kwargs=kwargs, + ) + return { + "supported": False, + "todo": True, + "reason": f"{client.__class__.__name__}.collective_rpc_async is not implemented yet", + } + except Exception as exc: + logger.exception( + "[StagePool] collective_rpc failed: stage=%s replica=%s method=%s", + self.stage_id, + replica_id, + method, + ) + return { + "supported": False, + "error": str(exc), + } + + def shutdown_replica(self, replica_id: int) -> None: + """Shutdown one backend handle in this stage pool.""" + client = self.clients[replica_id] + try: + client.shutdown() + logger.info( + "[StagePool] Stage %d replica %d shut down", + self.stage_id, + replica_id, + ) + except Exception as e: + logger.warning( + "[StagePool] Failed to shutdown stage %d replica %d: %s", + self.stage_id, + replica_id, + e, + ) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 8e74901dd..75b21139e 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -187,6 +187,12 @@ class OmniServeCommand(CLISubcommand): default=None, help="Select and launch a single stage by stage_id.", ) + omni_config_group.add_argument( + "--replica-id", + type=int, + default=0, + help="Replica id to register when launching a single headless stage.", + ) omni_config_group.add_argument( "--stage-init-timeout", type=int, @@ -529,14 +535,18 @@ def run_headless(args: argparse.Namespace) -> None: model = args.model stage_id: int | None = args.stage_id + replica_id: int = args.replica_id omni_master_address: str | None = args.omni_master_address omni_master_port: int | None = args.omni_master_port if stage_id is None: raise ValueError("--stage-id is required in headless mode") + if replica_id < 0: + raise ValueError("--replica-id must be >= 0 in headless mode") if omni_master_address is None or omni_master_port is None: raise ValueError("--omni-master-address and --omni-master-port are required in headless mode") - if getattr(args, "api_server_count", 0) and args.api_server_count > 1: + api_server_count = args.api_server_count or 0 + if api_server_count > 1: raise ValueError("api_server_count can't be set in headless mode") if args.worker_backend != "multi_process": raise ValueError("headless mode requires worker_backend=multi_process") @@ -552,20 +562,19 @@ def run_headless(args: argparse.Namespace) -> None: # Locate the stage config that matches stage_id. stage_cfg = None for cfg in stage_configs: - if getattr(cfg, "stage_id", None) == stage_id: + if cfg.stage_id == stage_id: stage_cfg = cfg break if stage_cfg is None: raise ValueError( - f"No stage config found for stage_id={stage_id}. " - f"Available stage ids: {[getattr(c, 'stage_id', None) for c in stage_configs]}" + f"No stage config found for stage_id={stage_id}. Available stage ids: {[c.stage_id for c in stage_configs]}" ) prepare_engine_environment() omni_transfer_config = load_omni_transfer_config_for_model(model, config_path) omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) - if getattr(stage_cfg, "stage_type", "llm") == "diffusion": + if stage_cfg.stage_type == "diffusion": metadata = extract_stage_metadata(stage_cfg) if omni_conn_cfg: inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) @@ -573,8 +582,9 @@ def run_headless(args: argparse.Namespace) -> None: od_config = build_diffusion_config(model, stage_cfg, metadata) logger.info( - "[Headless] Launching diffusion stage %d via OmniMasterServer at %s:%d", + "[Headless] Launching diffusion stage %d replica %d via OmniMasterServer at %s:%d", stage_id, + replica_id, omni_master_address, omni_master_port, ) @@ -587,6 +597,7 @@ def run_headless(args: argparse.Namespace) -> None: omni_stage_id=stage_id, omni_stage_config=stage_cfg, return_addresses=True, + replica_id=replica_id, ) proc, _, _, _ = spawn_diffusion_proc( model, @@ -595,13 +606,13 @@ def run_headless(args: argparse.Namespace) -> None: request_address=request_address, response_address=response_address, ) - complete_diffusion_handshake(proc, handshake_address) + complete_diffusion_handshake(proc, handshake_address, args.stage_init_timeout) proc.join() if proc.exitcode not in (None, 0): - raise RuntimeError(f"Diffusion stage {stage_id} exited with code {proc.exitcode}") + raise RuntimeError(f"Diffusion stage {stage_id} replica {replica_id} exited with code {proc.exitcode}") return finally: - logger.info("[Headless] Shutting down stage %d.", stage_id) + logger.info("[Headless] Shutting down stage %d replica %d.", stage_id, replica_id) if proc is not None and proc.is_alive(): terminate_alive_proc(proc) @@ -677,15 +688,17 @@ def run_headless(args: argparse.Namespace) -> None: enable_wave_coordination=vllm_config.model_config.is_moe, ) logger.info( - "[Headless] Started DP Coordinator process for stage %d (PID: %d)", + "[Headless] Started DP Coordinator process for stage %d replica %d (PID: %d)", stage_id, + replica_id, coordinator.proc.pid, ) logger.info( - "[Headless] Launching %d engine core(s) for stage %d via OmniMasterServer at %s:%d", + "[Headless] Launching %d engine core(s) for stage %d replica %d via OmniMasterServer at %s:%d", local_engine_count, stage_id, + replica_id, omni_master_address, omni_master_port, ) @@ -700,11 +713,12 @@ def run_headless(args: argparse.Namespace) -> None: omni_stage_id=stage_id, omni_stage_config=stage_cfg, coordinator=coordinator, + replica_id=replica_id, ) engine_manager = None - log_stats = bool(getattr(args, "log_stats", False)) - if getattr(args, "disable_log_stats", False): + log_stats = bool(args.log_stats) + if args.disable_log_stats: log_stats = False try: diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index dc1e12dfe..377d8d19b 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -63,15 +63,11 @@ def _decode_additional_information(raw_info: Any) -> dict[str, Any]: def text2flow( - stage_list: list[Any], - engine_input_source: list[int], + source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt = None, - requires_multimodal_data: bool = True, + _requires_multimodal_data: bool = True, ): """Build stage-1 inputs by prefixing stage-0 prompt ids to its outputs.""" - source_stage_id = engine_input_source[0] - source_outputs = stage_list[source_stage_id].engine_outputs - engine_inputs: list[OmniTokensPrompt] = [] for source_output in source_outputs: output = source_output.outputs[0] diff --git a/vllm_omni/model_executor/stage_input_processors/dynin_omni.py b/vllm_omni/model_executor/stage_input_processors/dynin_omni.py index 6ce8881a9..87cecc103 100644 --- a/vllm_omni/model_executor/stage_input_processors/dynin_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/dynin_omni.py @@ -87,22 +87,10 @@ def _decode_runtime_bridge_info(value: Any) -> dict[str, Any]: def _bridge_tokens( - stage_list, - engine_input_source, + source_outputs, prompt: OmniTokensPrompt | TextPrompt = None, requires_multimodal_data: bool = False, ): - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - - source_stage_id = engine_input_source[0] - if source_stage_id >= len(stage_list): - raise IndexError(f"Invalid stage_id: {source_stage_id}") - - if stage_list[source_stage_id].engine_outputs is None: - raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") - - source_outputs = stage_list[source_stage_id].engine_outputs next_inputs = [] if not isinstance(prompt, list): prompt = [prompt] @@ -119,9 +107,7 @@ def _bridge_tokens( if not token_ids: token_ids = list(output.cumulative_token_ids or []) if not token_ids: - raise RuntimeError( - f"Stage {source_stage_id} output for request {source_output.request_id} has no token_ids" - ) + raise RuntimeError(f"Stage output for request {source_output.request_id} has no token_ids") detok_id = _to_int(mm_out.get("detok_id"), default=0) src_prompt = prompt_meta_by_reqid.get(source_output.request_id, {}) @@ -147,18 +133,16 @@ def _bridge_tokens( def token2text_to_token2image( - stage_list, - engine_input_source, + source_outputs, prompt: OmniTokensPrompt | TextPrompt = None, requires_multimodal_data: bool = False, ): - return _bridge_tokens(stage_list, engine_input_source, prompt, requires_multimodal_data) + return _bridge_tokens(source_outputs, prompt, requires_multimodal_data) def token2image_to_token2audio( - stage_list, - engine_input_source, + source_outputs, prompt: OmniTokensPrompt | TextPrompt = None, requires_multimodal_data: bool = False, ): - return _bridge_tokens(stage_list, engine_input_source, prompt, requires_multimodal_data) + return _bridge_tokens(source_outputs, prompt, requires_multimodal_data) diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py index 365b303be..542ebefd4 100644 --- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py +++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py @@ -24,16 +24,14 @@ def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None: def slow_ar_to_dac_decoder( - stage_list: list[Any], - engine_input_source: list[int], - prompt: Any = None, - requires_multimodal_data: bool = False, + source_outputs: list[Any], + _prompt: Any = None, + _requires_multimodal_data: bool = False, ) -> list[Any]: """Non-async processor: wait for Slow AR to finish, then pass all codes to DAC decoder.""" from vllm_omni.inputs.data import OmniTokensPrompt - from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _validate_stage_inputs - slow_ar_outputs = _validate_stage_inputs(stage_list, engine_input_source) + slow_ar_outputs = source_outputs dac_inputs: list[OmniTokensPrompt] = [] for output in slow_ar_outputs: diff --git a/vllm_omni/model_executor/stage_input_processors/glm_image.py b/vllm_omni/model_executor/stage_input_processors/glm_image.py index 828f7f218..39533913e 100644 --- a/vllm_omni/model_executor/stage_input_processors/glm_image.py +++ b/vllm_omni/model_executor/stage_input_processors/glm_image.py @@ -212,25 +212,20 @@ def _parse_generated_tokens( def ar2diffusion( - stage_list: list[Any], - engine_input_source: list[int], + source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | list | None = None, requires_multimodal_data: bool = False, + streaming_context: Any | None = None, ) -> list[dict[str, Any]]: - """Process AR stage outputs to create Diffusion stage inputs.""" + """Process AR stage outputs to create Diffusion stage inputs. + + This processor accepts the stage-pool transition interface: + ``ar2diffusion(source_outputs, prompt, requires_multimodal_data)``. + """ + del streaming_context + _t_total = time.perf_counter() - - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - - source_stage_id = engine_input_source[0] - if source_stage_id >= len(stage_list): - raise IndexError(f"Invalid stage_id: {source_stage_id}") - - if stage_list[source_stage_id].engine_outputs is None: - raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") - - ar_outputs = stage_list[source_stage_id].engine_outputs + ar_outputs = source_outputs diffusion_inputs = [] # Normalize prompt to list diff --git a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py b/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py index 0c0e6d7b3..a0273774f 100644 --- a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py +++ b/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py @@ -22,33 +22,20 @@ logger = init_logger(__name__) def ar2diffusion( - stage_list: list[Any], - engine_input_source: list[int], + source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | list | None = None, requires_multimodal_data: bool = False, ) -> list[dict[str, Any]]: """Process AR stage outputs to create Diffusion stage inputs. Args: - stage_list: List of stage clients (set by orchestrator). - engine_input_source: List of source stage IDs (from YAML). prompt: Original user prompt (may contain multimodal data). requires_multimodal_data: Whether to forward multimodal data. Returns: List of dicts, each consumable by the HunyuanImage3 diffusion pipeline. """ - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - - source_stage_id = engine_input_source[0] - if source_stage_id >= len(stage_list): - raise IndexError(f"Invalid source stage_id: {source_stage_id}") - - if stage_list[source_stage_id].engine_outputs is None: - raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") - - ar_outputs = stage_list[source_stage_id].engine_outputs + ar_outputs = source_outputs diffusion_inputs = [] # Normalize prompt to list diff --git a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py index 5eec3d045..6b699008e 100644 --- a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py +++ b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py @@ -9,15 +9,12 @@ from vllm_omni.inputs.data import OmniTokensPrompt def ar2dit( - stage_list: list[Any], - engine_input_source: list[int], + source_outputs: list[Any], prompts: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, # noqa: ARG001 — interface param, unused for ar2dit + _requires_multimodal_data: bool = False, ) -> list[OmniTokensPrompt]: """Convert AR stage outputs to DiT stage inputs.""" - - source_stage_id = engine_input_source[0] - ar_outputs = stage_list[source_stage_id].engine_outputs + ar_outputs = source_outputs dit_inputs: list[OmniTokensPrompt] = [] for ar_output, prompt in zip(ar_outputs, prompts): diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py index 96680b2dd..39071f760 100644 --- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py +++ b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py @@ -182,8 +182,7 @@ def llm2code2wav_async_chunk( def llm2code2wav( - stage_list: list[Any], - engine_input_source: list[int], + source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | None = None, requires_multimodal_data: bool = False, ) -> list[OmniTokensPrompt]: @@ -196,25 +195,13 @@ def llm2code2wav( 3. Package for code2wav stage Args: - stage_list: List of stage objects - engine_input_source: Source stage IDs (typically [1] for talker) prompt: Original prompt data requires_multimodal_data: Whether multimodal data is required Returns: List of OmniTokensPrompt for code2wav stage """ - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - - source_stage_id = engine_input_source[0] - if source_stage_id >= len(stage_list): - raise IndexError(f"Invalid stage_id: {source_stage_id}") - - if stage_list[source_stage_id].engine_outputs is None: - raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") - - talker_outputs = stage_list[source_stage_id].engine_outputs + talker_outputs = source_outputs code2wav_inputs = [] # Process each talker output diff --git a/vllm_omni/model_executor/stage_input_processors/omnivoice.py b/vllm_omni/model_executor/stage_input_processors/omnivoice.py index b866e9779..ce863c135 100644 --- a/vllm_omni/model_executor/stage_input_processors/omnivoice.py +++ b/vllm_omni/model_executor/stage_input_processors/omnivoice.py @@ -10,22 +10,15 @@ from vllm_omni.inputs.data import OmniTokensPrompt def tokens2audio( - stage_list: list[Any], - engine_input_source: list[int], - prompt: OmniTokensPrompt | TextPrompt = None, - requires_multimodal_data: bool = True, + source_outputs: list[Any], + _prompt: OmniTokensPrompt | TextPrompt = None, + _requires_multimodal_data: bool = True, ): """Build stage-1 (decoder) inputs from stage-0 (generator) outputs. Takes the 8-codebook audio tokens from the generator and packages them for the HiggsAudioV2 decoder. """ - source_stage_id = engine_input_source[0] - source_outputs = stage_list[source_stage_id].engine_outputs - - if not isinstance(prompt, list): - prompt = [prompt] - source_output = source_outputs[0] output = source_output.outputs[0] diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 072869eaf..c6f2de53a 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -10,19 +10,11 @@ TALKER_CODEC_END_TOKEN_ID = 8294 def thinker2talker( - stage_list, - engine_input_source, + source_outputs, prompt: OmniTokensPrompt | TextPrompt = None, requires_multimodal_data: bool = False, ): - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - source_stage_id = engine_input_source[0] - if source_stage_id >= len(stage_list): - raise IndexError(f"Invalid stage_id: {source_stage_id}") - if stage_list[source_stage_id].engine_outputs is None: - raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") - thinker_outputs = stage_list[source_stage_id].engine_outputs + thinker_outputs = source_outputs talker_inputs = [] if not isinstance(prompt, list): prompt = [prompt] diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index f7da44067..0eea1c2c0 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -90,36 +90,23 @@ def _ensure_list(x): return list(x) -def _validate_stage_inputs(stage_list, engine_input_source): - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - - stage_id = engine_input_source[0] - if stage_id >= len(stage_list): - raise IndexError(f"Invalid stage_id: {stage_id}") - - stage = stage_list[stage_id] - if stage.engine_outputs is None: - raise RuntimeError(f"Stage {stage_id} has no outputs yet") - - return stage.engine_outputs - - # ========================= # PD disaggregation helpers # ========================= -def _get_prefill_stage(stage_list: list[Any], source_stage_id: int) -> Any | None: - if source_stage_id <= 0: +def _get_prefill_multimodal_output( + request_id: str, + streaming_context: Any | None, +) -> dict[str, Any] | None: + bridge_states = getattr(streaming_context, "bridge_states", None) + if not isinstance(bridge_states, dict): return None - source_stage = stage_list[source_stage_id] - if not getattr(source_stage, "is_decode_only", False): + by_req = bridge_states.get("pd_prefill_multimodal_output_by_req") + if not isinstance(by_req, dict): return None - prev_stage = stage_list[source_stage_id - 1] - if getattr(prev_stage, "is_prefill_only", False) and prev_stage.engine_outputs is not None: - return prev_stage - return None + prefill_mm = by_req.get(request_id) + return prefill_mm if isinstance(prefill_mm, dict) else None def _merge_pd_embeddings( @@ -167,16 +154,6 @@ def _merge_pd_embeddings( return merged_emb, merged_hid -def _get_prefill_multimodal_output(prefill_stage: Any, output_index: int) -> dict[str, Any] | None: - """Return multimodal_output dict from the PD prefill stage for a given batch index.""" - try: - prefill_eos = prefill_stage.engine_outputs - prefill_eo = prefill_eos[min(output_index, len(prefill_eos) - 1)] - return prefill_eo.outputs[0].multimodal_output - except Exception: - return None - - def _resolve_tts_token_embedding( key: str, *, @@ -407,8 +384,7 @@ def thinker2talker_async_chunk( def thinker2talker( - stage_list: list[Any], - engine_input_source: list[int], + source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | None = None, requires_multimodal_data: bool = False, streaming_context: Any | None = None, @@ -425,23 +401,17 @@ def thinker2talker( decode-stage generated embeddings before handing off to the talker. Args: - stage_list: List of stage objects - engine_input_source: Source stage IDs (typically [0] for thinker) prompt: Original prompt data requires_multimodal_data: Whether multimodal data is required Returns: List of OmniTokensPrompt for talker stage """ - thinker_outputs = _validate_stage_inputs(stage_list, engine_input_source) + thinker_outputs = source_outputs talker_inputs: list[OmniTokensPrompt] = [] device = torch.device(current_platform.device_type) - # PD disaggregation: look up the preceding prefill stage (if any) - source_stage_id = engine_input_source[0] - prefill_stage = _get_prefill_stage(stage_list, source_stage_id) - # Process each thinker output for i, thinker_output in enumerate(thinker_outputs): output = thinker_output.outputs[0] @@ -478,8 +448,7 @@ def thinker2talker( thinker_hid = hid_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:] prefill_mm: dict[str, Any] | None = None - if prefill_stage is not None: - prefill_mm = _get_prefill_multimodal_output(prefill_stage, i) + prefill_mm = _get_prefill_multimodal_output(req_id, streaming_context) if prefill_mm is not None: expected_total = len(prompt_token_ids) + len(output_ids) @@ -610,10 +579,9 @@ def talker2code2wav_async_chunk( def talker2code2wav( - stage_list: list[Any], - engine_input_source: list[int], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, + source_outputs: list[Any], + _prompt: OmniTokensPrompt | TextPrompt | None = None, + _requires_multimodal_data: bool = False, streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: """ @@ -625,15 +593,10 @@ def talker2code2wav( 3. Package for code2wav stage Args: - stage_list: List of stage objects - engine_input_source: Source stage IDs (typically [1] for talker) - prompt: Original prompt data - requires_multimodal_data: Whether multimodal data is required - Returns: List of OmniTokensPrompt for code2wav stage """ - talker_outputs = _validate_stage_inputs(stage_list, engine_input_source) + talker_outputs = source_outputs code2wav_inputs: list[OmniTokensPrompt] = [] # Process each talker output for i, talker_output in enumerate(talker_outputs): diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 95a771534..3488eb5b0 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -20,16 +20,14 @@ logger = init_logger(__name__) def talker2code2wav( - stage_list: list[Any], - engine_input_source: list[int], + source_outputs: list[Any], prompt: Any = None, - requires_multimodal_data: bool = False, + _requires_multimodal_data: bool = False, ) -> list[Any]: """Non-async: collect all talker codes, then pass to code2wav at once.""" from vllm_omni.inputs.data import OmniTokensPrompt - from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _validate_stage_inputs - talker_outputs = _validate_stage_inputs(stage_list, engine_input_source) + talker_outputs = source_outputs code2wav_inputs: list[OmniTokensPrompt] = [] for i, talker_output in enumerate(talker_outputs): if not talker_output.finished: diff --git a/vllm_omni/model_executor/stage_input_processors/voxcpm.py b/vllm_omni/model_executor/stage_input_processors/voxcpm.py index c2fcf521b..c6bff820c 100644 --- a/vllm_omni/model_executor/stage_input_processors/voxcpm.py +++ b/vllm_omni/model_executor/stage_input_processors/voxcpm.py @@ -42,24 +42,10 @@ def _coerce_finished_flag(value: Any) -> bool: def latent2vae( - stage_list: list[Any], - engine_input_source: list[int], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, + source_outputs: list[Any], + _prompt: OmniTokensPrompt | TextPrompt | None = None, + _requires_multimodal_data: bool = False, ) -> list[OmniTokensPrompt]: - del prompt, requires_multimodal_data - - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - - source_stage_id = engine_input_source[0] - if source_stage_id >= len(stage_list): - raise IndexError(f"Invalid stage_id: {source_stage_id}") - - source_outputs = stage_list[source_stage_id].engine_outputs - if source_outputs is None: - raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") - vae_inputs: list[OmniTokensPrompt] = [] for source_output in source_outputs: output = source_output.outputs[0] diff --git a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py index e1a58b6f1..fc32ce6b1 100644 --- a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py @@ -10,22 +10,12 @@ logger = init_logger(__name__) def generator2tokenizer( - stage_list, - engine_input_source, - prompt: OmniTokensPrompt | TextPrompt = None, - requires_multimodal_data: bool = False, + source_outputs: list[Any], + _prompt: OmniTokensPrompt | TextPrompt = None, + _requires_multimodal_data: bool = False, ): - if not engine_input_source: - raise ValueError("engine_input_source cannot be empty") - source_stage_id = engine_input_source[0] - if source_stage_id >= len(stage_list): - raise IndexError(f"Invalid stage_id: {source_stage_id}") - if stage_list[source_stage_id].engine_outputs is None: - raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") - - generator_outputs = stage_list[source_stage_id].engine_outputs tokenizer_inputs = [] - for generator_output in generator_outputs: + for generator_output in source_outputs: output = generator_output.outputs[0] audio_tokens = torch.cat(output.multimodal_output["audio"], dim=-1).flatten().detach().cpu().tolist() tokenizer_inputs.append(