From 7c17111ff54f93d054cbcb95edd66a7d54e89807 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:53:03 +0800 Subject: [PATCH] [Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951) Signed-off-by: wangyu <410167048@qq.com> --- .buildkite/amd/test-amd-merge.yml | 2 +- .buildkite/amd/test-amd-ready.yml | 2 +- .buildkite/cuda/test-merge.yml | 72 +- .buildkite/cuda/test-nightly.yml | 6 + .buildkite/cuda/test-ready.yml | 119 +- docs/contributing/ci/test_system_overview.md | 76 +- docs/contributing/ci/test_writing_guide.md | 165 +- docs/design/feature/tensor_parallel.md | 2 +- examples/model_display_names.yml | 1 - recipes/JD/JoyAI-VL-Interaction.md | 2 +- recipes/NVIDIA/GR00T-N1.7.md | 9 +- .../hunyuan}/hunyuan_image_ref.png | Bin .../qwen_image_edit_2511_test1.png | Bin .../qwen_image_edit_2511_test2.png | Bin tests/{ => config}/test_config_factory.py | 51 +- tests/conftest.py | 1 + .../test_omni_scheduler_mixin_timeouts.py | 2 + tests/core/test_prefix_cache.py | 3 + .../{ => diffusion}/ar_diffusion/__init__.py | 0 .../ar_diffusion/test_capability_runner.py | 0 .../ar_diffusion/test_engine_routing.py | 2 + .../ar_diffusion/test_kv_cache.py | 3 + .../ar_diffusion/test_kv_orchestration.py | 1 + .../ar_diffusion/test_paged_attention.py | 26 +- .../ar_diffusion/test_pipeline_bridge.py | 3 + .../ar_diffusion/test_slot_mapping.py | 3 + .../attention/test_attention_config.py | 7 + .../diffusion/attention/test_attention_sp.py | 2 + tests/diffusion/attention/test_flash_attn.py | 3 + .../attention/test_piecewise_attn.py | 3 + tests/diffusion/attention/test_sage_attn3.py | 3 + tests/diffusion/attention/test_ulysses_uaa.py | 2 + .../distributed/test_cfg_parallel.py | 1334 ++++++++++------- tests/diffusion/distributed/test_comm.py | 609 ++++---- .../distributed/test_pipeline_parallel.py | 551 ++++--- ...parallelism.py => test_tensor_parallel.py} | 130 +- .../test_vae_decode_parallelism.py | 276 ++-- tests/diffusion/hooks/test_hook_registry.py | 3 + .../layers/test_rotary_emb_equivalence.py | 2 + .../models/dreamzero}/__init__.py | 0 .../models}/dreamzero/test_crossattn_cache.py | 0 .../models}/dreamzero/test_deploy_configs.py | 0 .../models}/dreamzero/test_dreamzero_utils.py | 0 .../dreamzero/test_fused_qk_rms_norm.py | 0 .../models}/dreamzero/test_pipeline_state.py | 0 .../models}/dreamzero/test_qkv_fusion.py | 0 .../test_session_state_accounting.py | 0 .../test_session_state_equivalence.py | 0 .../dreamzero/test_session_state_lifecycle.py | 0 .../models/ltx2/test_ltx2_transformer.py | 2 + .../sensenova_u1/test_sensenova_u1_cfg.py | 2 + .../offloader/test_diffusion_cpu_offload.py | 24 +- .../test_diffusion_layerwise_offload.py | 28 +- .../quantization}/test_gguf_diffusion.py | 5 +- .../quantization/test_quantization_quality.py | 2 + .../test_trajectory_similarity_tool.py | 2 + .../test_diffusion_config_fields.py | 2 +- .../test_diffusion_config_propagation.py | 0 tests/diffusion/test_diffusion_engine.py | 33 +- .../test_diffusion_engine_metrics.py | 5 + tests/docs/test_generate_examples.py | 3 + tests/dreamzero/openpi_client_helper.py | 307 ---- tests/dreamzero/test_openpi_client_helper.py | 46 - .../upstream/openpi_test_client_ar.py | 381 ----- .../upstream/test_openpi_e2e_source_parity.py | 358 ----- tests/e2e/accuracy/conftest.py | 225 --- tests/e2e/accuracy/helpers.py | 347 ++++- tests/e2e/accuracy/test_dreamzero.py | 63 + tests/e2e/accuracy/test_gebench_h100_smoke.py | 97 -- .../accuracy/test_gedit_bench_h100_smoke.py | 115 -- tests/e2e/accuracy/test_hunyuan_image3.py | 10 +- .../test_hunyuan_image3_pixel_accuracy.py | 3 +- tests/e2e/accuracy/test_qwen_image_edit.py | 6 +- .../upstream_socket_server_no_compile.py | 2 +- tests/{ => e2e/features}/comfyui/conftest.py | 2 +- .../comfyui/test_comfyui_integration.py | 4 +- .../test_async_omni_collective_rpc.py | 10 +- .../test_async_omni_qwen_image_generate.py | 8 +- .../custom_pipeline/test_worker_extension.py | 22 +- .../engine/test_duplex_control_client.py | 0 .../engine/test_duplex_control_plane.py | 0 .../engine/test_duplex_deploy_config.py | 0 .../engine/test_duplex_intermediate.py | 0 .../fullduplex/engine/test_duplex_lease.py | 0 .../fullduplex/engine/test_duplex_runtime.py | 0 .../fullduplex/minicpmo45/__init__.py | 0 .../minicpmo45/test_commit_policy.py | 0 .../fullduplex/minicpmo45/test_input.py | 0 .../fullduplex/openai/test_duplex_audio.py | 0 .../openai/test_runtime_adapter_boundary.py | 0 .../fullduplex/openai/test_websocket_actor.py | 0 .../features}/fullduplex/test_client.py | 0 .../features}/fullduplex/test_joyvl_brain.py | 0 .../features}/fullduplex/test_joyvl_policy.py | 0 .../features}/fullduplex/test_runtime.py | 0 .../helpers/custom_pipeline.py} | 288 +++- .../helpers/verl_omni_server.py} | 286 +--- .../features/rlhf_test/test_verl_omni_e2e.py | 217 +++ .../custom_pipeline/__init__.py | 14 - .../flow_match_sde_scheduler.py | 220 --- .../custom_pipeline/worker_extension.py | 77 - ...inger.py => test_soulxsinger_expansion.py} | 2 +- .../test_step_audio2_expansion.py | 4 +- .../test_dreamzero_expansion.py | 83 +- ...enpi.py => test_gr00t_openpi_expansion.py} | 78 +- .../test_omnivoice_expansion.py | 2 +- .../test_arg_utils_shared_fields.py} | 0 tests/{ => engine}/test_data_entry_keys.py | 0 tests/{ => engine}/test_omni_request.py | 0 .../test_omni_request_output.py} | 0 .../test_orchestrator_error_handling.py | 10 +- .../test_speaker_metadata_persistence.py | 0 .../openai_api/test_speech_usage.py | 2 + .../openai_api/test_stage_params.py | 3 + .../openai_api/test_tts_adapter.py | 3 + .../openai_api/test_video_frame_filter.py | 2 + .../openai_api/test_video_stream_handler.py | 3 + .../openai_api/test_video_stream_session.py | 3 + tests/entrypoints/test_pd_disaggregation.py | 3 + tests/entrypoints/test_stage_utils.py | 2 + tests/gr00t/__init__.py | 2 - tests/gr00t/openpi_client_helper.py | 183 --- tests/helpers/__init__.py | 9 +- tests/helpers/env.py | 27 + tests/helpers/runtime.py | 278 ++++ tests/helpers/{ => tests}/test_assertions.py | 2 + .../tests/test_media_pydub_removal.py} | 39 +- .../{ => tests}/test_runtime_stage_cli.py | 4 + .../cosyvoice3/test_cosyvoice3_components.py | 2 + .../dynin_omni/test_dynin_omni_token2audio.py | 18 + ...est_fish_speech_fast_ar_runtime_buffers.py | 3 + .../higgs_audio_v3/test_higgs_audio_v3.py | 44 +- .../test_mammoth_moda2_config.py | 3 - .../models/moss_tts/test_codec_cudagraph.py | 16 +- .../models/moss_tts/test_per_request_codes.py | 3 + .../qwen3_tts/test_cuda_graph_decoder.py | 24 +- .../test_qwen3_tts_talker_preprocess.py | 2 + .../test_tts_custom_voice_profiles.py | 0 ...st_runner_assisted_unified_graph_static.py | 5 + .../voxcpm2/test_talker_state_eviction.py | 3 + .../test_cosyvoice3_stage_input_processors.py | 3 + .../stage_input_processors/test_indextts2.py | 4 + .../test_qwen2_5_omni_thinker2talker.py | 2 + tests/plugin/__init__.py | 1 - tests/profile/test_omni_torch_profiler.py | 2 + tests/quantization/test_quack_fp8_gate.py | 3 + .../test_step_audio_reasoning_parser.py | 11 +- .../test_generate_nightly_perf_excel.py | 4 +- .../test_generate_nightly_perf_html.py | 4 +- tests/unit/higgs_audio_v3/__init__.py | 0 tests/utils/test_audio.py | 3 + tests/{ => utils}/test_speaker_cache.py | 0 .../test_speaker_cache_integration.py | 0 tests/utils/test_tracking_parser.py | 3 + tests/{ => utils}/test_version.py | 2 + .../nightly/buildkite_testcase_statistics.py | 4 +- .../deploy/step_audio2_thinker_ci.yaml | 2 +- vllm_omni/engine/arg_utils.py | 2 +- 158 files changed, 3479 insertions(+), 4130 deletions(-) rename tests/{e2e/accuracy/assets => assets/hunyuan}/hunyuan_image_ref.png (100%) rename tests/{e2e/accuracy/assets => assets/qwen_image_edit}/qwen_image_edit_2511_test1.png (100%) rename tests/{e2e/accuracy/assets => assets/qwen_image_edit}/qwen_image_edit_2511_test2.png (100%) rename tests/{ => config}/test_config_factory.py (97%) rename tests/{ => diffusion}/ar_diffusion/__init__.py (100%) rename tests/{ => diffusion}/ar_diffusion/test_capability_runner.py (100%) rename tests/{ => diffusion}/ar_diffusion/test_engine_routing.py (98%) rename tests/{ => diffusion}/ar_diffusion/test_kv_cache.py (99%) rename tests/{ => diffusion}/ar_diffusion/test_kv_orchestration.py (98%) rename tests/{ => diffusion}/ar_diffusion/test_paged_attention.py (94%) rename tests/{ => diffusion}/ar_diffusion/test_pipeline_bridge.py (99%) rename tests/{ => diffusion}/ar_diffusion/test_slot_mapping.py (97%) rename tests/diffusion/distributed/{test_zimage_parallelism.py => test_tensor_parallel.py} (53%) rename tests/{e2e/offline_inference/rlhf_test => diffusion/models/dreamzero}/__init__.py (100%) rename tests/{ => diffusion/models}/dreamzero/test_crossattn_cache.py (100%) rename tests/{ => diffusion/models}/dreamzero/test_deploy_configs.py (100%) rename tests/{ => diffusion/models}/dreamzero/test_dreamzero_utils.py (100%) rename tests/{ => diffusion/models}/dreamzero/test_fused_qk_rms_norm.py (100%) rename tests/{ => diffusion/models}/dreamzero/test_pipeline_state.py (100%) rename tests/{ => diffusion/models}/dreamzero/test_qkv_fusion.py (100%) rename tests/{ => diffusion/models}/dreamzero/test_session_state_accounting.py (100%) rename tests/{ => diffusion/models}/dreamzero/test_session_state_equivalence.py (100%) rename tests/{ => diffusion/models}/dreamzero/test_session_state_lifecycle.py (100%) rename tests/{plugin => diffusion/quantization}/test_gguf_diffusion.py (97%) rename tests/{ => diffusion}/test_diffusion_config_fields.py (97%) rename tests/{ => diffusion}/test_diffusion_config_propagation.py (100%) delete mode 100644 tests/dreamzero/openpi_client_helper.py delete mode 100644 tests/dreamzero/test_openpi_client_helper.py delete mode 100644 tests/dreamzero/upstream/openpi_test_client_ar.py delete mode 100644 tests/dreamzero/upstream/test_openpi_e2e_source_parity.py create mode 100644 tests/e2e/accuracy/test_dreamzero.py delete mode 100644 tests/e2e/accuracy/test_gebench_h100_smoke.py delete mode 100644 tests/e2e/accuracy/test_gedit_bench_h100_smoke.py rename tests/{dreamzero/upstream => e2e/accuracy/upstream_mock}/upstream_socket_server_no_compile.py (98%) rename tests/{ => e2e/features}/comfyui/conftest.py (97%) rename tests/{ => e2e/features}/comfyui/test_comfyui_integration.py (99%) rename tests/e2e/{offline_inference => features}/custom_pipeline/test_async_omni_collective_rpc.py (96%) rename tests/e2e/{offline_inference => features}/custom_pipeline/test_async_omni_qwen_image_generate.py (96%) rename tests/e2e/{offline_inference => features}/custom_pipeline/test_worker_extension.py (59%) rename tests/{ => e2e/features}/fullduplex/engine/test_duplex_control_client.py (100%) rename tests/{ => e2e/features}/fullduplex/engine/test_duplex_control_plane.py (100%) rename tests/{ => e2e/features}/fullduplex/engine/test_duplex_deploy_config.py (100%) rename tests/{ => e2e/features}/fullduplex/engine/test_duplex_intermediate.py (100%) rename tests/{ => e2e/features}/fullduplex/engine/test_duplex_lease.py (100%) rename tests/{ => e2e/features}/fullduplex/engine/test_duplex_runtime.py (100%) rename tests/{ => e2e/features}/fullduplex/minicpmo45/__init__.py (100%) rename tests/{ => e2e/features}/fullduplex/minicpmo45/test_commit_policy.py (100%) rename tests/{ => e2e/features}/fullduplex/minicpmo45/test_input.py (100%) rename tests/{ => e2e/features}/fullduplex/openai/test_duplex_audio.py (100%) rename tests/{ => e2e/features}/fullduplex/openai/test_runtime_adapter_boundary.py (100%) rename tests/{ => e2e/features}/fullduplex/openai/test_websocket_actor.py (100%) rename tests/{ => e2e/features}/fullduplex/test_client.py (100%) rename tests/{ => e2e/features}/fullduplex/test_joyvl_brain.py (100%) rename tests/{ => e2e/features}/fullduplex/test_joyvl_policy.py (100%) rename tests/{ => e2e/features}/fullduplex/test_runtime.py (100%) rename tests/e2e/{offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py => features/helpers/custom_pipeline.py} (59%) rename tests/e2e/{offline_inference/rlhf_test/test_verl_omni_e2e.py => features/helpers/verl_omni_server.py} (52%) create mode 100644 tests/e2e/features/rlhf_test/test_verl_omni_e2e.py delete mode 100644 tests/e2e/offline_inference/custom_pipeline/__init__.py delete mode 100644 tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py delete mode 100644 tests/e2e/offline_inference/custom_pipeline/worker_extension.py rename tests/e2e/offline_inference/{test_soulxsinger.py => test_soulxsinger_expansion.py} (99%) rename tests/e2e/online_serving/{test_gr00t_openpi.py => test_gr00t_openpi_expansion.py} (66%) rename tests/{test_arg_utils.py => engine/test_arg_utils_shared_fields.py} (100%) rename tests/{ => engine}/test_data_entry_keys.py (100%) rename tests/{ => engine}/test_omni_request.py (100%) rename tests/{test_outputs.py => engine/test_omni_request_output.py} (100%) rename tests/{ => entrypoints/openai_api}/test_speaker_metadata_persistence.py (100%) delete mode 100644 tests/gr00t/__init__.py delete mode 100644 tests/gr00t/openpi_client_helper.py rename tests/helpers/{ => tests}/test_assertions.py (92%) rename tests/{test_pydub_removal.py => helpers/tests/test_media_pydub_removal.py} (85%) rename tests/helpers/{ => tests}/test_runtime_stage_cli.py (95%) rename tests/{unit => model_executor/models}/higgs_audio_v3/test_higgs_audio_v3.py (97%) rename tests/{unit => model_executor/models}/mammoth_moda2/test_mammoth_moda2_config.py (98%) rename tests/{ => model_executor/models/qwen3_tts}/test_tts_custom_voice_profiles.py (100%) delete mode 100644 tests/plugin/__init__.py rename tests/{ => tools}/test_generate_nightly_perf_excel.py (96%) rename tests/{ => tools}/test_generate_nightly_perf_html.py (95%) delete mode 100644 tests/unit/higgs_audio_v3/__init__.py rename tests/{ => utils}/test_speaker_cache.py (100%) rename tests/{ => utils}/test_speaker_cache_integration.py (100%) rename tests/{ => utils}/test_version.py (98%) rename tests/e2e/offline_inference/stage_configs/step_audio2_ci.yaml => vllm_omni/deploy/step_audio2_thinker_ci.yaml (93%) 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