diff --git a/.buildkite/amd/test-amd-merge.yml b/.buildkite/amd/test-amd-merge.yml
index 63433cc12..4e80d1968 100644
--- a/.buildkite/amd/test-amd-merge.yml
+++ b/.buildkite/amd/test-amd-merge.yml
@@ -89,7 +89,7 @@ steps:
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
+ - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py
- label: "Diffusion Model Test"
agent_pool: mi325_1
diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml
index 8c645fff0..8f99ce909 100644
--- a/.buildkite/amd/test-amd-ready.yml
+++ b/.buildkite/amd/test-amd-ready.yml
@@ -35,7 +35,7 @@ steps:
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
+ - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
- label: "Voxtral TTS CUDA Unit Test"
diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml
index d72b8f287..e971acb75 100644
--- a/.buildkite/cuda/test-merge.yml
+++ b/.buildkite/cuda/test-merge.yml
@@ -8,14 +8,24 @@ steps:
- group: ":card_index_dividers: Simple Test"
depends_on: upload-merge-pipeline
steps:
- - label: "Simple · Diffusion & Model Executor Test"
+ - label: "Simple · Model Executor Test"
commands:
- - "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+ - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+ mirror_hardwares: l4_1
+
+ - label: "Simple · Diffusion Test"
+ commands:
+ - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+ mirror_hardwares: l4_1
+
+ - label: "Simple · Engine&Entrypoints Test"
+ commands:
+ - "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
mirror_hardwares: l4_1
- label: "Simple · Other Test"
commands:
- - "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+ - "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
mirror_hardwares: l4_1
- group: ":card_index_dividers: Diffusion Test"
@@ -27,12 +37,6 @@ steps:
- pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model"
mirror_hardwares: l4_1
- - label: "Diffusion · Offloader Test"
- timeout_in_minutes: 40
- commands:
- - timeout 40m pytest -s -v tests/diffusion/offloader -m "advanced_model and cuda"
- mirror_hardwares: l4_1
-
- label: "Distributed Test"
depends_on: upload-merge-pipeline
timeout_in_minutes: 60
@@ -135,39 +139,7 @@ steps:
timeout_in_minutes: 50
commands:
- pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'advanced_model and cuda' --run-level 'advanced_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: 2
- 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
+ mirror_hardwares: h100_2
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
source_file_dependencies:
@@ -210,22 +182,6 @@ steps:
"
mirror_hardwares: l4_1
- - label: "TTS · SoulX-Singer Test"
- source_file_dependencies:
- - tests/e2e/offline_inference/test_soulxsinger.py
- - vllm_omni/diffusion/models/soulx_singer/
- - vllm_omni/deploy/soulxsinger_svc.yaml
- - vllm_omni/deploy/soulxsinger_svs.yaml
- timeout_in_minutes: 30
- commands:
- - |
- timeout 30m bash -c "
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -m 'advanced_model and cuda' --run-level 'advanced_model'
- "
- mirror_hardwares: l4_1
-
- label: "Diffusion · Qwen Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_edit.py
diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml
index 72b31d0ce..377b78d1a 100644
--- a/.buildkite/cuda/test-nightly.yml
+++ b/.buildkite/cuda/test-nightly.yml
@@ -458,6 +458,11 @@ steps:
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
steps:
+ - group: ":full_moon: Diffusion · Distributed Test"
+ timeout_in_minutes: 30
+ commands:
+ - timeout 30m pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda' --run-level "full_model"
+ mirror_hardwares: l4_4
- label: ":full_moon: Diffusion · Offloader Test"
timeout_in_minutes: 40
@@ -473,6 +478,7 @@ steps:
- label: ":full_moon: Diffusion · Quantization Test with L4"
timeout_in_minutes: 60
commands:
+ - pip install "vllm-gguf-plugin>=0.0.3"
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
mirror_hardwares: l4_1
diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml
index b0ea85900..68149036d 100644
--- a/.buildkite/cuda/test-ready.yml
+++ b/.buildkite/cuda/test-ready.yml
@@ -8,22 +8,26 @@ steps:
- group: ":card_index_dividers: Simple Test"
depends_on: upload-ready-pipeline
steps:
- - label: "Simple · Diffusion & Model Executor Test"
+ - label: "Simple · Model Executor Test"
commands:
- - "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+ - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+ mirror_hardwares: l4_1
+
+ - label: "Simple · Diffusion Test"
+ commands:
+ - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+ mirror_hardwares: l4_1
+
+ - label: "Simple · Engine&Entrypoints Test"
+ commands:
+ - "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
mirror_hardwares: l4_1
- label: "Simple · Other Test"
commands:
- - "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+ - "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
mirror_hardwares: l4_1
- - label: "Custom Pipeline Test"
- depends_on: upload-ready-pipeline
- commands:
- - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
- mirror_hardwares: l4_1
-
- group: ":card_index_dividers: Diffusion Test"
depends_on: upload-ready-pipeline
steps:
@@ -104,11 +108,34 @@ steps:
- group: ":card_index_dividers: E2E Test"
depends_on: upload-ready-pipeline
steps:
+ - label: "Custom Pipeline Test"
+ source_file_dependencies:
+ - tests/e2e/features/custom_pipeline/
+ - tests/e2e/features/helpers/custom_pipeline.py
+ - vllm_omni/diffusion/
+ - vllm_omni/engine/
+ - vllm_omni/entrypoints/
+ - vllm_omni/inputs/
+ - vllm_omni/outputs/
+ commands:
+ - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
+ agents:
+ queue: "gpu_1_queue"
+ plugins:
+ - docker#v5.2.0:
+ image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+ always-pull: true
+ propagate-environment: true
+ environment:
+ - "HF_HOME=/fsx/hf_cache"
+ volumes:
+ - "/fsx/hf_cache:/fsx/hf_cache"
- label: "RLHF · VeRL-Omni E2E Test"
source_file_dependencies:
- - tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
- - tests/e2e/offline_inference/custom_pipeline/
+ - tests/e2e/features/rlhf_test/test_verl_omni_e2e.py
+ - tests/e2e/features/helpers/verl_omni_server.py
+ - tests/e2e/features/helpers/custom_pipeline.py
- vllm_omni/diffusion/
- vllm_omni/engine/
- vllm_omni/entrypoints/
@@ -117,7 +144,7 @@ steps:
commands:
- |
timeout 20m bash -c "
- pytest -s -v tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+ pytest -s -v tests/e2e/features/rlhf_test/test_verl_omni_e2e.py
"
mirror_hardwares: l4_1
@@ -143,34 +170,7 @@ steps:
timeout 20m bash -c "
pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.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: 2
- volumeMounts:
- - name: devshm
- mountPath: /dev/shm
- - name: hf-cache
- mountPath: /root/.cache/huggingface
- env:
- - name: HF_HOME
- value: /root/.cache/huggingface
- 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
+ mirror_hardwares: h100_2
- label: "Omni · MiniCPM-o 4.5 Online Test"
source_file_dependencies:
@@ -181,34 +181,7 @@ steps:
timeout 20m bash -c "
pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5.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: 2
- volumeMounts:
- - name: devshm
- mountPath: /dev/shm
- - name: hf-cache
- mountPath: /root/.cache/huggingface
- env:
- - name: HF_HOME
- value: /root/.cache/huggingface
- 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
+ mirror_hardwares: h100_2
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
source_file_dependencies:
@@ -371,13 +344,3 @@ steps:
pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
-
- - label: "Diffusion · GGUF Plugin Test"
- soft_fail: true
- source_file_dependencies:
- - tests/plugin/test_gguf_diffusion.py
- - vllm_omni/diffusion/model_loader
- commands:
- - pip install "vllm-gguf-plugin>=0.0.3"
- - timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py
- mirror_hardwares: l4_1
diff --git a/docs/contributing/ci/test_system_overview.md b/docs/contributing/ci/test_system_overview.md
index f1d4fb2c2..335ed5c53 100644
--- a/docs/contributing/ci/test_system_overview.md
+++ b/docs/contributing/ci/test_system_overview.md
@@ -56,7 +56,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
| / |
core_model and cpu |
<15min |
- /tests/{component_name}/test_xxx |
+
+ tests/{component}/… mirroring vllm_omni/{component}/
+ (e.g. tests/diffusion/, tests/engine/, tests/entrypoints/)
+ Do not add new top-level tests/ dirs unrelated to a component.
+ |
L1 & L2
Section 1 L1&L2: Purpose, Test Content, Directory Location, Example
@@ -68,16 +72,18 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
| L2
(E2E across models & GPU-required UT) |
Online (basic deployment scenarios):
dummy, normal inference function (output format, stream), some instance startup UT |
High-priority models + online basic scenarios; request success, non-empty output, format match (no Whisper/accuracy) |
- High-priority features (using random lightweight models) |
+ High-priority features and component tests that require launching instances |
High-priority interfaces (using random lightweight models) |
core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion |
- Model tests:
- /tests/e2e/online_serving/test_{model_name}.py
- Feature tests:
- /tests/{component_name}/test_xxx
- Interface tests:
- /tests/entrypoints/test_xxx
+ Model E2E:
+ tests/e2e/online_serving/test_{model}.py
+ tests/e2e/offline_inference/test_{model}.py
+ Feature integration:
+ tests/e2e/features/<feature>/
+ (e.g. fullduplex/, custom_pipeline/, rlhf_test/)
+ Component / interface Test (GPU):
+ tests/{component}/…, tests/entrypoints/…
|
L1 & L2
@@ -90,18 +96,19 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
| L3
(Important Perf & Integration & Accuracy) |
Online & Offline (multiple deployment scenarios):
real model, normal inference function, normal accuracy |
High/medium-priority models + key online/offline scenarios; real weights, Whisper/similarity, preset voice gender, basic accuracy |
- Medium-priority features (using random lightweight models) |
+ Medium-priority features and component tests that require launching instances |
Medium-priority interfaces (using random lightweight models) |
advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion |
<30min |
- Model tests:
- /tests/e2e/online_serving/test_{model_name}.py
- /tests/e2e/offline_inference/test_{model_name}.py
- Feature tests:
- /tests/{component_name}/test_xxx
- Interface tests:
- /tests/entrypoints/test_xxx
+ Model E2E:
+ tests/e2e/online_serving/test_{model}.py
+ tests/e2e/offline_inference/test_{model}.py
+ tests/e2e/accuracy/
+ Feature integration:
+ tests/e2e/features/<feature>/
+ Component / interface Test:
+ tests/{component}/…, tests/entrypoints/…
|
L3
@@ -114,26 +121,23 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
| L4
(Perf & Integration & Accuracy) |
Online: full functional scenarios + performance test + doc test + accuracy test |
High-priority models: function, performance, accuracy, and doc testing
Medium-priority models: function and doc testing |
- Low-priority features (using real weights) |
+ Low-priority features and component tests that require launching instances |
Low-priority interfaces (using real weights) |
full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion |
<3 hour |
- Model tests:
- /tests/e2e/online_serving/test_{model_name}_expansion.py
- Feature tests:
- /tests/{component_name}/test_xxx
- Interface tests:
- /tests/entrypoints/test_xxx
+ Model E2E:
+ tests/e2e/online_serving/test_{model}_expansion.py
+ tests/e2e/offline_inference/test_{model}_expansion.py
+ tests/e2e/accuracy/test_{model}.py
+ Feature integration:
+ tests/e2e/features/<feature>/
+ Component / interface Test:
+ tests/{component}/…, tests/entrypoints/…
Performance:
- /tests/dfx/perf/tests/test_qwen3_omni_*.json (Omni), test_tts.json (TTS),
- test_voxcpm2.json, test_higgs_audio_v3.json, and
- /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)
- Doc Test:
- tests/examples/online_serving/test_{model_name}.py
- tests/examples/offline_inference/test_{model_name}.py
- Accuracy Test:
- /tests/e2e/accuracy/test_{model_name}.py
+ tests/dfx/perf/tests/
+ Doc examples:
+ tests/examples/online_serving/, tests/examples/offline_inference/
|
L4
@@ -152,11 +156,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
| Depends on reality |
Stability:
- /tests/dfx/stability/tests/test_qwen3_omni.json
- /tests/dfx/stability/tests/test_wan22.json
+ tests/dfx/stability/tests/
Reliability:
- tests/dfx/reliability/test_reliability_{model_key}.py
- (e.g. test_reliability_qwen3_omni.py, test_reliability_wan22.py, test_reliability_hunyuan_image.py, test_reliability_voxcpm2.py)
+ tests/dfx/reliability/test_reliability_{model}.py
|
L5
@@ -168,7 +170,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
-For per-level test authoring (directories, markers, examples), see [Test Writing Guide](./test_writing_guide.md).
+**Test Dir placement (summary):** component / unit under `tests/{component}/` mirroring `vllm_omni/`; model E2E under `tests/e2e/online_serving/`, `tests/e2e/offline_inference/`, and `tests/e2e/accuracy/`; feature integration under `tests/e2e/features//`; doc example tests under `tests/examples/online_serving/` and `tests/examples/offline_inference/`; performance under `tests/dfx/perf/`; stability under `tests/dfx/stability/`; reliability under `tests/dfx/reliability/`. Do **not** add new top-level directories under `tests/` that are unrelated to a `vllm_omni` component (or to the established `e2e` / `dfx` / `helpers` / `examples` / `buildkite` layout).
+
+For per-level test authoring (markers, examples), see [Test Writing Guide](./test_writing_guide.md).
## Common Specifications
diff --git a/docs/contributing/ci/test_writing_guide.md b/docs/contributing/ci/test_writing_guide.md
index 131d20b83..397f5e83c 100644
--- a/docs/contributing/ci/test_writing_guide.md
+++ b/docs/contributing/ci/test_writing_guide.md
@@ -4,137 +4,20 @@ This document describes how to author tests for each CI level (L1–L5). For the
## Test Directory Structure
-Legend: `✅` = test exists, `⬜` = suggested to add.
-```
-vllm_omni/ tests/
-├── config/ → ├── config/
-│ ├── model.py │ └── test_model.py ⬜
-│ └── lora.py │ └── test_lora.py ⬜
-│
-├── core/ → ├── core/
-│ └── sched/ │ └── sched/
-│ ├── omni_ar_scheduler.py │ ├── test_omni_ar_scheduler.py ⬜
-│ ├── omni_generation_scheduler.py │ ├── test_omni_generation_scheduler.py ⬜
-│ └── output.py │ └── test_output.py ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData)
-│
-├── diffusion/ → ├── diffusion/
-│ ├── diffusion_engine.py │ ├── test_diffusion_engine.py ⬜
-│ ├── attention/ │ ├── attention/
-│ │ ├── layer.py │ │ ├── test_attention_sp.py ✅
-│ │ └── backends/ │ │ └── test_flash_attn.py ✅
-│ ├── distributed/ │ ├── distributed/
-│ │ └── ... │ │ ├── test_comm.py ✅
-│ │ │ │ ├── test_cfg_parallel.py ✅
-│ │ │ │ └── test_sp_plan_hooks.py ✅
-│ ├── lora/ │ ├── lora/
-│ │ └── ... │ │ ├── test_base_linear.py ✅
-│ │ │ │ └── test_lora_manager.py ✅
-│ ├── models/ │ ├── models/
-│ │ ├── qwen_image/ │ │ ├── qwen_image/ (e2e coverage)
-│ │ ├── ovis_image/ │ │ ├── ovis_image/
-│ │ │ └── ... │ │ │ └── test_ovis_image.py ✅
-│ │ ├── z_image/ │ │ └── z_image/
-│ │ └── ... │ │ └── test_zimage_tp_constraints.py ✅
-│ └── worker/ │ └── worker/
-│ ├── diffusion_worker.py │ └── test_diffusion_worker.py ✅ file at diffusion/test_diffusion_worker.py
-│ └── diffusion_model_runner.py │
-│
-├── distributed/ → ├── distributed/
-│ └── omni_connectors/ │ └── omni_connectors/
-│ ├── adapter.py │ ├── test_adapter_and_flow.py ✅
-│ ├── kv_transfer_manager.py │ ├── test_basic_connectors.py ✅
-│ ├── connectors/ │ ├── test_kv_flow.py ✅
-│ └── utils/ │ └── test_omni_connector_configs.py ✅
-│
-├── engine/ → ├── engine/
-│ ├── input_processor.py │ ├── test_input_processor.py ⬜ (no processor.py in source)
-│ ├── output_processor.py │ └── test_output_processor.py ⬜
-│ └── arg_utils.py │ └── test_arg_utils.py ⬜
-│
-├── entrypoints/ → ├── entrypoints/
-│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
-│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
-│ │ └── ... │ │ └── test_*.py ⬜
-│ └── openai/ │ └── openai_api/ # maps to entrypoints/openai/
-│ ├── api_server.py │ ├── test_api_server.py ⬜ (e2e indirect coverage)
-│ ├── serving_chat.py │ ├── test_serving_chat_sampling_params.py ✅
-│ ├── serving_speech.py │ ├── test_serving_speech.py ✅
-│ └── image_api_utils.py │ └── test_image_server.py ✅
-│
-├── inputs/ → ├── inputs/
-│ ├── data.py │ ├── test_data.py ⬜
-│ ├── parse.py │ ├── test_parse.py ⬜
-│ └── preprocess.py │ └── test_preprocess.py ✅ currently in entrypoints/test_omni_input_preprocessor.py
-│
-├── model_executor/ → ├── model_executor/
-│ ├── layers/ │ ├── layers/
-│ │ └── mrope.py │ │ └── test_mrope.py ⬜
-│ ├── model_loader/ │ ├── model_loader/
-│ │ └── weight_utils.py │ │ └── test_weight_utils.py ⬜
-│ ├── models/ │ ├── models/
-│ │ ├── qwen2_5_omni/ │ │ ├── qwen2_5_omni/
-│ │ │ ├── qwen2_5_omni_thinker.py │ │ │ ├── test_audio_length.py ✅
-│ │ │ ├── qwen2_5_omni_talker.py │ │ │ ├── test_qwen2_5_omni_thinker.py ⬜
-│ │ │ └── qwen2_5_omni_token2wav.py │ │ │ ├── test_qwen2_5_omni_talker.py ⬜
-│ │ └── qwen3_omni/ │ │ │ └── test_qwen2_5_omni_token2wav.py ⬜
-│ │ └── ... │ │ └── qwen3_omni/
-│ ├── stage_configs/ │ │ └── test_*.py ⬜
-│ │ └── *.yaml │ └── stage_configs/ (used by e2e, test_*.py can be added) ⬜
-│ └── stage_input_processors/ │ └── stage_input_processors/
-│ └── ... │ └── test_*.py ⬜
-│
-├── sample/ → ├── sample/
-│ └── __init__.py │ └── test_*.py ⬜
-│
-├── utils/ → ├── utils/
-│ └── __init__.py │ └── test_*.py ⬜ (no platform_utils.py currently)
-│
-├── worker/ → ├── worker/
-│ ├── gpu_ar_model_runner.py │ ├── test_gpu_ar_model_runner.py ⬜
-│ ├── gpu_ar_worker.py │ ├── test_gpu_ar_worker.py ⬜
-│ ├── gpu_generation_model_runner.py │ ├── test_gpu_generation_model_runner.py ✅
-│ ├── gpu_generation_worker.py │ ├── test_gpu_generation_worker.py ⬜
-│ ├── gpu_model_runner.py │ ├── test_omni_gpu_model_runner.py ✅
-│ └── mixins.py │ └── (npu under platforms/npu/worker/) # not worker/npu/
-│
-├── platforms/ → (no tests/platforms/, e2e and stage_configs provide indirect coverage)
-│ ├── cuda/
-│ ├── npu/worker/ # NPU worker here, not vllm_omni/worker/npu/
-│ ├── rocm/
-│ └── xpu/worker/
-│
-├── outputs.py → test_outputs.py ✅ (at tests root)
-├── (logger, patch, request, version) → (no corresponding unit test)
-│
-└── e2e (tests side only) → ├── e2e/
- ├── online_serving/ ✅ non-empty
- │ ├── test_qwen2_5_omni.py
- │ ├── test_async_omni.py
- │ ├── test_qwen3_omni.py
- │ ├── test_qwen3_omni_expansion.py
- │ ├── test_mimo_audio.py
- │ └── test_images_generations_lora.py
- └── offline_inference/ ✅
- ├── test_qwen2_5_omni.py
- ├── test_qwen3_omni.py
- ├── test_bagel_text2img.py
- ├── test_z_image.py
- ├── test_wan22.py
- ├── test_zimage_tensor_parallel.py
- ├── test_cache_dit.py
- ├── test_teacache.py
- ├── test_stable_audio_expansion.py
- ├── test_diffusion_cpu_offload.py
- ├── test_diffusion_layerwise_offload.py
- ├── test_diffusion_lora.py
- ├── test_sequence_parallel.py
- └── stage_configs/ (legacy schema, still
- ├── bagel_*.yaml present for unmigrated
- └── npu/, rocm/, etc. models)
+Canonical placement by CI level is in the **Test Dir** column of the level matrix in [Test System Overview](./test_system_overview.md).
-# Migrated models (qwen3_omni_moe, qwen2_5_omni, qwen3_tts) live under
-# vllm_omni/deploy/ instead — see docs/configuration/stage_configs.md.
-```
+Summary:
+
+- **Component / unit:** `tests/{component}/…` mirroring `vllm_omni/{component}/`
+- **Model E2E:** `tests/e2e/online_serving/`, `tests/e2e/offline_inference/`, `tests/e2e/accuracy/`
+- **Feature integration:** `tests/e2e/features//`
+- **Doc example tests:** `tests/examples/online_serving/`, `tests/examples/offline_inference/`
+- **Performance:** `tests/dfx/perf/tests/` (JSON configs; runners under `tests/dfx/perf/scripts/`)
+- **Stability:** `tests/dfx/stability/tests/`
+- **Reliability:** `tests/dfx/reliability/`
+- Do **not** add new top-level `tests/` directories unrelated to a component (or to `e2e` / `dfx` / `helpers` / `examples` / `buildkite`)
+
+Deploy YAMLs: `vllm_omni/deploy/` via `tests.helpers.stage_config.get_deploy_config_path` (see [stage configs](../../configuration/stage_configs.md)).
## Markers for Tests
@@ -373,13 +256,13 @@ L2 level testing builds upon L1 by introducing GPU resources and verifying that
#### 1.3 Test Directory and Execution Files
-A clear directory structure is key to managing test cases efficiently.
+A clear directory structure is key to managing test cases efficiently. See the **Test Dir** column in [Test System Overview](./test_system_overview.md) for the canonical placement rules.
-- ***L1 Test Directory***: `/tests/{component_name}/test_xxx.py`
-- - Here, `{component_name}` corresponds to modules in the source code, such as `distributed`, `entrypoints`, etc., and `test_xxx.py` is the specific test file.
-- ***L2 Test Directory***:
-- - Online Serving: `/tests/e2e/online_serving/test_{model_name}.py`
- - Offline Inference: `/tests/e2e/offline_inference/test_{model_name}.py`
+- ***L1 (component / unit)***: `tests/{component_name}/…` mirroring `vllm_omni/{component_name}/` (e.g. `tests/diffusion/`, `tests/entrypoints/`). Do **not** invent new top-level `tests/` folders unrelated to a component.
+- ***L2 (basic model E2E)***:
+ - Online serving: `tests/e2e/online_serving/test_{model_name}.py`
+ - Offline inference: `tests/e2e/offline_inference/test_{model_name}.py`
+- ***Feature-level integration*** (cross-cutting, not a single-model smoke): `tests/e2e/features//` (e.g. `fullduplex/`, `custom_pipeline/`, `rlhf_test/`).
#### 1.4 Execution Method and Example
@@ -485,10 +368,12 @@ L3 level testing executes after code is merged into the main branch. Its core pu
#### 2.3 Test Directory and Execution Files
-- ***Functional Testing***:
-- - Online Serving: `/tests/e2e/online_serving/test_{model_name}_expansion.py`
- - Offline Inference: `/tests/e2e/offline_inference/test_{model_name}_expansion.py`
- - (Note: `_expansion.py` likely means it contains more comprehensive scenario cases compared to L2 tests).
+- ***Functional / model E2E***:
+ - Online serving: `tests/e2e/online_serving/test_{model_name}_expansion.py`
+ - Offline inference: `tests/e2e/offline_inference/test_{model_name}_expansion.py`
+ - Accuracy / similarity: `tests/e2e/accuracy/`
+ - (`_expansion.py` holds broader scenario coverage than the L2 smoke file.)
+- ***Feature-level integration***: `tests/e2e/features//` (not under a new top-level `tests/` directory).
#### 2.4 Execution Method and Example
diff --git a/docs/design/feature/tensor_parallel.md b/docs/design/feature/tensor_parallel.md
index bcafde7e7..a01e2a0e6 100644
--- a/docs/design/feature/tensor_parallel.md
+++ b/docs/design/feature/tensor_parallel.md
@@ -264,7 +264,7 @@ Complete examples in the codebase:
| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation |
| **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams |
| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE |
-| **TP Tests** | `tests/e2e/offline_inference/test_zimage_parallelism.py` | E2E testing | TP correctness and performance |
+| **TP Tests** | `tests/diffusion/distributed/test_tensor_parallel.py` | E2E testing | TP correctness and performance |
| **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic |
---
diff --git a/examples/model_display_names.yml b/examples/model_display_names.yml
index be3f8e381..7a04ff8fa 100644
--- a/examples/model_display_names.yml
+++ b/examples/model_display_names.yml
@@ -8,5 +8,4 @@ mimo_audio: MiMo-Audio
ming_flash_omni: Ming-flash-omni 2.0
qwen2_5_omni: Qwen2.5-Omni
qwen3_omni: Qwen3-Omni
-sensenova_u1: SenseNova-U1
step_audio2: Step-Audio2
diff --git a/recipes/JD/JoyAI-VL-Interaction.md b/recipes/JD/JoyAI-VL-Interaction.md
index ed767830f..e09d5b16c 100644
--- a/recipes/JD/JoyAI-VL-Interaction.md
+++ b/recipes/JD/JoyAI-VL-Interaction.md
@@ -168,7 +168,7 @@ the model + orchestrator + WebUI together — point `WEBUI_DIR` at your clone.
python examples/online_serving/joyvl_interaction/cli/run_cli_demo.py \
path/to/video.mp4 --query "Alert me if a fire breaks out"
-pytest tests/fullduplex # framework + JoyVL unit tests
+pytest tests/e2e/features/fullduplex # framework + JoyVL unit tests
```
## Testing with an RTSP stream (optional)
diff --git a/recipes/NVIDIA/GR00T-N1.7.md b/recipes/NVIDIA/GR00T-N1.7.md
index 91f676a52..0a5270222 100644
--- a/recipes/NVIDIA/GR00T-N1.7.md
+++ b/recipes/NVIDIA/GR00T-N1.7.md
@@ -69,8 +69,13 @@ Notes:
## Verification
```python
-from tests.gr00t.openpi_client_helper import run_policy_session, validate_session_result
-validate_session_result(run_policy_session(host="127.0.0.1", port=8000))
+from tests.helpers.runtime import OpenAIClientHandler
+
+handler = OpenAIClientHandler(host="127.0.0.1", port=8000, log_stats=False)
+response = handler.send_robot_openpi_ws_request(
+ {"run_default_policy_session": True, "session_id": "gr00t-smoke"}
+)[0]
+assert response.operation_responses[-1]["status"] == "reset successful"
```
Or run the e2e test suite:
diff --git a/tests/e2e/accuracy/assets/hunyuan_image_ref.png b/tests/assets/hunyuan/hunyuan_image_ref.png
similarity index 100%
rename from tests/e2e/accuracy/assets/hunyuan_image_ref.png
rename to tests/assets/hunyuan/hunyuan_image_ref.png
diff --git a/tests/e2e/accuracy/assets/qwen_image_edit_2511_test1.png b/tests/assets/qwen_image_edit/qwen_image_edit_2511_test1.png
similarity index 100%
rename from tests/e2e/accuracy/assets/qwen_image_edit_2511_test1.png
rename to tests/assets/qwen_image_edit/qwen_image_edit_2511_test1.png
diff --git a/tests/e2e/accuracy/assets/qwen_image_edit_2511_test2.png b/tests/assets/qwen_image_edit/qwen_image_edit_2511_test2.png
similarity index 100%
rename from tests/e2e/accuracy/assets/qwen_image_edit_2511_test2.png
rename to tests/assets/qwen_image_edit/qwen_image_edit_2511_test2.png
diff --git a/tests/test_config_factory.py b/tests/config/test_config_factory.py
similarity index 97%
rename from tests/test_config_factory.py
rename to tests/config/test_config_factory.py
index 61d870e98..5a59dcf32 100644
--- a/tests/test_config_factory.py
+++ b/tests/config/test_config_factory.py
@@ -1145,7 +1145,7 @@ class TestResolveScheduler:
class TestDeployConfigLoading:
def test_load_minicpmo_duplex_deploy_config(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "minicpmo_4_5_duplex.yaml"
+ deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
deploy = load_deploy_config(deploy_path)
pipeline = resolve_pipeline_config("minicpmo_4_5")
@@ -1182,7 +1182,7 @@ class TestDeployConfigLoading:
stage2_devices: str,
stage1_replicas: int,
):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / filename
+ deploy_path = Path(get_deploy_config_path(filename))
deploy = load_deploy_config(deploy_path)
pipeline = resolve_pipeline_config("minicpmo_4_5")
@@ -1226,7 +1226,7 @@ stages:
assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)}
def test_load_qwen3_omni_moe_deploy_config(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
deploy = load_deploy_config(deploy_path)
assert len(deploy.stages) == 3
assert deploy.async_chunk is True
@@ -1234,7 +1234,7 @@ stages:
assert deploy.platforms is not None
def test_load_voxtral_tts_deploy_config_schema_fields(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "voxtral_tts.yaml"
+ deploy_path = Path(get_deploy_config_path("voxtral_tts.yaml"))
deploy = load_deploy_config(deploy_path)
assert deploy.stages[0].config_format == "mistral"
assert deploy.stages[0].load_format == "mistral"
@@ -1244,7 +1244,7 @@ stages:
)
def test_load_ming_flash_omni_deploy_config_schema_fields(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
+ deploy_path = Path(get_deploy_config_path("ming_flash_omni.yaml"))
deploy = load_deploy_config(deploy_path)
assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}}
assert "compilation_config" not in deploy.stages[0].engine_extras
@@ -1384,7 +1384,7 @@ stages:
assert not list(stage_config_dir.glob("*.yaml"))
def test_merge_pipeline_deploy(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not deploy_path.exists():
pytest.skip("Deploy config not found")
@@ -1411,7 +1411,7 @@ stages:
)
assert isinstance(pipeline, PipelineConfig)
- base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not base.exists():
pytest.skip("Deploy config not found")
@@ -1831,7 +1831,7 @@ class TestQwen3TTSPipeline:
def test_per_stage_model_arch_flows_through_merge(self, tmp_path):
"""Verify the new ps.model_arch override survives merge_pipeline_deploy."""
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_tts.yaml"
+ deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml"))
if not deploy_path.exists():
pytest.skip("qwen3_tts deploy yaml not found")
@@ -1955,7 +1955,7 @@ class TestMingFlashOmniPipeline:
def test_full_yaml_loads_and_merges(self):
"""deploy/ming_flash_omni.yaml parses and merges with the registered pipeline."""
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
+ deploy_path = Path(get_deploy_config_path("ming_flash_omni.yaml"))
if not deploy_path.exists():
pytest.skip("ming_flash_omni deploy yaml not found")
@@ -1975,7 +1975,7 @@ class TestMingFlashOmniPipeline:
def test_tts_yaml_loads_and_merges(self):
"""deploy/ming_flash_omni_tts.yaml parses and routes to the TTS-only pipeline."""
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_tts.yaml"
+ deploy_path = Path(get_deploy_config_path("ming_flash_omni_tts.yaml"))
if not deploy_path.exists():
pytest.skip("ming_flash_omni_tts deploy yaml not found")
@@ -2013,7 +2013,7 @@ class TestMingFlashOmniPipeline:
def test_thinker_only_yaml_loads_and_merges(self):
"""deploy/ming_flash_omni_thinker_only.yaml parses and routes to the thinker-only pipeline."""
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_thinker_only.yaml"
+ deploy_path = Path(get_deploy_config_path("ming_flash_omni_thinker_only.yaml"))
if not deploy_path.exists():
pytest.skip("ming_flash_omni_thinker_only deploy yaml not found")
@@ -2073,7 +2073,7 @@ class TestMingFlashOmniPipeline:
def test_image_yaml_loads_and_merges(self):
"""deploy/ming_flash_omni_image.yaml parses and routes to the image pipeline."""
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_image.yaml"
+ deploy_path = Path(get_deploy_config_path("ming_flash_omni_image.yaml"))
if not deploy_path.exists():
pytest.skip("ming_flash_omni_image deploy yaml not found")
@@ -2135,7 +2135,7 @@ class TestBaseConfigInheritance:
def test_pure_inheritance_overlay(self, tmp_path):
"""An overlay with only ``base_config`` inherits everything."""
- base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not base.exists():
pytest.skip("Base deploy config not found")
@@ -2148,7 +2148,7 @@ class TestBaseConfigInheritance:
def test_single_field_overlay(self, tmp_path):
"""An overlay overriding one stage field merges with the base."""
- base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not base.exists():
pytest.skip("Base deploy config not found")
@@ -2165,7 +2165,7 @@ class TestPlatformOverrides:
"""Test platform-specific deploy config overrides."""
def test_npu_overrides(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not deploy_path.exists():
pytest.skip("Deploy config not found")
@@ -2179,7 +2179,7 @@ class TestPlatformOverrides:
assert deploy.stages[2].enforce_eager is False
def test_xpu_overrides(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not deploy_path.exists():
pytest.skip("Deploy config not found")
@@ -2191,7 +2191,7 @@ class TestPlatformOverrides:
assert deploy.stages[0].engine_extras.get("max_cudagraph_capture_size") == 0
def test_unknown_platform_noop(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not deploy_path.exists():
pytest.skip("Deploy config not found")
@@ -2236,7 +2236,7 @@ class TestCLIOverrideFlow:
"""Test --stage-overrides JSON merge into StageConfig."""
def test_stage_overrides_merge(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not deploy_path.exists():
pytest.skip("Deploy config not found")
@@ -2254,7 +2254,7 @@ class TestCLIOverrideFlow:
assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5
def test_global_override_applies_to_all(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not deploy_path.exists():
pytest.skip("Deploy config not found")
@@ -2274,7 +2274,7 @@ class TestCLIOverrideFlow:
class TestAuraOmniDeploy:
def test_aura_omni_deploy_forces_pipeline_override(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"
+ deploy_path = Path(get_deploy_config_path("aura_omni.yaml"))
deploy = load_deploy_config(deploy_path)
assert deploy.pipeline == "aura_omni"
@@ -2285,7 +2285,7 @@ class TestAuraOmniDeploy:
stages, _ = StageConfigFactory._create_legacy_from_registry(
pipeline_cfg,
cli_overrides={},
- deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"),
+ deploy_config_path=get_deploy_config_path("aura_omni.yaml"),
)
assert [stage.model_stage for stage in stages] == [
@@ -2599,7 +2599,12 @@ class TestSentinelDefaultPrecedence:
from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy
pipeline_path = (
- Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "models" / "qwen3_omni" / "pipeline.py"
+ Path(__file__).parent.parent.parent
+ / "vllm_omni"
+ / "model_executor"
+ / "models"
+ / "qwen3_omni"
+ / "pipeline.py"
)
pipeline = runpy.run_path(str(pipeline_path))["QWEN3_OMNI_PIPELINE"]
@@ -2669,7 +2674,7 @@ class TestSamplingConstraintsPrecedence:
"""Test that pipeline sampling_constraints override deploy defaults."""
def test_constraints_win(self):
- deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+ deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not deploy_path.exists():
pytest.skip("Deploy config not found")
diff --git a/tests/conftest.py b/tests/conftest.py
index 11619b782..b453bbcbf 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,6 +3,7 @@ Root pytest entrypoint for the vLLM-Omni test suite.
- `tests/conftest.py` stays thin: plugin registration + compatibility re-exports.
- Importable utilities live under `tests/helpers/`.
+- Helper unit tests live under `tests/helpers/tests/`.
- Fixtures live under `tests/helpers/fixtures/` and are loaded via `pytest_plugins`.
"""
diff --git a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py
index 053c926cc..5fad4340c 100644
--- a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py
+++ b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py
@@ -18,6 +18,8 @@ import pytest
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
class _FakeCoordinator:
def __init__(self, timed_out_ids):
diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py
index 2111556f2..db7610496 100644
--- a/tests/core/test_prefix_cache.py
+++ b/tests/core/test_prefix_cache.py
@@ -3,6 +3,9 @@ import torch
from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
DEFAULT_SEQ_LEN = 15
NUM_BLOCKS = 10
BLOCK_SIZE = 4
diff --git a/tests/ar_diffusion/__init__.py b/tests/diffusion/ar_diffusion/__init__.py
similarity index 100%
rename from tests/ar_diffusion/__init__.py
rename to tests/diffusion/ar_diffusion/__init__.py
diff --git a/tests/ar_diffusion/test_capability_runner.py b/tests/diffusion/ar_diffusion/test_capability_runner.py
similarity index 100%
rename from tests/ar_diffusion/test_capability_runner.py
rename to tests/diffusion/ar_diffusion/test_capability_runner.py
diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/diffusion/ar_diffusion/test_engine_routing.py
similarity index 98%
rename from tests/ar_diffusion/test_engine_routing.py
rename to tests/diffusion/ar_diffusion/test_engine_routing.py
index 5972ffff4..19c4cb529 100644
--- a/tests/ar_diffusion/test_engine_routing.py
+++ b/tests/diffusion/ar_diffusion/test_engine_routing.py
@@ -21,6 +21,8 @@ from vllm_omni.experimental.ar_diffusion.engine import (
ARDiffusionEngine,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
def _cfg(backend):
return SimpleNamespace(engine_backend=backend)
diff --git a/tests/ar_diffusion/test_kv_cache.py b/tests/diffusion/ar_diffusion/test_kv_cache.py
similarity index 99%
rename from tests/ar_diffusion/test_kv_cache.py
rename to tests/diffusion/ar_diffusion/test_kv_cache.py
index 9ca0e547d..d09c1c1b6 100644
--- a/tests/ar_diffusion/test_kv_cache.py
+++ b/tests/diffusion/ar_diffusion/test_kv_cache.py
@@ -25,6 +25,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
)
from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
BLOCK = 16
diff --git a/tests/ar_diffusion/test_kv_orchestration.py b/tests/diffusion/ar_diffusion/test_kv_orchestration.py
similarity index 98%
rename from tests/ar_diffusion/test_kv_orchestration.py
rename to tests/diffusion/ar_diffusion/test_kv_orchestration.py
index b7c77390c..32603bf88 100644
--- a/tests/ar_diffusion/test_kv_orchestration.py
+++ b/tests/diffusion/ar_diffusion/test_kv_orchestration.py
@@ -7,6 +7,7 @@ import torch
from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
DIMS = dict(
num_layers=2,
num_kv_heads=4,
diff --git a/tests/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py
similarity index 94%
rename from tests/ar_diffusion/test_paged_attention.py
rename to tests/diffusion/ar_diffusion/test_paged_attention.py
index b5c5e5870..0b66c01a0 100644
--- a/tests/ar_diffusion/test_paged_attention.py
+++ b/tests/diffusion/ar_diffusion/test_paged_attention.py
@@ -22,6 +22,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
)
from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
BLOCK = 16
N_HEADS = 4
HEAD_DIM = 64
@@ -385,17 +388,22 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
return out
torch._dynamo.reset()
- from torch._dynamo.testing import CompileCounter
+ try:
+ from torch._dynamo.testing import CompileCounter
- counter = CompileCounter()
+ counter = CompileCounter()
- def fn(inputs, q, k, v):
- return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
+ def fn(inputs, q, k, v):
+ return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
- compiled = torch.compile(fn, backend=counter, fullgraph=True)
+ compiled = torch.compile(fn, backend=counter, fullgraph=True)
- run_one_forward(commit=True) # history grows between calls ->
- run_one_forward(commit=True) # block-table VALUES change, shapes don't
- run_one_forward(commit=False)
+ run_one_forward(commit=True) # history grows between calls ->
+ run_one_forward(commit=True) # block-table VALUES change, shapes don't
+ run_one_forward(commit=False)
- assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
+ assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
+ finally:
+ # Leave a clean dynamo state for later suites in the same pytest process
+ # (e.g. model_executor transformers models).
+ torch._dynamo.reset()
diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py
similarity index 99%
rename from tests/ar_diffusion/test_pipeline_bridge.py
rename to tests/diffusion/ar_diffusion/test_pipeline_bridge.py
index 92d490b92..ec60290b9 100644
--- a/tests/ar_diffusion/test_pipeline_bridge.py
+++ b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py
@@ -12,6 +12,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
)
from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
BLOCK = 16
N_HEADS = 4
HEAD_DIM = 64
diff --git a/tests/ar_diffusion/test_slot_mapping.py b/tests/diffusion/ar_diffusion/test_slot_mapping.py
similarity index 97%
rename from tests/ar_diffusion/test_slot_mapping.py
rename to tests/diffusion/ar_diffusion/test_slot_mapping.py
index 061df4d9c..4c1f82537 100644
--- a/tests/ar_diffusion/test_slot_mapping.py
+++ b/tests/diffusion/ar_diffusion/test_slot_mapping.py
@@ -13,6 +13,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
resident_block_ids,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
BLOCK = 16
diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py
index 26a93e901..adf13d51c 100644
--- a/tests/diffusion/attention/test_attention_config.py
+++ b/tests/diffusion/attention/test_attention_config.py
@@ -288,6 +288,13 @@ class TestBuildAttentionConfig:
class TestOmniDiffusionConfigAttentionParsing:
"""Test OmniDiffusionConfig attention shorthand and structured config."""
+ @pytest.fixture(autouse=True)
+ def _clear_diffusion_attention_backend_env(self, monkeypatch):
+ # OmniDiffusionConfig.__post_init__ applies DIFFUSION_ATTENTION_BACKEND via
+ # build_attention_config(); clear it so these tests assert config defaults,
+ # not whatever the process inherited from CI / sibling tests.
+ monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False)
+
def test_diffusion_attention_backend_sets_default(self):
config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN")
assert isinstance(config.diffusion_attention_config, AttentionConfig)
diff --git a/tests/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py
index e32c4c5e4..0f1faacd1 100644
--- a/tests/diffusion/attention/test_attention_sp.py
+++ b/tests/diffusion/attention/test_attention_sp.py
@@ -49,6 +49,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
from vllm_omni.diffusion.forward_context import set_forward_context
from vllm_omni.platforms import current_omni_platform
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
def update_environment_variables(envs_dict: dict[str, str]):
"""Update multiple environment variables with logging."""
diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py
index 5ffd6684a..fbe037274 100644
--- a/tests/diffusion/attention/test_flash_attn.py
+++ b/tests/diffusion/attention/test_flash_attn.py
@@ -18,6 +18,9 @@ from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
from vllm_omni.diffusion.attention.backends.utils import fa # noqa: E402
from vllm_omni.platforms import current_omni_platform
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
flash_attn_func = fa.flash_attn_func # noqa: N813
diff --git a/tests/diffusion/attention/test_piecewise_attn.py b/tests/diffusion/attention/test_piecewise_attn.py
index 00b98ac64..05967e28f 100644
--- a/tests/diffusion/attention/test_piecewise_attn.py
+++ b/tests/diffusion/attention/test_piecewise_attn.py
@@ -27,6 +27,9 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
piecewise_attn,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
DEVICE = torch.device("cpu")
diff --git a/tests/diffusion/attention/test_sage_attn3.py b/tests/diffusion/attention/test_sage_attn3.py
index e4d913974..5c46ea5d0 100644
--- a/tests/diffusion/attention/test_sage_attn3.py
+++ b/tests/diffusion/attention/test_sage_attn3.py
@@ -9,6 +9,9 @@ import pytest
import torch
from vllm.platforms import current_platform
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3"
diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py
index a9e2021c3..6bb588156 100644
--- a/tests/diffusion/attention/test_ulysses_uaa.py
+++ b/tests/diffusion/attention/test_ulysses_uaa.py
@@ -22,6 +22,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
from vllm_omni.platforms import current_omni_platform
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
def _find_free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
diff --git a/tests/diffusion/distributed/test_cfg_parallel.py b/tests/diffusion/distributed/test_cfg_parallel.py
index af81f00ef..8991ddec3 100644
--- a/tests/diffusion/distributed/test_cfg_parallel.py
+++ b/tests/diffusion/distributed/test_cfg_parallel.py
@@ -1,18 +1,23 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for CFG (Classifier-Free Guidance) parallel functionality.
+"""Tests for CFG (Classifier-Free Guidance) parallel functionality.
-This test verifies that predict_noise_maybe_with_cfg and
-predict_noise_with_multi_branch_cfg produce numerically equivalent results
-with and without CFG parallel using fixed random inputs.
+CPU tests mock the CFG process group and validate mixin logic without GPUs.
+Nightly GPU parity tests compare sequential CFG (cfg_parallel_size=1) against
+real multi-GPU CFG parallel under fixed seeds.
"""
+from __future__ import annotations
+
import os
+from typing import Any
import pytest
import torch
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from tests.helpers.mark import hardware_marks, hardware_test
+from vllm_omni.diffusion.distributed import parallel_state
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _dispatch_branches, _wrap
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env,
get_classifier_free_guidance_rank,
@@ -22,445 +27,105 @@ from vllm_omni.diffusion.distributed.parallel_state import (
)
from vllm_omni.platforms import current_omni_platform
+_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3)
-def update_environment_variables(envs_dict: dict[str, str]):
- """Update multiple environment variables."""
- for k, v in envs_dict.items():
- os.environ[k] = v
+
+def _set_random_seeds(seed: int) -> None:
+ torch.manual_seed(seed)
+ if torch.cuda.is_available():
+ torch.cuda.manual_seed_all(seed)
+
+
+def _init_model_weights(module: torch.nn.Module, seed: int) -> None:
+ _set_random_seeds(seed)
+ for param in module.parameters():
+ torch.nn.init.normal_(param, mean=0.0, std=0.02)
+
+
+def _cfg_close_tolerance(dtype: torch.dtype) -> tuple[float, float]:
+ if dtype == torch.float32:
+ return 1e-5, 1e-5
+ if dtype == torch.bfloat16:
+ return 1e-2, 1e-2
+ return 1e-3, 1e-3
+
+
+def _assert_cfg_outputs_close(
+ actual: torch.Tensor,
+ expected: torch.Tensor,
+ *,
+ dtype: torch.dtype,
+ msg: str,
+) -> None:
+ rtol, atol = _cfg_close_tolerance(dtype)
+ torch.testing.assert_close(actual, expected, rtol=rtol, atol=atol, msg=msg)
+
+
+def _update_environment_variables(envs_dict: dict[str, str]) -> None:
+ for key, value in envs_dict.items():
+ os.environ[key] = value
class SimpleTransformer(torch.nn.Module):
- """Simple transformer model for testing with random initialization.
-
- Contains:
- - Input projection (conv to hidden_dim)
- - QKV projection layers
- - Self-attention layer
- - Output projection
- """
-
def __init__(self, in_channels: int = 4, hidden_dim: int = 128, num_heads: int = 8):
super().__init__()
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.head_dim = hidden_dim // num_heads
+ assert hidden_dim % num_heads == 0
- assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads"
-
- # Input projection: (B, C, H, W) -> (B, hidden_dim, H, W)
self.input_proj = torch.nn.Conv2d(in_channels, hidden_dim, 1)
-
- # QKV projection layers
self.q_proj = torch.nn.Linear(hidden_dim, hidden_dim)
self.k_proj = torch.nn.Linear(hidden_dim, hidden_dim)
self.v_proj = torch.nn.Linear(hidden_dim, hidden_dim)
-
- # Output projection after attention
self.out_proj = torch.nn.Linear(hidden_dim, hidden_dim)
-
- # Final output projection: (B, hidden_dim, H, W) -> (B, C, H, W)
self.final_proj = torch.nn.Conv2d(hidden_dim, in_channels, 1)
-
- # Layer norm
self.norm1 = torch.nn.LayerNorm(hidden_dim)
self.norm2 = torch.nn.LayerNorm(hidden_dim)
def forward(self, x: torch.Tensor, **kwargs) -> tuple[torch.Tensor]:
- """Forward pass with self-attention.
-
- Args:
- x: Input tensor of shape (B, C, H, W)
-
- Returns:
- Output tensor of shape (B, C, H, W)
- """
- B, C, H, W = x.shape
-
- # Input projection
- x = self.input_proj(x) # (B, hidden_dim, H, W)
-
- # Reshape to sequence: (B, hidden_dim, H, W) -> (B, H*W, hidden_dim)
- x = x.flatten(2).transpose(1, 2) # (B, H*W, hidden_dim)
-
- # Self-attention with residual connection
+ del kwargs
+ batch_size, _channels, height, width = x.shape
+ x = self.input_proj(x)
+ x = x.flatten(2).transpose(1, 2)
residual = x
x = self.norm1(x)
-
- # QKV projection
- q = self.q_proj(x) # (B, H*W, hidden_dim)
- k = self.k_proj(x) # (B, H*W, hidden_dim)
- v = self.v_proj(x) # (B, H*W, hidden_dim)
-
- # Reshape for multi-head attention: (B, H*W, hidden_dim) -> (B, num_heads, H*W, head_dim)
- seq_len = H * W
- q = q.view(B, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
- k = k.view(B, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
- v = v.view(B, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-
- # Scaled dot-product attention
+ q = self.q_proj(x)
+ k = self.k_proj(x)
+ v = self.v_proj(x)
+ seq_len = height * width
+ q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+ k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+ v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
scale = self.head_dim**-0.5
- attn_scores = torch.matmul(q, k.transpose(-2, -1)) * scale # (B, num_heads, H*W, H*W)
+ attn_scores = torch.matmul(q, k.transpose(-2, -1)) * scale
attn_weights = torch.nn.functional.softmax(attn_scores, dim=-1)
attn_output = torch.matmul(attn_weights, v)
- attn_output = attn_output.transpose(1, 2).contiguous().view(B, seq_len, self.hidden_dim)
-
+ attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_dim)
attn_output = self.out_proj(attn_output)
-
x = residual + attn_output
residual = x
x = self.norm2(x)
x = residual + x
- x = x.transpose(1, 2).view(B, self.hidden_dim, H, W)
-
- out = self.final_proj(x)
-
- return (out,)
+ x = x.transpose(1, 2).view(batch_size, self.hidden_dim, height, width)
+ return (self.final_proj(x),)
class TestCFGPipeline(CFGParallelMixin):
- """Test pipeline using CFGParallelMixin."""
-
def __init__(self, in_channels: int = 4, hidden_dim: int = 128, seed: int = 42):
- # Set seed BEFORE creating transformer to ensure consistent layer initialization
- torch.manual_seed(seed)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(seed)
-
+ _set_random_seeds(seed)
self.transformer = SimpleTransformer(in_channels, hidden_dim)
-
- # Re-initialize all parameters with fixed seed for full reproducibility
- torch.manual_seed(seed)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(seed)
- for param in self.transformer.parameters():
- torch.nn.init.normal_(param, mean=0.0, std=0.02)
-
-
-def _test_cfg_parallel_worker(
- local_rank: int,
- world_size: int,
- cfg_parallel_size: int,
- dtype: torch.dtype,
- test_config: dict,
- result_queue: torch.multiprocessing.Queue,
-):
- """Worker function for CFG parallel test."""
- device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
- current_omni_platform.set_device(device)
-
- update_environment_variables(
- {
- "RANK": str(local_rank),
- "LOCAL_RANK": str(local_rank),
- "WORLD_SIZE": str(world_size),
- "MASTER_ADDR": "localhost",
- "MASTER_PORT": "29502",
- }
- )
-
- init_distributed_environment()
- initialize_model_parallel(cfg_parallel_size=cfg_parallel_size)
-
- cfg_rank = get_classifier_free_guidance_rank()
- cfg_world_size = get_classifier_free_guidance_world_size()
-
- assert cfg_world_size == cfg_parallel_size
-
- # Create pipeline with same seed to ensure identical model weights across all ranks
- # Note: model_seed is set inside TestCFGPipeline.__init__
- pipeline = TestCFGPipeline(
- in_channels=test_config["channels"],
- hidden_dim=test_config["hidden_dim"],
- seed=test_config["model_seed"],
- )
- pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
- pipeline.transformer.eval() # Set to eval mode for deterministic behavior
-
- # Create fixed inputs with explicit seed setting for reproducibility
- # Set both CPU and CUDA seeds to ensure identical inputs across all ranks
- torch.manual_seed(test_config["input_seed"])
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(test_config["input_seed"])
-
- batch_size = test_config["batch_size"]
- channels = test_config["channels"]
- height = test_config["height"]
- width = test_config["width"]
-
- # Positive input
- positive_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
-
- # Negative input with different seed
- torch.manual_seed(test_config["input_seed"] + 1)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
- negative_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
-
- # Prepare kwargs for predict_noise_maybe_with_cfg
- positive_kwargs = {"x": positive_input}
- negative_kwargs = {"x": negative_input}
-
- with torch.no_grad():
- # Call predict_noise_maybe_with_cfg
- noise_pred = pipeline.predict_noise_maybe_with_cfg(
- do_true_cfg=True,
- true_cfg_scale=test_config["cfg_scale"],
- positive_kwargs=positive_kwargs,
- negative_kwargs=negative_kwargs,
- cfg_normalize=test_config["cfg_normalize"],
- kwargs=test_config["kwargs"],
- )
-
- # CFG parallel returns the combined prediction on every rank.
- assert noise_pred is not None
- result_queue.put((cfg_rank, noise_pred.cpu()))
-
- destroy_distributed_env()
-
-
-def _test_cfg_sequential_worker(
- local_rank: int,
- world_size: int,
- dtype: torch.dtype,
- test_config: dict,
- result_queue: torch.multiprocessing.Queue,
-):
- """Worker function for sequential CFG test (baseline)."""
- device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
- current_omni_platform.set_device(device)
-
- update_environment_variables(
- {
- "RANK": str(local_rank),
- "LOCAL_RANK": str(local_rank),
- "WORLD_SIZE": str(world_size),
- "MASTER_ADDR": "localhost",
- "MASTER_PORT": "29503",
- }
- )
-
- init_distributed_environment()
- initialize_model_parallel(cfg_parallel_size=1) # No CFG parallel
-
- cfg_world_size = get_classifier_free_guidance_world_size()
- assert cfg_world_size == 1
-
- # Create pipeline with same seed to ensure identical model weights as CFG parallel
- # Note: model_seed is set inside TestCFGPipeline.__init__
- pipeline = TestCFGPipeline(
- in_channels=test_config["channels"],
- hidden_dim=test_config["hidden_dim"],
- seed=test_config["model_seed"],
- )
- pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
- pipeline.transformer.eval()
-
- # Create fixed inputs (same seed as CFG parallel to ensure identical inputs)
- # Set both CPU and CUDA seeds for full reproducibility
- torch.manual_seed(test_config["input_seed"])
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(test_config["input_seed"])
-
- batch_size = test_config["batch_size"]
- channels = test_config["channels"]
- height = test_config["height"]
- width = test_config["width"]
-
- # Positive input
- positive_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
-
- # Negative input with different seed
- torch.manual_seed(test_config["input_seed"] + 1)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
- negative_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
-
- positive_kwargs = {"x": positive_input}
- negative_kwargs = {"x": negative_input}
-
- with torch.no_grad():
- noise_pred = pipeline.predict_noise_maybe_with_cfg(
- do_true_cfg=True,
- true_cfg_scale=test_config["cfg_scale"],
- positive_kwargs=positive_kwargs,
- negative_kwargs=negative_kwargs,
- cfg_normalize=test_config["cfg_normalize"],
- kwargs=test_config["kwargs"],
- )
-
- # Sequential CFG always returns output
- assert noise_pred is not None
- result_queue.put(noise_pred.cpu())
-
- destroy_distributed_env()
-
-
-@pytest.mark.parametrize("cfg_parallel_size", [2])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("cfg_normalize", [False, True])
-@pytest.mark.parametrize("kwargs", [None, {"step_i": 1}])
-def test_predict_noise_maybe_with_cfg(
- cfg_parallel_size: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool, kwargs: dict | None
-):
- """
- Test that predict_noise_maybe_with_cfg produces identical results
- with and without CFG parallel.
-
- Args:
- cfg_parallel_size: Number of GPUs for CFG parallel
- dtype: Data type for computation
- batch_size: Batch size for testing
- cfg_normalize: Whether to normalize CFG output
- """
- available_gpus = current_omni_platform.get_device_count()
- if available_gpus < cfg_parallel_size:
- pytest.skip(f"Test requires {cfg_parallel_size} GPUs but only {available_gpus} available")
-
- test_config = {
- "batch_size": batch_size,
- "channels": 4,
- "height": 16,
- "width": 16,
- "hidden_dim": 128,
- "cfg_scale": 7.5,
- "cfg_normalize": cfg_normalize,
- "model_seed": 42, # Fixed seed for model initialization
- "input_seed": 123, # Fixed seed for input generation
- "kwargs": kwargs, # Additional kwargs to test passing through CFG parallel
- }
-
- mp_context = torch.multiprocessing.get_context("spawn")
-
- manager = mp_context.Manager()
- baseline_queue = manager.Queue()
- cfg_parallel_queue = manager.Queue()
-
- # Run baseline (sequential CFG) on single GPU
- torch.multiprocessing.spawn(
- _test_cfg_sequential_worker,
- args=(1, dtype, test_config, baseline_queue),
- nprocs=1,
- )
-
- # Run CFG parallel on multiple GPUs
- torch.multiprocessing.spawn(
- _test_cfg_parallel_worker,
- args=(cfg_parallel_size, cfg_parallel_size, dtype, test_config, cfg_parallel_queue),
- nprocs=cfg_parallel_size,
- )
-
- # Get results from queues
- baseline_output = baseline_queue.get()
- cfg_parallel_outputs = [cfg_parallel_queue.get() for _ in range(cfg_parallel_size)]
- cfg_parallel_outputs.sort(key=lambda item: item[0])
- cfg_parallel_output = cfg_parallel_outputs[0][1]
-
- for cfg_rank, rank_output in cfg_parallel_outputs[1:]:
- torch.testing.assert_close(
- rank_output,
- cfg_parallel_output,
- rtol=0,
- atol=0,
- msg=f"CFG parallel ranks produced different outputs (rank 0 vs rank {cfg_rank})",
- )
-
- # Verify shapes match
- assert baseline_output.shape == cfg_parallel_output.shape, (
- f"Shape mismatch: baseline {baseline_output.shape} vs CFG parallel {cfg_parallel_output.shape}"
- )
-
- # Verify numerical equivalence with appropriate tolerances
- if dtype == torch.float32:
- rtol, atol = 1e-5, 1e-5
- elif dtype == torch.bfloat16:
- rtol, atol = 1e-2, 1e-2
- else:
- rtol, atol = 1e-3, 1e-3
-
- torch.testing.assert_close(
- cfg_parallel_output,
- baseline_output,
- rtol=rtol,
- atol=atol,
- msg=(
- f"CFG parallel output differs from sequential CFG\n"
- f" dtype={dtype}, batch_size={batch_size}, cfg_normalize={cfg_normalize}\n"
- f" Max diff: {(cfg_parallel_output - baseline_output).abs().max().item():.6e}"
- ),
- )
-
- print(
- f"✓ Test passed: cfg_size={cfg_parallel_size}, dtype={dtype}, "
- f"batch_size={batch_size}, cfg_normalize={cfg_normalize}"
- )
-
-
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-def test_predict_noise_without_cfg(dtype: torch.dtype):
- """
- Test predict_noise_maybe_with_cfg when do_true_cfg=False.
-
- When CFG is disabled, only the positive branch should be computed.
- This test runs on a single GPU without distributed environment.
- """
- available_gpus = current_omni_platform.get_device_count()
- if available_gpus < 1:
- pytest.skip("Test requires at least 1 GPU")
-
- device = torch.device(f"{current_omni_platform.device_type}:0")
- current_omni_platform.set_device(device)
-
- # Create pipeline without distributed environment
- pipeline = TestCFGPipeline(in_channels=4, hidden_dim=128, seed=42)
- pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
- pipeline.transformer.eval()
-
- # Set seed for input generation
- torch.manual_seed(123)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(123)
- positive_input = torch.randn(1, 4, 16, 16, dtype=dtype, device=device)
-
- with torch.no_grad():
- noise_pred = pipeline.predict_noise_maybe_with_cfg(
- do_true_cfg=False, # No CFG
- true_cfg_scale=7.5,
- positive_kwargs={"x": positive_input},
- negative_kwargs=None,
- cfg_normalize=False,
- )
-
- # Should always return output when do_true_cfg=False
- assert noise_pred is not None
- assert noise_pred.shape == (1, 4, 16, 16)
-
- print(f"✓ Test passed: predict_noise without CFG (dtype={dtype})")
+ _init_model_weights(self.transformer, seed)
class MultiBranchTestPipeline(CFGParallelMixin):
- """Test pipeline with custom 3-branch combine logic (like OmniGen2)."""
-
def __init__(self, in_channels: int = 4, hidden_dim: int = 128, seed: int = 42):
- torch.manual_seed(seed)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(seed)
-
+ _set_random_seeds(seed)
self.transformer = SimpleTransformer(in_channels, hidden_dim)
-
- torch.manual_seed(seed)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(seed)
- for param in self.transformer.parameters():
- torch.nn.init.normal_(param, mean=0.0, std=0.02)
+ _init_model_weights(self.transformer, seed)
def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False):
- """N-branch combine with weighted sum for testing.
-
- - 2-branch: standard CFG formula (true_cfg_scale is float)
- - 3-branch: OmniGen2-style dual guidance scale (true_cfg_scale is dict)
- - 4-branch: DreamID-style weighted sum (true_cfg_scale is dict)
- """
if len(predictions) == 4:
text_scale = true_cfg_scale["text"]
image_scale = true_cfg_scale["image"]
@@ -486,6 +151,561 @@ class MultiBranchTestPipeline(CFGParallelMixin):
return combined
+class FakeCfgGroup:
+ def __init__(self, *, world_size: int, rank_in_group: int, rank_tensors: list[torch.Tensor]):
+ self.world_size = world_size
+ self.rank_in_group = rank_in_group
+ self._rank_tensors = rank_tensors
+
+ def all_gather(
+ self,
+ input_: torch.Tensor,
+ dim: int = 0,
+ separate_tensors: bool = False,
+ ) -> torch.Tensor | list[torch.Tensor]:
+ del dim, input_
+ assert separate_tensors
+ return [tensor.clone() for tensor in self._rank_tensors]
+
+
+class FakeMultiBranchCfgGroup:
+ def __init__(
+ self,
+ *,
+ world_size: int,
+ rank_in_group: int,
+ slots_per_rank: list[list[tuple[torch.Tensor, ...]]],
+ ):
+ self.world_size = world_size
+ self.rank_in_group = rank_in_group
+ self._slots_per_rank = slots_per_rank
+ # Production path re-runs predict_noise, so gathered tensors are value-equal
+ # but not identity-equal to the precomputed slot table. Index by call order
+ # instead: slot-major then element-major, matching _predict_multi_branch_parallel.
+ self._call_idx = 0
+
+ def all_gather(
+ self,
+ input_: torch.Tensor,
+ dim: int = 0,
+ separate_tensors: bool = False,
+ ) -> torch.Tensor | list[torch.Tensor]:
+ del dim, input_
+ assert separate_tensors
+ n_elems = len(self._slots_per_rank[0][0])
+ slot_idx, elem_idx = divmod(self._call_idx, n_elems)
+ self._call_idx += 1
+ assert slot_idx < len(self._slots_per_rank[self.rank_in_group]), (
+ f"all_gather call {self._call_idx} exceeds slot table "
+ f"(slot_idx={slot_idx}, n_slots={len(self._slots_per_rank[self.rank_in_group])})"
+ )
+ return [self._slots_per_rank[rank][slot_idx][elem_idx].clone() for rank in range(self.world_size)]
+
+
+def _patch_cfg_state(
+ monkeypatch: pytest.MonkeyPatch,
+ *,
+ world_size: int,
+ rank: int,
+ fake_group: FakeCfgGroup | FakeMultiBranchCfgGroup,
+) -> None:
+ # Patch both the source module and cfg_parallel's bound imports.
+ # cfg_parallel does `from ...parallel_state import get_cfg_group`, so
+ # patching parallel_state alone does not affect CFGParallelMixin.
+ from vllm_omni.diffusion.distributed import cfg_parallel as cfg_parallel_mod
+
+ for module in (parallel_state, cfg_parallel_mod):
+ monkeypatch.setattr(module, "get_classifier_free_guidance_world_size", lambda: world_size)
+ monkeypatch.setattr(module, "get_classifier_free_guidance_rank", lambda: rank)
+ monkeypatch.setattr(module, "get_cfg_group", lambda: fake_group)
+
+
+def _make_two_branch_inputs(
+ *,
+ batch_size: int,
+ channels: int,
+ height: int,
+ width: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ input_seed: int,
+) -> tuple[dict[str, Any], dict[str, Any]]:
+ _set_random_seeds(input_seed)
+ positive = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
+ _set_random_seeds(input_seed + 1)
+ negative = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
+ return {"x": positive}, {"x": negative}
+
+
+def _make_multi_branch_inputs(
+ *,
+ n_branches: int,
+ batch_size: int,
+ channels: int,
+ height: int,
+ width: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ input_seed: int,
+) -> list[dict[str, Any]]:
+ branches_kwargs: list[dict[str, Any]] = []
+ for branch_idx in range(n_branches):
+ _set_random_seeds(input_seed + branch_idx)
+ x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
+ branches_kwargs.append({"x": x})
+ return branches_kwargs
+
+
+def _multi_branch_cfg_scale(n_branches: int) -> float | dict[str, float]:
+ if n_branches == 2:
+ return 5.0
+ if n_branches == 3:
+ return {"text": 5.0, "image": 2.0}
+ return {"text": 5.0, "image": 2.0, "vid_ref": 1.5}
+
+
+def _build_multi_branch_slot_table(
+ *,
+ n_branches: int,
+ cfg_world_size: int,
+ branch_preds: list[torch.Tensor | tuple[torch.Tensor, ...]],
+) -> list[list[tuple[torch.Tensor, ...]]]:
+ assignments = _dispatch_branches(n_branches, cfg_world_size)
+ max_per_rank = max(len(branch_ids) for branch_ids in assignments)
+ ref_pred = _wrap(branch_preds[0])
+ slots_per_rank: list[list[tuple[torch.Tensor, ...]]] = []
+ for branch_ids in assignments:
+ slots = [_wrap(branch_preds[branch_id]) for branch_id in branch_ids]
+ while len(slots) < max_per_rank:
+ slots.append(tuple(torch.zeros_like(tensor) for tensor in ref_pred))
+ slots_per_rank.append(slots)
+ return slots_per_rank
+
+
+def _run_two_branch_cfg_parallel_mock(
+ pipeline: TestCFGPipeline,
+ monkeypatch: pytest.MonkeyPatch,
+ *,
+ cfg_world_size: int,
+ positive_kwargs: dict[str, Any],
+ negative_kwargs: dict[str, Any],
+ true_cfg_scale: float,
+ cfg_normalize: bool,
+ extra_kwargs: dict[str, Any] | None,
+) -> torch.Tensor:
+ with torch.no_grad():
+ pos_tensor = _wrap(pipeline.predict_noise(**positive_kwargs))[0]
+ neg_tensor = _wrap(pipeline.predict_noise(**negative_kwargs))[0]
+
+ outputs: list[torch.Tensor] = []
+ for cfg_rank in range(cfg_world_size):
+ fake_group = FakeCfgGroup(
+ world_size=cfg_world_size,
+ rank_in_group=cfg_rank,
+ rank_tensors=[pos_tensor, neg_tensor],
+ )
+ _patch_cfg_state(monkeypatch, world_size=cfg_world_size, rank=cfg_rank, fake_group=fake_group)
+ with torch.no_grad():
+ output = pipeline.predict_noise_maybe_with_cfg(
+ do_true_cfg=True,
+ true_cfg_scale=true_cfg_scale,
+ positive_kwargs=positive_kwargs,
+ negative_kwargs=negative_kwargs,
+ cfg_normalize=cfg_normalize,
+ kwargs=extra_kwargs,
+ )
+ assert output is not None
+ outputs.append(output)
+
+ for rank_output in outputs[1:]:
+ torch.testing.assert_close(
+ rank_output,
+ outputs[0],
+ rtol=0,
+ atol=0,
+ msg="CFG parallel ranks produced different outputs in mock path",
+ )
+ return outputs[0]
+
+
+def _run_multi_branch_cfg_parallel_mock(
+ pipeline: MultiBranchTestPipeline,
+ monkeypatch: pytest.MonkeyPatch,
+ *,
+ cfg_world_size: int,
+ branches_kwargs: list[dict[str, Any]],
+ true_cfg_scale: float | dict[str, float],
+ cfg_normalize: bool,
+) -> torch.Tensor:
+ with torch.no_grad():
+ branch_preds = [pipeline.predict_noise(**kwargs) for kwargs in branches_kwargs]
+ slots_per_rank = _build_multi_branch_slot_table(
+ n_branches=len(branches_kwargs),
+ cfg_world_size=cfg_world_size,
+ branch_preds=branch_preds,
+ )
+
+ outputs: list[torch.Tensor] = []
+ for cfg_rank in range(cfg_world_size):
+ fake_group = FakeMultiBranchCfgGroup(
+ world_size=cfg_world_size,
+ rank_in_group=cfg_rank,
+ slots_per_rank=slots_per_rank,
+ )
+ _patch_cfg_state(monkeypatch, world_size=cfg_world_size, rank=cfg_rank, fake_group=fake_group)
+ with torch.no_grad():
+ output = pipeline.predict_noise_with_multi_branch_cfg(
+ do_true_cfg=True,
+ true_cfg_scale=true_cfg_scale,
+ branches_kwargs=branches_kwargs,
+ cfg_normalize=cfg_normalize,
+ )
+ assert output is not None
+ outputs.append(output)
+
+ for rank_output in outputs[1:]:
+ torch.testing.assert_close(
+ rank_output,
+ outputs[0],
+ rtol=0,
+ atol=0,
+ msg="Multi-branch CFG parallel ranks produced different outputs in mock path",
+ )
+ return outputs[0]
+
+
+def _build_pipeline_on_device(
+ pipeline_cls: type,
+ *,
+ hidden_dim: int,
+ model_seed: int,
+ dtype: torch.dtype,
+ device: torch.device,
+):
+ pipeline = pipeline_cls(in_channels=4, hidden_dim=hidden_dim, seed=model_seed)
+ pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
+ pipeline.transformer.eval()
+ return pipeline
+
+
+# ---------------------------------------------------------------------------
+# CPU: mixin logic via mocked CFG groups (no GPU / no torch.distributed spawn)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("batch_size", [2])
+@pytest.mark.parametrize("cfg_normalize", [False, True])
+@pytest.mark.parametrize("extra_kwargs", [None, {"step_i": 1}])
+def test_predict_noise_with_cfg(
+ monkeypatch: pytest.MonkeyPatch,
+ batch_size: int,
+ cfg_normalize: bool,
+ extra_kwargs: dict[str, Any] | None,
+):
+ dtype = torch.float32
+ device = torch.device("cpu")
+ pipeline = _build_pipeline_on_device(TestCFGPipeline, hidden_dim=128, model_seed=42, dtype=dtype, device=device)
+ positive_kwargs, negative_kwargs = _make_two_branch_inputs(
+ batch_size=batch_size,
+ channels=4,
+ height=16,
+ width=16,
+ dtype=dtype,
+ device=device,
+ input_seed=123,
+ )
+
+ _patch_cfg_state(
+ monkeypatch,
+ world_size=1,
+ rank=0,
+ fake_group=FakeCfgGroup(world_size=1, rank_in_group=0, rank_tensors=[]),
+ )
+ with torch.no_grad():
+ baseline = pipeline.predict_noise_maybe_with_cfg(
+ do_true_cfg=True,
+ true_cfg_scale=7.5,
+ positive_kwargs=positive_kwargs,
+ negative_kwargs=negative_kwargs,
+ cfg_normalize=cfg_normalize,
+ kwargs=extra_kwargs,
+ )
+
+ parallel = _run_two_branch_cfg_parallel_mock(
+ pipeline,
+ monkeypatch,
+ cfg_world_size=2,
+ positive_kwargs=positive_kwargs,
+ negative_kwargs=negative_kwargs,
+ true_cfg_scale=7.5,
+ cfg_normalize=cfg_normalize,
+ extra_kwargs=extra_kwargs,
+ )
+ assert baseline is not None
+ _assert_cfg_outputs_close(
+ parallel,
+ baseline,
+ dtype=dtype,
+ msg=f"mock CFG parallel differs from sequential (cfg_normalize={cfg_normalize})",
+ )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_predict_noise_without_cfg():
+ dtype = torch.float32
+ device = torch.device("cpu")
+ pipeline = _build_pipeline_on_device(TestCFGPipeline, hidden_dim=128, model_seed=42, dtype=dtype, device=device)
+ positive_kwargs, _ = _make_two_branch_inputs(
+ batch_size=1,
+ channels=4,
+ height=16,
+ width=16,
+ dtype=dtype,
+ device=device,
+ input_seed=123,
+ )
+
+ with torch.no_grad():
+ noise_pred = pipeline.predict_noise_maybe_with_cfg(
+ do_true_cfg=False,
+ true_cfg_scale=7.5,
+ positive_kwargs=positive_kwargs,
+ negative_kwargs=None,
+ cfg_normalize=False,
+ )
+
+ assert noise_pred is not None
+ assert noise_pred.shape == (1, 4, 16, 16)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+ "cfg_parallel_size,n_branches",
+ [(2, 2), (2, 3), (3, 3), (2, 4)],
+)
+@pytest.mark.parametrize("batch_size", [2])
+@pytest.mark.parametrize("cfg_normalize", [False, True])
+def test_predict_noise_with_multi_branch_cfg(
+ monkeypatch: pytest.MonkeyPatch,
+ cfg_parallel_size: int,
+ n_branches: int,
+ batch_size: int,
+ cfg_normalize: bool,
+):
+ dtype = torch.float32
+ device = torch.device("cpu")
+ pipeline = _build_pipeline_on_device(
+ MultiBranchTestPipeline,
+ hidden_dim=128,
+ model_seed=42,
+ dtype=dtype,
+ device=device,
+ )
+ branches_kwargs = _make_multi_branch_inputs(
+ n_branches=n_branches,
+ batch_size=batch_size,
+ channels=4,
+ height=16,
+ width=16,
+ dtype=dtype,
+ device=device,
+ input_seed=123,
+ )
+ cfg_scale = _multi_branch_cfg_scale(n_branches)
+
+ _patch_cfg_state(
+ monkeypatch,
+ world_size=1,
+ rank=0,
+ fake_group=FakeCfgGroup(world_size=1, rank_in_group=0, rank_tensors=[]),
+ )
+ with torch.no_grad():
+ baseline = pipeline.predict_noise_with_multi_branch_cfg(
+ do_true_cfg=True,
+ true_cfg_scale=cfg_scale,
+ branches_kwargs=branches_kwargs,
+ cfg_normalize=cfg_normalize,
+ )
+
+ parallel = _run_multi_branch_cfg_parallel_mock(
+ pipeline,
+ monkeypatch,
+ cfg_world_size=cfg_parallel_size,
+ branches_kwargs=branches_kwargs,
+ true_cfg_scale=cfg_scale,
+ cfg_normalize=cfg_normalize,
+ )
+ assert baseline is not None
+ _assert_cfg_outputs_close(
+ parallel,
+ baseline,
+ dtype=dtype,
+ msg=(
+ f"mock multi-branch CFG parallel differs from sequential "
+ f"(n_branches={n_branches}, cfg_parallel_size={cfg_parallel_size})"
+ ),
+ )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_multi_branch_without_cfg():
+ dtype = torch.float32
+ device = torch.device("cpu")
+ pipeline = _build_pipeline_on_device(
+ MultiBranchTestPipeline,
+ hidden_dim=128,
+ model_seed=42,
+ dtype=dtype,
+ device=device,
+ )
+ branches_kwargs = _make_multi_branch_inputs(
+ n_branches=3,
+ batch_size=1,
+ channels=4,
+ height=16,
+ width=16,
+ dtype=dtype,
+ device=device,
+ input_seed=123,
+ )
+
+ with torch.no_grad():
+ noise_pred = pipeline.predict_noise_with_multi_branch_cfg(
+ do_true_cfg=False,
+ true_cfg_scale=5.0,
+ branches_kwargs=branches_kwargs,
+ cfg_normalize=False,
+ )
+
+ assert noise_pred is not None
+ assert noise_pred.shape == (1, 4, 16, 16)
+
+
+# ---------------------------------------------------------------------------
+# GPU nightly: same-seed sequential vs real multi-GPU CFG parallel parity
+# ---------------------------------------------------------------------------
+
+
+def _test_cfg_parallel_worker(
+ local_rank: int,
+ world_size: int,
+ cfg_parallel_size: int,
+ dtype: torch.dtype,
+ test_config: dict,
+ result_queue: torch.multiprocessing.Queue,
+):
+ device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+ current_omni_platform.set_device(device)
+ _update_environment_variables(
+ {
+ "RANK": str(local_rank),
+ "LOCAL_RANK": str(local_rank),
+ "WORLD_SIZE": str(world_size),
+ "MASTER_ADDR": "localhost",
+ "MASTER_PORT": "29502",
+ }
+ )
+
+ init_distributed_environment()
+ initialize_model_parallel(cfg_parallel_size=cfg_parallel_size)
+ assert get_classifier_free_guidance_world_size() == cfg_parallel_size
+
+ pipeline = _build_pipeline_on_device(
+ TestCFGPipeline,
+ hidden_dim=test_config["hidden_dim"],
+ model_seed=test_config["model_seed"],
+ dtype=dtype,
+ device=device,
+ )
+ positive_kwargs, negative_kwargs = _make_two_branch_inputs(
+ batch_size=test_config["batch_size"],
+ channels=test_config["channels"],
+ height=test_config["height"],
+ width=test_config["width"],
+ dtype=dtype,
+ device=device,
+ input_seed=test_config["input_seed"],
+ )
+
+ with torch.no_grad():
+ noise_pred = pipeline.predict_noise_maybe_with_cfg(
+ do_true_cfg=True,
+ true_cfg_scale=test_config["cfg_scale"],
+ positive_kwargs=positive_kwargs,
+ negative_kwargs=negative_kwargs,
+ cfg_normalize=test_config["cfg_normalize"],
+ kwargs=test_config["kwargs"],
+ )
+
+ assert noise_pred is not None
+ result_queue.put((get_classifier_free_guidance_rank(), noise_pred.cpu()))
+ destroy_distributed_env()
+
+
+def _test_cfg_sequential_worker(
+ local_rank: int,
+ world_size: int,
+ dtype: torch.dtype,
+ test_config: dict,
+ result_queue: torch.multiprocessing.Queue,
+):
+ device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+ current_omni_platform.set_device(device)
+ _update_environment_variables(
+ {
+ "RANK": str(local_rank),
+ "LOCAL_RANK": str(local_rank),
+ "WORLD_SIZE": str(world_size),
+ "MASTER_ADDR": "localhost",
+ "MASTER_PORT": "29503",
+ }
+ )
+
+ init_distributed_environment()
+ initialize_model_parallel(cfg_parallel_size=1)
+ assert get_classifier_free_guidance_world_size() == 1
+
+ pipeline = _build_pipeline_on_device(
+ TestCFGPipeline,
+ hidden_dim=test_config["hidden_dim"],
+ model_seed=test_config["model_seed"],
+ dtype=dtype,
+ device=device,
+ )
+ positive_kwargs, negative_kwargs = _make_two_branch_inputs(
+ batch_size=test_config["batch_size"],
+ channels=test_config["channels"],
+ height=test_config["height"],
+ width=test_config["width"],
+ dtype=dtype,
+ device=device,
+ input_seed=test_config["input_seed"],
+ )
+
+ with torch.no_grad():
+ noise_pred = pipeline.predict_noise_maybe_with_cfg(
+ do_true_cfg=True,
+ true_cfg_scale=test_config["cfg_scale"],
+ positive_kwargs=positive_kwargs,
+ negative_kwargs=negative_kwargs,
+ cfg_normalize=test_config["cfg_normalize"],
+ kwargs=test_config["kwargs"],
+ )
+
+ assert noise_pred is not None
+ result_queue.put(noise_pred.cpu())
+ destroy_distributed_env()
+
+
def _test_multi_branch_parallel_worker(
local_rank: int,
world_size: int,
@@ -494,11 +714,9 @@ def _test_multi_branch_parallel_worker(
test_config: dict,
result_queue: torch.multiprocessing.Queue,
):
- """Worker function for multi-branch CFG parallel test."""
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
-
- update_environment_variables(
+ _update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
@@ -510,33 +728,25 @@ def _test_multi_branch_parallel_worker(
init_distributed_environment()
initialize_model_parallel(cfg_parallel_size=cfg_parallel_size)
+ assert get_classifier_free_guidance_world_size() == cfg_parallel_size
- cfg_rank = get_classifier_free_guidance_rank()
- cfg_world_size = get_classifier_free_guidance_world_size()
- assert cfg_world_size == cfg_parallel_size
-
- pipeline = MultiBranchTestPipeline(
- in_channels=test_config["channels"],
+ pipeline = _build_pipeline_on_device(
+ MultiBranchTestPipeline,
hidden_dim=test_config["hidden_dim"],
- seed=test_config["model_seed"],
+ model_seed=test_config["model_seed"],
+ dtype=dtype,
+ device=device,
+ )
+ branches_kwargs = _make_multi_branch_inputs(
+ n_branches=test_config["n_branches"],
+ batch_size=test_config["batch_size"],
+ channels=test_config["channels"],
+ height=test_config["height"],
+ width=test_config["width"],
+ dtype=dtype,
+ device=device,
+ input_seed=test_config["input_seed"],
)
- pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
- pipeline.transformer.eval()
-
- n_branches = test_config["n_branches"]
- batch_size = test_config["batch_size"]
- channels = test_config["channels"]
- height = test_config["height"]
- width = test_config["width"]
-
- # Create N branch inputs with distinct seeds
- branches_kwargs = []
- for b in range(n_branches):
- torch.manual_seed(test_config["input_seed"] + b)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(test_config["input_seed"] + b)
- x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
- branches_kwargs.append({"x": x})
with torch.no_grad():
noise_pred = pipeline.predict_noise_with_multi_branch_cfg(
@@ -547,8 +757,7 @@ def _test_multi_branch_parallel_worker(
)
assert noise_pred is not None
- result_queue.put((cfg_rank, noise_pred.cpu()))
-
+ result_queue.put((get_classifier_free_guidance_rank(), noise_pred.cpu()))
destroy_distributed_env()
@@ -559,11 +768,9 @@ def _test_multi_branch_sequential_worker(
test_config: dict,
result_queue: torch.multiprocessing.Queue,
):
- """Worker function for sequential multi-branch CFG test (baseline)."""
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
-
- update_environment_variables(
+ _update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
@@ -575,31 +782,25 @@ def _test_multi_branch_sequential_worker(
init_distributed_environment()
initialize_model_parallel(cfg_parallel_size=1)
+ assert get_classifier_free_guidance_world_size() == 1
- cfg_world_size = get_classifier_free_guidance_world_size()
- assert cfg_world_size == 1
-
- pipeline = MultiBranchTestPipeline(
- in_channels=test_config["channels"],
+ pipeline = _build_pipeline_on_device(
+ MultiBranchTestPipeline,
hidden_dim=test_config["hidden_dim"],
- seed=test_config["model_seed"],
+ model_seed=test_config["model_seed"],
+ dtype=dtype,
+ device=device,
+ )
+ branches_kwargs = _make_multi_branch_inputs(
+ n_branches=test_config["n_branches"],
+ batch_size=test_config["batch_size"],
+ channels=test_config["channels"],
+ height=test_config["height"],
+ width=test_config["width"],
+ dtype=dtype,
+ device=device,
+ input_seed=test_config["input_seed"],
)
- pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
- pipeline.transformer.eval()
-
- n_branches = test_config["n_branches"]
- batch_size = test_config["batch_size"]
- channels = test_config["channels"]
- height = test_config["height"]
- width = test_config["width"]
-
- branches_kwargs = []
- for b in range(n_branches):
- torch.manual_seed(test_config["input_seed"] + b)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(test_config["input_seed"] + b)
- x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
- branches_kwargs.append({"x": x})
with torch.no_grad():
noise_pred = pipeline.predict_noise_with_multi_branch_cfg(
@@ -611,58 +812,86 @@ def _test_multi_branch_sequential_worker(
assert noise_pred is not None
result_queue.put(noise_pred.cpu())
-
destroy_distributed_env()
-@pytest.mark.parametrize(
- "cfg_parallel_size,n_branches",
- [
- (2, 2), # 2 branches on 2 GPUs: [[0],[1]]
- (2, 3), # 3 branches on 2 GPUs: [[0,2],[1]]
- (3, 3), # 3 branches on 3 GPUs: [[0],[1],[2]]
- (2, 4), # 4 branches on 2 GPUs: [[0,2],[1,3]]
- ],
-)
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("cfg_normalize", [False, True])
-def test_predict_noise_with_multi_branch_cfg(
+def _run_two_branch_gpu_parity(
+ *,
cfg_parallel_size: int,
- n_branches: int,
dtype: torch.dtype,
batch_size: int,
cfg_normalize: bool,
-):
- """
- Test that predict_noise_with_multi_branch_cfg produces identical results
- with and without CFG parallel for N-branch models.
-
- Args:
- cfg_parallel_size: Number of GPUs for CFG parallel
- n_branches: Number of CFG branches
- dtype: Data type for computation
- batch_size: Batch size for testing
- cfg_normalize: Whether to normalize CFG output
- """
- available_gpus = current_omni_platform.get_device_count()
- if available_gpus < cfg_parallel_size:
- pytest.skip(f"Test requires {cfg_parallel_size} GPUs but only {available_gpus} available")
-
- if n_branches == 2:
- cfg_scale = 5.0
- elif n_branches == 3:
- cfg_scale = {"text": 5.0, "image": 2.0}
- else:
- cfg_scale = {"text": 5.0, "image": 2.0, "vid_ref": 1.5}
-
+ extra_kwargs: dict[str, Any] | None,
+) -> None:
test_config = {
"batch_size": batch_size,
"channels": 4,
"height": 16,
"width": 16,
"hidden_dim": 128,
- "cfg_scale": cfg_scale,
+ "cfg_scale": 7.5,
+ "cfg_normalize": cfg_normalize,
+ "model_seed": 42,
+ "input_seed": 123,
+ "kwargs": extra_kwargs,
+ }
+
+ mp_context = torch.multiprocessing.get_context("spawn")
+ manager = mp_context.Manager()
+ baseline_queue = manager.Queue()
+ cfg_parallel_queue = manager.Queue()
+
+ torch.multiprocessing.spawn(
+ _test_cfg_sequential_worker,
+ args=(1, dtype, test_config, baseline_queue),
+ nprocs=1,
+ )
+ torch.multiprocessing.spawn(
+ _test_cfg_parallel_worker,
+ args=(cfg_parallel_size, cfg_parallel_size, dtype, test_config, cfg_parallel_queue),
+ nprocs=cfg_parallel_size,
+ )
+
+ baseline_output = baseline_queue.get()
+ cfg_parallel_outputs = [cfg_parallel_queue.get() for _ in range(cfg_parallel_size)]
+ cfg_parallel_outputs.sort(key=lambda item: item[0])
+ cfg_parallel_output = cfg_parallel_outputs[0][1]
+
+ for cfg_rank, rank_output in cfg_parallel_outputs[1:]:
+ torch.testing.assert_close(
+ rank_output,
+ cfg_parallel_output,
+ rtol=0,
+ atol=0,
+ msg=f"CFG parallel ranks differ (rank 0 vs rank {cfg_rank})",
+ )
+
+ _assert_cfg_outputs_close(
+ cfg_parallel_output,
+ baseline_output,
+ dtype=dtype,
+ msg=(
+ f"GPU CFG parallel output differs from sequential CFG "
+ f"(cfg_parallel_size={cfg_parallel_size}, cfg_normalize={cfg_normalize})"
+ ),
+ )
+
+
+def _run_multi_branch_gpu_parity(
+ *,
+ cfg_parallel_size: int,
+ n_branches: int,
+ dtype: torch.dtype,
+ batch_size: int,
+ cfg_normalize: bool,
+) -> None:
+ test_config = {
+ "batch_size": batch_size,
+ "channels": 4,
+ "height": 16,
+ "width": 16,
+ "hidden_dim": 128,
+ "cfg_scale": _multi_branch_cfg_scale(n_branches),
"cfg_normalize": cfg_normalize,
"model_seed": 42,
"input_seed": 123,
@@ -674,14 +903,11 @@ def test_predict_noise_with_multi_branch_cfg(
baseline_queue = manager.Queue()
cfg_parallel_queue = manager.Queue()
- # Run baseline (sequential, cfgp=1)
torch.multiprocessing.spawn(
_test_multi_branch_sequential_worker,
args=(1, dtype, test_config, baseline_queue),
nprocs=1,
)
-
- # Run CFG parallel
torch.multiprocessing.spawn(
_test_multi_branch_parallel_worker,
args=(cfg_parallel_size, cfg_parallel_size, dtype, test_config, cfg_parallel_queue),
@@ -693,7 +919,6 @@ def test_predict_noise_with_multi_branch_cfg(
cfg_parallel_outputs.sort(key=lambda item: item[0])
cfg_parallel_output = cfg_parallel_outputs[0][1]
- # All ranks should produce identical output
for cfg_rank, rank_output in cfg_parallel_outputs[1:]:
torch.testing.assert_close(
rank_output,
@@ -703,73 +928,62 @@ def test_predict_noise_with_multi_branch_cfg(
msg=f"Multi-branch CFG parallel ranks differ (rank 0 vs rank {cfg_rank})",
)
- assert baseline_output.shape == cfg_parallel_output.shape, (
- f"Shape mismatch: baseline {baseline_output.shape} vs CFG parallel {cfg_parallel_output.shape}"
- )
-
- if dtype == torch.float32:
- rtol, atol = 1e-5, 1e-5
- elif dtype == torch.bfloat16:
- rtol, atol = 1e-2, 1e-2
- else:
- rtol, atol = 1e-3, 1e-3
-
- torch.testing.assert_close(
+ _assert_cfg_outputs_close(
cfg_parallel_output,
baseline_output,
- rtol=rtol,
- atol=atol,
+ dtype=dtype,
msg=(
- f"Multi-branch CFG parallel output differs from sequential\n"
- f" n_branches={n_branches}, cfg_parallel_size={cfg_parallel_size}\n"
- f" dtype={dtype}, cfg_normalize={cfg_normalize}\n"
- f" Max diff: {(cfg_parallel_output - baseline_output).abs().max().item():.6e}"
+ f"GPU multi-branch CFG parallel differs from sequential "
+ f"(n_branches={n_branches}, cfg_parallel_size={cfg_parallel_size})"
),
)
- print(
- f"✓ Test passed: multi_branch n_branches={n_branches}, "
- f"cfg_size={cfg_parallel_size}, dtype={dtype}, cfg_normalize={cfg_normalize}"
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.parametrize("batch_size", [2])
+@pytest.mark.parametrize("cfg_normalize", [False, True])
+@pytest.mark.parametrize("extra_kwargs", [None, {"step_i": 1}])
+def test_predict_noise_with_cfg_parity(
+ batch_size: int,
+ cfg_normalize: bool,
+ extra_kwargs: dict[str, Any] | None,
+):
+ _run_two_branch_gpu_parity(
+ cfg_parallel_size=2,
+ dtype=torch.bfloat16,
+ batch_size=batch_size,
+ cfg_normalize=cfg_normalize,
+ extra_kwargs=extra_kwargs,
)
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-def test_multi_branch_without_cfg(dtype: torch.dtype):
- """
- Test predict_noise_with_multi_branch_cfg when do_true_cfg=False.
-
- When CFG is disabled, only the first branch (positive) should be computed.
- This test runs on a single GPU without distributed environment.
- """
- available_gpus = current_omni_platform.get_device_count()
- if available_gpus < 1:
- pytest.skip("Test requires at least 1 GPU")
-
- device = torch.device(f"{current_omni_platform.device_type}:0")
- current_omni_platform.set_device(device)
-
- pipeline = MultiBranchTestPipeline(in_channels=4, hidden_dim=128, seed=42)
- pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
- pipeline.transformer.eval()
-
- # Create 3 branch inputs (only first should be used)
- branches_kwargs = []
- for b in range(3):
- torch.manual_seed(123 + b)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(123 + b)
- x = torch.randn(1, 4, 16, 16, dtype=dtype, device=device)
- branches_kwargs.append({"x": x})
-
- with torch.no_grad():
- noise_pred = pipeline.predict_noise_with_multi_branch_cfg(
- do_true_cfg=False, # No CFG
- true_cfg_scale=5.0,
- branches_kwargs=branches_kwargs,
- cfg_normalize=False,
- )
-
- assert noise_pred is not None
- assert noise_pred.shape == (1, 4, 16, 16)
-
- print(f"✓ Test passed: multi_branch predict_noise without CFG (dtype={dtype})")
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+ "cfg_parallel_size,n_branches",
+ [
+ pytest.param(2, 2, marks=_L4_TWO_GPU),
+ pytest.param(2, 3, marks=_L4_TWO_GPU),
+ pytest.param(3, 3, marks=_L4_THREE_GPU),
+ pytest.param(2, 4, marks=_L4_TWO_GPU),
+ ],
+)
+@pytest.mark.parametrize("batch_size", [2])
+@pytest.mark.parametrize("cfg_normalize", [False, True])
+def test_predict_noise_with_multi_branch_cfg_parity(
+ cfg_parallel_size: int,
+ n_branches: int,
+ batch_size: int,
+ cfg_normalize: bool,
+):
+ _run_multi_branch_gpu_parity(
+ cfg_parallel_size=cfg_parallel_size,
+ n_branches=n_branches,
+ dtype=torch.bfloat16,
+ batch_size=batch_size,
+ cfg_normalize=cfg_normalize,
+ )
diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py
index 3e408c10a..647389048 100644
--- a/tests/diffusion/distributed/test_comm.py
+++ b/tests/diffusion/distributed/test_comm.py
@@ -1,12 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for SeqAllToAll4D and SeqAllToAll5D communication primitives."""
+"""Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives.
+
+CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same
+checks on real multi-GPU NCCL collectives.
+"""
+
+from __future__ import annotations
import os
+from typing import Literal
import pytest
import torch
+from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env,
@@ -16,56 +24,51 @@ from vllm_omni.diffusion.distributed.parallel_state import (
)
from vllm_omni.platforms import current_omni_platform
+_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
-def update_environment_variables(envs_dict: dict[str, str]):
- """Update multiple environment variables with logging."""
- for k, v in envs_dict.items():
- os.environ[k] = v
+DeviceKind = Literal["cpu", "cuda"]
-@pytest.mark.parametrize("world_size", [2, 4])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("seq_len_per_rank", [8])
-@pytest.mark.parametrize("num_heads", [8])
-@pytest.mark.parametrize("head_size", [32])
-@pytest.mark.parametrize("use_sync", [False, True])
-def test_4d_identity(
+def _update_environment_variables(envs_dict: dict[str, str]) -> None:
+ for key, value in envs_dict.items():
+ os.environ[key] = value
+
+
+def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
+ if device_kind == "cpu":
+ return torch.device("cpu")
+ return torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+
+
+def _init_worker(
+ local_rank: int,
world_size: int,
- dtype: torch.dtype,
- batch_size: int,
- seq_len_per_rank: int,
- num_heads: int,
- head_size: int,
- use_sync: bool,
-):
- """Test that two consecutive all-to-all operations return the original input."""
- # Skip if not enough GPUs available
- available_gpus = current_omni_platform.get_device_count()
- if available_gpus < world_size:
- pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
-
- # Ensure num_heads is divisible by world_size
- if num_heads % world_size != 0:
- pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
-
- # Run test with multiprocessing spawn
- torch.multiprocessing.spawn(
- _test_4d_identity_worker,
- args=(
- world_size,
- dtype,
- batch_size,
- seq_len_per_rank,
- num_heads,
- head_size,
- use_sync,
- ),
- nprocs=world_size,
+ master_port: int,
+ device_kind: DeviceKind,
+) -> None:
+ _update_environment_variables(
+ {
+ "RANK": str(local_rank),
+ "LOCAL_RANK": str(local_rank),
+ "WORLD_SIZE": str(world_size),
+ "MASTER_ADDR": "localhost",
+ "MASTER_PORT": str(master_port),
+ }
)
+ backend = "gloo" if device_kind == "cpu" else None
+ init_distributed_environment(backend=backend)
-def _test_4d_identity_worker(
+def _close_tolerance(dtype: torch.dtype) -> tuple[float, float]:
+ if dtype == torch.bfloat16:
+ return 1e-3, 1e-3
+ if dtype == torch.float16:
+ return 1e-3, 1e-3
+ return 1e-5, 1e-5
+
+
+def _run_4d_identity(
local_rank: int,
world_size: int,
dtype: torch.dtype,
@@ -74,29 +77,17 @@ def _test_4d_identity_worker(
num_heads: int,
head_size: int,
use_sync: bool,
-):
- """Worker function for test_4d_identity."""
- # Set device
- device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
- current_omni_platform.set_device(device)
+ device_kind: DeviceKind,
+ master_port: int,
+) -> None:
+ device = _worker_device(local_rank, device_kind)
+ if device_kind == "cuda":
+ current_omni_platform.set_device(device)
- # Set environment variables for distributed training
- update_environment_variables(
- {
- "RANK": str(local_rank),
- "LOCAL_RANK": str(local_rank),
- "WORLD_SIZE": str(world_size),
- "MASTER_ADDR": "localhost",
- "MASTER_PORT": "29500",
- }
- )
+ _init_worker(local_rank, world_size, master_port, device_kind)
+ initialize_model_parallel(ulysses_degree=world_size)
+ sp_group = get_sp_group().ulysses_group
- # Initialize distributed environment
- init_distributed_environment()
- initialize_model_parallel(ulysses_degree=world_size) # test ulysses sp by default
- sp_group = get_sp_group().ulysses_group # get ulysses sp group not ring sp group
-
- # Create input tensor: (bs, seqlen/P, hc, hs)
torch.manual_seed(42 + local_rank)
input_tensor = torch.randn(
batch_size,
@@ -106,100 +97,36 @@ def _test_4d_identity_worker(
dtype=dtype,
device=device,
)
-
- # Save original input for comparison
original_input = input_tensor.clone()
- # First all-to-all: (bs, seqlen/P, hc, hs) -> (bs, seqlen, hc/P, hs)
- intermediate = SeqAllToAll4D.apply(
- sp_group,
- input_tensor,
- 2, # scatter head dimension
- 1, # gather sequence dimension
- use_sync,
- )
-
- # Verify intermediate shape
- expected_shape = (
+ intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, use_sync)
+ expected_intermediate_shape = (
batch_size,
seq_len_per_rank * world_size,
num_heads // world_size,
head_size,
)
- assert intermediate.shape == expected_shape, (
- f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}"
+ assert intermediate.shape == expected_intermediate_shape, (
+ f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
)
- # Second all-to-all: (bs, seqlen, hc/P, hs) -> (bs, seqlen/P, hc, hs)
- output = SeqAllToAll4D.apply(
- sp_group,
- intermediate,
- 1, # scatter sequence dimension
- 2, # gather head dimension
- use_sync,
- )
-
- # Verify output shape matches input
+ output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, use_sync)
assert output.shape == original_input.shape, (
f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
)
- # Verify output matches original input
+ rtol, atol = _close_tolerance(dtype)
torch.testing.assert_close(
output,
original_input,
- rtol=1e-5,
- atol=1e-5,
+ rtol=rtol,
+ atol=atol,
msg="Output does not match original input after two all-to-all operations",
)
-
- # Cleanup distributed environment
destroy_distributed_env()
-@pytest.mark.parametrize("world_size", [2, 4])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("seq_len_per_rank", [8])
-@pytest.mark.parametrize("num_heads", [8])
-@pytest.mark.parametrize("head_size", [32])
-@pytest.mark.parametrize("use_sync", [False, True])
-def test_5d_identity(
- world_size: int,
- dtype: torch.dtype,
- batch_size: int,
- seq_len_per_rank: int,
- num_heads: int,
- head_size: int,
- use_sync: bool,
-):
- """Test that two consecutive all-to-all operations return the original input."""
- # Skip if not enough GPUs available
- available_gpus = current_omni_platform.get_device_count()
- if available_gpus < world_size:
- pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
-
- # Ensure num_heads is divisible by world_size
- if num_heads % world_size != 0:
- pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
-
- # Run test with multiprocessing spawn
- torch.multiprocessing.spawn(
- _test_5d_identity_worker,
- args=(
- world_size,
- dtype,
- batch_size,
- seq_len_per_rank,
- num_heads,
- head_size,
- use_sync,
- ),
- nprocs=world_size,
- )
-
-
-def _test_5d_identity_worker(
+def _run_5d_identity(
local_rank: int,
world_size: int,
dtype: torch.dtype,
@@ -208,199 +135,329 @@ def _test_5d_identity_worker(
num_heads: int,
head_size: int,
use_sync: bool,
-):
- """Worker function for test_5d_identity."""
- # Set device
- device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
- current_omni_platform.set_device(device)
+ device_kind: DeviceKind,
+ master_port: int,
+) -> None:
+ device = _worker_device(local_rank, device_kind)
+ if device_kind == "cuda":
+ current_omni_platform.set_device(device)
- # Set environment variables for distributed training
- update_environment_variables(
- {
- "RANK": str(local_rank),
- "LOCAL_RANK": str(local_rank),
- "WORLD_SIZE": str(world_size),
- "MASTER_ADDR": "localhost",
- "MASTER_PORT": "29500",
- }
- )
+ _init_worker(local_rank, world_size, master_port, device_kind)
+ initialize_model_parallel(ulysses_degree=world_size)
+ sp_group = get_sp_group().ulysses_group
- # Initialize distributed environment
- init_distributed_environment()
- initialize_model_parallel(ulysses_degree=world_size) # test ulysses sp by default
- sp_group = get_sp_group().ulysses_group # get ulysses sp group not ring sp group
-
- # Create input tensor: (bs, seqlen/P, 3, hc, hs)
- # The '3' dimension is for Q, K, V
torch.manual_seed(42 + local_rank)
input_tensor = torch.randn(
batch_size,
seq_len_per_rank,
- 3, # Q, K, V
+ 3,
num_heads,
head_size,
dtype=dtype,
device=device,
)
-
- # Save original input for comparison
original_input = input_tensor.clone()
- # First all-to-all: (bs, seqlen/P, 3, hc, hs) -> (bs, seqlen, 3, hc/P, hs)
- intermediate = SeqAllToAll5D.apply(
- sp_group,
- input_tensor,
- 3, # scatter head dimension
- 1, # gather sequence dimension
- use_sync,
- )
-
- # Verify intermediate shape
- expected_shape = (
+ intermediate = SeqAllToAll5D.apply(sp_group, input_tensor, 3, 1, use_sync)
+ expected_intermediate_shape = (
batch_size,
seq_len_per_rank * world_size,
3,
num_heads // world_size,
head_size,
)
- assert intermediate.shape == expected_shape, (
- f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}"
+ assert intermediate.shape == expected_intermediate_shape, (
+ f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
)
- # Second all-to-all: (bs, seqlen, 3, hc/P, hs) -> (bs, seqlen/P, 3, hc, hs)
- output = SeqAllToAll5D.apply(
- sp_group,
- intermediate,
- 1, # scatter sequence dimension
- 3, # gather head dimension
- use_sync,
- )
-
- # Verify output shape matches input
+ output = SeqAllToAll5D.apply(sp_group, intermediate, 1, 3, use_sync)
assert output.shape == original_input.shape, (
f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
)
- # Verify output matches original input
+ rtol, atol = _close_tolerance(dtype)
torch.testing.assert_close(
output,
original_input,
- rtol=1e-5,
- atol=1e-5,
+ rtol=rtol,
+ atol=atol,
msg="Output does not match original input after two all-to-all operations",
)
-
- # Cleanup distributed environment
destroy_distributed_env()
-@pytest.mark.parametrize("world_size", [2, 4])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("num_heads", [8])
-@pytest.mark.parametrize("head_size", [128])
-def test_ring_p2p(
- world_size: int,
- dtype: torch.dtype,
- batch_size: int,
- num_heads: int,
- head_size: int,
-):
- """Test Ring P2P communication (send_recv)."""
- # Skip if not enough GPUs available
- available_gpus = current_omni_platform.get_device_count()
- if available_gpus < world_size:
- pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
-
- torch.multiprocessing.spawn(
- _test_ring_p2p_worker,
- args=(world_size, dtype, batch_size, num_heads, head_size),
- nprocs=world_size,
- )
-
-
-def _test_ring_p2p_worker(
+def _run_ring_p2p(
local_rank: int,
world_size: int,
dtype: torch.dtype,
batch_size: int,
num_heads: int,
head_size: int,
-):
- """Worker for Ring P2P test."""
- import sys
+ device_kind: DeviceKind,
+ master_port: int,
+) -> None:
+ device = _worker_device(local_rank, device_kind)
+ if device_kind == "cuda":
+ current_omni_platform.set_device(device)
- # Set device
- device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
- current_omni_platform.set_device(device)
-
- # Set env vars
- # Use a different port to avoid conflict with other tests if run in parallel
- update_environment_variables(
- {
- "RANK": str(local_rank),
- "LOCAL_RANK": str(local_rank),
- "WORLD_SIZE": str(world_size),
- "MASTER_ADDR": "localhost",
- "MASTER_PORT": "29501",
- }
- )
-
- # Init distributed
try:
- init_distributed_environment()
- # Ring degree = world_size to test ring group
+ _init_worker(local_rank, world_size, master_port, device_kind)
initialize_model_parallel(ring_degree=world_size)
sp_group = get_sp_group()
-
- print(f"[Rank {local_rank}] Initialized. Ring group size: {sp_group.ring_group.size()}")
- sys.stdout.flush()
-
- # Create RingComm
comm = RingComm(sp_group.ring_group)
- # Create tensor: rank-specific data
- # (batch, num_heads, head_size)
- # Fill with rank value + 1 to avoid 0 and make verification easy
input_tensor = torch.full(
- (batch_size, num_heads, head_size), fill_value=float(local_rank + 1), dtype=dtype, device=device
+ (batch_size, num_heads, head_size),
+ fill_value=float(local_rank + 1),
+ dtype=dtype,
+ device=device,
)
-
- print(f"[Rank {local_rank}] Input sum: {input_tensor.sum().item()}")
- sys.stdout.flush()
-
- # Send input, receive from prev
- # RingComm.send_recv sends to next, receives from prev
- t0 = __import__("time").time()
recv_tensor = comm.send_recv(input_tensor)
comm.commit()
comm.wait()
- t1 = __import__("time").time()
- print(f"[Rank {local_rank}] Communication done in {t1 - t0:.4f}s")
-
- # Verify
- # Expected value: from (rank - 1) % world_size
prev_rank = (local_rank - 1 + world_size) % world_size
expected_value = float(prev_rank + 1)
-
- recv_sum = recv_tensor.sum().item()
- print(f"[Rank {local_rank}] Received sum: {recv_sum}, Expected value: {expected_value}")
- sys.stdout.flush()
-
expected_tensor = torch.full_like(recv_tensor, fill_value=expected_value)
-
- # Use a slightly loose tolerance for bfloat16
+ rtol, atol = _close_tolerance(dtype)
torch.testing.assert_close(
- recv_tensor, expected_tensor, rtol=1e-3, atol=1e-3, msg=f"[Rank {local_rank}] Data mismatch!"
+ recv_tensor,
+ expected_tensor,
+ rtol=rtol,
+ atol=atol,
+ msg=f"[Rank {local_rank}] Ring P2P data mismatch",
)
- print(f"[Rank {local_rank}] Verification PASSED")
-
- except Exception as e:
- print(f"[Rank {local_rank}] FAILED with error: {e}")
- import traceback
-
- traceback.print_exc()
- raise e
finally:
destroy_distributed_env()
+
+
+def _spawn_4d_identity(
+ *,
+ world_size: int,
+ dtype: torch.dtype,
+ batch_size: int,
+ seq_len_per_rank: int,
+ num_heads: int,
+ head_size: int,
+ use_sync: bool,
+ device_kind: DeviceKind,
+ master_port: int,
+) -> None:
+ torch.multiprocessing.spawn(
+ _run_4d_identity,
+ args=(
+ world_size,
+ dtype,
+ batch_size,
+ seq_len_per_rank,
+ num_heads,
+ head_size,
+ use_sync,
+ device_kind,
+ master_port,
+ ),
+ nprocs=world_size,
+ )
+
+
+def _spawn_5d_identity(
+ *,
+ world_size: int,
+ dtype: torch.dtype,
+ batch_size: int,
+ seq_len_per_rank: int,
+ num_heads: int,
+ head_size: int,
+ use_sync: bool,
+ device_kind: DeviceKind,
+ master_port: int,
+) -> None:
+ torch.multiprocessing.spawn(
+ _run_5d_identity,
+ args=(
+ world_size,
+ dtype,
+ batch_size,
+ seq_len_per_rank,
+ num_heads,
+ head_size,
+ use_sync,
+ device_kind,
+ master_port,
+ ),
+ nprocs=world_size,
+ )
+
+
+def _spawn_ring_p2p(
+ *,
+ world_size: int,
+ dtype: torch.dtype,
+ batch_size: int,
+ num_heads: int,
+ head_size: int,
+ device_kind: DeviceKind,
+ master_port: int,
+) -> None:
+ torch.multiprocessing.spawn(
+ _run_ring_p2p,
+ args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
+ nprocs=world_size,
+ )
+
+
+def _require_heads_divisible(num_heads: int, world_size: int) -> None:
+ if num_heads % world_size != 0:
+ pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
+
+
+def _require_gpus(world_size: int) -> None:
+ available_gpus = current_omni_platform.get_device_count()
+ if available_gpus < world_size:
+ pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
+
+# ---------------------------------------------------------------------------
+# CPU: gloo collectives on CPU tensors (no GPU)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_4d_identity(world_size: int, use_sync: bool):
+ _require_heads_divisible(8, world_size)
+ _spawn_4d_identity(
+ world_size=world_size,
+ dtype=torch.float32,
+ batch_size=2,
+ seq_len_per_rank=8,
+ num_heads=8,
+ head_size=32,
+ use_sync=use_sync,
+ device_kind="cpu",
+ master_port=29610,
+ )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_5d_identity(world_size: int, use_sync: bool):
+ _require_heads_divisible(8, world_size)
+ _spawn_5d_identity(
+ world_size=world_size,
+ dtype=torch.float32,
+ batch_size=2,
+ seq_len_per_rank=8,
+ num_heads=8,
+ head_size=32,
+ use_sync=use_sync,
+ device_kind="cpu",
+ master_port=29611,
+ )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_ring_p2p(world_size: int):
+ _spawn_ring_p2p(
+ world_size=world_size,
+ dtype=torch.float32,
+ batch_size=2,
+ num_heads=8,
+ head_size=128,
+ device_kind="cpu",
+ master_port=29612,
+ )
+
+
+# ---------------------------------------------------------------------------
+# Nightly: same checks on real multi-GPU NCCL collectives
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+ "world_size",
+ [
+ pytest.param(2, marks=_L4_TWO_GPU),
+ pytest.param(4, marks=_L4_FOUR_GPU),
+ ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_4d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
+ _require_gpus(world_size)
+ _require_heads_divisible(8, world_size)
+ _spawn_4d_identity(
+ world_size=world_size,
+ dtype=dtype,
+ batch_size=2,
+ seq_len_per_rank=8,
+ num_heads=8,
+ head_size=32,
+ use_sync=use_sync,
+ device_kind="cuda",
+ master_port=29500,
+ )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+ "world_size",
+ [
+ pytest.param(2, marks=_L4_TWO_GPU),
+ pytest.param(4, marks=_L4_FOUR_GPU),
+ ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_5d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
+ _require_gpus(world_size)
+ _require_heads_divisible(8, world_size)
+ _spawn_5d_identity(
+ world_size=world_size,
+ dtype=dtype,
+ batch_size=2,
+ seq_len_per_rank=8,
+ num_heads=8,
+ head_size=32,
+ use_sync=use_sync,
+ device_kind="cuda",
+ master_port=29502,
+ )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+ "world_size",
+ [
+ pytest.param(2, marks=_L4_TWO_GPU),
+ pytest.param(4, marks=_L4_FOUR_GPU),
+ ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
+ _require_gpus(world_size)
+ _spawn_ring_p2p(
+ world_size=world_size,
+ dtype=dtype,
+ batch_size=2,
+ num_heads=8,
+ head_size=128,
+ device_kind="cuda",
+ master_port=29501,
+ )
diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py
index b76bf8c89..0b48ed266 100644
--- a/tests/diffusion/distributed/test_pipeline_parallel.py
+++ b/tests/diffusion/distributed/test_pipeline_parallel.py
@@ -1,8 +1,17 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for PipelineParallelMixin and related PP communication helpers.
+
+CPU unit tests exercise mixin logic without a process group. Distributed CPU
+tests use gloo on CPU tensors; nightly parity tests run the same checks on
+real multi-GPU NCCL collectives.
+"""
+
+from __future__ import annotations
import os
import socket
+from typing import Literal
import pytest
import torch
@@ -11,6 +20,7 @@ from vllm.sequence import IntermediateTensors
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env,
@@ -22,7 +32,12 @@ from vllm_omni.diffusion.distributed.parallel_state import (
from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
from vllm_omni.platforms import current_omni_platform
-pytestmark = [pytest.mark.parallel]
+_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+
+DeviceKind = Literal["cpu", "cuda"]
+_UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _find_free_port() -> str:
@@ -36,6 +51,22 @@ def update_environment_variables(envs_dict: dict[str, str]) -> None:
os.environ[k] = v
+def _cleanup_distributed() -> None:
+ """Destroy omni distributed state and clear test-process leftovers.
+
+ ``destroy_distributed_env`` leaves ``vllm.distributed.parallel_state._PP``
+ aliased to the destroyed group. Parent-process baselines must clear that
+ dangling reference (and RANK/MASTER_* env) so later unit tests that call
+ ``vllm.distributed.initialize_model_parallel`` do not see stale state.
+ """
+ destroy_distributed_env()
+ import vllm.distributed.parallel_state as vllm_parallel_state
+
+ vllm_parallel_state._PP = None
+ for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
+ os.environ.pop(key, None)
+
+
# ---------------------------------------------------------------------------
# Shared stubs used by both unit and distributed tests
# ---------------------------------------------------------------------------
@@ -143,7 +174,7 @@ class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
class TestAsyncLatents:
"""Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
- pytestmark = [pytest.mark.cpu]
+ pytestmark = _UNIT_MARKS
def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
@@ -214,7 +245,7 @@ class TestAsyncLatents:
class TestSyncPPSend:
"""Verifies PipelineParallelMixin's internal PP-send flush."""
- pytestmark = [pytest.mark.cpu]
+ pytestmark = _UNIT_MARKS
@staticmethod
def _make_pipeline() -> PipelineParallelMixin:
@@ -247,7 +278,7 @@ class TestSyncPPSend:
class TestDiffuseWrapper:
"""Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
- pytestmark = [pytest.mark.cpu]
+ pytestmark = _UNIT_MARKS
def test_diffuse_flushes_pending_sends_on_success(self):
work = FakeWork()
@@ -289,7 +320,7 @@ class TestDiffuseWrapper:
class TestVaeDecodeGuard:
- pytestmark = [pytest.mark.cpu]
+ pytestmark = _UNIT_MARKS
@staticmethod
def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
@@ -351,6 +382,8 @@ class TestVaeDecodeGuard:
assert pipeline.vae.decode.__doc__ == "Original decode docstring."
+@pytest.mark.core_model
+@pytest.mark.diffusion
@pytest.mark.cpu
def test_pipeline_parallel_requires_cfg_mixin():
with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
@@ -359,6 +392,8 @@ def test_pipeline_parallel_requires_cfg_mixin():
pass
+@pytest.mark.core_model
+@pytest.mark.diffusion
@pytest.mark.cpu
def test_pipeline_parallel_requires_mro_before_cfg_mixin():
with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
@@ -372,10 +407,22 @@ def test_pipeline_parallel_requires_mro_before_cfg_mixin():
# ---------------------------------------------------------------------------
-def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
+def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
+ if device_kind == "cpu":
+ return torch.device("cpu")
+ return torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+
+
+def init_dist(
+ local_rank: int,
+ world_size: int,
+ master_port: str,
+ device_kind: DeviceKind,
+) -> torch.device:
"""Initialise the distributed environment for a spawned worker."""
- device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
- current_omni_platform.set_device(device)
+ device = _worker_device(local_rank, device_kind)
+ if device_kind == "cuda":
+ current_omni_platform.set_device(device)
update_environment_variables(
{
"RANK": str(local_rank),
@@ -385,13 +432,31 @@ def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.devic
"MASTER_PORT": master_port,
}
)
- init_distributed_environment()
+ backend = "gloo" if device_kind == "cpu" else None
+ init_distributed_environment(backend=backend)
return device
+def _mp_backend(device_kind: DeviceKind) -> str | None:
+ """Backend for model-parallel process groups.
+
+ ``GroupCoordinator.all_gather`` always uses ``device_group``. CPU tensors
+ therefore require gloo device groups; otherwise a CUDA host's default NCCL
+ backend breaks PP+CFG CPU tests at CFG all_gather (pp2-cfg2).
+ isend/irecv of CPU tensors can already fall back to ``cpu_group``.
+ """
+ return "gloo" if device_kind == "cpu" else None
+
+
+def _seed_device(input_seed: int, device: torch.device) -> None:
+ torch.manual_seed(input_seed)
+ if device.type == "cuda":
+ torch.cuda.manual_seed_all(input_seed)
+
+
def make_pipeline_and_inputs(
test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
-) -> tuple["MockPipelineParallel", dict, dict | None]:
+) -> tuple[MockPipelineParallel, dict, dict | None]:
"""Create a MockPipelineParallel and seeded inputs from a test_config dict.
Must be called after ``initialize_model_parallel`` so that ``make_layers``
@@ -408,16 +473,12 @@ def make_pipeline_and_inputs(
dtype=dtype,
)
- torch.manual_seed(test_config["input_seed"])
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(test_config["input_seed"])
+ _seed_device(test_config["input_seed"], device)
pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
negative_kwargs = None
if do_true_cfg:
- torch.manual_seed(test_config["input_seed"] + 1)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
+ _seed_device(test_config["input_seed"] + 1, device)
neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
negative_kwargs = {"x": neg_x}
@@ -425,19 +486,23 @@ def make_pipeline_and_inputs(
# ---------------------------------------------------------------------------
-# 3. isend_tensor_dict / irecv_tensor_dict (2 GPUs)
+# 3. isend_tensor_dict / irecv_tensor_dict
# ---------------------------------------------------------------------------
-def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
- device = init_dist(local_rank, world_size, master_port)
- initialize_model_parallel(pipeline_parallel_size=world_size)
+def isend_irecv_worker(
+ local_rank: int,
+ world_size: int,
+ master_port: str,
+ device_kind: DeviceKind,
+ result_queue,
+):
+ device = init_dist(local_rank, world_size, master_port, device_kind)
+ initialize_model_parallel(pipeline_parallel_size=world_size, backend=_mp_backend(device_kind))
pp_group = get_pp_group()
if pp_group.is_first_rank:
- torch.manual_seed(77)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(77)
+ _seed_device(77, device)
tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
handles = pp_group.isend_tensor_dict({"t": tensor})
for h in handles:
@@ -449,27 +514,42 @@ def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, resul
if torch.distributed.is_initialized():
torch.distributed.barrier()
- destroy_distributed_env()
+ _cleanup_distributed()
-@pytest.mark.gpu
-@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
-@pytest.mark.parametrize("pp_size", [2])
-def test_isend_irecv_tensor_dict(pp_size: int):
- """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None:
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
q = manager.Queue()
-
- port = _find_free_port()
- torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)
-
+ torch.multiprocessing.spawn(
+ isend_irecv_worker,
+ args=(pp_size, master_port, device_kind, q),
+ nprocs=pp_size,
+ )
results = {label: tensor for label, tensor in [q.get(), q.get()]}
torch.testing.assert_close(
results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
)
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("pp_size", [2])
+def test_isend_irecv_tensor_dict(pp_size: int):
+ """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+ _run_isend_irecv(pp_size, device_kind="cpu", master_port=_find_free_port())
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize("pp_size", [pytest.param(2, marks=_L4_TWO_GPU)])
+def test_isend_irecv_tensor_dict_parity(pp_size: int):
+ """Nightly: isend/irecv on real multi-GPU NCCL."""
+ _run_isend_irecv(pp_size, device_kind="cuda", master_port=_find_free_port())
+
+
# ---------------------------------------------------------------------------
# 4. predict_noise_maybe_with_cfg
# ---------------------------------------------------------------------------
@@ -477,12 +557,13 @@ def test_isend_irecv_tensor_dict(pp_size: int):
_baseline_cache: dict[tuple, torch.Tensor] = {}
-def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
- """Compute expected single-GPU output using the same MockPipelineParallel.
-
- Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
- Results are cached so identical configs are only computed once.
- """
+def compute_single_gpu_baseline(
+ test_config: dict,
+ dtype: torch.dtype,
+ do_true_cfg: bool,
+ device_kind: DeviceKind,
+) -> torch.Tensor:
+ """Compute expected single-process output using the same MockPipelineParallel."""
key = (
test_config["num_layers"],
test_config["dim"],
@@ -492,27 +573,29 @@ def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_c
test_config["cfg_scale"],
dtype,
do_true_cfg,
+ device_kind,
)
if key in _baseline_cache:
return _baseline_cache[key]
- device = init_dist(0, 1, _find_free_port())
- initialize_model_parallel(pipeline_parallel_size=1)
+ device = init_dist(0, 1, _find_free_port(), device_kind)
+ try:
+ initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
- pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
- test_config, dtype, device, do_true_cfg=do_true_cfg
- )
-
- with torch.inference_mode():
- noise_pred = pipeline.predict_noise_maybe_with_cfg(
- do_true_cfg=do_true_cfg,
- true_cfg_scale=test_config["cfg_scale"],
- positive_kwargs=positive_kwargs,
- negative_kwargs=negative_kwargs,
- cfg_normalize=False,
+ pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+ test_config, dtype, device, do_true_cfg=do_true_cfg
)
- destroy_distributed_env()
+ with torch.inference_mode():
+ noise_pred = pipeline.predict_noise_maybe_with_cfg(
+ do_true_cfg=do_true_cfg,
+ true_cfg_scale=test_config["cfg_scale"],
+ positive_kwargs=positive_kwargs,
+ negative_kwargs=negative_kwargs,
+ cfg_normalize=False,
+ )
+ finally:
+ _cleanup_distributed()
_baseline_cache[key] = noise_pred.cpu()
return _baseline_cache[key]
@@ -527,11 +610,16 @@ def predict_noise_worker(
do_true_cfg: bool,
dtype: torch.dtype,
test_config: dict,
+ device_kind: DeviceKind,
result_queue,
):
"""Generic predict-noise worker parameterized by PP and CFG topology."""
- device = init_dist(local_rank, world_size, master_port)
- initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+ device = init_dist(local_rank, world_size, master_port, device_kind)
+ initialize_model_parallel(
+ pipeline_parallel_size=pp_size,
+ cfg_parallel_size=cfg_size,
+ backend=_mp_backend(device_kind),
+ )
pp_group = get_pp_group()
cfg_rank = get_classifier_free_guidance_rank()
@@ -548,8 +636,6 @@ def predict_noise_worker(
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
- # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
- # Flush the non-last PP rank's async send before barrier / process teardown.
pipeline._sync_pp_send()
if pp_group.is_last_rank:
@@ -561,77 +647,22 @@ def predict_noise_worker(
if torch.distributed.is_initialized():
torch.distributed.barrier()
- destroy_distributed_env()
+ _cleanup_distributed()
-@pytest.mark.gpu
-@pytest.mark.parametrize(
- "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
- [
- pytest.param(
- 2,
- 1,
- False,
- torch.float32,
- 4,
- 100,
- 1e-5,
- 1e-5,
- marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
- id="pp2-no_cfg-float32",
- ),
- pytest.param(
- 2,
- 1,
- False,
- torch.bfloat16,
- 4,
- 100,
- 1e-2,
- 1e-2,
- marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
- id="pp2-no_cfg-bfloat16",
- ),
- pytest.param(
- 2,
- 1,
- True,
- torch.bfloat16,
- 4,
- 100,
- 1e-2,
- 1e-2,
- marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
- id="pp2-seq_cfg-bfloat16",
- ),
- pytest.param(
- 2,
- 2,
- True,
- torch.bfloat16,
- 4,
- 100,
- 1e-2,
- 1e-2,
- marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
- id="pp2-cfg2-bfloat16",
- ),
- pytest.param(
- 3,
- 1,
- False,
- torch.bfloat16,
- 6,
- 100,
- 1e-2,
- 1e-2,
- marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
- id="pp3-no_cfg-bfloat16",
- ),
- ],
-)
-def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
- """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
+def _run_predict_noise(
+ *,
+ pp_size: int,
+ cfg_size: int,
+ do_true_cfg: bool,
+ dtype: torch.dtype,
+ num_layers: int,
+ input_seed: int,
+ rtol: float,
+ atol: float,
+ device_kind: DeviceKind,
+ master_port: str,
+) -> None:
test_config = {
"num_layers": num_layers,
"dim": 64,
@@ -641,22 +672,20 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_
"input_seed": input_seed,
}
- baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
+ baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
pp_q = manager.Queue()
world_size = pp_size * cfg_size
- port = _find_free_port()
torch.multiprocessing.spawn(
predict_noise_worker,
- args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
+ args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
nprocs=world_size,
)
pp_out = pp_q.get()
-
assert baseline_out.shape == pp_out.shape
torch.testing.assert_close(
pp_out,
@@ -667,35 +696,157 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_
)
+_CPU_PREDICT_NOISE_CASES = [
+ pytest.param(2, 1, False, 4, 100, 1e-5, 1e-5, id="pp2-no_cfg"),
+ pytest.param(2, 1, True, 4, 100, 1e-5, 1e-5, id="pp2-seq_cfg"),
+ pytest.param(2, 2, True, 4, 100, 1e-5, 1e-5, id="pp2-cfg2"),
+ pytest.param(3, 1, False, 6, 100, 1e-5, 1e-5, id="pp3-no_cfg"),
+]
+
+_GPU_PREDICT_NOISE_CASES = [
+ pytest.param(
+ 2,
+ 1,
+ False,
+ torch.float32,
+ 4,
+ 100,
+ 1e-5,
+ 1e-5,
+ marks=_L4_TWO_GPU,
+ id="pp2-no_cfg-float32",
+ ),
+ pytest.param(
+ 2,
+ 1,
+ False,
+ torch.bfloat16,
+ 4,
+ 100,
+ 1e-2,
+ 1e-2,
+ marks=_L4_TWO_GPU,
+ id="pp2-no_cfg-bfloat16",
+ ),
+ pytest.param(
+ 2,
+ 1,
+ True,
+ torch.bfloat16,
+ 4,
+ 100,
+ 1e-2,
+ 1e-2,
+ marks=_L4_TWO_GPU,
+ id="pp2-seq_cfg-bfloat16",
+ ),
+ pytest.param(
+ 2,
+ 2,
+ True,
+ torch.bfloat16,
+ 4,
+ 100,
+ 1e-2,
+ 1e-2,
+ marks=_L4_FOUR_GPU,
+ id="pp2-cfg2-bfloat16",
+ ),
+ pytest.param(
+ 3,
+ 1,
+ False,
+ torch.bfloat16,
+ 6,
+ 100,
+ 1e-2,
+ 1e-2,
+ marks=_L4_THREE_GPU,
+ id="pp3-no_cfg-bfloat16",
+ ),
+]
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+ "pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol",
+ _CPU_PREDICT_NOISE_CASES,
+)
+def test_predict_noise(pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol):
+ """predict_noise_maybe_with_cfg output matches baseline across PP / CFG topologies."""
+ _run_predict_noise(
+ pp_size=pp_size,
+ cfg_size=cfg_size,
+ do_true_cfg=do_true_cfg,
+ dtype=torch.float32,
+ num_layers=num_layers,
+ input_seed=input_seed,
+ rtol=rtol,
+ atol=atol,
+ device_kind="cpu",
+ master_port=_find_free_port(),
+ )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+ "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
+ _GPU_PREDICT_NOISE_CASES,
+)
+def test_predict_noise_parity(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
+ """Nightly: predict_noise parity on real multi-GPU NCCL."""
+ _run_predict_noise(
+ pp_size=pp_size,
+ cfg_size=cfg_size,
+ do_true_cfg=do_true_cfg,
+ dtype=dtype,
+ num_layers=num_layers,
+ input_seed=input_seed,
+ rtol=rtol,
+ atol=atol,
+ device_kind="cuda",
+ master_port=_find_free_port(),
+ )
+
+
# ---------------------------------------------------------------------------
# 5. scheduler_step_maybe_with_cfg
# ---------------------------------------------------------------------------
-def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
- """Single-GPU reference: predict_noise + scheduler_step."""
- device = init_dist(0, 1, _find_free_port())
- initialize_model_parallel(pipeline_parallel_size=1)
+def compute_scheduler_step_baseline(
+ test_config: dict,
+ do_true_cfg: bool,
+ device_kind: DeviceKind,
+) -> torch.Tensor:
+ """Single-process reference: predict_noise + scheduler_step."""
+ device = init_dist(0, 1, _find_free_port(), device_kind)
+ try:
+ initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
- pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
- test_config, torch.float32, device, do_true_cfg=do_true_cfg
- )
- latents = positive_kwargs["x"]
- t = torch.tensor(500, device=device)
-
- with torch.inference_mode():
- noise_pred = pipeline.predict_noise_maybe_with_cfg(
- do_true_cfg=do_true_cfg,
- true_cfg_scale=test_config["cfg_scale"],
- positive_kwargs=positive_kwargs,
- negative_kwargs=negative_kwargs,
- cfg_normalize=False,
- )
- result = pipeline.scheduler_step_maybe_with_cfg(
- noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+ pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+ test_config, torch.float32, device, do_true_cfg=do_true_cfg
)
+ latents = positive_kwargs["x"]
+ t = torch.tensor(500, device=device)
- destroy_distributed_env()
+ with torch.inference_mode():
+ noise_pred = pipeline.predict_noise_maybe_with_cfg(
+ do_true_cfg=do_true_cfg,
+ true_cfg_scale=test_config["cfg_scale"],
+ positive_kwargs=positive_kwargs,
+ negative_kwargs=negative_kwargs,
+ cfg_normalize=False,
+ )
+ result = pipeline.scheduler_step_maybe_with_cfg(
+ noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+ )
+ finally:
+ _cleanup_distributed()
return result.cpu()
@@ -707,10 +858,15 @@ def scheduler_step_worker(
cfg_size: int,
do_true_cfg: bool,
test_config: dict,
+ device_kind: DeviceKind,
result_queue,
):
- device = init_dist(local_rank, world_size, master_port)
- initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+ device = init_dist(local_rank, world_size, master_port, device_kind)
+ initialize_model_parallel(
+ pipeline_parallel_size=pp_size,
+ cfg_parallel_size=cfg_size,
+ backend=_mp_backend(device_kind),
+ )
pp_group = get_pp_group()
cfg_rank = get_classifier_free_guidance_rank()
@@ -732,43 +888,25 @@ def scheduler_step_worker(
latents = pipeline.scheduler_step_maybe_with_cfg(
noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
)
- # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
- # Flush the last PP rank's async latent send before barrier / process teardown.
pipeline._sync_pp_send()
if pp_group.is_first_rank and cfg_rank == 0:
- resolved = latents.contiguous()
- result_queue.put(resolved.cpu())
+ result_queue.put(latents.contiguous().cpu())
if torch.distributed.is_initialized():
torch.distributed.barrier()
- destroy_distributed_env()
+ _cleanup_distributed()
-@pytest.mark.gpu
-@pytest.mark.parametrize(
- "pp_size, cfg_size, do_true_cfg, input_seed",
- [
- pytest.param(
- 2,
- 1,
- False,
- 300,
- marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
- id="pp2-no_cfg",
- ),
- pytest.param(
- 2,
- 2,
- True,
- 600,
- marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
- id="pp2-cfg2-true_cfg",
- ),
- ],
-)
-def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
- """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
+def _run_scheduler_step(
+ *,
+ pp_size: int,
+ cfg_size: int,
+ do_true_cfg: bool,
+ input_seed: int,
+ device_kind: DeviceKind,
+ master_port: str,
+) -> None:
test_config = {
"num_layers": 4,
"dim": 64,
@@ -778,17 +916,16 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
"input_seed": input_seed,
}
- baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
+ baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
q = manager.Queue()
- port = _find_free_port()
world_size = pp_size * cfg_size
torch.multiprocessing.spawn(
scheduler_step_worker,
- args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
+ args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
nprocs=world_size,
)
@@ -800,3 +937,47 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
atol=0,
msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+ "pp_size, cfg_size, do_true_cfg, input_seed",
+ [
+ pytest.param(2, 1, False, 300, id="pp2-no_cfg"),
+ pytest.param(2, 2, True, 600, id="pp2-cfg2-true_cfg"),
+ ],
+)
+def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
+ """Rank 0 latents after scheduler_step match baseline across PP / CFG topologies."""
+ _run_scheduler_step(
+ pp_size=pp_size,
+ cfg_size=cfg_size,
+ do_true_cfg=do_true_cfg,
+ input_seed=input_seed,
+ device_kind="cpu",
+ master_port=_find_free_port(),
+ )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+ "pp_size, cfg_size, do_true_cfg, input_seed",
+ [
+ pytest.param(2, 1, False, 300, marks=_L4_TWO_GPU, id="pp2-no_cfg"),
+ pytest.param(2, 2, True, 600, marks=_L4_FOUR_GPU, id="pp2-cfg2-true_cfg"),
+ ],
+)
+def test_scheduler_step_parity(pp_size, cfg_size, do_true_cfg, input_seed):
+ """Nightly: scheduler_step parity on real multi-GPU NCCL."""
+ _run_scheduler_step(
+ pp_size=pp_size,
+ cfg_size=cfg_size,
+ do_true_cfg=do_true_cfg,
+ input_seed=input_seed,
+ device_kind="cuda",
+ master_port=_find_free_port(),
+ )
diff --git a/tests/diffusion/distributed/test_zimage_parallelism.py b/tests/diffusion/distributed/test_tensor_parallel.py
similarity index 53%
rename from tests/diffusion/distributed/test_zimage_parallelism.py
rename to tests/diffusion/distributed/test_tensor_parallel.py
index ab330ee9a..9f39112ae 100644
--- a/tests/diffusion/distributed/test_zimage_parallelism.py
+++ b/tests/diffusion/distributed/test_tensor_parallel.py
@@ -1,16 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Z-Image end-to-end tests for diffusion parallelism.
+"""DiT tensor-parallel parity and perf checks for diffusion models.
-This file currently covers:
-- DiT tensor parallelism (TP=2) vs TP=1.
-- VAE patch parallelism (vae_patch_parallel_size=2) vs baseline on TP=2.
-
-Note: CUDA-only (>=2 GPUs). We use `enforce_eager=False` (default) to enable
-`torch.compile`.
+Compares TP=2 against TP=1 on output correctness, latency, and peak memory.
+Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
"""
+from __future__ import annotations
+
import os
import time
from pathlib import Path
@@ -31,22 +29,15 @@ from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
PROMPT = "a photo of a cat sitting on a laptop keyboard"
-
-
-def _get_zimage_model() -> str:
- # Allow overriding the model for local/offline environments.
- # Can be either a HuggingFace repo id or a local path.
- return os.environ.get("VLLM_TEST_ZIMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo")
+DEFAULT_MODEL = "Tongyi-MAI/Z-Image-Turbo"
def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
- """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
return torch.from_numpy(arr)
def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
- """Return (mean_abs_diff, p99_abs_diff) over RGB pixels in [0, 1]."""
ta = _pil_to_float_rgb_tensor(a)
tb = _pil_to_float_rgb_tensor(b)
assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
@@ -71,7 +62,7 @@ def _extract_single_image(outputs) -> Image.Image:
return images[0]
-def _run_zimage_generate(
+def _run_generate(
*,
tp_size: int,
height: int,
@@ -79,8 +70,6 @@ def _run_zimage_generate(
num_inference_steps: int,
seed: int,
enforce_eager: bool,
- vae_use_tiling: bool = False,
- vae_patch_parallel_size: int = 1,
num_requests: int = 4,
) -> tuple[Image.Image, float, float]:
if num_requests < 2:
@@ -91,25 +80,11 @@ def _run_zimage_generate(
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()
try:
- # Each run needs a distinct DiffusionParallelConfig; use OmniRunner per call (not the
- # parametrized omni_runner fixture, which is fixed per module).
with OmniRunner(
- _get_zimage_model(),
- parallel_config=DiffusionParallelConfig(
- tensor_parallel_size=tp_size,
- vae_patch_parallel_size=vae_patch_parallel_size,
- ),
+ DEFAULT_MODEL,
+ parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
enforce_eager=enforce_eager,
- vae_use_tiling=vae_use_tiling,
) as runner:
- # NOTE: Omni closes itself when a generate() call is exhausted.
- # To avoid measuring teardown time (process shutdown, memory cleanup),
- # we measure the latency to produce *subsequent* outputs within a single
- # generator run.
- #
- # This also serves as a warmup: the first output may include extra
- # compilation/caching overhead, while later outputs are closer to
- # steady-state inference.
gen = runner.omni.generate(
[PROMPT] * num_requests,
OmniDiffusionSamplingParams(
@@ -134,14 +109,11 @@ def _run_zimage_generate(
per_request_times_s.append(t_now - t_prev)
t_prev = t_now
- # Ensure the generator is fully consumed so it can clean up.
for _ in gen:
pass
median_time_s = float(np.median(per_request_times_s))
-
peak_memory_mb = monitor.peak_used_mb
-
return _extract_single_image([last_output]), median_time_s, peak_memory_mb
finally:
monitor.stop()
@@ -151,20 +123,20 @@ def _run_zimage_generate(
@pytest.mark.diffusion
@pytest.mark.parallel
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
-def test_zimage_tensor_parallel_tp2(tmp_path: Path):
+def test_tensor_parallel_tp2(tmp_path: Path):
if current_omni_platform.is_npu():
- pytest.skip("Z-Image TP e2e test is only supported on CUDA and ROCm for now.")
+ pytest.skip("Tensor parallel parity test is only supported on CUDA and ROCm for now.")
if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
- pytest.skip("Z-Image TP=2 requires >= 2 devices.")
+ pytest.skip("Tensor parallel TP=2 requires >= 2 devices.")
+ model_name = DEFAULT_MODEL
enforce_eager = False
-
height = 512
width = 512
num_inference_steps = 2
seed = 42
- tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
+ tp1_img, tp1_time_s, tp1_peak_mem = _run_generate(
tp_size=1,
height=height,
width=width,
@@ -172,7 +144,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
seed=seed,
enforce_eager=enforce_eager,
)
- tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
+ tp2_img, tp2_time_s, tp2_peak_mem = _run_generate(
tp_size=2,
height=height,
width=width,
@@ -181,8 +153,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
enforce_eager=enforce_eager,
)
- tp1_path = tmp_path / "zimage_tp1.png"
- tp2_path = tmp_path / "zimage_tp2.png"
+ tp1_path = tmp_path / "tp1.png"
+ tp2_path = tmp_path / "tp2.png"
tp1_img.save(tp1_path)
tp2_img.save(tp2_path)
@@ -193,7 +165,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
mean_threshold = 3e-2
p99_threshold = 2.5e-1
print(
- "Z-Image TP image diff stats (TP=1 vs TP=2): "
+ f"{model_name} TP image diff (TP=1 vs TP=2): "
f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
f"tp1_img={tp1_path}, tp2_img={tp2_path}"
@@ -203,73 +175,11 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
)
- print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
- # ROCm is not optimized TP2 can be slower than TP1
+ print(f"TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
if not current_omni_platform.is_rocm():
assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})"
- print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
+ print(f"TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
assert tp2_peak_mem < tp1_peak_mem, (
f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})"
)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
-def test_zimage_vae_patch_parallel_tp2(tmp_path: Path):
- if current_omni_platform.is_npu():
- pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA and ROCm for now.")
- if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
- pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.")
-
- enforce_eager = False
-
- # Use a larger image to ensure there are multiple VAE tiles.
- height = 1152
- width = 1152
- num_inference_steps = 2
- seed = 42
-
- baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate(
- tp_size=2,
- height=height,
- width=width,
- num_inference_steps=num_inference_steps,
- seed=seed,
- enforce_eager=enforce_eager,
- vae_use_tiling=True,
- vae_patch_parallel_size=1,
- num_requests=2,
- )
- pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate(
- tp_size=2,
- height=height,
- width=width,
- num_inference_steps=num_inference_steps,
- seed=seed,
- enforce_eager=enforce_eager,
- vae_use_tiling=True,
- vae_patch_parallel_size=2,
- num_requests=2,
- )
-
- baseline_path = tmp_path / "zimage_tp2_vae_pp1.png"
- pp2_path = tmp_path / "zimage_tp2_vae_pp2.png"
- baseline_img.save(baseline_path)
- pp2_img.save(pp2_path)
-
- mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
- mean_threshold = 5e-3
- p99_threshold = 1e-1
- print(
- "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): "
- f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
- f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
- f"pp1_img={baseline_path}, pp2_img={pp2_path}"
- )
- assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
- f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
- f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
- )
diff --git a/tests/diffusion/distributed/test_vae_decode_parallelism.py b/tests/diffusion/distributed/test_vae_decode_parallelism.py
index 2fd531e3e..57fe1e073 100644
--- a/tests/diffusion/distributed/test_vae_decode_parallelism.py
+++ b/tests/diffusion/distributed/test_vae_decode_parallelism.py
@@ -1,151 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""VAE patch-parallel decode parity tests for diffusion models.
+
+Compares ``vae_patch_parallel_size=2`` against ``vae_patch_parallel_size=1`` on
+TP=2 with VAE tiling enabled. Covers one video model (Wan T2V) and one image
+model (Qwen-Image). CUDA/ROCm-only (>=2 devices).
+"""
+
+from __future__ import annotations
+
import os
-import sys
from pathlib import Path
+from typing import Any
import numpy as np
import pytest
import torch
from PIL import Image
-from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
-if str(REPO_ROOT) not in sys.path:
- sys.path.insert(0, str(REPO_ROOT))
-
-import time
-
-from tests.helpers.runtime import OmniRunner
from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-all_model_configs = [
- {
- "model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
- "out_width": 1280,
- "out_height": 720,
- "out_frames": 16,
- },
- {
- "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
- "out_width": 1280,
- "out_height": 720,
- "out_frames": 16,
- },
- {
- "model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
- "out_width": 1280,
- "out_height": 720,
- "out_frames": 16,
- },
- {
- "model_name": "Qwen/Qwen-Image",
- "out_width": 3072,
- "out_height": 3072,
- "out_frames": 1,
- },
- # {
- # "model_name": "stepfun-ai/NextStep-1.1",
- # "out_width": 3072,
- # "out_height": 3072,
- # "out_frames": 1,
- # },
+PROMPT = "A cat sitting on a table"
+
+MODEL_CASES = [
+ pytest.param(
+ {
+ "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+ "height": 512,
+ "width": 512,
+ "num_frames": 5,
+ "needs_image": False,
+ "is_moe": False,
+ "mean_threshold": 3e-2,
+ "p99_threshold": 3e-2,
+ },
+ id="wan22_t2v_a14b",
+ ),
+ pytest.param(
+ {
+ "model_name": "Qwen/Qwen-Image",
+ "height": 1152,
+ "width": 1152,
+ "num_frames": 1,
+ "needs_image": False,
+ "is_moe": False,
+ "mean_threshold": 5e-3,
+ "p99_threshold": 1e-1,
+ },
+ id="qwen_image",
+ ),
]
-def is_nextstep_model(model_name: str) -> bool:
- """Check if the model is a NextStep model by reading its config."""
- from vllm.transformers_utils.config import get_hf_file_to_dict
-
- try:
- cfg = get_hf_file_to_dict("config.json", model_name)
- if cfg and cfg.get("model_type") == "nextstep":
- return True
- except Exception:
- pass
- return False
+def _to_float_array(data: Any) -> np.ndarray:
+ if isinstance(data, Image.Image):
+ return np.asarray(data.convert("RGB"), dtype=np.float32) / 255.0
+ if isinstance(data, torch.Tensor):
+ arr = data.detach().cpu().numpy()
+ else:
+ arr = np.asarray(data)
+ if arr.dtype == np.uint8:
+ return arr.astype(np.float32) / 255.0
+ return arr.astype(np.float32)
-def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
+def _extract_output_array(output: Any) -> np.ndarray:
+ req_out = output.request_output
+ payload = req_out.images[0]
+ return _to_float_array(payload)
+
+
+def _diff_metrics(a: np.ndarray, b: np.ndarray) -> tuple[float, float]:
+ assert a.shape == b.shape, f"Output shapes differ: {a.shape} vs {b.shape}"
+ abs_diff = np.abs(a - b)
+ return float(abs_diff.mean()), float(np.quantile(abs_diff, 0.99))
+
+
+def _save_preview(arr: np.ndarray, path: Path) -> None:
+ """Save a single RGB frame preview; video tensors use the first frame."""
+ preview = np.asarray(arr)
+ # (B, T, H, W, C) video
+ if preview.ndim == 5:
+ preview = preview[0, 0]
+ # (T, H, W, C) video or (B, H, W, C) image batch
+ elif preview.ndim == 4:
+ preview = preview[0]
+ if preview.ndim != 3 or preview.shape[-1] not in (3, 4):
+ raise ValueError(f"Expected HWC RGB(A) preview, got shape {arr.shape} -> {preview.shape}")
+ preview = np.clip(preview[..., :3], 0.0, 1.0)
+ Image.fromarray((preview * 255.0).astype(np.uint8)).save(path)
+
+
+def _build_sampling_params(model_case: dict[str, Any], seed: int) -> OmniDiffusionSamplingParams:
+ kwargs: dict[str, Any] = {
+ "height": model_case["height"],
+ "width": model_case["width"],
+ "num_frames": model_case["num_frames"],
+ "num_inference_steps": 2,
+ "seed": seed,
+ }
+ if model_case["model_name"].startswith("Wan-AI/"):
+ kwargs["guidance_scale"] = 4.0
+ if model_case["is_moe"]:
+ kwargs["guidance_scale_2"] = 1.0
+ kwargs["boundary_ratio"] = 0.5
+ return OmniDiffusionSamplingParams(**kwargs)
+
+
+def _run_generate(
+ *,
+ model_case: dict[str, Any],
+ vae_patch_parallel_size: int,
+ seed: int,
+) -> np.ndarray:
+ current_omni_platform.empty_cache()
parallel_config = DiffusionParallelConfig(
- tensor_parallel_size=tp,
+ tensor_parallel_size=2,
vae_patch_parallel_size=vae_patch_parallel_size,
)
-
- omni_kwargs = {
- "vae_use_tiling": using_tile,
- "parallel_config": parallel_config,
- }
- use_nextstep = is_nextstep_model(model_configs["model_name"])
- if use_nextstep:
- # NextStep-1.1 requires explicit pipeline class
- omni_kwargs["model_class_name"] = "NextStep11Pipeline"
-
- with OmniRunner(model_configs["model_name"], **omni_kwargs) as runner:
- m = runner.omni
- image = Image.new("RGB", (out_width, out_height), (0, 0, 0))
- start = time.perf_counter()
- outputs = m.generate(
- {
- "prompt": "A cat sitting on a table",
- "multi_modal_data": {"image": image},
- },
- sampling_params_list=OmniDiffusionSamplingParams(
- height=model_configs["out_height"],
- width=model_configs["out_width"],
- num_frames=model_configs["out_frames"],
- num_inference_steps=2,
- generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
- ),
+ with OmniRunner(
+ model_case["model_name"],
+ parallel_config=parallel_config,
+ vae_use_tiling=True,
+ enforce_eager=False,
+ ) as runner:
+ request: dict[str, Any] = {"prompt": model_case.get("prompt", PROMPT)}
+ if model_case["needs_image"]:
+ request["multi_modal_data"] = {
+ "image": Image.new("RGB", (model_case["width"], model_case["height"]), (0, 0, 0))
+ }
+ outputs = runner.omni.generate(
+ request,
+ sampling_params_list=_build_sampling_params(model_case, seed),
)
- end = time.perf_counter()
- first_output = outputs[0]
- req_out = first_output.request_output
- frames = req_out.images[0]
- if isinstance(frames, torch.Tensor):
- frames = frames.detach().cpu().numpy()
- elif isinstance(frames, Image.Image):
- frames = np.array(frames)
- # frames shape: (batch, num_frames, height, width, channels)
- cost = (end - start) * 1000
- return frames, cost
- cleanup_dist_env_and_memory()
+ return _extract_output_array(outputs[0])
-@pytest.mark.parametrize("model_configs", all_model_configs)
-def test_vae_parallel_model(model_configs: str):
- parallel_result, parallel_time = model_run(
- model_configs=model_configs,
- tp=2,
- out_width=1280,
- out_height=704,
- out_frames=5,
- using_tile=True,
- vae_patch_parallel_size=2,
- )
- non_parallel_result, non_parallel_time = model_run(
- model_configs=model_configs,
- tp=2,
- out_width=1280,
- out_height=704,
- out_frames=5,
- using_tile=True,
- vae_patch_parallel_size=1,
- )
- result_diff = np.abs(non_parallel_result - parallel_result)
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+@pytest.mark.parametrize("model_case", MODEL_CASES)
+def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path):
+ if current_omni_platform.is_npu():
+ pytest.skip("VAE patch parallel test is only supported on CUDA and ROCm for now.")
+ if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+ pytest.skip("VAE patch parallel TP=2 requires >= 2 devices.")
- mean_threshold = 3e-2
- max_threshold = 3e-2 # they should be totally same
+ seed = 42
+ model_id = model_case["model_name"].split("/")[-1]
+
+ baseline = _run_generate(model_case=model_case, vae_patch_parallel_size=1, seed=seed)
+ parallel = _run_generate(model_case=model_case, vae_patch_parallel_size=2, seed=seed)
+
+ baseline_path = tmp_path / f"{model_id}_tp2_vae_pp1.png"
+ parallel_path = tmp_path / f"{model_id}_tp2_vae_pp2.png"
+ _save_preview(baseline, baseline_path)
+ _save_preview(parallel, parallel_path)
+
+ mean_abs_diff, p99_abs_diff = _diff_metrics(baseline, parallel)
+ mean_threshold = model_case["mean_threshold"]
+ p99_threshold = model_case["p99_threshold"]
print(
- f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): "
- f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; "
- f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
- f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms"
+ f"{model_case['model_name']} VAE patch parallel diff (TP=2, pp=1 vs pp=2): "
+ f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+ f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
+ f"pp1_preview={baseline_path}, pp2_preview={parallel_path}"
+ )
+ assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+ f"Output diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, "
+ f"p99_abs_diff={p99_abs_diff:.6e} "
+ f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
)
- # assert non_parallel_time > parallel_time
- assert result_diff.mean() < mean_threshold
- assert result_diff.max() < max_threshold
diff --git a/tests/diffusion/hooks/test_hook_registry.py b/tests/diffusion/hooks/test_hook_registry.py
index 6c8535cfe..3ee93a073 100644
--- a/tests/diffusion/hooks/test_hook_registry.py
+++ b/tests/diffusion/hooks/test_hook_registry.py
@@ -12,6 +12,9 @@ from torch import nn
from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
DEFAULT_OUT = "ECHO"
OVERRIDE_OUT = "OVERRIDE"
INPUT_KWARG = "inp"
diff --git a/tests/diffusion/layers/test_rotary_emb_equivalence.py b/tests/diffusion/layers/test_rotary_emb_equivalence.py
index 2fbb7a31f..09aef0373 100644
--- a/tests/diffusion/layers/test_rotary_emb_equivalence.py
+++ b/tests/diffusion/layers/test_rotary_emb_equivalence.py
@@ -13,6 +13,8 @@ from __future__ import annotations
import pytest
import torch
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
def _apply_rotary_emb_helios_original(
hidden_states: torch.Tensor,
diff --git a/tests/e2e/offline_inference/rlhf_test/__init__.py b/tests/diffusion/models/dreamzero/__init__.py
similarity index 100%
rename from tests/e2e/offline_inference/rlhf_test/__init__.py
rename to tests/diffusion/models/dreamzero/__init__.py
diff --git a/tests/dreamzero/test_crossattn_cache.py b/tests/diffusion/models/dreamzero/test_crossattn_cache.py
similarity index 100%
rename from tests/dreamzero/test_crossattn_cache.py
rename to tests/diffusion/models/dreamzero/test_crossattn_cache.py
diff --git a/tests/dreamzero/test_deploy_configs.py b/tests/diffusion/models/dreamzero/test_deploy_configs.py
similarity index 100%
rename from tests/dreamzero/test_deploy_configs.py
rename to tests/diffusion/models/dreamzero/test_deploy_configs.py
diff --git a/tests/dreamzero/test_dreamzero_utils.py b/tests/diffusion/models/dreamzero/test_dreamzero_utils.py
similarity index 100%
rename from tests/dreamzero/test_dreamzero_utils.py
rename to tests/diffusion/models/dreamzero/test_dreamzero_utils.py
diff --git a/tests/dreamzero/test_fused_qk_rms_norm.py b/tests/diffusion/models/dreamzero/test_fused_qk_rms_norm.py
similarity index 100%
rename from tests/dreamzero/test_fused_qk_rms_norm.py
rename to tests/diffusion/models/dreamzero/test_fused_qk_rms_norm.py
diff --git a/tests/dreamzero/test_pipeline_state.py b/tests/diffusion/models/dreamzero/test_pipeline_state.py
similarity index 100%
rename from tests/dreamzero/test_pipeline_state.py
rename to tests/diffusion/models/dreamzero/test_pipeline_state.py
diff --git a/tests/dreamzero/test_qkv_fusion.py b/tests/diffusion/models/dreamzero/test_qkv_fusion.py
similarity index 100%
rename from tests/dreamzero/test_qkv_fusion.py
rename to tests/diffusion/models/dreamzero/test_qkv_fusion.py
diff --git a/tests/dreamzero/test_session_state_accounting.py b/tests/diffusion/models/dreamzero/test_session_state_accounting.py
similarity index 100%
rename from tests/dreamzero/test_session_state_accounting.py
rename to tests/diffusion/models/dreamzero/test_session_state_accounting.py
diff --git a/tests/dreamzero/test_session_state_equivalence.py b/tests/diffusion/models/dreamzero/test_session_state_equivalence.py
similarity index 100%
rename from tests/dreamzero/test_session_state_equivalence.py
rename to tests/diffusion/models/dreamzero/test_session_state_equivalence.py
diff --git a/tests/dreamzero/test_session_state_lifecycle.py b/tests/diffusion/models/dreamzero/test_session_state_lifecycle.py
similarity index 100%
rename from tests/dreamzero/test_session_state_lifecycle.py
rename to tests/diffusion/models/dreamzero/test_session_state_lifecycle.py
diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py
index de037fef8..73e10463e 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py
@@ -9,6 +9,8 @@ from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransforme
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
def test_ltx_rms_norm_no_affine_matches_official_torch_module():
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py
index 818d3d7c0..276e9e059 100644
--- a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py
+++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py
@@ -6,6 +6,8 @@ import torch
from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
def _pipeline_without_init() -> SenseNovaU1Pipeline:
return object.__new__(SenseNovaU1Pipeline)
diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py
index 555a62ca0..549a98142 100644
--- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py
+++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py
@@ -5,7 +5,6 @@ import pytest
import torch
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-from tests.helpers import skip_if_gated_repo_inaccessible
from tests.helpers.env import DeviceMemoryMonitor
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
@@ -16,32 +15,26 @@ AUDIO_MODEL = {
"stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None},
}
-IMAGE_VIDEO_MODELS = {
+IMAGE_MODELS = {
"riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100},
- "Tongyi-MAI/Z-Image-Turbo": {"cuda": 2500, "rocm": 2100},
- "OmniGen2/OmniGen2": {"cuda": 2500, "rocm": 2100},
}
-MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
+MODELS = {**AUDIO_MODEL, **IMAGE_MODELS}
MODEL_MARKS = {
"riverclouds/qwen_image_random": pytest.mark.core_model,
- "Tongyi-MAI/Z-Image-Turbo": pytest.mark.advanced_model,
"stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
- "OmniGen2/OmniGen2": pytest.mark.full_model,
}
_GATED_MODELS = {"stabilityai/stable-audio-open-1.0"}
-# Aliased for backward compatibility (imported by test_diffusion_layerwise_offload.py).
-_skip_if_gated_repo_inaccessible = skip_if_gated_repo_inaccessible
AUDIO_MODEL_PARAMS = {
"runner_params": {},
"sampler_params": {},
}
-IMAGE_VIDEO_MODELS_PARAMS = {
+IMAGE_MODEL_PARAMS = {
"runner_params": {},
"sampler_params": {
"height": 256,
@@ -60,7 +53,7 @@ def inference(model_name: str, offload: bool = True):
if model_name in AUDIO_MODEL:
params = AUDIO_MODEL_PARAMS
else:
- params = IMAGE_VIDEO_MODELS_PARAMS
+ params = IMAGE_MODEL_PARAMS
with OmniRunner(
model_name,
@@ -111,19 +104,12 @@ def check_audio_determinism(audio1, audio2, atol=1e-2):
[pytest.param(name, marks=MODEL_MARKS[name]) for name in MODELS],
)
def test_cpu_offload_diffusion_model(model_name: str):
- if model_name == "OmniGen2/OmniGen2":
- pytest.skip("issue #4537")
- if model_name in _GATED_MODELS:
- _skip_if_gated_repo_inaccessible(model_name)
try:
offload_peak_memory, output_offload = inference(model_name, offload=True)
cleanup_dist_env_and_memory()
no_offload_peak_memory, output_no_offload = inference(model_name, offload=False)
except ValueError as exc:
- # omni_snapshot_download wraps GatedRepoError in a ValueError.
- # If the pre-flight guard above did not catch it (e.g. partial
- # HF_TOKEN where config.json is accessible but weight shards are
- # blocked), skip instead of failing.
+ # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing.
if "Access to model" in str(exc) and "is restricted" in str(exc):
pytest.skip(
f"Skipping: gated HF repo {model_name!r} inaccessible "
diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py
index f02da6fab..342c9833b 100644
--- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py
@@ -3,11 +3,8 @@ import pytest
import torch
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-from tests.diffusion.offloader.test_diffusion_cpu_offload import (
- _GATED_MODELS,
- _skip_if_gated_repo_inaccessible,
-)
from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
@@ -23,6 +20,19 @@ IMAGE_VIDEO_MODELS = {
MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
+LAYERWISE_MODEL_CASES = [
+ pytest.param(
+ "stabilityai/stable-audio-open-1.0",
+ marks=[pytest.mark.full_model, pytest.mark.diffusion],
+ id="stable_audio",
+ ),
+ pytest.param(
+ "riverclouds/qwen_image_random",
+ marks=[pytest.mark.core_model, pytest.mark.diffusion],
+ id="qwen_image_random",
+ ),
+]
+
AUDIO_MODEL_PARAMS = {
"runner_params": {},
"sampler_params": {},
@@ -91,7 +101,8 @@ def run_inference(
return peak, output
-@pytest.mark.parametrize("model_name", list(MODELS.keys()))
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("model_name", LAYERWISE_MODEL_CASES)
def test_layerwise_offload_diffusion_model(model_name: str):
"""Test that layerwise offloading reduces GPU memory usage.
@@ -100,8 +111,6 @@ def test_layerwise_offload_diffusion_model(model_name: str):
offloader keeps only a single transformer block on GPU at a time, with
prefetching for compute-memory overlap.
"""
- if model_name in _GATED_MODELS:
- _skip_if_gated_repo_inaccessible(model_name)
try:
# Run without layerwise offloading (baseline)
no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False)
@@ -111,10 +120,7 @@ def test_layerwise_offload_diffusion_model(model_name: str):
layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True)
cleanup_dist_env_and_memory()
except ValueError as exc:
- # omni_snapshot_download wraps GatedRepoError in a ValueError.
- # If the pre-flight guard above did not catch it (e.g. partial
- # HF_TOKEN where config.json is accessible but weight shards are
- # blocked), skip instead of failing.
+ # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing.
if "Access to model" in str(exc) and "is restricted" in str(exc):
pytest.skip(
f"Skipping: gated HF repo {model_name!r} inaccessible "
diff --git a/tests/plugin/test_gguf_diffusion.py b/tests/diffusion/quantization/test_gguf_diffusion.py
similarity index 97%
rename from tests/plugin/test_gguf_diffusion.py
rename to tests/diffusion/quantization/test_gguf_diffusion.py
index 44de4a1fb..fa87e8a7c 100644
--- a/tests/plugin/test_gguf_diffusion.py
+++ b/tests/diffusion/quantization/test_gguf_diffusion.py
@@ -5,10 +5,10 @@
Validates that GGUF-quantized diffusion models generate valid images and
use less peak GPU memory than BF16 baseline.
-Requires vllm-omni to be installed alongside the plugin.
+Requires vllm-omni to be installed alongside the vllm-gguf-plugin package.
Usage:
- pytest tests/diffusion/test_gguf_memory.py -v
+ pytest tests/diffusion/quantization/test_gguf_diffusion.py -v
"""
from __future__ import annotations
@@ -112,7 +112,6 @@ def _generate_single_stage_image(
@hardware_test(res={"cuda": "L4"})
@pytest.mark.full_model
@pytest.mark.diffusion
-@pytest.mark.slow
@pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"])
def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None:
"""Z-Image-Turbo GGUF generates valid images and uses less memory than BF16."""
diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py
index c2807797d..d3cee6b6a 100644
--- a/tests/diffusion/quantization/test_quantization_quality.py
+++ b/tests/diffusion/quantization/test_quantization_quality.py
@@ -431,6 +431,8 @@ def _free_gpu_memory():
# ---------------------------------------------------------------------------
+@pytest.mark.core_model
+@pytest.mark.cpu
def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch):
from benchmarks.diffusion.quantization_quality import _generate_image as benchmark_generate_image
from vllm_omni.outputs import OmniRequestOutput
diff --git a/tests/diffusion/quantization/test_trajectory_similarity_tool.py b/tests/diffusion/quantization/test_trajectory_similarity_tool.py
index 252e6f1c1..963ed5685 100644
--- a/tests/diffusion/quantization/test_trajectory_similarity_tool.py
+++ b/tests/diffusion/quantization/test_trajectory_similarity_tool.py
@@ -22,6 +22,8 @@ from vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity import
summarize_output_image_metrics,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
def test_compute_tensor_metrics_identical_tensors():
metrics = compute_tensor_metrics(torch.ones(2, 3), torch.ones(2, 3))
diff --git a/tests/test_diffusion_config_fields.py b/tests/diffusion/test_diffusion_config_fields.py
similarity index 97%
rename from tests/test_diffusion_config_fields.py
rename to tests/diffusion/test_diffusion_config_fields.py
index b87ceec1d..d62e7bdc8 100644
--- a/tests/test_diffusion_config_fields.py
+++ b/tests/diffusion/test_diffusion_config_fields.py
@@ -29,7 +29,7 @@ def test_diffusion_stage_configs_only_contain_valid_fields():
Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
"""
# Scan both main configs and test configs
- repo_root = Path(__file__).parent.parent
+ repo_root = Path(__file__).parent.parent.parent
config_dirs = [
repo_root / "vllm_omni" / "model_executor" / "stage_configs",
]
diff --git a/tests/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py
similarity index 100%
rename from tests/test_diffusion_config_propagation.py
rename to tests/diffusion/test_diffusion_config_propagation.py
diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py
index e44766feb..192675908 100644
--- a/tests/diffusion/test_diffusion_engine.py
+++ b/tests/diffusion/test_diffusion_engine.py
@@ -14,6 +14,7 @@ import torch
from pytest_mock import MockerFixture
import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module
+from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.diffusion_engine import (
DiffusionEngine,
@@ -30,6 +31,8 @@ from vllm_omni.diffusion.sched.interface import (
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
@dataclass
class DiffusionSchedulerOutput:
@@ -56,6 +59,8 @@ class MockScheduler:
def __init__(self):
self._waiting_queue = []
self._step_id = 0
+ # Match RequestScheduler API used by request-batch admission wait.
+ self.max_num_running_reqs = 5
def add_request(self, request):
self._waiting_queue.append(request)
@@ -64,6 +69,12 @@ class MockScheduler:
def has_requests(self):
return len(self._waiting_queue) > 0
+ def num_waiting_requests(self) -> int:
+ return len(self._waiting_queue)
+
+ def num_running_requests(self) -> int:
+ return 0
+
def schedule(self) -> DiffusionSchedulerOutput:
if not self._waiting_queue:
return DiffusionSchedulerOutput(step_id=self._step_id)
@@ -582,9 +593,7 @@ def test_move_tensor_tree_returns_non_tensor_values_unchanged() -> None:
assert moved is value
-@pytest.mark.diffusion
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
tensor = torch.arange(8, dtype=torch.float32, device="cuda")
other = torch.arange(4, dtype=torch.int64, device="cuda")
@@ -610,6 +619,7 @@ async def _consume_final_output(generator):
return final_output
+@pytest.mark.cpu
@pytest.mark.asyncio
async def test_async_add_req_and_stream_response():
engine = object.__new__(DiffusionEngine)
@@ -621,13 +631,24 @@ async def test_async_add_req_and_stream_response():
engine._cv = threading.Condition(engine._rpc_lock)
engine._init_lock = asyncio.Lock()
engine._closed = False
- engine.od_config = SimpleNamespace(streaming_output=False)
+ # Enable admission wait so concurrent adds land in one schedule wave;
+ # otherwise the first request can execute alone (~1s) and the rest in a
+ # second wave (~2s), failing the latency-spread assertion.
+ engine.od_config = SimpleNamespace(
+ streaming_output=False,
+ request_batch_max_wait_ms=500.0,
+ )
engine._loop_started = False
engine.main_loop = None
- engine.supports_request_batch = False
+ engine.supports_request_batch = True
+ engine.step_execution = False
engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
- engine._finalize_finished_request = lambda rid, out, err: out.result
+ def _finalize(rid, out, err=None):
+ # Stream consumers stop on ``finished``; keep result_data for assertions.
+ return SimpleNamespace(result_data=out.result.result_data, finished=True)
+
+ engine._finalize_finished_request = _finalize
def mock_execute_batch(sched_output):
request_ids = sched_output.scheduled_request_ids
diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py
index 06ad2077f..8f434c0e2 100644
--- a/tests/diffusion/test_diffusion_engine_metrics.py
+++ b/tests/diffusion/test_diffusion_engine_metrics.py
@@ -12,6 +12,11 @@ from __future__ import annotations
import ast
import os
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
_ENGINE_PATH = os.path.normpath(
os.path.join(
os.path.dirname(__file__),
diff --git a/tests/docs/test_generate_examples.py b/tests/docs/test_generate_examples.py
index 0b29837a9..57d7d4832 100644
--- a/tests/docs/test_generate_examples.py
+++ b/tests/docs/test_generate_examples.py
@@ -6,6 +6,9 @@ from pathlib import Path
import pytest
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
ROOT_DIR = Path(__file__).parents[2]
HOOK_PATH = ROOT_DIR / "docs/mkdocs/hooks/generate_examples.py"
SPEC = importlib.util.spec_from_file_location("generate_examples", HOOK_PATH)
diff --git a/tests/dreamzero/openpi_client_helper.py b/tests/dreamzero/openpi_client_helper.py
deleted file mode 100644
index 8918e5c36..000000000
--- a/tests/dreamzero/openpi_client_helper.py
+++ /dev/null
@@ -1,307 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import uuid
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-
-try:
- import cv2
-except ImportError: # pragma: no cover - optional e2e dependency
- cv2 = None
-
-try:
- import websockets.sync.client as websockets_client
-except ImportError: # pragma: no cover - optional e2e dependency
- websockets_client = None
-
-try:
- from openpi_client import msgpack_numpy
-except ImportError: # pragma: no cover - optional e2e dependency
- msgpack_numpy = None
-
-PING_INTERVAL_SECS = 300
-PING_TIMEOUT_SECS = 3600
-DEFAULT_HOST = "127.0.0.1"
-DEFAULT_PORT = 8000
-DEFAULT_PATH = "/v1/realtime/robot/openpi"
-DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
-ACTION_HORIZON = 24
-DEFAULT_ACTION_DIM = 8
-RELATIVE_OFFSETS = [-23, -16, -8, 0]
-CAMERA_FILES = {
- "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
- "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
- "observation/wrist_image_left": "wrist_image_left.mp4",
-}
-
-
-def _decode_action_response(response: bytes | str) -> np.ndarray:
- if isinstance(response, str):
- raise RuntimeError(f"Inference failed: {response}")
- decoded = msgpack_numpy.unpackb(response)
- if isinstance(decoded, dict) and decoded.get("type") == "error":
- message = decoded.get("message", decoded)
- raise RuntimeError(f"Inference failed: {message}")
- return np.asarray(decoded, dtype=np.float32)
-
-
-def require_dependencies() -> None:
- missing = []
- if cv2 is None:
- missing.append("opencv-python")
- if websockets_client is None:
- missing.append("websockets")
- if msgpack_numpy is None:
- missing.append("openpi-client")
- if missing:
- raise ModuleNotFoundError(f"DreamZero OpenPI test dependencies are missing: {', '.join(missing)}")
-
-
-@dataclass(frozen=True)
-class DreamZeroServerMetadata:
- image_resolution: tuple[int, int]
- n_external_cameras: int
- needs_wrist_camera: bool
- needs_stereo_camera: bool
- needs_session_id: bool
- action_space: str
-
- @classmethod
- def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata:
- required_keys = (
- "image_resolution",
- "n_external_cameras",
- "needs_wrist_camera",
- "needs_stereo_camera",
- "needs_session_id",
- "action_space",
- )
- missing_keys = [key for key in required_keys if key not in payload]
- if missing_keys:
- raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}")
-
- image_resolution = payload["image_resolution"]
- if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2:
- raise ValueError(f"Invalid image_resolution: {image_resolution!r}")
-
- return cls(
- image_resolution=(int(image_resolution[0]), int(image_resolution[1])),
- n_external_cameras=int(payload["n_external_cameras"]),
- needs_wrist_camera=bool(payload["needs_wrist_camera"]),
- needs_stereo_camera=bool(payload["needs_stereo_camera"]),
- needs_session_id=bool(payload["needs_session_id"]),
- action_space=str(payload["action_space"]),
- )
-
-
-class OpenPIWebsocketClient:
- def __init__(
- self,
- *,
- host: str = DEFAULT_HOST,
- port: int = DEFAULT_PORT,
- path: str = DEFAULT_PATH,
- ) -> None:
- require_dependencies()
- self._uri = f"ws://{host}:{port}{path}"
- self._packer = msgpack_numpy.Packer()
- self._ws, self._server_metadata = self._connect()
-
- def _connect(self):
- conn = websockets_client.connect(
- self._uri,
- compression=None,
- max_size=None,
- ping_interval=PING_INTERVAL_SECS,
- ping_timeout=PING_TIMEOUT_SECS,
- )
- metadata = msgpack_numpy.unpackb(conn.recv())
- if not isinstance(metadata, dict):
- raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
- return conn, metadata
-
- def get_server_metadata(self) -> dict[str, Any]:
- return dict(self._server_metadata)
-
- def infer(self, obs: dict[str, Any]) -> np.ndarray:
- payload = dict(obs)
- payload["endpoint"] = "infer"
- self._ws.send(self._packer.pack(payload))
- response = self._ws.recv()
- return _decode_action_response(response)
-
- def reset(self, reset_info: dict[str, Any] | None = None) -> str:
- payload = dict(reset_info or {})
- payload["endpoint"] = "reset"
- self._ws.send(self._packer.pack(payload))
- response = self._ws.recv()
- if isinstance(response, str):
- return response
- decoded = msgpack_numpy.unpackb(response)
- if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
- raise RuntimeError(f"Unexpected reset response: {decoded!r}")
- return str(decoded["status"])
-
- def close(self) -> None:
- self._ws.close()
-
-
-def load_all_frames(video_path: Path) -> np.ndarray:
- require_dependencies()
- cap = cv2.VideoCapture(str(video_path))
- frames = []
- while True:
- ok, frame = cap.read()
- if not ok:
- break
- frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
- cap.release()
- if not frames:
- raise RuntimeError(f"No frames loaded from {video_path}")
- return np.stack(frames, axis=0)
-
-
-def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
- camera_frames: dict[str, np.ndarray] = {}
- for camera_key, file_name in CAMERA_FILES.items():
- video_path = video_dir / file_name
- if not video_path.exists():
- raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}")
- camera_frames[camera_key] = load_all_frames(video_path)
- return camera_frames
-
-
-def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
- chunks: list[list[int]] = []
- current_frame = 23
- for _ in range(num_chunks):
- indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
- if indices[-1] >= total_frames:
- break
- chunks.append(indices)
- current_frame += ACTION_HORIZON
- return chunks
-
-
-def make_obs_from_video(
- camera_frames: dict[str, np.ndarray],
- frame_indices: list[int],
- *,
- prompt: str,
- session_id: str,
-) -> dict[str, Any]:
- obs: dict[str, Any] = {}
- for camera_key, all_frames in camera_frames.items():
- selected = all_frames[frame_indices]
- obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
-
- obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
- obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
- obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
- obs["prompt"] = prompt
- obs["session_id"] = session_id
- return obs
-
-
-def build_demo_observations(
- camera_frames: dict[str, np.ndarray],
- *,
- prompt: str,
- session_id: str,
- num_chunks: int = 2,
-) -> list[dict[str, Any]]:
- if num_chunks < 1:
- raise ValueError("num_chunks must be at least 1")
-
- total_frames = min(frames.shape[0] for frames in camera_frames.values())
- observations = [
- make_obs_from_video(
- camera_frames,
- [0],
- prompt=prompt,
- session_id=session_id,
- )
- ]
- for indices in build_frame_schedule(total_frames, num_chunks - 1):
- observations.append(
- make_obs_from_video(
- camera_frames,
- indices,
- prompt=prompt,
- session_id=session_id,
- )
- )
- return observations
-
-
-def validate_session_result(
- result: dict[str, Any],
- *,
- expected_action_horizon: int = ACTION_HORIZON,
- expected_action_dim: int = DEFAULT_ACTION_DIM,
-) -> None:
- metadata = DreamZeroServerMetadata.from_dict(result["metadata"])
- if metadata.image_resolution != (180, 320):
- raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}")
- if metadata.n_external_cameras != 2:
- raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}")
- if not metadata.needs_wrist_camera:
- raise AssertionError("DreamZero test expects wrist camera metadata")
- if metadata.action_space != "joint_position":
- raise AssertionError(f"Unexpected action_space: {metadata.action_space}")
-
- actions = result["actions"]
- if len(actions) != 3:
- raise AssertionError(f"Expected 3 action tensors, got {len(actions)}")
- for index, action in enumerate(actions):
- if action.shape != (expected_action_horizon, expected_action_dim):
- raise AssertionError(
- f"Action {index} shape mismatch: expected "
- f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
- )
- if not np.isfinite(action).all():
- raise AssertionError(f"Action {index} contains non-finite values")
-
- if result["reset_status"] != "reset successful":
- raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}")
-
-
-def run_policy_session(
- *,
- host: str = DEFAULT_HOST,
- port: int = DEFAULT_PORT,
- path: str = DEFAULT_PATH,
- video_dir: Path,
- prompt: str = DEFAULT_PROMPT,
- session_id: str | None = None,
- num_chunks: int = 2,
-) -> dict[str, Any]:
- session_id = session_id or str(uuid.uuid4())
- camera_frames = load_camera_frames(video_dir)
- observations = build_demo_observations(
- camera_frames,
- prompt=prompt,
- session_id=session_id,
- num_chunks=num_chunks,
- )
-
- client = OpenPIWebsocketClient(host=host, port=port, path=path)
- try:
- metadata = client.get_server_metadata()
- actions = [client.infer(obs) for obs in observations]
- reset_status = client.reset({})
- actions.append(client.infer(observations[0]))
- return {
- "metadata": metadata,
- "actions": actions,
- "reset_status": reset_status,
- "session_id": session_id,
- }
- finally:
- client.close()
diff --git a/tests/dreamzero/test_openpi_client_helper.py b/tests/dreamzero/test_openpi_client_helper.py
deleted file mode 100644
index 76950a9ed..000000000
--- a/tests/dreamzero/test_openpi_client_helper.py
+++ /dev/null
@@ -1,46 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import numpy as np
-import pytest
-
-from tests.dreamzero import openpi_client_helper
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class FakeMsgpackNumpy:
- payloads = {}
-
- @staticmethod
- def packb(obj):
- key = str(len(FakeMsgpackNumpy.payloads)).encode()
- FakeMsgpackNumpy.payloads[key] = obj
- return key
-
- @staticmethod
- def unpackb(data):
- return FakeMsgpackNumpy.payloads[data]
-
-
-def test_decode_action_response_surfaces_structured_error(monkeypatch):
- monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy)
- payload = FakeMsgpackNumpy.packb(
- {
- "type": "error",
- "message": "Internal inference error",
- }
- )
-
- with pytest.raises(RuntimeError, match="Internal inference error"):
- openpi_client_helper._decode_action_response(payload)
-
-
-def test_decode_action_response_converts_action_payload_to_float32(monkeypatch):
- monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy)
- payload = FakeMsgpackNumpy.packb(np.asarray([[1.0, 2.0]], dtype=np.float64))
-
- actions = openpi_client_helper._decode_action_response(payload)
-
- assert actions.dtype == np.float32
- np.testing.assert_array_equal(actions, np.asarray([[1.0, 2.0]], dtype=np.float32))
diff --git a/tests/dreamzero/upstream/openpi_test_client_ar.py b/tests/dreamzero/upstream/openpi_test_client_ar.py
deleted file mode 100644
index 7c5aea1d9..000000000
--- a/tests/dreamzero/upstream/openpi_test_client_ar.py
+++ /dev/null
@@ -1,381 +0,0 @@
-#!/usr/bin/env python3
-"""
-Copied from the DreamZero repository's `test_client_AR.py`:
-https://github.com/dreamzero0/dreamzero/blob/main/test_client_AR.py
-
-Kept here for end-to-end compatibility / parity testing against the vLLM
-OpenPI server.
-
-Test client for AR_droid policy server using roboarena interface.
-
-Sends real video frames from debug_image/ directory instead of zero dummy images.
-
-Frame schedule (matching debug_inference.py):
- - Step 0 (initial): send frame [0] (1 frame, H W 3)
- - Step 1: send frames [0, 7, 15, 23] (4 frames, 4 H W 3)
- - Step 2: send frames [24, 31, 39, 47] (4 frames)
- - Step 3: send frames [48, 55, 63, 71] (4 frames)
- - ...
-
-Expected server configuration:
- - image_resolution: (180, 320)
- - n_external_cameras: 2
- - needs_wrist_camera: True
- - action_space: "joint_position"
-
-Usage:
- # Start server with roboarena interface:
- torchrun --nproc_per_node=8 socket_test_optimized_AR.py --port 8000
-
- # Run against the original DreamZero websocket server (default path):
- python test_client_AR.py --host --port 8000
-
- # Run against the vLLM OpenPI server:
- python test_client_AR.py --host --port 8000 --path /v1/realtime/robot/openpi
-
- # Use zero images instead of real video (old behavior):
- python test_client_AR.py --host --port 8000 --use-zero-images
-"""
-
-import argparse
-import logging
-import os
-import sys
-import time
-import uuid
-from pathlib import Path
-
-import cv2
-import numpy as np
-from openpi_client import msgpack_numpy
-
-_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
-DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None
-
-
-def _import_upstream_policy_modules():
- if DREAMZERO_REPO is None:
- raise ImportError("Set DREAMZERO_REPO to an upstream DreamZero checkout before using this helper.")
- if DREAMZERO_REPO.exists() and str(DREAMZERO_REPO) not in sys.path:
- sys.path.insert(0, str(DREAMZERO_REPO))
-
- import eval_utils.policy_server as policy_server
- from eval_utils.policy_client import WebsocketClientPolicy
-
- return policy_server, WebsocketClientPolicy
-
-
-policy_server, WebsocketClientPolicy = _import_upstream_policy_modules()
-
-VIDEO_DIR = os.environ.get(
- "DREAMZERO_VIDEO_DIR",
- str(DREAMZERO_REPO / "debug_image") if DREAMZERO_REPO is not None else "debug_image",
-)
-
-# roboarena key -> video filename
-CAMERA_FILES = {
- "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
- "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
- "observation/wrist_image_left": "wrist_image_left.mp4",
-}
-
-# Frame schedule constants (matching debug_inference.py)
-RELATIVE_OFFSETS = [-23, -16, -8, 0]
-ACTION_HORIZON = 24
-DEFAULT_WEBSOCKET_PATH = ""
-
-
-class OpenPIWebsocketClientPolicy(WebsocketClientPolicy):
- """DreamZero websocket client with a configurable path suffix.
-
- The original DreamZero client connects to ``ws://host:port``.
- vLLM serves the compatible robot policy endpoint at
- ``/v1/realtime/robot/openpi`` when ``path`` is set accordingly.
- """
-
- def __init__(
- self,
- host: str = "0.0.0.0",
- port: int = 8000,
- path: str = DEFAULT_WEBSOCKET_PATH,
- ) -> None:
- self._uri = f"ws://{host}:{port}{path}"
- self._packer = msgpack_numpy.Packer()
- self._ws, self._server_metadata = self._wait_for_server()
-
-
-def load_all_frames(video_path: str) -> np.ndarray:
- """Load all frames from a video file. Returns (N, H, W, 3) uint8 array (RGB)."""
- cap = cv2.VideoCapture(video_path)
- frames = []
- while True:
- ret, frame = cap.read()
- if not ret:
- break
- frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
- cap.release()
- if not frames:
- raise RuntimeError(f"No frames loaded from {video_path}")
- return np.stack(frames, axis=0)
-
-
-def load_camera_frames() -> dict[str, np.ndarray]:
- """Load all video frames for each camera from the debug_image/ directory.
-
- Returns:
- Dict mapping roboarena camera keys to (N, H, W, 3) uint8 arrays.
- """
- camera_frames: dict[str, np.ndarray] = {}
- for cam_key, fname in CAMERA_FILES.items():
- path = os.path.join(VIDEO_DIR, fname)
- camera_frames[cam_key] = load_all_frames(path)
- logging.info(f"Loaded {cam_key}: {camera_frames[cam_key].shape}")
- return camera_frames
-
-
-def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
- """Build the frame index schedule for multi-frame chunks.
-
- Returns a list of frame-index lists. Each inner list has 4 indices.
- """
- chunks: list[list[int]] = []
- current_frame = 23 # first anchor frame
- for _ in range(num_chunks):
- indices = [max(current_frame + off, 0) for off in RELATIVE_OFFSETS]
- if indices[-1] >= total_frames:
- logging.info(f"Frame {indices[-1]} >= {total_frames}, stopping at {len(chunks)} chunks")
- break
- chunks.append(indices)
- current_frame += ACTION_HORIZON
- return chunks
-
-
-def _make_obs_from_video(
- camera_frames: dict[str, np.ndarray],
- frame_indices: list[int],
- prompt: str,
- session_id: str,
-) -> dict:
- """Build an observation dict from real video frames.
-
- For 1 frame: each image key is (H, W, 3).
- For 4 frames: each image key is (4, H, W, 3).
- """
- obs: dict = {}
- for cam_key, all_frames in camera_frames.items():
- selected = all_frames[frame_indices] # (T, H, W, 3)
- if len(frame_indices) == 1:
- selected = selected[0] # (H, W, 3)
- obs[cam_key] = selected
-
- obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
- obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
- obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
- obs["prompt"] = prompt
- obs["session_id"] = session_id
- return obs
-
-
-def _make_zero_observation(
- server_config: policy_server.PolicyServerConfig,
- prompt: str = "pick up the object",
- session_id: str | None = None,
-) -> dict:
- """Create a dummy observation matching AR_droid expectations.
-
- AR_droid expects:
- - 2 external cameras (exterior_image_0_left, exterior_image_1_left)
- - 1 wrist camera (wrist_image_left)
- - Image resolution: 180x320 (H x W)
- - joint_position: 7 DoF
- - gripper_position: 1 DoF
- """
- obs = {}
-
- # Determine image resolution
- if server_config.image_resolution is not None:
- h, w = server_config.image_resolution
- else:
- # Default for AR_droid
- h, w = 180, 320
-
- # External cameras (0-indexed in roboarena)
- for i in range(server_config.n_external_cameras):
- obs[f"observation/exterior_image_{i}_left"] = np.zeros((h, w, 3), dtype=np.uint8)
- if server_config.needs_stereo_camera:
- obs[f"observation/exterior_image_{i}_right"] = np.zeros((h, w, 3), dtype=np.uint8)
-
- # Wrist camera
- if server_config.needs_wrist_camera:
- obs["observation/wrist_image_left"] = np.zeros((h, w, 3), dtype=np.uint8)
- if server_config.needs_stereo_camera:
- obs["observation/wrist_image_right"] = np.zeros((h, w, 3), dtype=np.uint8)
-
- # Session ID - should be passed in to ensure consistency within a session
- if server_config.needs_session_id:
- import uuid
-
- # Generate unique session ID if not provided
- obs["session_id"] = session_id if session_id else str(uuid.uuid4())
-
- # State observations (AR_droid: 7 DoF arm + 1 gripper)
- obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
- obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
- obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
-
- # Language prompt
- obs["prompt"] = prompt
-
- return obs
-
-
-def test_ar_droid_policy_server(
- host: str = "localhost",
- port: int = 8000,
- path: str = DEFAULT_WEBSOCKET_PATH,
- num_chunks: int = 15,
- prompt: str = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
- use_zero_images: bool = False,
-):
- """Test the AR_droid policy server with roboarena interface.
-
- When use_zero_images is False (default), loads real video frames from
- debug_image/ and follows the frame schedule from debug_inference.py.
- """
- logging.info(f"Connecting to AR_droid server at ws://{host}:{port}{path} ...")
-
- client = OpenPIWebsocketClientPolicy(host=host, port=port, path=path)
-
- # Validate server metadata
- metadata = client.get_server_metadata()
- logging.info(f"Server metadata: {metadata}")
- assert isinstance(metadata, dict), "Metadata should be a dict"
-
- try:
- server_config = policy_server.PolicyServerConfig(**metadata)
- except Exception as e:
- logging.error(f"Error parsing metadata: {e}")
- raise e
-
- # Validate expected AR_droid configuration
- logging.info(f"Server config: {server_config}")
- assert server_config.n_external_cameras == 2, f"Expected 2 external cameras, got {server_config.n_external_cameras}"
- assert server_config.needs_wrist_camera, "Expected wrist camera to be enabled"
- assert server_config.action_space == "joint_position", (
- f"Expected joint_position action space, got {server_config.action_space}"
- )
-
- logging.info("Server configuration validated for AR_droid")
-
- # Generate unique session ID for this test run
- session_id = str(uuid.uuid4())
- logging.info(f"Session ID: {session_id}")
-
- # ── Zero-image fallback mode ──────────────────────────────────────
- if use_zero_images:
- logging.info("Using ZERO dummy images (legacy mode)")
- for i in range(num_chunks):
- obs = _make_zero_observation(server_config, prompt=prompt, session_id=session_id)
- logging.info(f"Inference {i + 1}/{num_chunks}: prompt='{prompt}'")
- t0 = time.time()
- actions = client.infer(obs)
- dt = time.time() - t0
- _log_action(actions, dt)
-
- logging.info("Sending reset...")
- client.reset({})
- logging.info("Done (zero-image mode).")
- return
-
- # ── Real video frame mode ─────────────────────────────────────────
- logging.info("Loading real video frames from debug_image/ directory")
- camera_frames = load_camera_frames()
-
- total_frames = min(v.shape[0] for v in camera_frames.values())
- logging.info(f"Total frames available: {total_frames}")
-
- # Build frame schedule
- chunks = build_frame_schedule(total_frames, num_chunks)
-
- logging.info("Frame schedule:")
- logging.info(" Initial: [0]")
- for i, indices in enumerate(chunks):
- logging.info(f" Chunk {i}: {indices}")
-
- # Step 0: initial single frame
- logging.info("=== Initial: frame [0] ===")
- obs = _make_obs_from_video(camera_frames, [0], prompt, session_id)
- t0 = time.time()
- actions = client.infer(obs)
- dt = time.time() - t0
- _log_action(actions, dt)
-
- # Subsequent chunks: send 4 frames at a time
- for chunk_idx, frame_indices in enumerate(chunks):
- logging.info(f"=== Chunk {chunk_idx}: frames {frame_indices} ===")
- obs = _make_obs_from_video(camera_frames, frame_indices, prompt, session_id)
- t0 = time.time()
- actions = client.infer(obs)
- dt = time.time() - t0
- _log_action(actions, dt)
-
- # Reset triggers video save on the server
- logging.info("Sending reset to save video...")
- client.reset({})
-
- logging.info("Done.")
-
-
-def _log_action(actions: np.ndarray, dt: float) -> None:
- """Pretty-print action shape, range, and timing."""
- assert isinstance(actions, np.ndarray), f"Expected numpy array, got {type(actions)}"
- assert actions.ndim == 2, f"Expected 2D array, got shape {actions.shape}"
- assert actions.shape[-1] == 8, f"Expected 8 action dims (7 joints + 1 gripper), got {actions.shape[-1]}"
- logging.info(f" Action shape: {actions.shape}, range: [{actions.min():.4f}, {actions.max():.4f}], time: {dt:.2f}s")
-
-
-def main():
- parser = argparse.ArgumentParser(description="Test AR_droid policy server with real video frames from debug_image/")
- parser.add_argument("--host", default="localhost", help="Server hostname")
- parser.add_argument("--port", type=int, default=8000, help="Server port")
- parser.add_argument(
- "--path",
- default=DEFAULT_WEBSOCKET_PATH,
- help="WebSocket path suffix (default: empty string for the original DreamZero server)",
- )
- parser.add_argument(
- "--num-chunks",
- type=int,
- default=15,
- help="Number of 4-frame chunks to send after the initial frame (default: 15)",
- )
- parser.add_argument(
- "--prompt",
- default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
- help="Language prompt for the policy",
- )
- parser.add_argument(
- "--use-zero-images",
- action="store_true",
- help="Use zero dummy images instead of real video frames (legacy mode)",
- )
-
- args = parser.parse_args()
-
- logging.basicConfig(
- level=logging.INFO,
- format="%(asctime)s [%(levelname)s] %(message)s",
- )
-
- test_ar_droid_policy_server(
- host=args.host,
- port=args.port,
- path=args.path,
- num_chunks=args.num_chunks,
- prompt=args.prompt,
- use_zero_images=args.use_zero_images,
- )
-
-
-if __name__ == "__main__":
- main()
diff --git a/tests/dreamzero/upstream/test_openpi_e2e_source_parity.py b/tests/dreamzero/upstream/test_openpi_e2e_source_parity.py
deleted file mode 100644
index e41cff903..000000000
--- a/tests/dreamzero/upstream/test_openpi_e2e_source_parity.py
+++ /dev/null
@@ -1,358 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Formal OpenPI end-to-end parity: upstream DreamZero server vs `vllm serve`.
-
-This test uses DreamZero's own client-side observation builders from
-`${DREAMZERO_REPO}/test_client_AR.py`, and client-side websocket protocol from
-`${DREAMZERO_REPO}/eval_utils/policy_client.py`.
-
-The only client-side adaptation for vLLM is the websocket path:
-DreamZero's upstream server serves at `/`, while vLLM serves OpenPI at
-`/v1/realtime/robot/openpi`.
-
-Current scope for this test:
-- default two-GPU run (`nproc_per_node=2` on upstream, `--cfg-parallel-size 2` on `vllm serve`)
-- non-`torch.compile` (upstream launched through
- `upstream_socket_server_no_compile.py`, vLLM with `--enforce-eager`)
-- non-DiT-cache / non-skip-schedule (`NUM_DIT_STEPS=16`)
-
-Serving contract locked by this test:
-- upstream DreamZero still boots from the local checkpoint directory
-- vLLM boots from the official DreamZero HF repo name (`GEAR-Dreams/DreamZero-DROID`)
- rather than a prepared local bundle path
-"""
-
-from __future__ import annotations
-
-import os
-import shutil
-import subprocess
-import sys
-import time
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-from tests.helpers.runtime import get_open_port
-
-msgpack_numpy = pytest.importorskip("openpi_client.msgpack_numpy")
-
-_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
-DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None
-if DREAMZERO_REPO is not None and str(DREAMZERO_REPO) not in sys.path:
- sys.path.insert(0, str(DREAMZERO_REPO))
-
-try:
- import test_client_AR as dreamzero_client
- from eval_utils.policy_client import WebsocketClientPolicy
-except Exception: # pragma: no cover - guarded by pytest skip below
- dreamzero_client = None
- WebsocketClientPolicy = None
-_BaseWebsocketClientPolicy = WebsocketClientPolicy if WebsocketClientPolicy is not None else object
-
-CHECKPOINT_DIR = DREAMZERO_REPO / "checkpoints" / "dreamzero" if DREAMZERO_REPO is not None else None
-VLLM_MODEL = os.environ.get("VLLM_DREAMZERO_MODEL", "GEAR-Dreams/DreamZero-DROID")
-SERVICE_READY_TIMEOUT_S = int(os.environ.get("OPENPI_SERVICE_READY_TIMEOUT_S", "900"))
-PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
-SESSION_ID = "openpi-e2e-parity-session"
-
-pytestmark = [
- pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU required"),
- pytest.mark.skipif(
- dreamzero_client is None or WebsocketClientPolicy is None,
- reason="DreamZero client modules are required on PYTHONPATH",
- ),
- pytest.mark.skipif(
- DREAMZERO_REPO is None or not DREAMZERO_REPO.exists(),
- reason="DreamZero source repo is required at DREAMZERO_REPO",
- ),
- pytest.mark.skipif(
- CHECKPOINT_DIR is None or not CHECKPOINT_DIR.exists(), reason="DreamZero local checkpoint is required"
- ),
-]
-
-
-class OpenPIWebsocketClientPolicy(_BaseWebsocketClientPolicy):
- """DreamZero client protocol with an OpenPI websocket path suffix."""
-
- def __init__(
- self,
- host: str = "127.0.0.1",
- port: int = 8000,
- path: str = "/v1/realtime/robot/openpi",
- ) -> None:
- self._uri = f"ws://{host}:{port}{path}"
- self._packer = msgpack_numpy.Packer()
- self._ws, self._server_metadata = self._wait_for_server()
-
-
-def _vllm_executable() -> str:
- fallback = Path(sys.executable).with_name("vllm")
- if fallback.exists():
- return str(fallback)
- exe = shutil.which("vllm")
- if exe:
- return exe
- raise FileNotFoundError("Unable to locate `vllm` executable in current environment.")
-
-
-def _cfg_parallel_size() -> int:
- return int(os.environ.get("OPENPI_E2E_CFG_PARALLEL_SIZE", "2"))
-
-
-def _pick_test_gpus() -> list[str]:
- cfg_parallel_size = _cfg_parallel_size()
- override = os.environ.get("OPENPI_E2E_GPUS") or os.environ.get("OPENPI_E2E_GPU")
- if override is not None:
- gpus = [part.strip() for part in override.split(",") if part.strip()]
- if not gpus:
- raise ValueError("OPENPI_E2E_GPUS is set but empty.")
- if len(gpus) < cfg_parallel_size:
- raise RuntimeError(f"Need {cfg_parallel_size} GPUs, but OPENPI_E2E_GPUS only provided {gpus}.")
- return gpus
-
- query = subprocess.check_output(
- [
- "nvidia-smi",
- "--query-gpu=index,memory.used",
- "--format=csv,noheader,nounits",
- ],
- text=True,
- )
- gpu_rows = []
- for line in query.strip().splitlines():
- gpu_index, used_mb = [part.strip() for part in line.split(",", maxsplit=1)]
- gpu_rows.append((int(used_mb), gpu_index))
- gpu_rows.sort()
- gpus = [gpu_index for _, gpu_index in gpu_rows[: max(cfg_parallel_size, 1)]]
- if len(gpus) < cfg_parallel_size:
- raise RuntimeError(
- f"Need {cfg_parallel_size} GPUs for cfg_parallel_size={cfg_parallel_size}, "
- f"but found only {len(gpus)} candidates."
- )
- return gpus
-
-
-def _torchrun_argv(script: str, port: int) -> list[str]:
- return [
- sys.executable,
- "-m",
- "torch.distributed.run",
- "--standalone",
- "--nproc_per_node",
- str(_cfg_parallel_size()),
- script,
- "--port",
- str(port),
- "--model_path",
- str(CHECKPOINT_DIR),
- ]
-
-
-def _run_upstream_service(port: int, log_path: Path) -> subprocess.Popen[str]:
- env = os.environ.copy()
- env.setdefault("PYTHONPATH", "")
- env["PYTHONPATH"] = f"{Path.cwd()}:{DREAMZERO_REPO}:{env['PYTHONPATH']}".rstrip(":")
- env["CUDA_VISIBLE_DEVICES"] = ",".join(_pick_test_gpus())
- env.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")
- env["ATTENTION_BACKEND"] = "torch"
- env.setdefault("ENABLE_TENSORRT", "false")
- env["ENABLE_DIT_CACHE"] = "false"
- env["NUM_DIT_STEPS"] = "16"
- env["DYNAMIC_CACHE_SCHEDULE"] = "false"
- argv = _torchrun_argv(
- str(Path("tests/dreamzero/upstream/upstream_socket_server_no_compile.py")),
- port,
- )
- log_file = log_path.open("w")
- proc = subprocess.Popen(
- argv,
- stdout=log_file,
- stderr=subprocess.STDOUT,
- text=True,
- cwd=str(Path.cwd()),
- env=env,
- )
- proc._codex_log_file = log_file # type: ignore[attr-defined]
- return proc
-
-
-def _run_vllm_service(port: int, log_path: Path) -> subprocess.Popen[str]:
- env = os.environ.copy()
- gpus = _pick_test_gpus()
- cfg_parallel_size = _cfg_parallel_size()
- if cfg_parallel_size > len(gpus):
- raise RuntimeError(
- f"cfg_parallel_size={cfg_parallel_size} requires at least {cfg_parallel_size} GPUs, but only got {gpus}."
- )
- env["CUDA_VISIBLE_DEVICES"] = ",".join(gpus[:cfg_parallel_size])
- env.setdefault("ATTENTION_BACKEND", "torch")
- env.setdefault("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
- env.setdefault("MASTER_PORT", str(get_open_port()))
- argv = [
- _vllm_executable(),
- "serve",
- VLLM_MODEL,
- "--omni",
- "--host",
- "127.0.0.1",
- "--port",
- str(port),
- "--served-model-name",
- "dreamzero-droid",
- "--enforce-eager",
- ]
- if cfg_parallel_size > 1:
- argv.extend(["--cfg-parallel-size", str(cfg_parallel_size)])
- log_file = log_path.open("w")
- proc = subprocess.Popen(
- argv,
- stdout=log_file,
- stderr=subprocess.STDOUT,
- text=True,
- env=env,
- cwd=str(Path.cwd()),
- )
- proc._codex_log_file = log_file # type: ignore[attr-defined]
- return proc
-
-
-def _stop_process(proc: subprocess.Popen[str]) -> None:
- log_file = getattr(proc, "_codex_log_file", None)
- if proc.poll() is None:
- proc.terminate()
- try:
- proc.wait(timeout=30)
- except subprocess.TimeoutExpired: # pragma: no cover - cleanup path
- proc.kill()
- proc.wait(timeout=10)
- if log_file is not None:
- log_file.close()
-
-
-def _build_obs_sequence() -> tuple[dict, dict]:
- camera_frames = dreamzero_client.load_camera_frames()
- chunks = dreamzero_client.build_frame_schedule(
- min(v.shape[0] for v in camera_frames.values()),
- 1,
- )
- obs0 = dreamzero_client._make_obs_from_video(camera_frames, [0], PROMPT, SESSION_ID)
- obs1 = dreamzero_client._make_obs_from_video(camera_frames, chunks[0], PROMPT, SESSION_ID)
- return obs0, obs1
-
-
-def _wait_for_client_ready(client_factory, timeout_s: float, proc=None, log_path: Path | None = None):
- deadline = time.time() + timeout_s
- last_err: Exception | None = None
- while time.time() < deadline:
- if proc is not None and proc.poll() is not None:
- details = ""
- if log_path is not None and log_path.exists():
- details = log_path.read_text(errors="replace")[-8000:]
- raise RuntimeError(f"Service exited before becoming ready with code {proc.returncode}.\n{details}")
- try:
- return client_factory()
- except Exception as exc: # pragma: no cover - retry path
- last_err = exc
- time.sleep(1)
- raise TimeoutError(f"Timed out waiting for websocket service: {last_err}")
-
-
-def _collect_outputs_with_client(client) -> tuple[dict, list[np.ndarray]]:
- metadata = client.get_server_metadata()
- obs0, obs1 = _build_obs_sequence()
- outputs = [
- client.infer(dict(obs0)),
- client.infer(dict(obs1)),
- ]
- assert _normalize_reset_response(client.reset({})) == "reset successful"
- outputs.append(client.infer(dict(obs0)))
- client._ws.close()
- return metadata, outputs
-
-
-def _normalize_reset_response(response) -> str:
- if isinstance(response, str):
- return response
- decoded = msgpack_numpy.unpackb(response)
- if isinstance(decoded, dict):
- return str(decoded.get("status"))
- return str(decoded)
-
-
-def _normalize_metadata(metadata: dict) -> dict:
- normalized = dict(metadata)
- if isinstance(normalized.get("image_resolution"), tuple):
- normalized["image_resolution"] = list(normalized["image_resolution"])
- return normalized
-
-
-def _assert_logs_clean(log_path: Path) -> None:
- text = log_path.read_text(errors="replace")
- if "SignalException: Process" in text and "got signal: 15" in text:
- text = text.split("Traceback (most recent call last):", 1)[0]
- assert "Traceback" not in text, text
- assert "RuntimeError:" not in text, text
-
-
-def _assert_upstream_log_matches_vllm_baseline(log_path: Path) -> None:
- text = log_path.read_text(errors="replace")
- assert "DIT Compute Steps 8 steps" not in text, text
- assert "DIT Compute Steps 16 steps" in text, text
-
-
-def test_openpi_service_matches_upstream_server_noncompile(tmp_path: Path) -> None:
- expected_metadata = {
- "image_resolution": [180, 320],
- "n_external_cameras": 2,
- "needs_wrist_camera": True,
- "needs_stereo_camera": False,
- "needs_session_id": True,
- "action_space": "joint_position",
- }
-
- upstream_port = get_open_port()
- upstream_log = tmp_path / "dreamzero_upstream.log"
- upstream_proc = _run_upstream_service(upstream_port, upstream_log)
- try:
- upstream_client = _wait_for_client_ready(
- lambda: WebsocketClientPolicy(host="127.0.0.1", port=upstream_port),
- timeout_s=SERVICE_READY_TIMEOUT_S,
- proc=upstream_proc,
- log_path=upstream_log,
- )
- upstream_metadata, upstream_outputs = _collect_outputs_with_client(upstream_client)
- finally:
- _stop_process(upstream_proc)
- _assert_logs_clean(upstream_log)
- _assert_upstream_log_matches_vllm_baseline(upstream_log)
-
- vllm_port = get_open_port()
- vllm_log = tmp_path / "vllm_openpi.log"
- vllm_proc = _run_vllm_service(vllm_port, vllm_log)
- try:
- vllm_client = _wait_for_client_ready(
- lambda: OpenPIWebsocketClientPolicy(host="127.0.0.1", port=vllm_port),
- timeout_s=SERVICE_READY_TIMEOUT_S,
- proc=vllm_proc,
- log_path=vllm_log,
- )
- vllm_metadata, vllm_outputs = _collect_outputs_with_client(vllm_client)
- finally:
- _stop_process(vllm_proc)
- _assert_logs_clean(vllm_log)
-
- assert _normalize_metadata(upstream_metadata) == expected_metadata
- assert _normalize_metadata(vllm_metadata) == expected_metadata
-
- for idx, (actual, expected) in enumerate(zip(vllm_outputs, upstream_outputs, strict=True)):
- np.testing.assert_allclose(
- actual,
- expected,
- rtol=1e-2,
- atol=1e-3,
- err_msg=f"OpenPI step {idx} output mismatch",
- )
diff --git a/tests/e2e/accuracy/conftest.py b/tests/e2e/accuracy/conftest.py
index 6c0258aa2..750a81095 100644
--- a/tests/e2e/accuracy/conftest.py
+++ b/tests/e2e/accuracy/conftest.py
@@ -1,54 +1,15 @@
from __future__ import annotations
-import os
-import subprocess
-from contextlib import contextmanager
-from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
import pytest
import requests
-import torch
from PIL import Image
-from tests.helpers.runtime import OmniServer, OmniServerParams
-
def pytest_addoption(parser):
group = parser.getgroup("accuracy-e2e")
- group.addoption("--gebench-root", action="store", default=None, help="Local GEBench dataset root")
- group.addoption("--gedit-root", action="store", default=None, help="Local GEdit-Bench dataset root")
- group.addoption(
- "--gebench-model", action="store", default="Qwen/Qwen-Image-2512", help="Generate model for GEBench smoke"
- )
- group.addoption(
- "--gedit-model", action="store", default="Qwen/Qwen-Image-Edit", help="Generate model for GEdit-Bench smoke"
- )
- group.addoption(
- "--accuracy-judge-model",
- action="store",
- default="QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ",
- help="Judge model path",
- )
- group.addoption("--accuracy-gpu", action="store", default="0", help="Single GPU id used sequentially")
- group.addoption("--gebench-port", action="store", type=int, default=8093, help="Generate port for GEBench")
- group.addoption("--gedit-port", action="store", type=int, default=8093, help="Generate port for GEdit-Bench")
- group.addoption(
- "--gebench-samples-per-type",
- action="store",
- type=int,
- default=10,
- help="Balanced sample count per GEBench type",
- )
- group.addoption(
- "--gedit-samples-per-group",
- action="store",
- type=int,
- default=20,
- help="Balanced sample count per GEdit task group",
- )
- group.addoption("--accuracy-workers", action="store", type=int, default=1, help="Worker count for accuracy benches")
group.addoption(
"--wan22-i2v-image-source",
action="store",
@@ -77,99 +38,6 @@ def pytest_addoption(parser):
)
-def _hf_cache_root() -> Path:
- return Path(os.environ.get("HF_HOME", "/root/.cache/huggingface"))
-
-
-def _dataset_cache_dirs(dataset_id: str) -> list[Path]:
- cache_root = _hf_cache_root() / "hub" / f"datasets--{dataset_id.replace('/', '--')}" / "snapshots"
- if not cache_root.exists():
- return []
- return sorted(
- (path for path in cache_root.iterdir() if path.is_dir()), key=lambda path: path.stat().st_mtime, reverse=True
- )
-
-
-def _ensure_dataset_snapshot(dataset_id: str) -> Path:
- candidates = _dataset_cache_dirs(dataset_id)
- if candidates:
- return candidates[0]
-
- subprocess.run(
- ["hf", "download", "--repo-type", "dataset", dataset_id],
- check=True,
- )
- candidates = _dataset_cache_dirs(dataset_id)
- if not candidates:
- raise FileNotFoundError(
- f"Dataset {dataset_id} was downloaded but no snapshot was found under {_hf_cache_root()}"
- )
- return candidates[0]
-
-
-def _resolve_dataset_root(request: pytest.FixtureRequest, option_name: str, dataset_id: str) -> Path:
- value = request.config.getoption(option_name)
- if value:
- path = Path(value)
- if not path.exists():
- pytest.skip(f"Dataset path does not exist: {path}")
- return path
- return _ensure_dataset_snapshot(dataset_id)
-
-
-@dataclass
-class AccuracyServerConfig:
- generate_params: OmniServerParams
- judge_params: OmniServerParams
- run_level: str
- model_prefix: str
-
- @contextmanager
- def generate_server(self):
- params = self.generate_params
- model = self.model_prefix + params.model
- server_args = params.server_args or []
- if params.use_omni and params.stage_init_timeout is not None:
- server_args = ["--stage-init-timeout", str(params.stage_init_timeout), *server_args]
- with OmniServer(
- model,
- server_args,
- port=params.port,
- env_dict=params.env_dict,
- use_omni=params.use_omni,
- ) as server:
- yield server
-
- @contextmanager
- def judge_server(self):
- params = self.judge_params
- model = self.model_prefix + params.model
- server_args = params.server_args or []
- with OmniServer(
- model,
- server_args,
- port=params.port,
- env_dict=params.env_dict,
- use_omni=params.use_omni,
- ) as server:
- yield server
-
-
-@pytest.fixture(scope="session")
-def gebench_dataset_root(request: pytest.FixtureRequest) -> Path:
- return _resolve_dataset_root(request, "gebench_root", "stepfun-ai/GEBench")
-
-
-@pytest.fixture(scope="session")
-def gedit_dataset_root(request: pytest.FixtureRequest) -> Path:
- return _resolve_dataset_root(request, "gedit_root", "stepfun-ai/GEdit-Bench")
-
-
-@pytest.fixture(scope="session")
-def accuracy_workers(request: pytest.FixtureRequest) -> int:
- return int(request.config.getoption("accuracy_workers"))
-
-
@pytest.fixture(scope="session")
def wan22_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
value = request.config.getoption("wan22_i2v_image_source")
@@ -192,16 +60,6 @@ def hunyuanvideo15_online_timeout_seconds(request: pytest.FixtureRequest) -> int
return int(request.config.getoption("hunyuanvideo15_online_timeout_seconds"))
-@pytest.fixture(scope="session")
-def gebench_samples_per_type(request: pytest.FixtureRequest) -> int:
- return int(request.config.getoption("gebench_samples_per_type"))
-
-
-@pytest.fixture(scope="session")
-def gedit_samples_per_group(request: pytest.FixtureRequest) -> int:
- return int(request.config.getoption("gedit_samples_per_group"))
-
-
@pytest.fixture(scope="session")
def accuracy_artifact_root() -> Path:
root = Path(__file__).resolve().parent / "artifacts"
@@ -209,12 +67,6 @@ def accuracy_artifact_root() -> Path:
return root
-@pytest.fixture(scope="session")
-def accuracy_assets_root() -> Path:
- root = Path(__file__).resolve().parent / "assets"
- return root
-
-
@pytest.fixture(scope="session")
def qwen_bear_image(accuracy_artifact_root: Path):
"""Download the Qwen bear image from the URL and save it to the accuracy artifact root."""
@@ -249,80 +101,3 @@ def rabbit_image(accuracy_artifact_root: Path):
image.save(image_path)
yield image
image.close()
-
-
-def _build_accuracy_server_config(
- *,
- generate_model: str,
- judge_model: str,
- shared_gpu: str,
- port: int,
- run_level: str,
- model_prefix: str,
-) -> AccuracyServerConfig:
- if torch.accelerator.device_count() < 1:
- pytest.skip("Need at least 1 CUDA GPU for accuracy benchmark smoke tests.")
-
- if not generate_model:
- pytest.skip("No generate model configured for accuracy benchmark test.")
- generate_server_args = ["--num-gpus", "1"]
- judge_server_args = [
- "--max-model-len",
- "32768",
- "--gpu-memory-utilization",
- "0.8",
- ]
-
- judge_env = {"CUDA_VISIBLE_DEVICES": shared_gpu}
-
- return AccuracyServerConfig(
- generate_params=OmniServerParams(
- model=generate_model,
- port=port,
- server_args=generate_server_args,
- env_dict={"CUDA_VISIBLE_DEVICES": shared_gpu},
- use_omni=True,
- stage_init_timeout=300,
- ),
- judge_params=OmniServerParams(
- model=judge_model,
- port=port,
- server_args=judge_server_args,
- env_dict=judge_env,
- use_omni=False,
- ),
- run_level=run_level,
- model_prefix=model_prefix,
- )
-
-
-@pytest.fixture
-def gebench_accuracy_servers(
- request: pytest.FixtureRequest,
- run_level: str,
- model_prefix: str,
-) -> AccuracyServerConfig:
- return _build_accuracy_server_config(
- generate_model=request.config.getoption("gebench_model"),
- judge_model=request.config.getoption("accuracy_judge_model"),
- shared_gpu=str(request.config.getoption("accuracy_gpu")),
- port=int(request.config.getoption("gebench_port")),
- run_level=run_level,
- model_prefix=model_prefix,
- )
-
-
-@pytest.fixture
-def gedit_accuracy_servers(
- request: pytest.FixtureRequest,
- run_level: str,
- model_prefix: str,
-) -> AccuracyServerConfig:
- return _build_accuracy_server_config(
- generate_model=request.config.getoption("gedit_model"),
- judge_model=request.config.getoption("accuracy_judge_model"),
- shared_gpu=str(request.config.getoption("accuracy_gpu")),
- port=int(request.config.getoption("gedit_port")),
- run_level=run_level,
- model_prefix=model_prefix,
- )
diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py
index ebe635ea8..f3816136c 100644
--- a/tests/e2e/accuracy/helpers.py
+++ b/tests/e2e/accuracy/helpers.py
@@ -3,8 +3,12 @@ import os
import re
import shutil
import subprocess
+import sys
import time
+from abc import ABC, abstractmethod
from base64 import b64decode, b64encode
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
from fractions import Fraction
from hashlib import sha1
from io import BytesIO
@@ -53,8 +57,6 @@ def env_to_apply_ftfy_mock_in_subproc(env: dict[str, str] | None = None) -> dict
def reset_artifact_dir(path: Path) -> Path:
- import shutil
-
if path.exists():
shutil.rmtree(path)
path.mkdir(parents=True, exist_ok=True)
@@ -609,3 +611,344 @@ def download_images(urls: list[str]) -> list[Image.Image]:
images.append(image)
return images
+
+
+_DREAMZERO_UPSTREAM_SERVER_SCRIPT = (
+ Path(__file__).resolve().parent / "upstream_mock" / "upstream_socket_server_no_compile.py"
+)
+_DREAMZERO_DEFAULT_MODEL = "GEAR-Dreams/DreamZero-DROID"
+_DREAMZERO_DEFAULT_READY_TIMEOUT_S = int(os.environ.get("OPENPI_SERVICE_READY_TIMEOUT_S", "900"))
+_DREAMZERO_DEFAULT_PROMPT = (
+ "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
+)
+_DREAMZERO_DEFAULT_SESSION_ID = "openpi-e2e-parity-session"
+_DREAMZERO_OPENPI_VLLM_PATH = "/v1/realtime/robot/openpi"
+
+
+def require_dreamzero_parity_env() -> tuple[Path, Path]:
+ repo_env = os.environ.get("DREAMZERO_REPO")
+ if repo_env is None:
+ raise ValueError("DREAMZERO_REPO environment variable is required.")
+ repo = Path(repo_env).expanduser()
+ if not repo.is_dir():
+ raise FileNotFoundError(f"DREAMZERO_REPO does not exist: {repo}")
+ checkpoint_dir = repo / "checkpoints" / "dreamzero"
+ if not checkpoint_dir.is_dir():
+ raise FileNotFoundError(f"DreamZero checkpoint not found: {checkpoint_dir}")
+ if str(repo) not in sys.path:
+ sys.path.insert(0, str(repo))
+ return repo, checkpoint_dir
+
+
+def load_dreamzero_upstream_modules():
+ require_dreamzero_parity_env()
+ import test_client_AR as dreamzero_upstream_client
+ from eval_utils.policy_client import WebsocketClientPolicy
+
+ return dreamzero_upstream_client, WebsocketClientPolicy
+
+
+def dreamzero_cfg_parallel_size() -> int:
+ return int(os.environ.get("OPENPI_E2E_CFG_PARALLEL_SIZE", "2"))
+
+
+def _dreamzero_parity_vllm_env() -> dict[str, str]:
+ from tests.helpers.env import pick_least_used_device_indices
+ from tests.helpers.runtime import get_open_port
+
+ parallel_size = dreamzero_cfg_parallel_size()
+ gpus = [str(index) for index in pick_least_used_device_indices(parallel_size)]
+ if parallel_size > len(gpus):
+ raise RuntimeError(
+ f"cfg_parallel_size={parallel_size} requires at least {parallel_size} GPUs, but only got {gpus}."
+ )
+ return {
+ "CUDA_VISIBLE_DEVICES": ",".join(gpus[:parallel_size]),
+ "ATTENTION_BACKEND": "torch",
+ "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
+ "VLLM_DISABLE_COMPILE_CACHE": "1",
+ "MASTER_PORT": str(get_open_port()),
+ }
+
+
+def dreamzero_parity_vllm_params(model: str | None = None):
+ from tests.helpers.runtime import OmniServerParams
+ from tests.helpers.stage_config import get_deploy_config_path
+
+ parallel_size = dreamzero_cfg_parallel_size()
+ server_args = [
+ "--served-model-name",
+ "dreamzero-droid",
+ "--enforce-eager",
+ ]
+ if parallel_size > 1:
+ server_args.extend(["--cfg-parallel-size", str(parallel_size)])
+ return OmniServerParams(
+ model=model or os.environ.get("VLLM_DREAMZERO_MODEL", _DREAMZERO_DEFAULT_MODEL),
+ stage_config_path=get_deploy_config_path("dreamzero_tp1_cfg2.yaml"),
+ server_args=server_args,
+ env_dict=_dreamzero_parity_vllm_env(),
+ )
+
+
+def _dreamzero_omni_server_from_params(params):
+ from tests.helpers.runtime import OmniServer
+
+ server_args = list(params.server_args or [])
+ if params.stage_config_path is not None:
+ server_args += ["--stage-configs-path", params.stage_config_path]
+ return OmniServer(
+ params.model,
+ server_args,
+ port=params.port,
+ env_dict=params.env_dict,
+ use_omni=params.use_omni,
+ )
+
+
+@contextmanager
+def dreamzero_parity_vllm_client(
+ params=None,
+ *,
+ ready_timeout_s: float = _DREAMZERO_DEFAULT_READY_TIMEOUT_S,
+) -> Iterator["DreamZeroParityClient"]:
+ """Start vLLM via ``OmniServer`` and return a ready OpenPI parity client."""
+ with _dreamzero_omni_server_from_params(params or dreamzero_parity_vllm_params()) as server:
+ client = _dreamzero_wait_for_client_ready(
+ lambda: _dreamzero_create_openpi_vllm_client(server.host, server.port),
+ timeout_s=ready_timeout_s,
+ proc=server.proc,
+ )
+ yield DreamZeroParityClient(client)
+
+
+def _dreamzero_create_openpi_vllm_client(host: str, port: int) -> Any:
+ from openpi_client import msgpack_numpy
+
+ _, WebsocketClientPolicy = load_dreamzero_upstream_modules()
+
+ class _OpenPIWebsocketClientPolicy(WebsocketClientPolicy):
+ def __init__(
+ self,
+ host: str = "127.0.0.1",
+ port: int = 8000,
+ path: str = _DREAMZERO_OPENPI_VLLM_PATH,
+ ) -> None:
+ self._uri = f"ws://{host}:{port}{path}"
+ self._packer = msgpack_numpy.Packer()
+ self._ws, self._server_metadata = self._wait_for_server()
+
+ return _OpenPIWebsocketClientPolicy(host=host, port=port)
+
+
+def _dreamzero_stop_process(proc: subprocess.Popen[str]) -> None:
+ log_file = getattr(proc, "_codex_log_file", None)
+ if proc.poll() is None:
+ proc.terminate()
+ try:
+ proc.wait(timeout=30)
+ except subprocess.TimeoutExpired: # pragma: no cover - cleanup path
+ proc.kill()
+ proc.wait(timeout=10)
+ if log_file is not None:
+ log_file.close()
+
+
+def _dreamzero_wait_for_client_ready(
+ client_factory: Callable[[], Any],
+ *,
+ timeout_s: float,
+ proc: subprocess.Popen[str] | None = None,
+ log_path: Path | None = None,
+) -> Any:
+ deadline = time.time() + timeout_s
+ last_err: Exception | None = None
+ while time.time() < deadline:
+ if proc is not None and proc.poll() is not None:
+ details = ""
+ if log_path is not None and log_path.exists():
+ details = log_path.read_text(errors="replace")[-8000:]
+ raise RuntimeError(f"Service exited before becoming ready with code {proc.returncode}.\n{details}")
+ try:
+ return client_factory()
+ except Exception as exc: # pragma: no cover - retry path
+ last_err = exc
+ time.sleep(1)
+ raise TimeoutError(f"Timed out waiting for websocket service: {last_err}")
+
+
+def _dreamzero_normalize_reset_response(response: Any) -> str:
+ from openpi_client import msgpack_numpy
+
+ if isinstance(response, str):
+ return response
+ decoded = msgpack_numpy.unpackb(response)
+ if isinstance(decoded, dict):
+ return str(decoded.get("status"))
+ return str(decoded)
+
+
+def normalize_dreamzero_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
+ normalized = dict(metadata)
+ if isinstance(normalized.get("image_resolution"), tuple):
+ normalized["image_resolution"] = list(normalized["image_resolution"])
+ return normalized
+
+
+def assert_dreamzero_logs_clean(log_path: Path) -> None:
+ text = log_path.read_text(errors="replace")
+ if "SignalException: Process" in text and "got signal: 15" in text:
+ text = text.split("Traceback (most recent call last):", 1)[0]
+ assert "Traceback" not in text, text
+ assert "RuntimeError:" not in text, text
+
+
+def assert_dreamzero_upstream_log_matches_vllm_baseline(log_path: Path) -> None:
+ text = log_path.read_text(errors="replace")
+ assert "DIT Compute Steps 8 steps" not in text, text
+ assert "DIT Compute Steps 16 steps" in text, text
+
+
+class DreamZeroParityClient:
+ """Connected OpenPI websocket client for DreamZero parity inference."""
+
+ def __init__(self, client: Any) -> None:
+ self._client = client
+
+ @property
+ def raw_client(self) -> Any:
+ return self._client
+
+ def build_obs_sequence(self) -> tuple[dict[str, Any], dict[str, Any]]:
+ dreamzero_upstream_client, _ = load_dreamzero_upstream_modules()
+ camera_frames = dreamzero_upstream_client.load_camera_frames()
+ chunks = dreamzero_upstream_client.build_frame_schedule(
+ min(v.shape[0] for v in camera_frames.values()),
+ 1,
+ )
+ obs0 = dreamzero_upstream_client._make_obs_from_video(
+ camera_frames,
+ [0],
+ _DREAMZERO_DEFAULT_PROMPT,
+ _DREAMZERO_DEFAULT_SESSION_ID,
+ )
+ obs1 = dreamzero_upstream_client._make_obs_from_video(
+ camera_frames,
+ chunks[0],
+ _DREAMZERO_DEFAULT_PROMPT,
+ _DREAMZERO_DEFAULT_SESSION_ID,
+ )
+ return obs0, obs1
+
+ def collect_parity_outputs(self) -> tuple[dict[str, Any], list[np.ndarray]]:
+ metadata = self._client.get_server_metadata()
+ obs0, obs1 = self.build_obs_sequence()
+ outputs = [
+ self._client.infer(dict(obs0)),
+ self._client.infer(dict(obs1)),
+ ]
+ assert _dreamzero_normalize_reset_response(self._client.reset({})) == "reset successful"
+ outputs.append(self._client.infer(dict(obs0)))
+ self._client._ws.close()
+ return metadata, outputs
+
+
+class _DreamZeroParityServer(ABC):
+ host = "127.0.0.1"
+
+ def __init__(
+ self,
+ *,
+ log_path: Path,
+ port: int | None = None,
+ ready_timeout_s: float = _DREAMZERO_DEFAULT_READY_TIMEOUT_S,
+ ) -> None:
+ from tests.helpers.runtime import get_open_port
+
+ self.log_path = log_path
+ self.port = get_open_port() if port is None else port
+ self.ready_timeout_s = ready_timeout_s
+ self.proc: subprocess.Popen[str] | None = None
+ self._log_file = None
+
+ @abstractmethod
+ def _build_argv(self) -> list[str]:
+ raise NotImplementedError
+
+ @abstractmethod
+ def _build_env(self) -> dict[str, str]:
+ raise NotImplementedError
+
+ @abstractmethod
+ def _create_client(self) -> Any:
+ raise NotImplementedError
+
+ def _start(self) -> None:
+ self.log_path.parent.mkdir(parents=True, exist_ok=True)
+ self._log_file = self.log_path.open("w")
+ self.proc = subprocess.Popen(
+ self._build_argv(),
+ stdout=self._log_file,
+ stderr=subprocess.STDOUT,
+ text=True,
+ cwd=str(Path.cwd()),
+ env=self._build_env(),
+ )
+ self.proc._codex_log_file = self._log_file # type: ignore[attr-defined]
+
+ def __enter__(self) -> DreamZeroParityClient:
+ self._start()
+ client = _dreamzero_wait_for_client_ready(
+ self._create_client,
+ timeout_s=self.ready_timeout_s,
+ proc=self.proc,
+ log_path=self.log_path,
+ )
+ return DreamZeroParityClient(client)
+
+ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+ if self.proc is not None:
+ _dreamzero_stop_process(self.proc)
+ self.proc = None
+
+
+class DreamZeroUpstreamServer(_DreamZeroParityServer):
+ """Official DreamZero upstream reference server launched via torchrun."""
+
+ def _build_argv(self) -> list[str]:
+ _, checkpoint_dir = require_dreamzero_parity_env()
+ return [
+ sys.executable,
+ "-m",
+ "torch.distributed.run",
+ "--standalone",
+ "--nproc_per_node",
+ str(dreamzero_cfg_parallel_size()),
+ str(_DREAMZERO_UPSTREAM_SERVER_SCRIPT),
+ "--port",
+ str(self.port),
+ "--model_path",
+ str(checkpoint_dir),
+ ]
+
+ def _build_env(self) -> dict[str, str]:
+ from tests.helpers.env import pick_least_used_device_indices
+
+ dreamzero_repo, _ = require_dreamzero_parity_env()
+ env = os.environ.copy()
+ env.setdefault("PYTHONPATH", "")
+ env["PYTHONPATH"] = f"{Path.cwd()}:{dreamzero_repo}:{env['PYTHONPATH']}".rstrip(":")
+ env["CUDA_VISIBLE_DEVICES"] = ",".join(
+ str(index) for index in pick_least_used_device_indices(max(dreamzero_cfg_parallel_size(), 1))
+ )
+ env.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")
+ env["ATTENTION_BACKEND"] = "torch"
+ env.setdefault("ENABLE_TENSORRT", "false")
+ env["ENABLE_DIT_CACHE"] = "false"
+ env["NUM_DIT_STEPS"] = "16"
+ env["DYNAMIC_CACHE_SCHEDULE"] = "false"
+ return env
+
+ def _create_client(self) -> Any:
+ _, WebsocketClientPolicy = load_dreamzero_upstream_modules()
+ return WebsocketClientPolicy(host=self.host, port=self.port)
diff --git a/tests/e2e/accuracy/test_dreamzero.py b/tests/e2e/accuracy/test_dreamzero.py
new file mode 100644
index 000000000..2507530b3
--- /dev/null
+++ b/tests/e2e/accuracy/test_dreamzero.py
@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero OpenPI upstream parity accuracy test."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from tests.e2e.accuracy.helpers import (
+ DreamZeroUpstreamServer,
+ assert_dreamzero_logs_clean,
+ assert_dreamzero_upstream_log_matches_vllm_baseline,
+ dreamzero_parity_vllm_client,
+ load_dreamzero_upstream_modules,
+ normalize_dreamzero_metadata,
+)
+from tests.helpers.mark import hardware_test
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+
+EXPECTED_METADATA = {
+ "image_resolution": [180, 320],
+ "n_external_cameras": 2,
+ "needs_wrist_camera": True,
+ "needs_stereo_camera": False,
+ "needs_session_id": True,
+ "action_space": "joint_position",
+}
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_dreamzero_openpi_upstream_parity(tmp_path: Path) -> None:
+ """Assert vLLM OpenPI actions match the upstream DreamZero reference server."""
+ load_dreamzero_upstream_modules()
+
+ upstream_log = tmp_path / "dreamzero_upstream.log"
+ with DreamZeroUpstreamServer(log_path=upstream_log) as upstream_client:
+ upstream_metadata, upstream_outputs = upstream_client.collect_parity_outputs()
+ assert_dreamzero_logs_clean(upstream_log)
+ assert_dreamzero_upstream_log_matches_vllm_baseline(upstream_log)
+
+ with dreamzero_parity_vllm_client() as vllm_client:
+ vllm_metadata, vllm_outputs = vllm_client.collect_parity_outputs()
+
+ assert normalize_dreamzero_metadata(upstream_metadata) == EXPECTED_METADATA
+ assert normalize_dreamzero_metadata(vllm_metadata) == EXPECTED_METADATA
+
+ for idx, (vllm_action, upstream_action) in enumerate(zip(vllm_outputs, upstream_outputs, strict=True)):
+ np.testing.assert_allclose(
+ vllm_action,
+ upstream_action,
+ rtol=1e-2,
+ atol=1e-3,
+ err_msg=f"OpenPI step {idx} output mismatch",
+ )
diff --git a/tests/e2e/accuracy/test_gebench_h100_smoke.py b/tests/e2e/accuracy/test_gebench_h100_smoke.py
deleted file mode 100644
index b17c951e5..000000000
--- a/tests/e2e/accuracy/test_gebench_h100_smoke.py
+++ /dev/null
@@ -1,97 +0,0 @@
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import pytest
-
-from benchmarks.accuracy.text_to_image.gbench import main as gbench_main
-from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir
-from tests.helpers.mark import hardware_test
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
-
-
-@pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(reason="#3256")
-def test_gebench_h100_smoke(
- gebench_accuracy_servers,
- accuracy_artifact_root: Path,
- gebench_dataset_root: Path,
- accuracy_workers: int,
-) -> None:
- model_label = infer_model_label(gebench_accuracy_servers.generate_params.model).lower()
- output_root = reset_artifact_dir(accuracy_artifact_root / f"gebench_{model_label}")
-
- with gebench_accuracy_servers.generate_server() as generate_server:
- for data_type in ("type3", "type4"):
- assert (
- gbench_main(
- [
- "generate",
- "--dataset-root",
- str(gebench_dataset_root),
- "--output-root",
- str(output_root),
- "--base-url",
- f"http://{generate_server.host}:{generate_server.port}",
- "--model",
- generate_server.model,
- "--data-type",
- data_type,
- "--width",
- "768",
- "--height",
- "576",
- "--output-compression",
- "98",
- "--num-inference-steps",
- "8",
- "--workers",
- str(accuracy_workers),
- ]
- )
- == 0
- )
-
- with gebench_accuracy_servers.judge_server() as judge_server:
- for data_type in ("type3", "type4"):
- assert (
- gbench_main(
- [
- "evaluate",
- "--dataset-root",
- str(gebench_dataset_root),
- "--output-root",
- str(output_root),
- "--data-type",
- data_type,
- "--judge-base-url",
- f"http://{judge_server.host}:{judge_server.port}",
- "--judge-model",
- judge_server.model,
- "--judge-api-key",
- "EMPTY",
- "--workers",
- str(accuracy_workers),
- ]
- )
- == 0
- )
-
- assert gbench_main(["summarize", "--output-root", str(output_root)]) == 0
-
- summary = json.loads((output_root / "summary.json").read_text(encoding="utf-8"))
- assert "generation" in summary
- assert "evaluation" in summary
-
- for data_type in ("type3", "type4"):
- assert data_type in summary["generation"]["by_type"]
- assert summary["generation"]["by_type"][data_type]["count"] > 0
- assert data_type in summary["evaluation"]["by_type"]
- assert summary["evaluation"]["by_type"][data_type]["count"] > 0
-
- assert summary["evaluation"]["overall_mean"] >= 0.45
- assert summary["evaluation"]["by_type"]["type3"]["overall_mean"] >= 0.45
- assert summary["evaluation"]["by_type"]["type4"]["overall_mean"] >= 0.45
diff --git a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py b/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py
deleted file mode 100644
index a7831fac8..000000000
--- a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py
+++ /dev/null
@@ -1,115 +0,0 @@
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import pytest
-
-from benchmarks.accuracy.image_to_image.gedit_bench import GROUPS
-from benchmarks.accuracy.image_to_image.gedit_bench import main as gedit_main
-from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir
-from tests.helpers.mark import hardware_test
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
-
-
-@pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(reason="#3257")
-def test_gedit_bench_h100_smoke(
- gedit_accuracy_servers,
- accuracy_artifact_root: Path,
- gedit_dataset_root: Path,
- gedit_samples_per_group: int,
- accuracy_workers: int,
-) -> None:
- model_label = infer_model_label(gedit_accuracy_servers.generate_params.model).lower()
- output_root = reset_artifact_dir(accuracy_artifact_root / f"gedit_results_{model_label}")
- score_root = reset_artifact_dir(accuracy_artifact_root / f"gedit_scores_{model_label}")
- model_name = model_label
-
- with gedit_accuracy_servers.generate_server() as generate_server:
- assert (
- gedit_main(
- [
- "generate",
- "--dataset-ref",
- str(gedit_dataset_root),
- "--output-root",
- str(output_root),
- "--base-url",
- f"http://{generate_server.host}:{generate_server.port}",
- "--model",
- generate_server.model,
- "--model-name",
- model_name,
- "--task-type",
- "all",
- "--instruction-language",
- "all",
- "--samples-per-group",
- str(gedit_samples_per_group),
- "--workers",
- str(accuracy_workers),
- ]
- )
- == 0
- )
-
- with gedit_accuracy_servers.judge_server() as judge_server:
- assert (
- gedit_main(
- [
- "evaluate",
- "--dataset-ref",
- str(gedit_dataset_root),
- "--output-root",
- str(output_root),
- "--model-name",
- model_name,
- "--save-dir",
- str(score_root),
- "--task-type",
- "all",
- "--instruction-language",
- "all",
- "--judge-base-url",
- f"http://{judge_server.host}:{judge_server.port}",
- "--judge-model",
- judge_server.model,
- "--judge-api-key",
- "EMPTY",
- "--samples-per-group",
- str(gedit_samples_per_group),
- "--workers",
- str(accuracy_workers),
- ]
- )
- == 0
- )
-
- csv_path = score_root / f"{model_name}_all_all_vie_score.csv"
- assert gedit_main(["summarize", "--csv-path", str(csv_path), "--language", "all"]) == 0
-
- summary_path = score_root / f"{model_name}_all_all_summary.json"
- summary = json.loads(summary_path.read_text(encoding="utf-8"))
- assert set(summary["languages"]) == {"en", "cn"}
-
- for language in ["en", "cn"]:
- language_summary = summary["languages"][language]
- assert language_summary["overall"]["count"] is not None
- assert language_summary["intersection"]["count"] is not None
- assert language_summary["overall"]["Q_SC"] is not None
- assert language_summary["overall"]["Q_PQ"] is not None
- assert language_summary["overall"]["Q_O"] is not None
-
- for group in GROUPS:
- group_summary = language_summary["by_group"][group]
- assert set(group_summary) == {"count", "Q_SC", "Q_PQ", "Q_O"}
-
- assert summary["languages"]["en"]["overall"]["Q_SC"] >= 6.95
- assert summary["languages"]["en"]["overall"]["Q_PQ"] >= 5.8
- assert summary["languages"]["en"]["overall"]["Q_O"] >= 6.15
- assert summary["languages"]["cn"]["overall"]["Q_SC"] >= 6.9
- assert summary["languages"]["cn"]["overall"]["Q_PQ"] >= 5.7
- assert summary["languages"]["cn"]["overall"]["Q_O"] >= 6.1
diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py
index e2d84b13d..8c4570fe3 100644
--- a/tests/e2e/accuracy/test_hunyuan_image3.py
+++ b/tests/e2e/accuracy/test_hunyuan_image3.py
@@ -28,6 +28,7 @@ from tests.e2e.accuracy.helpers import (
model_output_dir,
)
from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniRunner, OmniServer
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
build_prompt_tokens,
@@ -94,6 +95,7 @@ def _default_ar_dit_devices() -> tuple[str, str]:
AR_DEVICES, DIT_DEVICES = _default_ar_dit_devices()
MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
+HUNYUAN_IMAGE_REF_PATH = get_asset_path("hunyuan/hunyuan_image_ref.png")
NUM_INFERENCE_STEPS = 50
GUIDANCE_SCALE = 2.5
@@ -414,7 +416,7 @@ def _run_online(stage_configs_path: str, output_path: Path) -> tuple[Image.Image
torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE,
reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)",
)
-def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
+def test_image_to_image_alignment_online(accuracy_artifact_root: Path) -> None:
"""Online API test: same pipeline, same seed as offline → PSNR >= 10 dB."""
if importlib.util.find_spec("FlagEmbedding") is None:
raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")
@@ -431,7 +433,7 @@ def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_
scorer = SemanticSimilarityScorer()
clip_scorer = CLIPScorer()
cot_results = scorer.text_similarity(online_cot, COT_REF)
- image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB")
+ image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB")
image_clip_score = clip_scorer.image_image_score(online_image, image_ref)
ssim_value, psnr_value = compute_image_ssim_psnr(prediction=online_image, reference=image_ref, compare_mode="RGB")
@@ -529,7 +531,7 @@ def _run_dit_model(
torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE,
reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)",
)
-def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
+def test_image_to_image_alignment(accuracy_artifact_root: Path) -> None:
if importlib.util.find_spec("FlagEmbedding") is None:
raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")
from tabulate import tabulate # lazy import
@@ -545,7 +547,7 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_
scorer = SemanticSimilarityScorer()
clip_scorer = CLIPScorer()
cot_results = scorer.text_similarity(omni_cot, COT_REF)
- image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB")
+ image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB")
image_clip_score = clip_scorer.image_image_score(omni_image, image_ref)
ssim_value, psnr_value = compute_image_ssim_psnr(prediction=omni_image, reference=image_ref, compare_mode="RGB")
diff --git a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py
index 197044ef3..34131dc5a 100644
--- a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py
+++ b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py
@@ -15,6 +15,7 @@ from PIL import Image
from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServer
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
@@ -33,8 +34,8 @@ P99_THRESHOLD = 3e-1
SSIM_THRESHOLD = 0.97
PSNR_THRESHOLD = 30.0
+BASELINE_PATH = get_asset_path("hunyuan/hunyuan_baseline.png")
_REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
-BASELINE_PATH = _REPO_ROOT / "tests" / "assets" / "hunyuan" / "hunyuan_baseline.png"
_OFFLINE_SCRIPT = _REPO_ROOT / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
# DiT-only deploy config with trust_remote_code (based on hunyuan_image3_dit.yaml).
diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py
index 338758cc0..9ee9c2c87 100644
--- a/tests/e2e/accuracy/test_qwen_image_edit.py
+++ b/tests/e2e/accuracy/test_qwen_image_edit.py
@@ -14,6 +14,7 @@ from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServer
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
@@ -46,10 +47,9 @@ EDIT_2511_PROMPT = (
"（帽子紧贴头顶，未贴合部分可调整帽子摆放角度达到完全贴合）， 头部与身体过渡必须自然，无拼接痕迹，"
"无额外元素出现在头部四周，纯白背景，不要出现绿色圆圈"
)
-EDIT_2511_ASSET_DIR = Path(__file__).resolve().parent / "assets"
EDIT_2511_IMAGE_PATHS = [
- EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test1.png",
- EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test2.png",
+ get_asset_path("qwen_image_edit/qwen_image_edit_2511_test1.png"),
+ get_asset_path("qwen_image_edit/qwen_image_edit_2511_test2.png"),
]
PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure."
diff --git a/tests/dreamzero/upstream/upstream_socket_server_no_compile.py b/tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py
similarity index 98%
rename from tests/dreamzero/upstream/upstream_socket_server_no_compile.py
rename to tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py
index a9ef7a3d0..e6fccdd23 100644
--- a/tests/dreamzero/upstream/upstream_socket_server_no_compile.py
+++ b/tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py
@@ -24,7 +24,7 @@ PyTorch SDPA (`ATTENTION_BACKEND=torch`) for this subprocess only.
Usage:
PYTHONPATH="${DREAMZERO_REPO}" \\
.venv/bin/python -m torch.distributed.run --standalone --nproc_per_node=2 \\
- tests/dreamzero/upstream/upstream_socket_server_no_compile.py --port 18081 \\
+ tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py --port 18081 \\
--model_path "${DREAMZERO_REPO}/checkpoints/dreamzero"
"""
diff --git a/tests/comfyui/conftest.py b/tests/e2e/features/comfyui/conftest.py
similarity index 97%
rename from tests/comfyui/conftest.py
rename to tests/e2e/features/comfyui/conftest.py
index 4280d3506..c3ac99bf2 100644
--- a/tests/comfyui/conftest.py
+++ b/tests/e2e/features/comfyui/conftest.py
@@ -25,7 +25,7 @@ def _setup_comfyui_test_environment():
"""Set up the test environment for ComfyUI plugin testing."""
# === Add ComfyUI plugin path to allow importing comfyui_vllm_omni ===
_COMFYUI_PLUGIN_PATH = os.path.abspath(
- os.path.join(os.path.dirname(__file__), "..", "..", "apps", "ComfyUI-vLLM-Omni")
+ os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "apps", "ComfyUI-vLLM-Omni")
)
if not os.path.isdir(_COMFYUI_PLUGIN_PATH):
raise FileNotFoundError(
diff --git a/tests/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py
similarity index 99%
rename from tests/comfyui/test_comfyui_integration.py
rename to tests/e2e/features/comfyui/test_comfyui_integration.py
index e15098b4e..2024cce8b 100644
--- a/tests/comfyui/test_comfyui_integration.py
+++ b/tests/e2e/features/comfyui/test_comfyui_integration.py
@@ -6,6 +6,8 @@ It ensures that
2. The sampling parameters are correctly passed from the node to AsyncOmni through the API layer.
"""
+from __future__ import annotations
+
import multiprocessing
import time
import traceback
@@ -52,7 +54,7 @@ class ServerCase(NamedTuple):
class SamplingCase(NamedTuple):
"""Parametrizing the input sampling parameters."""
- kind: "SamplingKind"
+ kind: SamplingKind
sampling_params: dict | list[dict] | None
lora: dict | None = None
diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py
similarity index 96%
rename from tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py
rename to tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py
index dc701af7a..04e83caeb 100644
--- a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py
+++ b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py
@@ -16,7 +16,7 @@ exercises when it calls ``await self.engine.list_loras()`` before
dispatching a generation request.
Usage:
- pytest tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py -v -s
+ pytest tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py -v -s
"""
from __future__ import annotations
@@ -34,12 +34,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
MODEL = "tiny-random/Qwen-Image"
-CUSTOM_PIPELINE_CLASS = (
- "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
-)
-WORKER_EXTENSION_CLASS = (
- "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
-)
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
def _create_inline_engine() -> AsyncOmni:
diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
similarity index 96%
rename from tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py
rename to tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
index 02519f90f..b24f364e4 100644
--- a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py
+++ b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
@@ -24,12 +24,8 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
-CUSTOM_PIPELINE_CLASS = (
- "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
-)
-WORKER_EXTENSION_CLASS = (
- "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
-)
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
# Use your specified HF repo name/path directly here
MODEL = "tiny-random/Qwen-Image"
diff --git a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py b/tests/e2e/features/custom_pipeline/test_worker_extension.py
similarity index 59%
rename from tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py
rename to tests/e2e/features/custom_pipeline/test_worker_extension.py
index 16dddde40..da7b0ac17 100644
--- a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py
+++ b/tests/e2e/features/custom_pipeline/test_worker_extension.py
@@ -7,37 +7,19 @@ from contextlib import ExitStack
import pytest
-from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import (
+from tests.e2e.features.helpers.custom_pipeline import (
QwenImagePipelineWithLogProbForTest,
)
-from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
- vLLMOmniColocateWorkerExtensionForTest,
-)
from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
from vllm_omni.entrypoints.async_omni import AsyncOmni
CUSTOM_PIPELINE_CLASS = QwenImagePipelineWithLogProbForTest
-WORKER_EXTENSION_CLASS = (
- "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
-)
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
MODEL = "tiny-random/Qwen-Image"
pytestmark = [pytest.mark.core_model]
-@pytest.mark.cpu
-def test_worker_extension_inheritance():
- assert issubclass(vLLMOmniColocateWorkerExtensionForTest, CustomPipelineWorkerExtension)
-
-
-@pytest.mark.cpu
-def test_worker_extension_test_function():
- assert (
- vLLMOmniColocateWorkerExtensionForTest.test_extension_name() == "vllm-omni-colocate-worker-extension-for-test"
- )
-
-
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.asyncio
diff --git a/tests/fullduplex/engine/test_duplex_control_client.py b/tests/e2e/features/fullduplex/engine/test_duplex_control_client.py
similarity index 100%
rename from tests/fullduplex/engine/test_duplex_control_client.py
rename to tests/e2e/features/fullduplex/engine/test_duplex_control_client.py
diff --git a/tests/fullduplex/engine/test_duplex_control_plane.py b/tests/e2e/features/fullduplex/engine/test_duplex_control_plane.py
similarity index 100%
rename from tests/fullduplex/engine/test_duplex_control_plane.py
rename to tests/e2e/features/fullduplex/engine/test_duplex_control_plane.py
diff --git a/tests/fullduplex/engine/test_duplex_deploy_config.py b/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py
similarity index 100%
rename from tests/fullduplex/engine/test_duplex_deploy_config.py
rename to tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py
diff --git a/tests/fullduplex/engine/test_duplex_intermediate.py b/tests/e2e/features/fullduplex/engine/test_duplex_intermediate.py
similarity index 100%
rename from tests/fullduplex/engine/test_duplex_intermediate.py
rename to tests/e2e/features/fullduplex/engine/test_duplex_intermediate.py
diff --git a/tests/fullduplex/engine/test_duplex_lease.py b/tests/e2e/features/fullduplex/engine/test_duplex_lease.py
similarity index 100%
rename from tests/fullduplex/engine/test_duplex_lease.py
rename to tests/e2e/features/fullduplex/engine/test_duplex_lease.py
diff --git a/tests/fullduplex/engine/test_duplex_runtime.py b/tests/e2e/features/fullduplex/engine/test_duplex_runtime.py
similarity index 100%
rename from tests/fullduplex/engine/test_duplex_runtime.py
rename to tests/e2e/features/fullduplex/engine/test_duplex_runtime.py
diff --git a/tests/fullduplex/minicpmo45/__init__.py b/tests/e2e/features/fullduplex/minicpmo45/__init__.py
similarity index 100%
rename from tests/fullduplex/minicpmo45/__init__.py
rename to tests/e2e/features/fullduplex/minicpmo45/__init__.py
diff --git a/tests/fullduplex/minicpmo45/test_commit_policy.py b/tests/e2e/features/fullduplex/minicpmo45/test_commit_policy.py
similarity index 100%
rename from tests/fullduplex/minicpmo45/test_commit_policy.py
rename to tests/e2e/features/fullduplex/minicpmo45/test_commit_policy.py
diff --git a/tests/fullduplex/minicpmo45/test_input.py b/tests/e2e/features/fullduplex/minicpmo45/test_input.py
similarity index 100%
rename from tests/fullduplex/minicpmo45/test_input.py
rename to tests/e2e/features/fullduplex/minicpmo45/test_input.py
diff --git a/tests/fullduplex/openai/test_duplex_audio.py b/tests/e2e/features/fullduplex/openai/test_duplex_audio.py
similarity index 100%
rename from tests/fullduplex/openai/test_duplex_audio.py
rename to tests/e2e/features/fullduplex/openai/test_duplex_audio.py
diff --git a/tests/fullduplex/openai/test_runtime_adapter_boundary.py b/tests/e2e/features/fullduplex/openai/test_runtime_adapter_boundary.py
similarity index 100%
rename from tests/fullduplex/openai/test_runtime_adapter_boundary.py
rename to tests/e2e/features/fullduplex/openai/test_runtime_adapter_boundary.py
diff --git a/tests/fullduplex/openai/test_websocket_actor.py b/tests/e2e/features/fullduplex/openai/test_websocket_actor.py
similarity index 100%
rename from tests/fullduplex/openai/test_websocket_actor.py
rename to tests/e2e/features/fullduplex/openai/test_websocket_actor.py
diff --git a/tests/fullduplex/test_client.py b/tests/e2e/features/fullduplex/test_client.py
similarity index 100%
rename from tests/fullduplex/test_client.py
rename to tests/e2e/features/fullduplex/test_client.py
diff --git a/tests/fullduplex/test_joyvl_brain.py b/tests/e2e/features/fullduplex/test_joyvl_brain.py
similarity index 100%
rename from tests/fullduplex/test_joyvl_brain.py
rename to tests/e2e/features/fullduplex/test_joyvl_brain.py
diff --git a/tests/fullduplex/test_joyvl_policy.py b/tests/e2e/features/fullduplex/test_joyvl_policy.py
similarity index 100%
rename from tests/fullduplex/test_joyvl_policy.py
rename to tests/e2e/features/fullduplex/test_joyvl_policy.py
diff --git a/tests/fullduplex/test_runtime.py b/tests/e2e/features/fullduplex/test_runtime.py
similarity index 100%
rename from tests/fullduplex/test_runtime.py
rename to tests/e2e/features/fullduplex/test_runtime.py
diff --git a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py b/tests/e2e/features/helpers/custom_pipeline.py
similarity index 59%
rename from tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py
rename to tests/e2e/features/helpers/custom_pipeline.py
index 193d1b1e2..63cb21056 100644
--- a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py
+++ b/tests/e2e/features/helpers/custom_pipeline.py
@@ -1,27 +1,299 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Detailed custom Qwen-Image pipeline for E2E tests.
+"""Test doubles for custom-pipeline E2E and RLHF integration.
-This pipeline follows the structure of the user's reference implementation:
-- supports pre-tokenized prompt IDs via OmniCustomPrompt-style dict input
-- uses an SDE scheduler that can return step logprobs
-- returns structured trajectory_* fields (latents, timesteps, log_probs)
- consistent with the BAGEL trajectory recording design
+Three layers:
+
+* ``FlowMatchSDEDiscreteSchedulerForTest`` — SDE scheduler with per-step logprobs.
+* ``vLLMOmniColocateWorkerExtensionForTest`` — verl-omni colocate worker extension mirror.
+* ``QwenImagePipelineWithLogProbForTest`` — Qwen-Image pipeline returning trajectory_* fields.
"""
from __future__ import annotations
+import ctypes
+import logging
+import math
import os
+import platform
+import signal
+from dataclasses import dataclass
from typing import Any, Literal
import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+from msgspec import field
-from tests.e2e.offline_inference.custom_pipeline.flow_match_sde_scheduler import FlowMatchSDEDiscreteSchedulerForTest
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.models.qwen_image import QwenImagePipeline
+from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.lora.request import LoRARequest as OmniLoRARequest
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Scheduler layer
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class FlowMatchSDEDiscreteSchedulerOutputForTest(BaseOutput):
+ """
+ Output class for the scheduler's `step` function output.
+
+ Args:
+ prev_sample (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
+ Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+ denoising loop.
+ log_prob (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
+ The log probability of the previous sample.
+ prev_sample_mean (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
+ The mean of the computed sample of previous timestep.
+ std_dev_t (`torch.FloatTensor` of shape `(batch_size, 1, 1)`):
+ The standard deviation used to compute `prev_sample`.
+ """
+
+ prev_sample: torch.FloatTensor
+ log_prob: torch.FloatTensor | None
+ prev_sample_mean: torch.FloatTensor
+ std_dev_t: torch.FloatTensor
+
+
+class FlowMatchSDEDiscreteSchedulerForTest(FlowMatchEulerDiscreteScheduler):
+ """SDE version of the FlowMatchEulerDiscreteScheduler.
+ The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470)
+ and diffusers v0.37 branch.
+ """
+
+ def step(
+ self,
+ model_output: torch.FloatTensor,
+ timestep: float | torch.FloatTensor,
+ sample: torch.FloatTensor,
+ s_churn: float = 0.0,
+ s_tmin: float = 0.0,
+ s_tmax: float = float("inf"),
+ s_noise: float = 1.0,
+ generator: torch.Generator | None = None,
+ per_token_timesteps: torch.Tensor | None = None,
+ return_dict: bool = True,
+ noise_level: float = 0.7,
+ prev_sample: torch.FloatTensor | None = None,
+ sde_type: Literal["sde", "cps"] = "sde",
+ return_logprobs: bool = True,
+ ) -> FlowMatchSDEDiscreteSchedulerOutputForTest | tuple:
+ """
+ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+ process from the learned model outputs (most often the predicted noise).
+
+ Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py
+
+ Args:
+ model_output (`torch.FloatTensor`):
+ The direct output from learned diffusion model.
+ timestep (`float`):
+ The current discrete timestep in the diffusion chain.
+ sample (`torch.FloatTensor`):
+ A current instance of a sample created by the diffusion process.
+ s_churn (`float`):
+ s_tmin (`float`):
+ s_tmax (`float`):
+ s_noise (`float`, defaults to 1.0):
+ Scaling factor for noise added to the sample.
+ generator (`torch.Generator`, *optional*):
+ A random number generator.
+ per_token_timesteps (`torch.Tensor`, *optional*):
+ The timesteps for each token in the sample.
+ return_dict (`bool`):
+ Whether or not to return a
+ [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput`] or tuple.
+ noise_level (`float`, *optional*, defaults to 0.7):
+ The noise level used in the SDE.
+ prev_sample (`torch.FloatTensor`, *optional*):
+ The sample from the previous timestep. If not provided, it will be sampled inside the function.
+ sde_type (`str`, *optional*, defaults to "sde"):
+ The type of SDE to use. Choose between "sde" and "cps".
+ return_logprobs (`bool`, *optional*, defaults to True):
+ Whether to return log probabilities of the previous sample.
+ """
+
+ if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+ raise ValueError(
+ (
+ "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+ " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+ " one of the `scheduler.timesteps` as a timestep."
+ ),
+ )
+
+ if self.step_index is None:
+ self._init_step_index(timestep)
+
+ # Upcast to avoid precision issues when computing prev_sample
+ sample = sample.to(torch.float32)
+ if prev_sample is not None:
+ prev_sample = prev_sample.to(torch.float32)
+
+ prev_sample, log_prob, prev_sample_mean, std_dev_t = self.sample_previous_step(
+ sample=sample,
+ model_output=model_output,
+ generator=generator,
+ per_token_timesteps=per_token_timesteps,
+ noise_level=noise_level,
+ prev_sample=prev_sample,
+ sde_type=sde_type,
+ return_logprobs=return_logprobs,
+ )
+
+ # upon completion increase step index by one
+ self._step_index += 1
+ if per_token_timesteps is None:
+ # Cast sample back to model compatible dtype
+ prev_sample = prev_sample.to(model_output.dtype)
+
+ if not return_dict:
+ return (prev_sample, log_prob, prev_sample_mean, std_dev_t)
+
+ return FlowMatchSDEDiscreteSchedulerOutputForTest(
+ prev_sample=prev_sample, log_prob=log_prob, prev_sample_mean=prev_sample_mean, std_dev_t=std_dev_t
+ )
+
+ def sample_previous_step(
+ self,
+ sample: torch.Tensor,
+ model_output: torch.Tensor,
+ timestep: torch.FloatTensor | None = None,
+ generator: torch.Generator | None = None,
+ per_token_timesteps: torch.Tensor | None = None,
+ noise_level: float = 0.7,
+ prev_sample: torch.Tensor | None = None,
+ sde_type: Literal["cps", "sde"] = "sde",
+ return_logprobs: bool = True,
+ ):
+ assert sde_type in ["sde", "cps"]
+ assert sample.dtype == torch.float32
+ if prev_sample is not None:
+ assert prev_sample.dtype == torch.float32
+
+ if per_token_timesteps is not None:
+ raise NotImplementedError("per_token_timesteps is not supported yet for FlowMatchSDEDiscreteScheduler.")
+ else:
+ if timestep is None:
+ sigma_idx = self.step_index
+ sigma = self.sigmas[sigma_idx]
+ sigma_prev = self.sigmas[sigma_idx + 1]
+ else:
+ sigma_idx = torch.tensor([self.index_for_timestep(t) for t in timestep])
+ sigma = self.sigmas[sigma_idx].view(-1, *([1] * (len(sample.shape) - 1)))
+ sigma_prev = self.sigmas[sigma_idx + 1].view(-1, *([1] * (len(sample.shape) - 1)))
+
+ sigma_max = self.sigmas[1]
+ dt = sigma_prev - sigma
+
+ if sde_type == "sde":
+ std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level
+
+ prev_sample_mean = (
+ sample * (1 + std_dev_t**2 / (2 * sigma) * dt)
+ + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
+ )
+
+ if prev_sample is None:
+ variance_noise = randn_tensor(
+ model_output.shape,
+ generator=generator,
+ device=model_output.device,
+ dtype=model_output.dtype,
+ )
+ prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise
+
+ if return_logprobs:
+ log_prob = (
+ -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2))
+ - torch.log(std_dev_t * torch.sqrt(-1 * dt))
+ - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
+ )
+ else:
+ log_prob = None
+
+ elif sde_type == "cps":
+ std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2)
+ pred_original_sample = sample - sigma * model_output
+ noise_estimate = sample + model_output * (1 - sigma)
+ prev_sample_mean = pred_original_sample * (1 - sigma_prev) + noise_estimate * torch.sqrt(
+ sigma_prev**2 - std_dev_t**2
+ )
+
+ if prev_sample is None:
+ variance_noise = randn_tensor(
+ model_output.shape,
+ generator=generator,
+ device=model_output.device,
+ dtype=model_output.dtype,
+ )
+ prev_sample = prev_sample_mean + std_dev_t * variance_noise
+
+ if return_logprobs:
+ log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2)
+ else:
+ log_prob = None
+
+ # mean along all but batch dimension
+ log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) if log_prob is not None else None
+ return prev_sample, log_prob, prev_sample_mean, std_dev_t
+
+
+# ---------------------------------------------------------------------------
+# Worker extension layer
+# ---------------------------------------------------------------------------
+
+
+class OmniTensorLoRARequestForTest(OmniLoRARequest):
+ peft_config: dict = field(default=None)
+ lora_tensors: dict = field(default=None)
+
+
+def set_death_signal() -> None:
+ """Verbatim port of ``verl.workers.rollout.vllm_rollout.utils.set_death_signal``."""
+ if platform.system() != "Linux":
+ return
+ try:
+ libc = ctypes.CDLL("libc.so.6")
+ libc.prctl(1, signal.SIGKILL)
+ if os.getppid() == 1:
+ os.kill(os.getpid(), signal.SIGKILL)
+ except Exception: # noqa: BLE001
+ # libc.so.6 may not be present (musl, alpine) — best-effort only.
+ pass
+
+
+class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension):
+ """Mirror of ``vLLMOmniColocateWorkerExtension`` (verl-omni).
+
+ The production ``__new__`` runs ``set_death_signal`` + ``VLLMOmniHijack.hijack``
+ on every vLLM worker process. Replicating both ensures this test
+ reproduces the same monkey-patched environment.
+ """
+
+ def __new__(cls, **kwargs):
+ set_death_signal()
+ return super().__new__(cls)
+
+ @staticmethod
+ def test_extension_name() -> str:
+ """Return a stable identifier for assertions in unit tests."""
+ return "vllm-omni-colocate-worker-extension-for-test"
+
+
+# ---------------------------------------------------------------------------
+# Pipeline layer
+# ---------------------------------------------------------------------------
def _maybe_to_cpu(v):
@@ -30,8 +302,6 @@ def _maybe_to_cpu(v):
return v
-# Custom pipeline class for QwenImage that returns log probabilities during the diffusion process.
-# This is for test
class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
supports_request_batch = False
diff --git a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/features/helpers/verl_omni_server.py
similarity index 52%
rename from tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
rename to tests/e2e/features/helpers/verl_omni_server.py
index 420c9bb68..d53e0936f 100644
--- a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+++ b/tests/e2e/features/helpers/verl_omni_server.py
@@ -1,30 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""E2E test that follows the EXACT same flow as
-``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py``
-but inlines all the ``verl`` / ``verl_omni`` symbols the flow depends on
-(stripped down to only what the diffusion-mode test actually exercises).
+"""Ray-serializable verl-omni HTTP server test double for RLHF E2E tests.
-Flow (identical to the verl-omni test):
-
-1. Build two ``OmegaConf`` configs (``DiffusionRolloutConfig`` /
- ``DiffusionModelConfig`` shape).
-2. ``ServerCls = ray.remote(vLLMOmniHttpServerLocal)`` then
- ``ServerCls.options(...).remote(config=..., model_config=...,
- rollout_mode=RolloutMode.STANDALONE, workers=[], replica_rank=0,
- node_rank=0, gpus_per_node=1, nnodes=1, cuda_visible_devices="0")``.
-3. ``ray.get(server.launch_server.remote())``.
-4. ``server.generate.remote(prompt_ids=..., sampling_params=..., request_id=...)``
- and assert the returned ``DiffusionOutput``.
-
-The inlined classes / helpers mirror their ``verl`` / ``verl_omni``
-counterparts; bits not reached by this test (AR mode, LoRA,
-quantization, run_uvicorn, headless, profiler, port wiring) are removed.
-
-Usage:
- pytest tests/e2e/offline_inference/rlhf_test/\
-test_ray_async_omni_qwen_image_generate.py -v -s
+The actor class must live in an importable module (not a pytest test file) so
+Ray workers can unpickle ``vLLMOmniHttpServerLocal`` via
+``tests.e2e.features.helpers.verl_omni_server``.
"""
from __future__ import annotations
@@ -35,23 +16,15 @@ import logging
import os
from dataclasses import asdict
from enum import Enum
-from functools import lru_cache
from typing import Any
-from uuid import uuid4
-import pytest
-import ray
-import torch
import torchvision.transforms as T
-from huggingface_hub import snapshot_download
from omegaconf import DictConfig, OmegaConf
from pydantic import BaseModel, ConfigDict
-from transformers import AutoTokenizer
from vllm.entrypoints.openai.api_server import build_app
from vllm.utils.argparse_utils import FlexibleArgumentParser
import vllm_omni.entrypoints.cli.serve
-from tests.helpers.mark import hardware_test
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.entrypoints.openai.api_server import omni_init_app_state
@@ -60,33 +33,8 @@ from vllm_omni.outputs import OmniRequestOutput
logger = logging.getLogger(__name__)
-CUSTOM_PIPELINE_CLASS = (
- "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
-)
-WORKER_EXTENSION_CLASS = (
- "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
-)
-
-# Use your specified HF repo name/path directly here
-MODEL = "tiny-random/Qwen-Image"
-
-TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct"
-
-
-def _resolve_model_path(repo_id: str) -> str:
- """Resolve an HF repo ID to a local snapshot path, downloading if needed."""
- # Allow overriding with a pre-existing local path (skips download).
- if os.path.isdir(repo_id):
- return repo_id
- return snapshot_download(repo_id=repo_id)
-
-
-_MIN_PROMPT_TOKENS = 35
-
-
-# ---------------------------------------------------------------------
-# Inlined: verl.workers.rollout.replica.RolloutMode
-# ---------------------------------------------------------------------
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
class RolloutMode(Enum):
@@ -95,11 +43,6 @@ class RolloutMode(Enum):
STANDALONE = "standalone"
-# ---------------------------------------------------------------------
-# Inlined: verl_omni.workers.rollout.replica.DiffusionOutput
-# ---------------------------------------------------------------------
-
-
class DiffusionOutput(BaseModel):
"""Pydantic mirror of ``verl_omni.workers.rollout.replica.DiffusionOutput``."""
@@ -112,11 +55,6 @@ class DiffusionOutput(BaseModel):
extra_fields: dict[str, Any] = {}
-# ---------------------------------------------------------------------
-# Inlined: verl.utils.tokenizer.normalize_token_ids
-# ---------------------------------------------------------------------
-
-
def normalize_token_ids(tokenized_output) -> list[int]:
token_ids = tokenized_output
if isinstance(tokenized_output, dict):
@@ -142,11 +80,6 @@ def normalize_token_ids(tokenized_output) -> list[int]:
return out
-# ---------------------------------------------------------------------
-# Inlined: verl.workers.rollout.vllm_rollout.utils.build_cli_args_from_config
-# ---------------------------------------------------------------------
-
-
def build_cli_args_from_config(config: dict[str, Any]) -> list[str]:
cli_args: list[str] = []
for k, v in config.items():
@@ -166,34 +99,6 @@ def build_cli_args_from_config(config: dict[str, Any]) -> list[str]:
return cli_args
-# ---------------------------------------------------------------------
-# Tokenizer helper (test-side)
-# ---------------------------------------------------------------------
-
-
-@lru_cache(maxsize=1)
-def _get_tokenizer():
- return AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True)
-
-
-def _tokenize_prompt(text: str) -> list[int]:
- tokenizer = _get_tokenizer()
- messages = [{"role": "user", "content": text}]
- token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False))
- assert len(token_ids) > _MIN_PROMPT_TOKENS, (
- f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). "
- "The pipeline drops the first 34 chat-template prefix tokens; "
- "use a longer prompt so content tokens remain after the drop."
- )
- return token_ids
-
-
-# ---------------------------------------------------------------------
-# Inlined: vLLMOmniHttpServer (verl-omni) + relevant vLLMHttpServer (verl)
-# bits, stripped to the diffusion-mode path the test exercises.
-# ---------------------------------------------------------------------
-
-
def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any:
"""Uniform getter that works for both DictConfig and dataclass-like objects."""
if isinstance(cfg, DictConfig):
@@ -202,13 +107,8 @@ def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any:
class vLLMOmniHttpServerLocal:
- """Self-contained simulation of ``vLLMOmniHttpServer`` for diffusion mode.
+ """Self-contained simulation of ``vLLMOmniHttpServer`` for diffusion mode."""
- Constructor signature matches ``vLLMHttpServer.__init__`` exactly so the
- Ray-actor call site in the test is identical to the verl-omni one.
- """
-
- # ------------------------------------------------------------------ init
def __init__(
self,
config,
@@ -224,8 +124,6 @@ class vLLMOmniHttpServerLocal:
os.environ.setdefault("CUDA_VISIBLE_DEVICES", cuda_visible_devices)
os.environ["VERL_REPLICA_RANK"] = str(replica_rank)
- # Resolve configs. verl uses omega_conf_to_dataclass; here we keep
- # them as DictConfig — every read site below uses ``_cfg_get``.
self.config = OmegaConf.create(config) if not isinstance(config, DictConfig) else config
self.model_config = OmegaConf.create(model_config) if not isinstance(model_config, DictConfig) else model_config
@@ -237,8 +135,6 @@ class vLLMOmniHttpServerLocal:
self.nnodes = nnodes
self.global_steps = 0
self.engine: AsyncOmni | None = None
-
- # Diffusion-mode parity with ``vLLMOmniHttpServer._post_init``.
self._to_tensor = T.PILToTensor()
logger.info(
@@ -249,11 +145,7 @@ class vLLMOmniHttpServerLocal:
cuda_visible_devices,
)
- # ---------------------------------------------------------- launch_server
async def launch_server(self) -> None:
- """Diffusion-only mirror of ``vLLMHttpServer.launch_server``."""
-
- # vLLMOmniHttpServer._get_engine_kwargs_key() returns "vllm_omni".
engine_kwargs_root = _cfg_get(self.config, "engine_kwargs", {}) or {}
if isinstance(engine_kwargs_root, (dict, DictConfig)):
engine_kwargs = engine_kwargs_root.get("vllm_omni", {}) or {}
@@ -261,10 +153,8 @@ class vLLMOmniHttpServerLocal:
engine_kwargs = {}
engine_kwargs = {k: v for k, v in dict(engine_kwargs).items() if v is not None}
- # vLLMOmniHttpServer._preprocess_engine_kwargs (diffusion path).
engine_kwargs.pop("output_mode", None)
- # ``compilation_config`` handling — verbatim from parent launch_server.
compilation_config = engine_kwargs.pop("compilation_config", None) or {}
if isinstance(compilation_config, str):
compilation_config = json.loads(compilation_config)
@@ -273,8 +163,6 @@ class vLLMOmniHttpServerLocal:
compilation_config.setdefault("cudagraph_mode", "FULL_AND_PIECEWISE")
compilation_config = json.dumps(compilation_config)
- # Args dict (subset of parent's; the diffusion override drops AR-only
- # entries via ``_get_override_generation_config`` returning ``{}``).
args = {
"dtype": _cfg_get(self.config, "dtype", "bfloat16"),
"load_format": _cfg_get(self.config, "load_format", "auto"),
@@ -303,8 +191,6 @@ class vLLMOmniHttpServerLocal:
model_path = str(_cfg_get(self.model_config, "local_path", None) or _cfg_get(self.model_config, "path"))
server_args_list = ["serve", model_path] + build_cli_args_from_config(args)
- # Same parser pipeline as parent launch_server (with vllm-omni's
- # cmd_init injected, matching ``_get_cli_modules`` + ``_get_cli_description``).
parser = FlexibleArgumentParser(description="vLLM-Omni CLI")
subparsers = parser.add_subparsers(required=False, dest="subparser")
cmds: dict[str, Any] = {}
@@ -316,19 +202,12 @@ class vLLMOmniHttpServerLocal:
if parsed.subparser in cmds:
cmds[parsed.subparser].validate(parsed)
- # node_rank == 0 → run_server (the test never sets node_rank != 0).
await self.run_server(parsed)
- # ------------------------------------------------------------- run_server
async def run_server(self, args: argparse.Namespace) -> None:
- """Diffusion branch of ``vLLMOmniHttpServer.run_server``."""
-
engine_args = OmniEngineArgs.from_cli_args(args)
engine_args = asdict(engine_args)
- # Diffusion path: pin the in-repo test pipeline. The production code
- # calls ``VllmOmniPipelineBase.get_pipeline_path(architecture)``; we
- # bypass the registry and bind the test pipeline directly.
engine_args["enable_dummy_pipeline"] = True
engine_args["custom_pipeline_args"] = {"pipeline_class": CUSTOM_PIPELINE_CLASS}
@@ -337,10 +216,7 @@ class vLLMOmniHttpServerLocal:
await omni_init_app_state(engine_client, app.state, args)
self.engine = engine_client
- # ``run_uvicorn`` intentionally skipped — the test drives the engine
- # via Ray RPC, not HTTP.
- # ---------------------------------------------------------------- generate
async def generate(
self,
prompt_ids: list[int],
@@ -351,8 +227,6 @@ class vLLMOmniHttpServerLocal:
negative_prompt_ids: list[int] | None = None,
priority: int = 0, # noqa: ARG002 (signature parity)
) -> DiffusionOutput:
- # Verbatim copy of ``vLLMOmniHttpServer._generate_diffusion`` minus
- # the LoRA branch (the test does not enable ``lora_as_adapter``).
prompt_ids = normalize_token_ids(prompt_ids)
multi_modal_data: dict[str, Any] = {}
@@ -444,149 +318,3 @@ class vLLMOmniHttpServerLocal:
num_preempted=num_preempted,
extra_fields=extra_fields,
)
-
-
-# ---------------------------------------------------------------------
-# Fixtures & tests
-# ---------------------------------------------------------------------
-
-
-@pytest.fixture
-def init_server():
- """Create and launch a vLLMOmniHttpServerLocal Ray actor with Qwen-Image."""
- model_path = _resolve_model_path(MODEL)
-
- ray.init(
- runtime_env={
- "env_vars": {
- "TOKENIZERS_PARALLELISM": "true",
- "NCCL_DEBUG": "WARN",
- "VLLM_LOGGING_LEVEL": "INFO",
- }
- },
- ignore_reinit_error=True,
- )
-
- rollout_cfg = OmegaConf.create(
- {
- "name": "vllm_omni",
- "mode": "async",
- "tensor_model_parallel_size": 1,
- "data_parallel_size": 1,
- "pipeline_model_parallel_size": 1,
- "gpu_memory_utilization": 0.8,
- "max_num_batched_tokens": 8192,
- # The custom log-prob pipeline only supports serial requests.
- "max_num_seqs": 1,
- "max_model_len": 1058,
- "dtype": "bfloat16",
- "load_format": "auto",
- "enforce_eager": True,
- "enable_chunked_prefill": False,
- "enable_prefix_caching": False,
- "enable_sleep_mode": False,
- "free_cache_engine": True,
- "disable_log_stats": True,
- "logprobs_mode": "processed_logprobs",
- "scheduling_policy": "fcfs",
- "seed": 0,
- "n": 4,
- "pipeline": {
- "height": 512,
- "width": 512,
- "num_inference_steps": 10,
- },
- }
- )
-
- model_cfg = OmegaConf.create(
- {
- "path": str(model_path),
- "local_path": str(model_path),
- "tokenizer_path": TOKENIZER_MODEL,
- "trust_remote_code": True,
- "load_tokenizer": True,
- }
- )
-
- ServerCls = ray.remote(vLLMOmniHttpServerLocal)
- server = ServerCls.options(
- runtime_env={
- "env_vars": {
- "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
- "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
- "NCCL_CUMEM_ENABLE": "0",
- }
- },
- max_concurrency=10,
- num_gpus=1,
- ).remote(
- config=rollout_cfg,
- model_config=model_cfg,
- rollout_mode=RolloutMode.STANDALONE,
- workers=[],
- replica_rank=0,
- node_rank=0,
- gpus_per_node=1,
- nnodes=1,
- cuda_visible_devices="0",
- )
-
- ray.get(server.launch_server.remote())
-
- yield server
-
- ray.shutdown()
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_generate(init_server):
- """Concurrent generate() calls covering basic output, logprobs, and multi-request correctness."""
- server = init_server
-
- prompts = [
- "a beautiful sunset over the ocean with vibrant orange and purple clouds "
- "reflecting on the calm water surface near a rocky coastline",
- "a fluffy orange cat sitting on a wooden windowsill looking outside at "
- "a garden full of colorful flowers on a bright sunny afternoon",
- "a majestic mountain landscape covered with fresh white snow under a "
- "clear blue sky with pine trees in the foreground and a frozen lake",
- "a futuristic city at night with neon lights glowing on tall glass "
- "skyscrapers and flying vehicles soaring between the buildings",
- ]
-
- refs = []
- for i, prompt in enumerate(prompts):
- rid = f"test_{i}_{uuid4().hex[:8]}"
- ref = server.generate.remote(
- prompt_ids=_tokenize_prompt(prompt),
- sampling_params={
- "num_inference_steps": 10,
- "true_cfg_scale": 4.0,
- "height": 512,
- "width": 512,
- "logprobs": i == 0, # first request includes logprobs
- },
- request_id=rid,
- )
- refs.append(ref)
-
- results = ray.get(refs, timeout=600)
-
- for i, output in enumerate(results):
- assert isinstance(output, DiffusionOutput), f"Request {i}: expected DiffusionOutput"
- img = output.diffusion_output
- assert isinstance(img, torch.Tensor) and img.ndim == 3, (
- f"Request {i}: expected CHW torch.Tensor, got {type(img).__name__} ndim={getattr(img, 'ndim', None)}"
- )
- c, h, w = img.shape
- assert c == 3, f"Request {i}: expected 3 channels (CHW), got {c}"
- assert h > 0 and w > 0, f"Request {i}: image dimensions must be positive"
- assert output.stop_reason in ("completed", "aborted", None), (
- f"Request {i}: unexpected stop_reason {output.stop_reason!r}"
- )
- assert 0.0 <= float(img[0, 0, 0]) <= 1.0, f"Request {i}: pixel values must be in [0, 1]"
-
- print(f"All {len(prompts)} concurrent requests returned valid DiffusionOutput")
diff --git a/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py
new file mode 100644
index 000000000..dab5feef1
--- /dev/null
+++ b/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py
@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E test that follows the EXACT same flow as
+``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py``
+but inlines all the ``verl`` / ``verl_omni`` symbols the flow depends on
+(stripped down to only what the diffusion-mode test actually exercises).
+
+Flow (identical to the verl-omni test):
+
+1. Build two ``OmegaConf`` configs (``DiffusionRolloutConfig`` /
+ ``DiffusionModelConfig`` shape).
+2. ``ServerCls = ray.remote(vLLMOmniHttpServerLocal)`` then
+ ``ServerCls.options(...).remote(config=..., model_config=...,
+ rollout_mode=RolloutMode.STANDALONE, workers=[], replica_rank=0,
+ node_rank=0, gpus_per_node=1, nnodes=1, cuda_visible_devices="0")``.
+3. ``ray.get(server.launch_server.remote())``.
+4. ``server.generate.remote(prompt_ids=..., sampling_params=..., request_id=...)``
+ and assert the returned ``DiffusionOutput``.
+
+The Ray actor lives in ``tests.e2e.features.helpers.verl_omni_server`` so
+workers can import it after directory migration (pytest test modules are not
+importable as ``test_verl_omni_e2e`` in Ray child processes).
+
+Usage:
+ pytest tests/e2e/features/rlhf_test/test_verl_omni_e2e.py -v -s
+"""
+
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+from uuid import uuid4
+
+import pytest
+import ray
+import torch
+from huggingface_hub import snapshot_download
+from omegaconf import OmegaConf
+from transformers import AutoTokenizer
+
+from tests.e2e.features.helpers.verl_omni_server import (
+ DiffusionOutput,
+ RolloutMode,
+ normalize_token_ids,
+ vLLMOmniHttpServerLocal,
+)
+from tests.helpers.mark import hardware_test
+
+MODEL = "tiny-random/Qwen-Image"
+TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct"
+_MIN_PROMPT_TOKENS = 35
+
+
+def _resolve_model_path(repo_id: str) -> str:
+ if os.path.isdir(repo_id):
+ return repo_id
+ return snapshot_download(repo_id=repo_id)
+
+
+@lru_cache(maxsize=1)
+def _get_tokenizer():
+ return AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True)
+
+
+def _tokenize_prompt(text: str) -> list[int]:
+ tokenizer = _get_tokenizer()
+ messages = [{"role": "user", "content": text}]
+ token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False))
+ assert len(token_ids) > _MIN_PROMPT_TOKENS, (
+ f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). "
+ "The pipeline drops the first 34 chat-template prefix tokens; "
+ "use a longer prompt so content tokens remain after the drop."
+ )
+ return token_ids
+
+
+@pytest.fixture
+def init_server():
+ """Create and launch a vLLMOmniHttpServerLocal Ray actor with Qwen-Image."""
+ model_path = _resolve_model_path(MODEL)
+
+ ray.init(
+ runtime_env={
+ "env_vars": {
+ "TOKENIZERS_PARALLELISM": "true",
+ "NCCL_DEBUG": "WARN",
+ "VLLM_LOGGING_LEVEL": "INFO",
+ }
+ },
+ ignore_reinit_error=True,
+ )
+
+ rollout_cfg = OmegaConf.create(
+ {
+ "name": "vllm_omni",
+ "mode": "async",
+ "tensor_model_parallel_size": 1,
+ "data_parallel_size": 1,
+ "pipeline_model_parallel_size": 1,
+ "gpu_memory_utilization": 0.8,
+ "max_num_batched_tokens": 8192,
+ # QwenImagePipeline (and the logprob custom subclass) only supports
+ # serial request execution; max_num_seqs>1 enables request-level batching.
+ "max_num_seqs": 1,
+ "max_model_len": 1058,
+ "dtype": "bfloat16",
+ "load_format": "auto",
+ "enforce_eager": True,
+ "enable_chunked_prefill": False,
+ "enable_prefix_caching": False,
+ "enable_sleep_mode": False,
+ "free_cache_engine": True,
+ "disable_log_stats": True,
+ "logprobs_mode": "processed_logprobs",
+ "scheduling_policy": "fcfs",
+ "seed": 0,
+ "n": 4,
+ "pipeline": {
+ "height": 512,
+ "width": 512,
+ "num_inference_steps": 10,
+ },
+ }
+ )
+
+ model_cfg = OmegaConf.create(
+ {
+ "path": str(model_path),
+ "local_path": str(model_path),
+ "tokenizer_path": TOKENIZER_MODEL,
+ "trust_remote_code": True,
+ "load_tokenizer": True,
+ }
+ )
+
+ ServerCls = ray.remote(vLLMOmniHttpServerLocal)
+ server = ServerCls.options(
+ runtime_env={
+ "env_vars": {
+ "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
+ "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+ "NCCL_CUMEM_ENABLE": "0",
+ }
+ },
+ max_concurrency=10,
+ num_gpus=1,
+ ).remote(
+ config=rollout_cfg,
+ model_config=model_cfg,
+ rollout_mode=RolloutMode.STANDALONE,
+ workers=[],
+ replica_rank=0,
+ node_rank=0,
+ gpus_per_node=1,
+ nnodes=1,
+ cuda_visible_devices="0",
+ )
+
+ ray.get(server.launch_server.remote())
+
+ yield server
+
+ ray.shutdown()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_generate(init_server):
+ """Concurrent generate() calls covering basic output, logprobs, and multi-request correctness."""
+ server = init_server
+
+ prompts = [
+ "a beautiful sunset over the ocean with vibrant orange and purple clouds "
+ "reflecting on the calm water surface near a rocky coastline",
+ "a fluffy orange cat sitting on a wooden windowsill looking outside at "
+ "a garden full of colorful flowers on a bright sunny afternoon",
+ "a majestic mountain landscape covered with fresh white snow under a "
+ "clear blue sky with pine trees in the foreground and a frozen lake",
+ "a futuristic city at night with neon lights glowing on tall glass "
+ "skyscrapers and flying vehicles soaring between the buildings",
+ ]
+
+ refs = []
+ for i, prompt in enumerate(prompts):
+ rid = f"test_{i}_{uuid4().hex[:8]}"
+ ref = server.generate.remote(
+ prompt_ids=_tokenize_prompt(prompt),
+ sampling_params={
+ "num_inference_steps": 10,
+ "true_cfg_scale": 4.0,
+ "height": 512,
+ "width": 512,
+ "logprobs": i == 0,
+ },
+ request_id=rid,
+ )
+ refs.append(ref)
+
+ results = ray.get(refs, timeout=600)
+
+ for i, output in enumerate(results):
+ assert isinstance(output, DiffusionOutput), f"Request {i}: expected DiffusionOutput"
+ img = output.diffusion_output
+ assert isinstance(img, torch.Tensor) and img.ndim == 3, (
+ f"Request {i}: expected CHW torch.Tensor, got {type(img).__name__} ndim={getattr(img, 'ndim', None)}"
+ )
+ c, h, w = img.shape
+ assert c == 3, f"Request {i}: expected 3 channels (CHW), got {c}"
+ assert h > 0 and w > 0, f"Request {i}: image dimensions must be positive"
+ assert output.stop_reason in ("completed", "aborted", None), (
+ f"Request {i}: unexpected stop_reason {output.stop_reason!r}"
+ )
+ assert 0.0 <= float(img[0, 0, 0]) <= 1.0, f"Request {i}: pixel values must be in [0, 1]"
+
+ print(f"All {len(prompts)} concurrent requests returned valid DiffusionOutput")
diff --git a/tests/e2e/offline_inference/custom_pipeline/__init__.py b/tests/e2e/offline_inference/custom_pipeline/__init__.py
deleted file mode 100644
index 99663a46a..000000000
--- a/tests/e2e/offline_inference/custom_pipeline/__init__.py
+++ /dev/null
@@ -1,14 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import (
- QwenImagePipelineWithLogProbForTest,
-)
-from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
- vLLMOmniColocateWorkerExtensionForTest,
-)
-
-__all__ = [
- "QwenImagePipelineWithLogProbForTest",
- "vLLMOmniColocateWorkerExtensionForTest",
-]
diff --git a/tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py b/tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py
deleted file mode 100644
index f85eb6602..000000000
--- a/tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py
+++ /dev/null
@@ -1,220 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import math
-from dataclasses import dataclass
-from typing import Literal
-
-import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import BaseOutput
-from diffusers.utils.torch_utils import randn_tensor
-
-
-@dataclass
-class FlowMatchSDEDiscreteSchedulerOutputForTest(BaseOutput):
- """
- Output class for the scheduler's `step` function output.
-
- Args:
- prev_sample (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
- Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
- denoising loop.
- log_prob (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
- The log probability of the previous sample.
- prev_sample_mean (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
- The mean of the computed sample of previous timestep.
- std_dev_t (`torch.FloatTensor` of shape `(batch_size, 1, 1)`):
- The standard deviation used to compute `prev_sample`.
- """
-
- prev_sample: torch.FloatTensor
- log_prob: torch.FloatTensor | None
- prev_sample_mean: torch.FloatTensor
- std_dev_t: torch.FloatTensor
-
-
-class FlowMatchSDEDiscreteSchedulerForTest(FlowMatchEulerDiscreteScheduler):
- """SDE version of the FlowMatchEulerDiscreteScheduler.
- The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470)
- and diffusers v0.37 branch.
- """
-
- def step(
- self,
- model_output: torch.FloatTensor,
- timestep: float | torch.FloatTensor,
- sample: torch.FloatTensor,
- s_churn: float = 0.0,
- s_tmin: float = 0.0,
- s_tmax: float = float("inf"),
- s_noise: float = 1.0,
- generator: torch.Generator | None = None,
- per_token_timesteps: torch.Tensor | None = None,
- return_dict: bool = True,
- noise_level: float = 0.7,
- prev_sample: torch.FloatTensor | None = None,
- sde_type: Literal["sde", "cps"] = "sde",
- return_logprobs: bool = True,
- ) -> FlowMatchSDEDiscreteSchedulerOutputForTest | tuple:
- """
- Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
- process from the learned model outputs (most often the predicted noise).
-
- Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py
-
- Args:
- model_output (`torch.FloatTensor`):
- The direct output from learned diffusion model.
- timestep (`float`):
- The current discrete timestep in the diffusion chain.
- sample (`torch.FloatTensor`):
- A current instance of a sample created by the diffusion process.
- s_churn (`float`):
- s_tmin (`float`):
- s_tmax (`float`):
- s_noise (`float`, defaults to 1.0):
- Scaling factor for noise added to the sample.
- generator (`torch.Generator`, *optional*):
- A random number generator.
- per_token_timesteps (`torch.Tensor`, *optional*):
- The timesteps for each token in the sample.
- return_dict (`bool`):
- Whether or not to return a
- [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput`] or tuple.
- noise_level (`float`, *optional*, defaults to 0.7):
- The noise level used in the SDE.
- prev_sample (`torch.FloatTensor`, *optional*):
- The sample from the previous timestep. If not provided, it will be sampled inside the function.
- sde_type (`str`, *optional*, defaults to "sde"):
- The type of SDE to use. Choose between "sde" and "cps".
- return_logprobs (`bool`, *optional*, defaults to True):
- Whether to return log probabilities of the previous sample.
- """
-
- if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
- raise ValueError(
- (
- "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
- " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
- " one of the `scheduler.timesteps` as a timestep."
- ),
- )
-
- if self.step_index is None:
- self._init_step_index(timestep)
-
- # Upcast to avoid precision issues when computing prev_sample
- sample = sample.to(torch.float32)
- if prev_sample is not None:
- prev_sample = prev_sample.to(torch.float32)
-
- prev_sample, log_prob, prev_sample_mean, std_dev_t = self.sample_previous_step(
- sample=sample,
- model_output=model_output,
- generator=generator,
- per_token_timesteps=per_token_timesteps,
- noise_level=noise_level,
- prev_sample=prev_sample,
- sde_type=sde_type,
- return_logprobs=return_logprobs,
- )
-
- # upon completion increase step index by one
- self._step_index += 1
- if per_token_timesteps is None:
- # Cast sample back to model compatible dtype
- prev_sample = prev_sample.to(model_output.dtype)
-
- if not return_dict:
- return (prev_sample, log_prob, prev_sample_mean, std_dev_t)
-
- return FlowMatchSDEDiscreteSchedulerOutputForTest(
- prev_sample=prev_sample, log_prob=log_prob, prev_sample_mean=prev_sample_mean, std_dev_t=std_dev_t
- )
-
- def sample_previous_step(
- self,
- sample: torch.Tensor,
- model_output: torch.Tensor,
- timestep: torch.FloatTensor | None = None,
- generator: torch.Generator | None = None,
- per_token_timesteps: torch.Tensor | None = None,
- noise_level: float = 0.7,
- prev_sample: torch.Tensor | None = None,
- sde_type: Literal["cps", "sde"] = "sde",
- return_logprobs: bool = True,
- ):
- assert sde_type in ["sde", "cps"]
- assert sample.dtype == torch.float32
- if prev_sample is not None:
- assert prev_sample.dtype == torch.float32
-
- if per_token_timesteps is not None:
- raise NotImplementedError("per_token_timesteps is not supported yet for FlowMatchSDEDiscreteScheduler.")
- else:
- if timestep is None:
- sigma_idx = self.step_index
- sigma = self.sigmas[sigma_idx]
- sigma_prev = self.sigmas[sigma_idx + 1]
- else:
- sigma_idx = torch.tensor([self.index_for_timestep(t) for t in timestep])
- sigma = self.sigmas[sigma_idx].view(-1, *([1] * (len(sample.shape) - 1)))
- sigma_prev = self.sigmas[sigma_idx + 1].view(-1, *([1] * (len(sample.shape) - 1)))
-
- sigma_max = self.sigmas[1]
- dt = sigma_prev - sigma
-
- if sde_type == "sde":
- std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level
-
- prev_sample_mean = (
- sample * (1 + std_dev_t**2 / (2 * sigma) * dt)
- + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
- )
-
- if prev_sample is None:
- variance_noise = randn_tensor(
- model_output.shape,
- generator=generator,
- device=model_output.device,
- dtype=model_output.dtype,
- )
- prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise
-
- if return_logprobs:
- log_prob = (
- -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2))
- - torch.log(std_dev_t * torch.sqrt(-1 * dt))
- - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
- )
- else:
- log_prob = None
-
- elif sde_type == "cps":
- std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2)
- pred_original_sample = sample - sigma * model_output
- noise_estimate = sample + model_output * (1 - sigma)
- prev_sample_mean = pred_original_sample * (1 - sigma_prev) + noise_estimate * torch.sqrt(
- sigma_prev**2 - std_dev_t**2
- )
-
- if prev_sample is None:
- variance_noise = randn_tensor(
- model_output.shape,
- generator=generator,
- device=model_output.device,
- dtype=model_output.dtype,
- )
- prev_sample = prev_sample_mean + std_dev_t * variance_noise
-
- if return_logprobs:
- log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2)
- else:
- log_prob = None
-
- # mean along all but batch dimension
- log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) if log_prob is not None else None
- return prev_sample, log_prob, prev_sample_mean, std_dev_t
diff --git a/tests/e2e/offline_inference/custom_pipeline/worker_extension.py b/tests/e2e/offline_inference/custom_pipeline/worker_extension.py
deleted file mode 100644
index 6ed29d5a8..000000000
--- a/tests/e2e/offline_inference/custom_pipeline/worker_extension.py
+++ /dev/null
@@ -1,77 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Test worker extension that mirrors
-``verl_omni.workers.rollout.vllm_rollout.utils.vLLMOmniColocateWorkerExtension``
-
-Fidelity goals:
-
-* ``__new__`` calls ``set_death_signal()`` (Linux ``PR_SET_PDEATHSIG``) so
- spawned vLLM workers die with the parent.
-* ``__new__`` calls ``VLLMOmniHijackForTest.hijack()`` which monkey-patches
- ``vllm_omni.diffusion.lora.manager.DiffusionLoRAManager._load_adapter``
- to accept in-memory LoRA tensors (verbatim port of
- ``verl_omni.utils.vllm_omni.utils.VLLMOmniHijack.hijack``).
-
-The ``verl.utils.vllm.VLLMHijack.hijack()`` call (which patches vLLM's
-``LRUCacheWorkerLoRAManager._load_adapter`` for AR/text mode) is
-intentionally skipped because (a) ``verl`` is not importable from this
-repo and (b) the diffusion test path never reaches that manager.
-"""
-
-from __future__ import annotations
-
-import ctypes
-import logging
-import os
-import platform
-import signal
-
-from msgspec import field
-
-from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
-from vllm_omni.lora.request import LoRARequest as OmniLoRARequest
-
-logger = logging.getLogger(__name__)
-
-
-class OmniTensorLoRARequestForTest(OmniLoRARequest):
- peft_config: dict = field(default=None)
- lora_tensors: dict = field(default=None)
-
-
-def set_death_signal() -> None:
- """Verbatim port of ``verl.workers.rollout.vllm_rollout.utils.set_death_signal``."""
- if platform.system() != "Linux":
- return
- try:
- libc = ctypes.CDLL("libc.so.6")
- libc.prctl(1, signal.SIGKILL)
- if os.getppid() == 1:
- os.kill(os.getpid(), signal.SIGKILL)
- except Exception: # noqa: BLE001
- # libc.so.6 may not be present (musl, alpine) — best-effort only.
- pass
-
-
-# ---------------------------------------------------------------------
-# The worker extension itself
-# ---------------------------------------------------------------------
-
-
-class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension):
- """Mirror of ``vLLMOmniColocateWorkerExtension`` (verl-omni).
-
- The production ``__new__`` runs ``set_death_signal`` + ``VLLMOmniHijack.hijack``
- on every vLLM worker process. Replicating both ensures this test
- reproduces the same monkey-patched environment.
- """
-
- def __new__(cls, **kwargs):
- set_death_signal()
- return super().__new__(cls)
-
- @staticmethod
- def test_extension_name() -> str:
- """Return a stable identifier for assertions in unit tests."""
- return "vllm-omni-colocate-worker-extension-for-test"
diff --git a/tests/e2e/offline_inference/test_soulxsinger.py b/tests/e2e/offline_inference/test_soulxsinger_expansion.py
similarity index 99%
rename from tests/e2e/offline_inference/test_soulxsinger.py
rename to tests/e2e/offline_inference/test_soulxsinger_expansion.py
index 465f5547b..c3a6091d6 100644
--- a/tests/e2e/offline_inference/test_soulxsinger.py
+++ b/tests/e2e/offline_inference/test_soulxsinger_expansion.py
@@ -38,7 +38,7 @@ if not PROMPT_AUDIO.is_file() or not TARGET_AUDIO.is_file():
allow_module_level=True,
)
-pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.tts]
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.tts]
_CASES = (
pytest.param(
diff --git a/tests/e2e/offline_inference/test_step_audio2_expansion.py b/tests/e2e/offline_inference/test_step_audio2_expansion.py
index 362f2f33f..00fc7bbe2 100644
--- a/tests/e2e/offline_inference/test_step_audio2_expansion.py
+++ b/tests/e2e/offline_inference/test_step_audio2_expansion.py
@@ -6,7 +6,6 @@ E2E expansion tests for Step-Audio2 offline inference (nightly CI).
Full-model inference with audio input and audio output.
"""
-from pathlib import Path
from typing import Any
import numpy as np
@@ -15,10 +14,11 @@ import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.media import generate_synthetic_audio
from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.outputs import OmniRequestOutput
MODEL = "stepfun-ai/Step-Audio-2-mini"
-STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml")
+STAGE_CONFIG = get_deploy_config_path("step_audio2_thinker_ci.yaml")
TEST_PARAMS = [(MODEL, STAGE_CONFIG)]
pytestmark = [pytest.mark.slow, pytest.mark.tts]
diff --git a/tests/e2e/online_serving/test_dreamzero_expansion.py b/tests/e2e/online_serving/test_dreamzero_expansion.py
index 38e7d4b16..66d558120 100644
--- a/tests/e2e/online_serving/test_dreamzero_expansion.py
+++ b/tests/e2e/online_serving/test_dreamzero_expansion.py
@@ -1,34 +1,37 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""E2E online serving test for DreamZero OpenPI websocket serving."""
+"""E2E online serving tests for DreamZero OpenPI websocket serving."""
from __future__ import annotations
import os
from pathlib import Path
+import cv2
import numpy as np
import pytest
-from tests.dreamzero import openpi_client_helper as openpi_client
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServerParams, get_open_port
+from tests.helpers.runtime import (
+ DREAMZERO_ACTION_DIM,
+ DREAMZERO_ACTION_HORIZON,
+ DREAMZERO_CAMERA_FILES,
+ OmniServerParams,
+ OpenPIWebSocketResponse,
+ get_open_port,
+)
+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 = "GEAR-Dreams/DreamZero-DROID"
-
test_params = [
OmniServerParams(
model=MODEL,
- port=8091,
- server_args=[
- "--deploy-config",
- "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml",
- ],
+ stage_config_path=get_deploy_config_path("dreamzero_tp1_cfg2.yaml"),
env_dict={
"ATTENTION_BACKEND": "torch",
"DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
@@ -39,10 +42,10 @@ test_params = [
]
-def _write_synthetic_video(path: Path, cv2_module, *, channel: int) -> None:
+def _write_synthetic_video(path: Path, *, channel: int) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
height, width, num_frames = 180, 320, 24
- writer = cv2_module.VideoWriter(str(path), cv2_module.VideoWriter_fourcc(*"mp4v"), 15.0, (width, height))
+ writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 15.0, (width, height))
if not writer.isOpened():
raise RuntimeError(f"Failed to open video writer for {path}")
try:
@@ -50,39 +53,47 @@ def _write_synthetic_video(path: Path, cv2_module, *, channel: int) -> None:
frame = np.zeros((height, width, 3), dtype=np.uint8)
frame[..., channel] = (frame_idx * 7) % 255
frame[..., (channel + 1) % 3] = 64
- writer.write(cv2_module.cvtColor(frame, cv2_module.COLOR_RGB2BGR))
+ writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
finally:
writer.release()
-def _write_synthetic_dreamzero_videos(client_mod, video_dir: Path) -> None:
- for channel, file_name in enumerate(client_mod.CAMERA_FILES.values()):
- _write_synthetic_video(video_dir / file_name, client_mod.cv2, channel=channel)
+def _write_synthetic_dreamzero_videos(video_dir: Path) -> None:
+ for channel, file_name in enumerate(DREAMZERO_CAMERA_FILES.values()):
+ _write_synthetic_video(video_dir / file_name, channel=channel)
+
+
+def _validate_dreamzero_openpi_session(response: OpenPIWebSocketResponse) -> None:
+ metadata = response.server_metadata
+ assert tuple(metadata["image_resolution"]) == (180, 320)
+ assert metadata["n_external_cameras"] == 2
+ assert metadata["needs_wrist_camera"] is True
+ assert metadata["needs_stereo_camera"] is False
+ assert metadata["needs_session_id"] is True
+ assert metadata["action_space"] == "joint_position"
+ assert response.operation_responses[-1]["status"] == "reset successful"
+
+ action_tensors = response.action_tensors
+ assert action_tensors is not None
+ assert len(action_tensors) == 3
+ for index, action in enumerate(action_tensors):
+ assert action.shape == (DREAMZERO_ACTION_HORIZON, DREAMZERO_ACTION_DIM)
+ assert action.dtype == np.float32
+ assert np.isfinite(action).all()
@pytest.mark.slow
@pytest.mark.diffusion
-@pytest.mark.distributed_cuda
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_dreamzero_openpi_online(omni_server, tmp_path: Path) -> None:
- try:
- openpi_client.require_dependencies()
- except ModuleNotFoundError as exc:
- pytest.skip(str(exc))
-
+def test_dreamzero_openpi_online(omni_server, openai_client, tmp_path: Path) -> None:
video_dir = tmp_path / "dreamzero_videos"
- _write_synthetic_dreamzero_videos(openpi_client, video_dir)
- result = openpi_client.run_policy_session(
- host=omni_server.host,
- port=omni_server.port,
- video_dir=video_dir,
- session_id="dreamzero-online-e2e",
- )
-
- openpi_client.validate_session_result(result)
-
- metadata = result["metadata"]
- assert metadata["needs_session_id"] is True
- assert metadata["needs_stereo_camera"] is False
- assert tuple(metadata["image_resolution"]) == (180, 320)
+ _write_synthetic_dreamzero_videos(video_dir)
+ response = openai_client.send_robot_openpi_ws_request(
+ {
+ "run_dreamzero_policy_session": True,
+ "video_dir": video_dir,
+ "session_id": "dreamzero-online-e2e",
+ }
+ )[0]
+ _validate_dreamzero_openpi_session(response)
diff --git a/tests/e2e/online_serving/test_gr00t_openpi.py b/tests/e2e/online_serving/test_gr00t_openpi_expansion.py
similarity index 66%
rename from tests/e2e/online_serving/test_gr00t_openpi.py
rename to tests/e2e/online_serving/test_gr00t_openpi_expansion.py
index 40b8293d6..70e11088d 100644
--- a/tests/e2e/online_serving/test_gr00t_openpi.py
+++ b/tests/e2e/online_serving/test_gr00t_openpi_expansion.py
@@ -7,14 +7,19 @@ import os
import numpy as np
import pytest
-from tests.gr00t import openpi_client_helper as openpi_client
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServerParams
+from tests.helpers.runtime import (
+ OmniServerParams,
+ OpenPIWebSocketResponse,
+ build_openpi_droid_observation,
+)
from tests.helpers.stage_config import get_deploy_config_path
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
MODEL = "nvidia/GR00T-N1.7-3B"
+ACTION_KEYS = frozenset({"eef_9d", "gripper_position", "joint_position"})
+EXPECTED_EMBODIMENT_TAG = "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT"
pytest.importorskip("websockets")
pytest.importorskip("openpi_client.msgpack_numpy")
@@ -34,7 +39,7 @@ test_params = [
]
# Reference values captured from the Isaac-GR00T ZMQ server (nvidia/GR00T-N1.7-3B,
-# OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT embodiment) using build_droid_observation()
+# OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT embodiment) using build_openpi_droid_observation()
# inputs: zero images (256×256), identity eef_9d, zero gripper/joint, "pick up the object".
# Outputs are bit-reproducible across resets and across runs (max_diff=0.0).
_REF_EEF_9D = np.array(
@@ -112,32 +117,65 @@ _REF_JOINT = np.array(
) # rows = step 0, step 4, step 39
-@pytest.mark.full_model
+def _validate_gr00t_openpi_session(response: OpenPIWebSocketResponse) -> None:
+ metadata = response.server_metadata
+ for key in ("action_horizon", "action_keys", "embodiment_tag", "needs_session_id"):
+ assert key in metadata, f"Missing GR00T metadata key: {key}"
+
+ assert metadata["embodiment_tag"] == EXPECTED_EMBODIMENT_TAG
+ assert metadata["needs_session_id"] is True
+ assert frozenset(metadata["action_keys"]) == ACTION_KEYS
+ assert response.operation_responses[-1]["status"] == "reset successful"
+
+ actions = response.actions
+ assert actions is not None
+ assert frozenset(actions) == ACTION_KEYS
+
+ action_horizon = int(metadata["action_horizon"])
+ expected_shapes = {
+ "eef_9d": (1, action_horizon, 9),
+ "gripper_position": (1, action_horizon, 1),
+ "joint_position": (1, action_horizon, 7),
+ }
+ for key, expected_shape in expected_shapes.items():
+ assert actions[key].shape == expected_shape
+ assert actions[key].dtype == np.float32
+ assert np.isfinite(actions[key]).all()
+
+
+@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_gr00t_n1d7_openpi_online(omni_server) -> None:
- result = openpi_client.run_policy_session(
- host=omni_server.host,
- port=omni_server.port,
- session_id="gr00t-online-e2e",
- )
- openpi_client.validate_session_result(result)
+def test_gr00t_n1d7_openpi_online(omni_server, openai_client) -> None:
+ response = openai_client.send_robot_openpi_ws_request(
+ {
+ "run_default_policy_session": True,
+ "session_id": "gr00t-online-e2e",
+ }
+ )[0]
+ _validate_gr00t_openpi_session(response)
-@pytest.mark.full_model
+@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_gr00t_n1d7_openpi_precision(omni_server) -> None:
+def test_gr00t_n1d7_openpi_precision(omni_server, openai_client) -> None:
"""Assert actions match Isaac-GR00T reference (GR00T_NOISE_SEED=42, zero inputs)."""
- client = openpi_client.OpenPIWebsocketClient(host=omni_server.host, port=omni_server.port)
- try:
- obs = openpi_client.build_droid_observation(session_id="gr00t-precision-e2e")
- client.reset({})
- actions = client.infer(obs)
- finally:
- client.close()
+ response = openai_client.send_robot_openpi_ws_request(
+ {
+ "operations": [
+ {"endpoint": "reset", "payload": {}},
+ {
+ "endpoint": "infer",
+ "payload": build_openpi_droid_observation(session_id="gr00t-precision-e2e"),
+ },
+ ],
+ }
+ )[0]
+ actions = response.actions
+ assert actions is not None
steps = [0, 4, 39]
np.testing.assert_allclose(
diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py
index ff41ee413..ab35b74ea 100644
--- a/tests/e2e/online_serving/test_omnivoice_expansion.py
+++ b/tests/e2e/online_serving/test_omnivoice_expansion.py
@@ -25,7 +25,7 @@ try:
except ImportError:
_HAS_VOICE_CLONE = False
-pytestmark = [pytest.mark.full_model, pytest.mark.tts]
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
MODEL = "k2-fsa/OmniVoice"
diff --git a/tests/test_arg_utils.py b/tests/engine/test_arg_utils_shared_fields.py
similarity index 100%
rename from tests/test_arg_utils.py
rename to tests/engine/test_arg_utils_shared_fields.py
diff --git a/tests/test_data_entry_keys.py b/tests/engine/test_data_entry_keys.py
similarity index 100%
rename from tests/test_data_entry_keys.py
rename to tests/engine/test_data_entry_keys.py
diff --git a/tests/test_omni_request.py b/tests/engine/test_omni_request.py
similarity index 100%
rename from tests/test_omni_request.py
rename to tests/engine/test_omni_request.py
diff --git a/tests/test_outputs.py b/tests/engine/test_omni_request_output.py
similarity index 100%
rename from tests/test_outputs.py
rename to tests/engine/test_omni_request_output.py
diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py
index 0553c03ff..ed32732ba 100644
--- a/tests/engine/test_orchestrator_error_handling.py
+++ b/tests/engine/test_orchestrator_error_handling.py
@@ -73,9 +73,17 @@ def orchestrator_factory():
class FakeDeadLLMStageClient(FakeStageClient):
- """LLM stage client that raises EngineDeadError on get_output_async."""
+ """LLM stage client that raises EngineDeadError once a request is in-flight.
+
+ Raising on the first poll races with ``add_request``: under suite load the
+ orchestrator can die and drain before the test enqueues anything, so no
+ fatal ErrorMessage is ever published. Keep returning empty outputs until
+ ``add_request_async`` has been called, then fail on the next poll.
+ """
async def get_output_async(self):
+ if not self.add_request_calls:
+ return SimpleNamespace(outputs=[])
raise EngineDeadError("Stage-0 engine core is dead")
diff --git a/tests/test_speaker_metadata_persistence.py b/tests/entrypoints/openai_api/test_speaker_metadata_persistence.py
similarity index 100%
rename from tests/test_speaker_metadata_persistence.py
rename to tests/entrypoints/openai_api/test_speaker_metadata_persistence.py
diff --git a/tests/entrypoints/openai_api/test_speech_usage.py b/tests/entrypoints/openai_api/test_speech_usage.py
index 8211dab88..2fba10516 100644
--- a/tests/entrypoints/openai_api/test_speech_usage.py
+++ b/tests/entrypoints/openai_api/test_speech_usage.py
@@ -16,6 +16,8 @@ from vllm_omni.entrypoints.openai.speech_usage import (
qwen3_tts_input_token_details,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
# A deterministic stand-in for a real tokenizer: 1 token per whitespace word.
def word_count(text: str) -> int:
diff --git a/tests/entrypoints/openai_api/test_stage_params.py b/tests/entrypoints/openai_api/test_stage_params.py
index f9fd40349..5db768b30 100644
--- a/tests/entrypoints/openai_api/test_stage_params.py
+++ b/tests/entrypoints/openai_api/test_stage_params.py
@@ -3,6 +3,7 @@
from types import SimpleNamespace
+import pytest
from vllm import SamplingParams
from vllm_omni.entrypoints.openai.stage_params import (
@@ -12,6 +13,8 @@ from vllm_omni.entrypoints.openai.stage_params import (
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def test_resolve_stage_sampling_params_clones_stage_default():
default = SamplingParams(temperature=0.2, seed=11)
diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py
index 2930a204d..5f6d5930a 100644
--- a/tests/entrypoints/openai_api/test_tts_adapter.py
+++ b/tests/entrypoints/openai_api/test_tts_adapter.py
@@ -15,6 +15,9 @@ from vllm_omni.entrypoints.openai.tts_adapters import (
)
from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
# Every dedicated TTS model-type must have an adapter so the orchestrator's
# uniform ``self._adapter.build(...)`` dispatch covers it.
EXPECTED_MODEL_TYPES = {
diff --git a/tests/entrypoints/openai_api/test_video_frame_filter.py b/tests/entrypoints/openai_api/test_video_frame_filter.py
index 70df6e34c..ecada7a55 100644
--- a/tests/entrypoints/openai_api/test_video_frame_filter.py
+++ b/tests/entrypoints/openai_api/test_video_frame_filter.py
@@ -12,6 +12,8 @@ from tests.entrypoints.openai_api.conftest_video import (
)
from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
class TestFrameSimilarityFilter:
def test_first_frame_always_retained(self):
diff --git a/tests/entrypoints/openai_api/test_video_stream_handler.py b/tests/entrypoints/openai_api/test_video_stream_handler.py
index a5ce3d616..e382b5e32 100644
--- a/tests/entrypoints/openai_api/test_video_stream_handler.py
+++ b/tests/entrypoints/openai_api/test_video_stream_handler.py
@@ -22,6 +22,9 @@ from vllm_omni.entrypoints.openai.video_stream_session import (
VideoStreamHandler,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
diff --git a/tests/entrypoints/openai_api/test_video_stream_session.py b/tests/entrypoints/openai_api/test_video_stream_session.py
index b0a12b1e4..b91723c26 100644
--- a/tests/entrypoints/openai_api/test_video_stream_session.py
+++ b/tests/entrypoints/openai_api/test_video_stream_session.py
@@ -15,6 +15,9 @@ from vllm_omni.entrypoints.openai.video_stream_session import (
VideoStreamSession,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
# ---------------------------------------------------------------------------
# VideoStreamConfig
# ---------------------------------------------------------------------------
diff --git a/tests/entrypoints/test_pd_disaggregation.py b/tests/entrypoints/test_pd_disaggregation.py
index 774e6d1a8..c896ccb88 100644
--- a/tests/entrypoints/test_pd_disaggregation.py
+++ b/tests/entrypoints/test_pd_disaggregation.py
@@ -22,6 +22,9 @@ from vllm import SamplingParams
from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
pytestmark = pytest.mark.skip(reason="Temporarily skip PD entrypoint tests while PD config is being removed.")
# Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies.
diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py
index c4c0667df..e2890797f 100644
--- a/tests/entrypoints/test_stage_utils.py
+++ b/tests/entrypoints/test_stage_utils.py
@@ -6,6 +6,8 @@ from pytest_mock import MockerFixture
from vllm_omni.entrypoints.stage_utils import _map_device_list, resolve_stage_physical_devices, set_stage_devices
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def _make_dummy_torch(call_log):
class _Props:
diff --git a/tests/gr00t/__init__.py b/tests/gr00t/__init__.py
deleted file mode 100644
index 208f01a7c..000000000
--- a/tests/gr00t/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/tests/gr00t/openpi_client_helper.py b/tests/gr00t/openpi_client_helper.py
deleted file mode 100644
index e5943ca33..000000000
--- a/tests/gr00t/openpi_client_helper.py
+++ /dev/null
@@ -1,183 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from dataclasses import dataclass
-from typing import Any
-
-import numpy as np
-
-try:
- import websockets.sync.client as websockets_client
-except ImportError: # pragma: no cover - optional e2e dependency
- websockets_client = None
-
-try:
- from openpi_client import msgpack_numpy
-except ImportError: # pragma: no cover - optional e2e dependency
- msgpack_numpy = None
-
-PING_INTERVAL_SECS = 300
-PING_TIMEOUT_SECS = 3600
-DEFAULT_HOST = "127.0.0.1"
-DEFAULT_PORT = 8000
-DEFAULT_PATH = "/v1/realtime/robot/openpi"
-DEFAULT_SESSION_ID = "gr00t-smoke"
-ACTION_KEYS = {"eef_9d", "gripper_position", "joint_position"}
-LANGUAGE_KEY = "annotation.language.language_instruction"
-
-
-def _identity_eef_9d_state() -> np.ndarray:
- state = np.zeros((1, 1, 9), dtype=np.float32)
- state[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32)
- return state
-
-
-def require_dependencies() -> None:
- missing = []
- if websockets_client is None:
- missing.append("websockets")
- if msgpack_numpy is None:
- missing.append("openpi-client")
- if missing:
- raise ModuleNotFoundError(f"GR00T OpenPI test dependencies are missing: {', '.join(missing)}")
-
-
-@dataclass(frozen=True)
-class Gr00tServerMetadata:
- action_horizon: int
- action_keys: set[str]
- embodiment_tag: str
- needs_session_id: bool
-
- @classmethod
- def from_dict(cls, payload: dict[str, Any]):
- required_keys = ("action_horizon", "action_keys", "embodiment_tag", "needs_session_id")
- missing_keys = [key for key in required_keys if key not in payload]
- if missing_keys:
- raise ValueError(f"Missing GR00T metadata keys: {missing_keys}")
-
- return cls(
- action_horizon=int(payload["action_horizon"]),
- action_keys={str(key) for key in payload["action_keys"]},
- embodiment_tag=str(payload["embodiment_tag"]),
- needs_session_id=bool(payload["needs_session_id"]),
- )
-
-
-class OpenPIWebsocketClient:
- def __init__(
- self,
- *,
- host: str = DEFAULT_HOST,
- port: int = DEFAULT_PORT,
- path: str = DEFAULT_PATH,
- ) -> None:
- require_dependencies()
- self._uri = f"ws://{host}:{port}{path}"
- self._packer = msgpack_numpy.Packer()
- self._ws, self._server_metadata = self._connect()
-
- def _connect(self):
- conn = websockets_client.connect(
- self._uri,
- compression=None,
- max_size=None,
- ping_interval=PING_INTERVAL_SECS,
- ping_timeout=PING_TIMEOUT_SECS,
- )
- metadata = msgpack_numpy.unpackb(conn.recv())
- if not isinstance(metadata, dict):
- raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
- return conn, metadata
-
- def get_server_metadata(self) -> dict[str, Any]:
- return dict(self._server_metadata)
-
- def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]:
- payload = dict(obs)
- payload["endpoint"] = "infer"
- self._ws.send(self._packer.pack(payload))
- response = msgpack_numpy.unpackb(self._ws.recv())
- if isinstance(response, dict) and response.get("type") == "error":
- raise RuntimeError(f"Inference failed: {response['message']}")
- if not isinstance(response, dict):
- raise RuntimeError(f"Expected dict actions from GR00T OpenPI endpoint, got {type(response)!r}")
- return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
-
- def reset(self, reset_info: dict[str, Any] | None = None) -> str:
- payload = dict(reset_info or {})
- payload["endpoint"] = "reset"
- self._ws.send(self._packer.pack(payload))
- response = msgpack_numpy.unpackb(self._ws.recv())
- if not isinstance(response, dict) or response.get("status") != "reset successful":
- raise RuntimeError(f"Unexpected reset response: {response!r}")
- return str(response["status"])
-
- def close(self) -> None:
- self._ws.close()
-
-
-def build_droid_observation(*, session_id: str = DEFAULT_SESSION_ID) -> dict[str, Any]:
- return {
- "session_id": session_id,
- "video": {
- "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
- "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
- },
- "state": {
- "eef_9d": _identity_eef_9d_state(),
- "gripper_position": np.zeros((1, 1, 1), dtype=np.float32),
- "joint_position": np.zeros((1, 1, 7), dtype=np.float32),
- },
- "language": {LANGUAGE_KEY: [["pick up the object"]]},
- }
-
-
-def run_policy_session(
- *,
- host: str = DEFAULT_HOST,
- port: int = DEFAULT_PORT,
- path: str = DEFAULT_PATH,
- session_id: str = DEFAULT_SESSION_ID,
-) -> dict[str, Any]:
- client = OpenPIWebsocketClient(host=host, port=port, path=path)
- try:
- metadata = client.get_server_metadata()
- actions = client.infer(build_droid_observation(session_id=session_id))
- reset_status = client.reset({})
- return {
- "metadata": metadata,
- "actions": actions,
- "reset_status": reset_status,
- "session_id": session_id,
- }
- finally:
- client.close()
-
-
-def validate_session_result(result: dict[str, Any]) -> None:
- metadata = Gr00tServerMetadata.from_dict(result["metadata"])
- if metadata.embodiment_tag != "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT":
- raise AssertionError(f"Unexpected embodiment_tag: {metadata.embodiment_tag}")
- if not metadata.needs_session_id:
- raise AssertionError("GR00T test expects needs_session_id metadata")
- if metadata.action_keys != ACTION_KEYS:
- raise AssertionError(f"Unexpected action keys: {metadata.action_keys}")
- if result["reset_status"] != "reset successful":
- raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}")
-
- actions = result["actions"]
- if set(actions) != ACTION_KEYS:
- raise AssertionError(f"Unexpected action keys: {set(actions)}")
- expected_shapes = {
- "eef_9d": (1, metadata.action_horizon, 9),
- "gripper_position": (1, metadata.action_horizon, 1),
- "joint_position": (1, metadata.action_horizon, 7),
- }
- for key, expected_shape in expected_shapes.items():
- if actions[key].shape != expected_shape:
- raise AssertionError(f"Action {key} shape mismatch: expected {expected_shape}, got {actions[key].shape}")
- if actions[key].dtype != np.float32:
- raise AssertionError(f"Action {key} dtype mismatch: expected float32, got {actions[key].dtype}")
- if not np.isfinite(actions[key]).all():
- raise AssertionError(f"Action {key} contains non-finite values")
diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py
index 04d51fb90..9929beab8 100644
--- a/tests/helpers/__init__.py
+++ b/tests/helpers/__init__.py
@@ -1,10 +1,11 @@
"""Shared, importable test helper utilities.
Submodules (``assertions``, ``env``, ``media``, ``runtime``, …) are imported
-explicitly by callers. Avoid star-importing everything here: that ran before
-refactor only inside the old monolithic ``conftest``; a greedy ``__init__``
-changes import order and can affect in-process Omni (``OmniRunner`` / offline
-e2e) vs subprocess-based ``OmniServer`` tests.
+explicitly by callers. Unit tests for these helpers live under ``tests/``.
+Avoid star-importing everything here: that ran before refactor only inside the
+old monolithic ``conftest``; a greedy ``__init__`` changes import order and can
+affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based
+``OmniServer`` tests.
"""
import pytest
diff --git a/tests/helpers/env.py b/tests/helpers/env.py
index dcf2484cd..39cec30e7 100644
--- a/tests/helpers/env.py
+++ b/tests/helpers/env.py
@@ -27,6 +27,32 @@ def get_physical_device_indices(devices):
return [index_mapping[i] for i in devices if i in index_mapping]
+def pick_least_used_device_indices(num_devices: int) -> list[int]:
+ """Return physical device indices for the least-used accelerators.
+
+ Args:
+ num_devices: Number of device indices to return.
+
+ Returns:
+ Physical device indices sorted from least to most used memory.
+ """
+ if num_devices <= 0:
+ raise ValueError(f"num_devices must be positive, got {num_devices}.")
+
+ logical_count = current_omni_platform.get_device_count()
+ if logical_count < num_devices:
+ raise RuntimeError(f"Need {num_devices} devices, but only {logical_count} are available.")
+
+ device_usage: list[tuple[int, int]] = []
+ for device_index in range(logical_count):
+ with current_omni_platform.device(device_index):
+ free_bytes, total_bytes = current_omni_platform.get_device_memory()
+ device_usage.append((total_bytes - free_bytes, device_index))
+ device_usage.sort()
+ logical_indices = [device_index for _, device_index in device_usage[:num_devices]]
+ return get_physical_device_indices(logical_indices)
+
+
def wait_for_gpu_memory_to_clear(
*,
devices: list[int],
@@ -264,6 +290,7 @@ class DeviceMemoryMonitor:
__all__ = [
"DeviceMemoryMonitor",
"get_physical_device_indices",
+ "pick_least_used_device_indices",
"run_post_test_cleanup",
"run_pre_test_cleanup",
"wait_for_gpu_memory_to_clear",
diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py
index 39baf1ce0..f8a5416e4 100644
--- a/tests/helpers/runtime.py
+++ b/tests/helpers/runtime.py
@@ -22,6 +22,7 @@ from pathlib import Path
from typing import Any, NamedTuple, cast
from urllib.parse import quote
+import numpy as np
import psutil
import requests
import soundfile as sf
@@ -802,6 +803,189 @@ class WebSocketJsonResponse:
json_body: dict[str, Any]
+@dataclass
+class OpenPIWebSocketResponse:
+ """Msgpack WebSocket session against ``/v1/realtime/robot/openpi``."""
+
+ server_metadata: dict[str, Any]
+ operation_responses: list[Any]
+ actions: dict[str, np.ndarray] | None = None
+ action_tensors: list[np.ndarray] | None = None
+
+
+def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[str, Any]:
+ """Build a minimal DROID-style observation payload for the OpenPI robot endpoint."""
+ identity_eef_9d = np.zeros((1, 1, 9), dtype=np.float32)
+ identity_eef_9d[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32)
+ return {
+ "session_id": session_id,
+ "video": {
+ "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
+ "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
+ },
+ "state": {
+ "eef_9d": identity_eef_9d,
+ "gripper_position": np.zeros((1, 1, 1), dtype=np.float32),
+ "joint_position": np.zeros((1, 1, 7), dtype=np.float32),
+ },
+ "language": {"annotation.language.language_instruction": [["pick up the object"]]},
+ }
+
+
+DREAMZERO_DEFAULT_PROMPT = (
+ "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
+)
+DREAMZERO_ACTION_HORIZON = 24
+DREAMZERO_ACTION_DIM = 8
+DREAMZERO_CAMERA_FILES = {
+ "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+ "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+ "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+
+def _require_opencv() -> Any:
+ try:
+ import cv2
+ except ImportError as exc: # pragma: no cover - optional e2e dependency
+ raise ModuleNotFoundError("DreamZero OpenPI test dependencies are missing: opencv-python") from exc
+ return cv2
+
+
+def load_dreamzero_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
+ cv2 = _require_opencv()
+ camera_frames: dict[str, np.ndarray] = {}
+ for camera_key, file_name in DREAMZERO_CAMERA_FILES.items():
+ video_path = video_dir / file_name
+ if not video_path.exists():
+ raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}")
+ cap = cv2.VideoCapture(str(video_path))
+ frames = []
+ while True:
+ ok, frame = cap.read()
+ if not ok:
+ break
+ frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+ cap.release()
+ if not frames:
+ raise RuntimeError(f"No frames loaded from {video_path}")
+ camera_frames[camera_key] = np.stack(frames, axis=0)
+ return camera_frames
+
+
+def build_dreamzero_demo_observations(
+ camera_frames: dict[str, np.ndarray],
+ *,
+ prompt: str,
+ session_id: str,
+ num_chunks: int = 2,
+) -> list[dict[str, Any]]:
+ if num_chunks < 1:
+ raise ValueError("num_chunks must be at least 1")
+
+ relative_offsets = [-23, -16, -8, 0]
+ total_frames = min(frames.shape[0] for frames in camera_frames.values())
+ frame_schedules = [[0]]
+ current_frame = 23
+ for _ in range(num_chunks - 1):
+ indices = [max(current_frame + offset, 0) for offset in relative_offsets]
+ if indices[-1] >= total_frames:
+ break
+ frame_schedules.append(indices)
+ current_frame += DREAMZERO_ACTION_HORIZON
+
+ observations: list[dict[str, Any]] = []
+ for frame_indices in frame_schedules:
+ obs: dict[str, Any] = {}
+ for camera_key, all_frames in camera_frames.items():
+ selected = all_frames[frame_indices]
+ obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
+ obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+ obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+ obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+ obs["prompt"] = prompt
+ obs["session_id"] = session_id
+ observations.append(obs)
+ return observations
+
+
+class OpenPIWebSocketSession:
+ """Persistent msgpack session for ``/v1/realtime/robot/openpi``."""
+
+ DEFAULT_PING_INTERVAL_SECS = 300
+ DEFAULT_PING_TIMEOUT_SECS = 3600
+
+ @staticmethod
+ def _require_dependencies() -> tuple[Any, Any]:
+ try:
+ import websockets.sync.client as websockets_client
+ except ImportError as exc: # pragma: no cover - optional e2e dependency
+ raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: websockets") from exc
+ try:
+ from openpi_client import msgpack_numpy
+ except ImportError as exc: # pragma: no cover - optional e2e dependency
+ raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: openpi-client") from exc
+ return websockets_client, msgpack_numpy
+
+ def __init__(
+ self,
+ uri: str,
+ *,
+ ping_interval: float = DEFAULT_PING_INTERVAL_SECS,
+ ping_timeout: float = DEFAULT_PING_TIMEOUT_SECS,
+ open_timeout: float = 120.0,
+ close_timeout: float = 120.0,
+ ) -> None:
+ websockets_client, msgpack_numpy = self._require_dependencies()
+ self._msgpack_numpy = msgpack_numpy
+ self._packer = msgpack_numpy.Packer()
+ self._conn = websockets_client.connect(
+ uri,
+ compression=None,
+ max_size=None,
+ ping_interval=ping_interval,
+ ping_timeout=ping_timeout,
+ open_timeout=open_timeout,
+ close_timeout=close_timeout,
+ )
+ server_metadata = msgpack_numpy.unpackb(self._conn.recv())
+ if not isinstance(server_metadata, dict):
+ raise TypeError(f"Expected dict metadata from {uri}, got {type(server_metadata)!r}")
+ self._server_metadata = server_metadata
+
+ def get_server_metadata(self) -> dict[str, Any]:
+ return dict(self._server_metadata)
+
+ def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]:
+ response = self._send_operation("infer", obs)
+ if not isinstance(response, dict):
+ raise TypeError(f"Expected dict infer response, got {type(response)!r}")
+ return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
+
+ def reset(self, reset_info: dict[str, Any] | None = None) -> str:
+ response = self._send_operation("reset", dict(reset_info or {}))
+ return str(response["status"])
+
+ def close(self) -> None:
+ self._conn.close()
+
+ def _send_operation(self, endpoint: str, payload: dict[str, Any]) -> Any:
+ body = dict(payload)
+ body["endpoint"] = endpoint
+ self._conn.send(self._packer.pack(body))
+ raw = self._conn.recv()
+ if isinstance(raw, str):
+ raise RuntimeError(f"OpenPI {endpoint!r} failed: {raw}")
+ decoded = self._msgpack_numpy.unpackb(raw)
+ if isinstance(decoded, dict) and decoded.get("type") == "error":
+ raise RuntimeError(f"OpenPI {endpoint!r} failed: {decoded.get('message')}")
+ if endpoint == "reset":
+ if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
+ raise RuntimeError(f"Unexpected OpenPI reset response: {decoded!r}")
+ return decoded
+ return decoded
+
+
def _merge_http_expectation_kwargs(
base: dict[str, Any] | None,
*,
@@ -1593,6 +1777,92 @@ class OpenAIClientHandler:
)
return self._send_websocket_first_json_request("/v1/realtime", cfg)
+ def send_robot_openpi_ws_request(
+ self,
+ request_config: dict[str, Any] | None = None,
+ ) -> list[OpenPIWebSocketResponse]:
+ """WebSocket ``/v1/realtime/robot/openpi`` — msgpack metadata plus optional infer/reset ops.
+
+ ``request_config`` keys:
+
+ - ``operations``: optional sequence of ``{"endpoint": "infer"|"reset", "payload": {...}}``.
+ - ``run_default_policy_session``: when true and ``operations`` is omitted, run infer then reset
+ using :func:`build_openpi_droid_observation`.
+ - ``run_dreamzero_policy_session``: when true and ``operations`` is omitted, run the DreamZero
+ infer/reset/infer sequence using :func:`build_dreamzero_demo_observations`.
+ - ``video_dir``: required for ``run_dreamzero_policy_session``.
+ - ``prompt``: language instruction for DreamZero observations (optional).
+ - ``session_id``: used by the default policy session observation builder.
+ - ``num_chunks``: DreamZero chunk count (default ``2``).
+ - ``timeout``: seconds to wait for each inbound msgpack frame (default ``120``).
+ - ``ping_interval`` / ``ping_timeout``: forwarded to :func:`websockets.sync.client.connect`.
+ """
+ cfg = dict(request_config or {})
+ operations_cfg = cfg.get("operations")
+ if operations_cfg is not None:
+ operations = [dict(op) for op in operations_cfg]
+ elif cfg.get("run_dreamzero_policy_session"):
+ video_dir = cfg.get("video_dir")
+ if video_dir is None:
+ raise ValueError("run_dreamzero_policy_session requires video_dir")
+ session_id = str(cfg.get("session_id", "dreamzero-smoke"))
+ prompt = str(cfg.get("prompt", DREAMZERO_DEFAULT_PROMPT))
+ num_chunks = int(cfg.get("num_chunks", 2))
+ observations = build_dreamzero_demo_observations(
+ load_dreamzero_camera_frames(Path(video_dir)),
+ prompt=prompt,
+ session_id=session_id,
+ num_chunks=num_chunks,
+ )
+ operations = [{"endpoint": "infer", "payload": obs} for obs in observations]
+ operations.append({"endpoint": "reset", "payload": {}})
+ operations.append({"endpoint": "infer", "payload": observations[0]})
+ elif cfg.get("run_default_policy_session"):
+ session_id = str(cfg.get("session_id", "gr00t-smoke"))
+ operations = [
+ {"endpoint": "infer", "payload": build_openpi_droid_observation(session_id=session_id)},
+ {"endpoint": "reset", "payload": {}},
+ ]
+ else:
+ operations = []
+ timeout = float(cfg.get("timeout", 120.0))
+ ping_interval = float(cfg.get("ping_interval", OpenPIWebSocketSession.DEFAULT_PING_INTERVAL_SECS))
+ ping_timeout = float(cfg.get("ping_timeout", OpenPIWebSocketSession.DEFAULT_PING_TIMEOUT_SECS))
+
+ uri = self._build_ws_url("/v1/realtime/robot/openpi")
+ session = OpenPIWebSocketSession(
+ uri,
+ ping_interval=ping_interval,
+ ping_timeout=ping_timeout,
+ open_timeout=timeout,
+ close_timeout=timeout,
+ )
+ try:
+ operation_responses: list[Any] = []
+ actions: dict[str, np.ndarray] | None = None
+ action_tensors: list[np.ndarray] = []
+ for operation in operations:
+ endpoint = str(operation["endpoint"])
+ payload = dict(operation.get("payload") or {})
+ response = session._send_operation(endpoint, payload)
+ operation_responses.append(response)
+ if endpoint == "infer":
+ if isinstance(response, dict):
+ actions = {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
+ else:
+ action_tensors.append(np.asarray(response, dtype=np.float32))
+ finally:
+ session.close()
+
+ return [
+ OpenPIWebSocketResponse(
+ server_metadata=session.get_server_metadata(),
+ operation_responses=operation_responses,
+ actions=actions,
+ action_tensors=action_tensors or None,
+ )
+ ]
+
def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
"""Chat completions via the OpenAI Python SDK (not raw HTTP)."""
responses: list[OmniResponse] = []
@@ -2946,6 +3216,14 @@ __all__ = [
"DiffusionResponse",
"HttpResponse",
"WebSocketJsonResponse",
+ "OpenPIWebSocketResponse",
+ "build_openpi_droid_observation",
+ "build_dreamzero_demo_observations",
+ "DREAMZERO_ACTION_DIM",
+ "DREAMZERO_ACTION_HORIZON",
+ "DREAMZERO_CAMERA_FILES",
+ "load_dreamzero_camera_frames",
+ "OpenPIWebSocketSession",
"OmniResponse",
"OmniRunner",
"OmniRunnerHandler",
diff --git a/tests/helpers/test_assertions.py b/tests/helpers/tests/test_assertions.py
similarity index 92%
rename from tests/helpers/test_assertions.py
rename to tests/helpers/tests/test_assertions.py
index 0bfa7f593..8293349fa 100644
--- a/tests/helpers/test_assertions.py
+++ b/tests/helpers/tests/test_assertions.py
@@ -5,6 +5,8 @@ import pytest
from tests.helpers.assertions import _assert_transcript_matches
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def test_short_transcript_repeat_passes_containment_fallback():
_assert_transcript_matches(
diff --git a/tests/test_pydub_removal.py b/tests/helpers/tests/test_media_pydub_removal.py
similarity index 85%
rename from tests/test_pydub_removal.py
rename to tests/helpers/tests/test_media_pydub_removal.py
index ae65bdab2..60e5cab65 100644
--- a/tests/test_pydub_removal.py
+++ b/tests/helpers/tests/test_media_pydub_removal.py
@@ -7,13 +7,19 @@ identical results to the old pydub-based code for all three paths:
3. TTS WER PCM export: resample to 24000Hz mono s16
"""
+import base64
import io
import wave
import av
import numpy as np
+import pytest
import soundfile as sf
+from tests.helpers.media import _merge_base64_audio_to_segment
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=1.0):
"""Generate a WAV file with a sine wave, return (wav_bytes, pcm_bytes, params)."""
@@ -69,8 +75,6 @@ def test_wav_path_duration_and_frames():
f"sr={sr} ch={ch} dur={dur}: frames {audio_frames} != {expected_frames}"
)
- print("[PASS] test_wav_path_duration_and_frames")
-
def test_non_wav_path_duration_av():
"""Path 2: non-WAV audio → duration & frames via av."""
@@ -86,8 +90,6 @@ def test_non_wav_path_duration_av():
# OGG encoding may add padding, allow small tolerance
assert abs(audio_duration_sec - dur) < 0.05, f"dur={dur}: got {audio_duration_sec}"
- print("[PASS] test_non_wav_path_duration_av")
-
def test_tts_pcm_export_wav_path():
"""Path 3a: WAV PCM → resample to 24000Hz mono s16."""
@@ -114,8 +116,6 @@ def test_tts_pcm_export_wav_path():
assert abs(actual_samples - expected_samples) < 10, f"expected ~{expected_samples} samples, got {actual_samples}"
assert len(pcm_array.shape) == 1, "should be mono (1D array)"
- print("[PASS] test_tts_pcm_export_wav_path")
-
def test_tts_pcm_export_non_wav_path():
"""Path 3b: non-WAV → resample to 24000Hz mono s16 via av.AudioResampler."""
@@ -135,22 +135,13 @@ def test_tts_pcm_export_non_wav_path():
actual_samples = len(result_bytes) // 2
assert abs(actual_samples - expected_samples) < 500, f"expected ~{expected_samples} samples, got {actual_samples}"
- print("[PASS] test_tts_pcm_export_non_wav_path")
-
def test_media_merge_and_export():
"""media.py: _merge_base64_audio_to_segment decode + concat + export WAV."""
- import base64 as b64mod
- import sys
-
- sys.path.insert(0, ".")
- from tests.helpers.media import _merge_base64_audio_to_segment
-
- # Create two WAV chunks as base64
chunks_b64 = []
for dur in [0.5, 0.3]:
wav_bytes, _, _ = _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=dur)
- chunks_b64.append(b64mod.b64encode(wav_bytes).decode())
+ chunks_b64.append(base64.b64encode(wav_bytes).decode())
merged = _merge_base64_audio_to_segment(chunks_b64)
@@ -164,19 +155,3 @@ def test_media_merge_and_export():
out_data, out_sr = sf.read(io.BytesIO(wav_out))
actual_dur = len(out_data) / out_sr
assert abs(actual_dur - 0.8) < 0.05, f"expected ~0.8s merged duration, got {actual_dur}"
-
- print(f"[PASS] test_media_merge_and_export (duration={actual_dur:.3f}s, wav_size={len(wav_out)} bytes)")
-
-
-if __name__ == "__main__":
- test_wav_path_duration_and_frames()
- test_non_wav_path_duration_av()
- import importlib.util
-
- if importlib.util.find_spec("torch") and importlib.util.find_spec("torchaudio"):
- test_tts_pcm_export_wav_path()
- else:
- print("[SKIP] test_tts_pcm_export_wav_path (torch/torchaudio not available)")
- test_tts_pcm_export_non_wav_path()
- test_media_merge_and_export()
- print("\nAll tests passed.")
diff --git a/tests/helpers/test_runtime_stage_cli.py b/tests/helpers/tests/test_runtime_stage_cli.py
similarity index 95%
rename from tests/helpers/test_runtime_stage_cli.py
rename to tests/helpers/tests/test_runtime_stage_cli.py
index d9b50c85f..b127eac13 100644
--- a/tests/helpers/test_runtime_stage_cli.py
+++ b/tests/helpers/tests/test_runtime_stage_cli.py
@@ -4,8 +4,12 @@ from __future__ import annotations
from pathlib import Path
+import pytest
+
from tests.helpers.runtime import OmniServerStageCli
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def _write_stage_config(tmp_path: Path) -> str:
path = tmp_path / "stages.yaml"
diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
index bf2261cb9..2a4cb2f1e 100644
--- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
+++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
@@ -10,6 +10,8 @@ import torch.nn as nn
from tests.helpers.mark import hardware_test
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
class TestPreLookaheadLayer:
"""Tests for PreLookaheadLayer."""
diff --git a/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py b/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py
index d77b28cd5..d2138541f 100644
--- a/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py
+++ b/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py
@@ -1,16 +1,34 @@
from __future__ import annotations
+import pytest
from transformers import PreTrainedModel
from vllm_omni.model_executor.models.dynin_omni.dynin_omni_token2audio import (
_ensure_transformers_tied_weights_compat,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def test_transformers_tied_weights_compat(monkeypatch):
+ """Compat shim must not permanently mutate ``PreTrainedModel``.
+
+ ``monkeypatch.delattr(..., raising=False)`` is a no-op undo when the attr
+ was absent. ``_ensure_transformers_tied_weights_compat`` then assigns a
+ read-only ``property`` on the class that later PreTrainedModel subclasses
+ (e.g. qwen3 TTS mask-cache) can trip over with
+ ``AttributeError: property ... has no setter``.
+ """
monkeypatch.delattr(PreTrainedModel, "all_tied_weights_keys", raising=False)
_ensure_transformers_tied_weights_compat()
+ installed = PreTrainedModel.__dict__.get("all_tied_weights_keys")
+ assert installed is not None
+
+ # Re-own the installed descriptor through monkeypatch so teardown always
+ # removes/restores it, even when the attr did not exist before this test.
+ delattr(PreTrainedModel, "all_tied_weights_keys")
+ monkeypatch.setattr(PreTrainedModel, "all_tied_weights_keys", installed, raising=False)
model = PreTrainedModel.__new__(PreTrainedModel)
model._tied_weights_keys = ["decoder.weight"]
diff --git a/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py b/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py
index 41a29ca63..9586ad053 100644
--- a/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py
+++ b/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
import torch
import torch.nn as nn
@@ -13,6 +14,8 @@ from vllm_omni.model_executor.models.fish_speech.fish_speech_slow_ar import (
FishSpeechSlowARForConditionalGeneration,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def test_fast_ar_reuses_dense_position_id_buffer(monkeypatch):
fast_config = FishSpeechFastARConfig(
diff --git a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py
similarity index 97%
rename from tests/unit/higgs_audio_v3/test_higgs_audio_v3.py
rename to tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py
index 1d4bdb48d..6d81f300a 100755
--- a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py
+++ b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py
@@ -8,6 +8,7 @@ without requiring the actual checkpoint or GPU.
import asyncio
import json
+import os
import time
from collections import OrderedDict, defaultdict
from types import SimpleNamespace
@@ -15,6 +16,11 @@ from types import SimpleNamespace
import pytest
import torch
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
# ---- AC-1: Configuration ----
@@ -240,6 +246,37 @@ class TestDelayPatternBehavior:
class TestSamplerMethods:
"""Test the actual sampler and batched delay masking methods."""
+ @pytest.fixture(autouse=True)
+ def _stub_apply_top_k_top_p_on_cpu(self, monkeypatch):
+ # File is marked cpu; vLLM's apply_top_k_top_p uses Triton and rejects
+ # CPU tensors. Use a pure-PyTorch fallback that still honors top-k/top-p.
+ def _cpu_apply_top_k_top_p(logits, top_ks, top_ps):
+ out = logits.clone()
+ vocab = out.shape[-1]
+ top_ks = top_ks.reshape(-1)
+ top_ps = top_ps.reshape(-1)
+ for row in range(out.shape[0]):
+ k = int(top_ks[row].item())
+ p = float(top_ps[row].item())
+ if 0 < k < vocab:
+ threshold = torch.topk(out[row], k).values[-1]
+ out[row].masked_fill_(out[row] < threshold, float("-inf"))
+ if 0.0 < p < 1.0:
+ sorted_logits, sorted_idx = torch.sort(out[row], descending=True)
+ probs = torch.softmax(sorted_logits, dim=-1)
+ cumprobs = torch.cumsum(probs, dim=-1)
+ # Keep the first token that exceeds p; drop the rest.
+ to_remove = cumprobs > p
+ to_remove[1:] = to_remove[:-1].clone()
+ to_remove[0] = False
+ out[row, sorted_idx[to_remove]] = float("-inf")
+ return out
+
+ monkeypatch.setattr(
+ "vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker.apply_top_k_top_p",
+ _cpu_apply_top_k_top_p,
+ )
+
def _make_minimal_talker(self):
"""Create a minimal talker-like object with sampler/masking methods."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
@@ -1290,17 +1327,12 @@ class TestRegistry:
assert "higgs_multimodal_qwen3" in OMNI_PIPELINES
def test_deploy_yaml_exists(self):
- import os
-
- # __file__ = tests/unit/higgs_audio_v3/test_*.py → 4 dirnames to repo root
- repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
- deploy_dir = os.path.join(repo_root, "vllm_omni", "deploy")
for name in (
"higgs_multimodal_qwen3.yaml",
"higgs_multimodal_qwen3_high_throughput.yaml",
"higgs_multimodal_qwen3_low_latency.yaml",
):
- yaml_path = os.path.join(deploy_dir, name)
+ yaml_path = get_deploy_config_path(name)
assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"
diff --git a/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py b/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py
similarity index 98%
rename from tests/unit/mammoth_moda2/test_mammoth_moda2_config.py
rename to tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py
index 6e35a94f6..05d5dbdde 100644
--- a/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py
+++ b/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py
@@ -23,7 +23,6 @@ from vllm_omni.transformers_utils.configs.mammoth_moda2 import (
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-@pytest.mark.cpu
class TestMammothmoda2Config:
"""Tests for Mammothmoda2Config class."""
@@ -75,7 +74,6 @@ class TestMammothmoda2Config:
assert config.llm_config is not None
-@pytest.mark.cpu
class TestMammothmoda2Qwen2_5_VLConfig:
"""Tests for Mammothmoda2Qwen2_5_VLConfig class."""
@@ -95,7 +93,6 @@ class TestMammothmoda2Qwen2_5_VLConfig:
assert hasattr(text_config, "video_token_id")
-@pytest.mark.cpu
class TestMammothmoda2Qwen3VLConfig:
def test_dev_config_uses_qwen3_vl_subconfigs(self):
config = Mammothmoda2Config(
diff --git a/tests/model_executor/models/moss_tts/test_codec_cudagraph.py b/tests/model_executor/models/moss_tts/test_codec_cudagraph.py
index 49d6fa411..91a853d57 100644
--- a/tests/model_executor/models/moss_tts/test_codec_cudagraph.py
+++ b/tests/model_executor/models/moss_tts/test_codec_cudagraph.py
@@ -1,14 +1,15 @@
"""Tests for MossTTSCUDAGraphCodecWrapper numerical equivalence.
Verifies that CUDA Graph-accelerated decoding produces results equivalent
-to eager mode, with special attention to the two-argument _decode interface
-(codes [NQ, 1, T] + lengths [1]) and the NQ-first input convention.
+to eager mode, with special attention to the two-argument _decode
+interface (codes [NQ, 1, T] + lengths [1]) and the NQ-first input convention.
"""
import pytest
import torch
import torch.nn as nn
+from tests.helpers.mark import hardware_marks
from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import (
MossAudioTokenizerDecoderOutput,
)
@@ -16,8 +17,10 @@ from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
MossTTSCUDAGraphCodecWrapper,
)
-pytestmark = [pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")]
-
+pytestmark = [
+ pytest.mark.core_model,
+ *hardware_marks(res={"cuda": "L4"}, num_cards=1),
+]
DEVICE = torch.device("cuda:0")
NUM_QUANTIZERS = 8
DOWNSAMPLE_RATE = 4 # synthetic; real checkpoint uses 1920
@@ -31,8 +34,9 @@ DOWNSAMPLE_RATE = 4 # synthetic; real checkpoint uses 1920
class SyntheticCodecModel(nn.Module):
"""Minimal stand-in for MossAudioTokenizerModel.
- Exposes the same two-argument _decode(codes, lengths) interface and
- returns a MossAudioTokenizerDecoderOutput, mirroring the real model.
+ Exposes the same two-argument ``_decode(codes, lengths)`` interface
+ used by ``MossTTSCUDAGraphCodecWrapper`` and returns a
+ MossAudioTokenizerDecoderOutput.
Input:
codes: [NQ, B, T] long — RVQ codes
diff --git a/tests/model_executor/models/moss_tts/test_per_request_codes.py b/tests/model_executor/models/moss_tts/test_per_request_codes.py
index b35346c2f..91804ee15 100644
--- a/tests/model_executor/models/moss_tts/test_per_request_codes.py
+++ b/tests/model_executor/models/moss_tts/test_per_request_codes.py
@@ -23,6 +23,9 @@ import functools
import pytest
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
torch = pytest.importorskip("torch")
N_VQ = 4
diff --git a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py
index 0f77455d7..f3e8d0166 100644
--- a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py
+++ b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py
@@ -20,8 +20,9 @@ import pytest
import torch
import torch.nn as nn
-pytestmark = [pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")]
+from tests.helpers.mark import hardware_test
+pytestmark = [pytest.mark.core_model]
DEVICE = torch.device("cuda:0")
NUM_QUANTIZERS = 8
TOTAL_UPSAMPLE = 4
@@ -112,6 +113,7 @@ def _random_codes(seq_len, batch_size=1, device=DEVICE):
# ──────────────────────────────────────────────────────────────────
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("seq_len", [25, 50, 100])
def test_exact_size_numerical_equivalence(decoder, wrapper, seq_len):
"""When input exactly matches a capture size, output must be bit-identical."""
@@ -127,6 +129,7 @@ def test_exact_size_numerical_equivalence(decoder, wrapper, seq_len):
# ──────────────────────────────────────────────────────────────────
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
def test_padded_output_shape_and_length(decoder, wrapper, seq_len):
"""Padded decode must return output trimmed to actual input length."""
@@ -139,6 +142,7 @@ def test_padded_output_shape_and_length(decoder, wrapper, seq_len):
assert graph_out.shape[-1] == expected_len
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
def test_padded_interior_positions_close(decoder, wrapper, seq_len):
"""Interior positions (away from padding boundary) should be very close.
@@ -160,6 +164,7 @@ def test_padded_interior_positions_close(decoder, wrapper, seq_len):
torch.testing.assert_close(interior_graph, interior_eager, atol=1e-5, rtol=1e-5)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
def test_padded_output_bounded(decoder, wrapper, seq_len):
"""Padded output values must remain in [-1, 1] and max diff should be bounded."""
@@ -175,6 +180,7 @@ def test_padded_output_bounded(decoder, wrapper, seq_len):
assert max_diff < 1.0, f"Max diff {max_diff} exceeds bound"
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
def test_padded_short_output_decode_length_matches_eager(seq_len):
"""Streaming ``decode()`` of a shorter-than-nominal decoder must match eager length.
@@ -210,6 +216,7 @@ def test_padded_short_output_decode_length_matches_eager(seq_len):
torch.testing.assert_close(graph_out[..., :-boundary], eager_out[..., :-boundary], atol=1e-5, rtol=1e-5)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_compiled_padded_short_output_length_matches_eager(monkeypatch):
"""Compiled replay path must also trim short-output decoders to eager length.
@@ -251,6 +258,7 @@ def test_compiled_padded_short_output_length_matches_eager(monkeypatch):
# ──────────────────────────────────────────────────────────────────
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("seq_len", [101, 150, 200])
def test_fallback_eager_exact_match(decoder, wrapper, seq_len):
"""Input larger than all capture sizes falls back to eager -> bit-identical."""
@@ -266,6 +274,7 @@ def test_fallback_eager_exact_match(decoder, wrapper, seq_len):
# ──────────────────────────────────────────────────────────────────
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("total_len", [60, 100, 150, 250])
def test_chunked_decode_shape_match(decoder, wrapper, total_len):
"""Chunked decode output shape must match between eager and graph modes."""
@@ -279,6 +288,7 @@ def test_chunked_decode_shape_match(decoder, wrapper, total_len):
assert eager_out.shape == graph_out.shape
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("total_len", [50, 100])
def test_chunked_decode_exact_size_equivalence(decoder, wrapper, total_len):
"""Chunked decode with chunks matching capture sizes should be bit-identical."""
@@ -293,6 +303,7 @@ def test_chunked_decode_exact_size_equivalence(decoder, wrapper, total_len):
torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_chunked_decode_output_survives_later_replay(wrapper):
"""Chunked output must not alias graph static buffers overwritten by later replays."""
codes = _random_codes(100)
@@ -307,6 +318,7 @@ def test_chunked_decode_output_survives_later_replay(wrapper):
torch.testing.assert_close(graph_out, expected, atol=0, rtol=0)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_chunked_decode_preserves_short_chunk_concat_semantics():
"""Chunked decode must allow shorter-than-nominal chunk outputs.
@@ -336,6 +348,7 @@ def test_chunked_decode_preserves_short_chunk_concat_semantics():
torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_batched_chunked_decode_variable_lengths_matches_per_request_eager(decoder, wrapper):
"""Variable-length chunk batching should match independent chunked decodes."""
long_codes = _random_codes(100)
@@ -380,6 +393,7 @@ def _eager_chunked(decoder, codes, chunk_size, left_context_size):
# ──────────────────────────────────────────────────────────────────
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_single_frame(decoder, wrapper):
"""Single-frame input (seq_len=1) should work with padding."""
codes = _random_codes(1)
@@ -390,6 +404,7 @@ def test_single_frame(decoder, wrapper):
assert graph_out.shape[-1] == TOTAL_UPSAMPLE
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_disabled_wrapper_matches_eager(decoder, wrapper):
"""Disabled wrapper should produce bit-identical output to eager."""
codes = _random_codes(30)
@@ -401,6 +416,7 @@ def test_disabled_wrapper_matches_eager(decoder, wrapper):
torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_batch_size_gt1_uses_matching_graph(decoder, wrapper):
"""Captured batch size > 1 should replay a matching graph."""
assert (2, 25) in wrapper.graphs
@@ -411,6 +427,7 @@ def test_batch_size_gt1_uses_matching_graph(decoder, wrapper):
torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_uncaptured_batch_size_falls_back(decoder, wrapper):
"""Uncaptured batch sizes should fall back to eager."""
assert (3, 25) not in wrapper.graphs
@@ -421,6 +438,7 @@ def test_uncaptured_batch_size_falls_back(decoder, wrapper):
torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_extra_capture_shape_uses_sparse_graph(decoder):
"""Extra capture shapes should not expand to a full batch x size product."""
sparse_wrapper = CUDAGraphDecoderWrapper(
@@ -444,6 +462,7 @@ def test_extra_capture_shape_uses_sparse_graph(decoder):
torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_compile_shape_supports_exact_and_padded_buckets(decoder, monkeypatch):
"""Configured torch.compile shapes should replay exact and padded CUDA Graph buckets."""
@@ -490,6 +509,7 @@ def test_compile_shape_supports_exact_and_padded_buckets(decoder, monkeypatch):
assert compile_kwargs["dynamic"] is False
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_deterministic_across_calls(decoder, wrapper):
"""Same input should produce identical CUDA graph output across calls."""
codes = _random_codes(30)
@@ -504,6 +524,7 @@ def test_deterministic_across_calls(decoder, wrapper):
# ──────────────────────────────────────────────────────────────────
+@pytest.mark.cpu
@pytest.mark.parametrize(
"kwargs,expected_in,not_expected",
[
@@ -545,6 +566,7 @@ def test_compute_capture_sizes(kwargs, expected_in, not_expected):
# ──────────────────────────────────────────────────────────────────
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize(
"batch,channels,seq_len",
[(2, 64, 1000), (1, 32, 1), (1, 32, 7), (1, 32, 128), (1, 32, 1024), (1, 32, 4096)],
diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py
index 7358c7893..c76583a4e 100644
--- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py
+++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py
@@ -17,6 +17,8 @@ from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
Qwen3TTSTalkerForConditionalGeneration,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def _make_minimal_talker(
tts_pad_embed: torch.Tensor | None = None,
diff --git a/tests/test_tts_custom_voice_profiles.py b/tests/model_executor/models/qwen3_tts/test_tts_custom_voice_profiles.py
similarity index 100%
rename from tests/test_tts_custom_voice_profiles.py
rename to tests/model_executor/models/qwen3_tts/test_tts_custom_voice_profiles.py
diff --git a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py
index 8154ddd65..0007f021a 100644
--- a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py
+++ b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py
@@ -8,6 +8,11 @@ lightweight local checkout. Runtime audio quality still requires a CUDA test.
from pathlib import Path
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
REPO_ROOT = Path(__file__).resolve().parents[4]
RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py"
RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py"
diff --git a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py
index f5a94376a..3938051d3 100644
--- a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py
+++ b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py
@@ -9,6 +9,9 @@ from types import SimpleNamespace
import pytest
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
torch = pytest.importorskip("torch")
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 5309c7675..d5971efe7 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
@@ -4,6 +4,7 @@
from collections import defaultdict
from types import SimpleNamespace
+import pytest
import torch
from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
@@ -12,6 +13,8 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
text2flow_token_only,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], mm: dict, finished: bool = True):
return SimpleNamespace(
diff --git a/tests/model_executor/stage_input_processors/test_indextts2.py b/tests/model_executor/stage_input_processors/test_indextts2.py
index 11468603e..1380b99ff 100644
--- a/tests/model_executor/stage_input_processors/test_indextts2.py
+++ b/tests/model_executor/stage_input_processors/test_indextts2.py
@@ -3,6 +3,7 @@
from types import SimpleNamespace
+import pytest
import torch
from vllm_omni.model_executor.stage_input_processors.indextts2 import (
@@ -11,6 +12,9 @@ from vllm_omni.model_executor.stage_input_processors.indextts2 import (
talker2s2mel_token_only,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
STOP_MEL_TOKEN = 8193
LATENT_DIM = 16
diff --git a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py
index 0cb61a972..3ecae3d56 100644
--- a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py
+++ b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py
@@ -20,6 +20,8 @@ from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
thinker2talker_full_payload,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def _make_request(
prompt_token_ids,
diff --git a/tests/plugin/__init__.py b/tests/plugin/__init__.py
deleted file mode 100644
index 988131360..000000000
--- a/tests/plugin/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
diff --git a/tests/profile/test_omni_torch_profiler.py b/tests/profile/test_omni_torch_profiler.py
index 3920078af..7bfc8f8cc 100644
--- a/tests/profile/test_omni_torch_profiler.py
+++ b/tests/profile/test_omni_torch_profiler.py
@@ -12,6 +12,8 @@ from openpyxl import load_workbook
import vllm_omni.profiler.omni_torch_profiler as profiler_mod
from vllm_omni.profiler.omni_torch_profiler import OmniTorchProfilerWrapper
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
@pytest.fixture(autouse=True)
def patch_worker_profiler_init(monkeypatch):
diff --git a/tests/quantization/test_quack_fp8_gate.py b/tests/quantization/test_quack_fp8_gate.py
index 4a8848ffa..843f4ff5f 100644
--- a/tests/quantization/test_quack_fp8_gate.py
+++ b/tests/quantization/test_quack_fp8_gate.py
@@ -20,6 +20,9 @@ import torch
from vllm_omni.quantization import quack_fp8
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
_ENV = "VLLM_OMNI_USE_QUACK_FP8"
diff --git a/tests/reasoning/test_step_audio_reasoning_parser.py b/tests/reasoning/test_step_audio_reasoning_parser.py
index 7eb9fa17a..11105265c 100644
--- a/tests/reasoning/test_step_audio_reasoning_parser.py
+++ b/tests/reasoning/test_step_audio_reasoning_parser.py
@@ -21,6 +21,9 @@ from unittest.mock import MagicMock
import pytest
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
# ---------------------------------------------------------------------------
# Bootstrap minimal vLLM stubs so the parser can be imported
# ---------------------------------------------------------------------------
@@ -784,8 +787,8 @@ class TestExtractReasoningStreaming:
"""When buffered text turns out not to be a marker prefix, it is
emitted as reasoning on the next call."""
# "reason<" is buffered because "<" could start a marker.
- # Next delta ">" doesn't form a marker, so the buffer is
- # flushed as reasoning.
+ # Next delta ">" doesn't form a marker, so the buffer ("<") is
+ # prepended to the next delta and flushed as reasoning: "<> more".
parser = multi_token_parser
deltas_and_ids = [
(TS, [100]),
@@ -794,7 +797,7 @@ class TestExtractReasoningStreaming:
("> more", [8]), # ">" doesn't complete any marker
]
r, c = _stream_extract(parser, deltas_and_ids)
- assert r == "reason> more"
+ assert r == "reason<> more"
assert c == ""
def test_multi_token_end_with_reasoning_before_buffer(self, multi_token_parser):
@@ -837,7 +840,7 @@ class TestExtractReasoningStreaming:
("> not a marker", [8]),
]
r, c = _stream_extract(parser, deltas_and_ids)
- assert r == "reason> not a marker"
+ assert r == "reason<> not a marker"
assert c == ""
diff --git a/tests/test_generate_nightly_perf_excel.py b/tests/tools/test_generate_nightly_perf_excel.py
similarity index 96%
rename from tests/test_generate_nightly_perf_excel.py
rename to tests/tools/test_generate_nightly_perf_excel.py
index 6658b89b8..80b27fbf5 100644
--- a/tests/test_generate_nightly_perf_excel.py
+++ b/tests/tools/test_generate_nightly_perf_excel.py
@@ -10,7 +10,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _load_excel_module():
- repo_root = Path(__file__).resolve().parents[1]
+ repo_root = Path(__file__).resolve().parents[2]
module_path = repo_root / "tools" / "nightly" / "generate_nightly_perf_excel.py"
spec = importlib.util.spec_from_file_location("generate_nightly_perf_excel", module_path)
module = importlib.util.module_from_spec(spec)
@@ -27,7 +27,7 @@ def _cell_value_by_header(ws, header_name: str, row_idx: int = 2):
def test_generate_excel_report_with_perf_templates(tmp_path: Path):
module = _load_excel_module()
- repo_root = Path(__file__).resolve().parents[1]
+ repo_root = Path(__file__).resolve().parents[2]
perf_scripts_dir = repo_root / "tests" / "dfx" / "perf" / "scripts"
omni_template_path = perf_scripts_dir / "result_omni_template.json"
diff --git a/tests/test_generate_nightly_perf_html.py b/tests/tools/test_generate_nightly_perf_html.py
similarity index 95%
rename from tests/test_generate_nightly_perf_html.py
rename to tests/tools/test_generate_nightly_perf_html.py
index 7a05165da..4592c2183 100644
--- a/tests/test_generate_nightly_perf_html.py
+++ b/tests/tools/test_generate_nightly_perf_html.py
@@ -9,7 +9,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _load_html_module():
- repo_root = Path(__file__).resolve().parents[1]
+ repo_root = Path(__file__).resolve().parents[2]
module_path = repo_root / "tools" / "nightly" / "generate_nightly_perf_html.py"
spec = importlib.util.spec_from_file_location("generate_nightly_perf_html", module_path)
module = importlib.util.module_from_spec(spec)
@@ -20,7 +20,7 @@ def _load_html_module():
def test_generate_html_report_with_perf_templates(tmp_path: Path):
module = _load_html_module()
- repo_root = Path(__file__).resolve().parents[1]
+ repo_root = Path(__file__).resolve().parents[2]
perf_scripts_dir = repo_root / "tests" / "dfx" / "perf" / "scripts"
omni_template_path = perf_scripts_dir / "result_omni_template.json"
diff --git a/tests/unit/higgs_audio_v3/__init__.py b/tests/unit/higgs_audio_v3/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/tests/utils/test_audio.py b/tests/utils/test_audio.py
index 0e483e646..ff29b2d04 100644
--- a/tests/utils/test_audio.py
+++ b/tests/utils/test_audio.py
@@ -9,6 +9,9 @@ import torch
from vllm_omni.utils.audio import mel_filter_bank, peak_normalize
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
# Parameter combinations used across the codebase.
_PARAM_SETS = [
# Qwen3-TTS talker / speaker encoder (sr=24000)
diff --git a/tests/test_speaker_cache.py b/tests/utils/test_speaker_cache.py
similarity index 100%
rename from tests/test_speaker_cache.py
rename to tests/utils/test_speaker_cache.py
diff --git a/tests/test_speaker_cache_integration.py b/tests/utils/test_speaker_cache_integration.py
similarity index 100%
rename from tests/test_speaker_cache_integration.py
rename to tests/utils/test_speaker_cache_integration.py
diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py
index 00882744c..13c405af7 100644
--- a/tests/utils/test_tracking_parser.py
+++ b/tests/utils/test_tracking_parser.py
@@ -23,6 +23,9 @@ from vllm_omni.utils.tracking_parser import (
build_shadow_kwargs,
)
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
### Fake pipeline/deploy config for integration tests
_TEST_MODEL = "_test_tracking"
diff --git a/tests/test_version.py b/tests/utils/test_version.py
similarity index 98%
rename from tests/test_version.py
rename to tests/utils/test_version.py
index 911a79611..0c666133c 100644
--- a/tests/test_version.py
+++ b/tests/utils/test_version.py
@@ -9,6 +9,8 @@ import pytest
from vllm_omni.version import warn_if_misaligned_vllm_version
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
@mock.patch("vllm_omni.version.__version_tuple__", (0, 19, 0))
@mock.patch("vllm_omni.version.__version__", "0.19.0")
diff --git a/tools/nightly/buildkite_testcase_statistics.py b/tools/nightly/buildkite_testcase_statistics.py
index 7a4f9bff3..740f896f5 100644
--- a/tools/nightly/buildkite_testcase_statistics.py
+++ b/tools/nightly/buildkite_testcase_statistics.py
@@ -903,11 +903,11 @@ def _sample_preview_stats() -> list[tuple]:
"test-ready",
None,
"Custom Pipeline Test",
- ["tests/e2e/offline_inference/custom_pipeline/test_cp.py"],
+ ["tests/e2e/features/custom_pipeline/test_cp.py"],
False,
[],
1,
- ["tests/e2e/offline_inference/custom_pipeline/test_cp.py::test_cp_one"],
+ ["tests/e2e/features/custom_pipeline/test_cp.py::test_cp_one"],
[""],
set(),
),
diff --git a/tests/e2e/offline_inference/stage_configs/step_audio2_ci.yaml b/vllm_omni/deploy/step_audio2_thinker_ci.yaml
similarity index 93%
rename from tests/e2e/offline_inference/stage_configs/step_audio2_ci.yaml
rename to vllm_omni/deploy/step_audio2_thinker_ci.yaml
index f9dd4f9af..7856ebda9 100644
--- a/tests/e2e/offline_inference/stage_configs/step_audio2_ci.yaml
+++ b/vllm_omni/deploy/step_audio2_thinker_ci.yaml
@@ -1,4 +1,4 @@
-# Stage config for Step-Audio2 CI tests
+# Stage config for Step-Audio2 CI tests (offline inference, thinker-only).
# Single Thinker stage (audio/text understanding → text). ASR/text-output
# tests don't need Token2Wav; a two-stage config would deadlock on stage-1 skip.
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 924f41c3f..ed3f5464d 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -395,7 +395,7 @@ class OmniAsyncEngineArgs(AsyncEngineArgs, OmniEngineArgs):
# Fields in ``SHARED_FIELDS`` (e.g. ``model``, ``log_stats``) flow to BOTH
# orchestrator and engine by design.
#
-# Invariants enforced by ``tests/test_arg_utils.py``:
+# Invariants enforced by ``tests/engine/test_arg_utils_shared_fields.py``:
#
# 1. ``OrchestratorArgs`` ∩ ``OmniEngineArgs`` ⊆ ``SHARED_FIELDS``
# 2. Every CLI flag is classifiable into one of the three buckets
|